diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 8f8330603893..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; @@ -204,13 +205,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 +226,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 +257,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 +272,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 +281,18 @@ 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 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/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..9b83205451a7 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp @@ -21,11 +21,83 @@ #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 + // 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; +} + +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} @@ -54,9 +126,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 +138,16 @@ void BaseTransBufferManager::freeBufferIndexForSend(std::optional bufferId) freeBufferIndex(mConcurrenceSendResource, bufferId, mSendBufferCount, mOnlyUseDynamicBuffer); } -std::optional BaseTransBufferManager::assignBufferIndexForRecv() +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) { - return assignBufferIndex(mConcurrenceRecvResource, mRecvBufferCount, mOnlyUseDynamicBuffer); + return assignBufferIndex(mConcurrenceRecvResource, mRecvBufferCount, mOnlyUseDynamicBuffer, perRequestCancel, + waitSliceMs, requestIdForLog); } void BaseTransBufferManager::freeBufferIndexForRecv(std::optional bufferId) @@ -74,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) @@ -225,16 +311,46 @@ 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) { + 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 + // 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 resource.mPoisoned.load(std::memory_order_relaxed) + || 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()); + } + } + } + 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++) { @@ -264,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--; @@ -271,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 1efeb89ccc04..ad39cdcd7f19 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h @@ -46,6 +46,128 @@ 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(true) + , 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; + + /// @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{}; + 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,21 +179,51 @@ 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 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 + /// 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. 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. @@ -125,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, @@ -132,9 +285,13 @@ 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); + 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 e05e8d6f76fc..e697c6774ae6 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 { @@ -468,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++) @@ -494,7 +494,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 @@ -573,12 +583,26 @@ 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 + // 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..15090eb765f0 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()) { - senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); + // 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()) + { + 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,198 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR { toCompleteIdSet.insert(requestId); } - if (useMPI()) - { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), " checkGenTransferStatus freqVec requestId: %zu,freq:%d ", - requestId, freq); - } - 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)); + TLLM_LOG_DEBUG( + currentRankForLog(), " checkGenTransferStatus freqVec requestId: %zu,freq:%d ", requestId, freq); } - else + 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) { - TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), - " checkGenTransferStatus toCompleteIdSet size: %zu, atLeastRequestNum: %d ", toCompleteIdSet.size(), - atLeastRequestNum.value_or(0)); - } + request->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); + if (!common::getEnvKVCacheTimeOutputPath().empty()) + { + auto transferSyncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; + updateKVCacheTransferBW(transferSyncComm, request.get()); + } + 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()) { - try + auto elapsed = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()); + auto elapsedMs = static_cast(elapsed.count()); + if (elapsedMs > kvTransferTimeoutMs.value()) { - 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()) + bool const firstTimeout = mTimedOutRequesterIds.insert(request->mRequestId).second; + if (firstTimeout) { - auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; - updateKVCacheTransferBW(syncComm, it->first); + 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; } - 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()) + } + if (blockAll || toCompleteIdSet.find(request->mRequestId) != toCompleteIdSet.end()) + { + try { - 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()); + // 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) + { + future.get(); + 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); + } + 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(); + 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); + } + } + else + { + 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; + } } - else + catch (std::exception const& e) { - 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); + mTimedOutRequesterIds.erase(request->mRequestId); + 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 +1011,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..1e26f43ee964 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,14 @@ class CacheReceiver::Impl return isCancelled; } - bool receiveReadySignal(TransferSession& session) + enum class ReadySignalResult + { + kReady, + kNotReady, + kCancelled, + }; + + ReadySignalResult receiveReadySignalDetailed(TransferSession& session, std::atomic const& perRequestCancel) { bool isReadyFinal = true; bool isReady = false; @@ -1013,28 +1292,74 @@ 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 ReadySignalResult::kCancelled; + } auto* agentConnectionManager = dynamic_cast(mManager); if (agentConnectionManager) { auto* agentConnection = dynamic_cast(connections.at(i)); TLLM_CHECK(agentConnection); - isReady = agentConnection->recvReadySignal( - executor::kv_cache::DataContext{TransceiverTag::kREADY_SIGNAL_TAG, mTerminate}); + // 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(). + 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 + // 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,74 +1372,160 @@ 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()); - TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); - auto session = sendRequestInfo(llmRequest); - session.setTime(TransferSession::kTimeRequestInfo); - bool isReady = receiveReadySignal(session); - if (!isReady) + // Early-out if cancel fired between queueing and dequeue. + if (perRequestCancel.load() || mTerminate.load()) { - // Reuse the error state for the cancelled request. llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); return; } - receiveSync(session); - llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); + TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); + + // 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 poisonRecvHolders = [&recvHolders]() + { + 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()); + + // 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(); + } + } + catch (...) + { + if (agentConnectionManagerForAcq) + { + poisonRecvHolders(); + } + llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); + throw; + } 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 +1562,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 +1589,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 +1644,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 +1667,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 +1710,7 @@ CacheReceiver::CacheReceiver( { } -std::future CacheReceiver::receiveAsync(LlmRequest& llmRequest) const +std::future CacheReceiver::receiveAsync(std::shared_ptr const& llmRequest) const { return mImpl->requestAndReceiveAsyncMultiThreads(llmRequest); } diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h index 4d072a04521b..aa4849495cc7 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h @@ -257,9 +257,12 @@ class CacheSender /// @brief Asynchronously respond to the request and send data. /// @param llmRequest Request object. Its data should be ready when called, and the data for this request - /// should remain valid until future synchronization. + /// should remain valid until future synchronization. Passed as shared_ptr + /// so the async worker keeps the LlmRequest alive past any Python-side + /// _terminate_request; see CacheTransceiver::mSenderFutures for the full + /// lifetime invariant. /// @return Once the data is fully sent, the future object will become valid. - [[nodiscard]] virtual std::future sendAsync(LlmRequest& llmRequest) const; + [[nodiscard]] virtual std::future sendAsync(std::shared_ptr const& llmRequest) const; /// @brief Return the internal communicator status. /// @return The communicator status. @@ -314,9 +317,12 @@ class CacheReceiver /// @brief Asynchronously send a request to receive data. /// @param llmRequest Request object. Its data should be in an allocated but unwritten state when called, and the - /// data for this request should remain intact only after future synchronization. + /// data for this request should remain intact only after future synchronization. Passed as shared_ptr so the + /// async worker keeps the LlmRequest alive past any Python-side + /// _terminate_request; see CacheTransceiver::mSenderFutures for the full + /// lifetime invariant. /// @return Once the data is fully received, the future object will become valid. - [[nodiscard]] virtual std::future receiveAsync(LlmRequest& llmRequest) const; + [[nodiscard]] virtual std::future receiveAsync(std::shared_ptr const& llmRequest) const; virtual TransferSession sendRequestInfo(LlmRequest const& llmRequest); diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp index c72090867f29..47fc8d400c10 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp @@ -30,6 +30,7 @@ #include "tensorrt_llm/runtime/cudaEvent.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" +#include #include #include #include @@ -253,7 +254,17 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses return bufferSizeForTarget; }; auto bufferEleSizes = getBufferSizeForTarget(); - auto cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForSend(); + // RAII wrapper mirrors the receiver-side pattern. Happy-path calls + // sendHolder.release(); other exits auto-release via + // ~BufferIndexHolder. Send-pool CV wait observes the per-request + // cancel flag via the session's DataContext and throws on cancel + // (parity with recv side). + auto const sendReqIdForLog = std::make_optional(static_cast(llmRequest.mRequestId)); + auto const* sendCancelFlag = &session.getDataContext().getTransferTerminate(); + auto cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForSend( + sendCancelFlag, kBufferAcquireSliceMs, sendReqIdForLog); + BufferIndexHolder sendHolder( + *mCacheTransBufferManagers[transferIndexerKCache], cacheBufferId, /*isRecv=*/false, sendReqIdForLog); auto result = mCacheTransBufferManagers[transferIndexerKCache]->getOrAllocateSendBuffers( cacheBufferId, static_cast(pPDomainSize * cPDomainSize), bufferEleSizes, bufferManager); auto& outputSplitCaches = std::get<0>(result); @@ -380,7 +391,8 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses { sendBufferFun(deviceId, pickUpConnections[0]); } - mCacheTransBufferManagers[transferIndexerKCache]->freeBufferIndexForSend(cacheBufferId); + // Atomic happy-path release — silent free + disarm in one noexcept call. + sendHolder.release(); } session.setTime(TransferSession::kTimeTransmissions); session.setTime(TransferSession::kTimePostprocess); diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index 8bb2c0e2ba88..94679a54f022 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -915,7 +915,7 @@ void TrtGptModelInflightBatching::forwardSync() TLLM_CHECK_WITH_INFO(mCacheTransceiver, "Disaggregated serving is not enabled, please check the configuration of " "cacheTransceiverConfig."); - mCacheTransceiver->respondAndSendAsync(llmReq.get()); + mCacheTransceiver->respondAndSendAsync(llmReq); } mSeqSlotManager->freeSequenceSlot(llmReq->mRequestId); } @@ -1596,11 +1596,11 @@ void TrtGptModelInflightBatching::prepareDisaggGenInitRequests( mCacheTransceiver, "Disaggregated serving is not enabled, please check the configuration."); if (common::getEnvDisableKVCacheTransferOverlap()) { - mCacheTransceiver->requestAndReceiveSync(newGenReq.get()); + mCacheTransceiver->requestAndReceiveSync(newGenReq); } else { - mCacheTransceiver->requestAndReceiveAsync(newGenReq.get()); + mCacheTransceiver->requestAndReceiveAsync(newGenReq); } } if (!common::getEnvDisableKVCacheTransferOverlap()) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp index d46defdf50ad..cbf6245e9e65 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()); } @@ -245,9 +302,17 @@ 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}; - mAgentConnectionManager->waitForReadySignal(mRemoteAgentName, readySignalInfo, ctx.getTransferTerminate()); + if (!mAgentConnectionManager->waitForReadySignal(mRemoteAgentName, readySignalInfo, ctx.getTransferTerminate())) + { + return std::nullopt; + } return readySignalInfo.mIsReady; } @@ -582,7 +647,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 +655,7 @@ void AgentConnectionManager::waitForNotification( if (!mIsRunning) { - return; + return false; } updateUnhandledNotifications(); std::scoped_lock lock(mNotificationMutex); @@ -623,7 +688,7 @@ void AgentConnectionManager::waitForNotification( { it = mUnhandledNotifications.erase(it); } - return; + return true; } } } @@ -643,7 +708,7 @@ void AgentConnectionManager::waitForNotification( { it = mUnhandledNotifications.erase(it); } - return; + return true; } } } @@ -663,24 +728,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..c912c557f52c 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h @@ -259,14 +259,21 @@ 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); [[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; @@ -320,11 +327,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/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/disaggregation/base/agent.py b/tensorrt_llm/_torch/disaggregation/base/agent.py index 1aac24d98e75..4c5548bd2472 100644 --- a/tensorrt_llm/_torch/disaggregation/base/agent.py +++ b/tensorrt_llm/_torch/disaggregation/base/agent.py @@ -89,6 +89,9 @@ def is_completed(self) -> bool: ... @abstractmethod def wait(self, timeout_ms: int | None = None) -> bool: ... + def release(self) -> bool: + return False + class BaseTransferAgent(ABC): @abstractmethod diff --git a/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py b/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py index 84f85d41b325..1dd85e177a13 100644 --- a/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py +++ b/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py @@ -30,6 +30,10 @@ def wait(self, timeout_ms=None) -> bool: timeout_ms = -1 return self._cpp_status.wait(timeout_ms) == TransferState.SUCCESS + def release(self) -> bool: + """Release the transfer handle, requesting backend cancel if still active.""" + return self._cpp_status.release() + class BindingsNixlTransferAgent(BaseTransferAgent): """NixlTransferAgent using C++ bindings with GIL release support. diff --git a/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py b/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py index dcfa28210f81..703dc42f33a1 100644 --- a/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py +++ b/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py @@ -26,6 +26,14 @@ def is_completed(self): status = TransferState(self.agent.check_xfer_state(self.handle)) return status == TransferState.DONE + def release(self): + try: + self.handle.release() + return True + except Exception: + logger.exception("Failed to release NIXL transfer handle (agent=%s).", self.agent.name) + return False + def wait(self, timeout_ms=None): start_time = time.time() status = TransferState.PENDING diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 98cf521d1ff4..8d25021a2c9d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -366,6 +366,15 @@ 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() + 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) @@ -554,7 +563,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 @@ -3147,9 +3156,6 @@ def _pad_attention_dp_dummy_request(self): @nvtx_range("_prepare_disagg_gen_init") def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests): if fitting_disagg_gen_init_requests: - disagg_gen_init_to_prepare = ScheduledRequests() - disagg_gen_init_to_prepare.context_requests_last_chunk = fitting_disagg_gen_init_requests - for resource_mgr_type in ( ResourceManagerType.KV_CACHE_MANAGER, ResourceManagerType.SPEC_RESOURCE_MANAGER, @@ -3157,9 +3163,22 @@ def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests): if (resource_mgr_type in self.resource_manager.resource_managers and self.resource_manager. resource_managers[resource_mgr_type] is not None): + prepared_ids = self._disagg_gen_init_prepared_ids[ + resource_mgr_type] + resources_to_prepare = [ + req for req in fitting_disagg_gen_init_requests + if req.py_request_id not in prepared_ids + ] + if not resources_to_prepare: + continue + + disagg_gen_init_to_prepare = ScheduledRequests() + disagg_gen_init_to_prepare.context_requests_last_chunk = resources_to_prepare self.resource_manager.resource_managers[ resource_mgr_type].prepare_resources( disagg_gen_init_to_prepare) + for req in resources_to_prepare: + prepared_ids.add(req.py_request_id) # Trigger KV cache exchange for new disagg_gen_init_requests self._recv_disagg_gen_cache(fitting_disagg_gen_init_requests) @@ -3242,15 +3261,36 @@ def _recv_disagg_gen_cache(self, new_gen_reqs): req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE return + recv_reqs = [] + for req in new_gen_reqs: + if req.py_request_id in self._disagg_gen_kv_recv_started_ids: + continue + recv_reqs.append(req) + + if not recv_reqs: + return + if os.getenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP") == "1": - for req in new_gen_reqs: - self.kv_cache_transceiver.request_and_receive_sync(req) + for req in recv_reqs: + self._disagg_gen_kv_recv_started_ids.add(req.py_request_id) + try: + self.kv_cache_transceiver.request_and_receive_sync(req) + except Exception: + self._disagg_gen_kv_recv_started_ids.discard( + req.py_request_id) + raise else: - for req in new_gen_reqs: - self.kv_cache_transceiver.request_and_receive_async(req) + for req in recv_reqs: + self._disagg_gen_kv_recv_started_ids.add(req.py_request_id) + try: + self.kv_cache_transceiver.request_and_receive_async(req) + except Exception: + self._disagg_gen_kv_recv_started_ids.discard( + req.py_request_id) + raise if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None: - for req in new_gen_reqs: + for req in recv_reqs: if req.state == LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS: req.py_kv_transfer_start_time = time.time() @@ -3349,14 +3389,19 @@ 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: - 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") @@ -3549,26 +3594,71 @@ def _handle_errors(self, requests: Optional[List[LlmRequest]] = None): error_responses: Dict[int, LlmResponse] = {} error_msg = error_msg or "error" - failed_requests = requests if requests is not None else self.active_requests + failed_requests = list( + requests if requests is not None else self.active_requests) + 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: + 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, 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: self._terminate_request(request) - def _terminate_request(self, request: LlmRequest): + 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: + 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: + logger.warning( + f"Deferring termination for request {request.py_request_id} " + "because generation KV transfer is still in progress") + 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, @@ -3578,6 +3668,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) @@ -3585,6 +3676,17 @@ 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) + 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.""" return (request.state @@ -3727,13 +3829,19 @@ 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. 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]) + if (request.py_request_id + 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_gen_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: