From 2875dfc4345a57b81291e8308b12f05609013c74 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Sat, 2 May 2026 23:17:38 -0700 Subject: [PATCH 1/9] [https://nvbugs/6104831][fix] Disagg request cancellation fix Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 33 +- .../tensorrt_llm/executor/transferAgent.h | 9 + .../batch_manager/baseTransBuffer.cpp | 69 ++- .../batch_manager/baseTransBuffer.h | 148 ++++- .../batch_manager/cacheFormatter.cpp | 25 +- .../batch_manager/cacheTransceiver.cpp | 391 +++++++++--- .../batch_manager/dataTransceiver.cpp | 581 +++++++++++++++--- .../batch_manager/dataTransceiver.h | 14 +- .../batch_manager/mlaCacheFormatter.cpp | 16 +- .../trtGptModelInflightBatching.cpp | 6 +- .../agent_utils/connection.cpp | 89 ++- .../agent_utils/connection.h | 14 +- .../nixl_utils/agentBindings.cpp | 6 +- .../nixl_utils/transferAgent.cpp | 35 ++ .../nixl_utils/transferAgent.h | 3 + .../batch_manager/cacheTransceiver.cpp | 8 +- .../_torch/disaggregation/base/agent.py | 3 + .../_torch/disaggregation/nixl/_agent_cpp.py | 4 + .../_torch/disaggregation/nixl/_agent_py.py | 8 + tensorrt_llm/_torch/pyexecutor/py_executor.py | 91 ++- 20 files changed, 1309 insertions(+), 244 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 8f8330603893..c0af00e9bf4d 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -204,13 +204,17 @@ class BaseCacheTransceiver { public: virtual ~BaseCacheTransceiver() = default; - virtual void respondAndSendAsync(LlmRequest* llmRequest) = 0; + // These methods take std::shared_ptr so the transceiver and + // its async workers can hold a strong reference for the duration of the + // transfer. See the comment on CacheTransceiver::mSenderFutures for the + // lifetime invariant (kept in one place to avoid drift). + virtual void respondAndSendAsync(std::shared_ptr llmRequest) = 0; virtual void respondAndSendLayerWise( RequestVector const& requests, std::shared_ptr const& progress) = 0; - virtual void requestAndReceiveSync(LlmRequest* llmRequest) = 0; - virtual void requestAndReceiveAsync(LlmRequest* llmRequest) = 0; + virtual void requestAndReceiveSync(std::shared_ptr llmRequest) = 0; + virtual void requestAndReceiveAsync(std::shared_ptr llmRequest) = 0; /// Check all requests transferring context, and return the requests that have completed or encountered an error. virtual RequestStatuses checkContextTransferStatus( @@ -221,7 +225,7 @@ class BaseCacheTransceiver [[nodiscard]] virtual bool checkGenTransferComplete() const = 0; - virtual bool cancelRequest(LlmRequest* llmRequest) = 0; + virtual bool cancelRequest(std::shared_ptr llmRequest) = 0; }; class CacheTransceiver : public BaseCacheTransceiver @@ -252,13 +256,13 @@ class CacheTransceiver : public BaseCacheTransceiver virtual ~CacheTransceiver(); - void respondAndSendAsync(LlmRequest* llmRequest) override; + void respondAndSendAsync(std::shared_ptr llmRequest) override; void respondAndSendLayerWise( RequestVector const& requests, std::shared_ptr const& progress) override; - void requestAndReceiveSync(LlmRequest* llmRequest) override; - void requestAndReceiveAsync(LlmRequest* llmRequest) override; + void requestAndReceiveSync(std::shared_ptr llmRequest) override; + void requestAndReceiveAsync(std::shared_ptr llmRequest) override; RequestStatuses checkContextTransferStatus( std::optional const& atLeastRequestNum = std::nullopt, bool markComplete = false) override; @@ -267,7 +271,7 @@ class CacheTransceiver : public BaseCacheTransceiver [[nodiscard]] bool checkGenTransferComplete() const override; - virtual bool cancelRequest(LlmRequest* llmRequest) override; + virtual bool cancelRequest(std::shared_ptr llmRequest) override; private: void initializeCommState(); @@ -276,8 +280,17 @@ class CacheTransceiver : public BaseCacheTransceiver std::unique_ptr mCacheSender; std::unique_ptr mCacheReceiver; - std::vector>> mSenderFutures; - std::vector>> mRequesterFutures; + // Store shared_ptr rather than raw LlmRequest* so the futures map holds a + // strong reference for the duration of the transfer. Otherwise Python's + // _terminate_request can drop its pybind shared_ptr while the C++ side's + // raw pointer is still dereferenced by checkGenTransferStatus / + // checkContextTransferStatus (the UAF forensically confirmed via + // MALLOC_PERTURB_=85 producing mRequestId=0x5555555555555555). Entries are + // erased only in the normal eviction paths (completion, deadline, + // exception), ensuring the request outlives every C++ access but doesn't + // leak past the transfer lifecycle. + std::vector, std::future>> mSenderFutures; + std::vector, std::future>> mRequesterFutures; mpi::MpiComm const* mMpiWorldComm{nullptr}; std::shared_ptr mGroupComm; diff --git a/cpp/include/tensorrt_llm/executor/transferAgent.h b/cpp/include/tensorrt_llm/executor/transferAgent.h index 8d6a46107675..73cd1829cf73 100644 --- a/cpp/include/tensorrt_llm/executor/transferAgent.h +++ b/cpp/include/tensorrt_llm/executor/transferAgent.h @@ -288,6 +288,15 @@ class TransferStatus virtual ~TransferStatus() = default; [[nodiscard]] virtual bool isCompleted() const = 0; virtual TransferState wait(int64_t timeout_ms = -1) const = 0; + + /// Release the backend transfer request. If the request is still active, + /// backends may attempt to cancel it. A true return only means the backend + /// accepted release of the transfer handle; callers must still treat remote + /// memory quiescence as backend-specific. + [[nodiscard]] virtual bool release() + { + return false; + } }; struct BaseAgentConfig diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp index 58092897ebbe..2928c87bbd06 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp @@ -26,6 +26,36 @@ namespace tensorrt_llm::batch_manager { +void BufferIndexHolder::release() noexcept +{ + // Happy-path release: frees the slot and disarms the holder in one + // noexcept call. Used in place of an older detach() + explicit + // freeBufferIndex*() sequence so a throw between the two calls cannot + // leave the holder in a partially-released state. + if (!mHeld || mMgr == nullptr) + { + return; + } + try + { + if (mIsRecv) + { + mMgr->freeBufferIndexForRecv(mIndex); + } + else + { + mMgr->freeBufferIndexForSend(mIndex); + } + } + catch (...) + { + // Swallow; the destructor must be noexcept and any exit path that + // failed to release explicitly relies on this fallback to free the + // slot. + } + mHeld = false; +} + BaseTransBufferManager::BaseTransBufferManager( size_t transferBufferSize, nvinfer1::DataType dataType, std::optional maxNumTokens) : mDataType{dataType} @@ -54,9 +84,11 @@ BaseTransBufferManager::BaseTransBufferManager( allocateBuffer(); } -std::optional BaseTransBufferManager::assignBufferIndexForSend() +std::optional BaseTransBufferManager::assignBufferIndexForSend( + std::atomic const* perRequestCancel, int64_t waitSliceMs, std::optional requestIdForLog) { - return assignBufferIndex(mConcurrenceSendResource, mSendBufferCount, mOnlyUseDynamicBuffer); + return assignBufferIndex(mConcurrenceSendResource, mSendBufferCount, mOnlyUseDynamicBuffer, perRequestCancel, + waitSliceMs, requestIdForLog); } void BaseTransBufferManager::freeBufferIndexForSend(std::optional bufferId) @@ -64,9 +96,11 @@ void BaseTransBufferManager::freeBufferIndexForSend(std::optional bufferId) freeBufferIndex(mConcurrenceSendResource, bufferId, mSendBufferCount, mOnlyUseDynamicBuffer); } -std::optional BaseTransBufferManager::assignBufferIndexForRecv() +std::optional BaseTransBufferManager::assignBufferIndexForRecv( + std::atomic const* perRequestCancel, int64_t waitSliceMs, std::optional requestIdForLog) { - return assignBufferIndex(mConcurrenceRecvResource, mRecvBufferCount, mOnlyUseDynamicBuffer); + return assignBufferIndex(mConcurrenceRecvResource, mRecvBufferCount, mOnlyUseDynamicBuffer, perRequestCancel, + waitSliceMs, requestIdForLog); } void BaseTransBufferManager::freeBufferIndexForRecv(std::optional bufferId) @@ -225,16 +259,35 @@ void BaseTransBufferManager::allocateBuffer() } } -std::optional BaseTransBufferManager::assignBufferIndex( - ConcurrenceResource& resource, size_t bufferCount, bool onlyUseDynamicBuffer) +std::optional BaseTransBufferManager::assignBufferIndex(ConcurrenceResource& resource, size_t bufferCount, + bool onlyUseDynamicBuffer, std::atomic const* perRequestCancel, int64_t waitSliceMs, + std::optional requestIdForLog) { if (onlyUseDynamicBuffer) { return std::nullopt; } + // Bounded wait_for loop so a cancel fired on this request while parked + // here can interrupt the wait via the per-request cancel atomic, and so + // mTerminate (flipped between slices) keeps the drain worker responsive + // to shutdown. std::unique_lock lk(resource.mBuffersMutex); - resource.mBuffersCV.wait( - lk, [&resource, bufferCount]() { return static_cast(resource.mConcurrence) < bufferCount; }); + auto const predicate + = [&resource, bufferCount]() { return static_cast(resource.mConcurrence) < bufferCount; }; + if (!predicate()) + { + auto const slice = std::chrono::milliseconds{waitSliceMs}; + while (!predicate()) + { + resource.mBuffersCV.wait_for(lk, slice); + if (perRequestCancel != nullptr && perRequestCancel->load(std::memory_order_relaxed)) + { + auto const reqIdStr + = requestIdForLog.has_value() ? std::to_string(requestIdForLog.value()) : std::string{"?"}; + TLLM_THROW("assignBufferIndex cancelled via perRequestCancel (reqId=%s)", reqIdStr.c_str()); + } + } + } int bufferId = -1; for (size_t i = 0; i < bufferCount; i++) { diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h index 1efeb89ccc04..de66d3e593f8 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h @@ -46,6 +46,122 @@ enum class BufferKind : uint8_t kRNN = 2 }; +class BaseTransBufferManager; + +/// @brief RAII scoped holder for a buffer index acquired from +/// BaseTransBufferManager::assignBufferIndexForRecv / +/// assignBufferIndexForSend. Releases the index on destruction, +/// including stack unwind from exceptions. +/// +/// Motivation: CacheReceiver::Impl::requestSync has at least six exit +/// paths (normal, early-cancel, not-ready, cancel-after-ready, +/// receiveReadySignal cancelled, exception from requestSync). Pre-fix, +/// the buffer-index release lived inside receiveSync's formatter, so any +/// exit path that skipped receiveSync leaked one index. Under saturation +/// even a single leaked index permanently wedged the (default size-1) +/// pool, so every subsequent request waited forever for an index that +/// would never be released. +/// +/// This holder closes that class of bug rather than patching the one +/// observed branch. Move-only so ownership is unambiguous; `detach()` +/// hands off ownership when the formatter inside receiveSync takes the +/// buffer's release responsibility on the happy path. +class BufferIndexHolder +{ +public: + BufferIndexHolder() = default; + + BufferIndexHolder(BaseTransBufferManager& mgr, std::optional index, bool isRecv, + std::optional requestIdForLog = std::nullopt) noexcept + : mMgr(&mgr) + , mIndex(index) + , mHeld(index.has_value()) + , mIsRecv(isRecv) + , mRequestIdForLog(requestIdForLog) + { + } + + // The destructor is inline but `release()` is out-of-line in + // baseTransBuffer.cpp because it dereferences BaseTransBufferManager and + // needs the full definition; defining it inline would create an include + // cycle between this header and the manager definition. + ~BufferIndexHolder() + { + release(); + } + + BufferIndexHolder(BufferIndexHolder const&) = delete; + BufferIndexHolder& operator=(BufferIndexHolder const&) = delete; + + BufferIndexHolder(BufferIndexHolder&& other) noexcept + : mMgr(other.mMgr) + , mIndex(other.mIndex) + , mHeld(other.mHeld) + , mIsRecv(other.mIsRecv) + , mRequestIdForLog(other.mRequestIdForLog) + { + other.mHeld = false; + } + + BufferIndexHolder& operator=(BufferIndexHolder&& other) noexcept + { + if (this != &other) + { + release(); + mMgr = other.mMgr; + mIndex = other.mIndex; + mHeld = other.mHeld; + mIsRecv = other.mIsRecv; + mRequestIdForLog = other.mRequestIdForLog; + other.mHeld = false; + } + return *this; + } + + [[nodiscard]] std::optional index() const noexcept + { + return mIndex; + } + + [[nodiscard]] bool held() const noexcept + { + return mHeld; + } + + /// @brief Relinquish ownership without releasing. Use when a downstream + /// owner (e.g. the formatter inside receiveSync) takes over the + /// release responsibility on the happy path. + std::optional detach() noexcept + { + mHeld = false; + return mIndex; + } + + /// @brief Happy-path release. Frees the slot immediately and disarms the + /// destructor. Use this on any path where the caller has confirmed + /// the slot is no longer needed and the release is the expected + /// outcome (e.g. the sender formatter after sendAllBuffers + /// returns). After this call, the holder owns nothing; subsequent + /// destructor or move-assignment is a no-op. + /// + /// If the holder goes out of scope with mHeld still true (exception + /// or early return that forgot to call release/detach), the + /// destructor calls release() to free the slot. + void release() noexcept; + +private: + BaseTransBufferManager* mMgr{nullptr}; + std::optional mIndex{}; + bool mHeld{false}; + bool mIsRecv{true}; + std::optional mRequestIdForLog{}; +}; + +/// @brief Default per-iteration slice for the buffer-acquire CV wait. Chosen +/// small enough to keep cancellation latency under ~100 ms and large +/// enough to avoid spinning under saturation. +inline constexpr int64_t kBufferAcquireSliceMs = 100; + /// @brief Base class for cache transfer buffer management. /// Handles buffer pool allocation, index assignment, and slicing. /// Derived classes provide cache-specific size calculations. @@ -57,16 +173,40 @@ class BaseTransBufferManager [[nodiscard]] virtual BufferKind getBufferKind() const = 0; /// @brief Assign a buffer index for sending. + /// @param perRequestCancel Optional per-request cancel flag. When non-null + /// and flipped true while this call is parked on the pool-exhausted + /// CV wait, the function throws so the caller (sender worker) can + /// unwind instead of blocking indefinitely. Checked every + /// `waitSliceMs` during the wait. Parity with + /// assignBufferIndexForRecv. + /// @param waitSliceMs Per-iteration timeout for the internal condition + /// variable wait (ms). Defaults to kBufferAcquireSliceMs. + /// @param requestIdForLog Optional request id used to tag any + /// diagnostic log lines so a pool-exhausted wedge on the send + /// side can be attributed to a specific reqId. /// @return Assigned buffer index, or nullopt if using dynamic buffers. - std::optional assignBufferIndexForSend(); + std::optional assignBufferIndexForSend(std::atomic const* perRequestCancel = nullptr, + int64_t waitSliceMs = kBufferAcquireSliceMs, std::optional requestIdForLog = std::nullopt); /// @brief Free a buffer index used for sending. /// @param bufferId The buffer index to free. void freeBufferIndexForSend(std::optional bufferId); /// @brief Assign a buffer index for receiving. + /// @param perRequestCancel Optional per-request cancel flag. When non-null + /// and flipped true while this call is parked on the pool-exhausted + /// CV wait, the function throws so the caller (drain worker) can + /// unwind instead of blocking indefinitely. Checked every + /// `waitSliceMs` during the wait; also bounds the wait by polling + /// for recovery even without an explicit cancel. + /// @param waitSliceMs Per-iteration timeout for the internal condition + /// variable wait (ms). Defaults to kBufferAcquireSliceMs. + /// @param requestIdForLog Optional request id used to tag any + /// diagnostic log lines so a pool-exhausted wedge can be + /// attributed to a specific reqId. /// @return Assigned buffer index, or nullopt if using dynamic buffers. - std::optional assignBufferIndexForRecv(); + std::optional assignBufferIndexForRecv(std::atomic const* perRequestCancel = nullptr, + int64_t waitSliceMs = kBufferAcquireSliceMs, std::optional requestIdForLog = std::nullopt); /// @brief Free a buffer index used for receiving. /// @param bufferId The buffer index to free. @@ -132,7 +272,9 @@ class BaseTransBufferManager runtime::BufferManager const& bufferManagerToUse, ConcurrenceResource& concurrenceResource); void allocateBuffer(); - std::optional assignBufferIndex(ConcurrenceResource& resource, size_t bufferCount, bool onlyUseDynamicBuffer); + std::optional assignBufferIndex(ConcurrenceResource& resource, size_t bufferCount, bool onlyUseDynamicBuffer, + std::atomic const* perRequestCancel = nullptr, int64_t waitSliceMs = kBufferAcquireSliceMs, + std::optional requestIdForLog = std::nullopt); void freeBufferIndex( ConcurrenceResource& resource, std::optional bufferId, size_t bufferCount, bool onlyUseDynamicBuffer); diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index e05e8d6f76fc..7b531165f127 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -34,6 +34,7 @@ #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include +#include #include #include #include @@ -63,7 +64,6 @@ void sendBuffer(TransferSession& session, int deviceId, size_t localIdx, executor::kv_cache::TargetRanksInfo const& targetInfo, std::vector const& pickUpConnections) { size_t connIdx = pickUpConnections[localIdx]; - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), " send localIdx: %ld connIdx: %ld", localIdx, connIdx); NVTX3_SCOPED_RANGE(sendBuffer); TLLM_CUDA_CHECK(cudaSetDevice(deviceId)); TLLM_CHECK(session.getConnections().size() > (connIdx / targetInfo.mPeerDupHeadFactor)); @@ -76,11 +76,7 @@ void sendBuffer(TransferSession& session, int deviceId, size_t localIdx, if (bufferIdx < bufferCoverTargetNum) { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), " send connIdx: %ld bufferIdx: %ld size:%ld", connIdx, - bufferIdx, outputBuffers[bufferIdx]->getSizeInBytes()); session.send(connIdx, outputBuffers[bufferIdx]->data(), outputBuffers[bufferIdx]->getSizeInBytes()); - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), " end send connIdx: %ld bufferIdx: %ld size:%ld", connIdx, - bufferIdx, outputBuffers[bufferIdx]->getSizeInBytes()); } else { @@ -494,7 +490,17 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio // cache blocks to the corresponding buffer. // 5. send the buffer to the corresponding target. Ideally, we send only once (one buffer) for each target. - auto cacheBufferId = mCacheTransBufferManager->assignBufferIndexForSend(); + // RAII wrapper mirrors the receiver-side pattern. Happy-path calls + // sendHolder.release() after sendAllBuffers returns; any other exit + // between acquire and release auto-releases via ~BufferIndexHolder. + // The send-pool CV wait inside assignBufferIndexForSend observes the + // session's per-request cancel flag and throws on cancel (parity + // with the recv side). + auto const sendReqIdForLog = std::make_optional(static_cast(llmRequest.mRequestId)); + auto const* sendCancelFlag = &session.getDataContext().getTransferTerminate(); + auto cacheBufferId = mCacheTransBufferManager->assignBufferIndexForSend( + sendCancelFlag, kBufferAcquireSliceMs, sendReqIdForLog); + BufferIndexHolder sendHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/false, sendReqIdForLog); int peerDuplicateHeadFactor = targetInfo.mPeerDupHeadFactor; auto bufferTargetNum = targetNum / peerDuplicateHeadFactor; auto ppRank = selfIdx @@ -576,9 +582,12 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio sendAllBuffers(session, deviceId, outputSplitCaches, bufferCoverTargetNum, preAllocSendBuffer, bufferManager, targetInfo, pickUpConnections); + // Happy-path release — frees the slot and disarms the holder in + // one noexcept call. Placed immediately after sendAllBuffers so any + // subsequent throw (e.g. from setTime) does not turn a non-leak into + // a destructor-driven release. + sendHolder.release(); session.setTime(TransferSession::kTimeTransmissions); - - mCacheTransBufferManager->freeBufferIndexForSend(cacheBufferId); session.setTime(TransferSession::kTimePostprocess); } TLLM_LOG_DEBUG( diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 2e4bf1f06667..b109fe8793cb 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -60,6 +60,18 @@ namespace tensorrt_llm::batch_manager { +namespace +{ +/// @brief Source the rank used as the prefix of TLLM_LOG_DEBUG / TLLM_LOG_WARNING +/// calls. We call useMPI() at log time so the same call site works +/// regardless of whether the world communicator is initialized via +/// MPI or via the torch process group. +inline int currentRankForLog() +{ + return useMPI() ? mpi::MpiComm::world().getRank() : tensorrt_llm::pg_utils::get_world_pg()->getRank(); +} +} // namespace + std::mutex CacheTransceiver::mDllMutex; std::unique_ptr CacheTransceiverFactory::createCacheTransceiver( @@ -323,7 +335,7 @@ void CacheTransceiver::setContextState(LlmRequest* llmRequest) } } -void CacheTransceiver::respondAndSendAsync(LlmRequest* llmRequest) +void CacheTransceiver::respondAndSendAsync(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isContextOnlyRequest()); llmRequest->setState(LlmRequestState::kDISAGG_CONTEXT_TRANS_IN_PROGRESS); @@ -337,9 +349,11 @@ void CacheTransceiver::respondAndSendAsync(LlmRequest* llmRequest) } return; } - setContextState(llmRequest); - auto future = mCacheSender->sendAsync(*llmRequest); - mSenderFutures.emplace_back(llmRequest, std::move(future)); + setContextState(llmRequest.get()); + auto future = mCacheSender->sendAsync(llmRequest); + TLLM_LOG_DEBUG("respondAndSendAsync: adding request %ld to mSenderFutures (size=%zu)", llmRequest->mRequestId, + mSenderFutures.size() + 1); + mSenderFutures.emplace_back(std::move(llmRequest), std::move(future)); } void CacheTransceiver::respondAndSendLayerWise( @@ -354,22 +368,22 @@ void CacheTransceiver::respondAndSendLayerWise( llmRequest->setState(LlmRequestState::kDISAGG_CONTEXT_INIT_AND_TRANS); setContextState(llmRequest.get()); - auto future = mCacheSender->sendAsync(*llmRequest); - mSenderFutures.emplace_back(llmRequest.get(), std::move(future)); + auto future = mCacheSender->sendAsync(llmRequest); + mSenderFutures.emplace_back(llmRequest, std::move(future)); } } -void CacheTransceiver::requestAndReceiveSync(LlmRequest* llmRequest) +void CacheTransceiver::requestAndReceiveSync(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isGenerationOnlyRequest()); { - auto future = mCacheReceiver->receiveAsync(*llmRequest); + auto future = mCacheReceiver->receiveAsync(llmRequest); future.get(); } llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); } -void CacheTransceiver::requestAndReceiveAsync(LlmRequest* llmRequest) +void CacheTransceiver::requestAndReceiveAsync(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isGenerationOnlyRequest()); @@ -381,9 +395,29 @@ void CacheTransceiver::requestAndReceiveAsync(LlmRequest* llmRequest) return; } - auto future = mCacheReceiver->receiveAsync(*llmRequest); - mRequesterFutures.emplace_back(llmRequest, std::move(future)); + // Record the transfer start synchronously, before pushing the future. + // CacheReceiver::receiveAsync() spawns a background thread that does the + // same thing inside requestSync(), but until that thread runs, + // getKvCacheTransferStart() would return the default (epoch) time point. + // checkGenTransferStatus's elapsed-time deadline check would then see a + // huge elapsed value and falsely evict the entry. A preemptive set + // here keeps the invariant "every entry in mRequesterFutures has a + // valid transfer start time" and is overwritten by requestSync() with a + // slightly later time — harmless for the deadline check. + llmRequest->setKvCacheTransferStart(LlmRequest::getSteadyClockNow()); + auto future = mCacheReceiver->receiveAsync(llmRequest); + TLLM_LOG_DEBUG("requestAndReceiveAsync: adding request %ld to mRequesterFutures (size=%zu)", llmRequest->mRequestId, + mRequesterFutures.size() + 1); + // Order: setKvCacheTransferStart -> receiveAsync -> setState -> emplace. + // setState is intentionally moved AFTER receiveAsync (the original code + // set it before). This is safe today because the async worker spawned + // by receiveAsync does not read llmRequest state at entry; the new + // ordering tightens the IN_PROGRESS<->mRequesterFutures invariant and + // ensures a throw from receiveAsync leaves the request out of both + // sets atomically. Do not revert without verifying the worker still + // does not depend on the IN_PROGRESS state being set at entry. llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS); + mRequesterFutures.emplace_back(std::move(llmRequest), std::move(future)); } std::vector gatherRequestIds( @@ -485,10 +519,27 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( { bool blockAll = !atLeastRequestNum.has_value(); std::optional senderFutureTimeoutMs = std::nullopt; - // If blockAll is true, we want to block and not use a timeout - if (!blockAll && mCacheTransceiverConfig.has_value()) + std::optional kvTransferTimeoutMs = std::nullopt; + if (mCacheTransceiverConfig.has_value()) + { + // The overall transfer deadline applies in both blockAll and polling modes; a + // stuck sender must not hang indefinitely regardless of how callers invoke this. + kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); + // The per-iteration poll timeout is only relevant when the caller wants to + // return after checking at least `atLeastRequestNum` entries. + if (!blockAll) + { + senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); + } + } + + // Log mSenderFutures state for diagnosing dangling pointer issues. + // Each entry's pointer address and request ID are logged so we can detect + // when a pointer's underlying memory is freed (reqId changes to 0). + if (!mSenderFutures.empty()) { - senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); + TLLM_LOG_DEBUG("checkContextTransferStatus: mSenderFutures.size()=%zu, blockAll=%d, kvTransferTimeoutMs=%d", + mSenderFutures.size(), blockAll ? 1 : 0, kvTransferTimeoutMs.value_or(-1)); } auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupTPInDPComm : mGroupTensorParaComm; @@ -546,13 +597,46 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( for (auto it = mSenderFutures.begin(); it != mSenderFutures.end();) { auto& [request, future] = *it; + // Enforce the overall deadline for every entry on every invocation, + // independent of the readiness gate below. `toCompleteIdSet` is + // populated from the consensus freqVec (entries that were ready on + // every rank) plus up to `atLeastRequestNum` insertion-order entries. + // A stuck sender whose future never becomes ready is not in the + // consensus set, so with `atLeastRequestNum=0` it would otherwise + // escape the deadline check entirely and pin KV blocks forever. + // mSenderFutures holds shared_ptr, so the request is kept + // alive for every C++ access here regardless of Python-side + // termination timing. + if (kvTransferTimeoutMs.has_value()) + { + auto elapsed = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()); + auto elapsedMs = static_cast(elapsed.count()); + if (elapsedMs > kvTransferTimeoutMs.value()) + { + TLLM_LOG_WARNING( + "Context KV cache transfer for request %ld exceeded total timeout: " + "elapsed %ld ms > limit %d ms. Marking as error.", + request->mRequestId, elapsedMs, kvTransferTimeoutMs.value()); + // Defense-in-depth sender-side cancel. Sender zombies empirically + // unwind on peer teardown (decode-pod restart), but in general + // CacheSender::cancelRequest clears mReadyResponses / + // mCancelledRequests bookkeeping so a subsequent re-enqueue + // or telemetry path doesn't see the stale request. + mCacheSender->cancelRequest(*request); + request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + requestsStatus.errorRequestIds.insert(request->mRequestId); + it = mSenderFutures.erase(it); + continue; + } + } if (blockAll || (toCompleteIdSet.find(request->mRequestId) != toCompleteIdSet.end())) { try { - // Wait for up to a specified timeout + // Wait for up to a specified timeout (0 means a single non-blocking poll in blockAll mode). auto status = future.wait_for(std::chrono::milliseconds(senderFutureTimeoutMs.value_or(0))); - if (status == std::future_status::ready || !senderFutureTimeoutMs.has_value()) + if (status == std::future_status::ready) { future.get(); requestsStatus.completedRequestIds.insert(request->mRequestId); @@ -564,9 +648,28 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } else if (status == std::future_status::timeout) { - TLLM_LOG_WARNING("Timed out waiting for context KV cache transfer after %d milliseconds.", - senderFutureTimeoutMs.value()); - ++it; + // The overall deadline was already enforced unconditionally + // at the top of this loop, so if we reach here the deadline + // has not yet passed for this entry. + if (senderFutureTimeoutMs.has_value()) + { + TLLM_LOG_WARNING("Timed out waiting for context KV cache transfer after %d milliseconds.", + senderFutureTimeoutMs.value()); + ++it; + } + else + { + // blockAll mode with no deadline exceeded: block on get() as the + // caller intends, but the deadline will be re-checked on each entry + // of subsequent outer invocations. + future.get(); + requestsStatus.completedRequestIds.insert(request->mRequestId); + if (markComplete) + { + request->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); + } + it = mSenderFutures.erase(it); + } } else { @@ -580,8 +683,12 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } catch (std::exception const& e) { - TLLM_LOG_ERROR( - "Error occurred during context transfer for request %ld: %s", request->mRequestId, e.what()); + // mSenderFutures holds shared_ptr, so the request + // is alive here regardless of Python-side termination timing. + // Report as error so Python can call end_transfer to unpin + // blocks. + TLLM_LOG_WARNING("Error during context transfer for request %ld: %s. Marking as error.", + request->mRequestId, e.what()); request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); requestsStatus.errorRequestIds.insert(request->mRequestId); it = mSenderFutures.erase(it); @@ -593,12 +700,48 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } } + if (!requestsStatus.completedRequestIds.empty() || !requestsStatus.errorRequestIds.empty()) + { + TLLM_LOG_DEBUG( + "checkContextTransferStatus done: completed=%zu, errors=%zu, " + "mSenderFutures.size()=%zu", + requestsStatus.completedRequestIds.size(), requestsStatus.errorRequestIds.size(), mSenderFutures.size()); + } + return requestsStatus; } void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastRequestNum) { bool blockAll = !atLeastRequestNum.has_value(); + std::optional senderFutureTimeoutMs = std::nullopt; + std::optional kvTransferTimeoutMs = std::nullopt; + if (mCacheTransceiverConfig.has_value()) + { + // The overall transfer deadline applies in both blockAll and polling modes; a + // stuck receiver must not hang indefinitely regardless of how callers invoke + // this. Without this, mRequesterFutures would accumulate stuck entries, + // pinning generation-side KV blocks and eventually exhausting the pool. + kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); + // The per-iteration poll timeout is only relevant when the caller wants to + // return after checking at least `atLeastRequestNum` entries. + if (!blockAll) + { + senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); + } + } + + // Log mRequesterFutures state for diagnosing dangling pointer / stuck + // receiver issues. Mirrors the diagnostic added to checkContextTransferStatus. + if (!mRequesterFutures.empty()) + { + TLLM_LOG_DEBUG( + "checkGenTransferStatus: mRequesterFutures.size()=%zu, blockAll=%d, kvTransferTimeoutMs=%d, " + "senderFutureTimeoutMs=%d", + mRequesterFutures.size(), blockAll ? 1 : 0, kvTransferTimeoutMs.value_or(-1), + senderFutureTimeoutMs.value_or(-1)); + } + std::vector genTransferReadyRequestIds; for (auto&& [request, future] : mRequesterFutures) { @@ -641,16 +784,8 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR break; } toCompleteIdSet.insert(freqVec.at(idx).first); - if (useMPI()) - { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - " checkGenTransferStatus at least from freqVec requestId: %zu ", freqVec.at(idx).first); - } - else - { - TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), - " checkGenTransferStatus at least from freqVec requestId: %zu ", freqVec.at(idx).first); - } + TLLM_LOG_DEBUG(currentRankForLog(), " checkGenTransferStatus at least from freqVec requestId: %zu ", + freqVec.at(idx).first); idx++; } idx = 0; @@ -665,18 +800,9 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR if (toCompleteIdSet.find(mRequesterFutures.at(idx).first->mRequestId) == toCompleteIdSet.end()) { toCompleteIdSet.insert(mRequesterFutures.at(idx).first->mRequestId); - if (useMPI()) - { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - " checkGenTransferStatus at least from RequesterFuture requestId: %zu atLeastRequestNum:%d", - mRequesterFutures.at(idx).first->mRequestId, atLeastRequestNum.value_or(0)); - } - else - { - TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), - " checkGenTransferStatus at least from RequesterFuture requestId: %zu atLeastRequestNum:%d", - mRequesterFutures.at(idx).first->mRequestId, atLeastRequestNum.value_or(0)); - } + TLLM_LOG_DEBUG(currentRankForLog(), + " checkGenTransferStatus at least from RequesterFuture requestId: %zu atLeastRequestNum:%d", + mRequesterFutures.at(idx).first->mRequestId, atLeastRequestNum.value_or(0)); } idx++; } @@ -686,70 +812,157 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR { toCompleteIdSet.insert(requestId); } - if (useMPI()) + TLLM_LOG_DEBUG( + currentRankForLog(), " checkGenTransferStatus freqVec requestId: %zu,freq:%d ", requestId, freq); + } + TLLM_LOG_DEBUG(currentRankForLog(), " checkGenTransferStatus toCompleteIdSet size: %zu, atLeastRequestNum: %d ", + toCompleteIdSet.size(), atLeastRequestNum.value_or(0)); + // Helper: finalize a generation-side transfer on the happy path. + // Called in two places: status == ready, and blockAll fall-through + // from the timeout branch. + auto const completeEntry = [this](std::shared_ptr const& request) + { + request->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); + if (!common::getEnvKVCacheTimeOutputPath().empty()) { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), " checkGenTransferStatus freqVec requestId: %zu,freq:%d ", - requestId, freq); + auto transferSyncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; + updateKVCacheTransferBW(transferSyncComm, request.get()); } - else - { - TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), - " checkGenTransferStatus freqVec requestId: %zu,freq:%d ", requestId, freq); - } - } - if (useMPI()) - { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - " checkGenTransferStatus toCompleteIdSet size: %zu, atLeastRequestNum: %d ", toCompleteIdSet.size(), - atLeastRequestNum.value_or(0)); - } - else - { - TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), - " checkGenTransferStatus toCompleteIdSet size: %zu, atLeastRequestNum: %d ", toCompleteIdSet.size(), - atLeastRequestNum.value_or(0)); - } + TLLM_LOG_DEBUG(currentRankForLog(), + "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***", request->mRequestId, + request->getContextPhaseParams().value().getReqId()); + }; + + std::size_t numCompleted = 0; + std::size_t numErrored = 0; for (auto it = mRequesterFutures.begin(); it != mRequesterFutures.end();) { - if (blockAll || toCompleteIdSet.find(it->first->mRequestId) != toCompleteIdSet.end()) + auto& request = it->first; // std::shared_ptr const& — our own strong ref + auto& future = it->second; + // Enforce the overall deadline for every entry on every invocation, + // independent of the readiness gate below. In polling mode (blockAll == + // false) `toCompleteIdSet` is populated from the initial + // `wait_for(0)` sweep and only contains receivers whose futures were + // already ready. A receiver whose peer sender evicted the transfer + // (see checkContextTransferStatus above) never becomes ready, so + // gating the deadline check behind `toCompleteIdSet` would let stuck + // entries accumulate in `mRequesterFutures` and pin generation-side + // KV blocks forever. + // + // Lifetime: mRequesterFutures holds shared_ptr, so the + // request is kept alive for every C++ access here regardless of + // whether Python has already dropped its active_requests reference. + // This is the load-bearing invariant for the lifetime of all + // accesses below; the older raw-pointer design relied on a Python + // guard to delay _terminate_request and produced an orphan-induced + // KV-block leak when the guard interacted with stuck transfers. + if (kvTransferTimeoutMs.has_value()) + { + auto elapsed = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()); + auto elapsedMs = static_cast(elapsed.count()); + if (elapsedMs > kvTransferTimeoutMs.value()) + { + TLLM_LOG_WARNING( + "Generation KV cache transfer for request %ld exceeded total timeout: " + "elapsed %ld ms > limit %d ms. Marking as error.", + request->mRequestId, elapsedMs, kvTransferTimeoutMs.value()); + // Erasing the std::future from mRequesterFutures only + // destroys the Python-facing handle; the std::async task + // spawned by CacheReceiver::receiveAsync is still running, + // blocked inside recvReadySignal waiting on a peer signal + // that will never arrive. Over a saturation window these + // zombies accumulate and pin NIXL/UCX per-request state, so + // subsequent receives queue behind a frozen pool and every + // new transfer deterministically waits kv_transfer_timeout_ms + // even after load stops (the "no-recovery" bug). Call + // cancelRequest to flip the per-request cancel flag in + // CacheReceiver::Impl, which the notification polling loop + // observes via the DataContext and returns early with + // isReady=false. requestSync then takes the + // kDISAGG_TRANS_ERROR branch and the task exits cleanly. + mCacheReceiver->cancelRequest(*request); + request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + it = mRequesterFutures.erase(it); + ++numErrored; + continue; + } + } + if (blockAll || toCompleteIdSet.find(request->mRequestId) != toCompleteIdSet.end()) { try { - it->second.get(); - it->first->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); - - // Gather the kv cache transfer time from all workers and update to leader rank - if (!common::getEnvKVCacheTimeOutputPath().empty()) + // Wait for up to a specified timeout (0 means a single non-blocking + // poll in blockAll mode). + auto status = future.wait_for(std::chrono::milliseconds(senderFutureTimeoutMs.value_or(0))); + if (status == std::future_status::ready) { - auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; - updateKVCacheTransferBW(syncComm, it->first); + future.get(); + completeEntry(request); + it = mRequesterFutures.erase(it); + ++numCompleted; + } + else if (status == std::future_status::timeout) + { + // The overall deadline was already enforced unconditionally at + // the top of this loop, so if we reach here the deadline has + // not yet passed for this entry. + if (senderFutureTimeoutMs.has_value()) + { + TLLM_LOG_WARNING( + "Timed out waiting for generation KV cache transfer for request %ld " + "after %d milliseconds (per-iteration).", + request->mRequestId, senderFutureTimeoutMs.value()); + ++it; + } + else + { + // blockAll mode with deadline not yet exceeded: block on get() + // as the caller requested. A hard cap across a stuck receiver + // only applies on subsequent outer invocations — callers that + // need an unconditional hard cap should pass atLeastRequestNum + // so the per-iteration poll timeout can drive the deadline + // check each tick. + future.get(); + completeEntry(request); + it = mRequesterFutures.erase(it); + ++numCompleted; + } + } + else + { + TLLM_LOG_ERROR("Future returned unexpected status for generation request %ld. Marking as error.", + request->mRequestId); + request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + it = mRequesterFutures.erase(it); + ++numErrored; } } catch (std::exception const& e) { - TLLM_LOG_ERROR( - "Error occurred during generation transfer for request %ld: %s", it->first->mRequestId, e.what()); - it->first->setState(LlmRequestState::kDISAGG_TRANS_ERROR); - } - if (useMPI()) - { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***", - it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId()); - } - else - { - TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), - "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***", - it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId()); + // mRequesterFutures holds shared_ptr, so the + // request is alive here regardless of Python-side + // termination timing. Report as error so Python's + // _check_cache_transfer_errors picks it up (via state == + // kDISAGG_TRANS_ERROR) and cleanly unwinds the request. + TLLM_LOG_WARNING("Error during generation transfer for request %ld: %s. Marking as error.", + request->mRequestId, e.what()); + request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + it = mRequesterFutures.erase(it); + ++numErrored; } - it = mRequesterFutures.erase(it); } else { ++it; } } + + if (numCompleted > 0 || numErrored > 0) + { + TLLM_LOG_DEBUG("checkGenTransferStatus done: completed=%zu, errored=%zu, mRequesterFutures.size()=%zu", + numCompleted, numErrored, mRequesterFutures.size()); + } } bool CacheTransceiver::checkGenTransferComplete() const @@ -757,7 +970,7 @@ bool CacheTransceiver::checkGenTransferComplete() const return mRequesterFutures.empty(); } -bool CacheTransceiver::cancelRequest(LlmRequest* llmRequest) +bool CacheTransceiver::cancelRequest(std::shared_ptr llmRequest) { if (llmRequest->isContextOnlyRequest()) { diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 3ecceb9f3f2c..e535b6887027 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -296,16 +296,27 @@ class CacheSender::Impl } } - [[nodiscard]] std::future sendAsync(LlmRequest& llmRequest) + [[nodiscard]] std::future sendAsync(std::shared_ptr const& llmRequest) { std::promise promise; auto future = promise.get_future(); - llmRequest.setKvCacheTransferStart(LlmRequest::getSteadyClockNow()); + llmRequest->setKvCacheTransferStart(LlmRequest::getSteadyClockNow()); + // Register a per-request cancel flag for this sender-side request, + // symmetric to CacheReceiver::Impl's flag registry. This allows + // CacheSender::cancelRequest to flip the flag when the request is + // already being processed (i.e. when it is the mCurrentRequest and + // therefore not cancellable via the queue-drain path). The flag is + // consumed by the session's DataContext (built in recvRequestInfo) + // and by AgentConnection::send's poll-wait loop, which breaks out + // on cancel. + (void) getOrCreateInFlightCancelFlag(llmRequest->mRequestId); { { std::scoped_lock lkResp(mSenderMutex); - mReadyResponses.emplace( - llmRequest.mRequestId, Response{std::addressof(llmRequest), std::move(promise)}); + // Worker holds shared_ptr so Python-side _terminate_request + // cannot drop the LlmRequest out from under the async-send + // worker's dereferences. + mReadyResponses.emplace(llmRequest->mRequestId, Response{llmRequest, std::move(promise)}); } std::unique_lock lkCond(mCondMutex); mAnyReady = true; @@ -314,6 +325,19 @@ class CacheSender::Impl return future; } + std::shared_ptr> getOrCreateInFlightCancelFlag(RequestIdType requestId) + { + std::lock_guard lg(mInFlightCancelMutex); + auto it = mInFlightCancelFlags.find(requestId); + if (it != mInFlightCancelFlags.end()) + { + return it->second; + } + auto flag = std::make_shared>(false); + mInFlightCancelFlags.emplace(requestId, flag); + return flag; + } + [[nodiscard]] executor::kv_cache::CommState const& getCommState() const { return mSelfState.getCommState().value(); @@ -334,21 +358,37 @@ class CacheSender::Impl void release(LlmRequest::RequestIdType requestId) { - std::unique_lock lk(mMtxForMap); - auto it = mRequestToSession.find(requestId); - TLLM_CHECK(it != mRequestToSession.end()); - if (!common::getEnvKVCacheTimeOutputPath().empty()) { - if (!mMeasuresFile.is_open()) + std::unique_lock lk(mMtxForMap); + auto it = mRequestToSession.find(requestId); + TLLM_CHECK(it != mRequestToSession.end()); + if (!common::getEnvKVCacheTimeOutputPath().empty()) { - auto outputPath = getTransferOutputPath("send"); - mMeasuresFile.open(outputPath); - TLLM_CHECK_WITH_INFO( - mMeasuresFile.is_open(), "Failed to open transfer output file: %s", outputPath.string().c_str()); + if (!mMeasuresFile.is_open()) + { + auto outputPath = getTransferOutputPath("send"); + mMeasuresFile.open(outputPath); + TLLM_CHECK_WITH_INFO(mMeasuresFile.is_open(), "Failed to open transfer output file: %s", + outputPath.string().c_str()); + } + it->second.exportMeasure(mMeasuresFile, true); } - it->second.exportMeasure(mMeasuresFile, true); + // Erase the session first so its DataContext (which references the + // per-request cancel atomic) is destroyed before we drop the + // flag's shared_ptr below. This ordering guarantees no dangling + // reference from DataContext into a freed atomic. + mRequestToSession.erase(it); + } + // Drop the per-request cancel flag now that the session is gone. + // This is the single point where flags are reclaimed during normal + // operation; cancel paths that skip release() (sendSync threw or + // mCancelledRequests fired) intentionally leak the flag shared_ptr + // until ~Impl / terminate(), matching the existing session-leak + // behavior on those error paths. + { + std::lock_guard lg(mInFlightCancelMutex); + mInFlightCancelFlags.erase(requestId); } - mRequestToSession.erase(it); } [[nodiscard]] RequestInfo recvRequestInfo() @@ -391,13 +431,27 @@ class CacheSender::Impl TLLM_CHECK_WITH_INFO(peerIdx < static_cast(allCounterparts.size()), "Peer rank %d not found in expected counterparts", peerSelfIdx); + // Get or create the per-request cancel flag for this requestId. If + // sendAsync already ran for this reqId the flag is registered; if + // recvRequestInfo races ahead, we create it here so the session + // DataContext below holds a live reference. The flag's shared_ptr + // lifetime is tied to the mInFlightCancelFlags map; it's erased + // only after sendSync finishes (in removeResponse / cancel + // cleanup), so the DataContext reference never dangles. + auto cancelFlag = getOrCreateInFlightCancelFlag(requestId); { std::unique_lock lk(mMtxForMap); auto it = mRequestToSession.find(requestId); if (it == mRequestToSession.end()) { + // Use the per-request cancel flag (not mTerminate) so a + // cancel fired by CacheSender::cancelRequest mid-send is + // observed by AgentConnection::send's poll-wait loop via + // ctx.getTransferTerminate(). Shutdown still works because + // ~Impl flips every registered per-request flag before + // joining the response worker. auto session = TransferSession(std::vector(allCounterparts.size(), nullptr), - DataContext{tagFromRequestId(requestId), mTerminate}, allCounterparts, mSelfState, + DataContext{tagFromRequestId(requestId), *cancelFlag}, allCounterparts, mSelfState, info.getTransState(), mBufferManager, info.getIndexFromEnd(), info.getLastBlockKey(), nullptr, !common::getEnvKVCacheTimeOutputPath().empty()); session.setTime(TransferSession::kTimeRequestInfo); @@ -425,18 +479,39 @@ class CacheSender::Impl bool cancelRequest(LlmRequest const& llmRequest) { bool isCancelled = false; - std::scoped_lock lkResp(mSenderMutex); - auto it = mReadyResponses.find(llmRequest.mRequestId); - // If the request is not the current request and already in the ready queue, we can cancel it. - if (it != mReadyResponses.end() - && (!mCurrentRequest.has_value() || getCurrentRequestId() != llmRequest.mRequestId)) { - mCancelledRequests.insert(llmRequest.mRequestId); - isCancelled = true; + std::scoped_lock lkResp(mSenderMutex); + auto it = mReadyResponses.find(llmRequest.mRequestId); + // If the request is not the current request and already in the ready queue, we can cancel it. + if (it != mReadyResponses.end() + && (!mCurrentRequest.has_value() || getCurrentRequestId() != llmRequest.mRequestId)) + { + mCancelledRequests.insert(llmRequest.mRequestId); + isCancelled = true; + } } - else + if (!isCancelled) { - TLLM_LOG_WARNING("Cannot cancel request %zu", llmRequest.mRequestId); + // Request is the mCurrentRequest (or not even in mReadyResponses) + // — the queue-drain branch can't abort it. Flip the per-request + // cancel flag registered at sendAsync time; AgentConnection::send + // observes ctx.getTransferTerminate() in its poll-wait loop and + // throws, which unwinds sendSync and lets the response worker + // resume. Symmetric to CacheReceiver::cancelRequest's in-flight + // cancel-flag branch below. + std::lock_guard lg(mInFlightCancelMutex); + auto flagIt = mInFlightCancelFlags.find(llmRequest.mRequestId); + if (flagIt != mInFlightCancelFlags.end()) + { + flagIt->second->store(true); + isCancelled = true; + TLLM_LOG_DEBUG("Flipped in-flight sender cancel flag for request %zu (not in queue or is current).", + llmRequest.mRequestId); + } + else + { + TLLM_LOG_WARNING("Cannot cancel request %zu", llmRequest.mRequestId); + } } return isCancelled; } @@ -477,7 +552,10 @@ class CacheSender::Impl private: struct Response { - LlmRequest* mRequest; + // Store shared_ptr rather than raw pointer so the async-send worker's + // dereferences stay safe past Python-side _terminate_request. Same + // UAF mitigation as RequestAndPromise on the receiver side. + std::shared_ptr mRequest; std::promise mPromise; }; @@ -511,7 +589,9 @@ class CacheSender::Impl resp = std::move(resource.mSendQueue.front()); resource.mSendQueue.pop_front(); } - sendAndRemoveResponse(resp.mRequest->mRequestId, std::move(resp)); + TLLM_CHECK(resp.mRequest != nullptr); + auto const reqId = resp.mRequest->mRequestId; + sendAndRemoveResponse(reqId, std::move(resp)); } } @@ -584,14 +664,22 @@ class CacheSender::Impl { // TODO: if the generation does not require the kv cache, the request will // not be removed from mCancelledRequests. This should be handled by timeout. - auto it = mReadyResponses.find(mCurrentRequest.value()); - TLLM_CHECK(it != mReadyResponses.end()); + auto const cancelledReqId = mCurrentRequest.value(); + Response cancelledResponse; { std::scoped_lock lkResp(mSenderMutex); + auto it = mReadyResponses.find(cancelledReqId); + TLLM_CHECK(it != mReadyResponses.end()); + cancelledResponse = std::move(it->second); mReadyResponses.erase(it); - mCancelledRequests.erase(mCurrentRequest.value()); - mRemainSendCount.erase(mCurrentRequest.value()); + mCancelledRequests.erase(cancelledReqId); + mRemainSendCount.erase(cancelledReqId); } + // Intentionally do NOT erase mInFlightCancelFlags[cancelledReqId] + // here — see removeResponse for rationale. The session for + // this reqId remains in mRequestToSession (release() is not + // called on this path), so its DataContext still references + // the atomic. The flag shared_ptr is reclaimed at ~Impl. mCurrentRequest = std::nullopt; if (mReadyResponses.empty()) @@ -599,6 +687,9 @@ class CacheSender::Impl std::unique_lock lk(mCondMutex); mAnyReady = false; } + cancelledResponse.mPromise.set_exception(std::make_exception_ptr( + TLLM_REQUEST_EXCEPTION(cancelledReqId, common::RequestErrorCode::kNETWORK_ERROR, + "KV cache transfer for request %zu was cancelled", cancelledReqId))); } } mCurrentRequest = std::nullopt; @@ -670,6 +761,31 @@ class CacheSender::Impl it.second.mPromise.set_exception(std::current_exception()); } } + catch (...) + { + // Non-std::exception escape (integer throw, custom type throw + // from NIXL/UCX backends, C++ ABI edge cases). response() is + // noexcept, so an uncaught non-std throw would call + // std::terminate. Catch here to resolve pending promises with + // the exception so callers see a failure instead of the + // process aborting; the symmetric catch is in + // CacheReceiver::Impl::request(). The worker thread still + // exits after this catch — sender is then dead for this + // process, but fail-closed via the promises is strictly + // better than terminate(). + TLLM_LOG_ERROR("[CacheSender] UNKNOWN (non-std::exception) escape in response() — worker exiting"); + for (auto& it : mReadyResponses) + { + try + { + it.second.mPromise.set_exception(std::current_exception()); + } + catch (...) + { + // promise already satisfied + } + } + } } void terminate() @@ -678,6 +794,19 @@ class CacheSender::Impl std::unique_lock lk(mCondMutex); mTerminate = true; } + // Flip every registered per-request cancel flag so any AgentConnection::send + // currently in its poll-wait observes shutdown via the same atomic it + // polls for per-request cancellation. Needed because recvRequestInfo's + // TransferSession DataContext references per-request flags (not + // mTerminate), so without this, shutdown would not interrupt an + // in-flight sender wedge. + { + std::lock_guard lg(mInFlightCancelMutex); + for (auto& [id, flag] : mInFlightCancelFlags) + { + flag->store(true); + } + } // We don't have to wait for the future. If another thread is sending data, it won't pay attention // to the terminate flag. mSenderCv.notify_all(); @@ -699,6 +828,13 @@ class CacheSender::Impl std::scoped_lock lkResp(mSenderMutex); mReadyResponses.erase(it); } + // Flag erase is intentionally NOT here. The session's DataContext + // references the atomic held by mInFlightCancelFlags; dropping the + // map entry while the session still exists in mRequestToSession + // would dangle that reference. Flag lifetime is tied to session + // lifetime — reclaimed in release() (called by + // sendAndRemoveResponse's success path after sendSync completes) + // or at ~Impl for leaked sessions on error paths. if (mReadyResponses.empty()) { std::unique_lock lkCond(mCondMutex); @@ -737,6 +873,13 @@ class CacheSender::Impl std::mutex mMtxForMap; runtime::BufferManager mBufferManager; std::ofstream mMeasuresFile; + // Per-request cancel-flag registry (sender-side parity with + // CacheReceiver::Impl). Registered at sendAsync time, referenced by + // the TransferSession's DataContext in recvRequestInfo, flipped by + // cancelRequest on the non-queue-drainable case, and erased after + // removeResponse / cancel-cleanup. ~terminate flips all for shutdown. + std::mutex mInFlightCancelMutex; + std::unordered_map>> mInFlightCancelFlags; }; class CacheReceiver::Impl @@ -753,23 +896,27 @@ class CacheReceiver::Impl TLLM_CUDA_CHECK(cudaGetDevice(&mDeviceId)); } - [[nodiscard]] std::future receiveAsync(LlmRequest& llmRequest) + [[nodiscard]] std::future receiveAsync(std::shared_ptr const& llmRequest) { // TODO: Modify the implementation here to avoid frequent thread creation. - return std::async(std::launch::async, &CacheReceiver::Impl::requestSync, this, std::ref(llmRequest)); + // Lambda captures the shared_ptr so the request is kept alive until the + // async task completes — closes the raw-pointer UAF that the old + // `[this, &llmRequest]` capture was vulnerable to. + auto llmRequestCopy = llmRequest; + return std::async(std::launch::async, [this, llmRequestCopy]() { requestSync(*llmRequestCopy); }); } - [[nodiscard]] std::future requestAndReceiveAsyncMultiThreads(LlmRequest& llmRequest) + [[nodiscard]] std::future requestAndReceiveAsyncMultiThreads(std::shared_ptr const& llmRequest) { try { auto promise = std::make_unique>(); auto future = promise->get_future(); - TLLM_CHECK(llmRequest.getDataTransceiverState().getCommState().has_value()); + TLLM_CHECK(llmRequest->getDataTransceiverState().getCommState().has_value()); std::string processInfo = kDefaultProcessInfo; if (common::getEnvRequestKVCacheConcurrent()) { - processInfo = llmRequest.getDataTransceiverState().getCommState()->toString(); + processInfo = llmRequest->getDataTransceiverState().getCommState()->toString(); } if (mInstanceToAsyncResource.find(processInfo) == mInstanceToAsyncResource.end()) { @@ -780,9 +927,25 @@ class CacheReceiver::Impl mRequestFutures.emplace_back(std::move(requestFuture)); } auto& asyncResource = mInstanceToAsyncResource.at(processInfo); + // Register a per-request cancel flag so cancelRequest() can abort + // the receive even after it has been dequeued and is blocked inside + // requestSync (e.g. waiting on recvReadySignal for a peer that + // will never respond — the ghost-UCX scenario described in the + // gen-side no-recovery investigation). The shared_ptr keeps the + // atomic alive for the duration of both the RequestAndPromise + // (held on the queue / in the worker's local) and any downstream + // DataContext references. + auto cancelFlag = std::make_shared>(false); + { + std::lock_guard lg(mInFlightCancelMutex); + mInFlightCancelFlags[llmRequest->mRequestId] = cancelFlag; + } { std::unique_lock lck(asyncResource->mMtxForQueue); - asyncResource->mRequestsQueue.emplace_back(std::addressof(llmRequest), std::move(promise)); + // Worker holds shared_ptr so Python-side _terminate_request + // cannot drop the LlmRequest out from under the worker's + // dereferences — closes the raw-pointer UAF. + asyncResource->mRequestsQueue.emplace_back(llmRequest, std::move(promise), cancelFlag); } asyncResource->mCVforQueue.notify_all(); return future; @@ -810,7 +973,29 @@ class CacheReceiver::Impl } } + // Overload kept for the public CacheReceiver::sendRequestInfo wrapper + // and any direct caller without a per-request cancel flag. Threads + // mTerminate through so process shutdown still interrupts the poll loop. TransferSession sendRequestInfo(LlmRequest const& llmRequest) + { + return sendRequestInfo(llmRequest, mTerminate); + } + + TransferSession sendRequestInfo(LlmRequest const& llmRequest, std::atomic const& perRequestCancel) + { + // Legacy / public-wrapper path: no pre-acquired ids; acquires + // internally with no RAII protection. Not on the drain-worker path. + return sendRequestInfo(llmRequest, perRequestCancel, /*preAcquiredCacheBufferIds=*/{}); + } + + // Overload that takes caller-acquired buffer indices (wrapped in + // BufferIndexHolders at the call site). requestSync uses this variant so + // RAII covers all non-happy-path exits (not-ready, cancel, throw). The + // preAcquiredCacheBufferIds vector must have one entry per cache + // transfer buffer manager (aligned with + // AgentConnectionManager::getCacheTransBufferManagers()). + TransferSession sendRequestInfo(LlmRequest const& llmRequest, std::atomic const& perRequestCancel, + std::vector> preAcquiredCacheBufferIds) { uint64_t requestId = llmRequest.getContextPhaseParams().value().getReqId(); auto const& contextState = llmRequest.getDataTransceiverState(); @@ -849,11 +1034,28 @@ class CacheReceiver::Impl std::vector> cacheBufferIds; if (agentConnectionManager) { - for (auto& cacheTransBufferManager : agentConnectionManager->getCacheTransBufferManagers()) + if (!preAcquiredCacheBufferIds.empty()) + { + // requestSync already acquired these indices under RAII + // holders — use them as-is so receiveSync's formatter + // releases via the existing path while any non-happy-path + // exit from requestSync releases via ~BufferIndexHolder. + cacheBufferIds = std::move(preAcquiredCacheBufferIds); + TLLM_CHECK(!cacheBufferIds.empty()); + } + else { - cacheBufferIds.push_back(cacheTransBufferManager->assignBufferIndexForRecv()); + // Legacy path: no RAII protection. Acquire internally, + // matching pre-RAII behavior for callers that don't go + // through requestSync. + auto const reqIdForLog = std::make_optional(static_cast(llmRequest.mRequestId)); + for (auto& cacheTransBufferManager : agentConnectionManager->getCacheTransBufferManagers()) + { + cacheBufferIds.push_back(cacheTransBufferManager->assignBufferIndexForRecv( + &perRequestCancel, kBufferAcquireSliceMs, reqIdForLog)); + } + TLLM_CHECK(!cacheBufferIds.empty()); } - TLLM_CHECK(!cacheBufferIds.empty()); } auto allCounterparts @@ -882,6 +1084,17 @@ class CacheReceiver::Impl for (size_t ci = 0; ci < allCounterparts.size(); ci++) { + // Honor perRequestCancel between per-peer notify iterations. If + // the drain worker's cancelRequest has fired (e.g. the gen-side + // kv_transfer_timeout_ms hit), bail out of this loop instead of + // continuing to notify peers that have already abandoned their + // side of the transfer. Without this check the for-body can + // block on notifySyncMessage to a stuck peer without any + // opportunity to observe the cancel flag. + if (perRequestCancel.load(std::memory_order_relaxed)) + { + TLLM_THROW("sendRequestInfo cancelled via perRequestCancel for request %zu", llmRequest.mRequestId); + } auto rank = allCounterparts[ci]; auto const* connection = connections.at(rank); @@ -932,15 +1145,20 @@ class CacheReceiver::Impl TLLM_CHECK(agentConnection != nullptr); const_cast(agentConnection) - ->sendRequestAndBufferInfo(requestInfo, idsForRank, validConnectionIdx); + ->sendRequestAndBufferInfo(requestInfo, idsForRank, validConnectionIdx, &perRequestCancel); } else { - sendRequestInfo(connection, requestInfo); + sendRequestInfo(connection, requestInfo, perRequestCancel); } } auto const& resource = getReceiveCacheResource(llmRequest); - return TransferSession(std::move(allConnections), DataContext{tagFromRequestId(requestId), mTerminate}, + // The TransferSession's DataContext is used by downstream data-phase + // operations (receiveSync -> unformat -> per-connection send/recv). + // Wire perRequestCancel through so those calls observe the same + // cancel flag the per-iteration loops here honor. AgentConnection::send + // reads ctx.getTransferTerminate() in its wait-poll loop. + return TransferSession(std::move(allConnections), DataContext{tagFromRequestId(requestId), perRequestCancel}, std::move(allCounterparts), mSelfState, contextState, resource->mBufferManager, requestInfo.getIndexFromEnd(), requestInfo.getLastBlockKey(), &llmRequest, !common::getEnvKVCacheTimeOutputPath().empty()); @@ -964,16 +1182,21 @@ class CacheReceiver::Impl return mProcessToResources.at(processString); } - void sendRequestInfo(executor::kv_cache::Connection const* connection, RequestInfo const& info) + void sendRequestInfo(executor::kv_cache::Connection const* connection, RequestInfo const& info, + std::atomic const& perRequestCancel) { std::ostringstream oss; RequestInfo::serialize(info, oss); auto const& serializedInfo = oss.str(); std::size_t const infoSize = serializedInfo.size(); TransceiverTag::Id id{TransceiverTag::Id::REQUEST_SEND}; - connection->send(DataContext{TransceiverTag::kID_TAG}, &id, sizeof(id)); - connection->send(DataContext{TransceiverTag::kINFO_SIZE_TAG}, &infoSize, sizeof(infoSize)); - connection->send(DataContext{TransceiverTag::kINFO_TAG}, serializedInfo.data(), infoSize); + // Propagate perRequestCancel via each DataContext so the underlying + // connection->send implementation (e.g. AgentConnection::send's + // poll-wait on submitted transfer) can observe a cancel fired + // asynchronously by CacheReceiver::cancelRequest. + connection->send(DataContext{TransceiverTag::kID_TAG, perRequestCancel}, &id, sizeof(id)); + connection->send(DataContext{TransceiverTag::kINFO_SIZE_TAG, perRequestCancel}, &infoSize, sizeof(infoSize)); + connection->send(DataContext{TransceiverTag::kINFO_TAG, perRequestCancel}, serializedInfo.data(), infoSize); } bool cancelRequest(LlmRequest const& llmRequest) @@ -986,6 +1209,7 @@ class CacheReceiver::Impl } bool isCancelled = false; + std::optional queuedCancelledReqId; auto& asyncResource = mInstanceToAsyncResource.at(processInfo); { std::unique_lock lck(asyncResource->mMtxForQueue); @@ -994,9 +1218,57 @@ class CacheReceiver::Impl { return requestAndPromise.mRequest->mRequestId == llmRequest.mRequestId; }); if (it != asyncResource->mRequestsQueue.end()) { + // Fulfil the queued promise with a structured cancellation + // exception before erasing the entry. Without this, dropping the + // unique_ptr> destroys it unfulfilled, and any + // consumer awaiting the corresponding future via + // mRequesterFutures observes std::future_error: Broken promise + // instead of a clean per-request kNETWORK_ERROR. + queuedCancelledReqId = it->mRequest->mRequestId; + if (it->mPromise) + { + try + { + it->mPromise->set_exception(std::make_exception_ptr( + TLLM_REQUEST_EXCEPTION(*queuedCancelledReqId, common::RequestErrorCode::kNETWORK_ERROR, + "Generation KV cache request cancelled before send for request %zu", + *queuedCancelledReqId))); + } + catch (std::future_error const&) + { + // Promise already satisfied (or no associated future); + // nothing else to do. + } + } asyncResource->mRequestsQueue.erase(it); isCancelled = true; } + } + if (queuedCancelledReqId.has_value()) + { + // The worker will never dequeue this entry, so remove the cancel + // flag here instead of relying on the normal request() cleanup. + std::lock_guard lg(mInFlightCancelMutex); + mInFlightCancelFlags.erase(*queuedCancelledReqId); + } + if (!isCancelled) + { + // Request already dequeued past mRequestsQueue (worker thread has + // picked it up; most likely blocked inside requestSync -> + // recvReadySignal waiting on a peer). Flip the per-request cancel + // flag so the notification polling loop in + // AgentConnectionManager::waitForNotification observes it and + // returns early with isReady=false. requestSync then falls into + // the kDISAGG_TRANS_ERROR branch and the future resolves, freeing + // the NIXL/UCX per-request state instead of leaking it. + std::lock_guard lg(mInFlightCancelMutex); + auto flagIt = mInFlightCancelFlags.find(llmRequest.mRequestId); + if (flagIt != mInFlightCancelFlags.end()) + { + flagIt->second->store(true); + isCancelled = true; + TLLM_LOG_DEBUG("Flipped in-flight cancel flag for request %zu (not in queue).", llmRequest.mRequestId); + } else { TLLM_LOG_WARNING("Cannot cancel request %zu", llmRequest.mRequestId); @@ -1005,7 +1277,7 @@ class CacheReceiver::Impl return isCancelled; } - bool receiveReadySignal(TransferSession& session) + bool receiveReadySignal(TransferSession& session, std::atomic const& perRequestCancel) { bool isReadyFinal = true; bool isReady = false; @@ -1013,13 +1285,26 @@ class CacheReceiver::Impl for (size_t i = 0; i < connections.size(); i++) { + // Honor perRequestCancel between per-peer ready-signal waits so + // a cancel fired during the multi-rank wait sequence bails out + // rather than continuing to wait on subsequent peers. + if (perRequestCancel.load(std::memory_order_relaxed)) + { + return false; + } auto* agentConnectionManager = dynamic_cast(mManager); if (agentConnectionManager) { auto* agentConnection = dynamic_cast(connections.at(i)); TLLM_CHECK(agentConnection); + // Pass the per-request cancel flag as the DataContext's + // transferTerminate. The notification polling loop in + // AgentConnectionManager::waitForNotification checks this + // atomic and returns early when it flips — either on + // process shutdown (all per-request flags flipped in + // ~Impl()) or on per-request cancelRequest(). isReady = agentConnection->recvReadySignal( - executor::kv_cache::DataContext{TransceiverTag::kREADY_SIGNAL_TAG, mTerminate}); + executor::kv_cache::DataContext{TransceiverTag::kREADY_SIGNAL_TAG, perRequestCancel}); } else { @@ -1032,9 +1317,28 @@ class CacheReceiver::Impl return isReadyFinal; } + // Overload preserved for the (currently unused) std::async-based receiveAsync + // path and for any callers that don't plumb a per-request cancel flag. Uses + // the process-level mTerminate only. + bool receiveReadySignal(TransferSession& session) + { + return receiveReadySignal(session, mTerminate); + } + ~Impl() { mTerminate.store(true); + // Flip every per-request cancel flag so recvReadySignal's polling + // loop can observe shutdown via the same atomic it uses for + // per-request cancellation. The shared_ptr aliasing in DataContext + // keeps these alive for any still-in-flight recvReadySignal calls. + { + std::lock_guard lg(mInFlightCancelMutex); + for (auto& [id, flag] : mInFlightCancelFlags) + { + flag->store(true); + } + } for (auto&& [processInfo, asyncResource] : mInstanceToAsyncResource) { asyncResource->mTerminate = true; @@ -1047,19 +1351,70 @@ class CacheReceiver::Impl } private: - void requestSync(LlmRequest& llmRequest) + void requestSync(LlmRequest& llmRequest, std::atomic const& perRequestCancel) { TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "Start calling requestSync for request ID: %zu, context request ID: %zu.", llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId()); llmRequest.setKvCacheTransferStart(std::chrono::steady_clock::now()); + // Early-out if cancel fired between queueing and dequeue. + if (perRequestCancel.load() || mTerminate.load()) + { + llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); + return; + } TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); - auto session = sendRequestInfo(llmRequest); + + // Pre-acquire receive-buffer indices under RAII holders BEFORE entering + // sendRequestInfo. This is the core of the RAII fix: the formatter + // inside receiveSync releases the indices on the happy path (hence + // the `detach()` after a successful receiveSync below), but every + // OTHER exit path from requestSync used to leak one index per exit + // (e.g. an early return on `(not-ready or cancel after ready)` would + // permanently wedge the size-1 pool). With the holder, any return or + // throw between acquisition and the final detach releases the index + // via the holder's destructor, closing the class of bug rather than + // any specific branch. + std::vector recvHolders; + std::vector> cacheBufferIds; + auto* agentConnectionManagerForAcq = dynamic_cast(mManager); + if (agentConnectionManagerForAcq) + { + auto const reqIdForLog = std::make_optional(static_cast(llmRequest.mRequestId)); + auto const& managers = agentConnectionManagerForAcq->getCacheTransBufferManagers(); + recvHolders.reserve(managers.size()); + cacheBufferIds.reserve(managers.size()); + for (auto* cacheTransBufferManager : managers) + { + auto rawIdx = cacheTransBufferManager->assignBufferIndexForRecv( + &perRequestCancel, kBufferAcquireSliceMs, reqIdForLog); + recvHolders.emplace_back(*cacheTransBufferManager, rawIdx, /*isRecv=*/true, reqIdForLog); + if (rawIdx.has_value()) + { + cacheBufferIds.push_back(static_cast(rawIdx.value())); + } + else + { + cacheBufferIds.push_back(std::nullopt); + } + } + } + + auto session = sendRequestInfo(llmRequest, perRequestCancel, std::move(cacheBufferIds)); session.setTime(TransferSession::kTimeRequestInfo); - bool isReady = receiveReadySignal(session); - if (!isReady) + // receiveReadySignal blocks inside AgentConnectionManager::waitForNotification's + // polling loop until the peer sends the ready notification OR the + // perRequestCancel flag flips. That's the one path that can actually + // release a zombie receive when a timeout-eviction fires on the C++ + // side; without it the std::async/worker task would stay blocked + // indefinitely and hold NIXL/UCX per-request state. + bool isReady = receiveReadySignal(session, perRequestCancel); + if (!isReady || perRequestCancel.load()) { - // Reuse the error state for the cancelled request. + // Either the peer never sent the ready signal (timeout / cancel + // fired) or the cancel was observed after ready. Either way, + // reuse the error state — AsyncTransferManager cleanup will run. llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); return; @@ -1067,54 +1422,65 @@ class CacheReceiver::Impl receiveSync(session); llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); + // Happy path only: the formatter invoked inside receiveSync already + // released each buffer index via freeBufferIndexForRecv. Detach the + // holders so they don't double-release when this stack frame + // unwinds. If control is about to leave requestSync via any OTHER + // path below (none today, but safe against future edits), the + // holders would still release automatically — detaching is strictly + // a correctness requirement for the just-completed receiveSync path. + for (auto& h : recvHolders) + { + (void) h.detach(); + } + TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "End calling requestSync for request ID: %zu, context request ID: %zu.", llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId()); } + // Overload used by the (currently unused) std::async-based + // Impl::receiveAsync path, where no per-request cancel flag is plumbed. + // Threads the process-wide mTerminate through the same polling path so + // shutdown still interrupts it. + void requestSync(LlmRequest& llmRequest) + { + requestSync(llmRequest, mTerminate); + } + struct RequestAndPromise { - LlmRequest* mRequest; + // Store shared_ptr rather than a raw pointer so the async worker's + // dereferences stay safe even after Python's _terminate_request drops + // its own pybind shared_ptr. See CacheTransceiver::mSenderFutures for + // the full lifetime invariant. + std::shared_ptr mRequest; std::unique_ptr> mPromise; + // Per-request cancel flag. Flipped by CacheReceiver::Impl::cancelRequest + // when a timeout-eviction (in checkGenTransferStatus) or similar + // wants to abort an in-flight receive that has already been dequeued + // past the mRequestsQueue stage. + std::shared_ptr> mCancelFlag; RequestAndPromise() : mRequest(nullptr) , mPromise(nullptr) + , mCancelFlag(nullptr) { } - RequestAndPromise(LlmRequest* request, std::unique_ptr>&& promise) - : mRequest(request) + RequestAndPromise(std::shared_ptr request, std::unique_ptr>&& promise, + std::shared_ptr> cancelFlag) + : mRequest(std::move(request)) , mPromise(std::move(promise)) + , mCancelFlag(std::move(cancelFlag)) { } RequestAndPromise(RequestAndPromise const&) = delete; - RequestAndPromise(RequestAndPromise&& other) noexcept - : mRequest(other.mRequest) - , mPromise(std::move(other.mPromise)) - { - other.mRequest = nullptr; - } - - RequestAndPromise& operator=(RequestAndPromise&& other) noexcept - { - if (this != &other) - { - mRequest = nullptr; - if (mPromise) - { - mPromise.reset(); - } - - mRequest = other.mRequest; - mPromise = std::move(other.mPromise); - - other.mRequest = nullptr; - } - return *this; - } + RequestAndPromise(RequestAndPromise&& other) noexcept = default; + RequestAndPromise& operator=(RequestAndPromise&& other) noexcept = default; }; struct AsyncResource @@ -1151,11 +1517,15 @@ class CacheReceiver::Impl requestAndPromise = std::move(resource.mRequestsQueue.front()); resource.mRequestsQueue.pop_front(); } + auto const reqId = requestAndPromise.mRequest != nullptr ? requestAndPromise.mRequest->mRequestId + : static_cast(0); { try { TLLM_CHECK_WITH_INFO(requestAndPromise.mRequest != nullptr, "requestAndPromise.mRequest is null"); - requestSync(*requestAndPromise.mRequest); + TLLM_CHECK_WITH_INFO( + requestAndPromise.mCancelFlag != nullptr, "requestAndPromise.mCancelFlag is null"); + requestSync(*requestAndPromise.mRequest, *requestAndPromise.mCancelFlag); requestAndPromise.mPromise->set_value(); } catch (tensorrt_llm::common::RequestSpecificException const& err) @@ -1174,6 +1544,44 @@ class CacheReceiver::Impl requestAndPromise.mRequest->getContextPhaseParams().value().getReqId(), err.what()); requestAndPromise.mPromise->set_exception(std::current_exception()); } + catch (...) + { + // Non-std::exception escapes (e.g. throws of integer / + // custom types from NIXL / UCX backends, or the rare C++ + // ABI abort path) would otherwise kill this worker thread + // and leave the mRequestsQueue unserviced forever. When + // that happens, the in-flight cancel flag registry keeps + // inserting entries, cancelRequest keeps satisfying them + // via the queue-removal branch and silently drops them, + // and the in-flight cancel-flag branch is never reached — + // producing a post-saturation no-recovery state where + // every new receive just waits kv_transfer_timeout_ms. + // Swallow and continue so the drain loop remains alive; + // set the promise so the caller's future resolves with an + // error rather than hanging. + TLLM_LOG_ERROR( + "Non-std::exception escape in CacheReceiver request() loop for reqId=%zu; continuing.", reqId); + if (requestAndPromise.mPromise) + { + try + { + requestAndPromise.mPromise->set_exception(std::current_exception()); + } + catch (...) + { + // promise already satisfied; nothing else to do + } + } + } + // Deregister the per-request cancel flag regardless of + // success / exception. The shared_ptr in requestAndPromise + // still keeps the atomic alive for any in-flight DataContext + // references that haven't unwound yet. + if (requestAndPromise.mRequest != nullptr) + { + std::lock_guard lg(mInFlightCancelMutex); + mInFlightCancelFlags.erase(requestAndPromise.mRequest->mRequestId); + } } } } @@ -1191,6 +1599,11 @@ class CacheReceiver::Impl std::ofstream mMeasuresFile; std::mutex mMeasuresFileMutex; std::atomic mTerminate{false}; + // Per-request cancel flags for in-flight receives that have been dequeued + // past mRequestsQueue. Registered on enqueue, looked up by cancelRequest + // for in-flight cancellation, and unregistered after requestSync returns. + std::mutex mInFlightCancelMutex; + std::unordered_map>> mInFlightCancelFlags; }; void CacheSender::ImplDeleter::operator()(Impl* ptr) @@ -1209,7 +1622,7 @@ CacheSender::CacheSender( { } -std::future CacheSender::sendAsync(LlmRequest& llmRequest) const +std::future CacheSender::sendAsync(std::shared_ptr const& llmRequest) const { return mImpl->sendAsync(llmRequest); } @@ -1252,7 +1665,7 @@ CacheReceiver::CacheReceiver( { } -std::future CacheReceiver::receiveAsync(LlmRequest& llmRequest) const +std::future CacheReceiver::receiveAsync(std::shared_ptr const& llmRequest) const { return mImpl->requestAndReceiveAsyncMultiThreads(llmRequest); } diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h index 4d072a04521b..aa4849495cc7 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h @@ -257,9 +257,12 @@ class CacheSender /// @brief Asynchronously respond to the request and send data. /// @param llmRequest Request object. Its data should be ready when called, and the data for this request - /// should remain valid until future synchronization. + /// should remain valid until future synchronization. Passed as shared_ptr + /// so the async worker keeps the LlmRequest alive past any Python-side + /// _terminate_request; see CacheTransceiver::mSenderFutures for the full + /// lifetime invariant. /// @return Once the data is fully sent, the future object will become valid. - [[nodiscard]] virtual std::future sendAsync(LlmRequest& llmRequest) const; + [[nodiscard]] virtual std::future sendAsync(std::shared_ptr const& llmRequest) const; /// @brief Return the internal communicator status. /// @return The communicator status. @@ -314,9 +317,12 @@ class CacheReceiver /// @brief Asynchronously send a request to receive data. /// @param llmRequest Request object. Its data should be in an allocated but unwritten state when called, and the - /// data for this request should remain intact only after future synchronization. + /// data for this request should remain intact only after future synchronization. Passed as shared_ptr so the + /// async worker keeps the LlmRequest alive past any Python-side + /// _terminate_request; see CacheTransceiver::mSenderFutures for the full + /// lifetime invariant. /// @return Once the data is fully received, the future object will become valid. - [[nodiscard]] virtual std::future receiveAsync(LlmRequest& llmRequest) const; + [[nodiscard]] virtual std::future receiveAsync(std::shared_ptr const& llmRequest) const; virtual TransferSession sendRequestInfo(LlmRequest const& llmRequest); diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp index c72090867f29..47fc8d400c10 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp @@ -30,6 +30,7 @@ #include "tensorrt_llm/runtime/cudaEvent.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" +#include #include #include #include @@ -253,7 +254,17 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses return bufferSizeForTarget; }; auto bufferEleSizes = getBufferSizeForTarget(); - auto cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForSend(); + // RAII wrapper mirrors the receiver-side pattern. Happy-path calls + // sendHolder.release(); other exits auto-release via + // ~BufferIndexHolder. Send-pool CV wait observes the per-request + // cancel flag via the session's DataContext and throws on cancel + // (parity with recv side). + auto const sendReqIdForLog = std::make_optional(static_cast(llmRequest.mRequestId)); + auto const* sendCancelFlag = &session.getDataContext().getTransferTerminate(); + auto cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForSend( + sendCancelFlag, kBufferAcquireSliceMs, sendReqIdForLog); + BufferIndexHolder sendHolder( + *mCacheTransBufferManagers[transferIndexerKCache], cacheBufferId, /*isRecv=*/false, sendReqIdForLog); auto result = mCacheTransBufferManagers[transferIndexerKCache]->getOrAllocateSendBuffers( cacheBufferId, static_cast(pPDomainSize * cPDomainSize), bufferEleSizes, bufferManager); auto& outputSplitCaches = std::get<0>(result); @@ -380,7 +391,8 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses { sendBufferFun(deviceId, pickUpConnections[0]); } - mCacheTransBufferManagers[transferIndexerKCache]->freeBufferIndexForSend(cacheBufferId); + // Atomic happy-path release — silent free + disarm in one noexcept call. + sendHolder.release(); } session.setTime(TransferSession::kTimeTransmissions); session.setTime(TransferSession::kTimePostprocess); diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index 8bb2c0e2ba88..94679a54f022 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -915,7 +915,7 @@ void TrtGptModelInflightBatching::forwardSync() TLLM_CHECK_WITH_INFO(mCacheTransceiver, "Disaggregated serving is not enabled, please check the configuration of " "cacheTransceiverConfig."); - mCacheTransceiver->respondAndSendAsync(llmReq.get()); + mCacheTransceiver->respondAndSendAsync(llmReq); } mSeqSlotManager->freeSequenceSlot(llmReq->mRequestId); } @@ -1596,11 +1596,11 @@ void TrtGptModelInflightBatching::prepareDisaggGenInitRequests( mCacheTransceiver, "Disaggregated serving is not enabled, please check the configuration."); if (common::getEnvDisableKVCacheTransferOverlap()) { - mCacheTransceiver->requestAndReceiveSync(newGenReq.get()); + mCacheTransceiver->requestAndReceiveSync(newGenReq); } else { - mCacheTransceiver->requestAndReceiveAsync(newGenReq.get()); + mCacheTransceiver->requestAndReceiveAsync(newGenReq); } } if (!common::getEnvDisableKVCacheTransferOverlap()) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp index d46defdf50ad..2d52b34fd6aa 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp @@ -143,8 +143,48 @@ void AgentConnection::send(DataContext const& ctx, void const* data, size_t size NotificationInfo notificationInfo{syncInfo}; std::stringstream ss; NotificationInfo::serialize(notificationInfo, ss); - TransferState transferState = status->wait(); + // Poll in bounded chunks so we can observe ctx.getTransferTerminate() + // (which may be a per-request cancel flag plumbed in from the drain + // worker's requestSync) between polls. Without this, a stuck peer + // would block submitTransferRequests's completion indefinitely and + // the corresponding kv_transfer_timeout_ms eviction in + // CacheTransceiver::checkGenTransferStatus would never recover. + static constexpr int64_t kCancelPollTimeoutMs = 100; + TransferState transferState = TransferState::kIN_PROGRESS; + while (transferState == TransferState::kIN_PROGRESS) + { + transferState = status->wait(kCancelPollTimeoutMs); + if (transferState == TransferState::kIN_PROGRESS && ctx.getTransferTerminate().load(std::memory_order_relaxed)) + { + // Cancel observed mid-send. Release the underlying backend + // transfer handle before unwinding so NIXL does not retain + // stale in-flight state below TRT-LLM's request cleanup. + bool const released = status->release(); + TLLM_LOG_WARNING( + "AgentConnection::send cancelled while transfer was in progress (ctx tag=%d, remote=%s, " + "releaseAccepted=%d)", + ctx.getTag(), mRemoteAgentName.c_str(), released); + TLLM_CHECK_WITH_INFO( + released, "AgentConnection::send cancel could not release the backend transfer handle"); + TLLM_THROW("AgentConnection::send cancelled mid-transfer"); + } + } TLLM_CHECK_WITH_INFO(transferState == TransferState::kSUCCESS, "AgentConnection::send failed"); + // One more cancel check between the successful wait and the (potentially + // blocking) notifySyncMessage. notifySyncMessage goes through the NIXL + // backend and is not interruptible from this layer; if cancel fires in + // the gap between wait() returning and notify starting, bail out now. + if (ctx.getTransferTerminate().load(std::memory_order_relaxed)) + { + bool const released = status->release(); + TLLM_LOG_WARNING( + "AgentConnection::send cancelled after transfer completed but before notify (ctx tag=%d, remote=%s, " + "releaseAccepted=%d)", + ctx.getTag(), mRemoteAgentName.c_str(), released); + TLLM_CHECK_WITH_INFO( + released, "AgentConnection::send pre-notify cancel could not release the backend transfer handle"); + TLLM_THROW("AgentConnection::send cancelled pre-notify"); + } // TODO: there is a bug in request_with_notify https://github.com/ai-dynamo/nixl/pull/252 mAgentConnectionManager->getAgent()->notifySyncMessage(mRemoteAgentName, ss.str()); } @@ -153,11 +193,16 @@ void AgentConnection::recv(DataContext const& ctx, void* data, size_t size) cons { NotificationSyncInfo syncInfo{mAgentName, ctx}; - mAgentConnectionManager->waitForSyncInfo(mRemoteAgentName, syncInfo, ctx.getTransferTerminate()); + bool const received + = mAgentConnectionManager->waitForSyncInfo(mRemoteAgentName, syncInfo, ctx.getTransferTerminate()); + TLLM_CHECK_WITH_INFO(received, + "AgentConnection::recv ended before receiving sync notification (ctx tag=%d, remote=%s)", ctx.getTag(), + mRemoteAgentName.c_str()); } void AgentConnection::sendRequestAndBufferInfo(batch_manager::RequestInfo& requestInfo, - std::vector> const& cacheBufferIds, int connectionIdx) + std::vector> const& cacheBufferIds, int connectionIdx, + std::atomic const* perRequestCancel) { TLLM_CHECK(!common::getEnvTryZCopyForKVCacheTransfer()); @@ -209,6 +254,18 @@ void AgentConnection::sendRequestAndBufferInfo(batch_manager::RequestInfo& reque std::stringstream ss; NotificationInfo notificationInfo{requestAndBufferInfo}; NotificationInfo::serialize(notificationInfo, ss); + // Pre-notify cancel check: notifySyncMessage enters the NIXL backend and + // is not interruptible from this layer. The caller's for-loop in + // CacheReceiver::Impl::sendRequestInfo also checks perRequestCancel + // between iterations; this check narrows the window where a cancel + // fired right before this peer's notify is suppressed. + if (perRequestCancel != nullptr && perRequestCancel->load(std::memory_order_relaxed)) + { + TLLM_LOG_WARNING( + "AgentConnection::sendRequestAndBufferInfo cancelled via perRequestCancel before notify (remote=%s)", + mRemoteAgentName.c_str()); + TLLM_THROW("sendRequestAndBufferInfo cancelled pre-notify"); + } mAgentConnectionManager->getAgent()->notifySyncMessage(mRemoteAgentName, ss.str()); } @@ -247,7 +304,10 @@ void AgentConnection::sendReadySignal(DataContext const& ctx, bool isReady) cons bool AgentConnection::recvReadySignal(DataContext const& ctx) const { ReadySignalInfo readySignalInfo{mAgentName, ctx, false}; - mAgentConnectionManager->waitForReadySignal(mRemoteAgentName, readySignalInfo, ctx.getTransferTerminate()); + if (!mAgentConnectionManager->waitForReadySignal(mRemoteAgentName, readySignalInfo, ctx.getTransferTerminate())) + { + return false; + } return readySignalInfo.mIsReady; } @@ -582,7 +642,7 @@ int AgentConnectionManager::getDeviceId() const } template -void AgentConnectionManager::waitForNotification( +bool AgentConnectionManager::waitForNotification( std::string const& remoteAgentName, NotificationType& expectedInfo, std::atomic const& terminateFlag) { while (!terminateFlag.load()) @@ -590,7 +650,7 @@ void AgentConnectionManager::waitForNotification( if (!mIsRunning) { - return; + return false; } updateUnhandledNotifications(); std::scoped_lock lock(mNotificationMutex); @@ -623,7 +683,7 @@ void AgentConnectionManager::waitForNotification( { it = mUnhandledNotifications.erase(it); } - return; + return true; } } } @@ -643,7 +703,7 @@ void AgentConnectionManager::waitForNotification( { it = mUnhandledNotifications.erase(it); } - return; + return true; } } } @@ -663,24 +723,25 @@ void AgentConnectionManager::waitForNotification( } } } + return false; } // Explicit template instantiations -template void AgentConnectionManager::waitForNotification( +template bool AgentConnectionManager::waitForNotification( std::string const& remoteAgentName, NotificationSyncInfo& expectedInfo, std::atomic const& terminateFlag); -template void AgentConnectionManager::waitForNotification( +template bool AgentConnectionManager::waitForNotification( std::string const& remoteAgentName, ReadySignalInfo& expectedInfo, std::atomic const& terminateFlag); -void AgentConnectionManager::waitForSyncInfo( +bool AgentConnectionManager::waitForSyncInfo( std::string const& remoteAgentName, NotificationSyncInfo& syncInfo, std::atomic const& terminateFlag) { - waitForNotification(remoteAgentName, syncInfo, terminateFlag); + return waitForNotification(remoteAgentName, syncInfo, terminateFlag); } -void AgentConnectionManager::waitForReadySignal( +bool AgentConnectionManager::waitForReadySignal( std::string const& remoteAgentName, ReadySignalInfo& readySignalInfo, std::atomic const& terminateFlag) { - waitForNotification(remoteAgentName, readySignalInfo, terminateFlag); + return waitForNotification(remoteAgentName, readySignalInfo, terminateFlag); } std::string const& AgentConnectionManager::getAgentName() const diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h index 8ec948cfafeb..e32f4440b096 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h @@ -259,8 +259,14 @@ class AgentConnection : public Connection std::string mAgentName, std::string mRemoteAgentName, AgentConnectionManager* mAgentConnectionManager); void send(DataContext const& ctx, void const* data, size_t size) const override; void recv(DataContext const& ctx, void* data, size_t size) const override; + // `perRequestCancel` is an optional pointer to a per-request cancel + // atomic (matches the flag registry in CacheReceiver::Impl). When + // non-null and flipped true before notifySyncMessage runs, the function + // throws so requestSync can unwind instead of blocking inside a NIXL + // backend notify call with no visibility into cancellation. void sendRequestAndBufferInfo(batch_manager::RequestInfo& requestInfo, - std::vector> const& cacheBufferIds, int validConnectionIdx); + std::vector> const& cacheBufferIds, int validConnectionIdx, + std::atomic const* perRequestCancel = nullptr); void setSenderState(std::vector cacheReceiverBufferDescs, int valideSegmentIdx, std::vector> offsetRatios, std::vector bufferKinds); void setHasLoadRemoteAgent(bool hasLoadRemoteAgent); @@ -320,11 +326,11 @@ class AgentConnectionManager : public ConnectionManager [[nodiscard]] std::string const& getAgentName() const; template - void waitForNotification( + [[nodiscard]] bool waitForNotification( std::string const& remoteAgentName, NotificationType& expectedInfo, std::atomic const& terminateFlag); - void waitForSyncInfo( + [[nodiscard]] bool waitForSyncInfo( std::string const& remoteAgentName, NotificationSyncInfo& syncInfo, std::atomic const& terminateFlag); - void waitForReadySignal( + [[nodiscard]] bool waitForReadySignal( std::string const& remoteAgentName, ReadySignalInfo& readySignalInfo, std::atomic const& terminateFlag); [[nodiscard]] bool isRunning() const override; diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp index 2f5eb9342bf6..f9157d59426e 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp @@ -175,7 +175,8 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m) // subclass type is not directly registered (e.g., agents created via factory). nb::class_(m, "TransferStatus") .def("is_completed", &kvc::TransferStatus::isCompleted, nb::call_guard()) - .def("wait", &kvc::TransferStatus::wait, nb::arg("timeout_ms") = -1, nb::call_guard()); + .def("wait", &kvc::TransferStatus::wait, nb::arg("timeout_ms") = -1, nb::call_guard()) + .def("release", &kvc::TransferStatus::release, nb::call_guard()); // BaseAgentConfig struct nb::class_(m, "BaseAgentConfig") @@ -228,7 +229,8 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m) nb::class_(m, "NixlTransferStatus") .def("is_completed", &kvc::NixlTransferStatus::isCompleted, nb::call_guard()) .def("wait", &kvc::NixlTransferStatus::wait, nb::arg("timeout_ms") = -1, - nb::call_guard()); + nb::call_guard()) + .def("release", &kvc::NixlTransferStatus::release, nb::call_guard()); // NixlTransferAgent class nb::class_(m, "NixlTransferAgent") diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp index bad3e184f983..23f061037353 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp @@ -324,6 +324,14 @@ NixlTransferStatus::NixlTransferStatus(nixlAgent* agent, nixlXferReqH* handle) TLLM_CHECK(mHandle); } +NixlTransferStatus::~NixlTransferStatus() +{ + if (!release()) + { + TLLM_LOG_WARNING("NIXL transfer handle release failed during destruction; backend handle may remain active"); + } +} + [[nodiscard]] MemoryDescs NixlHelper::coalesceMemoryDescs(MemoryDescs const& descs) { auto const& descVec = descs.getDescs(); @@ -484,6 +492,11 @@ NixlTransferStatus::NixlTransferStatus(nixlAgent* agent, nixlXferReqH* handle) TransferState NixlTransferStatus::wait(int64_t timeout_ms) const { + if (mHandle == nullptr) + { + return TransferState::kFAILURE; + } + auto startTime = std::chrono::steady_clock::now(); while (true) @@ -520,9 +533,31 @@ TransferState NixlTransferStatus::wait(int64_t timeout_ms) const [[nodiscard]] bool NixlTransferStatus::isCompleted() const { + if (mHandle == nullptr) + { + return false; + } return mRawAgent->getXferStatus(mHandle) == NIXL_SUCCESS; } +[[nodiscard]] bool NixlTransferStatus::release() +{ + if (mHandle == nullptr) + { + return true; + } + + auto status = mRawAgent->releaseXferReq(mHandle); + if (status == NIXL_SUCCESS) + { + mHandle = nullptr; + return true; + } + + TLLM_LOG_WARNING("NIXL releaseXferReq failed with status: %s", nixlEnumStrings::statusStr(status).c_str()); + return false; +} + [[nodiscard]] MemoryDescs NixlHelper::splitVmmDescs(MemoryDescs const& descs, size_t& detectedChunkSize) { detectedChunkSize = 0; diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h index b35c5deb182b..26504e1ecf11 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h @@ -63,11 +63,14 @@ class NixlTransferStatus final : public TransferStatus { public: NixlTransferStatus(nixlAgent* agent, nixlXferReqH* handle); + ~NixlTransferStatus() override; [[nodiscard]] bool isCompleted() const override; [[nodiscard]] TransferState wait(int64_t timeout_ms = -1) const override; + [[nodiscard]] bool release() override; + private: nixlAgent* mRawAgent{}; nixlXferReqH* mHandle{}; diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index 61cac5df4c7e..b0df67ba965a 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -46,17 +46,17 @@ class PyCacheTransceiver : public tb::BaseCacheTransceiver // using BaseCacheTransceiver::BaseCacheTransceiver; // Inherit constructors NB_TRAMPOLINE(tb::BaseCacheTransceiver, 6); - void respondAndSendAsync(tb::LlmRequest* llmRequest) override + void respondAndSendAsync(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(respondAndSendAsync, llmRequest); } - void requestAndReceiveSync(tb::LlmRequest* llmRequest) override + void requestAndReceiveSync(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(requestAndReceiveSync, llmRequest); } - void requestAndReceiveAsync(tb::LlmRequest* llmRequest) override + void requestAndReceiveAsync(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(requestAndReceiveAsync, llmRequest); } @@ -77,7 +77,7 @@ class PyCacheTransceiver : public tb::BaseCacheTransceiver NB_OVERRIDE_PURE(checkGenTransferComplete); } - bool cancelRequest(tb::LlmRequest* llmRequest) override + bool cancelRequest(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(cancelRequest, llmRequest); } diff --git a/tensorrt_llm/_torch/disaggregation/base/agent.py b/tensorrt_llm/_torch/disaggregation/base/agent.py index 1aac24d98e75..4c5548bd2472 100644 --- a/tensorrt_llm/_torch/disaggregation/base/agent.py +++ b/tensorrt_llm/_torch/disaggregation/base/agent.py @@ -89,6 +89,9 @@ def is_completed(self) -> bool: ... @abstractmethod def wait(self, timeout_ms: int | None = None) -> bool: ... + def release(self) -> bool: + return False + class BaseTransferAgent(ABC): @abstractmethod diff --git a/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py b/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py index 84f85d41b325..1dd85e177a13 100644 --- a/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py +++ b/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py @@ -30,6 +30,10 @@ def wait(self, timeout_ms=None) -> bool: timeout_ms = -1 return self._cpp_status.wait(timeout_ms) == TransferState.SUCCESS + def release(self) -> bool: + """Release the transfer handle, requesting backend cancel if still active.""" + return self._cpp_status.release() + class BindingsNixlTransferAgent(BaseTransferAgent): """NixlTransferAgent using C++ bindings with GIL release support. diff --git a/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py b/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py index dcfa28210f81..703dc42f33a1 100644 --- a/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py +++ b/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py @@ -26,6 +26,14 @@ def is_completed(self): status = TransferState(self.agent.check_xfer_state(self.handle)) return status == TransferState.DONE + def release(self): + try: + self.handle.release() + return True + except Exception: + logger.exception("Failed to release NIXL transfer handle (agent=%s).", self.agent.name) + return False + def wait(self, timeout_ms=None): start_time = time.time() status = TransferState.PENDING diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 98cf521d1ff4..420ff9bea535 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -366,6 +366,13 @@ def __init__( self.num_fetch_requests_cur_rank = 0 self.num_fetch_requests = 0 self.shutdown_event = threading.Event() + self._disagg_gen_init_prepared_ids: Dict[ + ResourceManagerType, set[int]] = { + ResourceManagerType.KV_CACHE_MANAGER: set(), + ResourceManagerType.SPEC_RESOURCE_MANAGER: set(), + ResourceManagerType.DRAFT_KV_CACHE_MANAGER: set(), + } + self._disagg_gen_kv_recv_started_ids: set[int] = set() # Rolling acceptance tracking for spec decode (disable speculation if rolling acceptance is below threshold) spec_config = getattr(self.model_engine, 'spec_config', None) @@ -3147,9 +3154,6 @@ def _pad_attention_dp_dummy_request(self): @nvtx_range("_prepare_disagg_gen_init") def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests): if fitting_disagg_gen_init_requests: - disagg_gen_init_to_prepare = ScheduledRequests() - disagg_gen_init_to_prepare.context_requests_last_chunk = fitting_disagg_gen_init_requests - for resource_mgr_type in ( ResourceManagerType.KV_CACHE_MANAGER, ResourceManagerType.SPEC_RESOURCE_MANAGER, @@ -3157,9 +3161,22 @@ def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests): if (resource_mgr_type in self.resource_manager.resource_managers and self.resource_manager. resource_managers[resource_mgr_type] is not None): + prepared_ids = self._disagg_gen_init_prepared_ids[ + resource_mgr_type] + resources_to_prepare = [ + req for req in fitting_disagg_gen_init_requests + if req.py_request_id not in prepared_ids + ] + if not resources_to_prepare: + continue + + disagg_gen_init_to_prepare = ScheduledRequests() + disagg_gen_init_to_prepare.context_requests_last_chunk = resources_to_prepare self.resource_manager.resource_managers[ resource_mgr_type].prepare_resources( disagg_gen_init_to_prepare) + for req in resources_to_prepare: + prepared_ids.add(req.py_request_id) # Trigger KV cache exchange for new disagg_gen_init_requests self._recv_disagg_gen_cache(fitting_disagg_gen_init_requests) @@ -3242,15 +3259,36 @@ def _recv_disagg_gen_cache(self, new_gen_reqs): req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE return + recv_reqs = [] + for req in new_gen_reqs: + if req.py_request_id in self._disagg_gen_kv_recv_started_ids: + continue + recv_reqs.append(req) + + if not recv_reqs: + return + if os.getenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP") == "1": - for req in new_gen_reqs: - self.kv_cache_transceiver.request_and_receive_sync(req) + for req in recv_reqs: + self._disagg_gen_kv_recv_started_ids.add(req.py_request_id) + try: + self.kv_cache_transceiver.request_and_receive_sync(req) + except Exception: + self._disagg_gen_kv_recv_started_ids.discard( + req.py_request_id) + raise else: - for req in new_gen_reqs: - self.kv_cache_transceiver.request_and_receive_async(req) + for req in recv_reqs: + self._disagg_gen_kv_recv_started_ids.add(req.py_request_id) + try: + self.kv_cache_transceiver.request_and_receive_async(req) + except Exception: + self._disagg_gen_kv_recv_started_ids.discard( + req.py_request_id) + raise if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None: - for req in new_gen_reqs: + for req in recv_reqs: if req.state == LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS: req.py_kv_transfer_start_time = time.time() @@ -3349,6 +3387,15 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): for request_id in list(requests_in_transfer.keys()): request = requests_in_transfer[request_id] if request.py_kv_transfer_timed_out and request_id not in completed_req_ids: + # The previous "defer Python cleanup while + # DISAGG_CONTEXT_TRANS_IN_PROGRESS" guard existed to avoid a + # UAF on CacheTransceiver::mSenderFutures's raw LlmRequest*. + # mSenderFutures now holds std::shared_ptr, so the + # LlmRequest outlives every C++ access regardless of when + # Python drops its pybind reference — the UAF class is gone + # and the guard is no longer needed. Running cancel_request + + # end_transfer here recovers the KV blocks promptly even when + # the C++ deadline check cannot reach the entry (orphan class). is_cancelled = self.kv_cache_transceiver.cancel_request(request) # If cancel is successful, mark as complete so it can be cleaned up # Otherwise, try at next iteration @@ -3585,6 +3632,13 @@ def _do_terminate_request(self, request: LlmRequest): if self.gather_all_responses or self.dist.rank == 0: self.result_wait_queues.pop(request.py_request_id, None) + # Drop the per-request idempotency markers used by + # _prepare_disagg_gen_init / _recv_disagg_gen_cache so the tracking + # sets do not grow unboundedly across the deployment lifetime. + for prepared_ids in self._disagg_gen_init_prepared_ids.values(): + prepared_ids.discard(request.py_request_id) + self._disagg_gen_kv_recv_started_ids.discard(request.py_request_id) + def _is_request_in_transmission(self, request) -> bool: """Check if a request is currently in transmission state.""" return (request.state @@ -3727,7 +3781,17 @@ def _handle_responses(self): requests_to_terminate.append(request) continue - # Check if generation request needs cleanup due to KV cache transfer timeout + # Check if generation request needs cleanup due to KV cache transfer timeout. + # + # The C++ transceiver's async workers and mSenderFutures / + # mRequesterFutures now hold std::shared_ptr (not raw + # pointers), so Python-side _terminate_request can safely run + # regardless of C++ state — the LlmRequest stays alive until all + # strong references (Python active_requests, C++ futures map, + # async workers) drop. The previous "defer cleanup until C++ + # evicts first" guard was what turned C++ tracker orphans into + # permanent KV-block leaks because the C++ deadline check could + # not reach orphaned entries. if request.py_kv_transfer_timed_out: is_cancelled = self.kv_cache_transceiver.cancel_request(request) if is_cancelled: @@ -3801,6 +3865,15 @@ def _handle_responses(self): # If partial reuse is enabled, and the KV cache manager is not VSWA, and the PP size is 1, # then we need to terminate the request. TODO: Remove this once disagg support from KVCache reuse # path is fixed. + # Note: the previous "elif request.is_disagg_context_transmission_state: + # pass" guard existed to avoid a UAF on the C++ + # CacheTransceiver::mSenderFutures's raw LlmRequest*. Since + # mSenderFutures now holds std::shared_ptr, the + # LlmRequest outlives every C++ access regardless of when + # Python drops its pybind reference, so the guard is no + # longer needed and was actively turning C++ tracker + # orphans into permanent KV-block leaks (the deadline check + # cannot reach orphaned entries). force_terminate_for_partial_reuse = ( self.enable_partial_reuse_for_disagg and not self.kv_cache_manager.is_vswa From e33c47da18dad0e2079c5641b49e7538f05c3226 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Wed, 29 Apr 2026 20:27:00 -0700 Subject: [PATCH 2/9] [https://nvbugs/6104831][test] Add reproducers for broken-promise on disagg cancellation. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> (cherry picked from commit 944561dd5c71bbd7d2ce4b374cea9a6a8a483009) --- .../others/test_kv_cache_transceiver.py | 83 +++++++++++++++++++ 1 file changed, 83 insertions(+) diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index 6960b1e28d07..0ab919dfd5d3 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -233,6 +233,89 @@ def test_cancel_request_in_transmission(attention_type): assert gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR +@pytest.mark.timeout(120) +@pytest.mark.parametrize("attention_type", + [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA], + ids=["mha", "mla"]) +def test_cancel_request_in_transmission_does_not_break_sender_future( + attention_type, capfd): + tensorrt_llm.logger.set_level("info") + mapping = Mapping(world_size=1, rank=0) + dist = Distributed.get(mapping) + ctx_kv_cache_dtype, gen_kv_cache_dtype = DataType.HALF, DataType.HALF + kv_cache_manager_ctx = create_kv_cache_manager(mapping, ctx_kv_cache_dtype) + kv_cache_manager_gen = create_kv_cache_manager(mapping, gen_kv_cache_dtype) + + cache_transceiver_config = CacheTransceiverConfig(backend="DEFAULT", + max_tokens_in_buffer=512) + + kv_cache_transceiver_ctx = create_kv_cache_transceiver( + mapping, dist, kv_cache_manager_ctx, attention_type, + cache_transceiver_config) + + kv_cache_transceiver_gen = create_kv_cache_transceiver( + mapping, dist, kv_cache_manager_gen, attention_type, + cache_transceiver_config) + + fill_kv_cache_buffer(kv_cache_manager_ctx) + + sampling_params = SamplingParams() + ctx_request = LlmRequest( + request_id=0, + max_new_tokens=1, + input_tokens=list(range(256)), + sampling_config=tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config()), + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) + + kv_cache_manager_ctx.impl.add_sequence_batch( + [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], [ctx_request]) + kv_cache_transceiver_ctx.respond_and_send_async(ctx_request) + + time.sleep(2) + is_cancelled = kv_cache_transceiver_ctx.cancel_request(ctx_request) + assert is_cancelled + + gen_request = LlmRequest( + request_id=0, + max_new_tokens=1, + input_tokens=list(range(256)), + sampling_config=tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config()), + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + context_phase_params=ctx_request.context_phase_params) + + kv_cache_manager_gen.impl.add_sequence_batch( + [(gen_request.py_request_id, gen_request.prompt_len, 1)], [gen_request]) + kv_cache_transceiver_gen.request_and_receive_async(gen_request) + + completed_ids, error_ids = [], [] + deadline = time.time() + 10 + while time.time() < deadline and not error_ids: + completed_ids, error_ids = kv_cache_transceiver_ctx.check_context_transfer_status( + 1) + if error_ids: + break + time.sleep(0.1) + + assert ctx_request.py_request_id not in completed_ids + assert ctx_request.py_request_id in error_ids + + deadline = time.time() + 10 + while time.time( + ) < deadline and gen_request.state != LlmRequestState.DISAGG_TRANS_ERROR: + kv_cache_transceiver_gen.check_gen_transfer_status(1) + time.sleep(0.1) + + assert gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR + + captured = capfd.readouterr() + merged = captured.out + captured.err + assert "Broken promise" not in merged + + def create_hybrid_cache_manager(mapping, dtype, mamba_conv_dtype=torch.float16, From a61fa1c7f1e55646d45112c0b8e35f93b77b4f2a Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 4 May 2026 12:59:29 -0700 Subject: [PATCH 3/9] [https://nvbugs/6104831][test] Add disagg cancellation regression tests Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../others/test_kv_cache_transceiver.py | 542 ++++++++++++++++-- 1 file changed, 482 insertions(+), 60 deletions(-) diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index 0ab919dfd5d3..243f345cfabd 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -1,3 +1,4 @@ +import threading import time import uuid @@ -51,6 +52,127 @@ def fill_kv_cache_buffer(kv_cache_manager): buffer.copy_(random_values) +# create_kv_cache_manager() configures a 256-token KV pool. Tests that need +# two concurrent generation requests must keep prompts small enough that both +# fit, hence this default. Single-request tests can override via prompt_len. +_DEFAULT_PROMPT_LEN = 64 + + +def _make_request(request_id, + llm_request_type, + context_phase_params=None, + prompt_len=_DEFAULT_PROMPT_LEN): + """Construct an ``LlmRequest`` with the shared disagg test boilerplate. + + ``prompt_len`` controls the synthetic input-token range and must fit in + the small KV pool from :func:`create_kv_cache_manager`. + """ + sampling_params = SamplingParams() + kwargs = dict( + request_id=request_id, + max_new_tokens=1, + input_tokens=list(range(prompt_len)), + sampling_config=tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config()), + is_streaming=False, + llm_request_type=llm_request_type, + ) + if context_phase_params is not None: + kwargs["context_phase_params"] = context_phase_params + return LlmRequest(**kwargs) + + +def _add_sequence(kv_cache_manager, request): + """Bind ``request`` into ``kv_cache_manager``. + + This uses the call signature every disagg transceiver test relies on. + """ + kv_cache_manager.impl.add_sequence(request.py_request_id, + request.prompt_len, 1, request) + + +def _drive_template_handshake_to_completion(ctx_xcvr, + gen_xcvr, + ctx_request_id, + deadline_s=10): + """Poll both transceivers until the template transfer completes. + + ``ctx_request_id`` must land in the sender-side ``completed_ids`` AND the + gen-side ``mRequesterFutures`` must drain. + + The post-fix ``check_gen_transfer_status(1)`` is non-blocking (it skips + unready futures via ``wait_for(0)``), so a single call may leave the + template entry in ``mRequesterFutures`` if its data has not yet + arrived. Tests that stage a "blocked" request after a template + handshake must drain mRequesterFutures first so the blocked-side + assertions only see their own request. + """ + deadline = time.time() + deadline_s + ctx_done = False + while time.time() < deadline and ( + not ctx_done or not gen_xcvr.check_gen_transfer_complete()): + completed_ids, error_ids = ctx_xcvr.check_context_transfer_status(1) + assert ctx_request_id not in error_ids, ( + f"template ctx request {ctx_request_id} unexpectedly errored: " + f"{error_ids}") + if ctx_request_id in completed_ids: + ctx_done = True + gen_xcvr.check_gen_transfer_status(1) + time.sleep(0.05) + assert ctx_done, ( + f"template ctx request {ctx_request_id} did not complete within " + f"{deadline_s}s; sender side may be stuck before the test can " + f"stage its blocked request") + assert gen_xcvr.check_gen_transfer_complete(), ( + f"gen-side mRequesterFutures still non-empty within {deadline_s}s; " + f"the staged blocked-request assertions would otherwise observe " + f"a spurious lingering entry from the template handshake") + + +@pytest.fixture +def disagg_transceiver_pair(request): + """Build a single-process disagg ctx/gen transceiver pair. + + Used by the cancellation-flow regression tests. + + Yields ``(kv_cache_manager_ctx, kv_cache_manager_gen, + kv_cache_transceiver_ctx, kv_cache_transceiver_gen)``. + + Pass ``(attention_type, backend)`` as the parametrize argument; the + backend defaults to ``"DEFAULT"`` (UCX) for backend-agnostic fixes + and can be set to ``"NIXL"`` for tests that depend on NIXL-specific + paths (e.g. the recv buffer index manager). + """ + param = request.param + if isinstance(param, tuple): + attention_type = param[0] + backend = param[1] if len(param) > 1 else "DEFAULT" + else: + attention_type = param + backend = "DEFAULT" + + tensorrt_llm.logger.set_level("info") + mapping = Mapping(world_size=1, rank=0) + dist = Distributed.get(mapping) + kv_cache_manager_ctx = create_kv_cache_manager(mapping, DataType.HALF) + kv_cache_manager_gen = create_kv_cache_manager(mapping, DataType.HALF) + + cache_transceiver_config = CacheTransceiverConfig(backend=backend, + max_tokens_in_buffer=512) + + kv_cache_transceiver_ctx = create_kv_cache_transceiver( + mapping, dist, kv_cache_manager_ctx, attention_type, + cache_transceiver_config) + kv_cache_transceiver_gen = create_kv_cache_transceiver( + mapping, dist, kv_cache_manager_gen, attention_type, + cache_transceiver_config) + + fill_kv_cache_buffer(kv_cache_manager_ctx) + + yield (kv_cache_manager_ctx, kv_cache_manager_gen, kv_cache_transceiver_ctx, + kv_cache_transceiver_gen) + + @pytest.fixture(scope="function") def ctx_gen_kv_cache_dtype(request): if request.param == "ctx_fp8_gen_fp8": @@ -234,95 +356,396 @@ def test_cancel_request_in_transmission(attention_type): @pytest.mark.timeout(120) -@pytest.mark.parametrize("attention_type", +@pytest.mark.parametrize("disagg_transceiver_pair", [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA], - ids=["mha", "mla"]) + ids=["mha", "mla"], + indirect=True) def test_cancel_request_in_transmission_does_not_break_sender_future( - attention_type, capfd): - tensorrt_llm.logger.set_level("info") - mapping = Mapping(world_size=1, rank=0) - dist = Distributed.get(mapping) - ctx_kv_cache_dtype, gen_kv_cache_dtype = DataType.HALF, DataType.HALF - kv_cache_manager_ctx = create_kv_cache_manager(mapping, ctx_kv_cache_dtype) - kv_cache_manager_gen = create_kv_cache_manager(mapping, gen_kv_cache_dtype) - - cache_transceiver_config = CacheTransceiverConfig(backend="DEFAULT", - max_tokens_in_buffer=512) - - kv_cache_transceiver_ctx = create_kv_cache_transceiver( - mapping, dist, kv_cache_manager_ctx, attention_type, - cache_transceiver_config) + disagg_transceiver_pair, capfd): + """Reproduce the sender-side broken-promise on cancel-after-ready. - kv_cache_transceiver_gen = create_kv_cache_transceiver( - mapping, dist, kv_cache_manager_gen, attention_type, - cache_transceiver_config) - - fill_kv_cache_buffer(kv_cache_manager_ctx) - - sampling_params = SamplingParams() - ctx_request = LlmRequest( - request_id=0, - max_new_tokens=1, - input_tokens=list(range(256)), - sampling_config=tensorrt_llm.bindings.SamplingConfig( - sampling_params._get_sampling_config()), - is_streaming=False, - llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) + Pre-fix ``CacheSender::Impl::sendResponse`` erases the ready + ``(request, promise)`` pair on cancel without first fulfilling the + promise, and the destructor surfaces ``future_error: Broken promise`` + on the sender future returned by ``respond_and_send_async()``. + """ + (kv_cache_manager_ctx, kv_cache_manager_gen, kv_cache_transceiver_ctx, + kv_cache_transceiver_gen) = disagg_transceiver_pair - kv_cache_manager_ctx.impl.add_sequence_batch( - [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], [ctx_request]) + ctx_request = _make_request(0, + LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY, + prompt_len=256) + _add_sequence(kv_cache_manager_ctx, ctx_request) kv_cache_transceiver_ctx.respond_and_send_async(ctx_request) + # Wait for the sender to reach the ready state, then cancel it. time.sleep(2) is_cancelled = kv_cache_transceiver_ctx.cancel_request(ctx_request) - assert is_cancelled - - gen_request = LlmRequest( - request_id=0, - max_new_tokens=1, - input_tokens=list(range(256)), - sampling_config=tensorrt_llm.bindings.SamplingConfig( - sampling_params._get_sampling_config()), - is_streaming=False, - llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, - context_phase_params=ctx_request.context_phase_params) - - kv_cache_manager_gen.impl.add_sequence_batch( - [(gen_request.py_request_id, gen_request.prompt_len, 1)], [gen_request]) + assert is_cancelled, ( + "ctx_request must be cancellable while still ready in " + "mReadyResponses") + + gen_request = _make_request(0, + LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + ctx_request.context_phase_params, + prompt_len=256) + _add_sequence(kv_cache_manager_gen, gen_request) kv_cache_transceiver_gen.request_and_receive_async(gen_request) + # Sender side: the cancelled request must surface as an error rather + # than being marked complete. completed_ids, error_ids = [], [] deadline = time.time() + 10 while time.time() < deadline and not error_ids: - completed_ids, error_ids = kv_cache_transceiver_ctx.check_context_transfer_status( - 1) + completed_ids, error_ids = ( + kv_cache_transceiver_ctx.check_context_transfer_status(1)) if error_ids: break time.sleep(0.1) - assert ctx_request.py_request_id not in completed_ids - assert ctx_request.py_request_id in error_ids + assert ctx_request.py_request_id not in completed_ids, ( + "cancelled ctx request must not appear in completed_ids") + assert ctx_request.py_request_id in error_ids, ( + "cancelled ctx request must surface in error_ids") + # Receiver side: the matching gen request must observe the cancellation + # as a structured DISAGG_TRANS_ERROR rather than a Broken-promise + # exception. deadline = time.time() + 10 - while time.time( - ) < deadline and gen_request.state != LlmRequestState.DISAGG_TRANS_ERROR: + while time.time() < deadline and (gen_request.state + != LlmRequestState.DISAGG_TRANS_ERROR): kv_cache_transceiver_gen.check_gen_transfer_status(1) time.sleep(0.1) - assert gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR + assert gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR, ( + "gen request mirroring the cancelled ctx request must end in " + "DISAGG_TRANS_ERROR") captured = capfd.readouterr() merged = captured.out + captured.err - assert "Broken promise" not in merged + assert "Broken promise" not in merged, ( + "signature #1 reproduced: cancel-after-ready left the sender " + "promise unresolved and the destructor surfaced as Broken " + "promise on the sender future") + + kv_cache_manager_ctx.free_resources(ctx_request) + kv_cache_manager_gen.free_resources(gen_request) + + +_PROBE_TIMEOUT_S = 2.5 + + +@pytest.mark.timeout(120) +@pytest.mark.parametrize("disagg_transceiver_pair", + [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA], + ids=["mha", "mla"], + indirect=True) +def test_check_gen_transfer_status_at_least_one_does_not_block_on_unready_future( + disagg_transceiver_pair): + """Reproduce the gen-side blocking hang in checkGenTransferStatus(1). + + On stock ``rc11`` the polling path called from the PyExecutor disagg + loop unconditionally ``future.get()``s the first selected requester + future, even when its ``wait_for(0)`` is still ``timeout``. A single + in-flight generation request whose context-side ready signal has not + yet arrived therefore blocks the entire decoder event loop, which is + indistinguishable from a wedge. + + The test exercises the same shape as the wedge: drive one full + ctx/gen handshake to completion to capture an opaque comm/cache + state, then enqueue a generation request whose context counterpart + has not yet been ``respond_and_send_async()``-ed and call + ``check_gen_transfer_status(1)`` from a separate thread. The call + must return within a bounded probe timeout larger than the + production 1000ms sender-future wait instead of blocking indefinitely + on the unresolved future. + """ + (kv_cache_manager_ctx, kv_cache_manager_gen, kv_cache_transceiver_ctx, + kv_cache_transceiver_gen) = disagg_transceiver_pair + + # Complete one normal transfer first so we can reuse its opaque + # comm/cache state for the second (intentionally unresolved) generation + # request. + template_ctx_request = _make_request( + 100, LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY, prompt_len=256) + _add_sequence(kv_cache_manager_ctx, template_ctx_request) + kv_cache_transceiver_ctx.respond_and_send_async(template_ctx_request) + + template_gen_request = _make_request( + 100, + LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + template_ctx_request.context_phase_params, + prompt_len=256) + _add_sequence(kv_cache_manager_gen, template_gen_request) + kv_cache_transceiver_gen.request_and_receive_async(template_gen_request) + + _drive_template_handshake_to_completion(kv_cache_transceiver_ctx, + kv_cache_transceiver_gen, + template_ctx_request.py_request_id) + + opaque_state = template_ctx_request.context_phase_params.opaque_state + assert opaque_state is not None, ( + "template ctx_request must expose its opaque comm/cache state for " + "the staged blocked request below to reuse") + + kv_cache_manager_ctx.free_resources(template_ctx_request) + kv_cache_manager_gen.free_resources(template_gen_request) + + # Build a generation request for a different ctx_request_id before the + # sender has any matching ready response. This leaves a real unresolved + # future in the C++ transceiver and reproduces the blocking pattern + # behind signature #4. + blocked_request_id = 101 + blocked_ctx_request = _make_request( + blocked_request_id, + LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY, + prompt_len=256) + _add_sequence(kv_cache_manager_ctx, blocked_ctx_request) + + blocked_context_phase_params = trtllm.ContextPhaseParams( + list(template_ctx_request.context_phase_params.first_gen_tokens), + blocked_request_id, + bytes(opaque_state), + template_ctx_request.context_phase_params.draft_tokens, + template_ctx_request.context_phase_params.ctx_dp_rank, + template_ctx_request.context_phase_params.disagg_info_endpoint, + ) + blocked_gen_request = _make_request( + blocked_request_id, + LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + blocked_context_phase_params, + prompt_len=256) + _add_sequence(kv_cache_manager_gen, blocked_gen_request) + kv_cache_transceiver_gen.request_and_receive_async(blocked_gen_request) + + # Sanity: at_least_request_num=0 must not block under any circumstance. + start = time.time() + kv_cache_transceiver_gen.check_gen_transfer_status(0) + assert time.time() - start < 1.0, ( + "check_gen_transfer_status(0) must be non-blocking even when " + "futures are unresolved") + + # Real reproducer: at_least_request_num=1 must NOT hang on an + # unresolved future. Run it on a worker thread so that pre-fix this + # thread stays blocked past the bounded probe timeout (the failure + # signature), and post-fix it returns after at most the configured + # per-iteration sender-future wait because the unready future is skipped. + check_result = {"returned": False, "error": None} + + def call_blocking_check(): + try: + kv_cache_transceiver_gen.check_gen_transfer_status(1) + except BaseException as exc: # noqa: BLE001 + check_result["error"] = exc + finally: + check_result["returned"] = True + + blocked_check = threading.Thread(target=call_blocking_check, daemon=True) + blocked_check.start() + blocked_check.join(timeout=_PROBE_TIMEOUT_S) + blocked_during_probe = blocked_check.is_alive() + + # Allow the wedged context request to complete cleanly so the worker + # thread can finish (in either pre- or post-fix behaviour) and we can + # tear down the test without leaking threads. + kv_cache_transceiver_ctx.respond_and_send_async(blocked_ctx_request) + + deadline = time.time() + 10 + completed_ids, error_ids = [], [] + while time.time() < deadline and (blocked_ctx_request.py_request_id + not in completed_ids): + completed_ids, error_ids = ( + kv_cache_transceiver_ctx.check_context_transfer_status(1)) + assert blocked_ctx_request.py_request_id not in error_ids, ( + "blocked ctx request must not error during teardown polling") + if blocked_ctx_request.py_request_id in completed_ids: + break + time.sleep(0.1) + assert blocked_ctx_request.py_request_id in completed_ids, ( + "blocked ctx request must complete on the sender side once " + "respond_and_send_async unblocks the receiver") + + if blocked_during_probe: + # Pre-fix path: the thread is wedged inside + # check_gen_transfer_status(1). It will only return after the + # sender's respond_and_send_async unblocks the receive that the + # in-thread call is parked on. + blocked_check.join(timeout=10) + assert not blocked_check.is_alive(), ( + "blocked check thread must drain within 10s after the sender " + "supplies the ready signal") + else: + # Post-fix path: the in-thread call already returned because it + # skipped the unready future. Drain mRequesterFutures via + # additional polling so subsequent assertions see an empty queue. + deadline = time.time() + 10 + while time.time() < deadline and ( + not kv_cache_transceiver_gen.check_gen_transfer_complete()): + kv_cache_transceiver_gen.check_gen_transfer_status(1) + time.sleep(0.1) + + if check_result["error"] is not None: + raise check_result["error"] + assert check_result["returned"], ( + "blocked check thread must signal completion via check_result") + assert not blocked_during_probe, ( + "signature #4 reproduced: check_gen_transfer_status(1) blocked on " + "an unresolved generation future for longer than the bounded probe") + assert kv_cache_transceiver_gen.check_gen_transfer_complete(), ( + "gen-side mRequesterFutures must drain after the blocked request " + "is unblocked") + + assert torch.equal( + kv_cache_manager_gen.get_buffers(0), + kv_cache_manager_ctx.get_buffers(0)), "different kv-cache values" + + kv_cache_manager_ctx.free_resources(blocked_ctx_request) + kv_cache_manager_gen.free_resources(blocked_gen_request) + + +@pytest.mark.timeout(120) +@pytest.mark.parametrize("disagg_transceiver_pair", + [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA], + ids=["mha", "mla"], + indirect=True) +def test_cancel_queued_gen_request_fulfills_receiver_future( + disagg_transceiver_pair, capfd): + """Reproduce the receiver-side queued-cancel broken-promise. + + Mirrors the sender-side reproducer but on + ``CacheReceiver::Impl::cancelRequest()``. Pre-fix that function erases + the queued ``(request, promise)`` pair without first fulfilling the + promise, so the ``std::promise`` destructor sets + ``future_error: Broken promise`` on the future returned by + ``request_and_receive_async()``. + + The test holds the receiver worker thread busy with a first + generation request that has no matching context counterpart, + enqueues a second generation request behind it, then cancels the + queued request. Post-fix the queued promise is fulfilled with a + structured cancellation exception before the queue entry is erased. + """ + (kv_cache_manager_ctx, kv_cache_manager_gen, kv_cache_transceiver_ctx, + kv_cache_transceiver_gen) = disagg_transceiver_pair + + # Drive one full ctx/gen handshake to completion so we can reuse a + # real opaque comm/cache state for the orphan requests below. + template_ctx_request = _make_request( + 100, LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) + _add_sequence(kv_cache_manager_ctx, template_ctx_request) + kv_cache_transceiver_ctx.respond_and_send_async(template_ctx_request) + + template_gen_request = _make_request( + 100, LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + template_ctx_request.context_phase_params) + _add_sequence(kv_cache_manager_gen, template_gen_request) + kv_cache_transceiver_gen.request_and_receive_async(template_gen_request) + + _drive_template_handshake_to_completion(kv_cache_transceiver_ctx, + kv_cache_transceiver_gen, + template_ctx_request.py_request_id) + + opaque_state = template_ctx_request.context_phase_params.opaque_state + assert opaque_state is not None, ( + "template ctx_request must expose its opaque comm/cache state for " + "the staged orphan requests below to reuse") + + kv_cache_manager_ctx.free_resources(template_ctx_request) + kv_cache_manager_gen.free_resources(template_gen_request) + + def make_orphan_gen_request(request_id): + ctx_phase_params = trtllm.ContextPhaseParams( + list(template_ctx_request.context_phase_params.first_gen_tokens), + request_id, + bytes(opaque_state), + template_ctx_request.context_phase_params.draft_tokens, + template_ctx_request.context_phase_params.ctx_dp_rank, + template_ctx_request.context_phase_params.disagg_info_endpoint, + ) + gen_request = _make_request( + request_id, LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + ctx_phase_params) + _add_sequence(kv_cache_manager_gen, gen_request) + return gen_request + + # Submit a first orphan gen request whose context counterpart will + # never respond_and_send_async. The receiver worker dequeues it and + # parks inside requestSync() / sendRequestInfo(), tying up the worker + # thread. + blocking_gen_request = make_orphan_gen_request(101) + kv_cache_transceiver_gen.request_and_receive_async(blocking_gen_request) + + # Submit a second orphan gen request. The first one is still in + # requestSync(), so this one stays in mRequestsQueue and is the + # actual subject of the queued-cancel reproducer. + queued_gen_request = make_orphan_gen_request(102) + kv_cache_transceiver_gen.request_and_receive_async(queued_gen_request) + + # Wait briefly so the receiver worker has had time to dequeue the + # first request and block on it, leaving the second one queued. + time.sleep(1) + + # Cancel the queued request. Pre-fix this erases the (request, + # promise) pair without fulfilling the promise; post-fix it + # set_exception()s a structured kNETWORK_ERROR before erasing. + is_cancelled = kv_cache_transceiver_gen.cancel_request(queued_gen_request) + assert is_cancelled, ( + "queued_gen_request must still be in the receiver queue when we " + "call cancel_request(); if this fails, the receiver worker may " + "have dequeued faster than expected and the test setup needs to " + "be tightened") + + # Poll the gen-side polling loop and assert the cancelled request + # lands in DISAGG_TRANS_ERROR within a reasonable window. Pre-fix + # this returns via a Broken-promise exception with no useful + # diagnostic; post-fix it returns via the structured kNETWORK_ERROR + # set by the fix. + deadline = time.time() + 10 + while time.time() < deadline and (queued_gen_request.state + != LlmRequestState.DISAGG_TRANS_ERROR): + kv_cache_transceiver_gen.check_gen_transfer_status(1) + time.sleep(0.1) + + assert queued_gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR, ( + "queued gen request must land in DISAGG_TRANS_ERROR after " + "cancellation rather than surfacing a Broken-promise exception") + + # Clean up the intentionally blocked request before the transceiver + # is destroyed. Leaving it in requestSync() would make teardown race + # the worker thread that this test deliberately parked. + is_blocking_cancelled = ( + kv_cache_transceiver_gen.cancel_request(blocking_gen_request)) + assert is_blocking_cancelled, ( + "blocking_gen_request must be cancellable so teardown does not " + "race the parked receiver worker") + deadline = time.time() + 10 + while time.time() < deadline and (blocking_gen_request.state + != LlmRequestState.DISAGG_TRANS_ERROR): + kv_cache_transceiver_gen.check_gen_transfer_status(1) + time.sleep(0.1) + assert blocking_gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR, ( + "blocking_gen_request must drain to DISAGG_TRANS_ERROR before " + "transceiver destruction") + + kv_cache_manager_gen.free_resources(blocking_gen_request) + kv_cache_manager_gen.free_resources(queued_gen_request) + + captured = capfd.readouterr() + merged = captured.out + captured.err + assert "Broken promise" not in merged, ( + "signature #5 reproduced: cancelling a queued generation request " + "left its std::promise unresolved and the destructor surfaced " + "as Broken promise on the consumer side") def create_hybrid_cache_manager(mapping, dtype, mamba_conv_dtype=torch.float16, mamba_ssm_dtype=torch.float16): - """ - Create a MambaHybridCacheManager for testing hybrid models. - This manager handles both KV cache (attention layers) and Mamba cache (RNN layers). + """Create a MambaHybridCacheManager for testing hybrid models. + + This manager handles both KV cache (attention layers) and Mamba cache + (RNN layers). Args: mapping: The mapping configuration. @@ -390,8 +813,7 @@ def fill_hybrid_cache_buffers(hybrid_cache_manager): @pytest.fixture(scope="function") def hybrid_dtypes(request): - """ - Returns (kv_dtype, mamba_conv_dtype, mamba_ssm_dtype) based on the parametrized string. + """Return dtypes for the parametrized hybrid-cache test case. KV dtype: fp8, bf16 Conv dtype: fp8, bf16, fp32 From 4def7b2173bc0a3ebbba045d26cae666bd31479e Mon Sep 17 00:00:00 2001 From: Yifan Jiang Date: Mon, 4 May 2026 11:24:21 -0700 Subject: [PATCH 4/9] Fail closed on unquiesced disagg KV transfer Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> 1777930306 -0700 --- .../batch_manager/cacheTransceiver.h | 8 +- .../batch_manager/baseTransBuffer.cpp | 124 ++++++++++++- .../batch_manager/baseTransBuffer.h | 17 +- .../batch_manager/cacheFormatter.cpp | 19 +- .../batch_manager/cacheTransceiver.cpp | 93 +++++++--- .../batch_manager/dataTransceiver.cpp | 107 ++++++++---- .../agent_utils/connection.cpp | 7 +- .../agent_utils/connection.h | 1 + .../disagg-kv-transfer-cancel-safety.md | 126 ++++++++++++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 163 +++++++++++++++--- 10 files changed, 578 insertions(+), 87 deletions(-) create mode 100644 docs/source/features/disagg-kv-transfer-cancel-safety.md diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index c0af00e9bf4d..fb58378dd4d0 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -35,6 +35,7 @@ #include #include #include +#include #include using SizeType32 = tensorrt_llm::runtime::SizeType32; @@ -286,11 +287,12 @@ class CacheTransceiver : public BaseCacheTransceiver // raw pointer is still dereferenced by checkGenTransferStatus / // checkContextTransferStatus (the UAF forensically confirmed via // MALLOC_PERTURB_=85 producing mRequestId=0x5555555555555555). Entries are - // erased only in the normal eviction paths (completion, deadline, - // exception), ensuring the request outlives every C++ access but doesn't - // leak past the transfer lifecycle. + // erased only after completion, exception, or timeout plus worker-future + // readiness. A timeout/cancel request is not a quiescence proof, so the + // tracker must keep the shared_ptr until the worker future resolves. std::vector, std::future>> mSenderFutures; std::vector, std::future>> mRequesterFutures; + std::unordered_set mTimedOutRequesterIds; mpi::MpiComm const* mMpiWorldComm{nullptr}; std::shared_ptr mGroupComm; diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp index 2928c87bbd06..9b83205451a7 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp @@ -21,11 +21,28 @@ #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/opUtils.h" +#include #include namespace tensorrt_llm::batch_manager { +namespace +{ + +char const* bufferKindName(BufferKind kind) +{ + switch (kind) + { + case BufferKind::kKV: return "kv"; + case BufferKind::kKV_INDEXER: return "kv_indexer"; + case BufferKind::kRNN: return "rnn"; + } + return "unknown"; +} + +} // namespace + void BufferIndexHolder::release() noexcept { // Happy-path release: frees the slot and disarms the holder in one @@ -56,6 +73,31 @@ void BufferIndexHolder::release() noexcept mHeld = false; } +void BufferIndexHolder::poison() noexcept +{ + if (!mHeld || mMgr == nullptr) + { + return; + } + try + { + if (mIsRecv) + { + mMgr->poisonBufferIndexForRecv(mIndex); + } + else + { + mMgr->poisonBufferIndexForSend(mIndex); + } + } + catch (...) + { + // poisonBufferIndex is noexcept; keep this as belt-and-suspenders so + // fail-closed cleanup cannot throw from an exception path. + } + mHeld = false; +} + BaseTransBufferManager::BaseTransBufferManager( size_t transferBufferSize, nvinfer1::DataType dataType, std::optional maxNumTokens) : mDataType{dataType} @@ -96,6 +138,11 @@ void BaseTransBufferManager::freeBufferIndexForSend(std::optional bufferId) freeBufferIndex(mConcurrenceSendResource, bufferId, mSendBufferCount, mOnlyUseDynamicBuffer); } +void BaseTransBufferManager::poisonBufferIndexForSend(std::optional bufferId) noexcept +{ + poisonBufferIndex(mConcurrenceSendResource, bufferId, mSendBufferCount, mOnlyUseDynamicBuffer, "send"); +} + std::optional BaseTransBufferManager::assignBufferIndexForRecv( std::atomic const* perRequestCancel, int64_t waitSliceMs, std::optional requestIdForLog) { @@ -108,6 +155,11 @@ void BaseTransBufferManager::freeBufferIndexForRecv(std::optional bufferId) freeBufferIndex(mConcurrenceRecvResource, bufferId, mRecvBufferCount, mOnlyUseDynamicBuffer); } +void BaseTransBufferManager::poisonBufferIndexForRecv(std::optional bufferId) noexcept +{ + poisonBufferIndex(mConcurrenceRecvResource, bufferId, mRecvBufferCount, mOnlyUseDynamicBuffer, "recv"); +} + std::tuple, size_t, bool> BaseTransBufferManager::getOrAllocateSendBuffers( std::optional bufferId, int targetNum, std::vector const& requestedNumberOfElements, runtime::BufferManager const& bufferManagerToUse) @@ -265,6 +317,10 @@ std::optional BaseTransBufferManager::assignBufferIndex(ConcurrenceResource { if (onlyUseDynamicBuffer) { + TLLM_CHECK_WITH_INFO(!resource.mPoisoned.load(std::memory_order_relaxed), + "Cannot assign dynamic cache transfer buffer kind=%s because a previous transfer left dynamic transfer " + "memory poisoned. The process must restart before these memory ranges can be safely reused.", + bufferKindName(getBufferKind())); return std::nullopt; } // Bounded wait_for loop so a cancel fired on this request while parked @@ -272,8 +328,11 @@ std::optional BaseTransBufferManager::assignBufferIndex(ConcurrenceResource // mTerminate (flipped between slices) keeps the drain worker responsive // to shutdown. std::unique_lock lk(resource.mBuffersMutex); - auto const predicate - = [&resource, bufferCount]() { return static_cast(resource.mConcurrence) < bufferCount; }; + auto const predicate = [&resource, bufferCount]() + { + return resource.mPoisoned.load(std::memory_order_relaxed) + || static_cast(resource.mConcurrence) < bufferCount; + }; if (!predicate()) { auto const slice = std::chrono::milliseconds{waitSliceMs}; @@ -288,6 +347,10 @@ std::optional BaseTransBufferManager::assignBufferIndex(ConcurrenceResource } } } + TLLM_CHECK_WITH_INFO(!resource.mPoisoned.load(std::memory_order_relaxed), + "Cannot assign cache transfer buffer kind=%s because a previous transfer left the buffer pool poisoned. " + "The process must restart before these memory ranges can be safely reused.", + bufferKindName(getBufferKind())); int bufferId = -1; for (size_t i = 0; i < bufferCount; i++) { @@ -317,6 +380,12 @@ void BaseTransBufferManager::freeBufferIndex( TLLM_CHECK(static_cast(bufferId.value()) < bufferCount); { std::scoped_lock lk(resource.mBuffersMutex); + if (resource.mBufferIndexFlag[bufferId.value()] == 2) + { + TLLM_LOG_ERROR("Refusing to free poisoned cache transfer buffer kind=%s index=%d", + bufferKindName(getBufferKind()), bufferId.value()); + return; + } resource.mBufferIndexFlag[bufferId.value()] = 0; } resource.mConcurrence--; @@ -324,6 +393,57 @@ void BaseTransBufferManager::freeBufferIndex( } } +void BaseTransBufferManager::poisonBufferIndex(ConcurrenceResource& resource, std::optional bufferId, + size_t bufferCount, bool onlyUseDynamicBuffer, char const* direction) noexcept +{ + resource.mPoisoned.store(true, std::memory_order_relaxed); + + if (onlyUseDynamicBuffer) + { + TLLM_LOG_ERROR( + "Poisoned dynamic %s cache transfer buffer kind=%s. Dynamic transfer memory cannot be safely reused; " + "the process must restart.", + direction, bufferKindName(getBufferKind())); + resource.mBuffersCV.notify_all(); + return; + } + + if (!bufferId.has_value()) + { + TLLM_LOG_ERROR("Poisoned unknown %s cache transfer buffer kind=%s. The process must restart.", direction, + bufferKindName(getBufferKind())); + resource.mBuffersCV.notify_all(); + return; + } + + try + { + TLLM_CHECK(static_cast(bufferId.value()) < bufferCount); + { + std::scoped_lock lk(resource.mBuffersMutex); + if (resource.mBufferIndexFlag[bufferId.value()] == 1) + { + resource.mBufferIndexFlag[bufferId.value()] = 2; + } + } + TLLM_LOG_ERROR( + "Poisoned %s cache transfer buffer kind=%s index=%d. The slot will not be returned to the pool because " + "transport quiescence is unknown; restart the process before serving more KV transfers.", + direction, bufferKindName(getBufferKind()), bufferId.value()); + } + catch (std::exception const& e) + { + TLLM_LOG_ERROR("Exception while poisoning %s cache transfer buffer kind=%s index=%d: %s", direction, + bufferKindName(getBufferKind()), bufferId.value_or(-1), e.what()); + } + catch (...) + { + TLLM_LOG_ERROR("Unknown exception while poisoning %s cache transfer buffer kind=%s index=%d", direction, + bufferKindName(getBufferKind()), bufferId.value_or(-1)); + } + resource.mBuffersCV.notify_all(); +} + size_t BaseTransBufferManager::getRecvBufferCount() { return mRecvBufferCount; diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h index de66d3e593f8..ad39cdcd7f19 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h @@ -75,7 +75,7 @@ class BufferIndexHolder std::optional requestIdForLog = std::nullopt) noexcept : mMgr(&mgr) , mIndex(index) - , mHeld(index.has_value()) + , mHeld(true) , mIsRecv(isRecv) , mRequestIdForLog(requestIdForLog) { @@ -149,6 +149,12 @@ class BufferIndexHolder /// destructor calls release() to free the slot. void release() noexcept; + /// @brief Fail-closed release for paths where transport quiescence is not + /// known. The buffer slot is marked poisoned, is not returned to + /// the pool, and later assignments fail so the process must restart + /// before serving more transfer traffic. + void poison() noexcept; + private: BaseTransBufferManager* mMgr{nullptr}; std::optional mIndex{}; @@ -192,6 +198,9 @@ class BaseTransBufferManager /// @param bufferId The buffer index to free. void freeBufferIndexForSend(std::optional bufferId); + /// @brief Poison a send buffer index after an unquiesced transfer exit. + void poisonBufferIndexForSend(std::optional bufferId) noexcept; + /// @brief Assign a buffer index for receiving. /// @param perRequestCancel Optional per-request cancel flag. When non-null /// and flipped true while this call is parked on the pool-exhausted @@ -212,6 +221,9 @@ class BaseTransBufferManager /// @param bufferId The buffer index to free. void freeBufferIndexForRecv(std::optional bufferId); + /// @brief Poison a receive buffer index after an unquiesced transfer exit. + void poisonBufferIndexForRecv(std::optional bufferId) noexcept; + /// @brief Get or allocate send buffers for cache transfer. /// @param bufferId The assigned buffer ID. /// @param targetNum Number of target sequences. @@ -265,6 +277,7 @@ class BaseTransBufferManager std::mutex mBuffersMutex; std::condition_variable mBuffersCV; std::atomic mConcurrence{0}; + std::atomic mPoisoned{false}; }; std::tuple, size_t, bool> getOrAllocateBuffers(std::optional bufferId, @@ -277,6 +290,8 @@ class BaseTransBufferManager std::optional requestIdForLog = std::nullopt); void freeBufferIndex( ConcurrenceResource& resource, std::optional bufferId, size_t bufferCount, bool onlyUseDynamicBuffer); + void poisonBufferIndex(ConcurrenceResource& resource, std::optional bufferId, size_t bufferCount, + bool onlyUseDynamicBuffer, char const* direction) noexcept; size_t mPreAllocBufferSize; size_t mRecvBufferCount; diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index 7b531165f127..e697c6774ae6 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -464,6 +464,10 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio NVTX3_SCOPED_RANGE(sendBufferFun); TLLM_CHECK(pickUpConnections.size() == 1); + TLLM_CHECK_WITH_INFO(false, + "Zero-copy KV cache transfer sends directly from request-owned KV blocks. It is disabled for " + "cancellable disaggregated transfers until KV-block leases can prove the peer has stopped reading " + "from those blocks before the request is freed."); TLLM_CUDA_CHECK(cudaSetDevice(deviceId)); for (size_t i = 0; i < pickUpConnections.size(); i++) @@ -579,8 +583,19 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio == inputKvCacheBlocksPerWindow.begin()->second.front()->getDataType()); } - sendAllBuffers(session, deviceId, outputSplitCaches, bufferCoverTargetNum, preAllocSendBuffer, bufferManager, - targetInfo, pickUpConnections); + try + { + sendAllBuffers(session, deviceId, outputSplitCaches, bufferCoverTargetNum, preAllocSendBuffer, + bufferManager, targetInfo, pickUpConnections); + } + catch (...) + { + if (agentConnection != nullptr) + { + sendHolder.poison(); + } + throw; + } // Happy-path release — frees the slot and disarms the holder in // one noexcept call. Placed immediately after sendAllBuffers so any diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index b109fe8793cb..15090eb765f0 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -863,28 +863,45 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR auto elapsedMs = static_cast(elapsed.count()); if (elapsedMs > kvTransferTimeoutMs.value()) { - TLLM_LOG_WARNING( - "Generation KV cache transfer for request %ld exceeded total timeout: " - "elapsed %ld ms > limit %d ms. Marking as error.", - request->mRequestId, elapsedMs, kvTransferTimeoutMs.value()); - // Erasing the std::future from mRequesterFutures only - // destroys the Python-facing handle; the std::async task - // spawned by CacheReceiver::receiveAsync is still running, - // blocked inside recvReadySignal waiting on a peer signal - // that will never arrive. Over a saturation window these - // zombies accumulate and pin NIXL/UCX per-request state, so - // subsequent receives queue behind a frozen pool and every - // new transfer deterministically waits kv_transfer_timeout_ms - // even after load stops (the "no-recovery" bug). Call - // cancelRequest to flip the per-request cancel flag in - // CacheReceiver::Impl, which the notification polling loop - // observes via the DataContext and returns early with - // isReady=false. requestSync then takes the - // kDISAGG_TRANS_ERROR branch and the task exits cleanly. - mCacheReceiver->cancelRequest(*request); - request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); - it = mRequesterFutures.erase(it); - ++numErrored; + bool const firstTimeout = mTimedOutRequesterIds.insert(request->mRequestId).second; + if (firstTimeout) + { + TLLM_LOG_WARNING( + "Generation KV cache transfer for request %ld exceeded total timeout: " + "elapsed %ld ms > limit %d ms. Requesting cancellation.", + request->mRequestId, elapsedMs, kvTransferTimeoutMs.value()); + // cancelRequest requests worker unwind, but it is not a + // quiescence proof. Keep the future tracked until the + // worker future becomes ready, otherwise Python could + // free KV resources while the worker/transport may still + // reference the advertised buffers. + mCacheReceiver->cancelRequest(*request); + } + if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + { + try + { + future.get(); + } + catch (std::exception const& e) + { + TLLM_LOG_WARNING( + "Generation KV cache transfer for timed-out request %ld finished with error: %s", + request->mRequestId, e.what()); + } + request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + mTimedOutRequesterIds.erase(request->mRequestId); + it = mRequesterFutures.erase(it); + ++numErrored; + } + else + { + if (firstTimeout) + { + ++numErrored; + } + ++it; + } continue; } } @@ -898,9 +915,20 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR if (status == std::future_status::ready) { future.get(); - completeEntry(request); + if (request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR) + { + TLLM_LOG_WARNING( + "Generation KV cache transfer for request %ld finished after an error state was set.", + request->mRequestId); + ++numErrored; + } + else + { + completeEntry(request); + ++numCompleted; + } + mTimedOutRequesterIds.erase(request->mRequestId); it = mRequesterFutures.erase(it); - ++numCompleted; } else if (status == std::future_status::timeout) { @@ -924,9 +952,20 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR // so the per-iteration poll timeout can drive the deadline // check each tick. future.get(); - completeEntry(request); + if (request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR) + { + TLLM_LOG_WARNING( + "Generation KV cache transfer for request %ld finished after an error state was set.", + request->mRequestId); + ++numErrored; + } + else + { + completeEntry(request); + ++numCompleted; + } + mTimedOutRequesterIds.erase(request->mRequestId); it = mRequesterFutures.erase(it); - ++numCompleted; } } else @@ -934,6 +973,7 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR TLLM_LOG_ERROR("Future returned unexpected status for generation request %ld. Marking as error.", request->mRequestId); request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + mTimedOutRequesterIds.erase(request->mRequestId); it = mRequesterFutures.erase(it); ++numErrored; } @@ -948,6 +988,7 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR TLLM_LOG_WARNING("Error during generation transfer for request %ld: %s. Marking as error.", request->mRequestId, e.what()); request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + mTimedOutRequesterIds.erase(request->mRequestId); it = mRequesterFutures.erase(it); ++numErrored; } diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index e535b6887027..1e26f43ee964 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -1277,7 +1277,14 @@ class CacheReceiver::Impl return isCancelled; } - bool receiveReadySignal(TransferSession& session, std::atomic const& perRequestCancel) + enum class ReadySignalResult + { + kReady, + kNotReady, + kCancelled, + }; + + ReadySignalResult receiveReadySignalDetailed(TransferSession& session, std::atomic const& perRequestCancel) { bool isReadyFinal = true; bool isReady = false; @@ -1290,7 +1297,7 @@ class CacheReceiver::Impl // rather than continuing to wait on subsequent peers. if (perRequestCancel.load(std::memory_order_relaxed)) { - return false; + return ReadySignalResult::kCancelled; } auto* agentConnectionManager = dynamic_cast(mManager); if (agentConnectionManager) @@ -1303,18 +1310,32 @@ class CacheReceiver::Impl // atomic and returns early when it flips — either on // process shutdown (all per-request flags flipped in // ~Impl()) or on per-request cancelRequest(). - isReady = agentConnection->recvReadySignal( + auto readyResult = agentConnection->recvReadySignalWithStatus( executor::kv_cache::DataContext{TransceiverTag::kREADY_SIGNAL_TAG, perRequestCancel}); + if (!readyResult.has_value()) + { + return ReadySignalResult::kCancelled; + } + isReady = readyResult.value(); } else { connections.at(i)->recv( executor::kv_cache::DataContext{TransceiverTag::kREADY_SIGNAL_TAG}, &isReady, sizeof(isReady)); + if (perRequestCancel.load(std::memory_order_relaxed)) + { + return ReadySignalResult::kCancelled; + } } isReadyFinal &= isReady; } - return isReadyFinal; + return isReadyFinal ? ReadySignalResult::kReady : ReadySignalResult::kNotReady; + } + + bool receiveReadySignal(TransferSession& session, std::atomic const& perRequestCancel) + { + return receiveReadySignalDetailed(session, perRequestCancel) == ReadySignalResult::kReady; } // Overload preserved for the (currently unused) std::async-based receiveAsync @@ -1401,37 +1422,61 @@ class CacheReceiver::Impl } } - auto session = sendRequestInfo(llmRequest, perRequestCancel, std::move(cacheBufferIds)); - session.setTime(TransferSession::kTimeRequestInfo); - // receiveReadySignal blocks inside AgentConnectionManager::waitForNotification's - // polling loop until the peer sends the ready notification OR the - // perRequestCancel flag flips. That's the one path that can actually - // release a zombie receive when a timeout-eviction fires on the C++ - // side; without it the std::async/worker task would stay blocked - // indefinitely and hold NIXL/UCX per-request state. - bool isReady = receiveReadySignal(session, perRequestCancel); - if (!isReady || perRequestCancel.load()) + auto poisonRecvHolders = [&recvHolders]() { - // Either the peer never sent the ready signal (timeout / cancel - // fired) or the cancel was observed after ready. Either way, - // reuse the error state — AsyncTransferManager cleanup will run. - llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + for (auto& h : recvHolders) + { + h.poison(); + } + }; + + try + { + auto session = sendRequestInfo(llmRequest, perRequestCancel, std::move(cacheBufferIds)); + session.setTime(TransferSession::kTimeRequestInfo); + // receiveReadySignal blocks inside AgentConnectionManager::waitForNotification's + // polling loop until the peer sends the ready notification OR the + // perRequestCancel flag flips. The result is intentionally + // tri-state: explicit peer not-ready means no data phase will + // write into the advertised receive buffers, while local cancel / + // no-notification means TRT-LLM cannot prove the peer is not still + // writing and must poison the advertised slots. + auto readyResult = receiveReadySignalDetailed(session, perRequestCancel); + if (readyResult == ReadySignalResult::kCancelled) + { + poisonRecvHolders(); + llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); + return; + } + if (readyResult == ReadySignalResult::kNotReady) + { + llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); + return; + } + + receiveSync(session); llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); - return; + + // Happy path only: the formatter invoked inside receiveSync already + // released each buffer index via freeBufferIndexForRecv. Detach the + // holders so they don't double-release when this stack frame + // unwinds. + for (auto& h : recvHolders) + { + (void) h.detach(); + } } - receiveSync(session); - llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); - - // Happy path only: the formatter invoked inside receiveSync already - // released each buffer index via freeBufferIndexForRecv. Detach the - // holders so they don't double-release when this stack frame - // unwinds. If control is about to leave requestSync via any OTHER - // path below (none today, but safe against future edits), the - // holders would still release automatically — detaching is strictly - // a correctness requirement for the just-completed receiveSync path. - for (auto& h : recvHolders) + catch (...) { - (void) h.detach(); + if (agentConnectionManagerForAcq) + { + poisonRecvHolders(); + } + llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); + throw; } TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp index 2d52b34fd6aa..cbf6245e9e65 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp @@ -302,11 +302,16 @@ void AgentConnection::sendReadySignal(DataContext const& ctx, bool isReady) cons } bool AgentConnection::recvReadySignal(DataContext const& ctx) const +{ + return recvReadySignalWithStatus(ctx).value_or(false); +} + +std::optional AgentConnection::recvReadySignalWithStatus(DataContext const& ctx) const { ReadySignalInfo readySignalInfo{mAgentName, ctx, false}; if (!mAgentConnectionManager->waitForReadySignal(mRemoteAgentName, readySignalInfo, ctx.getTransferTerminate())) { - return false; + return std::nullopt; } return readySignalInfo.mIsReady; } diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h index e32f4440b096..c912c557f52c 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h @@ -273,6 +273,7 @@ class AgentConnection : public Connection [[nodiscard]] bool hasLoadRemoteAgent() const; void sendReadySignal(DataContext const& ctx, bool isReady) const; bool recvReadySignal(DataContext const& ctx) const; + std::optional recvReadySignalWithStatus(DataContext const& ctx) const; void activateBuffer(uint8_t kind) const override; [[nodiscard]] std::optional getPreAssignedBufferId(uint8_t kind) const override; diff --git a/docs/source/features/disagg-kv-transfer-cancel-safety.md b/docs/source/features/disagg-kv-transfer-cancel-safety.md new file mode 100644 index 000000000000..a05dd0084625 --- /dev/null +++ b/docs/source/features/disagg-kv-transfer-cancel-safety.md @@ -0,0 +1,126 @@ +# Disaggregated KV transfer cancellation safety + +This note documents the cancellation contract added on top of the rc11 +disaggregated-transfer recovery work. The contract is intentionally +fail-closed: when TRT-LLM cannot prove transport quiescence, it must not return +advertised transfer buffers or in-flight KV resources to normal serving. + +## Problem + +`cancelRequest()` and backend-handle release are not a global quiescence proof. +For Agent/NIXL transfers, the receiver may have already advertised a GPU +receive buffer to the sender. If a timeout or Python cleanup path frees that +request while the sender or NIXL worker may still write, the write can land in a +buffer slot that has been reused for another request. + +The same issue appears on the sender side for cancellation while `sendAllBuffers` +is still using the preallocated send slot: releasing the slot back to the pool +is only safe after the transfer has completed or failed with a known quiescent +state. + +## Main / unsafe lifetime + +```mermaid +sequenceDiagram + participant Py as "Python executor" + participant CT as "CacheTransceiver" + participant CR as "CacheReceiver" + participant W as "Receiver worker" + participant N as "NIXL or UCX" + participant B as "Recv buffer slot" + + Py->>CR: "request_and_receive_async(request)" + CR->>W: "queue shared request and promise" + W->>B: "assign receive buffer" + W->>N: "advertise buffer address" + Py->>CT: "timeout or broad error cleanup" + CT->>CR: "cancelRequest(request)" + Py->>Py: "_terminate_request frees resources" + B->>B: "slot can be reused" + N-->>B: "late write may still arrive" +``` + +## PR 13713 lifetime + +PR 13713 fixes the raw `LlmRequest*` lifetime issue by keeping shared +references in the C++ tracker and worker paths. It also makes several waits +cancellable and adds RAII release for buffer-index slots. That prevents the +observed OOM/leak class, but RAII `release()` still returns a slot to the pool +even when transport quiescence is unknown. + +```mermaid +sequenceDiagram + participant Py as "Python executor" + participant CT as "CacheTransceiver" + participant CR as "CacheReceiver" + participant W as "Receiver worker" + participant B as "Recv buffer slot" + + Py->>CR: "request_and_receive_async(shared request)" + CR->>W: "worker owns shared request" + W->>B: "RAII holder acquires slot" + CT->>CR: "cancelRequest flips cancel flag" + CT->>CT: "future erased on timeout" + W->>B: "holder releases slot on early return" + Py->>Py: "request resources may be freed" +``` + +## This PR lifetime + +This PR separates three outcomes: + +1. Explicit peer not-ready: no data phase will write; the receive slot can be + released normally. +2. Worker future ready: C++ has observed worker quiescence; the tracker can be + erased. +3. Local cancel, timeout, missing notification, or exception while a buffer is + advertised: quiescence is unknown; poison the buffer pool and fail the + executor closed so orchestration restarts the pod. + +```mermaid +sequenceDiagram + participant Py as "Python executor" + participant CT as "CacheTransceiver" + participant CR as "CacheReceiver" + participant W as "Receiver worker" + participant B as "Recv buffer slot" + participant O as "Orchestrator" + + Py->>CR: "request_and_receive_async(shared request)" + W->>B: "acquire recv holder" + W->>CR: "send request and buffer info" + CT->>CR: "cancelRequest on timeout" + CT->>CT: "keep future until ready" + W->>W: "ready result is cancelled" + W->>B: "poison slot instead of release" + Py->>Py: "fail closed without freeing active KV resources" + O->>Py: "restart unhealthy pod" +``` + +## Code-level contract + +- `BufferIndexHolder::release()` is only for proven safe paths. +- `BufferIndexHolder::poison()` marks the corresponding send or receive pool + poisoned and does not return the slot to the pool. +- `BaseTransBufferManager::assignBufferIndex*()` fails once a pool is poisoned. + This prevents continued serving on memory whose transport quiescence is + unknown. +- `CacheReceiver::Impl::receiveReadySignalDetailed()` distinguishes explicit + peer not-ready from local cancellation or missing notification. +- `CacheReceiver::Impl::requestSync()` poisons advertised receive slots on + local cancel or exception. +- `CacheFormatter::format()` poisons an Agent send slot if `sendAllBuffers` + exits through an exception. +- Direct zero-copy sends from request-owned KV blocks are disabled for + cancellable disaggregated transfers until KV-block leases can prove sender + quiescence before the request is freed. +- `CacheTransceiver::checkGenTransferStatus()` no longer erases a generation + receive future on timeout until the worker future becomes ready. +- Python `_handle_errors()` and timeout cleanup fail closed for in-flight + disaggregated transfers instead of freeing active request resources while + quiescence is unknown. + +The result is not an in-process NIXL/UCX deadlock repair. It is a memory-safety +and operability boundary: no unsafe buffer reuse, no unbounded leak from +continuing to serve on poisoned transfer memory, and a clear restart signal for +health checks. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 420ff9bea535..3feaff9096f7 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -561,7 +561,7 @@ def on_detected(): self._disagg_pp_termination_handler = None if self.dist.pp_size > 1 and self.enable_kv_cache_reuse and self.kv_cache_transceiver: self._disagg_pp_termination_handler = DisaggPPTerminationHandler( - self.dist, self._do_terminate_request) + self.dist, self._do_terminate_request_if_safe) if self.dist.pp_size > 1: self.event_loop = self._executor_loop_pp @@ -3400,10 +3400,10 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): # If cancel is successful, mark as complete so it can be cleaned up # Otherwise, try at next iteration if is_cancelled: - request.py_kv_transfer_start_time = None - request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE - - self._end_transfer_and_maybe_terminate(request) + self._fail_closed_for_unquiesced_disagg_transfer( + f"Context request {request.py_request_id} timed out " + "while KV transfer quiescence is unknown", [request]) + return self._check_cache_transfer_errors("context requests") @@ -3596,10 +3596,73 @@ def _handle_errors(self, requests: Optional[List[LlmRequest]] = None): error_responses: Dict[int, LlmResponse] = {} error_msg = error_msg or "error" - failed_requests = requests if requests is not None else self.active_requests + + is_fatal = (self._error_budget.consume(error_msg) + if charge_budget else False) + if is_fatal and self._error_budget.budget < 1e-9: + logger.error(f"Error budget exhausted " + f"(budget={self._error_budget.budget:.3f}), " + "treating as fatal") + + if is_fatal: + self._fatal_error = RuntimeError(f"Fatal error: {error_msg}") + self.is_shutdown = True + logger.error( + f"Fatal error detected, initiating shutdown: {error_msg}") + requests = None + + # Drain waiting_queue so that queued-but-not-yet-activated + # requests don't get picked up on the next iteration. + # These are RequestQueueItems (not yet LlmRequests), so we + # fail them via error responses. Buffer all responses and + # call _enqueue_responses once after the loop so every rank + # enters the same number of collectives (attention-DP / + # gather-all modes use collective gathers internally). + waiting_responses: List[Tuple[int, LlmResponse]] = [] + while self.waiting_queue: + item = self.waiting_queue.pop_request() + if (self.gather_all_responses + or self.dist.rank == 0) and item.request is not None: + waiting_responses.append( + (item.id, + LlmResponse(request_id=item.id, + error_msg=error_msg, + client_id=getattr(item.request, + 'client_id', None)))) + # Also drain executor_request_queue so items already queued + # but not yet fetched by the main loop are not scheduled + # after the CUDA context is corrupted. Safe to use empty() + # here because is_shutdown is True and the queue's active + # flag is about to be set False, so no new items arrive. + raw_queue = self.executor_request_queue.get_request_queue() + while not raw_queue.empty(): + item = raw_queue.get_nowait() + if item.is_shutdown_request: + continue + if ((self.gather_all_responses or self.dist.rank == 0) + and item.request is not None): + waiting_responses.append( + (item.id, + LlmResponse(request_id=item.id, + error_msg=error_msg, + client_id=getattr(item.request, + 'client_id', None)))) + + if waiting_responses: + self._enqueue_responses(waiting_responses) + logger.info(f"Drained {len(waiting_responses)} queued requests " + "on fatal error") + + failed_requests = (list(self.active_requests) + if requests is None else list(requests)) + if self._has_unquiesced_disagg_transfers(failed_requests): + self._fail_closed_for_unquiesced_disagg_transfer( + error_msg, failed_requests) + return for request in failed_requests: req_id = request.py_request_id - request.state = LlmRequestState.GENERATION_COMPLETE + if not request.is_disagg_context_transmission_state: + request.state = LlmRequestState.GENERATION_COMPLETE error_responses[req_id] = LlmResponse( request_id=req_id, error_msg=error_msg, @@ -3615,7 +3678,73 @@ def _handle_errors(self, for request in failed_requests: self._terminate_request(request) - def _terminate_request(self, request: LlmRequest): + if self._fatal_error is not None: + self.executor_request_queue.enqueue_shutdown_request() + + def _has_unquiesced_disagg_transfers(self, + requests: Iterable[LlmRequest] + ) -> bool: + return any(req.is_disagg_generation_transmission_in_progress + or req.is_disagg_context_transmission_state + for req in requests) + + def _fail_closed_for_unquiesced_disagg_transfer( + self, error_msg: str, requests: Iterable[LlmRequest]): + failed_requests = list(requests) + in_flight_req_ids = [ + req.py_request_id for req in failed_requests + if req.is_disagg_generation_transmission_in_progress + or req.is_disagg_context_transmission_state + ] + logger.error( + "Fatal executor error while disaggregated KV transfer is in flight. " + "Failing closed without freeing active request resources because " + "transport quiescence is unknown. " + f"in_flight_request_ids={in_flight_req_ids}, error={error_msg}") + + error_responses: Dict[int, LlmResponse] = {} + response_requests = { + request.py_request_id: request + for request in list(self.active_requests) + failed_requests + } + for request in response_requests.values(): + req_id = request.py_request_id + error_responses[req_id] = LlmResponse( + request_id=req_id, + error_msg=error_msg, + client_id=request.py_client_id) + + self.active_requests.clear() + self.waiting_queue.clear() + self.request_accumulated.clear() + self.control_requests.clear() + self.is_shutdown = True + self._enqueue_responses(list(error_responses.items())) + with self.response_cv: + self.response_cv.notify_all() + + def _can_terminate_request_now(self, request: LlmRequest) -> bool: + if request.is_disagg_context_transmission_state: + logger.warning( + f"Deferring termination for request {request.py_request_id} " + "because context KV transfer is still in progress") + return False + if request.is_disagg_generation_transmission_in_progress: + self._fail_closed_for_unquiesced_disagg_transfer( + f"Attempted to terminate generation request {request.py_request_id} " + "while KV transfer is still in progress", [request]) + return False + return True + + def _do_terminate_request_if_safe(self, request: LlmRequest) -> bool: + if not self._can_terminate_request_now(request): + return False + self._do_terminate_request(request) + return True + + def _terminate_request(self, request: LlmRequest) -> bool: + if not self._can_terminate_request_now(request): + return False # Dummy requests don't participate in disagg KV cache transfers, # so they must bypass the PP termination handler to avoid stale # sequences in the KV cache manager (the handler delays removal, @@ -3625,6 +3754,7 @@ def _terminate_request(self, request: LlmRequest): self._disagg_pp_termination_handler.terminate(request) else: self._do_terminate_request(request) + return True def _do_terminate_request(self, request: LlmRequest): self.resource_manager.free_resources(request) @@ -3782,22 +3912,13 @@ def _handle_responses(self): continue # Check if generation request needs cleanup due to KV cache transfer timeout. - # - # The C++ transceiver's async workers and mSenderFutures / - # mRequesterFutures now hold std::shared_ptr (not raw - # pointers), so Python-side _terminate_request can safely run - # regardless of C++ state — the LlmRequest stays alive until all - # strong references (Python active_requests, C++ futures map, - # async workers) drop. The previous "defer cleanup until C++ - # evicts first" guard was what turned C++ tracker orphans into - # permanent KV-block leaks because the C++ deadline check could - # not reach orphaned entries. if request.py_kv_transfer_timed_out: is_cancelled = self.kv_cache_transceiver.cancel_request(request) if is_cancelled: - self._handle_errors( - error_msg=f"Request {request.py_request_id} timed out", - requests=[request]) + self._fail_closed_for_unquiesced_disagg_transfer( + f"Generation request {request.py_request_id} timed out " + "while KV transfer quiescence is unknown", [request]) + return requests_to_terminate continue if request.is_generation_only_request() and not request.is_finished: From 0ff78dd296f3796aa724882c6cc6f3efcbc0f1b6 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 4 May 2026 18:59:42 -0700 Subject: [PATCH 5/9] [https://nvbugs/6104831][fix] Incorporate PR#13728's improvements Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 17 ++- .../batch_manager/baseTransBuffer.h | 6 + .../batch_manager/cacheFormatter.cpp | 13 +- .../batch_manager/dataTransceiver.cpp | 17 ++- .../batch_manager/mlaCacheFormatter.cpp | 87 +++++++----- .../agent_utils/connection.h | 10 ++ .../disagg-kv-transfer-cancel-safety.md | 126 ------------------ tensorrt_llm/_torch/pyexecutor/py_executor.py | 6 +- 8 files changed, 104 insertions(+), 178 deletions(-) delete mode 100644 docs/source/features/disagg-kv-transfer-cancel-safety.md diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index fb58378dd4d0..091124595c54 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -286,10 +286,19 @@ class CacheTransceiver : public BaseCacheTransceiver // _terminate_request can drop its pybind shared_ptr while the C++ side's // raw pointer is still dereferenced by checkGenTransferStatus / // checkContextTransferStatus (the UAF forensically confirmed via - // MALLOC_PERTURB_=85 producing mRequestId=0x5555555555555555). Entries are - // erased only after completion, exception, or timeout plus worker-future - // readiness. A timeout/cancel request is not a quiescence proof, so the - // tracker must keep the shared_ptr until the worker future resolves. + // MALLOC_PERTURB_=85 producing mRequestId=0x5555555555555555). + // + // Eviction policy is asymmetric: + // - mRequesterFutures (gen side): on timeout, keep the entry tracked + // via mTimedOutRequesterIds until the worker future resolves. A + // timeout/cancel is not a quiescence proof on the recv side, so the + // advertised receive buffers may still be written to until the worker + // unwinds. See checkGenTransferStatus. + // - mSenderFutures (ctx side): erased immediately on completion, + // exception, or timeout. Sender zombies empirically unwind on peer + // teardown (decode-pod restart), and CacheSender::cancelRequest is + // only required to clear bookkeeping for telemetry / re-enqueue + // paths. See checkContextTransferStatus. std::vector, std::future>> mSenderFutures; std::vector, std::future>> mRequesterFutures; std::unordered_set mTimedOutRequesterIds; diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h index ad39cdcd7f19..d74ca7b52535 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h @@ -75,6 +75,9 @@ class BufferIndexHolder std::optional requestIdForLog = std::nullopt) noexcept : mMgr(&mgr) , mIndex(index) + // mHeld means this object is responsible for either releasing a + // concrete preallocated slot or poisoning dynamic transfer memory on an + // unsafe exit. It is therefore true even when index == std::nullopt. , mHeld(true) , mIsRecv(isRecv) , mRequestIdForLog(requestIdForLog) @@ -123,6 +126,9 @@ class BufferIndexHolder return mIndex; } + /// @brief Whether this holder still owns release/poison responsibility. + /// For dynamic-buffer paths, mIndex can be std::nullopt while this + /// remains true so poison() can still fail closed on unsafe exits. [[nodiscard]] bool held() const noexcept { return mHeld; diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index e697c6774ae6..f6b47440c50c 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -495,11 +495,14 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio // 5. send the buffer to the corresponding target. Ideally, we send only once (one buffer) for each target. // RAII wrapper mirrors the receiver-side pattern. Happy-path calls - // sendHolder.release() after sendAllBuffers returns; any other exit - // between acquire and release auto-releases via ~BufferIndexHolder. - // The send-pool CV wait inside assignBufferIndexForSend observes the - // session's per-request cancel flag and throws on cancel (parity - // with the recv side). + // sendHolder.release() after sendAllBuffers returns. If sendAllBuffers + // exits through an exception while an AgentConnection (NIXL) is + // active, transport quiescence is unknown and the catch block below + // poisons the held slot instead of returning it to the pool; on the + // direct-UCX path (no AgentConnection) the holder's destructor falls + // back to release. The send-pool CV wait inside + // assignBufferIndexForSend observes the session's per-request cancel + // flag and throws on cancel (parity with the recv side). auto const sendReqIdForLog = std::make_optional(static_cast(llmRequest.mRequestId)); auto const* sendCancelFlag = &session.getDataContext().getTransferTerminate(); auto cacheBufferId = mCacheTransBufferManager->assignBufferIndexForSend( diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 1e26f43ee964..9de26a87b152 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -990,9 +990,11 @@ class CacheReceiver::Impl // Overload that takes caller-acquired buffer indices (wrapped in // BufferIndexHolders at the call site). requestSync uses this variant so - // RAII covers all non-happy-path exits (not-ready, cancel, throw). The - // preAcquiredCacheBufferIds vector must have one entry per cache - // transfer buffer manager (aligned with + // every non-happy-path exit is covered: explicit not-ready releases via + // ~BufferIndexHolder, while cancel and exception poison the holders + // before the stack unwinds (transport quiescence is unknown on those + // paths). The preAcquiredCacheBufferIds vector must have one entry per + // cache transfer buffer manager (aligned with // AgentConnectionManager::getCacheTransBufferManagers()). TransferSession sendRequestInfo(LlmRequest const& llmRequest, std::atomic const& perRequestCancel, std::vector> preAcquiredCacheBufferIds) @@ -1037,9 +1039,12 @@ class CacheReceiver::Impl if (!preAcquiredCacheBufferIds.empty()) { // requestSync already acquired these indices under RAII - // holders — use them as-is so receiveSync's formatter - // releases via the existing path while any non-happy-path - // exit from requestSync releases via ~BufferIndexHolder. + // holders — use them as-is. On the happy path receiveSync's + // formatter releases each slot via freeBufferIndexForRecv + // and requestSync detaches the holder. On cancel / exception + // requestSync explicitly poisons the holders before the + // stack unwinds; only the explicit not-ready path falls + // through to ~BufferIndexHolder for a normal release. cacheBufferIds = std::move(preAcquiredCacheBufferIds); TLLM_CHECK(!cacheBufferIds.empty()); } diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp index 47fc8d400c10..9e95ef336a97 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp @@ -210,6 +210,10 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses TLLM_LOG_DEBUG("Try using zero-copy for the KV cache."); NVTX3_SCOPED_RANGE(sendBufferFun); + TLLM_CHECK_WITH_INFO(false, + "Zero-copy KV cache transfer sends directly from request-owned KV blocks. It is disabled for " + "cancellable disaggregated transfers until KV-block leases can prove the peer has stopped reading " + "from those blocks before the request is freed."); TLLM_CUDA_CHECK(cudaSetDevice(deviceId)); for (size_t i = 0; i < pickUpConnections.size(); i++) @@ -255,10 +259,14 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses }; auto bufferEleSizes = getBufferSizeForTarget(); // RAII wrapper mirrors the receiver-side pattern. Happy-path calls - // sendHolder.release(); other exits auto-release via - // ~BufferIndexHolder. Send-pool CV wait observes the per-request - // cancel flag via the session's DataContext and throws on cancel - // (parity with recv side). + // sendHolder.release() after the send loop returns. If the send + // loop exits through an exception while an AgentConnection (NIXL) + // is active, transport quiescence is unknown and the catch block + // below poisons the held slot instead of returning it to the pool; + // on the direct-UCX path (no AgentConnection) the holder's + // destructor falls back to release. The send-pool CV wait inside + // assignBufferIndexForSend observes the session's per-request + // cancel flag and throws on cancel (parity with the recv side). auto const sendReqIdForLog = std::make_optional(static_cast(llmRequest.mRequestId)); auto const* sendCancelFlag = &session.getDataContext().getTransferTerminate(); auto cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForSend( @@ -348,48 +356,59 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses session.appendMeasure(startTime, endTime, outputSplitCaches.at(cacheIdx)->getSizeInBytes()); }; - if (pickUpConnections.size() > 1) + try { - if (!common::getEnvEnableReceiveKVCacheParallel()) + if (pickUpConnections.size() > 1) { - TLLM_LOG_DEBUG("Disable parallel receiving of the KV cache."); - for (size_t i = 0; i < pickUpConnections.size(); i++) + if (!common::getEnvEnableReceiveKVCacheParallel()) { - sendBufferFun(deviceId, pickUpConnections[i]); + TLLM_LOG_DEBUG("Disable parallel receiving of the KV cache."); + for (size_t i = 0; i < pickUpConnections.size(); i++) + { + sendBufferFun(deviceId, pickUpConnections[i]); + } } - } - else - { - // concurrency num - auto concurrencyNum - = std::min(std::max(static_cast(1), bufferCoverTargetNum), pPDomainSize * cPDomainSize); + else + { + // concurrency num + auto concurrencyNum + = std::min(std::max(static_cast(1), bufferCoverTargetNum), pPDomainSize * cPDomainSize); - auto remainSendNum = pickUpConnections.size(); + auto remainSendNum = pickUpConnections.size(); - while (remainSendNum > 0) - { - auto sendConcurrencyNum = std::min(remainSendNum, concurrencyNum); - std::vector> futures; - futures.reserve(sendConcurrencyNum); - for (size_t i = 0; i < sendConcurrencyNum; i++) - { - size_t idx = i + (pickUpConnections.size() - remainSendNum); - size_t connIdx = pickUpConnections[idx]; - TLLM_CHECK(idx < pickUpConnections.size()); - TLLM_CHECK(connIdx < session.getConnections().size()); - futures.push_back(std::async(std::launch::async, sendBufferFun, deviceId, connIdx)); - } - for (auto& future : futures) + while (remainSendNum > 0) { - future.get(); + auto sendConcurrencyNum = std::min(remainSendNum, concurrencyNum); + std::vector> futures; + futures.reserve(sendConcurrencyNum); + for (size_t i = 0; i < sendConcurrencyNum; i++) + { + size_t idx = i + (pickUpConnections.size() - remainSendNum); + size_t connIdx = pickUpConnections[idx]; + TLLM_CHECK(idx < pickUpConnections.size()); + TLLM_CHECK(connIdx < session.getConnections().size()); + futures.push_back(std::async(std::launch::async, sendBufferFun, deviceId, connIdx)); + } + for (auto& future : futures) + { + future.get(); + } + remainSendNum -= sendConcurrencyNum; } - remainSendNum -= sendConcurrencyNum; } } + else + { + sendBufferFun(deviceId, pickUpConnections[0]); + } } - else + catch (...) { - sendBufferFun(deviceId, pickUpConnections[0]); + if (agentConnection != nullptr) + { + sendHolder.poison(); + } + throw; } // Atomic happy-path release — silent free + disarm in one noexcept call. sendHolder.release(); diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h index c912c557f52c..ce43cf096094 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h @@ -272,7 +272,17 @@ class AgentConnection : public Connection void setHasLoadRemoteAgent(bool hasLoadRemoteAgent); [[nodiscard]] bool hasLoadRemoteAgent() const; void sendReadySignal(DataContext const& ctx, bool isReady) const; + /// @brief Wait for the peer's ready signal. Returns the peer's ready + /// bit, collapsing cancel/no-signal into `false`. Callers that + /// need to distinguish "peer said not-ready" from "wait was + /// cancelled" must use `recvReadySignalWithStatus`. bool recvReadySignal(DataContext const& ctx) const; + /// @brief Tri-state variant of `recvReadySignal`. Returns the peer's + /// ready bit when a signal arrived, or `std::nullopt` if the + /// wait unwound without a notification (e.g. the per-request + /// cancel flag flipped before the peer sent ready). Callers + /// treat `std::nullopt` as "transport quiescence is unknown" + /// and must fail closed (poison receive buffers, etc.). std::optional recvReadySignalWithStatus(DataContext const& ctx) const; void activateBuffer(uint8_t kind) const override; diff --git a/docs/source/features/disagg-kv-transfer-cancel-safety.md b/docs/source/features/disagg-kv-transfer-cancel-safety.md deleted file mode 100644 index a05dd0084625..000000000000 --- a/docs/source/features/disagg-kv-transfer-cancel-safety.md +++ /dev/null @@ -1,126 +0,0 @@ -# Disaggregated KV transfer cancellation safety - -This note documents the cancellation contract added on top of the rc11 -disaggregated-transfer recovery work. The contract is intentionally -fail-closed: when TRT-LLM cannot prove transport quiescence, it must not return -advertised transfer buffers or in-flight KV resources to normal serving. - -## Problem - -`cancelRequest()` and backend-handle release are not a global quiescence proof. -For Agent/NIXL transfers, the receiver may have already advertised a GPU -receive buffer to the sender. If a timeout or Python cleanup path frees that -request while the sender or NIXL worker may still write, the write can land in a -buffer slot that has been reused for another request. - -The same issue appears on the sender side for cancellation while `sendAllBuffers` -is still using the preallocated send slot: releasing the slot back to the pool -is only safe after the transfer has completed or failed with a known quiescent -state. - -## Main / unsafe lifetime - -```mermaid -sequenceDiagram - participant Py as "Python executor" - participant CT as "CacheTransceiver" - participant CR as "CacheReceiver" - participant W as "Receiver worker" - participant N as "NIXL or UCX" - participant B as "Recv buffer slot" - - Py->>CR: "request_and_receive_async(request)" - CR->>W: "queue shared request and promise" - W->>B: "assign receive buffer" - W->>N: "advertise buffer address" - Py->>CT: "timeout or broad error cleanup" - CT->>CR: "cancelRequest(request)" - Py->>Py: "_terminate_request frees resources" - B->>B: "slot can be reused" - N-->>B: "late write may still arrive" -``` - -## PR 13713 lifetime - -PR 13713 fixes the raw `LlmRequest*` lifetime issue by keeping shared -references in the C++ tracker and worker paths. It also makes several waits -cancellable and adds RAII release for buffer-index slots. That prevents the -observed OOM/leak class, but RAII `release()` still returns a slot to the pool -even when transport quiescence is unknown. - -```mermaid -sequenceDiagram - participant Py as "Python executor" - participant CT as "CacheTransceiver" - participant CR as "CacheReceiver" - participant W as "Receiver worker" - participant B as "Recv buffer slot" - - Py->>CR: "request_and_receive_async(shared request)" - CR->>W: "worker owns shared request" - W->>B: "RAII holder acquires slot" - CT->>CR: "cancelRequest flips cancel flag" - CT->>CT: "future erased on timeout" - W->>B: "holder releases slot on early return" - Py->>Py: "request resources may be freed" -``` - -## This PR lifetime - -This PR separates three outcomes: - -1. Explicit peer not-ready: no data phase will write; the receive slot can be - released normally. -2. Worker future ready: C++ has observed worker quiescence; the tracker can be - erased. -3. Local cancel, timeout, missing notification, or exception while a buffer is - advertised: quiescence is unknown; poison the buffer pool and fail the - executor closed so orchestration restarts the pod. - -```mermaid -sequenceDiagram - participant Py as "Python executor" - participant CT as "CacheTransceiver" - participant CR as "CacheReceiver" - participant W as "Receiver worker" - participant B as "Recv buffer slot" - participant O as "Orchestrator" - - Py->>CR: "request_and_receive_async(shared request)" - W->>B: "acquire recv holder" - W->>CR: "send request and buffer info" - CT->>CR: "cancelRequest on timeout" - CT->>CT: "keep future until ready" - W->>W: "ready result is cancelled" - W->>B: "poison slot instead of release" - Py->>Py: "fail closed without freeing active KV resources" - O->>Py: "restart unhealthy pod" -``` - -## Code-level contract - -- `BufferIndexHolder::release()` is only for proven safe paths. -- `BufferIndexHolder::poison()` marks the corresponding send or receive pool - poisoned and does not return the slot to the pool. -- `BaseTransBufferManager::assignBufferIndex*()` fails once a pool is poisoned. - This prevents continued serving on memory whose transport quiescence is - unknown. -- `CacheReceiver::Impl::receiveReadySignalDetailed()` distinguishes explicit - peer not-ready from local cancellation or missing notification. -- `CacheReceiver::Impl::requestSync()` poisons advertised receive slots on - local cancel or exception. -- `CacheFormatter::format()` poisons an Agent send slot if `sendAllBuffers` - exits through an exception. -- Direct zero-copy sends from request-owned KV blocks are disabled for - cancellable disaggregated transfers until KV-block leases can prove sender - quiescence before the request is freed. -- `CacheTransceiver::checkGenTransferStatus()` no longer erases a generation - receive future on timeout until the worker future becomes ready. -- Python `_handle_errors()` and timeout cleanup fail closed for in-flight - disaggregated transfers instead of freeing active request resources while - quiescence is unknown. - -The result is not an in-process NIXL/UCX deadlock repair. It is a memory-safety -and operability boundary: no unsafe buffer reuse, no unbounded leak from -continuing to serve on poisoned transfer memory, and a clear restart signal for -health checks. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 3feaff9096f7..d49bc2099dc7 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3681,9 +3681,8 @@ def _handle_errors(self, if self._fatal_error is not None: self.executor_request_queue.enqueue_shutdown_request() - def _has_unquiesced_disagg_transfers(self, - requests: Iterable[LlmRequest] - ) -> bool: + def _has_unquiesced_disagg_transfers( + self, requests: Iterable[LlmRequest]) -> bool: return any(req.is_disagg_generation_transmission_in_progress or req.is_disagg_context_transmission_state for req in requests) @@ -3719,6 +3718,7 @@ def _fail_closed_for_unquiesced_disagg_transfer( self.request_accumulated.clear() self.control_requests.clear() self.is_shutdown = True + self.shutdown_event.set() self._enqueue_responses(list(error_responses.items())) with self.response_cv: self.response_cv.notify_all() From a1b82a89b78258990b870a5a3b388c824df64539 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Wed, 6 May 2026 10:50:25 -0700 Subject: [PATCH 6/9] Narrow Python disagg transfer timeout handling Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 105 ++++++++---------- 1 file changed, 45 insertions(+), 60 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index d49bc2099dc7..ad7bf0021d99 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -373,6 +373,7 @@ def __init__( ResourceManagerType.DRAFT_KV_CACHE_MANAGER: set(), } self._disagg_gen_kv_recv_started_ids: set[int] = set() + self._disagg_timed_out_kv_cancelled_ids: set[int] = set() # Rolling acceptance tracking for spec decode (disable speculation if rolling acceptance is below threshold) spec_config = getattr(self.model_engine, 'spec_config', None) @@ -3400,10 +3401,10 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): # If cancel is successful, mark as complete so it can be cleaned up # Otherwise, try at next iteration if is_cancelled: - self._fail_closed_for_unquiesced_disagg_transfer( - f"Context request {request.py_request_id} timed out " - "while KV transfer quiescence is unknown", [request]) - return + request.py_kv_transfer_start_time = None + request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE + + self._end_transfer_and_maybe_terminate(request) self._check_cache_transfer_errors("context requests") @@ -3655,10 +3656,26 @@ def _handle_errors(self, failed_requests = (list(self.active_requests) if requests is None else list(requests)) - if self._has_unquiesced_disagg_transfers(failed_requests): - self._fail_closed_for_unquiesced_disagg_transfer( - error_msg, failed_requests) - return + deferred_requests = [ + request for request in failed_requests + if self._is_unquiesced_disagg_transfer(request) + ] + if deferred_requests: + logger.warning( + "Deferring error cleanup for disaggregated KV transfers still " + "in flight. C++ transfer status remains the source of truth " + "for final cleanup. request_ids=%s, error=%s", + [request.py_request_id for request in deferred_requests], + error_msg) + failed_requests = [ + request for request in failed_requests + if request not in deferred_requests + ] + if not failed_requests: + if self._fatal_error is not None: + self.executor_request_queue.enqueue_shutdown_request() + return + for request in failed_requests: req_id = request.py_request_id if not request.is_disagg_context_transmission_state: @@ -3668,11 +3685,11 @@ def _handle_errors(self, error_msg=error_msg, client_id=request.py_client_id) if requests is None: - self.active_requests.clear() + self.active_requests = deferred_requests else: self.active_requests = [ request for request in self.active_requests - if request not in requests + if request not in failed_requests ] self._enqueue_responses(list(error_responses.items())) for request in failed_requests: @@ -3681,47 +3698,9 @@ def _handle_errors(self, if self._fatal_error is not None: self.executor_request_queue.enqueue_shutdown_request() - def _has_unquiesced_disagg_transfers( - self, requests: Iterable[LlmRequest]) -> bool: - return any(req.is_disagg_generation_transmission_in_progress - or req.is_disagg_context_transmission_state - for req in requests) - - def _fail_closed_for_unquiesced_disagg_transfer( - self, error_msg: str, requests: Iterable[LlmRequest]): - failed_requests = list(requests) - in_flight_req_ids = [ - req.py_request_id for req in failed_requests - if req.is_disagg_generation_transmission_in_progress - or req.is_disagg_context_transmission_state - ] - logger.error( - "Fatal executor error while disaggregated KV transfer is in flight. " - "Failing closed without freeing active request resources because " - "transport quiescence is unknown. " - f"in_flight_request_ids={in_flight_req_ids}, error={error_msg}") - - error_responses: Dict[int, LlmResponse] = {} - response_requests = { - request.py_request_id: request - for request in list(self.active_requests) + failed_requests - } - for request in response_requests.values(): - req_id = request.py_request_id - error_responses[req_id] = LlmResponse( - request_id=req_id, - error_msg=error_msg, - client_id=request.py_client_id) - - self.active_requests.clear() - self.waiting_queue.clear() - self.request_accumulated.clear() - self.control_requests.clear() - self.is_shutdown = True - self.shutdown_event.set() - self._enqueue_responses(list(error_responses.items())) - with self.response_cv: - self.response_cv.notify_all() + def _is_unquiesced_disagg_transfer(self, request: LlmRequest) -> bool: + return (request.is_disagg_generation_transmission_in_progress + or request.is_disagg_context_transmission_state) def _can_terminate_request_now(self, request: LlmRequest) -> bool: if request.is_disagg_context_transmission_state: @@ -3730,9 +3709,9 @@ def _can_terminate_request_now(self, request: LlmRequest) -> bool: "because context KV transfer is still in progress") return False if request.is_disagg_generation_transmission_in_progress: - self._fail_closed_for_unquiesced_disagg_transfer( - f"Attempted to terminate generation request {request.py_request_id} " - "while KV transfer is still in progress", [request]) + logger.warning( + f"Deferring termination for request {request.py_request_id} " + "because generation KV transfer is still in progress") return False return True @@ -3768,6 +3747,7 @@ def _do_terminate_request(self, request: LlmRequest): for prepared_ids in self._disagg_gen_init_prepared_ids.values(): prepared_ids.discard(request.py_request_id) self._disagg_gen_kv_recv_started_ids.discard(request.py_request_id) + self._disagg_timed_out_kv_cancelled_ids.discard(request.py_request_id) def _is_request_in_transmission(self, request) -> bool: """Check if a request is currently in transmission state.""" @@ -3913,12 +3893,17 @@ def _handle_responses(self): # Check if generation request needs cleanup due to KV cache transfer timeout. if request.py_kv_transfer_timed_out: - is_cancelled = self.kv_cache_transceiver.cancel_request(request) - if is_cancelled: - self._fail_closed_for_unquiesced_disagg_transfer( - f"Generation request {request.py_request_id} timed out " - "while KV transfer quiescence is unknown", [request]) - return requests_to_terminate + if (request.py_request_id + not in self._disagg_timed_out_kv_cancelled_ids): + is_cancelled = self.kv_cache_transceiver.cancel_request( + request) + if is_cancelled: + self._disagg_timed_out_kv_cancelled_ids.add( + request.py_request_id) + logger.warning( + f"Cancelled timed-out generation KV transfer for " + f"request {request.py_request_id}; waiting for " + "C++ transfer status to report final cleanup") continue if request.is_generation_only_request() and not request.is_finished: From 3c1ac2c156a149efe439b2b2b62c5d621ba7b7d4 Mon Sep 17 00:00:00 2001 From: Yifan Jiang Date: Tue, 5 May 2026 11:13:41 -0700 Subject: [PATCH 7/9] Defer context transfer cleanup after timeout cancel Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> # Conflicts: # tensorrt_llm/_torch/pyexecutor/py_executor.py --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 48 +++++++------------ 1 file changed, 18 insertions(+), 30 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index ad7bf0021d99..00ea4d766983 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -373,7 +373,8 @@ def __init__( ResourceManagerType.DRAFT_KV_CACHE_MANAGER: set(), } self._disagg_gen_kv_recv_started_ids: set[int] = set() - self._disagg_timed_out_kv_cancelled_ids: set[int] = set() + self._disagg_timed_out_ctx_cancelled_ids: set[int] = set() + self._disagg_timed_out_gen_cancelled_ids: set[int] = set() # Rolling acceptance tracking for spec decode (disable speculation if rolling acceptance is below threshold) spec_config = getattr(self.model_engine, 'spec_config', None) @@ -3388,23 +3389,18 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): for request_id in list(requests_in_transfer.keys()): request = requests_in_transfer[request_id] if request.py_kv_transfer_timed_out and request_id not in completed_req_ids: - # The previous "defer Python cleanup while - # DISAGG_CONTEXT_TRANS_IN_PROGRESS" guard existed to avoid a - # UAF on CacheTransceiver::mSenderFutures's raw LlmRequest*. - # mSenderFutures now holds std::shared_ptr, so the - # LlmRequest outlives every C++ access regardless of when - # Python drops its pybind reference — the UAF class is gone - # and the guard is no longer needed. Running cancel_request + - # end_transfer here recovers the KV blocks promptly even when - # the C++ deadline check cannot reach the entry (orphan class). - is_cancelled = self.kv_cache_transceiver.cancel_request(request) - # If cancel is successful, mark as complete so it can be cleaned up - # Otherwise, try at next iteration - if is_cancelled: - request.py_kv_transfer_start_time = None - request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE - - self._end_transfer_and_maybe_terminate(request) + # Ask C++ to cancel once, but do not free Python-side + # resources until C++ reports the transfer in finished/error + # status. cancel_request() is not a transport-quiescence proof. + if request_id not in self._disagg_timed_out_ctx_cancelled_ids: + is_cancelled = self.kv_cache_transceiver.cancel_request( + request) + if is_cancelled: + self._disagg_timed_out_ctx_cancelled_ids.add(request_id) + logger.warning( + f"Cancelled timed-out context KV transfer for " + f"request {request.py_request_id}; waiting for " + "C++ transfer status to report final cleanup") self._check_cache_transfer_errors("context requests") @@ -3747,7 +3743,8 @@ def _do_terminate_request(self, request: LlmRequest): for prepared_ids in self._disagg_gen_init_prepared_ids.values(): prepared_ids.discard(request.py_request_id) self._disagg_gen_kv_recv_started_ids.discard(request.py_request_id) - self._disagg_timed_out_kv_cancelled_ids.discard(request.py_request_id) + self._disagg_timed_out_ctx_cancelled_ids.discard(request.py_request_id) + self._disagg_timed_out_gen_cancelled_ids.discard(request.py_request_id) def _is_request_in_transmission(self, request) -> bool: """Check if a request is currently in transmission state.""" @@ -3894,11 +3891,11 @@ def _handle_responses(self): # Check if generation request needs cleanup due to KV cache transfer timeout. if request.py_kv_transfer_timed_out: if (request.py_request_id - not in self._disagg_timed_out_kv_cancelled_ids): + not in self._disagg_timed_out_gen_cancelled_ids): is_cancelled = self.kv_cache_transceiver.cancel_request( request) if is_cancelled: - self._disagg_timed_out_kv_cancelled_ids.add( + self._disagg_timed_out_gen_cancelled_ids.add( request.py_request_id) logger.warning( f"Cancelled timed-out generation KV transfer for " @@ -3971,15 +3968,6 @@ def _handle_responses(self): # If partial reuse is enabled, and the KV cache manager is not VSWA, and the PP size is 1, # then we need to terminate the request. TODO: Remove this once disagg support from KVCache reuse # path is fixed. - # Note: the previous "elif request.is_disagg_context_transmission_state: - # pass" guard existed to avoid a UAF on the C++ - # CacheTransceiver::mSenderFutures's raw LlmRequest*. Since - # mSenderFutures now holds std::shared_ptr, the - # LlmRequest outlives every C++ access regardless of when - # Python drops its pybind reference, so the guard is no - # longer needed and was actively turning C++ tracker - # orphans into permanent KV-block leaks (the deadline check - # cannot reach orphaned entries). force_terminate_for_partial_reuse = ( self.enable_partial_reuse_for_disagg and not self.kv_cache_manager.is_vswa From 66746088b403d6cc2045d9dd7da571a4846a3dac Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Wed, 6 May 2026 22:58:05 -0700 Subject: [PATCH 8/9] [https://nvbugs/6104831][fix] Complete deferred disagg cleanup after transfer Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.cpp | 24 ---- tensorrt_llm/_torch/pyexecutor/py_executor.py | 88 +++++++++++--- .../executor/test_async_transfer_manager.py | 109 +++++++++++++++++- 3 files changed, 176 insertions(+), 45 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 15090eb765f0..f086a7acf9d8 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -351,8 +351,6 @@ void CacheTransceiver::respondAndSendAsync(std::shared_ptr llmReques } setContextState(llmRequest.get()); auto future = mCacheSender->sendAsync(llmRequest); - TLLM_LOG_DEBUG("respondAndSendAsync: adding request %ld to mSenderFutures (size=%zu)", llmRequest->mRequestId, - mSenderFutures.size() + 1); mSenderFutures.emplace_back(std::move(llmRequest), std::move(future)); } @@ -406,8 +404,6 @@ void CacheTransceiver::requestAndReceiveAsync(std::shared_ptr llmReq // slightly later time — harmless for the deadline check. llmRequest->setKvCacheTransferStart(LlmRequest::getSteadyClockNow()); auto future = mCacheReceiver->receiveAsync(llmRequest); - TLLM_LOG_DEBUG("requestAndReceiveAsync: adding request %ld to mRequesterFutures (size=%zu)", llmRequest->mRequestId, - mRequesterFutures.size() + 1); // Order: setKvCacheTransferStart -> receiveAsync -> setState -> emplace. // setState is intentionally moved AFTER receiveAsync (the original code // set it before). This is safe today because the async worker spawned @@ -533,15 +529,6 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } } - // Log mSenderFutures state for diagnosing dangling pointer issues. - // Each entry's pointer address and request ID are logged so we can detect - // when a pointer's underlying memory is freed (reqId changes to 0). - if (!mSenderFutures.empty()) - { - TLLM_LOG_DEBUG("checkContextTransferStatus: mSenderFutures.size()=%zu, blockAll=%d, kvTransferTimeoutMs=%d", - mSenderFutures.size(), blockAll ? 1 : 0, kvTransferTimeoutMs.value_or(-1)); - } - auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupTPInDPComm : mGroupTensorParaComm; std::vector contextCompleteRequestIds; for (auto&& [request, future] : mSenderFutures) @@ -731,17 +718,6 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR } } - // Log mRequesterFutures state for diagnosing dangling pointer / stuck - // receiver issues. Mirrors the diagnostic added to checkContextTransferStatus. - if (!mRequesterFutures.empty()) - { - TLLM_LOG_DEBUG( - "checkGenTransferStatus: mRequesterFutures.size()=%zu, blockAll=%d, kvTransferTimeoutMs=%d, " - "senderFutureTimeoutMs=%d", - mRequesterFutures.size(), blockAll ? 1 : 0, kvTransferTimeoutMs.value_or(-1), - senderFutureTimeoutMs.value_or(-1)); - } - std::vector genTransferReadyRequestIds; for (auto&& [request, future] : mRequesterFutures) { diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 00ea4d766983..4e04298cd0eb 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -162,10 +162,19 @@ class AsyncTransferManager: """ class RequestTransferMetadata: + """STOP-GAP metadata for rc13 partial-reuse cleanup. + + `termination_requested` and `resources_freed` bridge the current split + where AsyncTransferManager owns transfer pinning, while PyExecutor owns + final resource cleanup. Remove them with the follow-up KV reuse lease / + cleanup-session work that replaces should_store_blocks + pin=True. + """ def __init__(self, block_id: Optional[int]): self.block_id = block_id self.counter = 0 + self.termination_requested = False + self.resources_freed = False def start_transfer(self): self.counter += 1 @@ -178,6 +187,18 @@ def end_transfer(self) -> bool: self.counter -= 1 return self.counter == 0 + @dataclasses.dataclass + class EndTransferResult: + """STOP-GAP result for deferred partial-reuse termination. + + `needs_termination` carries a deferred cleanup obligation from the + transfer owner back to PyExecutor. Remove it with the follow-up KV + reuse lease / cleanup-session work. + """ + + completed: bool + needs_termination: bool = False + def __init__(self, resource_manager: "ResourceManager", should_store_blocks: bool = True): @@ -197,6 +218,20 @@ def __init__(self, def requests_in_transfer(self) -> Dict[int, LlmRequest]: return self._requests_in_transfer + def mark_termination_requested(self, request: LlmRequest): + metadata = self._request_transfer_metadata.get(request.py_request_id) + if metadata is not None: + metadata.termination_requested = True + else: + logger.debug( + f"Termination requested for request {request.py_request_id}, " + "but it is not tracked by AsyncTransferManager") + + def mark_resources_freed(self, request: LlmRequest): + metadata = self._request_transfer_metadata.get(request.py_request_id) + if metadata is not None: + metadata.resources_freed = True + def start_transfer(self, request: LlmRequest): """ Called when a Cache transceiver or connector transfer is started. @@ -230,14 +265,15 @@ def start_transfer(self, request: LlmRequest): self._request_transfer_metadata[req_id].start_transfer() - def end_transfer(self, request: LlmRequest) -> bool: + def end_transfer(self, request: LlmRequest) -> EndTransferResult: """ Called after a send of KV cache is complete. 1. Decrements counter for request. 2. If there are no more inflight transfers for this request, unpin the blocks and mark the request as complete. Returns: - bool: True if the request should be terminated after call to end_transfer + EndTransferResult: whether the transfer is complete and whether a + deferred termination must be retried now. """ try: transfer_metadata = self._request_transfer_metadata[ @@ -246,9 +282,11 @@ def end_transfer(self, request: LlmRequest) -> bool: logger.warning( f"Request {request.py_request_id} not found in transfer manager" ) - return False + return self.EndTransferResult(completed=False) if transfer_metadata.end_transfer(): + needs_termination = (transfer_metadata.termination_requested + and not transfer_metadata.resources_freed) self._requests_in_transfer.pop(request.py_request_id) self._request_transfer_metadata.pop(request.py_request_id) @@ -260,9 +298,10 @@ def end_transfer(self, request: LlmRequest) -> bool: if request.state != LlmRequestState.DISAGG_TRANS_ERROR: request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE - return True + return self.EndTransferResult(completed=True, + needs_termination=needs_termination) - return False + return self.EndTransferResult(completed=False) def has_any_inflight_requests(self) -> bool: return len(self._requests_in_transfer) > 0 @@ -643,17 +682,28 @@ def _end_transfer_and_maybe_terminate(self, request: LlmRequest): response = request.create_response(False, self.dist.rank) if response: response.result.cached_tokens = request.cached_tokens - self._enqueue_responses([(request.py_request_id, response)]) - if self.async_transfer_manager.end_transfer(request): + # Buffer the response instead of enqueueing immediately. + # With ADP, _enqueue_responses does a tp_gather collective. + # Calling it here would deadlock because only the owning DP + # rank reaches this point; the other DP rank never enters + # the matching collective. The buffer is flushed later at + # _flush_pending_transfer_responses where all ranks + # participate. + self._pending_transfer_responses.append( + (request.py_request_id, response)) + transfer_result = self.async_transfer_manager.end_transfer(request) + if transfer_result.completed: self.active_requests.remove(request) self._terminate_request(request) return - if self.async_transfer_manager.end_transfer(request): - # When should_store_blocks is True, _handle_responses already - # terminated this request via the early-termination path - # (enable_partial_reuse_for_disagg branch). Skip the redundant - # termination to avoid double free_resources calls. - if not self.async_transfer_manager.should_store_blocks: + transfer_result = self.async_transfer_manager.end_transfer(request) + if transfer_result.completed: + # When should_store_blocks is True, _handle_responses normally owns + # the early termination path. However, termination can be deferred + # while context transfer is still in progress; retry it now that + # C++ reports the transfer terminal. + if (not self.async_transfer_manager.should_store_blocks + or transfer_result.needs_termination): self._terminate_request(request) # Performance metrics methods are in PerfMetricsManager (self.perf_manager) @@ -3700,12 +3750,19 @@ def _is_unquiesced_disagg_transfer(self, request: LlmRequest) -> bool: def _can_terminate_request_now(self, request: LlmRequest) -> bool: if request.is_disagg_context_transmission_state: - logger.warning( + # STOP-GAP: partial reuse asks _handle_responses to terminate + # completed context requests before async KV transfer has always + # drained. Remember the failed attempt so + # _end_transfer_and_maybe_terminate retries exactly once after C++ + # reports transfer completion. Remove this handoff with the + # follow-up KV reuse lease / cleanup-session work. + self.async_transfer_manager.mark_termination_requested(request) + logger.debug( f"Deferring termination for request {request.py_request_id} " "because context KV transfer is still in progress") return False if request.is_disagg_generation_transmission_in_progress: - logger.warning( + logger.debug( f"Deferring termination for request {request.py_request_id} " "because generation KV transfer is still in progress") return False @@ -3733,6 +3790,7 @@ def _terminate_request(self, request: LlmRequest) -> bool: def _do_terminate_request(self, request: LlmRequest): self.resource_manager.free_resources(request) + self.async_transfer_manager.mark_resources_freed(request) if self.gather_all_responses or self.dist.rank == 0: self.result_wait_queues.pop(request.py_request_id, None) diff --git a/tests/unittest/_torch/executor/test_async_transfer_manager.py b/tests/unittest/_torch/executor/test_async_transfer_manager.py index 1f2f9013d903..7fbda2f88e73 100644 --- a/tests/unittest/_torch/executor/test_async_transfer_manager.py +++ b/tests/unittest/_torch/executor/test_async_transfer_manager.py @@ -15,7 +15,7 @@ from unittest.mock import MagicMock -from tensorrt_llm._torch.pyexecutor.py_executor import AsyncTransferManager +from tensorrt_llm._torch.pyexecutor.py_executor import AsyncTransferManager, PyExecutor from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType from tensorrt_llm.bindings import LlmRequestState @@ -81,7 +81,9 @@ def test_start_transfer_single_request(): # Check seq slot manager was called to free resources seq_slot_manager.free_resources.assert_called_once_with(request) - manager.end_transfer(request) + transfer_result = manager.end_transfer(request) + assert transfer_result.completed + assert not transfer_result.needs_termination kv_cache_manager.unpin_blocks_by_id.assert_called_once() @@ -105,10 +107,13 @@ def test_start_transfer_multiple_transfers_same_request(): kv_cache_manager.store_blocks_for_reuse.assert_called_once() for _ in range(2): - manager.end_transfer(request) + transfer_result = manager.end_transfer(request) + assert not transfer_result.completed kv_cache_manager.unpin_blocks_by_id.assert_not_called() - manager.end_transfer(request) + transfer_result = manager.end_transfer(request) + assert transfer_result.completed + assert not transfer_result.needs_termination kv_cache_manager.unpin_blocks_by_id.assert_called_once() @@ -135,7 +140,9 @@ def test_transfer_without_storing_blocks(): kv_cache_manager.store_blocks_for_reuse.assert_not_called() spec_resource_manager.free_resources.assert_called_once_with(request) - assert manager.end_transfer(request) + transfer_result = manager.end_transfer(request) + assert transfer_result.completed + assert not transfer_result.needs_termination kv_cache_manager.unpin_blocks_by_id.assert_not_called() @@ -153,7 +160,9 @@ def test_end_transfer_preserves_error_state(): # Set error state before end_transfer request.state = LlmRequestState.DISAGG_TRANS_ERROR - manager.end_transfer(request) + transfer_result = manager.end_transfer(request) + assert transfer_result.completed + assert not transfer_result.needs_termination # Error state should be preserved assert request.state == LlmRequestState.DISAGG_TRANS_ERROR @@ -180,3 +189,91 @@ def test_requests_in_transfer(): assert in_transfer[1] is request1 assert in_transfer[2] is request2 assert in_transfer[3] is request3 + + +def test_end_transfer_reports_deferred_termination_needed(): + """End transfer reports when deferred termination must be retried.""" + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [100] + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager, should_store_blocks=True) + request = create_mock_request(42) + manager.start_transfer(request) + manager.mark_termination_requested(request) + + transfer_result = manager.end_transfer(request) + + assert transfer_result.completed + assert transfer_result.needs_termination + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with([100]) + + +def test_deferred_termination_survives_until_transfer_completion(): + """Termination requested during transfer is retried after transfer ends. + + This mock-driven test covers the cleanup contract directly; the + cancellation timing that triggered the rc13 block-reuse wedge is covered by + the disaggregated serving stress harness. + """ + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [100] + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager, should_store_blocks=True) + request = create_mock_request(42) + manager.start_transfer(request) + manager.mark_termination_requested(request) + + executor = PyExecutor.__new__(PyExecutor) + executor.kv_cache_transceiver = None + executor.active_requests = [] + executor.async_transfer_manager = manager + executor._terminate_request = MagicMock(return_value=True) + + executor._end_transfer_and_maybe_terminate(request) + + executor._terminate_request.assert_called_once_with(request) + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with([100]) + + +def test_no_deferred_termination_preserves_partial_reuse_skip(): + """Partial-reuse path still skips termination if none was requested.""" + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [100] + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager, should_store_blocks=True) + request = create_mock_request(42) + manager.start_transfer(request) + + executor = PyExecutor.__new__(PyExecutor) + executor.kv_cache_transceiver = None + executor.active_requests = [] + executor.async_transfer_manager = manager + executor._terminate_request = MagicMock(return_value=True) + + executor._end_transfer_and_maybe_terminate(request) + + executor._terminate_request.assert_not_called() + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with([100]) + + +def test_resources_freed_clears_deferred_termination(): + """Successful early resource free prevents duplicate termination.""" + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [100] + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager, should_store_blocks=True) + request = create_mock_request(42) + manager.start_transfer(request) + manager.mark_termination_requested(request) + manager.mark_resources_freed(request) + + executor = PyExecutor.__new__(PyExecutor) + executor.kv_cache_transceiver = None + executor.active_requests = [] + executor.async_transfer_manager = manager + executor._terminate_request = MagicMock(return_value=True) + + executor._end_transfer_and_maybe_terminate(request) + + executor._terminate_request.assert_not_called() + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with([100]) From 1db0b1b96836f0e9c8fe5d111e6230230cf4e1fb Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Fri, 8 May 2026 12:17:46 -0700 Subject: [PATCH 9/9] [https://nvbugs/6104831][feat] Expose CacheTransceiver pool-poisoned state to Python Adds public accessors `isRecvPoolPoisoned()` / `isSendPoolPoisoned()` on BaseTransBufferManager and BaseCacheTransceiver, surfacing the existing `ConcurrenceResource::mPoisoned` flag. Once any underlying pool's mPoisoned flag is set, BaseTransBufferManager:: assignBufferIndex throws unconditionally for the lifetime of the process (see baseTransBuffer.cpp:350 and the matching message: "The process must restart before these memory ranges can be safely reused"). Until now the only signal surfaced to callers was a per-request RequestError carrying the C++ exception text, which forces higher layers to string-match the message to distinguish a permanent worker-fatal state from genuinely transient per-request KV-transfer errors. This commit gives Python (the dynamo request handler / readiness probe) a structured, contract-stable way to ask "is this worker still able to serve disagg KV transfers, or must it be restarted?" without parsing error strings. Changes: - baseTransBuffer.h: two inline `noexcept` accessors that read the per-direction `mPoisoned` atomics (relaxed load mirrors the producer in poisonBufferIndex). - cacheTransceiver.h: virtual `isRecvPoolPoisoned()` / `isSendPoolPoisoned()` on BaseCacheTransceiver with a `return false` default so non-disagg subclasses (test mocks, generation-first Python transceiver) need not override. CacheTransceiver overrides walk the KV + RNN buffer-manager pointers and OR the per-manager state. - nanobind: bind the two methods on BaseCacheTransceiver. No trampoline size change required (defaults are concrete). - _torch/pyexecutor/kv_cache_transceiver.py: matching `is_recv_pool_poisoned` / `is_send_pool_poisoned` on the Python ABC (default False) and delegating implementations on `BindKvCacheTransceiver`. - tests/unittest/others/test_kv_cache_transceiver.py: regression test asserting both accessors return False on a fresh transceiver pair (NIXL + UCX backends). Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 23 ++++++++++++++ .../batch_manager/baseTransBuffer.h | 17 ++++++++++ .../batch_manager/cacheTransceiver.cpp | 31 +++++++++++++++++++ .../batch_manager/cacheTransceiver.cpp | 4 ++- .../_torch/pyexecutor/kv_cache_transceiver.py | 28 +++++++++++++++++ .../others/test_kv_cache_transceiver.py | 22 +++++++++++++ 6 files changed, 124 insertions(+), 1 deletion(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 091124595c54..e5c4a152b142 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -227,6 +227,26 @@ class BaseCacheTransceiver [[nodiscard]] virtual bool checkGenTransferComplete() const = 0; virtual bool cancelRequest(std::shared_ptr llmRequest) = 0; + + /// @brief Returns true if any underlying receive buffer pool has been + /// poisoned and can no longer serve allocations until process restart. + /// Higher layers (e.g. dynamo's request handler / readiness probe) use + /// this to escalate the worker to permanently-unhealthy without parsing + /// exception text from the request-error path. + /// Default returns false so non-disagg subclasses do not need to + /// override. + [[nodiscard]] virtual bool isRecvPoolPoisoned() const + { + return false; + } + + /// @brief Symmetric counterpart of isRecvPoolPoisoned for the send-side + /// pool. Provided for completeness; the historically-observed wedge is + /// receive-side only. + [[nodiscard]] virtual bool isSendPoolPoisoned() const + { + return false; + } }; class CacheTransceiver : public BaseCacheTransceiver @@ -274,6 +294,9 @@ class CacheTransceiver : public BaseCacheTransceiver virtual bool cancelRequest(std::shared_ptr llmRequest) override; + [[nodiscard]] bool isRecvPoolPoisoned() const override; + [[nodiscard]] bool isSendPoolPoisoned() const override; + private: void initializeCommState(); diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h index d74ca7b52535..e98daeffcb34 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h @@ -268,6 +268,23 @@ class BaseTransBufferManager return mMaxNumTokens; } + /// @brief Returns true if the receive buffer pool has been poisoned. + /// A poisoned pool refuses every subsequent allocation in + /// assignBufferIndex (see baseTransBuffer.cpp:350) and the only + /// recovery is process restart. This accessor lets higher layers + /// query the state without parsing exception text. + [[nodiscard]] bool isRecvPoolPoisoned() const noexcept + { + return mConcurrenceRecvResource.mPoisoned.load(std::memory_order_relaxed); + } + + /// @brief Returns true if the send buffer pool has been poisoned. + /// Same semantics as isRecvPoolPoisoned but for the send-side pool. + [[nodiscard]] bool isSendPoolPoisoned() const noexcept + { + return mConcurrenceSendResource.mPoisoned.load(std::memory_order_relaxed); + } + protected: /// @brief Constructor - derived classes call this after computing buffer sizes. /// @param transferBufferSize Size of each transfer buffer in bytes. diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index f086a7acf9d8..5fbf873677db 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -1000,4 +1000,35 @@ bool CacheTransceiver::cancelRequest(std::shared_ptr llmRequest) return false; } +bool CacheTransceiver::isRecvPoolPoisoned() const +{ + // Walk every BaseTransBufferManager (KV pools + optional RNN pool). Once + // any underlying pool's mPoisoned flag is set, BaseTransBufferManager:: + // assignBufferIndex unconditionally throws on subsequent allocations and + // there is no recovery path short of restarting the process. Aggregating + // per-manager state into a single boolean lets callers (e.g. the Python + // request handler / readiness probe) escalate the worker to permanently + // unhealthy without depending on parsing exception text. + for (auto const* mgr : mCacheTransBufferManagerPtrs) + { + if (mgr != nullptr && mgr->isRecvPoolPoisoned()) + { + return true; + } + } + return false; +} + +bool CacheTransceiver::isSendPoolPoisoned() const +{ + for (auto const* mgr : mCacheTransBufferManagerPtrs) + { + if (mgr != nullptr && mgr->isSendPoolPoisoned()) + { + return true; + } + } + return false; +} + } // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index b0df67ba965a..13ec4dc17f7f 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -110,7 +110,9 @@ void tb::CacheTransceiverBindings::initBindings(nb::module_& m) .def("check_gen_transfer_status", &BaseCacheTransceiver::checkGenTransferStatus, nb::call_guard()) .def("check_gen_transfer_complete", &BaseCacheTransceiver::checkGenTransferComplete) - .def("cancel_request", &BaseCacheTransceiver::cancelRequest); + .def("cancel_request", &BaseCacheTransceiver::cancelRequest) + .def("is_recv_pool_poisoned", &BaseCacheTransceiver::isRecvPoolPoisoned) + .def("is_send_pool_poisoned", &BaseCacheTransceiver::isSendPoolPoisoned); nb::enum_(m, "AttentionType") .value("DEFAULT", executor::kv_cache::CacheState::AttentionType::kDEFAULT) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index e27290b472fa..fd39729a69e1 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -140,6 +140,28 @@ def get_disaggregated_params(self) -> Dict[str, Any]: """ ... + def is_recv_pool_poisoned(self) -> bool: + """Returns True if any underlying receive buffer pool has been poisoned. + + A poisoned pool refuses every subsequent allocation in the C++ + BaseTransBufferManager::assignBufferIndex guard, and the only recovery + is process restart. Higher layers (e.g. the dynamo request handler / + readiness probe) check this to escalate the worker to permanently + unhealthy without parsing exception text from the request-error path. + Default returns False so subclasses that don't have an underlying + C++ buffer pool (test mocks, generation-first Python transceiver) are + unaffected. + """ + return False + + def is_send_pool_poisoned(self) -> bool: + """Returns True if any underlying send buffer pool has been poisoned. + + Symmetric counterpart of is_recv_pool_poisoned; the historically + observed wedge is receive-side only. + """ + return False + def shutdown(self): """Shut down the transceiver and release registered resources.""" @@ -214,6 +236,12 @@ def get_disaggregated_params(self): # Only new py cache transceiver will support gen-first disagg return {} + def is_recv_pool_poisoned(self) -> bool: + return self.impl.is_recv_pool_poisoned() + + def is_send_pool_poisoned(self) -> bool: + return self.impl.is_send_pool_poisoned() + class CacheTransBufferManager: diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index 243f345cfabd..d8c26b22a284 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -1033,3 +1033,25 @@ def test_hybrid_cache_transceiver_cancel_request(backend, monkeypatch): # Block the main thread due to the async operation time.sleep(2) assert gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR + + +@pytest.mark.timeout(60) +@pytest.mark.parametrize("disagg_transceiver_pair", + [(AttentionTypeCpp.DEFAULT, "NIXL"), + (AttentionTypeCpp.DEFAULT, "UCX")], + ids=["NIXL", "UCX"], + indirect=True) +def test_pool_poisoned_default_false(disagg_transceiver_pair): + """A fresh transceiver pair reports both pools as not poisoned. + + Regression test for the is_recv_pool_poisoned / is_send_pool_poisoned + accessors that surface BaseTransBufferManager::mPoisoned to higher + layers (e.g. dynamo's request handler) so they can escalate the worker + to permanently unhealthy without parsing exception text from the + request-error path. + """ + _, _, ctx_transceiver, gen_transceiver = disagg_transceiver_pair + assert ctx_transceiver.is_recv_pool_poisoned() is False + assert ctx_transceiver.is_send_pool_poisoned() is False + assert gen_transceiver.is_recv_pool_poisoned() is False + assert gen_transceiver.is_send_pool_poisoned() is False