From 630fa3b4f220f250e9bd3482ddc09f8e6680a5bc 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 01/24] [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 100e96f03119..05ed827a9511 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -923,7 +923,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); } @@ -1604,11 +1604,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 fbadd440d611..5e33fd504299 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -368,6 +368,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) @@ -3306,9 +3313,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, @@ -3316,9 +3320,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) @@ -3402,15 +3419,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() @@ -3510,6 +3548,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 @@ -3847,6 +3894,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 @@ -3989,7 +4043,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: @@ -4064,6 +4128,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 0ad4d0adb5954311dde14cb0fa449c41cd5d2e2e 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 02/24] [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) Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../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 cf9d544a9ce0..68d7abff50a0 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -229,6 +229,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 e21e846105237241175a01672635fae9612c1283 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 03/24] [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 68d7abff50a0..c224056e83c4 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": @@ -230,95 +352,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) + disagg_transceiver_pair, capfd): + """Reproduce the sender-side broken-promise on cancel-after-ready. - 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) + 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. @@ -386,8 +809,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 1c321cc0b0cbe8f121a79eecd37d81c64a2fe59c Mon Sep 17 00:00:00 2001 From: Yifan Jiang Date: Mon, 4 May 2026 11:24:21 -0700 Subject: [PATCH 04/24] Fail closed on unquiesced disagg KV transfer Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> 1777930306 -0700 Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../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 | 103 +++++++++++--- 10 files changed, 517 insertions(+), 88 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 5e33fd504299..3113398fa0e5 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -577,7 +577,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 @@ -3561,10 +3561,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") @@ -3855,10 +3855,15 @@ def _handle_errors(self, "on fatal error") failed_requests = (list(self.active_requests) - if requests is None else 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, @@ -3877,7 +3882,70 @@ def _handle_errors(self, if self._fatal_error is not None: self.executor_request_queue.enqueue_shutdown_request() - def _terminate_request(self, request: LlmRequest): + 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, @@ -3887,6 +3955,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) @@ -4044,23 +4113,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], - charge_budget=False) + 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 9a1af0bfd4f9fefe71b9dd068a2991d0c2832bfa 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 05/24] [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 3113398fa0e5..df484638c7b0 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3882,9 +3882,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) @@ -3920,6 +3919,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 78fa9d712a308d758c7ccd5e046bdd127ea1fd4c 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 06/24] 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 df484638c7b0..c1023776ea71 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -375,6 +375,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) @@ -3561,10 +3562,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") @@ -3856,10 +3857,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: @@ -3869,11 +3886,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: @@ -3882,47 +3899,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: @@ -3931,9 +3910,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 @@ -3969,6 +3948,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.""" @@ -4114,12 +4094,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 89001e30627891626808d2755d53f876b1242fdc Mon Sep 17 00:00:00 2001 From: Yifan Jiang Date: Tue, 5 May 2026 11:13:41 -0700 Subject: [PATCH 07/24] 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 c1023776ea71..66517689d8b1 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -375,7 +375,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) @@ -3549,23 +3550,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") @@ -3948,7 +3944,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.""" @@ -4095,11 +4092,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 " @@ -4172,15 +4169,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 62102e831704741c5c49bc7934314e4245685760 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 08/24] [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 | 78 ++++++++++--- .../executor/test_async_transfer_manager.py | 109 +++++++++++++++++- 3 files changed, 167 insertions(+), 44 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 66517689d8b1..529c2242e66d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -164,10 +164,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 @@ -180,6 +189,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): @@ -199,6 +220,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. @@ -232,14 +267,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[ @@ -248,9 +284,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) @@ -262,9 +300,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 @@ -668,16 +707,19 @@ def _end_transfer_and_maybe_terminate(self, request: LlmRequest): # participate. self._pending_transfer_responses.append( (request.py_request_id, response)) - if self.async_transfer_manager.end_transfer(request): + 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) def _flush_pending_transfer_responses(self): @@ -3901,12 +3943,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 @@ -3934,6 +3983,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 a201be28a02455dbf2254c0487aaf2b324827317 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 8 May 2026 12:42:25 -0700 Subject: [PATCH 09/24] [https://nvbugs/6104831][fix]Cleanup PR for review Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.cpp | 37 +++++- .../batch_manager/dataTransceiver.cpp | 10 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 120 ++++++++++++++++-- .../executor/test_async_transfer_manager.py | 24 ++++ .../others/test_kv_cache_transceiver.py | 105 ++++++++++++++- 5 files changed, 266 insertions(+), 30 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index f086a7acf9d8..4a3aa42d42db 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -66,6 +66,16 @@ namespace /// 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. +/// +/// @note Logging policy for cacheTransceiver / dataTransceiver: +/// new logs are added only at entry/exit boundaries +/// (checkGenTransferStatus done / context-completion summaries) +/// or for cancel-flow transitions ("Flipped ... cancel flag"). +/// Per-iteration polling-loop logs are deliberately omitted to +/// keep TLLM_LOG_DEBUG output proportional to actual +/// state-changing events. See PR #13713 review for the +/// rationale; if you add a new TLLM_LOG_DEBUG here, audit +/// whether it fires per iteration vs per state transition. inline int currentRankForLog() { return useMPI() ? mpi::MpiComm::world().getRank() : tensorrt_llm::pg_utils::get_world_pg()->getRank(); @@ -518,8 +528,19 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( 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. + // Behaviour change vs upstream/main: the overall transfer deadline + // ``kvTransferTimeoutMs`` now applies in BOTH ``blockAll`` and polling + // modes, where it previously only ran in polling mode. A stuck sender + // whose future never becomes ready was otherwise able to pin KV + // blocks forever when callers invoked the function in ``blockAll`` + // mode. Today the only ``blockAll`` callers in the codebase are + // shutdown / drain paths that already complete quickly under + // healthy conditions, so this widening is purely defensive -- + // healthy transfers see no behaviour change because they finish + // long before ``kvTransferTimeoutMs`` (default 60s on the Python + // side; configurable). If a caller relied on unbounded blocking + // here, set ``kv_transfer_timeout_ms`` to a value comfortably + // larger than the expected worst-case transfer. kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); // The per-iteration poll timeout is only relevant when the caller wants to // return after checking at least `atLeastRequestNum` entries. @@ -705,10 +726,14 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR 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. + // Behaviour change vs upstream/main, mirroring the context side + // (see checkContextTransferStatus): ``kvTransferTimeoutMs`` now + // applies in both ``blockAll`` and polling modes. Without this, + // mRequesterFutures would accumulate stuck entries, pinning + // generation-side KV blocks and eventually exhausting the pool. + // Healthy transfers complete well under the configured timeout, + // so this widening is defensive; callers that need more headroom + // can raise ``kv_transfer_timeout_ms``. kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); // The per-iteration poll timeout is only relevant when the caller wants to // return after checking at least `atLeastRequestNum` entries. diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 9de26a87b152..6e40aec66b0f 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -1475,10 +1475,12 @@ class CacheReceiver::Impl } catch (...) { - if (agentConnectionManagerForAcq) - { - poisonRecvHolders(); - } + // recvHolders is only populated on the AgentConnectionManager + // path (see the if-guarded block above), so on direct-UCX this + // is a no-op loop. Calling it unconditionally keeps the catch + // symmetric with the happy-path detach loop and removes a + // redundant branch on the failure path. + poisonRecvHolders(); llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); throw; diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 529c2242e66d..8bebbf4d8ee0 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -86,6 +86,18 @@ # Default: "0" (only rank 0 prints, matching existing behavior). PROFILE_LOG_RANKS_ENV_VAR_NAME = "TLLM_PROFILE_LOG_RANKS" +# Defense-in-depth Python-side fallback deadline applied to every disagg +# KV transfer when ``CacheTransceiverConfig.kv_transfer_timeout_ms`` is +# unset. The C++ deadline check is gated on ``kvTransferTimeoutMs.has_value()`` +# (cacheTransceiver.cpp), so a deployment that explicitly passes ``None`` +# (or constructs the C++ config bypassing the Python defaults) would +# otherwise leave ``_handle_errors``-deferred requests waiting forever +# for a ``kDISAGG_TRANS_ERROR`` that never arrives -- the original +# wedge in disguise. Ten minutes is two orders of magnitude beyond +# any healthy disagg KV transfer; we surface the error rather than +# wait silently. +_DISAGG_KV_TRANSFER_FALLBACK_DEADLINE_MS = 10 * 60 * 1000 + class PPCommTag(IntEnum): """ @@ -221,13 +233,36 @@ def requests_in_transfer(self) -> Dict[int, LlmRequest]: return self._requests_in_transfer def mark_termination_requested(self, request: LlmRequest): + """Record that termination was attempted while transfer was in flight. + + The expected fast path is: ``_terminate_request`` -> + ``_can_terminate_request_now`` finds the request still in + ``is_disagg_context_transmission_state`` -> calls this -> returns False + -> ``_end_transfer_and_maybe_terminate`` later observes + ``needs_termination`` and retries termination once C++ reports the + transfer terminal. + + The "no metadata" branch only fires when ``end_transfer`` has already + run and popped the entry; that path also transitions state to + ``DISAGG_CONTEXT_COMPLETE`` (or ``DISAGG_TRANS_ERROR``), so by the + time ``_can_terminate_request_now`` is re-entered for the same + request its first branch is False and we never reach this method. + Any reachable miss means the state/metadata invariant has drifted + (e.g. a future refactor of ``end_transfer`` ordering) -- a debug log + is enough today, but anyone seeing this in field logs should treat + it as a regression of the order contract above and fix the caller, + not silence the message. + """ 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") + "but it is not tracked by AsyncTransferManager (likely a " + "harmless ordering window where end_transfer ran first; if " + "this fires repeatedly the state-vs-metadata invariant has " + "drifted)") def mark_resources_freed(self, request: LlmRequest): metadata = self._request_transfer_metadata.get(request.py_request_id) @@ -3227,9 +3262,20 @@ def _check_disagg_gen_transfer_status(self): def _check_kv_transfer_timeout(self): if not self.kv_cache_transceiver: return - timeout_ms = self.kv_cache_transceiver.kv_transfer_timeout_ms - if timeout_ms is None: - return + configured_timeout_ms = ( + self.kv_cache_transceiver.kv_transfer_timeout_ms) + # Always apply *some* deadline. When the user/admin disabled the + # configured timeout we still need a Python-side fallback because + # the deferral path in _handle_errors relies on a future + # DISAGG_TRANS_ERROR transition to surface the failure: a request + # that never times out on the C++ side never moves to error and + # would wedge here forever. See _DISAGG_KV_TRANSFER_FALLBACK_DEADLINE_MS. + if configured_timeout_ms is None: + timeout_ms = _DISAGG_KV_TRANSFER_FALLBACK_DEADLINE_MS + using_fallback = True + else: + timeout_ms = configured_timeout_ms + using_fallback = False def flag_if_kv_transfer_timed_out(req: LlmRequest, type: str) -> None: current_time = time.time() @@ -3237,9 +3283,17 @@ def flag_if_kv_transfer_timed_out(req: LlmRequest, type: str) -> None: return elapsed_time = (current_time - req.py_kv_transfer_start_time) * 1000 if elapsed_time > timeout_ms and not req.py_kv_transfer_timed_out: - logger.warning( - f"Terminating {type} request {req.py_request_id} due to KV cache transfer timeout" - ) + if using_fallback: + logger.warning( + f"Terminating {type} request {req.py_request_id} " + f"after Python fallback deadline ({timeout_ms} ms); " + "kv_transfer_timeout_ms is unset, so this is the " + "only escalation path. Configure " + "kv_transfer_timeout_ms to surface failures sooner.") + else: + logger.warning( + f"Terminating {type} request {req.py_request_id} due to KV cache transfer timeout" + ) req.py_kv_transfer_timed_out = True for req in self.async_transfer_manager.requests_in_transfer().values(): @@ -3491,10 +3545,13 @@ def _recv_disagg_gen_cache(self, new_gen_reqs): req.py_request_id) raise - if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None: - for req in recv_reqs: - if req.state == LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS: - req.py_kv_transfer_start_time = time.time() + # Record the transfer start time unconditionally so the Python + # fallback deadline in _check_kv_transfer_timeout can fire even + # when kv_transfer_timeout_ms is not configured. The check there + # picks the configured timeout or the fallback floor. + for req in recv_reqs: + if req.state == LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS: + req.py_kv_transfer_start_time = time.time() non_gen_first_active = [ req for req in self.active_requests @@ -3533,8 +3590,12 @@ def kv_connector_request_finished(req: LlmRequest): self.async_transfer_manager.start_transfer(req) self.kv_cache_transceiver.respond_and_send_async(req) - if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None: - req.py_kv_transfer_start_time = time.time() + # Always record the transfer start time. The Python + # fallback deadline in _check_kv_transfer_timeout + # falls back to a hard floor when + # kv_transfer_timeout_ms is unset, and that path + # needs a baseline timestamp on every request. + req.py_kv_transfer_start_time = time.time() if self.kv_connector_manager: if not self.disable_overlap_scheduler: @@ -3900,6 +3961,12 @@ def _handle_errors(self, if self._is_unquiesced_disagg_transfer(request) ] if deferred_requests: + # Deferred requests rely on a follow-up DISAGG_TRANS_ERROR + # transition to surface the failure. The transition is + # driven by the C++ deadline (kvTransferTimeoutMs) when + # configured and by _check_kv_transfer_timeout's Python + # fallback floor (_DISAGG_KV_TRANSFER_FALLBACK_DEADLINE_MS) + # otherwise; either way, this path cannot wedge. logger.warning( "Deferring error cleanup for disaggregated KV transfers still " "in flight. C++ transfer status remains the source of truth " @@ -3932,6 +3999,11 @@ def _handle_errors(self, ] self._enqueue_responses(list(error_responses.items())) for request in failed_requests: + # _terminate_request may return False for disagg requests that + # are still in an unquiesced KV-transfer state (see contract + # docstring); in that case the deferred retry runs from + # _end_transfer_and_maybe_terminate via the mark_termination_requested + # handoff installed inside _can_terminate_request_now. self._terminate_request(request) if self._fatal_error is not None: @@ -3968,6 +4040,23 @@ def _do_terminate_request_if_safe(self, request: LlmRequest) -> bool: return True def _terminate_request(self, request: LlmRequest) -> bool: + """Terminate ``request`` immediately or defer until KV transfer drains. + + Returns ``True`` when termination ran (either inline or scheduled + via the PP termination handler). Returns ``False`` when the + request is still in an unquiesced disagg KV-transfer state; in + that case ``_can_terminate_request_now`` records the deferred + intent on ``async_transfer_manager`` via + ``mark_termination_requested`` and ``_end_transfer_and_maybe_terminate`` + retries exactly once after C++ surfaces transfer completion. + + Callers that pass non-disagg requests (e.g. paused requests, ADP + dummy requests) can safely ignore the return value -- those + requests never enter the unquiesced state and termination always + runs synchronously. Callers that may pass disagg requests + (``_handle_errors``, ``_handle_responses``) must rely on the + deferred-retry handoff above. + """ if not self._can_terminate_request_now(request): return False # Dummy requests don't participate in disagg KV cache transfers, @@ -4239,6 +4328,11 @@ def _handle_responses(self): # Request should be terminated after enqueueing response to ensure we can enqueue response successfully. self._enqueue_responses(new_responses) for request in requests_to_terminate: + # The force_terminate_for_partial_reuse branch above can route + # disagg context requests still in DISAGG_CONTEXT_TRANS_IN_PROGRESS + # here; _terminate_request will return False for those and the + # retry runs from _end_transfer_and_maybe_terminate (see contract + # docstring on _terminate_request). self._terminate_request(request) return requests_to_terminate + requests_finished_by_transfer diff --git a/tests/unittest/_torch/executor/test_async_transfer_manager.py b/tests/unittest/_torch/executor/test_async_transfer_manager.py index 7fbda2f88e73..8fb6891bb5af 100644 --- a/tests/unittest/_torch/executor/test_async_transfer_manager.py +++ b/tests/unittest/_torch/executor/test_async_transfer_manager.py @@ -12,6 +12,30 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +"""Unit tests for ``AsyncTransferManager`` and the rc13 block-reuse stop-gap. + +Validation strategy for the stop-gap (NVBug 6104831 sig #8): + +1. **Unit tests (this file).** Cover the deferred-termination *cleanup + contract* by driving the manager API directly: + ``mark_termination_requested`` -> ``end_transfer`` -> retry termination + via ``EndTransferResult.needs_termination``. These tests run in + milliseconds and pin the contract against silent regressions. + +2. **Stress harness (out-of-tree).** Reproduces the *cancellation + timing* that originally manifested at high concurrency on rc13 with + ``should_store_blocks=True`` and an in-flight cancel landing on the + partial-reuse path. The harness lives with the disaggregated serving + reproducers and is run pre-merge for every PR touching the disagg + cleanup path; see PR #13713 description for run instructions and + recorded baselines. + + We deliberately do not land a unit test for the timing window because + constructing it without an actual disagg pipeline would either (a) + re-mock the same code path covered above, adding no signal, or (b) + spin up a real two-process disagg session, which belongs in the + stress harness rather than ``unittest``. +""" from unittest.mock import MagicMock diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index c224056e83c4..184d7e5862cc 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -85,10 +85,17 @@ def _make_request(request_id, def _add_sequence(kv_cache_manager, request): """Bind ``request`` into ``kv_cache_manager``. - This uses the call signature every disagg transceiver test relies on. + The production prepare_resources path always pairs add_sequence_batch + with refresh_blocks; without the second call the KV-block offsets + landed on the request are stale and the C++ disagg transceiver path + silently transfers into the wrong physical slots, leaving the + receiver-side pool buffers all-zero. Mirror that pairing here so + every test that uses this helper observes the same state the runtime + does. """ - kv_cache_manager.impl.add_sequence(request.py_request_id, - request.prompt_len, 1, request) + kv_cache_manager.impl.add_sequence_batch( + [(request.py_request_id, request.prompt_len, 1)], [request]) + kv_cache_manager.impl.refresh_blocks() def _drive_template_handshake_to_completion(ctx_xcvr, @@ -241,6 +248,11 @@ def test_kv_cache_transceiver_single_process(ctx_gen_kv_cache_dtype, kv_cache_manager_ctx.impl.add_sequence_batch( [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], [ctx_request]) + # add_sequence_batch must be paired with refresh_blocks so the C++ + # transceiver path observes the freshly-allocated KV-block offsets; + # without this the receiver-side pool buffers stay zero. See + # _add_sequence helper for the same pairing in cancellation tests. + kv_cache_manager_ctx.impl.refresh_blocks() # send ctx request kv_cache_transceiver_ctx.respond_and_send_async(ctx_request) @@ -271,6 +283,7 @@ def test_kv_cache_transceiver_single_process(ctx_gen_kv_cache_dtype, kv_cache_manager_gen.impl.add_sequence_batch( [(gen_request.py_request_id, gen_request.prompt_len, 1)], [gen_request]) + kv_cache_manager_gen.impl.refresh_blocks() # send gen request kv_cache_transceiver_gen.request_and_receive_async(gen_request) @@ -320,6 +333,9 @@ def test_cancel_request_in_transmission(attention_type): kv_cache_manager_ctx.impl.add_sequence_batch( [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], [ctx_request]) + # Pair with refresh_blocks; see _add_sequence helper docstring for + # why both calls are required for the C++ transceiver path. + kv_cache_manager_ctx.impl.refresh_blocks() # send ctx request kv_cache_transceiver_ctx.respond_and_send_async(ctx_request) @@ -343,6 +359,7 @@ def test_cancel_request_in_transmission(attention_type): kv_cache_manager_gen.impl.add_sequence_batch( [(gen_request.py_request_id, gen_request.prompt_len, 1)], [gen_request]) + kv_cache_manager_gen.impl.refresh_blocks() # send gen request kv_cache_transceiver_gen.request_and_receive_async(gen_request) @@ -351,6 +368,79 @@ def test_cancel_request_in_transmission(attention_type): assert gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR +@pytest.mark.timeout(120) +@pytest.mark.parametrize("disagg_transceiver_pair", + [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA], + ids=["mha", "mla"], + indirect=True) +def test_request_and_receive_async_state_ordering(disagg_transceiver_pair): + """Lock down the setState-after-receiveAsync ordering contract. + + ``cacheTransceiver.cpp::requestAndReceiveAsync`` was reordered so + ``llmRequest->setState(DISAGG_GENERATION_TRANS_IN_PROGRESS)`` runs + *after* ``mCacheReceiver->receiveAsync(llmRequest)`` returns (the + original code set the state first). The reorder is load-bearing: + if ``receiveAsync`` throws, the request is left out of both + ``mRequesterFutures`` and the IN_PROGRESS state set; if the spawned + worker ever read state at entry, the new ordering would race. + Today the worker does not read state at entry, but the contract is + fragile and easy to regress in a refactor. + + This test exercises a happy-path receive and asserts the observable + transition is ``DISAGG_GENERATION_INIT`` (pre-call) -> ``IN_PROGRESS`` + or already-``COMPLETE`` for a fast-completing transfer (post-call), + with the receiver-side bookkeeping reflecting the in-flight future. + """ + (kv_cache_manager_ctx, kv_cache_manager_gen, kv_cache_transceiver_ctx, + kv_cache_transceiver_gen) = disagg_transceiver_pair + + # Drive a real ctx handshake so the gen request has a live peer to + # receive from -- without one, request_and_receive_async would just + # park the future on an unresolvable peer. + ctx_request = _make_request(0, LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) + _add_sequence(kv_cache_manager_ctx, ctx_request) + kv_cache_transceiver_ctx.respond_and_send_async(ctx_request) + + gen_request = _make_request(0, + LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + ctx_request.context_phase_params) + _add_sequence(kv_cache_manager_gen, gen_request) + + # Pre-condition: gen-only requests built via _make_request land in + # DISAGG_GENERATION_INIT (set by the LlmRequest ctor for + # generation-only types). Locking this in protects the test from + # silent state-default drift. + assert gen_request.state == LlmRequestState.DISAGG_GENERATION_INIT, ( + f"gen_request.state should be DISAGG_GENERATION_INIT before " + f"request_and_receive_async, got {gen_request.state}") + assert kv_cache_transceiver_gen.check_gen_transfer_complete(), ( + "mRequesterFutures must be empty before request_and_receive_async") + + kv_cache_transceiver_gen.request_and_receive_async(gen_request) + + # Post-condition: setState and emplace happen together. The state + # must be IN_PROGRESS (or already TRANS_COMPLETE for a transfer that + # finished synchronously inside the call). It must NOT still be + # INIT -- if it is, the reordered setState was bypassed or didn't + # run. + assert gen_request.state in ( + LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS, + LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE, + ), (f"gen_request.state must be IN_PROGRESS or TRANS_COMPLETE after " + f"request_and_receive_async, got {gen_request.state}") + + # Drain the in-flight transfer so teardown is clean. + 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) + kv_cache_transceiver_ctx.check_context_transfer_status(1) + time.sleep(0.05) + + kv_cache_manager_ctx.free_resources(ctx_request) + kv_cache_manager_gen.free_resources(gen_request) + + @pytest.mark.timeout(120) @pytest.mark.parametrize("disagg_transceiver_pair", [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA], @@ -591,10 +681,6 @@ def call_blocking_check(): "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) @@ -767,6 +853,11 @@ def create_hybrid_cache_manager(mapping, mamba_layer_mask=mamba_layer_mask, mamba_cache_dtype=mamba_conv_dtype, mamba_ssm_cache_dtype=mamba_ssm_dtype, + # Disagg KV transfer is the whole point of this test file, and the + # MixedMambaHybridCacheManager (selected when is_disagg=True) is the + # only variant whose mamba state is exposed via the disagg + # transceivers exercised below. + is_disagg=True, kv_cache_config=KvCacheConfig( max_tokens=256, enable_block_reuse=False, From 36aef0405482b9bc3a6121f05b51a6e72175606c Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 11 May 2026 17:05:30 -0700 Subject: [PATCH 10/24] [https://nvbugs/6104831][fix] Address review feedback Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 29 +-- .../batch_manager/baseTransBuffer.cpp | 17 +- .../batch_manager/baseTransBuffer.h | 6 + .../batch_manager/cacheFormatter.cpp | 8 +- .../batch_manager/cacheFormatter.h | 9 + .../batch_manager/cacheTransceiver.cpp | 189 ++++++++++-------- .../batch_manager/mlaCacheFormatter.cpp | 4 - .../batch_manager/mlaCacheFormatter.h | 1 + .../batch_manager/rnnCacheFormatter.cpp | 1 + .../mooncake_utils/transferAgent.cpp | 11 +- .../mooncake_utils/transferAgent.h | 4 +- .../nixl_utils/transferAgent.cpp | 18 +- .../nixl_utils/transferAgent.h | 6 + .../batch_manager/cacheTransceiver.cpp | 3 +- .../_torch/pyexecutor/kv_cache_transceiver.py | 6 + tensorrt_llm/_torch/pyexecutor/py_executor.py | 15 +- 16 files changed, 195 insertions(+), 132 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 091124595c54..af2cb9b81395 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -205,10 +205,6 @@ class BaseCacheTransceiver { public: virtual ~BaseCacheTransceiver() = default; - // 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) @@ -227,6 +223,11 @@ class BaseCacheTransceiver [[nodiscard]] virtual bool checkGenTransferComplete() const = 0; virtual bool cancelRequest(std::shared_ptr llmRequest) = 0; + + [[nodiscard]] virtual bool hasPoisonedTransferBuffer() const + { + return false; + } }; class CacheTransceiver : public BaseCacheTransceiver @@ -274,6 +275,8 @@ class CacheTransceiver : public BaseCacheTransceiver virtual bool cancelRequest(std::shared_ptr llmRequest) override; + [[nodiscard]] bool hasPoisonedTransferBuffer() const override; + private: void initializeCommState(); @@ -281,24 +284,6 @@ class CacheTransceiver : public BaseCacheTransceiver std::unique_ptr mCacheSender; std::unique_ptr mCacheReceiver; - // 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). - // - // 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.cpp b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp index 9b83205451a7..f77d18f12ae8 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp @@ -333,18 +333,15 @@ std::optional BaseTransBufferManager::assignBufferIndex(ConcurrenceResource return resource.mPoisoned.load(std::memory_order_relaxed) || static_cast(resource.mConcurrence) < bufferCount; }; - if (!predicate()) + auto const slice = std::chrono::milliseconds{waitSliceMs}; + while (!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)) { - 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()); - } + auto const reqIdStr + = requestIdForLog.has_value() ? std::to_string(requestIdForLog.value()) : std::string{"?"}; + TLLM_THROW("assignBufferIndex cancelled via perRequestCancel (reqId=%s)", reqIdStr.c_str()); } } TLLM_CHECK_WITH_INFO(!resource.mPoisoned.load(std::memory_order_relaxed), diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h index d74ca7b52535..cf29efa8fc84 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h @@ -268,6 +268,12 @@ class BaseTransBufferManager return mMaxNumTokens; } + [[nodiscard]] bool hasPoisonedBuffer() const noexcept + { + return mConcurrenceSendResource.mPoisoned.load(std::memory_order_relaxed) + || mConcurrenceRecvResource.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/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index f6b47440c50c..2a1a443634de 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -464,10 +464,6 @@ 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++) @@ -595,6 +591,10 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio { if (agentConnection != nullptr) { + // AgentConnection::send can throw on cancel/error after the + // backend transfer has seen this send buffer. release() is + // not a quiescence proof, so do not return the slot to the + // pool; poison it and let Python restart the process. sendHolder.poison(); } throw; diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.h b/cpp/tensorrt_llm/batch_manager/cacheFormatter.h index 458cac8d4382..0f1f70337b9c 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.h +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.h @@ -232,6 +232,14 @@ class BaseCacheFormatter virtual ~BaseCacheFormatter() = default; }; +inline void checkZeroCopyDisabledForCancellableDisaggTransfers() +{ + TLLM_CHECK_WITH_INFO(!common::getEnvTryZCopyForKVCacheTransfer(), + "Zero-copy cache transfer sends/receives directly from request-owned blocks. It is disabled for cancellable " + "disaggregated transfers until KV-block leases can prove the peer has stopped accessing those blocks before " + "the request is freed."); +} + // Simple cache block copy. Because it does not involve data splitting or merging, it performs best when the // parallel topology is completely identical, making it the preferred method. class CacheFormatter final : public BaseCacheFormatter @@ -243,6 +251,7 @@ class CacheFormatter final : public BaseCacheFormatter { TLLM_CHECK(mCacheManager); TLLM_CHECK(mCacheTransBufferManager); + checkZeroCopyDisabledForCancellableDisaggTransfers(); } void format(tensorrt_llm::batch_manager::TransferSession& session) override; diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 4a3aa42d42db..93a4b398e8d0 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -62,20 +62,7 @@ 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. -/// -/// @note Logging policy for cacheTransceiver / dataTransceiver: -/// new logs are added only at entry/exit boundaries -/// (checkGenTransferStatus done / context-completion summaries) -/// or for cancel-flow transitions ("Flipped ... cancel flag"). -/// Per-iteration polling-loop logs are deliberately omitted to -/// keep TLLM_LOG_DEBUG output proportional to actual -/// state-changing events. See PR #13713 review for the -/// rationale; if you add a new TLLM_LOG_DEBUG here, audit -/// whether it fires per iteration vs per state transition. +/// @brief Rank prefix for TLLM_LOG_DEBUG/WARNING; resolves MPI vs torch process group at log time. inline int currentRankForLog() { return useMPI() ? mpi::MpiComm::world().getRank() : tensorrt_llm::pg_utils::get_world_pg()->getRank(); @@ -528,19 +515,10 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( std::optional kvTransferTimeoutMs = std::nullopt; if (mCacheTransceiverConfig.has_value()) { - // Behaviour change vs upstream/main: the overall transfer deadline - // ``kvTransferTimeoutMs`` now applies in BOTH ``blockAll`` and polling - // modes, where it previously only ran in polling mode. A stuck sender - // whose future never becomes ready was otherwise able to pin KV - // blocks forever when callers invoked the function in ``blockAll`` - // mode. Today the only ``blockAll`` callers in the codebase are - // shutdown / drain paths that already complete quickly under - // healthy conditions, so this widening is purely defensive -- - // healthy transfers see no behaviour change because they finish - // long before ``kvTransferTimeoutMs`` (default 60s on the Python - // side; configurable). If a caller relied on unbounded blocking - // here, set ``kv_transfer_timeout_ms`` to a value comfortably - // larger than the expected worst-case transfer. + // Apply the overall transfer deadline in both blockAll and polling + // modes so a stuck sender cannot pin KV blocks indefinitely. + // Healthy transfers complete well before the deadline (default + // 60s on the Python side, configurable via kv_transfer_timeout_ms). kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); // The per-iteration poll timeout is only relevant when the caller wants to // return after checking at least `atLeastRequestNum` entries. @@ -642,8 +620,26 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( { try { - // 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))); + // Bound the wait so a stuck sender cannot wedge the call: + // per-iteration timeout if set, else remaining overall + // deadline, else block (caller opted in). + int64_t effectiveSliceMs = 0; + if (senderFutureTimeoutMs.has_value()) + { + effectiveSliceMs = senderFutureTimeoutMs.value(); + } + else if (kvTransferTimeoutMs.has_value()) + { + auto const elapsedNow = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()) + .count(); + effectiveSliceMs = std::max(0, kvTransferTimeoutMs.value() - elapsedNow); + } + else + { + effectiveSliceMs = std::numeric_limits::max(); + } + auto status = future.wait_for(std::chrono::milliseconds(effectiveSliceMs)); if (status == std::future_status::ready) { future.get(); @@ -656,27 +652,38 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } 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()) + // Re-check the overall deadline inline; if reached, + // mark error here instead of blocking on a follow-up + // future.get(). + bool deadlineReached = false; + if (kvTransferTimeoutMs.has_value()) { - TLLM_LOG_WARNING("Timed out waiting for context KV cache transfer after %d milliseconds.", - senderFutureTimeoutMs.value()); - ++it; + auto const elapsedNow = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()) + .count(); + deadlineReached = elapsedNow >= kvTransferTimeoutMs.value(); + } + if (deadlineReached) + { + TLLM_LOG_WARNING( + "Context KV cache transfer for request %ld reached total timeout while waiting " + "for sender future. Marking as error.", + request->mRequestId); + mCacheSender->cancelRequest(*request); + request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + requestsStatus.errorRequestIds.insert(request->mRequestId); + it = mSenderFutures.erase(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) + if (senderFutureTimeoutMs.has_value()) { - request->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); + TLLM_LOG_WARNING( + "Timed out waiting for context KV cache transfer for request %ld after %d " + "milliseconds.", + request->mRequestId, senderFutureTimeoutMs.value()); } - it = mSenderFutures.erase(it); + ++it; } } else @@ -726,14 +733,11 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR std::optional kvTransferTimeoutMs = std::nullopt; if (mCacheTransceiverConfig.has_value()) { - // Behaviour change vs upstream/main, mirroring the context side - // (see checkContextTransferStatus): ``kvTransferTimeoutMs`` now - // applies in both ``blockAll`` and polling modes. Without this, - // mRequesterFutures would accumulate stuck entries, pinning - // generation-side KV blocks and eventually exhausting the pool. - // Healthy transfers complete well under the configured timeout, - // so this widening is defensive; callers that need more headroom - // can raise ``kv_transfer_timeout_ms``. + // Apply the overall transfer deadline in both blockAll and polling + // modes (mirrors checkContextTransferStatus) so a stuck receiver + // cannot pin generation-side KV blocks. Healthy transfers complete + // well under kv_transfer_timeout_ms; raise it if the workload needs + // more headroom. kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); // The per-iteration poll timeout is only relevant when the caller wants to // return after checking at least `atLeastRequestNum` entries. @@ -910,9 +914,25 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR { try { - // 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))); + // Bound the wait so a stuck receiver cannot wedge the + // call; see checkContextTransferStatus for the rationale. + int64_t effectiveSliceMs = 0; + if (senderFutureTimeoutMs.has_value()) + { + effectiveSliceMs = senderFutureTimeoutMs.value(); + } + else if (kvTransferTimeoutMs.has_value()) + { + auto const elapsedNow = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()) + .count(); + effectiveSliceMs = std::max(0, kvTransferTimeoutMs.value() - elapsedNow); + } + else + { + effectiveSliceMs = std::numeric_limits::max(); + } + auto status = future.wait_for(std::chrono::milliseconds(effectiveSliceMs)); if (status == std::future_status::ready) { future.get(); @@ -933,40 +953,43 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR } 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()) + // Re-check the overall deadline inline; if reached, + // mark error here instead of blocking on a follow-up + // future.get(). + bool deadlineReached = false; + if (kvTransferTimeoutMs.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; + auto const elapsedNow = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()) + .count(); + deadlineReached = elapsedNow >= kvTransferTimeoutMs.value(); } - else + if (deadlineReached) { - // 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(); - if (request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR) + bool const firstTimeout = mTimedOutRequesterIds.insert(request->mRequestId).second; + if (firstTimeout) { TLLM_LOG_WARNING( - "Generation KV cache transfer for request %ld finished after an error state was set.", + "Generation KV cache transfer for request %ld reached total timeout while " + "waiting for receiver future. Requesting cancellation and marking as error.", request->mRequestId); - ++numErrored; - } - else - { - completeEntry(request); - ++numCompleted; + mCacheReceiver->cancelRequest(*request); } + request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); mTimedOutRequesterIds.erase(request->mRequestId); it = mRequesterFutures.erase(it); + ++numErrored; + } + else + { + 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 @@ -1012,6 +1035,12 @@ bool CacheTransceiver::checkGenTransferComplete() const return mRequesterFutures.empty(); } +bool CacheTransceiver::hasPoisonedTransferBuffer() const +{ + return std::any_of(mCacheTransBufferManagerPtrs.begin(), mCacheTransBufferManagerPtrs.end(), + [](BaseTransBufferManager const* manager) { return manager != nullptr && manager->hasPoisonedBuffer(); }); +} + bool CacheTransceiver::cancelRequest(std::shared_ptr llmRequest) { if (llmRequest->isContextOnlyRequest()) diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp index 9e95ef336a97..898daf845728 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp @@ -210,10 +210,6 @@ 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++) diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.h b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.h index 1bc5daf87622..740f647c44ac 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.h +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.h @@ -32,6 +32,7 @@ class MLACacheFormatter final : public BaseCacheFormatter , mCacheTransBufferManagers{cacheTransBufferManagers} { TLLM_CHECK(mCacheManager); + checkZeroCopyDisabledForCancellableDisaggTransfers(); } void format(tensorrt_llm::batch_manager::TransferSession& session) override; diff --git a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp index 1fd1cbdc253f..53e7a61d1f0d 100644 --- a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp @@ -37,6 +37,7 @@ RnnCacheFormatter::RnnCacheFormatter(rnn_state_manager::RnnStateManager* rnnStat { TLLM_CHECK(mRnnStateManager != nullptr); TLLM_CHECK(mRnnCacheTransBufferManager != nullptr); + kv_cache_manager::checkZeroCopyDisabledForCancellableDisaggTransfers(); } void RnnCacheFormatter::format(TransferSession& session) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils/transferAgent.cpp index f16ece15526f..570e38073423 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils/transferAgent.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -166,6 +166,15 @@ TransferState MooncakeTransferStatus::wait(int64_t timeout_ms) const return true; } +[[nodiscard]] bool MooncakeTransferStatus::release() +{ + // Mooncake does not expose a per-request handle release/cancel primitive + // here. Report the release request as accepted so the generic cancel path + // can unwind; upper layers still decide whether transfer buffers must be + // quarantined when transport quiescence is unknown. + return true; +} + std::string const MooncakeBase64Helper::STANDARD_CHARS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" "abcdefghijklmnopqrstuvwxyz" diff --git a/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils/transferAgent.h b/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils/transferAgent.h index 42f870722e4f..92c99fca24c8 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils/transferAgent.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils/transferAgent.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -37,6 +37,8 @@ class MooncakeTransferStatus final : public TransferStatus TransferState wait(int64_t timeout_ms = -1) const override; + [[nodiscard]] bool release() override; + private: transfer_engine_t mEngine; uint64_t mBatchId; 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 23f061037353..f67cec1f0433 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp @@ -492,16 +492,18 @@ NixlTransferStatus::~NixlTransferStatus() TransferState NixlTransferStatus::wait(int64_t timeout_ms) const { - if (mHandle == nullptr) - { - return TransferState::kFAILURE; - } - auto startTime = std::chrono::steady_clock::now(); - while (true) { - auto status = mRawAgent->getXferStatus(mHandle); + nixl_status_t status; + { + std::lock_guard lock(mHandleMutex); + if (mHandle == nullptr) + { + return TransferState::kFAILURE; + } + status = mRawAgent->getXferStatus(mHandle); + } if (status == NIXL_SUCCESS) { return TransferState::kSUCCESS; @@ -533,6 +535,7 @@ TransferState NixlTransferStatus::wait(int64_t timeout_ms) const [[nodiscard]] bool NixlTransferStatus::isCompleted() const { + std::lock_guard lock(mHandleMutex); if (mHandle == nullptr) { return false; @@ -542,6 +545,7 @@ TransferState NixlTransferStatus::wait(int64_t timeout_ms) const [[nodiscard]] bool NixlTransferStatus::release() { + std::lock_guard lock(mHandleMutex); if (mHandle == nullptr) { return true; 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 26504e1ecf11..6b91a7745bd0 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h @@ -21,6 +21,7 @@ #include "tensorrt_llm/executor/transferAgent.h" #include #include +#include #include namespace tensorrt_llm::executor::kv_cache @@ -74,6 +75,11 @@ class NixlTransferStatus final : public TransferStatus private: nixlAgent* mRawAgent{}; nixlXferReqH* mHandle{}; + // Serializes all access to mHandle (and the NIXL handle it points to) + // across wait(), isCompleted(), and release(), because multiple Python + // threads holding the same BindingsNixlTransferStatus wrapper can call + // into the C++ methods concurrently. + mutable std::mutex mHandleMutex; }; class NixlTransferAgent final : public BaseTransferAgent diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index b0df67ba965a..4cbc576fcf1a 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -110,7 +110,8 @@ 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("has_poisoned_transfer_buffer", &BaseCacheTransceiver::hasPoisonedTransferBuffer); 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..2b3b381100c4 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -122,6 +122,9 @@ def check_gen_transfer_complete(self): def cancel_request(self, req: LlmRequest): raise NotImplementedError + def has_poisoned_transfer_buffer(self) -> bool: + return False + @abstractmethod def prepare_context_requests(self, requests: List[LlmRequest]): """ @@ -205,6 +208,9 @@ def check_gen_transfer_complete(self): def cancel_request(self, req: LlmRequest): return self.impl.cancel_request(req) + def has_poisoned_transfer_buffer(self) -> bool: + return self.impl.has_poisoned_transfer_buffer() + def prepare_context_requests(self, requests: List[LlmRequest]): # not implemented, an empty placeholder to allow being invoked unconditionally ... diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 8bebbf4d8ee0..17975e41a7cb 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3620,10 +3620,15 @@ def _check_cache_transfer_errors(self, error_msg_prefix: str): """Common helper to check for and handle cache transfer errors.""" error_requests = self._get_disagg_reqs_in_error_state() if error_requests: + poisoned_transfer_buffer = ( + self.kv_cache_transceiver is not None + and self.kv_cache_transceiver.has_poisoned_transfer_buffer()) self._handle_errors( - f"Error in kv cache transfer for {error_msg_prefix}", + f"Error in kv cache transfer for {error_msg_prefix}" + + ("; unrecoverable poisoned transfer buffer requires process restart" + if poisoned_transfer_buffer else ""), requests=error_requests, - charge_budget=False) + charge_budget=poisoned_transfer_buffer) @nvtx_range("_check_disagg_ctx_cache_transfer_status") def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): @@ -4241,6 +4246,12 @@ def _handle_responses(self): f"Cancelled timed-out generation KV transfer for " f"request {request.py_request_id}; waiting for " "C++ transfer status to report final cleanup") + # Keep the request in active_requests so the deferred-cleanup + # contract still holds: _check_disagg_gen_cache_transfer_status + # will eventually transition it to DISAGG_TRANS_ERROR once C++ + # surfaces the cancellation, and _check_cache_transfer_errors + # then enqueues the final error response and runs termination. + new_active_requests.append(request) continue if request.is_generation_only_request() and not request.is_finished: From 94ceef595b1f56ad105da2fa0966109b0c173703 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 12 May 2026 16:59:37 -0700 Subject: [PATCH 11/24] [https://nvbugs/6104831][fix] Update cacheTransceiverTest for shared_ptr async API Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../multi_gpu/cacheTransceiverTest.cpp | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp index 8f0c8c8dbe7e..be32be12dd91 100644 --- a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp @@ -348,11 +348,11 @@ class SymmetricalCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines- TLLM_CUDA_CHECK(cudaMemset(it->data(), llmRequest->getPromptLen(), it->getSizeInBytes())); } } - mFutures.emplace_back(mSender->sendAsync(*llmRequest)); + mFutures.emplace_back(mSender->sendAsync(llmRequest)); } else { - auto future = mRequester->receiveAsync(*llmRequest); + auto future = mRequester->receiveAsync(llmRequest); future.get(); TLLM_CUDA_CHECK(cudaDeviceSynchronize()); auto blockRange = BlockRange::fromAllBlockIds(*mManager, llmRequest->mRequestId); @@ -468,12 +468,12 @@ struct CPMetaData struct WrappedLlmRequest { - std::unique_ptr mLlmRequest; + std::shared_ptr mLlmRequest; std::optional mCPMetaData; using RequestIdType = LlmRequest::RequestIdType; - WrappedLlmRequest(std::unique_ptr llmRequest, std::optional cpMetaData) + WrappedLlmRequest(std::shared_ptr llmRequest, std::optional cpMetaData) : mLlmRequest(std::move(llmRequest)) , mCPMetaData(std::move(cpMetaData)) { @@ -887,7 +887,7 @@ class AsymmetricalCacheTest : public ::testing::TestWithParam(mRequestId++, std::move(request)); + auto llmRequestPtr = std::make_shared(mRequestId++, std::move(request)); return std::make_unique(std::move(llmRequestPtr), cpMetaData); } @@ -919,7 +919,7 @@ class AsymmetricalCacheTest : public ::testing::TestWithParamsetCacheState(cacheState); auto stats = texec::ContextPhaseParams({}, requestId, state.release(), std::nullopt); request.setContextPhaseParams(std::move(stats)); - auto llmRequestPtr = std::make_unique(requestId, std::move(request)); + auto llmRequestPtr = std::make_shared(requestId, std::move(request)); return std::make_unique(std::move(llmRequestPtr), cpMetaData); } @@ -973,7 +973,7 @@ class AsymmetricalCacheTest : public ::testing::TestWithParamsendAsync(*llmRequest); + auto future = mSender->sendAsync(llmRequest); return future; } @@ -984,7 +984,7 @@ class AsymmetricalCacheTest : public ::testing::TestWithParammLlmRequest; mManager->addSequenceBatch( {{{llmRequest->mRequestId, llmRequest->getNumTokens(beamIdx), beamWidth}}}, {std::ref(*llmRequest)}); - return mRequester->receiveAsync(*llmRequest); + return mRequester->receiveAsync(llmRequest); } void generationVerifyKVCache(std::shared_ptr const& request) From 0a2d5268e6d12972185f337ea6dc77796491ae33 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 14 May 2026 14:38:42 -0700 Subject: [PATCH 12/24] [https://nvbugs/6104831][feat] Gate disagg mid-flight cancel surface behind TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheFormatter.cpp | 10 +++- .../batch_manager/mlaCacheFormatter.cpp | 9 +++- cpp/tensorrt_llm/common/envUtils.cpp | 6 +++ cpp/tensorrt_llm/common/envUtils.h | 11 ++++ .../_torch/pyexecutor/kv_cache_transceiver.py | 52 +++++++++++++++++++ 5 files changed, 86 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index 2a1a443634de..4500e11aa5da 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -589,12 +589,20 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio } catch (...) { - if (agentConnection != nullptr) + if (agentConnection != nullptr && common::getEnvDisaggEnableInflightCancel()) { // AgentConnection::send can throw on cancel/error after the // backend transfer has seen this send buffer. release() is // not a quiescence proof, so do not return the slot to the // pool; poison it and let Python restart the process. + // + // Gated on TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL: poison the + // buffer only when the operator has explicitly opted in to + // the mid-flight cancellation surface. When opt-in is off + // the cancel chain is dormant by design (Python skips + // cancel_request), so this catch block can still execute + // from non-cancel exceptions but should not poison the + // buffer pool / drive Layer 5's fail-closed in py_executor. sendHolder.poison(); } throw; diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp index 898daf845728..81442d098089 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp @@ -400,8 +400,15 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses } catch (...) { - if (agentConnection != nullptr) + if (agentConnection != nullptr && common::getEnvDisaggEnableInflightCancel()) { + // Symmetric with cacheFormatter.cpp catch: AgentConnection::send + // can throw on cancel/error after the backend transfer has seen + // this send buffer; poison the holder so Layer 5 fail-closed + // takes the pod out of service rather than reusing a + // possibly-still-pinned buffer. + // Gated on TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL: see + // cacheFormatter.cpp for the full rationale. sendHolder.poison(); } throw; diff --git a/cpp/tensorrt_llm/common/envUtils.cpp b/cpp/tensorrt_llm/common/envUtils.cpp index bf3142a160dd..ab3bc1e04108 100644 --- a/cpp/tensorrt_llm/common/envUtils.cpp +++ b/cpp/tensorrt_llm/common/envUtils.cpp @@ -410,6 +410,12 @@ bool getEnvTryZCopyForKVCacheTransfer() return zcopyForSysmmetricKVCache; } +bool getEnvDisaggEnableInflightCancel() +{ + static bool const enabled = getBoolEnv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL"); + return enabled; +} + bool getEnvForceDeterministic() { static bool const forceDeterministic = getBoolEnv("FORCE_DETERMINISTIC"); diff --git a/cpp/tensorrt_llm/common/envUtils.h b/cpp/tensorrt_llm/common/envUtils.h index 82c56f300362..b2ec54a73633 100644 --- a/cpp/tensorrt_llm/common/envUtils.h +++ b/cpp/tensorrt_llm/common/envUtils.h @@ -115,6 +115,17 @@ std::string const& getEnvKVCacheTimeOutputPath(); bool getEnvTryZCopyForKVCacheTransfer(); +// Opt-in for the disagg mid-flight cancellation surface (Layer 1 send/recv +// poll loop, Layer 2b sendHolder.poison() catch-block call, Layer 5 Python +// fail-closed). Returns true when ``TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1``; +// defaults to false (cancel chain dormant, behavior matches the pre-PR-13713 +// baseline). The orthogonal fixes (BufferIndexHolder RAII, shared_ptr +// async lifetime, recv-side idempotency) remain always-on because they close +// baseline races that exist regardless of mid-flight cancellation. See the +// NVBug 6104831 investigation report section 10 for the empirical case for +// enabling this surface in production deployments. +bool getEnvDisaggEnableInflightCancel(); + // Force deterministic behavior for all kernels. bool getEnvForceDeterministic(); diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index 2b3b381100c4..d06704f2c7a2 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -18,6 +18,44 @@ CacheTransBufferManagerCpp = tensorrt_llm.bindings.internal.batch_manager.CacheTransBufferManager BackendTypeCpp = tensorrt_llm.bindings.executor.CacheTransceiverBackendType +# Opt-in for the PR #13713 disagg mid-flight cancellation surface +# (Layer 1 cancel propagation + Layer 2b sendHolder.poison() + Layer 5 +# fail-closed). Setting ``TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1`` arms the +# defensive surface: Python propagates cancel for timed-out transfers, the +# C++ poll loop unwinds workers on cancel, and the poisoned-buffer signal +# drives a graceful PyExecutor shutdown when transport quiescence cannot be +# proved. Defaults to "0" (cancel chain dormant; behavior matches the +# pre-PR-13713 baseline). The orthogonal fixes (BufferIndexHolder RAII, +# shared_ptr async lifetime, recv-side idempotency) remain +# always-on because they close baseline races that exist regardless of +# mid-flight cancellation. See the NVBug 6104831 investigation report +# section 10 for the empirical case for enabling this surface. +_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV = "TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL" +_disagg_inflight_cancel_enabled_cache: Optional[bool] = None + + +def is_disagg_inflight_cancel_enabled() -> bool: + """Return True iff the disagg mid-flight cancellation surface is opted + into via ``TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1``. Read once and + cached; logs a one-shot WARNING on the first call when the knob is set + so the operating mode is visible in logs. Defaults to False (cancel + chain dormant; behavior matches the pre-PR-13713 baseline). + """ + global _disagg_inflight_cancel_enabled_cache + if _disagg_inflight_cancel_enabled_cache is None: + _disagg_inflight_cancel_enabled_cache = (getenv( + _DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "0") == "1") + if _disagg_inflight_cancel_enabled_cache: + logger.warning( + f"{_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV}=1: disagg KV transfer " + "mid-flight cancellation and fail-closed memory-safety policy " + "are ENABLED. Timed-out transfers will be cancelled at the " + "C++ layer; on cancel-with-unknown-quiescence the PyExecutor " + "will be shut down via the fail-closed policy and the " + "orchestrator should restart the pod. See the NVBug 6104831 " + "investigation report section 10 for the rationale.") + return _disagg_inflight_cancel_enabled_cache + def mapping_to_world_config(mapping: Mapping) -> WorldConfig: @@ -206,9 +244,23 @@ def check_gen_transfer_complete(self): return self.impl.check_gen_transfer_complete() def cancel_request(self, req: LlmRequest): + # When TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL is not set (the default), + # the C++ cancel surface stays dormant — Python returns False so + # callers skip the "cancelled, waiting for C++ to surface final + # cleanup" dedup-set entry and the natural-completion path takes + # over. This matches the pre-PR-13713 baseline behavior; see + # is_disagg_inflight_cancel_enabled() for the opt-in details. + if not is_disagg_inflight_cancel_enabled(): + return False return self.impl.cancel_request(req) def has_poisoned_transfer_buffer(self) -> bool: + # Layer 5 fail-closed is gated by the same opt-in: when the cancel + # surface is not enabled no poison signal is ever set by Layer 2b + # (the C++ catch-block call is also gated), so reporting False + # unconditionally here is consistent + defensive. + if not is_disagg_inflight_cancel_enabled(): + return False return self.impl.has_poisoned_transfer_buffer() def prepare_context_requests(self, requests: List[LlmRequest]): From 6a9869b7519918163624fa4620f0a3ea2b675bb0 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 18 May 2026 15:48:46 -0700 Subject: [PATCH 13/24] [NVBUGS-6104831][refactor] bundle disagg KV transfer timeout enforcement under inflight-cancel flag Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 10 +- .../batch_manager/cacheTransceiver.cpp | 186 ++++++++++++------ cpp/tensorrt_llm/common/envUtils.h | 16 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 36 ++-- tests/unittest/llmapi/test_llm_pytorch.py | 139 ++++++++++++- 5 files changed, 304 insertions(+), 83 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index af2cb9b81395..9cd723cddbf3 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -286,9 +286,17 @@ class CacheTransceiver : public BaseCacheTransceiver std::unique_ptr mCacheReceiver; std::vector, std::future>> mSenderFutures; std::vector, std::future>> mRequesterFutures; - std::unordered_set mTimedOutRequesterIds; mpi::MpiComm const* mMpiWorldComm{nullptr}; + // First-timeout dedup sets for the kvTransferTimeoutMs enforcement, + // by role: mTimedOutSenderIds tracks context-side (mSenderFutures) + // and mTimedOutRequesterIds tracks generation-side (mRequesterFutures). + // Used to emit at most one timeout warning and at most one cancel + // request per request per role. Erased at every future-list cleanup + // site so the sets stay bounded. + std::unordered_set mTimedOutSenderIds; + std::unordered_set mTimedOutRequesterIds; + std::shared_ptr mGroupComm; std::shared_ptr mGroupTensorParaComm, mGroupPipeParaComm, mGroupDataComm, mGroupTPInDPComm; diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 93a4b398e8d0..6f5a9fc9486f 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -67,6 +67,16 @@ inline int currentRankForLog() { return useMPI() ? mpi::MpiComm::world().getRank() : tensorrt_llm::pg_utils::get_world_pg()->getRank(); } + +/// @brief Trailing-clause appended to kvTransferTimeoutMs warnings when the +/// in-flight cancellation surface is disabled. The detection itself runs +/// unconditionally so wedges are observable in logs even in flag-off mode; +/// this clause names the knob operators must flip to actually cancel. +/// See TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL in envUtils.h. +constexpr char const* kObserveOnlyTimeoutTail + = "TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=0; continuing to wait. " + "Set TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 to enable cancellation " + "and error surfacing."; } // namespace std::mutex CacheTransceiver::mDllMutex; @@ -600,20 +610,37 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( 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; + // Dedup is shared with the inline wait_for re-check below + // so a single hung request produces at most one warning + // regardless of which site observes the expired deadline. + bool const firstTimeout = mTimedOutSenderIds.insert(request->mRequestId).second; + bool const inflightCancelEnabled = common::getEnvDisaggEnableInflightCancel(); + if (firstTimeout) + { + TLLM_LOG_WARNING( + "Context KV cache transfer for request %ld exceeded total timeout: " + "elapsed %ld ms > limit %d ms. %s", + request->mRequestId, elapsedMs, kvTransferTimeoutMs.value(), + inflightCancelEnabled ? "Marking as error." : kObserveOnlyTimeoutTail); + } + if (inflightCancelEnabled) + { + // 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); + mTimedOutSenderIds.erase(request->mRequestId); + it = mSenderFutures.erase(it); + continue; + } + // Observe-only: cannot safely cancel the in-flight transfer + // (would risk a UAF on send buffers still touched by the + // network thread). Fall through to the readiness gate so a + // transfer that completes naturally is still detected. } } if (blockAll || (toCompleteIdSet.find(request->mRequestId) != toCompleteIdSet.end())) @@ -648,6 +675,9 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( { request->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); } + // Clear any observe-only timeout breadcrumb so the dedup + // set does not leak across the request's lifetime. + mTimedOutSenderIds.erase(request->mRequestId); it = mSenderFutures.erase(it); } else if (status == std::future_status::timeout) @@ -665,14 +695,29 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } if (deadlineReached) { - TLLM_LOG_WARNING( - "Context KV cache transfer for request %ld reached total timeout while waiting " - "for sender future. Marking as error.", - request->mRequestId); - mCacheSender->cancelRequest(*request); - request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); - requestsStatus.errorRequestIds.insert(request->mRequestId); - it = mSenderFutures.erase(it); + bool const firstTimeout = mTimedOutSenderIds.insert(request->mRequestId).second; + bool const inflightCancelEnabled = common::getEnvDisaggEnableInflightCancel(); + if (firstTimeout) + { + TLLM_LOG_WARNING( + "Context KV cache transfer for request %ld reached total timeout " + "while waiting for sender future. %s", + request->mRequestId, + inflightCancelEnabled ? "Marking as error." : kObserveOnlyTimeoutTail); + } + if (inflightCancelEnabled) + { + mCacheSender->cancelRequest(*request); + request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + requestsStatus.errorRequestIds.insert(request->mRequestId); + mTimedOutSenderIds.erase(request->mRequestId); + it = mSenderFutures.erase(it); + } + else + { + // Observe-only: continue waiting for natural completion. + ++it; + } } else { @@ -693,6 +738,7 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); requestsStatus.errorRequestIds.insert(request->mRequestId); + mTimedOutSenderIds.erase(request->mRequestId); it = mSenderFutures.erase(it); } } @@ -706,6 +752,7 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( request->mRequestId, e.what()); request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); requestsStatus.errorRequestIds.insert(request->mRequestId); + mTimedOutSenderIds.erase(request->mRequestId); it = mSenderFutures.erase(it); } } @@ -869,45 +916,58 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR if (elapsedMs > kvTransferTimeoutMs.value()) { bool const firstTimeout = mTimedOutRequesterIds.insert(request->mRequestId).second; + bool const inflightCancelEnabled = common::getEnvDisaggEnableInflightCancel(); 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); + "elapsed %ld ms > limit %d ms. %s", + request->mRequestId, elapsedMs, kvTransferTimeoutMs.value(), + inflightCancelEnabled ? "Requesting cancellation." : kObserveOnlyTimeoutTail); } - if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + if (inflightCancelEnabled) { - try + if (firstTimeout) { - future.get(); + // 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); } - catch (std::exception const& e) + if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) { - TLLM_LOG_WARNING( - "Generation KV cache transfer for timed-out request %ld finished with error: %s", - request->mRequestId, e.what()); + 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; } - request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); - mTimedOutRequesterIds.erase(request->mRequestId); - it = mRequesterFutures.erase(it); - ++numErrored; - } - else - { - if (firstTimeout) + else { - ++numErrored; + if (firstTimeout) + { + ++numErrored; + } + ++it; } - ++it; + continue; } - continue; + // Observe-only: memory safety forbids requesting a worker + // unwind here (the recv buffer pool may still be touched + // by the transport thread). Fall through to the readiness + // gate so a naturally completing transfer is still cleaned + // up normally. } } if (blockAll || toCompleteIdSet.find(request->mRequestId) != toCompleteIdSet.end()) @@ -967,18 +1027,32 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR if (deadlineReached) { bool const firstTimeout = mTimedOutRequesterIds.insert(request->mRequestId).second; + bool const inflightCancelEnabled = common::getEnvDisaggEnableInflightCancel(); if (firstTimeout) { TLLM_LOG_WARNING( - "Generation KV cache transfer for request %ld reached total timeout while " - "waiting for receiver future. Requesting cancellation and marking as error.", - request->mRequestId); - mCacheReceiver->cancelRequest(*request); + "Generation KV cache transfer for request %ld reached total timeout " + "while waiting for receiver future. %s", + request->mRequestId, + inflightCancelEnabled ? "Requesting cancellation and marking as error." + : kObserveOnlyTimeoutTail); + } + if (inflightCancelEnabled) + { + if (firstTimeout) + { + mCacheReceiver->cancelRequest(*request); + } + request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + mTimedOutRequesterIds.erase(request->mRequestId); + it = mRequesterFutures.erase(it); + ++numErrored; + } + else + { + // Observe-only: continue waiting for natural completion. + ++it; } - request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); - mTimedOutRequesterIds.erase(request->mRequestId); - it = mRequesterFutures.erase(it); - ++numErrored; } else { diff --git a/cpp/tensorrt_llm/common/envUtils.h b/cpp/tensorrt_llm/common/envUtils.h index b2ec54a73633..a19b30509cbb 100644 --- a/cpp/tensorrt_llm/common/envUtils.h +++ b/cpp/tensorrt_llm/common/envUtils.h @@ -115,15 +115,13 @@ std::string const& getEnvKVCacheTimeOutputPath(); bool getEnvTryZCopyForKVCacheTransfer(); -// Opt-in for the disagg mid-flight cancellation surface (Layer 1 send/recv -// poll loop, Layer 2b sendHolder.poison() catch-block call, Layer 5 Python -// fail-closed). Returns true when ``TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1``; -// defaults to false (cancel chain dormant, behavior matches the pre-PR-13713 -// baseline). The orthogonal fixes (BufferIndexHolder RAII, shared_ptr -// async lifetime, recv-side idempotency) remain always-on because they close -// baseline races that exist regardless of mid-flight cancellation. See the -// NVBug 6104831 investigation report section 10 for the empirical case for -// enabling this surface in production deployments. +// Opt-in for the disagg mid-flight cancellation surface. Gates the full +// detection + cancellation + mark-as-error triple, which must move together +// for memory safety, e.g., no UAF. +// Returns true when ``TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1``; defaults to +// false (cancel chain dormant: timeouts are still detected and logged once +// per request, but no in-flight cancellation is issued and the future is left +// to complete naturally). bool getEnvDisaggEnableInflightCancel(); // Force deterministic behavior for all kernels. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 17975e41a7cb..dc548ef2b483 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -56,7 +56,8 @@ from .handle_additional_outputs import HandleAdditionalOutputs from .handle_logits import HandleLogits from .hang_detector import HangDetector -from .kv_cache_transceiver import KvCacheTransceiver +from .kv_cache_transceiver import (KvCacheTransceiver, + is_disagg_inflight_cancel_enabled) from .llm_request import (ExecutorRequest, LlmRequest, LlmRequestState, LlmResponse, get_draft_token_length) from .mamba_cache_manager import MambaHybridCacheManager @@ -3277,24 +3278,33 @@ def _check_kv_transfer_timeout(self): timeout_ms = configured_timeout_ms using_fallback = False + cancel_enabled = is_disagg_inflight_cancel_enabled() + observe_only_tail = "" if cancel_enabled else ( + "TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=0; continuing to wait. " + "Set TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 to enable " + "cancellation and error surfacing.") + def flag_if_kv_transfer_timed_out(req: LlmRequest, type: str) -> None: current_time = time.time() if req.py_kv_transfer_start_time is None: return elapsed_time = (current_time - req.py_kv_transfer_start_time) * 1000 - if elapsed_time > timeout_ms and not req.py_kv_transfer_timed_out: - if using_fallback: - logger.warning( - f"Terminating {type} request {req.py_request_id} " - f"after Python fallback deadline ({timeout_ms} ms); " - "kv_transfer_timeout_ms is unset, so this is the " - "only escalation path. Configure " + if elapsed_time <= timeout_ms or req.py_kv_transfer_timed_out: + return + # Build the "{req-id} {deadline-source}" body once, then + # prepend the verb chosen by the cancel-enable flag. + if using_fallback: + body = (f"{type} request {req.py_request_id} after " + f"Python fallback deadline ({timeout_ms} ms); " + "kv_transfer_timeout_ms is unset, so this is " + "the only escalation path. Configure " "kv_transfer_timeout_ms to surface failures sooner.") - else: - logger.warning( - f"Terminating {type} request {req.py_request_id} due to KV cache transfer timeout" - ) - req.py_kv_transfer_timed_out = True + else: + body = (f"{type} request {req.py_request_id} due to " + "KV cache transfer timeout.") + verb = "Terminating" if cancel_enabled else "Observed timeout on" + logger.warning(f"{verb} {body} {observe_only_tail}") + req.py_kv_transfer_timed_out = True for req in self.async_transfer_manager.requests_in_transfer().values(): flag_if_kv_transfer_timed_out(req, "context") diff --git a/tests/unittest/llmapi/test_llm_pytorch.py b/tests/unittest/llmapi/test_llm_pytorch.py index 93b157e5c934..5fb86ffd135f 100644 --- a/tests/unittest/llmapi/test_llm_pytorch.py +++ b/tests/unittest/llmapi/test_llm_pytorch.py @@ -1506,7 +1506,16 @@ async def test_llm_rpc_get_stats_async(): @pytest.mark.part0 @skip_ray @pytest.mark.parametrize("transceiver_runtime", [None, "PYTHON"]) -def test_llm_context_only_timed_out(transceiver_runtime): +def test_llm_context_only_timed_out(transceiver_runtime, monkeypatch): + # PR #13713 moved the V1 KV-transfer timeout enforcement trio + # (detection + cancellation + mark-as-error + free-blocks) behind + # TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL. This test asserts the + # active path (blocks freed after timeout), so opt in explicitly. + # The default (off) is observe-only; see the sibling + # ``test_llm_context_only_timed_out_observe_only`` for the + # dormant-trio contract. + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") + tp_size = 1 use_overlap = False enable_iter_req_stats = False @@ -1576,6 +1585,107 @@ def test_llm_context_only_timed_out(transceiver_runtime): assert final_used_num_blocks == 0 +@pytest.mark.threadleak(enabled=False) +@pytest.mark.part0 +@skip_ray +def test_llm_context_only_timed_out_observe_only(monkeypatch): + """Pin down the observe-only (flag-off) V1 KV transfer timeout contract. + + Mirrors ``test_llm_context_only_timed_out`` but inverts the final + assertion: with ``TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=0`` (the new + default), the timeout detection still fires (one WARN per request + per role, with the observe-only tail naming the knob to flip), but + the cancellation + mark-as-error + free-blocks chain stays dormant. + The timed-out context-only request therefore keeps its KV blocks + until the transfer naturally completes -- preserving the + pre-PR-13713 hang surface in a memory-safe way and providing + visibility-only mode for operators who haven't opted into the + active cancellation surface. + + Restricted to the C++ transceiver (``transceiver_runtime=None``, + UCX backend) because the flag only gates that surface; + ``KvCacheTransceiverV2`` has independent always-on cancel/timeout + primitives and is orthogonal to this contract. + """ + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "0") + + tp_size = 1 + use_overlap = False + enable_iter_req_stats = False + + llm_args_extra = {} + llm_args_extra.update( + dict(enable_iter_perf_stats=True, + enable_iter_req_stats=enable_iter_req_stats, + disable_overlap_scheduler=not use_overlap)) + + llm = LLM(model=llama_model_path, + kv_cache_config=global_kvcache_config, + tensor_parallel_size=tp_size, + cache_transceiver_config=CacheTransceiverConfig( + backend="UCX", kv_transfer_timeout_ms=1000), + **llm_args_extra) + + try: + max_tokens = 1 + sampling_params = SamplingParams(max_tokens=max_tokens) + disaggregated_params = DisaggregatedParams(request_type="context_only") + prompts0 = ["What is your name?"] + prompts1 = ["Nvidia is awesome because"] + + # Send context-only request and wait for stats to populate. + for output in llm.generate(prompts1, + sampling_params=sampling_params, + disaggregated_params=disaggregated_params): + print(output) + + max_retries = 10 + for _ in range(max_retries): + results = llm.get_stats(2) + if len(results) == 1: + break + time.sleep(1) + else: + pytest.fail( + f"Failed to get stats with len==1 after {max_retries} retries") + + assert len(results) == 1 + context_only_used_num_blocks = ( + results[0]["kvCacheStats"]["usedNumBlocks"]) + assert context_only_used_num_blocks > 0, ( + "Context-only request should have allocated KV blocks before " + "the timeout window starts.") + print(f"Context only used num blocks: {context_only_used_num_blocks}") + + # Sleep past kv_transfer_timeout_ms (1000 ms) so the timeout + # detection path runs at the C++ side (one WARN per request). + time.sleep(5) + + # Send a regular request to drive additional executor iterations + # so the timeout-detection path is observably exercised, mirroring + # ``test_llm_context_only_timed_out``. + for output in llm.generate(prompts0, sampling_params=sampling_params): + print(output) + + results = llm.get_stats(2) + assert len(results) == 1 + final_used_num_blocks = results[0]["kvCacheStats"]["usedNumBlocks"] + + # Atomicity contract: with the flag off, the cancellation + + # mark-as-error chain is dormant, so KV blocks of the timed-out + # context-only request remain held. A ``final_used_num_blocks == 0`` + # observation here would mean the gate leaked the cleanup path, + # which we explicitly forbid (it would re-arm the trio behind the + # operator's back). + assert final_used_num_blocks > 0, ( + "Expected blocks to remain held in observe-only mode " + "(TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=0); got " + f"final_used_num_blocks={final_used_num_blocks}. " + "Atomic gate leaked: cleanup ran despite the flag being off.") + finally: + llm.shutdown() + + # This test is to verify that when the KV cache is exhausted and scheduled batch size is 0, the context only request will be aborted due to timeout. @@ -1587,11 +1697,17 @@ def test_llm_context_only_timed_out(transceiver_runtime): @pytest.mark.parametrize("transceiver_runtime", [None, "PYTHON"]) def test_llm_context_only_timed_out_kv_cache_exhausted(sender_future_timeout_ms, backend, - transceiver_runtime): + transceiver_runtime, + monkeypatch): # Python transceiver (V2) only supports NIXL/DEFAULT backends if transceiver_runtime == "PYTHON" and backend == "UCX": pytest.skip("Python transceiver (V2) does not support UCX backend") + # PR #13713: opt into the V1 active-path timeout enforcement so the + # ``final_used_num_blocks == 0`` assertion is reachable. See + # ``test_llm_context_only_timed_out`` for rationale. + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") + tp_size = 1 use_overlap = False enable_iter_req_stats = False @@ -1667,7 +1783,14 @@ def test_llm_context_only_timed_out_kv_cache_exhausted(sender_future_timeout_ms, @skip_ray @pytest.mark.asyncio @pytest.mark.parametrize("transceiver_runtime", [None, "PYTHON"]) -async def test_llm_disagg_gen_cancelled(transceiver_runtime): +async def test_llm_disagg_gen_cancelled(transceiver_runtime, monkeypatch): + # PR #13713: opt into the V1 active-path cancel surface so an abort + # that lands while the request is still in + # ``DISAGG_*_TRANS_IN_PROGRESS`` actually propagates through + # ``BindKvCacheTransceiver.cancel_request`` (gated when the flag is + # off). See ``test_llm_context_only_timed_out`` for rationale. + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") + tp_size = 1 use_overlap = False enable_iter_req_stats = False @@ -1849,7 +1972,15 @@ async def await_and_record(output, req_id: int): @skip_ray @pytest.mark.asyncio @pytest.mark.parametrize("transceiver_runtime", [None, "PYTHON"]) -async def test_llm_disagg_streaming_gen_cancelled(transceiver_runtime): +async def test_llm_disagg_streaming_gen_cancelled(transceiver_runtime, + monkeypatch): + # PR #13713: opt into the V1 active-path cancel surface so an abort + # that lands while the request is still in + # ``DISAGG_*_TRANS_IN_PROGRESS`` actually propagates through + # ``BindKvCacheTransceiver.cancel_request`` (gated when the flag is + # off). See ``test_llm_context_only_timed_out`` for rationale. + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") + tp_size = 1 use_overlap = False enable_iter_req_stats = False From 033534be5bc9054c5c0d670e03cdf387163919b1 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Wed, 20 May 2026 11:34:37 -0700 Subject: [PATCH 14/24] [NVBUGS-6104831][test] flag-gate cancel-class tests + add observe-only siblings Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tests/unittest/llmapi/test_llm_pytorch.py | 32 ++---- .../others/test_kv_cache_transceiver.py | 105 +++++++++++++++++- 2 files changed, 110 insertions(+), 27 deletions(-) diff --git a/tests/unittest/llmapi/test_llm_pytorch.py b/tests/unittest/llmapi/test_llm_pytorch.py index 5fb86ffd135f..8ab9e3c7ae99 100644 --- a/tests/unittest/llmapi/test_llm_pytorch.py +++ b/tests/unittest/llmapi/test_llm_pytorch.py @@ -1505,15 +1505,9 @@ async def test_llm_rpc_get_stats_async(): @pytest.mark.threadleak(enabled=False) @pytest.mark.part0 @skip_ray -@pytest.mark.parametrize("transceiver_runtime", [None, "PYTHON"]) +@pytest.mark.parametrize("transceiver_runtime", [None]) def test_llm_context_only_timed_out(transceiver_runtime, monkeypatch): - # PR #13713 moved the V1 KV-transfer timeout enforcement trio - # (detection + cancellation + mark-as-error + free-blocks) behind - # TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL. This test asserts the - # active path (blocks freed after timeout), so opt in explicitly. - # The default (off) is observe-only; see the sibling - # ``test_llm_context_only_timed_out_observe_only`` for the - # dormant-trio contract. + # Verify the active cancel path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1). monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") tp_size = 1 @@ -1694,7 +1688,7 @@ def test_llm_context_only_timed_out_observe_only(monkeypatch): @skip_ray @pytest.mark.parametrize("sender_future_timeout_ms", [100, 1000]) @pytest.mark.parametrize("backend", ["NIXL", "UCX"]) -@pytest.mark.parametrize("transceiver_runtime", [None, "PYTHON"]) +@pytest.mark.parametrize("transceiver_runtime", [None]) def test_llm_context_only_timed_out_kv_cache_exhausted(sender_future_timeout_ms, backend, transceiver_runtime, @@ -1703,9 +1697,7 @@ def test_llm_context_only_timed_out_kv_cache_exhausted(sender_future_timeout_ms, if transceiver_runtime == "PYTHON" and backend == "UCX": pytest.skip("Python transceiver (V2) does not support UCX backend") - # PR #13713: opt into the V1 active-path timeout enforcement so the - # ``final_used_num_blocks == 0`` assertion is reachable. See - # ``test_llm_context_only_timed_out`` for rationale. + # Verify the active cancel path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1). monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") tp_size = 1 @@ -1782,13 +1774,9 @@ def test_llm_context_only_timed_out_kv_cache_exhausted(sender_future_timeout_ms, @pytest.mark.part0 @skip_ray @pytest.mark.asyncio -@pytest.mark.parametrize("transceiver_runtime", [None, "PYTHON"]) +@pytest.mark.parametrize("transceiver_runtime", [None]) async def test_llm_disagg_gen_cancelled(transceiver_runtime, monkeypatch): - # PR #13713: opt into the V1 active-path cancel surface so an abort - # that lands while the request is still in - # ``DISAGG_*_TRANS_IN_PROGRESS`` actually propagates through - # ``BindKvCacheTransceiver.cancel_request`` (gated when the flag is - # off). See ``test_llm_context_only_timed_out`` for rationale. + # Verify the active cancel path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1). monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") tp_size = 1 @@ -1971,14 +1959,10 @@ async def await_and_record(output, req_id: int): @pytest.mark.part0 @skip_ray @pytest.mark.asyncio -@pytest.mark.parametrize("transceiver_runtime", [None, "PYTHON"]) +@pytest.mark.parametrize("transceiver_runtime", [None]) async def test_llm_disagg_streaming_gen_cancelled(transceiver_runtime, monkeypatch): - # PR #13713: opt into the V1 active-path cancel surface so an abort - # that lands while the request is still in - # ``DISAGG_*_TRANS_IN_PROGRESS`` actually propagates through - # ``BindKvCacheTransceiver.cancel_request`` (gated when the flag is - # off). See ``test_llm_context_only_timed_out`` for rationale. + # Verify the active cancel path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1). monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") tp_size = 1 diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index 184d7e5862cc..dae0b51b7ba8 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -136,6 +136,24 @@ def _drive_template_handshake_to_completion(ctx_xcvr, f"a spurious lingering entry from the template handshake") +@pytest.fixture(autouse=True) +def _reset_disagg_inflight_cancel_flag_cache(monkeypatch): + """Reset the module-level ``is_disagg_inflight_cancel_enabled`` cache + before each test so per-test ``TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL`` + settings are not poisoned by a prior test that initialized the cache. + + ``BindKvCacheTransceiver.cancel_request`` reads the env var once and + caches it in ``_disagg_inflight_cancel_enabled_cache`` (see + ``tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py``). Without + this fixture an active-path test (flag=1) leaves the cache as ``True`` + and a subsequent observe-only test (flag unset) reads stale ``True``; + or vice versa. ``monkeypatch.setattr`` auto-restores the cache at + test teardown so the reset is local to each test. + """ + from tensorrt_llm._torch.pyexecutor import kv_cache_transceiver as _kct + monkeypatch.setattr(_kct, "_disagg_inflight_cancel_enabled_cache", None) + + @pytest.fixture def disagg_transceiver_pair(request): """Build a single-process disagg ctx/gen transceiver pair. @@ -299,7 +317,10 @@ def test_kv_cache_transceiver_single_process(ctx_gen_kv_cache_dtype, @pytest.mark.parametrize("attention_type", [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA], ids=["mha", "mla"]) -def test_cancel_request_in_transmission(attention_type): +def test_cancel_request_in_transmission(attention_type, monkeypatch): + # Verify the active cancel path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1). + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") + # Init kv_cache manager and cache transceiver mapping = Mapping(world_size=1, rank=0) dist = Distributed.get(mapping) @@ -368,6 +389,76 @@ def test_cancel_request_in_transmission(attention_type): assert gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR +@pytest.mark.timeout(60) +@pytest.mark.parametrize("attention_type", + [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA], + ids=["mha", "mla"]) +def test_cancel_request_in_transmission_observe_only(attention_type, + monkeypatch): + """Pin down the flag-off contract for ``cancel_request``. + + Mirrors ``test_cancel_request_in_transmission`` but with + ``TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL`` unset (the default after + the PR #13713 atomicity refactor). The + ``BindKvCacheTransceiver.cancel_request`` shim must short-circuit + to return ``False`` without invoking the C++ cancel surface, so + the in-flight transfer is not interrupted and the ctx request + does not transition to ``DISAGG_TRANS_ERROR``. + + This regression-proofs the gate at the shim layer: any future + refactor that silently re-arms the cancel surface would surface + ``is_cancelled is True`` here and fail the assertion. + """ + # Verify the observe-only path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL unset). + monkeypatch.delenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", raising=False) + + mapping = Mapping(world_size=1, rank=0) + dist = Distributed.get(mapping) + ctx_kv_cache_dtype = DataType.HALF + kv_cache_manager_ctx = create_kv_cache_manager(mapping, ctx_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) + + 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_manager_ctx.impl.refresh_blocks() + kv_cache_transceiver_ctx.respond_and_send_async(ctx_request) + + # Mirror the active test's wait so the request is registered in + # ``mSenderFutures`` by the time we call cancel. + time.sleep(2) + + # Observe-only contract: + is_cancelled = kv_cache_transceiver_ctx.cancel_request(ctx_request) + assert is_cancelled is False, ( + "Flag-off cancel_request must return False without invoking " + "the C++ cancel surface. A True result here indicates the " + "BindKvCacheTransceiver.cancel_request gate has leaked and " + "the active cancel path is firing despite the opt-in being " + "unset.") + assert ctx_request.state != LlmRequestState.DISAGG_TRANS_ERROR, ( + "Flag-off cancel_request must not transition the ctx request " + "to DISAGG_TRANS_ERROR; no cancellation actually occurred at " + "the C++ layer.") + + @pytest.mark.timeout(120) @pytest.mark.parametrize("disagg_transceiver_pair", [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA], @@ -447,7 +538,7 @@ def test_request_and_receive_async_state_ordering(disagg_transceiver_pair): ids=["mha", "mla"], indirect=True) def test_cancel_request_in_transmission_does_not_break_sender_future( - disagg_transceiver_pair, capfd): + disagg_transceiver_pair, capfd, monkeypatch): """Reproduce the sender-side broken-promise on cancel-after-ready. Pre-fix ``CacheSender::Impl::sendResponse`` erases the ready @@ -455,6 +546,9 @@ def test_cancel_request_in_transmission_does_not_break_sender_future( promise, and the destructor surfaces ``future_error: Broken promise`` on the sender future returned by ``respond_and_send_async()``. """ + # Verify the active cancel path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1). + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") + (kv_cache_manager_ctx, kv_cache_manager_gen, kv_cache_transceiver_ctx, kv_cache_transceiver_gen) = disagg_transceiver_pair @@ -691,7 +785,7 @@ def call_blocking_check(): ids=["mha", "mla"], indirect=True) def test_cancel_queued_gen_request_fulfills_receiver_future( - disagg_transceiver_pair, capfd): + disagg_transceiver_pair, capfd, monkeypatch): """Reproduce the receiver-side queued-cancel broken-promise. Mirrors the sender-side reproducer but on @@ -707,6 +801,9 @@ def test_cancel_queued_gen_request_fulfills_receiver_future( queued request. Post-fix the queued promise is fulfilled with a structured cancellation exception before the queue entry is erased. """ + # Verify the active cancel path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1). + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") + (kv_cache_manager_ctx, kv_cache_manager_gen, kv_cache_transceiver_ctx, kv_cache_transceiver_gen) = disagg_transceiver_pair @@ -1055,6 +1152,8 @@ def test_hybrid_cache_transceiver_single_process(backend, hybrid_dtypes, @pytest.mark.parametrize("backend", ["NIXL", "UCX"], ids=["NIXL", "UCX"]) def test_hybrid_cache_transceiver_cancel_request(backend, monkeypatch): monkeypatch.setenv("TRTLLM_USE_CPP_MAMBA", "1") + # Verify the active cancel path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1). + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") mapping = Mapping(world_size=1, rank=0) dtype = DataType.HALF From 82e3b1f1e139c75ac17ed4adab410a3aec2c2a80 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Wed, 20 May 2026 11:55:34 -0700 Subject: [PATCH 15/24] [NVBUGS-6104831][test] fix ruff-legacy F821 in priority-request test closure Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tests/unittest/llmapi/test_llm_pytorch.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unittest/llmapi/test_llm_pytorch.py b/tests/unittest/llmapi/test_llm_pytorch.py index 8ab9e3c7ae99..842418fb5aaa 100644 --- a/tests/unittest/llmapi/test_llm_pytorch.py +++ b/tests/unittest/llmapi/test_llm_pytorch.py @@ -1916,7 +1916,7 @@ def test_priority_request_completes_before_low_priority(): prompt = "A B C D E F G H I J" completion_order = [] - async def run(): + async def run(llm): # Submit all three requests before the event loop can schedule any of # them – they all land in the waiting queue simultaneously. # Request 1 uses ignore_eos so it keeps the batch slot busy for the @@ -1943,7 +1943,7 @@ async def await_and_record(output, req_id: int): timeout=55, ) - asyncio.run(run()) + asyncio.run(run(llm)) del llm assert 2 in completion_order and 3 in completion_order, ( From db8db6913c6d7669ebd0f4ab1e80384ec0d5db47 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Wed, 20 May 2026 12:54:56 -0700 Subject: [PATCH 16/24] [NVBUGS-6104831][chore] update ruff-legacy baseline for line drift Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- ruff-legacy-baseline.json | 76 ++++++++------------------------------- 1 file changed, 15 insertions(+), 61 deletions(-) diff --git a/ruff-legacy-baseline.json b/ruff-legacy-baseline.json index 571197795369..2ef5a26d2a89 100644 --- a/ruff-legacy-baseline.json +++ b/ruff-legacy-baseline.json @@ -1,8 +1,8 @@ { "_meta": { "generated_by": "scripts/legacy_utils.py lint-update-violations", - "total_violations": 5341, - "total_files": 493 + "total_violations": 5252, + "total_files": 485 }, ".github/scripts/label_community_user.py": { "D212": 1 @@ -452,9 +452,6 @@ "D212": 1, "E741": 1 }, - "scripts/check_test_list.py": { - "D212": 1 - }, "scripts/dco_check.py": { "D212": 1 }, @@ -487,11 +484,8 @@ "D212": 1, "E402": 27 }, - "tensorrt_llm/_torch/attention_backend/interface.py": { - "F821": 2 - }, "tensorrt_llm/_torch/attention_backend/sparse/dsa.py": { - "F821": 3 + "F821": 2 }, "tensorrt_llm/_torch/attention_backend/sparse/kernel.py": { "E731": 3 @@ -521,9 +515,6 @@ "tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py": { "E741": 11 }, - "tensorrt_llm/_torch/custom_ops/torch_custom_ops.py": { - "F821": 2 - }, "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py": { "E741": 3 }, @@ -531,10 +522,6 @@ "E712": 1, "E731": 1 }, - "tensorrt_llm/_torch/model_config.py": { - "F811": 1, - "F821": 4 - }, "tensorrt_llm/_torch/models/checkpoints/auto_mapper.py": { "F821": 1 }, @@ -578,9 +565,6 @@ "E731": 1, "F821": 5 }, - "tensorrt_llm/_torch/modules/attention.py": { - "E731": 2 - }, "tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py": { "F821": 2 }, @@ -620,9 +604,6 @@ "tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py": { "F821": 1 }, - "tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py": { - "F821": 1 - }, "tensorrt_llm/_torch/pyexecutor/model_engine.py": { "F821": 1 }, @@ -673,9 +654,7 @@ "D411": 1 }, "tensorrt_llm/bench/build/dataclasses.py": { - "D205": 1, - "D208": 1, - "D210": 5 + "D210": 4 }, "tensorrt_llm/bench/build/tuning.py": { "D205": 2, @@ -836,18 +815,13 @@ "D411": 8, "D415": 9 }, - "tensorrt_llm/inputs/evs.py": { - "D205": 1, - "D212": 2 - }, "tensorrt_llm/inputs/multimodal.py": { - "D202": 1, "D415": 1 }, "tensorrt_llm/inputs/registry.py": { "D200": 3, "D205": 6, - "D212": 17, + "D212": 16, "D301": 1, "D405": 1, "D411": 1, @@ -920,15 +894,10 @@ "tensorrt_llm/llmapi/llm.py": { "D200": 2, "D202": 1, - "D205": 4, - "D207": 2, - "D212": 2, - "D300": 4, - "D410": 2, - "D411": 2 + "D205": 4 }, "tensorrt_llm/llmapi/llm_args.py": { - "D200": 24, + "D200": 23, "D202": 1, "D205": 11, "D208": 2, @@ -1587,8 +1556,7 @@ "D205": 1, "D212": 1, "D300": 1, - "D411": 1, - "D415": 8 + "D415": 7 }, "tensorrt_llm/serve/scripts/benchmark_dataset.py": { "D200": 4, @@ -1691,7 +1659,6 @@ }, "tests/integration/defs/accuracy/test_llm_api_autodeploy.py": { "D205": 1, - "D209": 1, "D415": 1, "F811": 2 }, @@ -1877,7 +1844,7 @@ "D212": 3 }, "tests/integration/defs/perf/open_search_db_utils.py": { - "D200": 2, + "D200": 1, "D212": 4, "E402": 1 }, @@ -1944,12 +1911,7 @@ }, "tests/integration/defs/test_e2e.py": { "D200": 5, - "D202": 3, "D205": 3, - "D210": 1, - "D212": 8, - "D300": 9, - "D403": 4, "D415": 10, "F811": 2 }, @@ -2074,9 +2036,6 @@ "tests/unittest/_torch/modeling/test_modeling_exaone4.py": { "E402": 12 }, - "tests/unittest/_torch/thop/parallel/test_mamba2_chunk_ss_update.py": { - "E731": 4 - }, "tests/unittest/api_stability/api_stability_core.py": { "D202": 1, "D205": 1, @@ -2088,7 +2047,7 @@ }, "tests/unittest/bindings/test_executor_bindings.py": { "D202": 1, - "E712": 56, + "E712": 54, "F403": 2, "F405": 3, "F811": 1 @@ -2149,8 +2108,7 @@ "F811": 1 }, "tests/unittest/llmapi/test_executor.py": { - "D205": 5, - "D209": 3 + "D205": 1 }, "tests/unittest/llmapi/test_llm.py": { "D200": 1, @@ -2164,9 +2122,7 @@ }, "tests/unittest/llmapi/test_llm_args.py": { "D200": 1, - "D202": 2, - "D212": 2, - "E712": 20, + "E712": 19, "F811": 1 }, "tests/unittest/llmapi/test_llm_multi_gpu.py": { @@ -2177,8 +2133,7 @@ "D202": 1, "D210": 2, "D212": 1, - "F811": 1, - "F821": 3 + "F811": 1 }, "tests/unittest/llmapi/test_llm_quant.py": { "D205": 1, @@ -2204,8 +2159,7 @@ "D415": 1 }, "tests/unittest/others/test_kv_cache_transceiver.py": { - "D205": 1, - "D212": 2 + "D205": 1 }, "tests/unittest/others/test_layer.py": { "D205": 1, @@ -2483,7 +2437,7 @@ }, "triton_backend/all_models/tests/test_python_backend.py": { "E402": 2, - "E712": 40, + "E712": 38, "F403": 1, "F405": 72 }, From 3259c8fb3a1744d595d5be022dad706d3fc718af Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 21 May 2026 15:50:02 -0700 Subject: [PATCH 17/24] [NVBUGS-6104831][fix] keep NIXL agent alive while its TransferStatus exists Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../executor/cache_transmission/nixl_utils/agentBindings.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 f9157d59426e..514b00e8c5be 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp @@ -217,7 +217,7 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m) "submit_transfer_requests", [](kvc::BaseTransferAgent& self, kvc::TransferRequest const& request) { return self.submitTransferRequests(request).release(); }, - nb::arg("request"), nb::rv_policy::take_ownership) + nb::arg("request"), nb::rv_policy::take_ownership, nb::keep_alive<0, 1>()) .def( "notify_sync_message", &kvc::BaseTransferAgent::notifySyncMessage, nb::arg("name"), nb::arg("sync_message")) .def("get_notified_sync_messages", &kvc::BaseTransferAgent::getNotifiedSyncMessages) @@ -251,7 +251,8 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m) "submit_transfer_requests", [](kvc::NixlTransferAgent& self, kvc::TransferRequest const& request) { return self.submitTransferRequests(request).release(); }, - nb::arg("request"), nb::rv_policy::take_ownership, nb::call_guard()) + nb::arg("request"), nb::rv_policy::take_ownership, nb::call_guard(), + nb::keep_alive<0, 1>()) .def( "notify_sync_message", &kvc::NixlTransferAgent::notifySyncMessage, nb::arg("name"), nb::arg("sync_message")) .def("get_notified_sync_messages", &kvc::NixlTransferAgent::getNotifiedSyncMessages) From f79d8b7a61ea0006e87efd56bdac77a049b65f8b Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 22 May 2026 01:27:37 -0700 Subject: [PATCH 18/24] [https://nvbugs/6104831][fix] gate deferred-cleanup paths under TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/dataTransceiver.cpp | 18 ++++++++++++++++-- cpp/tensorrt_llm/common/envUtils.h | 12 ++++++------ tensorrt_llm/_torch/pyexecutor/py_executor.py | 18 +++++++++++++----- 3 files changed, 35 insertions(+), 13 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 6e40aec66b0f..83cdaa18d3e2 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -490,7 +490,7 @@ class CacheSender::Impl isCancelled = true; } } - if (!isCancelled) + if (!isCancelled && common::getEnvDisaggEnableInflightCancel()) { // Request is the mCurrentRequest (or not even in mReadyResponses) // — the queue-drain branch can't abort it. Flip the per-request @@ -513,6 +513,13 @@ class CacheSender::Impl TLLM_LOG_WARNING("Cannot cancel request %zu", llmRequest.mRequestId); } } + else if (!isCancelled) + { + TLLM_LOG_WARNING( + "Cannot cancel request %zu (in-flight cancel disabled; " + "set TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 to opt in)", + llmRequest.mRequestId); + } return isCancelled; } @@ -1256,7 +1263,7 @@ class CacheReceiver::Impl std::lock_guard lg(mInFlightCancelMutex); mInFlightCancelFlags.erase(*queuedCancelledReqId); } - if (!isCancelled) + if (!isCancelled && common::getEnvDisaggEnableInflightCancel()) { // Request already dequeued past mRequestsQueue (worker thread has // picked it up; most likely blocked inside requestSync -> @@ -1279,6 +1286,13 @@ class CacheReceiver::Impl TLLM_LOG_WARNING("Cannot cancel request %zu", llmRequest.mRequestId); } } + else if (!isCancelled) + { + TLLM_LOG_WARNING( + "Cannot cancel request %zu (in-flight cancel disabled; " + "set TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 to opt in)", + llmRequest.mRequestId); + } return isCancelled; } diff --git a/cpp/tensorrt_llm/common/envUtils.h b/cpp/tensorrt_llm/common/envUtils.h index a19b30509cbb..0e2684a8cd03 100644 --- a/cpp/tensorrt_llm/common/envUtils.h +++ b/cpp/tensorrt_llm/common/envUtils.h @@ -115,13 +115,13 @@ std::string const& getEnvKVCacheTimeOutputPath(); bool getEnvTryZCopyForKVCacheTransfer(); -// Opt-in for the disagg mid-flight cancellation surface. Gates the full -// detection + cancellation + mark-as-error triple, which must move together -// for memory safety, e.g., no UAF. +// Opt-in for the disagg mid-flight cancellation + deferred-cleanup surface. // Returns true when ``TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1``; defaults to -// false (cancel chain dormant: timeouts are still detected and logged once -// per request, but no in-flight cancellation is issued and the future is left -// to complete naturally). +// false. When unset, timeouts are still detected and logged once per request, +// but no in-flight cancellation is issued, deferred-cleanup paths are +// inactive, and futures are left to complete naturally. The +// detection + cancellation + mark-as-error triple is atomic by design and +// must move together for memory safety. bool getEnvDisaggEnableInflightCancel(); // Force deterministic behavior for all kernels. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 6678fb40a188..7411b6e11d87 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -678,8 +678,11 @@ def on_detected(): # Initialize disagg PP termination handler if needed self._disagg_pp_termination_handler = None if self.dist.pp_size > 1 and self.enable_kv_cache_reuse and self.kv_cache_transceiver: + terminate_fn = (self._do_terminate_request_if_safe + if is_disagg_inflight_cancel_enabled() else + self._do_terminate_request) self._disagg_pp_termination_handler = DisaggPPTerminationHandler( - self.dist, self._do_terminate_request_if_safe) + self.dist, terminate_fn) if self.dist.pp_size > 1: self.event_loop = self._executor_loop_pp @@ -4319,10 +4322,13 @@ def _handle_errors(self, failed_requests = (list(self.active_requests) if requests is None else list(requests)) - deferred_requests = [ - request for request in failed_requests - if self._is_unquiesced_disagg_transfer(request) - ] + if is_disagg_inflight_cancel_enabled(): + deferred_requests = [ + request for request in failed_requests + if self._is_unquiesced_disagg_transfer(request) + ] + else: + deferred_requests = [] if deferred_requests: # Deferred requests rely on a follow-up DISAGG_TRANS_ERROR # transition to surface the failure. The transition is @@ -4377,6 +4383,8 @@ def _is_unquiesced_disagg_transfer(self, request: LlmRequest) -> bool: or request.is_disagg_context_transmission_state) def _can_terminate_request_now(self, request: LlmRequest) -> bool: + if not is_disagg_inflight_cancel_enabled(): + return True if request.is_disagg_context_transmission_state: # STOP-GAP: partial reuse asks _handle_responses to terminate # completed context requests before async KV transfer has always From bdfdf8be02cf2570455e61bb8acf557a0bbb64ed Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 22 May 2026 16:37:13 -0700 Subject: [PATCH 19/24] [https://nvbugs/6104831][fix] remove rank-asymmetric gates on disagg gen-side checkGenTransferStatus Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 97 +++++++++++-------- 1 file changed, 59 insertions(+), 38 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 9efe87ef6501..7165c7013c36 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3632,11 +3632,16 @@ def _schedule(self): @nvtx_range("_check_disagg_gen_transfer_status") def _check_disagg_gen_transfer_status(self): - - need_check = any([ - req.is_disagg_generation_transmission_in_progress - for req in self.active_requests - ]) + # The downstream C++ checkGenTransferStatus runs a cross-rank + # gatherRequestIds collective on the gen-side comm. Any + # rank-asymmetric gate here (e.g. the previous "if need_check" + # predicate, which reads rank-local LlmRequest state from + # self.active_requests) causes ABBA deadlock against the next + # unconditional collective on the same ranks (helix CI #39529 + # hang). All gen-side ranks must enter this call together, so + # the gate is intentionally absent. at_least_num is rank-local + # and only affects the per-request loop behavior inside the C++ + # function (blockAll vs polling), which is safe to diverge. non_gen_first_reqs = [ req for req in self.active_requests if req.py_disaggregated_params and req.py_disaggregated_params. @@ -3645,10 +3650,8 @@ def _check_disagg_gen_transfer_status(self): need_check_one = bool(non_gen_first_reqs) and all( req.is_disagg_generation_transmission_in_progress for req in non_gen_first_reqs) - - if need_check: - at_least_num = 1 if need_check_one else 0 - self._check_disagg_gen_cache_transfer_status(at_least_num) + at_least_num = 1 if need_check_one else 0 + self._check_disagg_gen_cache_transfer_status(at_least_num) return @@ -3838,8 +3841,16 @@ def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests): 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) + # Always trigger _recv_disagg_gen_cache, even when this rank has no + # newly fitting disagg-gen-init requests. The downstream C++ + # checkGenTransferStatus runs a cross-rank gatherRequestIds + # collective on the gen-side comm; skipping it on ranks with no + # local work creates the rank-asymmetric entry that produces an + # ABBA deadlock with the next unconditional collective on the + # same ranks. + # _recv_disagg_gen_cache handles the empty-input case as a no-op + # for the receive work while still entering the collective. + self._recv_disagg_gen_cache(fitting_disagg_gen_init_requests) @nvtx_range("_prepare_disagg_gen_transmission_complete") def _prepare_disagg_gen_transmission_complete(self, scheduled_batch): @@ -3928,35 +3939,43 @@ def _recv_disagg_gen_cache(self, new_gen_reqs): continue recv_reqs.append(req) - if not recv_reqs: - return + # NOTE: do NOT early-exit when recv_reqs is empty. + # + # The _check_disagg_gen_cache_transfer_status call below dispatches + # into C++ checkGenTransferStatus, whose body begins with a + # cross-rank gatherRequestIds collective on the gen-side comm + # (mGroupDataComm with attention-DP, otherwise mGroupComm). A + # rank-asymmetric early-exit here causes ABBA deadlock against + # any subsequent unconditional collective on the same ranks -- + # e.g. _can_queue's tp_allgather under attention-DP, or PP + # step-boundary collectives. + if recv_reqs: + if os.getenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP") == "1": + 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 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 os.getenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP") == "1": - 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: + # Record the transfer start time unconditionally so the Python + # fallback deadline in _check_kv_transfer_timeout can fire even + # when kv_transfer_timeout_ms is not configured. The check there + # picks the configured timeout or the fallback floor. 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 - - # Record the transfer start time unconditionally so the Python - # fallback deadline in _check_kv_transfer_timeout can fire even - # when kv_transfer_timeout_ms is not configured. The check there - # picks the configured timeout or the fallback floor. - for req in recv_reqs: - if req.state == LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS: - req.py_kv_transfer_start_time = time.time() + if req.state == LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS: + req.py_kv_transfer_start_time = time.time() non_gen_first_active = [ req for req in self.active_requests @@ -3966,6 +3985,8 @@ def _recv_disagg_gen_cache(self, new_gen_reqs): block_transfer = bool(non_gen_first_active) and all( req.is_disagg_generation_transmission_in_progress for req in non_gen_first_active) + # Always reach the C++ collective so every gen-side rank + # participates in lockstep with its peers. self._check_disagg_gen_cache_transfer_status(1 if block_transfer else 0) return From 53a0692aa4752b7736e58c9cc1ccc219f2eaf50c Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Sat, 23 May 2026 02:18:20 -0700 Subject: [PATCH 20/24] [https://nvbugs/6104831][fix] remove rank-asymmetric gates on disagg ctx-side checkContextTransferStatus Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 66 ++++++++++++------- 1 file changed, 42 insertions(+), 24 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 7165c7013c36..52e097ceaf25 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1838,17 +1838,34 @@ def _executor_loop_pp(self): and req.py_disaggregated_params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST for req in self.active_requests) - if num_fitting_reqs == 0 and not fitting_disagg_gen_init_requests: - if not all_gen_first: - logger.warning( - "num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache" - ) - self._check_disagg_ctx_cache_transfer_status(1) - elif self.async_transfer_manager.has_any_inflight_requests( - ): - # Non-blocking cleanup of completed/timed-out - # transfers to free KV blocks (see _executor_loop). - self._check_disagg_ctx_cache_transfer_status(0) + # The downstream C++ checkContextTransferStatus runs a + # cross-rank gatherRequestIds collective on the ctx-side + # comm (mGroupTPInDPComm with attention-DP, otherwise + # mGroupTensorParaComm). Any rank-asymmetric gate here + # (e.g. the previous `if num_fitting_reqs == 0` predicate, + # which reads rank-local scheduler / async-transfer-manager + # state and diverges in helix-CP setups by design) causes + # ABBA deadlock against the next unconditional collective + # on the same ranks (sampler NCCL ops, _can_queue's + # tp_allgather, PP step boundary). All ctx-side ranks must + # enter this call together, so the gate is intentionally + # absent. at_least_num is rank-local and only affects the + # per-request loop behavior inside the C++ function + # (blockAll vs polling), which is safe to diverge -- + # mirrors bdfdf8be02's fix for + # _check_disagg_gen_transfer_status. See PR #13713 CI + # build #39569 helix test hang for the failure signature + # on TestQwen3_8B::test_auto_dtype_with_helix + # [pp1tp1cp4, pp1tp2cp2]. + at_least_num = 0 + if (num_fitting_reqs == 0 + and not fitting_disagg_gen_init_requests + and not all_gen_first): + logger.warning( + "num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache" + ) + at_least_num = 1 + self._check_disagg_ctx_cache_transfer_status(at_least_num) self.num_scheduled_requests = scheduled_batch.batch_size @@ -2383,19 +2400,20 @@ def _prepare_and_schedule_batch(self): req.py_disaggregated_params and req.py_disaggregated_params. schedule_style == DisaggScheduleStyle.GENERATION_FIRST for req in self.active_requests) - if num_fitting_reqs == 0 and not fitting_disagg_gen_init_requests: - if not all_gen_first: - logger.warning( - "num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache" - ) - self._check_disagg_ctx_cache_transfer_status(1) - elif self.async_transfer_manager.has_any_inflight_requests(): - # Non-blocking cleanup of completed/timed-out transfers - # to free KV blocks. We avoid the blocking check because - # gen-first requests may be waiting for peer info (which - # would block indefinitely), but completed transfers must - # still be reaped so that KV cache can be reclaimed. - self._check_disagg_ctx_cache_transfer_status(0) + # Same rank-symmetric pattern as _executor_loop_pp above; see + # rationale comment there. The C++ gatherRequestIds collective + # inside checkContextTransferStatus must be entered by every + # ctx-side rank together. at_least_num is rank-local and safe + # to diverge (it only controls per-request loop behavior inside + # the C++ function, not collective entry). + at_least_num = 0 + if (num_fitting_reqs == 0 and not fitting_disagg_gen_init_requests + and not all_gen_first): + logger.warning( + "num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache" + ) + at_least_num = 1 + self._check_disagg_ctx_cache_transfer_status(at_least_num) # In gen-only benchmark mode, all requests must fit in KV cache # simultaneously because generation requests do not release their From dbaf7a11065fc78ad4e4e475c8b2181a998d56de Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Sat, 23 May 2026 21:24:43 -0700 Subject: [PATCH 21/24] [https://nvbugs/6104831][fix] make _pp_retry + cpp ctx-drain rank-symmetric Closes the remaining rank-asymmetric MPI collective entry points exposed by helix-CP timing changes in PR #13713. Build #39604 still hung in: - TestDeepSeekV3Lite::test_auto_dtype_with_helix [fifo_v2-cudagraph:with_padding-pp2tp1cp2] -> _pp_retry_until_can_schedule - TestNemotron3Super120B::test_auto_dtype[use_py_transceiver=False] -> _pp_retry_until_can_schedule - cpp test_asymmetric_executor[llama-4proc-mpi_kvcache-90] -> C++ forwardAsync's rank-local checkContextTransferStatus(1, true) All three failures share the same shape -- a Python/C++ rank-local predicate gating entry into a function that runs a cross-rank gatherRequestIds collective inside, ABBA-deadlocking against the next unconditional collective on the same ranks. Same class as bdfdf8be02 (gen-side) and 53a0692aa4 (ctx-side _executor_loop / _executor_loop_pp). Python changes (py_executor.py): - _pp_retry_until_can_schedule now runs on every rank, with cross-rank consensus (min-reduce for can_schedule, max-reduce for has_any_inflight). Rank 0 short-circuits when consensus says all ranks can schedule; otherwise it participates in lockstep through the retry body. Caller in _executor_loop_pp drops its `if dist.rank != 0` gate around the retry call (but keeps it around the subsequent schedule_request). - has_any_inflight_requests() is augmented with the gen-side equivalent `not check_gen_transfer_complete()` (mRequesterFutures non-empty in C++). Previously the gen-side process always saw inflight=False here, so non-zero gen ranks would raise on first entry when KV cache was exhausted by stuck recvs. The OR-consensus ensures the retry waits on EITHER ctx-side outbound transfers OR gen-side inbound recvs. C++ changes (trtGptModelInflightBatching.cpp): - checkDisaggGenTransferStatus drops its `if (needCheck)` gate around checkGenTransferStatus, mirroring bdfdf8be02 on the Python side. - forwardAsync moves the `checkContextTransferStatus(1, true)` call out of the rank-local `if (fittingRequests.empty() && ...)` conditional, using atLeastNum=0 by default and escalating to 1 on the nothing-fits-locally path. All ranks enter the collective; only the per-request blocking timer inside the C++ function diverges based on atLeastNum, which is safe. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../trtGptModelInflightBatching.cpp | 65 +++++++++----- tensorrt_llm/_torch/pyexecutor/py_executor.py | 88 +++++++++++++++---- 2 files changed, 114 insertions(+), 39 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index 05ed827a9511..a3bae10dcf4a 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -1022,18 +1022,30 @@ void TrtGptModelInflightBatching::forwardAsync(RequestList const& activeRequests { prepareDisaggGenInitRequests(activeRequests, fittingDisaggGenInitRequests); } - if (fittingRequests.empty() && fittingDisaggGenInitRequests.empty()) + // checkContextTransferStatus runs a cross-rank gatherRequestIds + // collective on the ctx-side comm. Entering it from a rank-local + // predicate (the previous `if (empty) ... checkContextTransferStatus(1, + // true)` gate) causes ABBA deadlock against the next unconditional + // collective on the same ranks. All ranks must enter this call + // together; ``atLeastNum`` is rank-local and only controls per-request + // blocking behavior inside the C++ function, which is safe to diverge. + // Mirrors 53a0692aa4 (Python ctx-side fix). See PR #13713 CI build + // #39604 failure of + // ``test_asymmetric_executor[llama-4proc-mpi_kvcache-90]``. + bool const nothingFitsLocally = fittingRequests.empty() && fittingDisaggGenInitRequests.empty(); + if (nothingFitsLocally) { TLLM_LOG_WARNING( "CapacityScheduler didn't schedule any requests in iteration %lu, " "probably because of insufficient resources such as KV cache, " "will try wait for KV cache transfer to complete", mIterCounter); - if (mCacheTransceiver) - { - mCacheTransceiver->checkContextTransferStatus(1, true); - // will free kvCache in next iteration. - } + } + if (mCacheTransceiver) + { + int const atLeastNum = nothingFitsLocally ? 1 : 0; + mCacheTransceiver->checkContextTransferStatus(atLeastNum, true); + // will free kvCache in next iteration when atLeastNum==1. } std::tie(currRequests.contextRequests, currRequests.generationRequests) = (*mMicroBatchScheduler)(fittingRequests, mInflightReqIds, mMaxBatchSizeRuntime, mMaxNumTokensRuntime); @@ -1642,28 +1654,33 @@ void TrtGptModelInflightBatching::checkDisaggGenTransferStatus(RequestList const auto timeStart = std::chrono::steady_clock::now(); - // TODO: - auto const needCheck = std::any_of(activeRequests.begin(), activeRequests.end(), - [](auto const& req) { return req->isDisaggGenerationTransmissionInProgress(); }); - - if (needCheck) - { - auto const needCheckOne = std::all_of(activeRequests.begin(), activeRequests.end(), + // The downstream checkGenTransferStatus runs a cross-rank + // gatherRequestIds collective on the gen-side comm (mGroupTPInDPComm + // with attention-DP, otherwise mGroupTensorParaComm). Any + // rank-asymmetric gate here (the previous `if (needCheck)` predicate + // reads rank-local LlmRequest state from activeRequests) causes ABBA + // deadlock against the next unconditional collective on the same ranks + // (asymmetric_executor[llama-4proc-mpi_kvcache-90] CI failure path). + // All gen-side ranks must enter this call together; ``atLeastNum`` is + // rank-local and only affects the per-request loop behavior inside the + // C++ function (blockAll vs polling), which is safe to diverge. + // Mirrors bdfdf8be02's fix on the Python side + // (_check_disagg_gen_transfer_status). See PR #13713 CI build #39604. + auto const needCheckOne = !activeRequests.empty() + && std::all_of(activeRequests.begin(), activeRequests.end(), [](auto const& req) { return req->isDisaggGenerationTransmissionInProgress(); }); - int atLeastNum = needCheckOne ? 1 : 0; - TLLM_LOG_DEBUG( - mpi::MpiComm::world().getRank(), "noPreppared requests, checkGenTransferStatus atLeastNum:%d", atLeastNum); + int atLeastNum = needCheckOne ? 1 : 0; + TLLM_LOG_DEBUG( + mpi::MpiComm::world().getRank(), "noPreppared requests, checkGenTransferStatus atLeastNum:%d", atLeastNum); - mCacheTransceiver->checkGenTransferStatus(atLeastNum); + mCacheTransceiver->checkGenTransferStatus(atLeastNum); - auto timeEnd = std::chrono::steady_clock::now(); - auto duration = std::chrono::duration(timeEnd - timeStart).count(); - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - "no Prepare checkDisaggGenTransferStatus time:%f ms, " - "needCheckOne:%d,needCheck:%ld,activeRequests.size():%ld", - duration, needCheckOne, needCheck, activeRequests.size()); - } + auto timeEnd = std::chrono::steady_clock::now(); + auto duration = std::chrono::duration(timeEnd - timeStart).count(); + TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), + "no Prepare checkDisaggGenTransferStatus time:%f ms, needCheckOne:%d, activeRequests.size():%ld", duration, + needCheckOne, activeRequests.size()); } void TrtGptModelInflightBatching::prepareDistGenBufferAndDecoder(RequestVector const& generationRequests) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 52e097ceaf25..bf87e51c72ba 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -42,6 +42,7 @@ get_global_profiler, host_profiler_context) from ..distributed import Distributed +from ..distributed.communicator import ReduceOp from ..expert_statistic import ExpertStatistic from ..models.modeling_llama import Llama4ForConditionalGeneration from ..models.modeling_utils import DecoderModelForCausalLM @@ -1745,41 +1746,91 @@ def _pp_schedule_and_propagate(self, microbatch_id: int): def _pp_retry_until_can_schedule(self, scheduled_batch): """ - If current rank cannot run the scheduled batch, it will retry following steps until it has enough KV cache resources or reach maximum retry count: + Wait until every rank can run the scheduled batch, retrying: 1. Wait for cache transceiver to finish at least one cache transmission. 2. Terminate requests that have finished context cache transmission. - 3. Check if current rank has enough KV cache resources to run the scheduled batch. + 3. Check if all ranks have enough KV cache resources to run the scheduled batch. + + Must be entered by EVERY rank on EVERY call. The function performs + cross-rank consensus on ``scheduler.can_schedule`` (rank 0 already + knows it can schedule, having just propagated the result, but + participates so collective entry stays rank-symmetric). Any + rank-asymmetric exit -- early-return, raise, or the inner break -- + causes ABBA deadlock against the next unconditional collective on + the same ranks (the ``_check_disagg_ctx_cache_transfer_status`` + gatherRequestIds call at the end of ``_prepare_and_schedule_batch``, + ``_can_queue``'s ``tp_allgather``, or PP step-boundary collectives). + Mirrors bdfdf8be02 (gen-side ``_check_disagg_gen_transfer_status``) + and 53a0692aa4 (ctx-side ``_check_disagg_ctx_cache_transfer_status`` + in ``_executor_loop`` / ``_executor_loop_pp``). See PR #13713 CI + build #39604 hang in + ``TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-...-pp2tp1cp2]`` + and ``TestNemotron3Super120B::test_auto_dtype[use_py_transceiver=False]``. """ scheduled_batch_requests = scheduled_batch.all_requests() - if self.scheduler.can_schedule(scheduled_batch_requests): + + def _all_ranks_can_schedule() -> bool: + # Bool AND via min-reduce: False on any rank -> consensus False. + local = self.scheduler.can_schedule(scheduled_batch_requests) + if self.dist.world_size <= 1: + return local + return bool(self.dist.allreduce(int(local), op=ReduceOp.MIN)) + + if _all_ranks_can_schedule(): return - logger.warning( - "Cannot run first PP's schedule result due to limited KV cache resources. This may cause bubbles in the PP pipeline. Please consider increasing the KV cache size by setting `free_gpu_memory_fraction` to a larger value." - ) + # Only non-leader ranks log the warning -- rank 0 is here purely to + # keep collective entry symmetric. + if self.dist.rank != 0: + logger.warning( + "Cannot run first PP's schedule result due to limited KV cache resources. This may cause bubbles in the PP pipeline. Please consider increasing the KV cache size by setting `free_gpu_memory_fraction` to a larger value." + ) + if self.kv_cache_transceiver is None: + # Config-level predicate identical across ranks. All raise. raise RuntimeError( "KV cache transceiver is not enabled, but current rank cannot run first PP's schedule result due to limited KV cache resources. This is not expected." ) - if not self.async_transfer_manager.has_any_inflight_requests(): + + # ``has_any_inflight_requests()`` tracks ctx-side outbound transfers + # only (``start_transfer`` is called from ``_send_kv_async``). On a + # gen-side process this is always False; the gen-side equivalent + # signal is ``not check_gen_transfer_complete()`` (mRequesterFutures + # non-empty in C++). Combine both, then OR-reduce across ranks so a + # gen-side process correctly waits on its own pending recvs instead + # of raising immediately. + local_has_inflight = ( + self.async_transfer_manager.has_any_inflight_requests() + or not self.kv_cache_transceiver.check_gen_transfer_complete()) + if self.dist.world_size > 1: + any_has_inflight = bool( + self.dist.allreduce(int(local_has_inflight), op=ReduceOp.MAX)) + else: + any_has_inflight = local_has_inflight + if not any_has_inflight: raise RuntimeError( - "No context cache transmission is in progress, but current rank cannot run first PP's schedule result due to limited KV cache resources. This is not expected." + "No context cache transmission is in progress and no generation cache transfer is pending, but at least one rank cannot run first PP's schedule result due to limited KV cache resources. This is not expected." ) + if self.enable_kv_cache_reuse and self._disagg_pp_termination_handler is not None: + # Config-level predicate identical across ranks. All raise. raise RuntimeError( "Cannot terminate requests in cache transmission and release their KV cache resources when block reuse is enabled. Please consider increasing the KV cache size." ) for retry_count in range(self.pp_scheduler_max_retry_count): - if self.scheduler.can_schedule(scheduled_batch_requests): + # Drain inside the loop BEFORE the consensus check so every rank + # enters the underlying gatherRequestIds collective on every + # iteration. 53a0692aa4 made this callsite rank-symmetric on + # collective entry, so the (1) is now safe for all ranks. + self._check_disagg_ctx_cache_transfer_status(1) + self._check_kv_transfer_timeout() + + if _all_ranks_can_schedule(): break logger.debug( f"Retrying to run first PP's schedule result ({retry_count + 1}/{self.pp_scheduler_max_retry_count})" ) - - # Let cache transceiver finish at least one cache transmission and release requests' KV cache resources - self._check_disagg_ctx_cache_transfer_status(1) - self._check_kv_transfer_timeout() else: raise RuntimeError( f"Reach maximum PP retry count ({self.pp_scheduler_max_retry_count}) but still cannot run first PP's schedule result. Please consider increasing the KV cache size by setting `free_gpu_memory_fraction` to a larger value. Or you can set `TLLM_PP_SCHEDULER_MAX_RETRY_COUNT` to a larger value to allow more retries." @@ -1821,9 +1872,16 @@ def _executor_loop_pp(self): # Stage 0: first PP rank schedules requests and propagates the result to all other PP ranks. scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs = self._pp_schedule_and_propagate( microbatch_id) + # Retry-until-all-ranks-can-schedule must run on EVERY rank + # so the cross-rank consensus inside (and the per-iteration + # _check_disagg_ctx_cache_transfer_status collective) stays + # rank-symmetric. Rank 0 already knows it can schedule + # (having just propagated the result); it short-circuits on + # the first consensus check when every other rank also + # agrees, otherwise it participates in lockstep with peers + # through the retry body. + self._pp_retry_until_can_schedule(scheduled_batch) if self.dist.rank != 0: - # Retry until current rank can run first PP's schedule result. - self._pp_retry_until_can_schedule(scheduled_batch) # Run scheduler locally because scheduler may change llm requests' state. self.scheduler.schedule_request(self.active_requests, self.inflight_req_ids) From e8f194f728d9b0c6d6dc91c58393df1dcb33982d Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Sun, 24 May 2026 09:14:23 -0700 Subject: [PATCH 22/24] Revert "[https://nvbugs/6104831][fix] make _pp_retry + cpp ctx-drain rank-symmetric" This reverts commit dbaf7a11065fc78ad4e4e475c8b2181a998d56de. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../trtGptModelInflightBatching.cpp | 65 +++++--------- tensorrt_llm/_torch/pyexecutor/py_executor.py | 88 ++++--------------- 2 files changed, 39 insertions(+), 114 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index a3bae10dcf4a..05ed827a9511 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -1022,30 +1022,18 @@ void TrtGptModelInflightBatching::forwardAsync(RequestList const& activeRequests { prepareDisaggGenInitRequests(activeRequests, fittingDisaggGenInitRequests); } - // checkContextTransferStatus runs a cross-rank gatherRequestIds - // collective on the ctx-side comm. Entering it from a rank-local - // predicate (the previous `if (empty) ... checkContextTransferStatus(1, - // true)` gate) causes ABBA deadlock against the next unconditional - // collective on the same ranks. All ranks must enter this call - // together; ``atLeastNum`` is rank-local and only controls per-request - // blocking behavior inside the C++ function, which is safe to diverge. - // Mirrors 53a0692aa4 (Python ctx-side fix). See PR #13713 CI build - // #39604 failure of - // ``test_asymmetric_executor[llama-4proc-mpi_kvcache-90]``. - bool const nothingFitsLocally = fittingRequests.empty() && fittingDisaggGenInitRequests.empty(); - if (nothingFitsLocally) + if (fittingRequests.empty() && fittingDisaggGenInitRequests.empty()) { TLLM_LOG_WARNING( "CapacityScheduler didn't schedule any requests in iteration %lu, " "probably because of insufficient resources such as KV cache, " "will try wait for KV cache transfer to complete", mIterCounter); - } - if (mCacheTransceiver) - { - int const atLeastNum = nothingFitsLocally ? 1 : 0; - mCacheTransceiver->checkContextTransferStatus(atLeastNum, true); - // will free kvCache in next iteration when atLeastNum==1. + if (mCacheTransceiver) + { + mCacheTransceiver->checkContextTransferStatus(1, true); + // will free kvCache in next iteration. + } } std::tie(currRequests.contextRequests, currRequests.generationRequests) = (*mMicroBatchScheduler)(fittingRequests, mInflightReqIds, mMaxBatchSizeRuntime, mMaxNumTokensRuntime); @@ -1654,33 +1642,28 @@ void TrtGptModelInflightBatching::checkDisaggGenTransferStatus(RequestList const auto timeStart = std::chrono::steady_clock::now(); - // The downstream checkGenTransferStatus runs a cross-rank - // gatherRequestIds collective on the gen-side comm (mGroupTPInDPComm - // with attention-DP, otherwise mGroupTensorParaComm). Any - // rank-asymmetric gate here (the previous `if (needCheck)` predicate - // reads rank-local LlmRequest state from activeRequests) causes ABBA - // deadlock against the next unconditional collective on the same ranks - // (asymmetric_executor[llama-4proc-mpi_kvcache-90] CI failure path). - // All gen-side ranks must enter this call together; ``atLeastNum`` is - // rank-local and only affects the per-request loop behavior inside the - // C++ function (blockAll vs polling), which is safe to diverge. - // Mirrors bdfdf8be02's fix on the Python side - // (_check_disagg_gen_transfer_status). See PR #13713 CI build #39604. - auto const needCheckOne = !activeRequests.empty() - && std::all_of(activeRequests.begin(), activeRequests.end(), + // TODO: + auto const needCheck = std::any_of(activeRequests.begin(), activeRequests.end(), + [](auto const& req) { return req->isDisaggGenerationTransmissionInProgress(); }); + + if (needCheck) + { + auto const needCheckOne = std::all_of(activeRequests.begin(), activeRequests.end(), [](auto const& req) { return req->isDisaggGenerationTransmissionInProgress(); }); - int atLeastNum = needCheckOne ? 1 : 0; - TLLM_LOG_DEBUG( - mpi::MpiComm::world().getRank(), "noPreppared requests, checkGenTransferStatus atLeastNum:%d", atLeastNum); + int atLeastNum = needCheckOne ? 1 : 0; + TLLM_LOG_DEBUG( + mpi::MpiComm::world().getRank(), "noPreppared requests, checkGenTransferStatus atLeastNum:%d", atLeastNum); - mCacheTransceiver->checkGenTransferStatus(atLeastNum); + mCacheTransceiver->checkGenTransferStatus(atLeastNum); - auto timeEnd = std::chrono::steady_clock::now(); - auto duration = std::chrono::duration(timeEnd - timeStart).count(); - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - "no Prepare checkDisaggGenTransferStatus time:%f ms, needCheckOne:%d, activeRequests.size():%ld", duration, - needCheckOne, activeRequests.size()); + auto timeEnd = std::chrono::steady_clock::now(); + auto duration = std::chrono::duration(timeEnd - timeStart).count(); + TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), + "no Prepare checkDisaggGenTransferStatus time:%f ms, " + "needCheckOne:%d,needCheck:%ld,activeRequests.size():%ld", + duration, needCheckOne, needCheck, activeRequests.size()); + } } void TrtGptModelInflightBatching::prepareDistGenBufferAndDecoder(RequestVector const& generationRequests) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index bf87e51c72ba..52e097ceaf25 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -42,7 +42,6 @@ get_global_profiler, host_profiler_context) from ..distributed import Distributed -from ..distributed.communicator import ReduceOp from ..expert_statistic import ExpertStatistic from ..models.modeling_llama import Llama4ForConditionalGeneration from ..models.modeling_utils import DecoderModelForCausalLM @@ -1746,91 +1745,41 @@ def _pp_schedule_and_propagate(self, microbatch_id: int): def _pp_retry_until_can_schedule(self, scheduled_batch): """ - Wait until every rank can run the scheduled batch, retrying: + If current rank cannot run the scheduled batch, it will retry following steps until it has enough KV cache resources or reach maximum retry count: 1. Wait for cache transceiver to finish at least one cache transmission. 2. Terminate requests that have finished context cache transmission. - 3. Check if all ranks have enough KV cache resources to run the scheduled batch. - - Must be entered by EVERY rank on EVERY call. The function performs - cross-rank consensus on ``scheduler.can_schedule`` (rank 0 already - knows it can schedule, having just propagated the result, but - participates so collective entry stays rank-symmetric). Any - rank-asymmetric exit -- early-return, raise, or the inner break -- - causes ABBA deadlock against the next unconditional collective on - the same ranks (the ``_check_disagg_ctx_cache_transfer_status`` - gatherRequestIds call at the end of ``_prepare_and_schedule_batch``, - ``_can_queue``'s ``tp_allgather``, or PP step-boundary collectives). - Mirrors bdfdf8be02 (gen-side ``_check_disagg_gen_transfer_status``) - and 53a0692aa4 (ctx-side ``_check_disagg_ctx_cache_transfer_status`` - in ``_executor_loop`` / ``_executor_loop_pp``). See PR #13713 CI - build #39604 hang in - ``TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-...-pp2tp1cp2]`` - and ``TestNemotron3Super120B::test_auto_dtype[use_py_transceiver=False]``. + 3. Check if current rank has enough KV cache resources to run the scheduled batch. """ scheduled_batch_requests = scheduled_batch.all_requests() - - def _all_ranks_can_schedule() -> bool: - # Bool AND via min-reduce: False on any rank -> consensus False. - local = self.scheduler.can_schedule(scheduled_batch_requests) - if self.dist.world_size <= 1: - return local - return bool(self.dist.allreduce(int(local), op=ReduceOp.MIN)) - - if _all_ranks_can_schedule(): + if self.scheduler.can_schedule(scheduled_batch_requests): return - # Only non-leader ranks log the warning -- rank 0 is here purely to - # keep collective entry symmetric. - if self.dist.rank != 0: - logger.warning( - "Cannot run first PP's schedule result due to limited KV cache resources. This may cause bubbles in the PP pipeline. Please consider increasing the KV cache size by setting `free_gpu_memory_fraction` to a larger value." - ) - + logger.warning( + "Cannot run first PP's schedule result due to limited KV cache resources. This may cause bubbles in the PP pipeline. Please consider increasing the KV cache size by setting `free_gpu_memory_fraction` to a larger value." + ) if self.kv_cache_transceiver is None: - # Config-level predicate identical across ranks. All raise. raise RuntimeError( "KV cache transceiver is not enabled, but current rank cannot run first PP's schedule result due to limited KV cache resources. This is not expected." ) - - # ``has_any_inflight_requests()`` tracks ctx-side outbound transfers - # only (``start_transfer`` is called from ``_send_kv_async``). On a - # gen-side process this is always False; the gen-side equivalent - # signal is ``not check_gen_transfer_complete()`` (mRequesterFutures - # non-empty in C++). Combine both, then OR-reduce across ranks so a - # gen-side process correctly waits on its own pending recvs instead - # of raising immediately. - local_has_inflight = ( - self.async_transfer_manager.has_any_inflight_requests() - or not self.kv_cache_transceiver.check_gen_transfer_complete()) - if self.dist.world_size > 1: - any_has_inflight = bool( - self.dist.allreduce(int(local_has_inflight), op=ReduceOp.MAX)) - else: - any_has_inflight = local_has_inflight - if not any_has_inflight: + if not self.async_transfer_manager.has_any_inflight_requests(): raise RuntimeError( - "No context cache transmission is in progress and no generation cache transfer is pending, but at least one rank cannot run first PP's schedule result due to limited KV cache resources. This is not expected." + "No context cache transmission is in progress, but current rank cannot run first PP's schedule result due to limited KV cache resources. This is not expected." ) - if self.enable_kv_cache_reuse and self._disagg_pp_termination_handler is not None: - # Config-level predicate identical across ranks. All raise. raise RuntimeError( "Cannot terminate requests in cache transmission and release their KV cache resources when block reuse is enabled. Please consider increasing the KV cache size." ) for retry_count in range(self.pp_scheduler_max_retry_count): - # Drain inside the loop BEFORE the consensus check so every rank - # enters the underlying gatherRequestIds collective on every - # iteration. 53a0692aa4 made this callsite rank-symmetric on - # collective entry, so the (1) is now safe for all ranks. - self._check_disagg_ctx_cache_transfer_status(1) - self._check_kv_transfer_timeout() - - if _all_ranks_can_schedule(): + if self.scheduler.can_schedule(scheduled_batch_requests): break logger.debug( f"Retrying to run first PP's schedule result ({retry_count + 1}/{self.pp_scheduler_max_retry_count})" ) + + # Let cache transceiver finish at least one cache transmission and release requests' KV cache resources + self._check_disagg_ctx_cache_transfer_status(1) + self._check_kv_transfer_timeout() else: raise RuntimeError( f"Reach maximum PP retry count ({self.pp_scheduler_max_retry_count}) but still cannot run first PP's schedule result. Please consider increasing the KV cache size by setting `free_gpu_memory_fraction` to a larger value. Or you can set `TLLM_PP_SCHEDULER_MAX_RETRY_COUNT` to a larger value to allow more retries." @@ -1872,16 +1821,9 @@ def _executor_loop_pp(self): # Stage 0: first PP rank schedules requests and propagates the result to all other PP ranks. scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs = self._pp_schedule_and_propagate( microbatch_id) - # Retry-until-all-ranks-can-schedule must run on EVERY rank - # so the cross-rank consensus inside (and the per-iteration - # _check_disagg_ctx_cache_transfer_status collective) stays - # rank-symmetric. Rank 0 already knows it can schedule - # (having just propagated the result); it short-circuits on - # the first consensus check when every other rank also - # agrees, otherwise it participates in lockstep with peers - # through the retry body. - self._pp_retry_until_can_schedule(scheduled_batch) if self.dist.rank != 0: + # Retry until current rank can run first PP's schedule result. + self._pp_retry_until_can_schedule(scheduled_batch) # Run scheduler locally because scheduler may change llm requests' state. self.scheduler.schedule_request(self.active_requests, self.inflight_req_ids) From 9a1a0329fbc0b121844f9f8a97ab634602819cff Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Wed, 27 May 2026 16:31:38 -0700 Subject: [PATCH 23/24] [https://nvbugs/6104831][fix] Cap disagg KV polling wait_for slice at 50ms Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.cpp | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 0860084af566..b59a712bb33c 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -722,6 +722,27 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( { effectiveSliceMs = std::numeric_limits::max(); } + // Cap the per-call wait slice so the polling thread always + // cycles back to drive MPI/UCX forward progress. Transports + // without a dedicated progress thread (notably mpi_kvcache + // over UCX-over-shm) rely on the app thread calling MPI + // functions to advance in-flight transfers. If the slice is + // unbounded (the fall-through above sets INT64_MAX when no + // timeout is configured), a single wait_for on a not-yet- + // ready future blocks the thread forever -- which means no + // MPI progress -- which means the future never becomes + // ready. The death spiral was empirically observed as + // test_asymmetric_executor[mpi_kvcache] wedging at its + // 300s gtest timeout despite the test's CacheTransceiverConfig + // having neither senderFutureTimeoutMs nor kvTransferTimeoutMs + // configured (nvbug-6104831). + // + // The deadline-check observability above is UNCHANGED: if + // kvTransferTimeoutMs is set, the deadline still fires within + // kMaxPollSliceMs of true expiry. Cancellation semantics + // (gated by TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL) unchanged. + static constexpr int64_t kMaxPollSliceMs = 50; + effectiveSliceMs = std::min(effectiveSliceMs, kMaxPollSliceMs); auto status = future.wait_for(std::chrono::milliseconds(effectiveSliceMs)); if (status == std::future_status::ready) { @@ -1048,6 +1069,13 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR { effectiveSliceMs = std::numeric_limits::max(); } + // Cap the per-call wait slice so the polling thread always + // cycles back to drive MPI/UCX forward progress -- see the + // long comment at the analogous site in + // checkContextTransferStatus for the rationale and the + // empirical mpi_kvcache wedge that motivated this cap. + static constexpr int64_t kMaxPollSliceMs = 50; + effectiveSliceMs = std::min(effectiveSliceMs, kMaxPollSliceMs); auto status = future.wait_for(std::chrono::milliseconds(effectiveSliceMs)); if (status == std::future_status::ready) { From f3622233c3953f69ba7ac522b748ee9497f28d4b Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 28 May 2026 23:56:19 -0700 Subject: [PATCH 24/24] [https://nvbugs/6104831][fix] explicit V2 transceiver shutdown in single-process test to prevent daemon thread leak Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../others/test_kv_cache_transceiver.py | 145 ++++++++++-------- 1 file changed, 81 insertions(+), 64 deletions(-) diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index dae0b51b7ba8..d458038a2423 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -245,72 +245,89 @@ def test_kv_cache_transceiver_single_process(ctx_gen_kv_cache_dtype, mapping, dist, kv_cache_manager_gen, attention_type, cache_transceiver_config) - fill_kv_cache_buffer(kv_cache_manager_ctx) + try: + fill_kv_cache_buffer(kv_cache_manager_ctx) + + # init ctx request + 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) + + if transceiver_runtime == "PYTHON": + disaggregated_params = tensorrt_llm.DisaggregatedParams( + request_type="context_only", + disagg_request_id=uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF) + ctx_request.py_disaggregated_params = disaggregated_params + + kv_cache_manager_ctx.impl.add_sequence_batch( + [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], + [ctx_request]) + # add_sequence_batch must be paired with refresh_blocks so the C++ + # transceiver path observes the freshly-allocated KV-block offsets; + # without this the receiver-side pool buffers stay zero. See + # _add_sequence helper for the same pairing in cancellation tests. + kv_cache_manager_ctx.impl.refresh_blocks() + # send ctx request + kv_cache_transceiver_ctx.respond_and_send_async(ctx_request) + + # init gen request + 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) + + if transceiver_runtime == "PYTHON": + disaggregated_params = tensorrt_llm.DisaggregatedParams( + request_type="generation_only", + disagg_request_id=ctx_request.py_disaggregated_params. + disagg_request_id, + ctx_request_id=ctx_request.request_id, + ctx_dp_rank=ctx_request.context_phase_params.ctx_dp_rank, + ctx_info_endpoint=ctx_request.context_phase_params. + disagg_info_endpoint, + first_gen_tokens=ctx_request.context_phase_params. + first_gen_tokens, + draft_tokens=ctx_request.context_phase_params.draft_tokens) + + gen_request.py_disaggregated_params = disaggregated_params + + kv_cache_manager_gen.impl.add_sequence_batch( + [(gen_request.py_request_id, gen_request.prompt_len, 1)], + [gen_request]) + kv_cache_manager_gen.impl.refresh_blocks() + # send gen request + kv_cache_transceiver_gen.request_and_receive_async(gen_request) - # init ctx request - 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) - - if transceiver_runtime == "PYTHON": - disaggregated_params = tensorrt_llm.DisaggregatedParams( - request_type="context_only", - disagg_request_id=uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF) - ctx_request.py_disaggregated_params = disaggregated_params - - kv_cache_manager_ctx.impl.add_sequence_batch( - [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], [ctx_request]) - # add_sequence_batch must be paired with refresh_blocks so the C++ - # transceiver path observes the freshly-allocated KV-block offsets; - # without this the receiver-side pool buffers stay zero. See - # _add_sequence helper for the same pairing in cancellation tests. - kv_cache_manager_ctx.impl.refresh_blocks() - # send ctx request - kv_cache_transceiver_ctx.respond_and_send_async(ctx_request) - - # init gen request - 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) - - if transceiver_runtime == "PYTHON": - disaggregated_params = tensorrt_llm.DisaggregatedParams( - request_type="generation_only", - disagg_request_id=ctx_request.py_disaggregated_params. - disagg_request_id, - ctx_request_id=ctx_request.request_id, - ctx_dp_rank=ctx_request.context_phase_params.ctx_dp_rank, - ctx_info_endpoint=ctx_request.context_phase_params. - disagg_info_endpoint, - first_gen_tokens=ctx_request.context_phase_params.first_gen_tokens, - draft_tokens=ctx_request.context_phase_params.draft_tokens) - - gen_request.py_disaggregated_params = disaggregated_params - - kv_cache_manager_gen.impl.add_sequence_batch( - [(gen_request.py_request_id, gen_request.prompt_len, 1)], [gen_request]) - kv_cache_manager_gen.impl.refresh_blocks() - # send gen request - kv_cache_transceiver_gen.request_and_receive_async(gen_request) - - kv_cache_transceiver_ctx.check_context_transfer_status(1) - kv_cache_transceiver_gen.check_gen_transfer_status(1) + kv_cache_transceiver_ctx.check_context_transfer_status(1) + kv_cache_transceiver_gen.check_gen_transfer_status(1) - assert torch.equal( - kv_cache_manager_gen.get_buffers(0), - kv_cache_manager_ctx.get_buffers(0)), "different kv-cache values" + assert torch.equal( + kv_cache_manager_gen.get_buffers(0), + kv_cache_manager_ctx.get_buffers(0)), "different kv-cache values" + finally: + # KvCacheTransceiverV2 (transceiver_runtime == "PYTHON") spawns + # daemon threads in its Messenger ("listener") and TransferWorker + # ("_process_task_queue") components. Those threads hold bound- + # method captures to self, creating a reference cycle that delays + # CPython's __del__ -- so pytest-threadleak fires at test teardown + # unless we call shutdown() explicitly. V1 (NIXL/UCX) cleans up + # via the C++ ~CacheTransceiver destructor and doesn't need this. + # shutdown() is idempotent via the self._shutdown guard, so safe + # to call even if the test body raised before send/receive. + if transceiver_runtime == "PYTHON": + kv_cache_transceiver_ctx.shutdown() + kv_cache_transceiver_gen.shutdown() @pytest.mark.timeout(120)