diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 452b90462448..49fde42bf52b 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,13 @@ class BaseCacheTransceiver { public: virtual ~BaseCacheTransceiver() = default; - virtual void respondAndSendAsync(LlmRequest* llmRequest) = 0; + 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 +222,12 @@ class BaseCacheTransceiver [[nodiscard]] virtual bool checkGenTransferComplete() const = 0; - virtual bool cancelRequest(LlmRequest* llmRequest) = 0; + virtual bool cancelRequest(std::shared_ptr llmRequest) = 0; + + [[nodiscard]] virtual bool hasPoisonedTransferBuffer() const + { + return false; + } }; class CacheTransceiver : public BaseCacheTransceiver @@ -252,13 +258,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 +273,9 @@ class CacheTransceiver : public BaseCacheTransceiver [[nodiscard]] bool checkGenTransferComplete() const override; - virtual bool cancelRequest(LlmRequest* llmRequest) override; + virtual bool cancelRequest(std::shared_ptr llmRequest) override; + + [[nodiscard]] bool hasPoisonedTransferBuffer() const override; private: void initializeCommState(); @@ -276,10 +284,19 @@ class CacheTransceiver : public BaseCacheTransceiver std::unique_ptr mCacheSender; std::unique_ptr mCacheReceiver; - std::vector>> mSenderFutures; - std::vector>> mRequesterFutures; + std::vector, std::future>> mSenderFutures; + std::vector, std::future>> mRequesterFutures; mpi::MpiComm const* mMpiWorldComm{nullptr}; + // First-timeout dedup sets for the kvTransferTimeoutMs enforcement, + // by role: mTimedOutSenderIds tracks context-side (mSenderFutures) + // and mTimedOutRequesterIds tracks generation-side (mRequesterFutures). + // Used to emit at most one timeout warning and at most one cancel + // request per request per role. Erased at every future-list cleanup + // site so the sets stay bounded. + std::unordered_set mTimedOutSenderIds; + std::unordered_set mTimedOutRequesterIds; + std::shared_ptr mGroupComm; std::shared_ptr mGroupTensorParaComm, mGroupPipeParaComm, mGroupDataComm, mGroupTPInDPComm; diff --git a/cpp/include/tensorrt_llm/executor/transferAgent.h b/cpp/include/tensorrt_llm/executor/transferAgent.h index 532f0ae70c44..0e6bb7a22d47 100644 --- a/cpp/include/tensorrt_llm/executor/transferAgent.h +++ b/cpp/include/tensorrt_llm/executor/transferAgent.h @@ -356,6 +356,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..f77d18f12ae8 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,43 @@ 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; + }; + 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 +377,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 +390,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..cf29efa8fc84 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h @@ -46,6 +46,134 @@ 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 means this object is responsible for either releasing a + // concrete preallocated slot or poisoning dynamic transfer memory on an + // unsafe exit. It is therefore true even when index == std::nullopt. + , mHeld(true) + , mIsRecv(isRecv) + , mRequestIdForLog(requestIdForLog) + { + } + + // 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; + } + + /// @brief Whether this holder still owns release/poison responsibility. + /// For dynamic-buffer paths, mIndex can be std::nullopt while this + /// remains true so poison() can still fail closed on unsafe exits. + [[nodiscard]] bool held() const noexcept + { + return mHeld; + } + + /// @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 +185,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. @@ -110,6 +268,12 @@ class BaseTransBufferManager return mMaxNumTokens; } + [[nodiscard]] bool hasPoisonedBuffer() const noexcept + { + return mConcurrenceSendResource.mPoisoned.load(std::memory_order_relaxed) + || mConcurrenceRecvResource.mPoisoned.load(std::memory_order_relaxed); + } + protected: /// @brief Constructor - derived classes call this after computing buffer sizes. /// @param transferBufferSize Size of each transfer buffer in bytes. @@ -125,6 +289,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 +297,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 3daf2735db96..e3c7a8a0b89a 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 { @@ -522,7 +518,20 @@ 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. If sendAllBuffers + // exits through an exception while an AgentConnection (NIXL) is + // active, transport quiescence is unknown and the catch block below + // poisons the held slot instead of returning it to the pool; on the + // direct-UCX path (no AgentConnection) the holder's destructor falls + // back to release. The send-pool CV wait inside + // assignBufferIndexForSend observes the session's per-request cancel + // flag and throws on cancel (parity with the recv side). + auto const sendReqIdForLog = std::make_optional(static_cast(llmRequest.mRequestId)); + auto const* sendCancelFlag = &session.getDataContext().getTransferTerminate(); + auto cacheBufferId = mCacheTransBufferManager->assignBufferIndexForSend( + sendCancelFlag, kBufferAcquireSliceMs, sendReqIdForLog); + BufferIndexHolder sendHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/false, sendReqIdForLog); int peerDuplicateHeadFactor = targetInfo.mPeerDupHeadFactor; auto bufferTargetNum = targetNum / peerDuplicateHeadFactor; auto ppRank = selfIdx @@ -601,12 +610,38 @@ 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 && common::getEnvDisaggEnableInflightCancel()) + { + // AgentConnection::send can throw on cancel/error after the + // backend transfer has seen this send buffer. release() is + // not a quiescence proof, so do not return the slot to the + // pool; poison it and let Python restart the process. + // + // Gated on TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL: poison the + // buffer only when the operator has explicitly opted in to + // the mid-flight cancellation surface. When opt-in is off + // the cancel chain is dormant by design (Python skips + // cancel_request), so this catch block can still execute + // from non-cancel exceptions but should not poison the + // buffer pool / drive Layer 5's fail-closed in py_executor. + sendHolder.poison(); + } + throw; + } + // 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/cacheFormatter.h b/cpp/tensorrt_llm/batch_manager/cacheFormatter.h index 458cac8d4382..0f1f70337b9c 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.h +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.h @@ -232,6 +232,14 @@ class BaseCacheFormatter virtual ~BaseCacheFormatter() = default; }; +inline void checkZeroCopyDisabledForCancellableDisaggTransfers() +{ + TLLM_CHECK_WITH_INFO(!common::getEnvTryZCopyForKVCacheTransfer(), + "Zero-copy cache transfer sends/receives directly from request-owned blocks. It is disabled for cancellable " + "disaggregated transfers until KV-block leases can prove the peer has stopped accessing those blocks before " + "the request is freed."); +} + // Simple cache block copy. Because it does not involve data splitting or merging, it performs best when the // parallel topology is completely identical, making it the preferred method. class CacheFormatter final : public BaseCacheFormatter @@ -243,6 +251,7 @@ class CacheFormatter final : public BaseCacheFormatter { TLLM_CHECK(mCacheManager); TLLM_CHECK(mCacheTransBufferManager); + checkZeroCopyDisabledForCancellableDisaggTransfers(); } void format(tensorrt_llm::batch_manager::TransferSession& session) override; diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 42fa3d9eac7c..b59a712bb33c 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -60,6 +60,25 @@ namespace tensorrt_llm::batch_manager { +namespace +{ +/// @brief Rank prefix for TLLM_LOG_DEBUG/WARNING; resolves MPI vs torch process group at log time. +inline int currentRankForLog() +{ + return useMPI() ? mpi::MpiComm::world().getRank() : tensorrt_llm::pg_utils::get_world_pg()->getRank(); +} + +/// @brief Trailing-clause appended to kvTransferTimeoutMs warnings when the +/// in-flight cancellation surface is disabled. The detection itself runs +/// unconditionally so wedges are observable in logs even in flag-off mode; +/// this clause names the knob operators must flip to actually cancel. +/// See TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL in envUtils.h. +constexpr char const* kObserveOnlyTimeoutTail + = "TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=0; continuing to wait. " + "Set TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 to enable cancellation " + "and error surfacing."; +} // namespace + std::mutex CacheTransceiver::mDllMutex; std::unique_ptr CacheTransceiverFactory::createCacheTransceiver( @@ -379,7 +398,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); @@ -393,9 +412,9 @@ 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); + mSenderFutures.emplace_back(std::move(llmRequest), std::move(future)); } void CacheTransceiver::respondAndSendLayerWise( @@ -410,22 +429,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()); @@ -437,9 +456,27 @@ 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); + // 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( @@ -541,10 +578,20 @@ 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(); + // Apply the overall transfer deadline in both blockAll and polling + // modes so a stuck sender cannot pin KV blocks indefinitely. + // Healthy transfers complete well before the deadline (default + // 60s on the Python side, configurable via kv_transfer_timeout_ms). + kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); + // The per-iteration poll timeout is only relevant when the caller wants to + // return after checking at least `atLeastRequestNum` entries. + if (!blockAll) + { + senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); + } } auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupTPInDPComm : mGroupTensorParaComm; @@ -602,13 +649,102 @@ 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()) + { + // Dedup is shared with the inline wait_for re-check below + // so a single hung request produces at most one warning + // regardless of which site observes the expired deadline. + bool const firstTimeout = mTimedOutSenderIds.insert(request->mRequestId).second; + bool const inflightCancelEnabled = common::getEnvDisaggEnableInflightCancel(); + if (firstTimeout) + { + TLLM_LOG_WARNING( + "Context KV cache transfer for request %ld exceeded total timeout: " + "elapsed %ld ms > limit %d ms. %s", + request->mRequestId, elapsedMs, kvTransferTimeoutMs.value(), + inflightCancelEnabled ? "Marking as error." : kObserveOnlyTimeoutTail); + } + if (inflightCancelEnabled) + { + // Defense-in-depth sender-side cancel. Sender zombies empirically + // unwind on peer teardown (decode-pod restart), but in general + // CacheSender::cancelRequest clears mReadyResponses / + // mCancelledRequests bookkeeping so a subsequent re-enqueue + // or telemetry path doesn't see the stale request. + mCacheSender->cancelRequest(*request); + request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + requestsStatus.errorRequestIds.insert(request->mRequestId); + mTimedOutSenderIds.erase(request->mRequestId); + it = mSenderFutures.erase(it); + continue; + } + // Observe-only: cannot safely cancel the in-flight transfer + // (would risk a UAF on send buffers still touched by the + // network thread). Fall through to the readiness gate so a + // transfer that completes naturally is still detected. + } + } if (blockAll || (toCompleteIdSet.find(request->mRequestId) != toCompleteIdSet.end())) { try { - // Wait for up to a specified timeout - auto status = future.wait_for(std::chrono::milliseconds(senderFutureTimeoutMs.value_or(0))); - if (status == std::future_status::ready || !senderFutureTimeoutMs.has_value()) + // Bound the wait so a stuck sender cannot wedge the call: + // per-iteration timeout if set, else remaining overall + // deadline, else block (caller opted in). + int64_t effectiveSliceMs = 0; + if (senderFutureTimeoutMs.has_value()) + { + effectiveSliceMs = senderFutureTimeoutMs.value(); + } + else if (kvTransferTimeoutMs.has_value()) + { + auto const elapsedNow = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()) + .count(); + effectiveSliceMs = std::max(0, kvTransferTimeoutMs.value() - elapsedNow); + } + else + { + effectiveSliceMs = std::numeric_limits::max(); + } + // Cap the per-call wait slice so the polling thread always + // cycles back to drive MPI/UCX forward progress. Transports + // without a dedicated progress thread (notably mpi_kvcache + // over UCX-over-shm) rely on the app thread calling MPI + // functions to advance in-flight transfers. If the slice is + // unbounded (the fall-through above sets INT64_MAX when no + // timeout is configured), a single wait_for on a not-yet- + // ready future blocks the thread forever -- which means no + // MPI progress -- which means the future never becomes + // ready. The death spiral was empirically observed as + // test_asymmetric_executor[mpi_kvcache] wedging at its + // 300s gtest timeout despite the test's CacheTransceiverConfig + // having neither senderFutureTimeoutMs nor kvTransferTimeoutMs + // configured (nvbug-6104831). + // + // The deadline-check observability above is UNCHANGED: if + // kvTransferTimeoutMs is set, the deadline still fires within + // kMaxPollSliceMs of true expiry. Cancellation semantics + // (gated by TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL) unchanged. + static constexpr int64_t kMaxPollSliceMs = 50; + effectiveSliceMs = std::min(effectiveSliceMs, kMaxPollSliceMs); + auto status = future.wait_for(std::chrono::milliseconds(effectiveSliceMs)); + if (status == std::future_status::ready) { future.get(); requestsStatus.completedRequestIds.insert(request->mRequestId); @@ -616,13 +752,61 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( { request->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); } + // Clear any observe-only timeout breadcrumb so the dedup + // set does not leak across the request's lifetime. + mTimedOutSenderIds.erase(request->mRequestId); it = mSenderFutures.erase(it); } else if (status == std::future_status::timeout) { - TLLM_LOG_WARNING("Timed out waiting for context KV cache transfer after %d milliseconds.", - senderFutureTimeoutMs.value()); - ++it; + // Re-check the overall deadline inline; if reached, + // mark error here instead of blocking on a follow-up + // future.get(). + bool deadlineReached = false; + if (kvTransferTimeoutMs.has_value()) + { + auto const elapsedNow = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()) + .count(); + deadlineReached = elapsedNow >= kvTransferTimeoutMs.value(); + } + if (deadlineReached) + { + bool const firstTimeout = mTimedOutSenderIds.insert(request->mRequestId).second; + bool const inflightCancelEnabled = common::getEnvDisaggEnableInflightCancel(); + if (firstTimeout) + { + TLLM_LOG_WARNING( + "Context KV cache transfer for request %ld reached total timeout " + "while waiting for sender future. %s", + request->mRequestId, + inflightCancelEnabled ? "Marking as error." : kObserveOnlyTimeoutTail); + } + if (inflightCancelEnabled) + { + mCacheSender->cancelRequest(*request); + request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + requestsStatus.errorRequestIds.insert(request->mRequestId); + mTimedOutSenderIds.erase(request->mRequestId); + it = mSenderFutures.erase(it); + } + else + { + // Observe-only: continue waiting for natural completion. + ++it; + } + } + else + { + if (senderFutureTimeoutMs.has_value()) + { + TLLM_LOG_WARNING( + "Timed out waiting for context KV cache transfer for request %ld after %d " + "milliseconds.", + request->mRequestId, senderFutureTimeoutMs.value()); + } + ++it; + } } else { @@ -631,15 +815,21 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); requestsStatus.errorRequestIds.insert(request->mRequestId); + mTimedOutSenderIds.erase(request->mRequestId); it = mSenderFutures.erase(it); } } 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); + mTimedOutSenderIds.erase(request->mRequestId); it = mSenderFutures.erase(it); } } @@ -649,12 +839,38 @@ 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()) + { + // Apply the overall transfer deadline in both blockAll and polling + // modes (mirrors checkContextTransferStatus) so a stuck receiver + // cannot pin generation-side KV blocks. Healthy transfers complete + // well under kv_transfer_timeout_ms; raise it if the workload needs + // more headroom. + kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); + // The per-iteration poll timeout is only relevant when the caller wants to + // return after checking at least `atLeastRequestNum` entries. + if (!blockAll) + { + senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); + } + } + std::vector genTransferReadyRequestIds; for (auto&& [request, future] : mRequesterFutures) { @@ -697,16 +913,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; @@ -721,18 +929,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++; } @@ -742,70 +941,251 @@ 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; + bool const inflightCancelEnabled = common::getEnvDisaggEnableInflightCancel(); + 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. %s", + request->mRequestId, elapsedMs, kvTransferTimeoutMs.value(), + inflightCancelEnabled ? "Requesting cancellation." : kObserveOnlyTimeoutTail); } + if (inflightCancelEnabled) + { + if (firstTimeout) + { + // 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; + } + // Observe-only: memory safety forbids requesting a worker + // unwind here (the recv buffer pool may still be touched + // by the transport thread). Fall through to the readiness + // gate so a naturally completing transfer is still cleaned + // up normally. } - 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()); + // Bound the wait so a stuck receiver cannot wedge the + // call; see checkContextTransferStatus for the rationale. + int64_t effectiveSliceMs = 0; + if (senderFutureTimeoutMs.has_value()) + { + effectiveSliceMs = senderFutureTimeoutMs.value(); + } + else if (kvTransferTimeoutMs.has_value()) + { + auto const elapsedNow = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()) + .count(); + effectiveSliceMs = std::max(0, kvTransferTimeoutMs.value() - elapsedNow); + } + else + { + effectiveSliceMs = std::numeric_limits::max(); + } + // Cap the per-call wait slice so the polling thread always + // cycles back to drive MPI/UCX forward progress -- see the + // long comment at the analogous site in + // checkContextTransferStatus for the rationale and the + // empirical mpi_kvcache wedge that motivated this cap. + static constexpr int64_t kMaxPollSliceMs = 50; + effectiveSliceMs = std::min(effectiveSliceMs, kMaxPollSliceMs); + auto status = future.wait_for(std::chrono::milliseconds(effectiveSliceMs)); + if (status == std::future_status::ready) + { + 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) + { + // Re-check the overall deadline inline; if reached, + // mark error here instead of blocking on a follow-up + // future.get(). + bool deadlineReached = false; + if (kvTransferTimeoutMs.has_value()) + { + auto const elapsedNow = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()) + .count(); + deadlineReached = elapsedNow >= kvTransferTimeoutMs.value(); + } + if (deadlineReached) + { + bool const firstTimeout = mTimedOutRequesterIds.insert(request->mRequestId).second; + bool const inflightCancelEnabled = common::getEnvDisaggEnableInflightCancel(); + if (firstTimeout) + { + TLLM_LOG_WARNING( + "Generation KV cache transfer for request %ld reached total timeout " + "while waiting for receiver future. %s", + request->mRequestId, + inflightCancelEnabled ? "Requesting cancellation and marking as error." + : kObserveOnlyTimeoutTail); + } + if (inflightCancelEnabled) + { + if (firstTimeout) + { + mCacheReceiver->cancelRequest(*request); + } + request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + mTimedOutRequesterIds.erase(request->mRequestId); + it = mRequesterFutures.erase(it); + ++numErrored; + } + else + { + // Observe-only: continue waiting for natural completion. + ++it; + } + } + else + { + if (senderFutureTimeoutMs.has_value()) + { + TLLM_LOG_WARNING( + "Timed out waiting for generation KV cache transfer for request %ld " + "after %d milliseconds (per-iteration).", + request->mRequestId, senderFutureTimeoutMs.value()); + } + ++it; + } + } + else + { + 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 @@ -813,7 +1193,13 @@ bool CacheTransceiver::checkGenTransferComplete() const return mRequesterFutures.empty(); } -bool CacheTransceiver::cancelRequest(LlmRequest* llmRequest) +bool CacheTransceiver::hasPoisonedTransferBuffer() const +{ + return std::any_of(mCacheTransBufferManagerPtrs.begin(), mCacheTransBufferManagerPtrs.end(), + [](BaseTransBufferManager const* manager) { return manager != nullptr && manager->hasPoisonedBuffer(); }); +} + +bool CacheTransceiver::cancelRequest(std::shared_ptr llmRequest) { if (llmRequest->isContextOnlyRequest()) { diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 13e95dd86e48..f873742484ca 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,46 @@ 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; + } + } + if (!isCancelled && common::getEnvDisaggEnableInflightCancel()) + { + // Request is the mCurrentRequest (or not even in mReadyResponses) + // — the queue-drain branch can't abort it. Flip the per-request + // 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); + } } - else + else if (!isCancelled) { - TLLM_LOG_WARNING("Cannot cancel request %zu", llmRequest.mRequestId); + TLLM_LOG_WARNING( + "Cannot cancel request %zu (in-flight cancel disabled; " + "set TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 to opt in)", + llmRequest.mRequestId); } return isCancelled; } @@ -477,7 +559,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 +596,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 +671,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 +694,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 +768,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 +801,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 +835,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); @@ -746,6 +889,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 @@ -762,23 +912,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()) { @@ -789,9 +943,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; @@ -819,7 +989,31 @@ 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 + // every non-happy-path exit is covered: explicit not-ready releases via + // ~BufferIndexHolder, while cancel and exception poison the holders + // before the stack unwinds (transport quiescence is unknown on those + // paths). The preAcquiredCacheBufferIds vector must have one entry per + // cache transfer buffer manager (aligned with + // AgentConnectionManager::getCacheTransBufferManagers()). + TransferSession sendRequestInfo(LlmRequest const& llmRequest, std::atomic const& perRequestCancel, + std::vector> preAcquiredCacheBufferIds) { uint64_t requestId = llmRequest.getContextPhaseParams().value().getReqId(); auto const& contextState = llmRequest.getDataTransceiverState(); @@ -858,11 +1052,31 @@ 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. On the happy path receiveSync's + // formatter releases each slot via freeBufferIndexForRecv + // and requestSync detaches the holder. On cancel / exception + // requestSync explicitly poisons the holders before the + // stack unwinds; only the explicit not-ready path falls + // through to ~BufferIndexHolder for a normal release. + cacheBufferIds = std::move(preAcquiredCacheBufferIds); + TLLM_CHECK(!cacheBufferIds.empty()); + } + 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 @@ -891,6 +1105,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); @@ -941,15 +1166,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()); @@ -973,16 +1203,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) @@ -995,6 +1230,7 @@ class CacheReceiver::Impl } bool isCancelled = false; + std::optional queuedCancelledReqId; auto& asyncResource = mInstanceToAsyncResource.at(processInfo); { std::unique_lock lck(asyncResource->mMtxForQueue); @@ -1003,18 +1239,80 @@ 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 && common::getEnvDisaggEnableInflightCancel()) + { + // 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); } } + else if (!isCancelled) + { + TLLM_LOG_WARNING( + "Cannot cancel request %zu (in-flight cancel disabled; " + "set TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 to opt in)", + llmRequest.mRequestId); + } return isCancelled; } - bool receiveReadySignal(TransferSession& session) + enum class ReadySignalResult + { + kReady, + kNotReady, + kCancelled, + }; + + ReadySignalResult receiveReadySignalDetailed(TransferSession& session, std::atomic const& perRequestCancel) { bool isReadyFinal = true; bool isReady = false; @@ -1022,28 +1320,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; @@ -1056,74 +1400,162 @@ 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 (...) + { + // recvHolders is only populated on the AgentConnectionManager + // path (see the if-guarded block above), so on direct-UCX this + // is a no-op loop. Calling it unconditionally keeps the catch + // symmetric with the happy-path detach loop and removes a + // redundant branch on the failure path. + poisonRecvHolders(); + llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); + throw; + } 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 @@ -1160,11 +1592,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) @@ -1183,6 +1619,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); + } } } } @@ -1210,6 +1684,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) @@ -1228,7 +1707,7 @@ CacheSender::CacheSender( { } -std::future CacheSender::sendAsync(LlmRequest& llmRequest) const +std::future CacheSender::sendAsync(std::shared_ptr const& llmRequest) const { return mImpl->sendAsync(llmRequest); } @@ -1277,7 +1756,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 14d0ba8c6e5d..e2f38c718464 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. @@ -320,9 +323,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..81442d098089 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,21 @@ 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() after the send loop returns. If the send + // loop exits through an exception while an AgentConnection (NIXL) + // is active, transport quiescence is unknown and the catch block + // below poisons the held slot instead of returning it to the pool; + // on the direct-UCX path (no AgentConnection) the holder's + // destructor falls back to release. The send-pool CV wait inside + // assignBufferIndexForSend observes the session's per-request + // cancel flag and throws on cancel (parity with the recv side). + auto const sendReqIdForLog = std::make_optional(static_cast(llmRequest.mRequestId)); + auto const* sendCancelFlag = &session.getDataContext().getTransferTerminate(); + auto cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForSend( + 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); @@ -337,50 +352,69 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses session.appendMeasure(startTime, endTime, outputSplitCaches.at(cacheIdx)->getSizeInBytes()); }; - if (pickUpConnections.size() > 1) + try { - if (!common::getEnvEnableReceiveKVCacheParallel()) + if (pickUpConnections.size() > 1) { - TLLM_LOG_DEBUG("Disable parallel receiving of the KV cache."); - for (size_t i = 0; i < pickUpConnections.size(); i++) + if (!common::getEnvEnableReceiveKVCacheParallel()) { - sendBufferFun(deviceId, pickUpConnections[i]); + TLLM_LOG_DEBUG("Disable parallel receiving of the KV cache."); + for (size_t i = 0; i < pickUpConnections.size(); i++) + { + sendBufferFun(deviceId, pickUpConnections[i]); + } } - } - else - { - // concurrency num - auto concurrencyNum - = std::min(std::max(static_cast(1), bufferCoverTargetNum), pPDomainSize * cPDomainSize); + else + { + // concurrency num + auto concurrencyNum + = std::min(std::max(static_cast(1), bufferCoverTargetNum), pPDomainSize * cPDomainSize); - auto remainSendNum = pickUpConnections.size(); + auto remainSendNum = pickUpConnections.size(); - while (remainSendNum > 0) - { - auto sendConcurrencyNum = std::min(remainSendNum, concurrencyNum); - std::vector> futures; - futures.reserve(sendConcurrencyNum); - for (size_t i = 0; i < sendConcurrencyNum; i++) - { - size_t idx = i + (pickUpConnections.size() - remainSendNum); - size_t connIdx = pickUpConnections[idx]; - TLLM_CHECK(idx < pickUpConnections.size()); - TLLM_CHECK(connIdx < session.getConnections().size()); - futures.push_back(std::async(std::launch::async, sendBufferFun, deviceId, connIdx)); - } - for (auto& future : futures) + while (remainSendNum > 0) { - future.get(); + auto sendConcurrencyNum = std::min(remainSendNum, concurrencyNum); + std::vector> futures; + futures.reserve(sendConcurrencyNum); + for (size_t i = 0; i < sendConcurrencyNum; i++) + { + size_t idx = i + (pickUpConnections.size() - remainSendNum); + size_t connIdx = pickUpConnections[idx]; + TLLM_CHECK(idx < pickUpConnections.size()); + TLLM_CHECK(connIdx < session.getConnections().size()); + futures.push_back(std::async(std::launch::async, sendBufferFun, deviceId, connIdx)); + } + for (auto& future : futures) + { + future.get(); + } + remainSendNum -= sendConcurrencyNum; } - remainSendNum -= sendConcurrencyNum; } } + else + { + sendBufferFun(deviceId, pickUpConnections[0]); + } } - else + catch (...) { - sendBufferFun(deviceId, pickUpConnections[0]); + if (agentConnection != nullptr && common::getEnvDisaggEnableInflightCancel()) + { + // Symmetric with cacheFormatter.cpp catch: AgentConnection::send + // can throw on cancel/error after the backend transfer has seen + // this send buffer; poison the holder so Layer 5 fail-closed + // takes the pod out of service rather than reusing a + // possibly-still-pinned buffer. + // Gated on TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL: see + // cacheFormatter.cpp for the full rationale. + sendHolder.poison(); + } + throw; } - 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/mlaCacheFormatter.h b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.h index 1bc5daf87622..740f647c44ac 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.h +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.h @@ -32,6 +32,7 @@ class MLACacheFormatter final : public BaseCacheFormatter , mCacheTransBufferManagers{cacheTransBufferManagers} { TLLM_CHECK(mCacheManager); + checkZeroCopyDisabledForCancellableDisaggTransfers(); } void format(tensorrt_llm::batch_manager::TransferSession& session) override; diff --git a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp index 05742af17a28..93db3a804eae 100644 --- a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp @@ -40,6 +40,7 @@ RnnCacheFormatter::RnnCacheFormatter(rnn_state_manager::RnnStateManager* rnnStat { TLLM_CHECK(mRnnStateManager != nullptr); TLLM_CHECK(mRnnCacheTransBufferManager != nullptr); + kv_cache_manager::checkZeroCopyDisabledForCancellableDisaggTransfers(); } RnnCacheFormatter::RnnCacheFormatter(kv_cache_manager::BaseKVCacheManager* kvCacheManager, diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index 100e96f03119..05ed827a9511 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -923,7 +923,7 @@ void TrtGptModelInflightBatching::forwardSync() TLLM_CHECK_WITH_INFO(mCacheTransceiver, "Disaggregated serving is not enabled, please check the configuration of " "cacheTransceiverConfig."); - mCacheTransceiver->respondAndSendAsync(llmReq.get()); + mCacheTransceiver->respondAndSendAsync(llmReq); } mSeqSlotManager->freeSequenceSlot(llmReq->mRequestId); } @@ -1604,11 +1604,11 @@ void TrtGptModelInflightBatching::prepareDisaggGenInitRequests( mCacheTransceiver, "Disaggregated serving is not enabled, please check the configuration."); if (common::getEnvDisableKVCacheTransferOverlap()) { - mCacheTransceiver->requestAndReceiveSync(newGenReq.get()); + mCacheTransceiver->requestAndReceiveSync(newGenReq); } else { - mCacheTransceiver->requestAndReceiveAsync(newGenReq.get()); + mCacheTransceiver->requestAndReceiveAsync(newGenReq); } } if (!common::getEnvDisableKVCacheTransferOverlap()) diff --git a/cpp/tensorrt_llm/common/envUtils.cpp b/cpp/tensorrt_llm/common/envUtils.cpp index 1a1b75e2f10e..ec1d6e0ddcde 100644 --- a/cpp/tensorrt_llm/common/envUtils.cpp +++ b/cpp/tensorrt_llm/common/envUtils.cpp @@ -410,6 +410,12 @@ bool getEnvTryZCopyForKVCacheTransfer() return zcopyForSysmmetricKVCache; } +bool getEnvDisaggEnableInflightCancel() +{ + static bool const enabled = getBoolEnv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL"); + return enabled; +} + bool getEnvForceDeterministic() { static bool const forceDeterministic = getBoolEnv("FORCE_DETERMINISTIC"); diff --git a/cpp/tensorrt_llm/common/envUtils.h b/cpp/tensorrt_llm/common/envUtils.h index 832421ec60ef..915283362aaa 100644 --- a/cpp/tensorrt_llm/common/envUtils.h +++ b/cpp/tensorrt_llm/common/envUtils.h @@ -115,6 +115,15 @@ std::string const& getEnvKVCacheTimeOutputPath(); bool getEnvTryZCopyForKVCacheTransfer(); +// Opt-in for the disagg mid-flight cancellation + deferred-cleanup surface. +// Returns true when ``TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1``; defaults to +// false. When unset, timeouts are still detected and logged once per request, +// but no in-flight cancellation is issued, deferred-cleanup paths are +// inactive, and futures are left to complete naturally. The +// detection + cancellation + mark-as-error triple is atomic by design and +// must move together for memory safety. +bool getEnvDisaggEnableInflightCancel(); + // Force deterministic behavior for all kernels. bool getEnvForceDeterministic(); 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..ce43cf096094 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,31 @@ 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; + /// @brief Wait for the peer's ready signal. Returns the peer's ready + /// bit, collapsing cancel/no-signal into `false`. Callers that + /// need to distinguish "peer said not-ready" from "wait was + /// cancelled" must use `recvReadySignalWithStatus`. bool recvReadySignal(DataContext const& ctx) const; + /// @brief Tri-state variant of `recvReadySignal`. Returns the peer's + /// ready bit when a signal arrived, or `std::nullopt` if the + /// wait unwound without a notification (e.g. the per-request + /// cancel flag flipped before the peer sent ready). Callers + /// treat `std::nullopt` as "transport quiescence is unknown" + /// and must fail closed (poison receive buffers, etc.). + std::optional recvReadySignalWithStatus(DataContext const& ctx) const; void activateBuffer(uint8_t kind) const override; [[nodiscard]] std::optional getPreAssignedBufferId(uint8_t kind) const override; @@ -320,11 +337,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/mooncake_utils/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils/transferAgent.cpp index f16ece15526f..570e38073423 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils/transferAgent.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -166,6 +166,15 @@ TransferState MooncakeTransferStatus::wait(int64_t timeout_ms) const return true; } +[[nodiscard]] bool MooncakeTransferStatus::release() +{ + // Mooncake does not expose a per-request handle release/cancel primitive + // here. Report the release request as accepted so the generic cancel path + // can unwind; upper layers still decide whether transfer buffers must be + // quarantined when transport quiescence is unknown. + return true; +} + std::string const MooncakeBase64Helper::STANDARD_CHARS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" "abcdefghijklmnopqrstuvwxyz" diff --git a/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils/transferAgent.h b/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils/transferAgent.h index 42f870722e4f..92c99fca24c8 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils/transferAgent.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils/transferAgent.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -37,6 +37,8 @@ class MooncakeTransferStatus final : public TransferStatus TransferState wait(int64_t timeout_ms = -1) const override; + [[nodiscard]] bool release() override; + private: transfer_engine_t mEngine; uint64_t mBatchId; diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp index 30ef42560c6e..f9867e24d599 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp @@ -189,7 +189,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") @@ -230,7 +231,7 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m) "submit_transfer_requests", [](kvc::BaseTransferAgent& self, kvc::TransferRequest const& request) { return self.submitTransferRequests(request).release(); }, - nb::arg("request"), nb::rv_policy::take_ownership) + nb::arg("request"), nb::rv_policy::take_ownership, nb::keep_alive<0, 1>()) .def( "notify_sync_message", &kvc::BaseTransferAgent::notifySyncMessage, nb::arg("name"), nb::arg("sync_message")) .def("get_notified_sync_messages", &kvc::BaseTransferAgent::getNotifiedSyncMessages) @@ -242,7 +243,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") @@ -263,7 +265,8 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m) "submit_transfer_requests", [](kvc::NixlTransferAgent& self, kvc::TransferRequest const& request) { return self.submitTransferRequests(request).release(); }, - nb::arg("request"), nb::rv_policy::take_ownership, nb::call_guard()) + nb::arg("request"), nb::rv_policy::take_ownership, nb::call_guard(), + nb::keep_alive<0, 1>()) .def( "notify_sync_message", &kvc::NixlTransferAgent::notifySyncMessage, nb::arg("name"), nb::arg("sync_message")) .def("get_notified_sync_messages", &kvc::NixlTransferAgent::getNotifiedSyncMessages) 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 09dd63b0a299..6542a308bd7d 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(); @@ -485,10 +493,17 @@ NixlTransferStatus::NixlTransferStatus(nixlAgent* agent, nixlXferReqH* handle) TransferState NixlTransferStatus::wait(int64_t timeout_ms) const { auto startTime = std::chrono::steady_clock::now(); - while (true) { - auto status = mRawAgent->getXferStatus(mHandle); + nixl_status_t status; + { + std::lock_guard lock(mHandleMutex); + if (mHandle == nullptr) + { + return TransferState::kFAILURE; + } + status = mRawAgent->getXferStatus(mHandle); + } if (status == NIXL_SUCCESS) { return TransferState::kSUCCESS; @@ -520,9 +535,33 @@ TransferState NixlTransferStatus::wait(int64_t timeout_ms) const [[nodiscard]] bool NixlTransferStatus::isCompleted() const { + std::lock_guard lock(mHandleMutex); + if (mHandle == nullptr) + { + return false; + } return mRawAgent->getXferStatus(mHandle) == NIXL_SUCCESS; } +[[nodiscard]] bool NixlTransferStatus::release() +{ + std::lock_guard lock(mHandleMutex); + 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; +} + NixlTransferAgent::NixlTransferAgent(BaseAgentConfig const& config) : mName{config.mName} { 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 50cf672400ce..998d0f5a9651 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h @@ -20,6 +20,7 @@ #include "nixl.h" #include "tensorrt_llm/executor/transferAgent.h" #include +#include #include namespace tensorrt_llm::executor::kv_cache @@ -57,14 +58,22 @@ 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{}; + // Serializes all access to mHandle (and the NIXL handle it points to) + // across wait(), isCompleted(), and release(), because multiple Python + // threads holding the same BindingsNixlTransferStatus wrapper can call + // into the C++ methods concurrently. + mutable std::mutex mHandleMutex; }; class NixlTransferAgent final : public BaseTransferAgent diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index 8adddf1c86c0..e191432d2ff5 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); } @@ -110,7 +110,8 @@ void tb::CacheTransceiverBindings::initBindings(nb::module_& m) .def("check_gen_transfer_status", &BaseCacheTransceiver::checkGenTransferStatus, nb::call_guard()) .def("check_gen_transfer_complete", &BaseCacheTransceiver::checkGenTransferComplete) - .def("cancel_request", &BaseCacheTransceiver::cancelRequest); + .def("cancel_request", &BaseCacheTransceiver::cancelRequest) + .def("has_poisoned_transfer_buffer", &BaseCacheTransceiver::hasPoisonedTransferBuffer); nb::enum_(m, "AttentionType") .value("DEFAULT", executor::kv_cache::CacheState::AttentionType::kDEFAULT) diff --git a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp index 4514dfe26030..0df239c04d9e 100644 --- a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp @@ -348,11 +348,11 @@ class SymmetricalCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines- TLLM_CUDA_CHECK(cudaMemset(it->data(), llmRequest->getPromptLen(), it->getSizeInBytes())); } } - mFutures.emplace_back(mSender->sendAsync(*llmRequest)); + mFutures.emplace_back(mSender->sendAsync(llmRequest)); } else { - auto future = mRequester->receiveAsync(*llmRequest); + auto future = mRequester->receiveAsync(llmRequest); future.get(); TLLM_CUDA_CHECK(cudaDeviceSynchronize()); auto blockRange = BlockRange::fromAllBlockIds(*mManager, llmRequest->mRequestId); @@ -468,12 +468,12 @@ struct CPMetaData struct WrappedLlmRequest { - std::unique_ptr mLlmRequest; + std::shared_ptr mLlmRequest; std::optional mCPMetaData; using RequestIdType = LlmRequest::RequestIdType; - WrappedLlmRequest(std::unique_ptr llmRequest, std::optional cpMetaData) + WrappedLlmRequest(std::shared_ptr llmRequest, std::optional cpMetaData) : mLlmRequest(std::move(llmRequest)) , mCPMetaData(std::move(cpMetaData)) { @@ -887,7 +887,7 @@ class AsymmetricalCacheTest : public ::testing::TestWithParam(mRequestId++, std::move(request)); + auto llmRequestPtr = std::make_shared(mRequestId++, std::move(request)); return std::make_unique(std::move(llmRequestPtr), cpMetaData); } @@ -919,7 +919,7 @@ class AsymmetricalCacheTest : public ::testing::TestWithParamsetCacheState(cacheState); auto stats = texec::ContextPhaseParams({}, requestId, state.release(), std::nullopt); request.setContextPhaseParams(std::move(stats)); - auto llmRequestPtr = std::make_unique(requestId, std::move(request)); + auto llmRequestPtr = std::make_shared(requestId, std::move(request)); return std::make_unique(std::move(llmRequestPtr), cpMetaData); } @@ -973,7 +973,7 @@ class AsymmetricalCacheTest : public ::testing::TestWithParamsendAsync(*llmRequest); + auto future = mSender->sendAsync(llmRequest); return future; } @@ -984,7 +984,7 @@ class AsymmetricalCacheTest : public ::testing::TestWithParammLlmRequest; mManager->addSequenceBatch( {{{llmRequest->mRequestId, llmRequest->getNumTokens(beamIdx), beamWidth}}}, {std::ref(*llmRequest)}); - return mRequester->receiveAsync(*llmRequest); + return mRequester->receiveAsync(llmRequest); } void generationVerifyKVCache(std::shared_ptr const& request) diff --git a/ruff-legacy-baseline.json b/ruff-legacy-baseline.json index cedd33dfadbe..616a6291338c 100644 --- a/ruff-legacy-baseline.json +++ b/ruff-legacy-baseline.json @@ -1,8 +1,8 @@ { "_meta": { "generated_by": "scripts/legacy_utils.py lint-update-violations", - "total_violations": 5338, - "total_files": 492 + "total_violations": 5172, + "total_files": 480 }, ".github/scripts/label_community_user.py": { "D212": 1 @@ -452,9 +452,6 @@ "D212": 1, "E741": 1 }, - "scripts/check_test_list.py": { - "D212": 1 - }, "scripts/dco_check.py": { "D212": 1 }, @@ -518,9 +515,6 @@ "tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py": { "E741": 11 }, - "tensorrt_llm/_torch/custom_ops/torch_custom_ops.py": { - "F821": 2 - }, "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py": { "E741": 3 }, @@ -528,10 +522,6 @@ "E712": 1, "E731": 1 }, - "tensorrt_llm/_torch/model_config.py": { - "F811": 1, - "F821": 4 - }, "tensorrt_llm/_torch/models/checkpoints/auto_mapper.py": { "F821": 1 }, @@ -562,12 +552,6 @@ "tensorrt_llm/_torch/models/modeling_nemotron.py": { "E731": 1 }, - "tensorrt_llm/_torch/models/modeling_qwen2vl.py": { - "F811": 1 - }, - "tensorrt_llm/_torch/models/modeling_qwen3_next.py": { - "F821": 1 - }, "tensorrt_llm/_torch/models/modeling_siglip.py": { "F811": 1 }, @@ -575,9 +559,6 @@ "E731": 1, "F821": 5 }, - "tensorrt_llm/_torch/modules/attention.py": { - "E731": 2 - }, "tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py": { "F821": 2 }, @@ -611,15 +592,9 @@ "tensorrt_llm/_torch/modules/mamba/ssd_state_passing.py": { "E731": 1 }, - "tensorrt_llm/_torch/pyexecutor/_util.py": { - "F811": 1 - }, "tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py": { "F821": 1 }, - "tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py": { - "F821": 1 - }, "tensorrt_llm/_torch/pyexecutor/model_engine.py": { "F821": 1 }, @@ -670,9 +645,7 @@ "D411": 1 }, "tensorrt_llm/bench/build/dataclasses.py": { - "D205": 1, - "D208": 1, - "D210": 5 + "D210": 4 }, "tensorrt_llm/bench/build/tuning.py": { "D205": 2, @@ -784,7 +757,6 @@ "D212": 5, "D300": 3, "F402": 1, - "F811": 1, "F821": 1 }, "tensorrt_llm/executor/rpc/rpc_client.py": { @@ -834,18 +806,13 @@ "D411": 8, "D415": 9 }, - "tensorrt_llm/inputs/evs.py": { - "D205": 1, - "D212": 2 - }, "tensorrt_llm/inputs/multimodal.py": { - "D202": 1, "D415": 1 }, "tensorrt_llm/inputs/registry.py": { "D200": 3, "D205": 6, - "D212": 17, + "D212": 16, "D301": 1, "D405": 1, "D411": 1, @@ -918,29 +885,22 @@ "tensorrt_llm/llmapi/llm.py": { "D200": 2, "D202": 1, - "D205": 4, - "D207": 2, - "D212": 2, - "D300": 4, - "D410": 2, - "D411": 2 + "D205": 4 }, "tensorrt_llm/llmapi/llm_args.py": { - "D200": 23, "D202": 1, - "D205": 11, + "D205": 10, "D208": 2, "D210": 7, - "D212": 37, + "D212": 14, "D300": 8, "D415": 2 }, "tensorrt_llm/llmapi/llm_utils.py": { "D200": 2, - "D205": 4, "D209": 1, "D210": 9, - "D212": 3, + "D212": 2, "D300": 12 }, "tensorrt_llm/llmapi/mgmn_leader_node.py": { @@ -979,9 +939,9 @@ "D202": 1, "D205": 5, "D209": 1, - "D210": 14, + "D210": 13, "D212": 5, - "D300": 8, + "D300": 7, "E722": 1 }, "tensorrt_llm/mapping.py": { @@ -1547,10 +1507,6 @@ "PLE0302": 1 }, "tensorrt_llm/scaffolding/task.py": { - "F811": 1, - "F821": 1 - }, - "tensorrt_llm/scaffolding/task_collection.py": { "F821": 1 }, "tensorrt_llm/scaffolding/worker.py": { @@ -1573,7 +1529,6 @@ }, "tensorrt_llm/serve/openai_server.py": { "D205": 1, - "D212": 2, "F821": 2 }, "tensorrt_llm/serve/responses_utils.py": { @@ -1674,10 +1629,6 @@ "D403": 9, "D415": 11 }, - "tests/integration/defs/accuracy/accuracy_core.py": { - "D205": 1, - "D212": 1 - }, "tests/integration/defs/accuracy/test_disaggregated_serving.py": { "D212": 1, "F601": 1 @@ -1688,15 +1639,12 @@ }, "tests/integration/defs/accuracy/test_llm_api_autodeploy.py": { "D205": 1, - "D209": 1, - "D415": 1, - "F811": 2 + "D415": 1 }, "tests/integration/defs/accuracy/test_llm_api_pytorch.py": { "D212": 1, "D300": 3, - "D415": 3, - "E402": 7 + "D415": 3 }, "tests/integration/defs/common.py": { "D200": 1, @@ -1730,7 +1678,6 @@ "E741": 2 }, "tests/integration/defs/disaggregated/test_disaggregated.py": { - "D202": 2, "D205": 3, "D209": 1, "D212": 5, @@ -1874,7 +1821,7 @@ "D212": 3 }, "tests/integration/defs/perf/open_search_db_utils.py": { - "D200": 2, + "D200": 1, "D212": 4, "E402": 1 }, @@ -1941,13 +1888,7 @@ }, "tests/integration/defs/test_e2e.py": { "D200": 5, - "D202": 3, - "D205": 3, - "D210": 1, - "D212": 8, - "D300": 9, - "D403": 4, - "D415": 10, + "D415": 8, "F811": 2 }, "tests/integration/defs/test_list_parser.py": { @@ -2069,10 +2010,7 @@ "E731": 1 }, "tests/unittest/_torch/modeling/test_modeling_exaone4.py": { - "E402": 12 - }, - "tests/unittest/_torch/thop/parallel/test_mamba2_chunk_ss_update.py": { - "E731": 4 + "E402": 11 }, "tests/unittest/api_stability/api_stability_core.py": { "D202": 1, @@ -2085,7 +2023,7 @@ }, "tests/unittest/bindings/test_executor_bindings.py": { "D202": 1, - "E712": 56, + "E712": 54, "F403": 2, "F405": 3, "F811": 1 @@ -2161,9 +2099,7 @@ }, "tests/unittest/llmapi/test_llm_args.py": { "D200": 1, - "D202": 2, - "D212": 2, - "E712": 20, + "E712": 19, "F811": 1 }, "tests/unittest/llmapi/test_llm_multi_gpu.py": { @@ -2174,8 +2110,7 @@ "D202": 1, "D210": 2, "D212": 1, - "F811": 1, - "F821": 3 + "F811": 1 }, "tests/unittest/llmapi/test_llm_quant.py": { "D205": 1, @@ -2201,8 +2136,7 @@ "D415": 1 }, "tests/unittest/others/test_kv_cache_transceiver.py": { - "D205": 1, - "D212": 2 + "D205": 1 }, "tests/unittest/others/test_layer.py": { "D205": 1, @@ -2434,9 +2368,6 @@ "D205": 1 }, "triton_backend/all_models/llmapi/tensorrt_llm/1/helpers.py": { - "D205": 3, - "D212": 3, - "D415": 3, "E722": 1 }, "triton_backend/all_models/llmapi/tensorrt_llm/1/model.py": { @@ -2480,7 +2411,7 @@ }, "triton_backend/all_models/tests/test_python_backend.py": { "E402": 2, - "E712": 40, + "E712": 38, "F403": 1, "F405": 72 }, 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 b7d1c6af0bf0..0d75024ec9c7 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/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index 6367829d55ce..750eafcc1220 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -19,6 +19,44 @@ CacheTransBufferManagerCpp = tensorrt_llm.bindings.internal.batch_manager.CacheTransBufferManager BackendTypeCpp = tensorrt_llm.bindings.executor.CacheTransceiverBackendType +# Opt-in for the PR #13713 disagg mid-flight cancellation surface +# (Layer 1 cancel propagation + Layer 2b sendHolder.poison() + Layer 5 +# fail-closed). Setting ``TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1`` arms the +# defensive surface: Python propagates cancel for timed-out transfers, the +# C++ poll loop unwinds workers on cancel, and the poisoned-buffer signal +# drives a graceful PyExecutor shutdown when transport quiescence cannot be +# proved. Defaults to "0" (cancel chain dormant; behavior matches the +# pre-PR-13713 baseline). The orthogonal fixes (BufferIndexHolder RAII, +# shared_ptr async lifetime, recv-side idempotency) remain +# always-on because they close baseline races that exist regardless of +# mid-flight cancellation. See the NVBug 6104831 investigation report +# section 10 for the empirical case for enabling this surface. +_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV = "TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL" +_disagg_inflight_cancel_enabled_cache: Optional[bool] = None + + +def is_disagg_inflight_cancel_enabled() -> bool: + """Return True iff the disagg mid-flight cancellation surface is opted + into via ``TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1``. Read once and + cached; logs a one-shot WARNING on the first call when the knob is set + so the operating mode is visible in logs. Defaults to False (cancel + chain dormant; behavior matches the pre-PR-13713 baseline). + """ + global _disagg_inflight_cancel_enabled_cache + if _disagg_inflight_cancel_enabled_cache is None: + _disagg_inflight_cancel_enabled_cache = (getenv( + _DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "0") == "1") + if _disagg_inflight_cancel_enabled_cache: + logger.warning( + f"{_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV}=1: disagg KV transfer " + "mid-flight cancellation and fail-closed memory-safety policy " + "are ENABLED. Timed-out transfers will be cancelled at the " + "C++ layer; on cancel-with-unknown-quiescence the PyExecutor " + "will be shut down via the fail-closed policy and the " + "orchestrator should restart the pod. See the NVBug 6104831 " + "investigation report section 10 for the rationale.") + return _disagg_inflight_cancel_enabled_cache + def mapping_to_world_config(mapping: Mapping) -> WorldConfig: @@ -123,6 +161,9 @@ def check_gen_transfer_complete(self): def cancel_request(self, req: LlmRequest): raise NotImplementedError + def has_poisoned_transfer_buffer(self) -> bool: + return False + @abstractmethod def prepare_context_requests(self, requests: List[LlmRequest]): """ @@ -225,8 +266,25 @@ def check_gen_transfer_complete(self): return self.impl.check_gen_transfer_complete() def cancel_request(self, req: LlmRequest): + # When TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL is not set (the default), + # the C++ cancel surface stays dormant — Python returns False so + # callers skip the "cancelled, waiting for C++ to surface final + # cleanup" dedup-set entry and the natural-completion path takes + # over. This matches the pre-PR-13713 baseline behavior; see + # is_disagg_inflight_cancel_enabled() for the opt-in details. + if not is_disagg_inflight_cancel_enabled(): + return False return self.impl.cancel_request(req) + def has_poisoned_transfer_buffer(self) -> bool: + # Layer 5 fail-closed is gated by the same opt-in: when the cancel + # surface is not enabled no poison signal is ever set by Layer 2b + # (the C++ catch-block call is also gated), so reporting False + # unconditionally here is consistent + defensive. + if not is_disagg_inflight_cancel_enabled(): + return False + return self.impl.has_poisoned_transfer_buffer() + def prepare_context_requests(self, requests: List[LlmRequest]): # not implemented, an empty placeholder to allow being invoked unconditionally ... diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 91484ee96a20..ae695c811ce9 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -58,7 +58,8 @@ from .handle_additional_outputs import HandleAdditionalOutputs from .handle_logits import HandleLogits from .hang_detector import HangDetector -from .kv_cache_transceiver import KvCacheTransceiver +from .kv_cache_transceiver import (KvCacheTransceiver, + is_disagg_inflight_cancel_enabled) from .llm_request import (ExecutorRequest, LlmRequest, LlmRequestState, LlmResponse, get_draft_token_length) from .mamba_cache_manager import (BaseMambaCacheManager, @@ -89,6 +90,18 @@ # Default: "0" (only rank 0 prints, matching existing behavior). PROFILE_LOG_RANKS_ENV_VAR_NAME = "TLLM_PROFILE_LOG_RANKS" +# Defense-in-depth Python-side fallback deadline applied to every disagg +# KV transfer when ``CacheTransceiverConfig.kv_transfer_timeout_ms`` is +# unset. The C++ deadline check is gated on ``kvTransferTimeoutMs.has_value()`` +# (cacheTransceiver.cpp), so a deployment that explicitly passes ``None`` +# (or constructs the C++ config bypassing the Python defaults) would +# otherwise leave ``_handle_errors``-deferred requests waiting forever +# for a ``kDISAGG_TRANS_ERROR`` that never arrives -- the original +# wedge in disguise. Ten minutes is two orders of magnitude beyond +# any healthy disagg KV transfer; we surface the error rather than +# wait silently. +_DISAGG_KV_TRANSFER_FALLBACK_DEADLINE_MS = 10 * 60 * 1000 + class PPCommTag(IntEnum): """ @@ -167,10 +180,19 @@ class AsyncTransferManager: """ class RequestTransferMetadata: + """STOP-GAP metadata for rc13 partial-reuse cleanup. + + `termination_requested` and `resources_freed` bridge the current split + where AsyncTransferManager owns transfer pinning, while PyExecutor owns + final resource cleanup. Remove them with the follow-up KV reuse lease / + cleanup-session work that replaces should_store_blocks + pin=True. + """ def __init__(self, block_id: Optional[int]): self.block_id = block_id self.counter = 0 + self.termination_requested = False + self.resources_freed = False def start_transfer(self): self.counter += 1 @@ -183,6 +205,18 @@ def end_transfer(self) -> bool: self.counter -= 1 return self.counter == 0 + @dataclasses.dataclass + class EndTransferResult: + """STOP-GAP result for deferred partial-reuse termination. + + `needs_termination` carries a deferred cleanup obligation from the + transfer owner back to PyExecutor. Remove it with the follow-up KV + reuse lease / cleanup-session work. + """ + + completed: bool + needs_termination: bool = False + def __init__(self, resource_manager: "ResourceManager", should_store_blocks: bool = True): @@ -202,6 +236,43 @@ def __init__(self, def requests_in_transfer(self) -> Dict[int, LlmRequest]: return self._requests_in_transfer + def mark_termination_requested(self, request: LlmRequest): + """Record that termination was attempted while transfer was in flight. + + The expected fast path is: ``_terminate_request`` -> + ``_can_terminate_request_now`` finds the request still in + ``is_disagg_context_transmission_state`` -> calls this -> returns False + -> ``_end_transfer_and_maybe_terminate`` later observes + ``needs_termination`` and retries termination once C++ reports the + transfer terminal. + + The "no metadata" branch only fires when ``end_transfer`` has already + run and popped the entry; that path also transitions state to + ``DISAGG_CONTEXT_COMPLETE`` (or ``DISAGG_TRANS_ERROR``), so by the + time ``_can_terminate_request_now`` is re-entered for the same + request its first branch is False and we never reach this method. + Any reachable miss means the state/metadata invariant has drifted + (e.g. a future refactor of ``end_transfer`` ordering) -- a debug log + is enough today, but anyone seeing this in field logs should treat + it as a regression of the order contract above and fix the caller, + not silence the message. + """ + metadata = self._request_transfer_metadata.get(request.py_request_id) + if metadata is not None: + metadata.termination_requested = True + else: + logger.debug( + f"Termination requested for request {request.py_request_id}, " + "but it is not tracked by AsyncTransferManager (likely a " + "harmless ordering window where end_transfer ran first; if " + "this fires repeatedly the state-vs-metadata invariant has " + "drifted)") + + def mark_resources_freed(self, request: LlmRequest): + metadata = self._request_transfer_metadata.get(request.py_request_id) + if metadata is not None: + metadata.resources_freed = True + def start_transfer(self, request: LlmRequest): """ Called when a Cache transceiver or connector transfer is started. @@ -235,14 +306,15 @@ def start_transfer(self, request: LlmRequest): self._request_transfer_metadata[req_id].start_transfer() - def end_transfer(self, request: LlmRequest) -> bool: + def end_transfer(self, request: LlmRequest) -> EndTransferResult: """ Called after a send of KV cache is complete. 1. Decrements counter for request. 2. If there are no more inflight transfers for this request, unpin the blocks and mark the request as complete. Returns: - bool: True if the request should be terminated after call to end_transfer + EndTransferResult: whether the transfer is complete and whether a + deferred termination must be retried now. """ try: transfer_metadata = self._request_transfer_metadata[ @@ -251,9 +323,11 @@ def end_transfer(self, request: LlmRequest) -> bool: logger.warning( f"Request {request.py_request_id} not found in transfer manager" ) - return False + return self.EndTransferResult(completed=False) if transfer_metadata.end_transfer(): + needs_termination = (transfer_metadata.termination_requested + and not transfer_metadata.resources_freed) self._requests_in_transfer.pop(request.py_request_id) self._request_transfer_metadata.pop(request.py_request_id) @@ -265,9 +339,10 @@ def end_transfer(self, request: LlmRequest) -> bool: if request.state != LlmRequestState.DISAGG_TRANS_ERROR: request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE - return True + return self.EndTransferResult(completed=True, + needs_termination=needs_termination) - return False + return self.EndTransferResult(completed=False) def has_any_inflight_requests(self) -> bool: return len(self._requests_in_transfer) > 0 @@ -374,6 +449,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) @@ -595,8 +679,11 @@ def on_detected(): # Initialize disagg PP termination handler if needed self._disagg_pp_termination_handler = None if self.dist.pp_size > 1 and self.enable_kv_cache_reuse and self.kv_cache_transceiver: + terminate_fn = (self._do_terminate_request_if_safe + if is_disagg_inflight_cancel_enabled() else + self._do_terminate_request) self._disagg_pp_termination_handler = DisaggPPTerminationHandler( - self.dist, self._do_terminate_request) + self.dist, terminate_fn) if self.dist.pp_size > 1: self.event_loop = self._executor_loop_pp @@ -730,16 +817,19 @@ def _end_transfer_and_maybe_terminate(self, request: LlmRequest): # participate. self._pending_transfer_responses.append( (request.py_request_id, response)) - if self.async_transfer_manager.end_transfer(request): + transfer_result = self.async_transfer_manager.end_transfer(request) + if transfer_result.completed: self.active_requests.remove(request) self._terminate_request(request) return - if self.async_transfer_manager.end_transfer(request): - # When should_store_blocks is True, _handle_responses already - # terminated this request via the early-termination path - # (enable_partial_reuse_for_disagg branch). Skip the redundant - # termination to avoid double free_resources calls. - if not self.async_transfer_manager.should_store_blocks: + transfer_result = self.async_transfer_manager.end_transfer(request) + if transfer_result.completed: + # When should_store_blocks is True, _handle_responses normally owns + # the early termination path. However, termination can be deferred + # while context transfer is still in progress; retry it now that + # C++ reports the transfer terminal. + if (not self.async_transfer_manager.should_store_blocks + or transfer_result.needs_termination): self._terminate_request(request) def _flush_pending_transfer_responses(self): @@ -1755,17 +1845,34 @@ def _executor_loop_pp(self): and req.py_disaggregated_params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST for req in self.active_requests) - if num_fitting_reqs == 0 and not fitting_disagg_gen_init_requests: - if not all_gen_first: - logger.warning( - "num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache" - ) - self._check_disagg_ctx_cache_transfer_status(1) - elif self.async_transfer_manager.has_any_inflight_requests( - ): - # Non-blocking cleanup of completed/timed-out - # transfers to free KV blocks (see _executor_loop). - self._check_disagg_ctx_cache_transfer_status(0) + # The downstream C++ checkContextTransferStatus runs a + # cross-rank gatherRequestIds collective on the ctx-side + # comm (mGroupTPInDPComm with attention-DP, otherwise + # mGroupTensorParaComm). Any rank-asymmetric gate here + # (e.g. the previous `if num_fitting_reqs == 0` predicate, + # which reads rank-local scheduler / async-transfer-manager + # state and diverges in helix-CP setups by design) causes + # ABBA deadlock against the next unconditional collective + # on the same ranks (sampler NCCL ops, _can_queue's + # tp_allgather, PP step boundary). All ctx-side ranks must + # enter this call together, so the gate is intentionally + # absent. at_least_num is rank-local and only affects the + # per-request loop behavior inside the C++ function + # (blockAll vs polling), which is safe to diverge -- + # mirrors bdfdf8be02's fix for + # _check_disagg_gen_transfer_status. See PR #13713 CI + # build #39569 helix test hang for the failure signature + # on TestQwen3_8B::test_auto_dtype_with_helix + # [pp1tp1cp4, pp1tp2cp2]. + at_least_num = 0 + if (num_fitting_reqs == 0 + and not fitting_disagg_gen_init_requests + and not all_gen_first): + logger.warning( + "num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache" + ) + at_least_num = 1 + self._check_disagg_ctx_cache_transfer_status(at_least_num) self.num_scheduled_requests = scheduled_batch.batch_size @@ -2300,19 +2407,20 @@ def _prepare_and_schedule_batch(self): req.py_disaggregated_params and req.py_disaggregated_params. schedule_style == DisaggScheduleStyle.GENERATION_FIRST for req in self.active_requests) - if num_fitting_reqs == 0 and not fitting_disagg_gen_init_requests: - if not all_gen_first: - logger.warning( - "num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache" - ) - self._check_disagg_ctx_cache_transfer_status(1) - elif self.async_transfer_manager.has_any_inflight_requests(): - # Non-blocking cleanup of completed/timed-out transfers - # to free KV blocks. We avoid the blocking check because - # gen-first requests may be waiting for peer info (which - # would block indefinitely), but completed transfers must - # still be reaped so that KV cache can be reclaimed. - self._check_disagg_ctx_cache_transfer_status(0) + # Same rank-symmetric pattern as _executor_loop_pp above; see + # rationale comment there. The C++ gatherRequestIds collective + # inside checkContextTransferStatus must be entered by every + # ctx-side rank together. at_least_num is rank-local and safe + # to diverge (it only controls per-request loop behavior inside + # the C++ function, not collective entry). + at_least_num = 0 + if (num_fitting_reqs == 0 and not fitting_disagg_gen_init_requests + and not all_gen_first): + logger.warning( + "num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache" + ) + at_least_num = 1 + self._check_disagg_ctx_cache_transfer_status(at_least_num) # In gen-only benchmark mode, all requests must fit in KV cache # simultaneously because generation requests do not release their @@ -3550,11 +3658,16 @@ def _schedule(self): @nvtx_range("_check_disagg_gen_transfer_status") def _check_disagg_gen_transfer_status(self): - - need_check = any([ - req.is_disagg_generation_transmission_in_progress - for req in self.active_requests - ]) + # The downstream C++ checkGenTransferStatus runs a cross-rank + # gatherRequestIds collective on the gen-side comm. Any + # rank-asymmetric gate here (e.g. the previous "if need_check" + # predicate, which reads rank-local LlmRequest state from + # self.active_requests) causes ABBA deadlock against the next + # unconditional collective on the same ranks (helix CI #39529 + # hang). All gen-side ranks must enter this call together, so + # the gate is intentionally absent. at_least_num is rank-local + # and only affects the per-request loop behavior inside the C++ + # function (blockAll vs polling), which is safe to diverge. non_gen_first_reqs = [ req for req in self.active_requests if req.py_disaggregated_params and req.py_disaggregated_params. @@ -3563,10 +3676,8 @@ def _check_disagg_gen_transfer_status(self): need_check_one = bool(non_gen_first_reqs) and all( req.is_disagg_generation_transmission_in_progress for req in non_gen_first_reqs) - - if need_check: - at_least_num = 1 if need_check_one else 0 - self._check_disagg_gen_cache_transfer_status(at_least_num) + at_least_num = 1 if need_check_one else 0 + self._check_disagg_gen_cache_transfer_status(at_least_num) return @@ -3574,20 +3685,48 @@ def _check_disagg_gen_transfer_status(self): def _check_kv_transfer_timeout(self): if not self.kv_cache_transceiver: return - timeout_ms = self.kv_cache_transceiver.kv_transfer_timeout_ms - if timeout_ms is None: - return + configured_timeout_ms = ( + self.kv_cache_transceiver.kv_transfer_timeout_ms) + # Always apply *some* deadline. When the user/admin disabled the + # configured timeout we still need a Python-side fallback because + # the deferral path in _handle_errors relies on a future + # DISAGG_TRANS_ERROR transition to surface the failure: a request + # that never times out on the C++ side never moves to error and + # would wedge here forever. See _DISAGG_KV_TRANSFER_FALLBACK_DEADLINE_MS. + if configured_timeout_ms is None: + timeout_ms = _DISAGG_KV_TRANSFER_FALLBACK_DEADLINE_MS + using_fallback = True + else: + timeout_ms = configured_timeout_ms + using_fallback = False + + cancel_enabled = is_disagg_inflight_cancel_enabled() + observe_only_tail = "" if cancel_enabled else ( + "TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=0; continuing to wait. " + "Set TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 to enable " + "cancellation and error surfacing.") def flag_if_kv_transfer_timed_out(req: LlmRequest, type: str) -> None: current_time = time.time() if req.py_kv_transfer_start_time is None: return elapsed_time = (current_time - req.py_kv_transfer_start_time) * 1000 - if elapsed_time > timeout_ms and not req.py_kv_transfer_timed_out: - logger.warning( - f"Terminating {type} request {req.py_request_id} due to KV cache transfer timeout" - ) - req.py_kv_transfer_timed_out = True + if elapsed_time <= timeout_ms or req.py_kv_transfer_timed_out: + return + # Build the "{req-id} {deadline-source}" body once, then + # prepend the verb chosen by the cancel-enable flag. + if using_fallback: + body = (f"{type} request {req.py_request_id} after " + f"Python fallback deadline ({timeout_ms} ms); " + "kv_transfer_timeout_ms is unset, so this is " + "the only escalation path. Configure " + "kv_transfer_timeout_ms to surface failures sooner.") + else: + body = (f"{type} request {req.py_request_id} due to " + "KV cache transfer timeout.") + verb = "Terminating" if cancel_enabled else "Observed timeout on" + logger.warning(f"{verb} {body} {observe_only_tail}") + req.py_kv_transfer_timed_out = True for req in self.async_transfer_manager.requests_in_transfer().values(): flag_if_kv_transfer_timed_out(req, "context") @@ -3704,9 +3843,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, @@ -3714,12 +3850,33 @@ 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) - - # Trigger KV cache exchange for new disagg_gen_init_requests - self._recv_disagg_gen_cache(fitting_disagg_gen_init_requests) + for req in resources_to_prepare: + prepared_ids.add(req.py_request_id) + + # Always trigger _recv_disagg_gen_cache, even when this rank has no + # newly fitting disagg-gen-init requests. The downstream C++ + # checkGenTransferStatus runs a cross-rank gatherRequestIds + # collective on the gen-side comm; skipping it on ranks with no + # local work creates the rank-asymmetric entry that produces an + # ABBA deadlock with the next unconditional collective on the + # same ranks. + # _recv_disagg_gen_cache handles the empty-input case as a no-op + # for the receive work while still entering the collective. + self._recv_disagg_gen_cache(fitting_disagg_gen_init_requests) @nvtx_range("_prepare_disagg_gen_transmission_complete") def _prepare_disagg_gen_transmission_complete(self, scheduled_batch): @@ -3802,15 +3959,47 @@ def _recv_disagg_gen_cache(self, new_gen_reqs): req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE 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) - else: - for req in new_gen_reqs: - self.kv_cache_transceiver.request_and_receive_async(req) + 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) + + # NOTE: do NOT early-exit when recv_reqs is empty. + # + # The _check_disagg_gen_cache_transfer_status call below dispatches + # into C++ checkGenTransferStatus, whose body begins with a + # cross-rank gatherRequestIds collective on the gen-side comm + # (mGroupDataComm with attention-DP, otherwise mGroupComm). A + # rank-asymmetric early-exit here causes ABBA deadlock against + # any subsequent unconditional collective on the same ranks -- + # e.g. _can_queue's tp_allgather under attention-DP, or PP + # step-boundary collectives. + if recv_reqs: + if os.getenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP") == "1": + for req in recv_reqs: + self._disagg_gen_kv_recv_started_ids.add(req.py_request_id) + try: + self.kv_cache_transceiver.request_and_receive_sync(req) + except Exception: + self._disagg_gen_kv_recv_started_ids.discard( + req.py_request_id) + raise + else: + for req in recv_reqs: + self._disagg_gen_kv_recv_started_ids.add(req.py_request_id) + try: + self.kv_cache_transceiver.request_and_receive_async(req) + except Exception: + self._disagg_gen_kv_recv_started_ids.discard( + req.py_request_id) + raise - if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None: - for req in new_gen_reqs: + # Record the transfer start time unconditionally so the Python + # fallback deadline in _check_kv_transfer_timeout can fire even + # when kv_transfer_timeout_ms is not configured. The check there + # picks the configured timeout or the fallback floor. + for req in recv_reqs: if req.state == LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS: req.py_kv_transfer_start_time = time.time() @@ -3822,6 +4011,8 @@ def _recv_disagg_gen_cache(self, new_gen_reqs): block_transfer = bool(non_gen_first_active) and all( req.is_disagg_generation_transmission_in_progress for req in non_gen_first_active) + # Always reach the C++ collective so every gen-side rank + # participates in lockstep with its peers. self._check_disagg_gen_cache_transfer_status(1 if block_transfer else 0) return @@ -3857,8 +4048,12 @@ def kv_connector_request_finished(req: LlmRequest): self.async_transfer_manager.start_transfer(req) self.kv_cache_transceiver.respond_and_send_async(req) - if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None: - req.py_kv_transfer_start_time = time.time() + # Always record the transfer start time. The Python + # fallback deadline in _check_kv_transfer_timeout + # falls back to a hard floor when + # kv_transfer_timeout_ms is unset, and that path + # needs a baseline timestamp on every request. + req.py_kv_transfer_start_time = time.time() if self.kv_connector_manager: if not self.disable_overlap_scheduler: @@ -3883,10 +4078,15 @@ def _check_cache_transfer_errors(self, error_msg_prefix: str): """Common helper to check for and handle cache transfer errors.""" error_requests = self._get_disagg_reqs_in_error_state() if error_requests: + poisoned_transfer_buffer = ( + self.kv_cache_transceiver is not None + and self.kv_cache_transceiver.has_poisoned_transfer_buffer()) self._handle_errors( - f"Error in kv cache transfer for {error_msg_prefix}", + f"Error in kv cache transfer for {error_msg_prefix}" + + ("; unrecoverable poisoned transfer buffer requires process restart" + if poisoned_transfer_buffer else ""), requests=error_requests, - charge_budget=False) + charge_budget=poisoned_transfer_buffer) @nvtx_range("_check_disagg_ctx_cache_transfer_status") def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): @@ -3916,14 +4116,18 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): for request_id in list(requests_in_transfer.keys()): request = requests_in_transfer[request_id] if request.py_kv_transfer_timed_out and request_id not in completed_req_ids: - 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") @@ -4217,29 +4421,115 @@ def _handle_errors(self, "on fatal error") failed_requests = (list(self.active_requests) - if requests is None else requests) + if requests is None else list(requests)) + if is_disagg_inflight_cancel_enabled(): + deferred_requests = [ + request for request in failed_requests + if self._is_unquiesced_disagg_transfer(request) + ] + else: + deferred_requests = [] + if deferred_requests: + # Deferred requests rely on a follow-up DISAGG_TRANS_ERROR + # transition to surface the failure. The transition is + # driven by the C++ deadline (kvTransferTimeoutMs) when + # configured and by _check_kv_transfer_timeout's Python + # fallback floor (_DISAGG_KV_TRANSFER_FALLBACK_DEADLINE_MS) + # otherwise; either way, this path cannot wedge. + logger.warning( + "Deferring error cleanup for disaggregated KV transfers still " + "in flight. C++ transfer status remains the source of truth " + "for final cleanup. request_ids=%s, error=%s", + [request.py_request_id for request in deferred_requests], + error_msg) + failed_requests = [ + request for request in failed_requests + if request not in deferred_requests + ] + if not failed_requests: + if self._fatal_error is not None: + self.executor_request_queue.enqueue_shutdown_request() + return + for request in failed_requests: req_id = request.py_request_id - 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: + # _terminate_request may return False for disagg requests that + # are still in an unquiesced KV-transfer state (see contract + # docstring); in that case the deferred retry runs from + # _end_transfer_and_maybe_terminate via the mark_termination_requested + # handoff installed inside _can_terminate_request_now. self._terminate_request(request) if self._fatal_error is not None: self.executor_request_queue.enqueue_shutdown_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 not is_disagg_inflight_cancel_enabled(): + return True + if request.is_disagg_context_transmission_state: + # STOP-GAP: partial reuse asks _handle_responses to terminate + # completed context requests before async KV transfer has always + # drained. Remember the failed attempt so + # _end_transfer_and_maybe_terminate retries exactly once after C++ + # reports transfer completion. Remove this handoff with the + # follow-up KV reuse lease / cleanup-session work. + self.async_transfer_manager.mark_termination_requested(request) + logger.debug( + f"Deferring termination for request {request.py_request_id} " + "because context KV transfer is still in progress") + return False + if request.is_disagg_generation_transmission_in_progress: + logger.debug( + 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: + """Terminate ``request`` immediately or defer until KV transfer drains. + + Returns ``True`` when termination ran (either inline or scheduled + via the PP termination handler). Returns ``False`` when the + request is still in an unquiesced disagg KV-transfer state; in + that case ``_can_terminate_request_now`` records the deferred + intent on ``async_transfer_manager`` via + ``mark_termination_requested`` and ``_end_transfer_and_maybe_terminate`` + retries exactly once after C++ surfaces transfer completion. + + Callers that pass non-disagg requests (e.g. paused requests, ADP + dummy requests) can safely ignore the return value -- those + requests never enter the unquiesced state and termination always + runs synchronously. Callers that may pass disagg requests + (``_handle_errors``, ``_handle_responses``) must rely on the + deferred-retry handoff above. + """ + if not self._can_terminate_request_now(request): + return False # Dummy requests don't participate in disagg KV cache transfers, # so they must bypass the PP termination handler to avoid stale # sequences in the KV cache manager (the handler delays removal, @@ -4249,13 +4539,24 @@ 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) + self.async_transfer_manager.mark_resources_freed(request) if self.gather_all_responses or self.dist.rank == 0: self.result_wait_queues.pop(request.py_request_id, None) + # 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 @@ -4443,14 +4744,25 @@ def _handle_responses(self, emit_first_iter: bool = True): 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], - charge_budget=False) + 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") + # Keep the request in active_requests so the deferred-cleanup + # contract still holds: _check_disagg_gen_cache_transfer_status + # will eventually transition it to DISAGG_TRANS_ERROR once C++ + # surfaces the cancellation, and _check_cache_transfer_errors + # then enqueues the final error response and runs termination. + new_active_requests.append(request) continue if request.is_generation_only_request() and not request.is_finished: @@ -4542,6 +4854,11 @@ def _handle_responses(self, emit_first_iter: bool = True): # Request should be terminated after enqueueing response to ensure we can enqueue response successfully. self._enqueue_responses(new_responses) for request in requests_to_terminate: + # The force_terminate_for_partial_reuse branch above can route + # disagg context requests still in DISAGG_CONTEXT_TRANS_IN_PROGRESS + # here; _terminate_request will return False for those and the + # retry runs from _end_transfer_and_maybe_terminate (see contract + # docstring on _terminate_request). self._terminate_request(request) return requests_to_terminate + requests_finished_by_transfer diff --git a/tests/unittest/_torch/executor/test_async_transfer_manager.py b/tests/unittest/_torch/executor/test_async_transfer_manager.py index 1f2f9013d903..8fb6891bb5af 100644 --- a/tests/unittest/_torch/executor/test_async_transfer_manager.py +++ b/tests/unittest/_torch/executor/test_async_transfer_manager.py @@ -12,10 +12,34 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +"""Unit tests for ``AsyncTransferManager`` and the rc13 block-reuse stop-gap. + +Validation strategy for the stop-gap (NVBug 6104831 sig #8): + +1. **Unit tests (this file).** Cover the deferred-termination *cleanup + contract* by driving the manager API directly: + ``mark_termination_requested`` -> ``end_transfer`` -> retry termination + via ``EndTransferResult.needs_termination``. These tests run in + milliseconds and pin the contract against silent regressions. + +2. **Stress harness (out-of-tree).** Reproduces the *cancellation + timing* that originally manifested at high concurrency on rc13 with + ``should_store_blocks=True`` and an in-flight cancel landing on the + partial-reuse path. The harness lives with the disaggregated serving + reproducers and is run pre-merge for every PR touching the disagg + cleanup path; see PR #13713 description for run instructions and + recorded baselines. + + We deliberately do not land a unit test for the timing window because + constructing it without an actual disagg pipeline would either (a) + re-mock the same code path covered above, adding no signal, or (b) + spin up a real two-process disagg session, which belongs in the + stress harness rather than ``unittest``. +""" from unittest.mock import MagicMock -from tensorrt_llm._torch.pyexecutor.py_executor import AsyncTransferManager +from tensorrt_llm._torch.pyexecutor.py_executor import AsyncTransferManager, PyExecutor from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType from tensorrt_llm.bindings import LlmRequestState @@ -81,7 +105,9 @@ def test_start_transfer_single_request(): # Check seq slot manager was called to free resources seq_slot_manager.free_resources.assert_called_once_with(request) - manager.end_transfer(request) + transfer_result = manager.end_transfer(request) + assert transfer_result.completed + assert not transfer_result.needs_termination kv_cache_manager.unpin_blocks_by_id.assert_called_once() @@ -105,10 +131,13 @@ def test_start_transfer_multiple_transfers_same_request(): kv_cache_manager.store_blocks_for_reuse.assert_called_once() for _ in range(2): - manager.end_transfer(request) + transfer_result = manager.end_transfer(request) + assert not transfer_result.completed kv_cache_manager.unpin_blocks_by_id.assert_not_called() - manager.end_transfer(request) + transfer_result = manager.end_transfer(request) + assert transfer_result.completed + assert not transfer_result.needs_termination kv_cache_manager.unpin_blocks_by_id.assert_called_once() @@ -135,7 +164,9 @@ def test_transfer_without_storing_blocks(): kv_cache_manager.store_blocks_for_reuse.assert_not_called() spec_resource_manager.free_resources.assert_called_once_with(request) - assert manager.end_transfer(request) + transfer_result = manager.end_transfer(request) + assert transfer_result.completed + assert not transfer_result.needs_termination kv_cache_manager.unpin_blocks_by_id.assert_not_called() @@ -153,7 +184,9 @@ def test_end_transfer_preserves_error_state(): # Set error state before end_transfer request.state = LlmRequestState.DISAGG_TRANS_ERROR - manager.end_transfer(request) + transfer_result = manager.end_transfer(request) + assert transfer_result.completed + assert not transfer_result.needs_termination # Error state should be preserved assert request.state == LlmRequestState.DISAGG_TRANS_ERROR @@ -180,3 +213,91 @@ def test_requests_in_transfer(): assert in_transfer[1] is request1 assert in_transfer[2] is request2 assert in_transfer[3] is request3 + + +def test_end_transfer_reports_deferred_termination_needed(): + """End transfer reports when deferred termination must be retried.""" + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [100] + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager, should_store_blocks=True) + request = create_mock_request(42) + manager.start_transfer(request) + manager.mark_termination_requested(request) + + transfer_result = manager.end_transfer(request) + + assert transfer_result.completed + assert transfer_result.needs_termination + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with([100]) + + +def test_deferred_termination_survives_until_transfer_completion(): + """Termination requested during transfer is retried after transfer ends. + + This mock-driven test covers the cleanup contract directly; the + cancellation timing that triggered the rc13 block-reuse wedge is covered by + the disaggregated serving stress harness. + """ + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [100] + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager, should_store_blocks=True) + request = create_mock_request(42) + manager.start_transfer(request) + manager.mark_termination_requested(request) + + executor = PyExecutor.__new__(PyExecutor) + executor.kv_cache_transceiver = None + executor.active_requests = [] + executor.async_transfer_manager = manager + executor._terminate_request = MagicMock(return_value=True) + + executor._end_transfer_and_maybe_terminate(request) + + executor._terminate_request.assert_called_once_with(request) + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with([100]) + + +def test_no_deferred_termination_preserves_partial_reuse_skip(): + """Partial-reuse path still skips termination if none was requested.""" + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [100] + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager, should_store_blocks=True) + request = create_mock_request(42) + manager.start_transfer(request) + + executor = PyExecutor.__new__(PyExecutor) + executor.kv_cache_transceiver = None + executor.active_requests = [] + executor.async_transfer_manager = manager + executor._terminate_request = MagicMock(return_value=True) + + executor._end_transfer_and_maybe_terminate(request) + + executor._terminate_request.assert_not_called() + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with([100]) + + +def test_resources_freed_clears_deferred_termination(): + """Successful early resource free prevents duplicate termination.""" + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = [100] + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager, should_store_blocks=True) + request = create_mock_request(42) + manager.start_transfer(request) + manager.mark_termination_requested(request) + manager.mark_resources_freed(request) + + executor = PyExecutor.__new__(PyExecutor) + executor.kv_cache_transceiver = None + executor.active_requests = [] + executor.async_transfer_manager = manager + executor._terminate_request = MagicMock(return_value=True) + + executor._end_transfer_and_maybe_terminate(request) + + executor._terminate_request.assert_not_called() + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with([100]) diff --git a/tests/unittest/llmapi/test_llm_pytorch.py b/tests/unittest/llmapi/test_llm_pytorch.py index 53b64a972898..7ac6dd311201 100644 --- a/tests/unittest/llmapi/test_llm_pytorch.py +++ b/tests/unittest/llmapi/test_llm_pytorch.py @@ -1505,8 +1505,11 @@ async def test_llm_rpc_get_stats_async(): @pytest.mark.threadleak(enabled=False) @pytest.mark.part0 @skip_ray -@pytest.mark.parametrize("transceiver_runtime", [None, "PYTHON"]) -def test_llm_context_only_timed_out(transceiver_runtime): +@pytest.mark.parametrize("transceiver_runtime", [None]) +def test_llm_context_only_timed_out(transceiver_runtime, monkeypatch): + # Verify the active cancel path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1). + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") + tp_size = 1 use_overlap = False enable_iter_req_stats = False @@ -1576,6 +1579,107 @@ def test_llm_context_only_timed_out(transceiver_runtime): assert final_used_num_blocks == 0 +@pytest.mark.threadleak(enabled=False) +@pytest.mark.part0 +@skip_ray +def test_llm_context_only_timed_out_observe_only(monkeypatch): + """Pin down the observe-only (flag-off) V1 KV transfer timeout contract. + + Mirrors ``test_llm_context_only_timed_out`` but inverts the final + assertion: with ``TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=0`` (the new + default), the timeout detection still fires (one WARN per request + per role, with the observe-only tail naming the knob to flip), but + the cancellation + mark-as-error + free-blocks chain stays dormant. + The timed-out context-only request therefore keeps its KV blocks + until the transfer naturally completes -- preserving the + pre-PR-13713 hang surface in a memory-safe way and providing + visibility-only mode for operators who haven't opted into the + active cancellation surface. + + Restricted to the C++ transceiver (``transceiver_runtime=None``, + UCX backend) because the flag only gates that surface; + ``KvCacheTransceiverV2`` has independent always-on cancel/timeout + primitives and is orthogonal to this contract. + """ + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "0") + + tp_size = 1 + use_overlap = False + enable_iter_req_stats = False + + llm_args_extra = {} + llm_args_extra.update( + dict(enable_iter_perf_stats=True, + enable_iter_req_stats=enable_iter_req_stats, + disable_overlap_scheduler=not use_overlap)) + + llm = LLM(model=llama_model_path, + kv_cache_config=global_kvcache_config, + tensor_parallel_size=tp_size, + cache_transceiver_config=CacheTransceiverConfig( + backend="UCX", kv_transfer_timeout_ms=1000), + **llm_args_extra) + + try: + max_tokens = 1 + sampling_params = SamplingParams(max_tokens=max_tokens) + disaggregated_params = DisaggregatedParams(request_type="context_only") + prompts0 = ["What is your name?"] + prompts1 = ["Nvidia is awesome because"] + + # Send context-only request and wait for stats to populate. + for output in llm.generate(prompts1, + sampling_params=sampling_params, + disaggregated_params=disaggregated_params): + print(output) + + max_retries = 10 + for _ in range(max_retries): + results = llm.get_stats(2) + if len(results) == 1: + break + time.sleep(1) + else: + pytest.fail( + f"Failed to get stats with len==1 after {max_retries} retries") + + assert len(results) == 1 + context_only_used_num_blocks = ( + results[0]["kvCacheStats"]["usedNumBlocks"]) + assert context_only_used_num_blocks > 0, ( + "Context-only request should have allocated KV blocks before " + "the timeout window starts.") + print(f"Context only used num blocks: {context_only_used_num_blocks}") + + # Sleep past kv_transfer_timeout_ms (1000 ms) so the timeout + # detection path runs at the C++ side (one WARN per request). + time.sleep(5) + + # Send a regular request to drive additional executor iterations + # so the timeout-detection path is observably exercised, mirroring + # ``test_llm_context_only_timed_out``. + for output in llm.generate(prompts0, sampling_params=sampling_params): + print(output) + + results = llm.get_stats(2) + assert len(results) == 1 + final_used_num_blocks = results[0]["kvCacheStats"]["usedNumBlocks"] + + # Atomicity contract: with the flag off, the cancellation + + # mark-as-error chain is dormant, so KV blocks of the timed-out + # context-only request remain held. A ``final_used_num_blocks == 0`` + # observation here would mean the gate leaked the cleanup path, + # which we explicitly forbid (it would re-arm the trio behind the + # operator's back). + assert final_used_num_blocks > 0, ( + "Expected blocks to remain held in observe-only mode " + "(TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=0); got " + f"final_used_num_blocks={final_used_num_blocks}. " + "Atomic gate leaked: cleanup ran despite the flag being off.") + finally: + llm.shutdown() + + # This test is to verify that when the KV cache is exhausted and scheduled batch size is 0, the context only request will be aborted due to timeout. @@ -1584,14 +1688,18 @@ def test_llm_context_only_timed_out(transceiver_runtime): @skip_ray @pytest.mark.parametrize("sender_future_timeout_ms", [100, 1000]) @pytest.mark.parametrize("backend", ["NIXL", "UCX"]) -@pytest.mark.parametrize("transceiver_runtime", [None, "PYTHON"]) +@pytest.mark.parametrize("transceiver_runtime", [None]) def test_llm_context_only_timed_out_kv_cache_exhausted(sender_future_timeout_ms, backend, - transceiver_runtime): + transceiver_runtime, + monkeypatch): # Python transceiver (V2) only supports NIXL/DEFAULT backends if transceiver_runtime == "PYTHON" and backend == "UCX": pytest.skip("Python transceiver (V2) does not support UCX backend") + # Verify the active cancel path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1). + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") + tp_size = 1 use_overlap = False enable_iter_req_stats = False @@ -1666,8 +1774,11 @@ def test_llm_context_only_timed_out_kv_cache_exhausted(sender_future_timeout_ms, @pytest.mark.part0 @skip_ray @pytest.mark.asyncio -@pytest.mark.parametrize("transceiver_runtime", [None, "PYTHON"]) -async def test_llm_disagg_gen_cancelled(transceiver_runtime): +@pytest.mark.parametrize("transceiver_runtime", [None]) +async def test_llm_disagg_gen_cancelled(transceiver_runtime, monkeypatch): + # Verify the active cancel path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1). + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") + tp_size = 1 use_overlap = False enable_iter_req_stats = False @@ -1805,7 +1916,7 @@ def test_priority_request_completes_before_low_priority(): prompt = "A B C D E F G H I J" completion_order = [] - async def run(): + async def run(llm): # Submit all three requests before the event loop can schedule any of # them – they all land in the waiting queue simultaneously. # Request 1 uses ignore_eos so it keeps the batch slot busy for the @@ -1832,7 +1943,7 @@ async def await_and_record(output, req_id: int): timeout=55, ) - asyncio.run(run()) + asyncio.run(run(llm)) del llm assert 2 in completion_order and 3 in completion_order, ( @@ -1848,8 +1959,12 @@ async def await_and_record(output, req_id: int): @pytest.mark.part0 @skip_ray @pytest.mark.asyncio -@pytest.mark.parametrize("transceiver_runtime", [None, "PYTHON"]) -async def test_llm_disagg_streaming_gen_cancelled(transceiver_runtime): +@pytest.mark.parametrize("transceiver_runtime", [None]) +async def test_llm_disagg_streaming_gen_cancelled(transceiver_runtime, + monkeypatch): + # Verify the active cancel path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1). + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") + tp_size = 1 use_overlap = False enable_iter_req_stats = False diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index cf9d544a9ce0..d458038a2423 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -1,3 +1,4 @@ +import threading import time import uuid @@ -51,6 +52,152 @@ def fill_kv_cache_buffer(kv_cache_manager): buffer.copy_(random_values) +# create_kv_cache_manager() configures a 256-token KV pool. Tests that need +# two concurrent generation requests must keep prompts small enough that both +# fit, hence this default. Single-request tests can override via prompt_len. +_DEFAULT_PROMPT_LEN = 64 + + +def _make_request(request_id, + llm_request_type, + context_phase_params=None, + prompt_len=_DEFAULT_PROMPT_LEN): + """Construct an ``LlmRequest`` with the shared disagg test boilerplate. + + ``prompt_len`` controls the synthetic input-token range and must fit in + the small KV pool from :func:`create_kv_cache_manager`. + """ + sampling_params = SamplingParams() + kwargs = dict( + request_id=request_id, + max_new_tokens=1, + input_tokens=list(range(prompt_len)), + sampling_config=tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config()), + is_streaming=False, + llm_request_type=llm_request_type, + ) + if context_phase_params is not None: + kwargs["context_phase_params"] = context_phase_params + return LlmRequest(**kwargs) + + +def _add_sequence(kv_cache_manager, request): + """Bind ``request`` into ``kv_cache_manager``. + + The production prepare_resources path always pairs add_sequence_batch + with refresh_blocks; without the second call the KV-block offsets + landed on the request are stale and the C++ disagg transceiver path + silently transfers into the wrong physical slots, leaving the + receiver-side pool buffers all-zero. Mirror that pairing here so + every test that uses this helper observes the same state the runtime + does. + """ + kv_cache_manager.impl.add_sequence_batch( + [(request.py_request_id, request.prompt_len, 1)], [request]) + kv_cache_manager.impl.refresh_blocks() + + +def _drive_template_handshake_to_completion(ctx_xcvr, + gen_xcvr, + ctx_request_id, + deadline_s=10): + """Poll both transceivers until the template transfer completes. + + ``ctx_request_id`` must land in the sender-side ``completed_ids`` AND the + gen-side ``mRequesterFutures`` must drain. + + The post-fix ``check_gen_transfer_status(1)`` is non-blocking (it skips + unready futures via ``wait_for(0)``), so a single call may leave the + template entry in ``mRequesterFutures`` if its data has not yet + arrived. Tests that stage a "blocked" request after a template + handshake must drain mRequesterFutures first so the blocked-side + assertions only see their own request. + """ + deadline = time.time() + deadline_s + ctx_done = False + while time.time() < deadline and ( + not ctx_done or not gen_xcvr.check_gen_transfer_complete()): + completed_ids, error_ids = ctx_xcvr.check_context_transfer_status(1) + assert ctx_request_id not in error_ids, ( + f"template ctx request {ctx_request_id} unexpectedly errored: " + f"{error_ids}") + if ctx_request_id in completed_ids: + ctx_done = True + gen_xcvr.check_gen_transfer_status(1) + time.sleep(0.05) + assert ctx_done, ( + f"template ctx request {ctx_request_id} did not complete within " + f"{deadline_s}s; sender side may be stuck before the test can " + f"stage its blocked request") + assert gen_xcvr.check_gen_transfer_complete(), ( + f"gen-side mRequesterFutures still non-empty within {deadline_s}s; " + f"the staged blocked-request assertions would otherwise observe " + f"a spurious lingering entry from the template handshake") + + +@pytest.fixture(autouse=True) +def _reset_disagg_inflight_cancel_flag_cache(monkeypatch): + """Reset the module-level ``is_disagg_inflight_cancel_enabled`` cache + before each test so per-test ``TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL`` + settings are not poisoned by a prior test that initialized the cache. + + ``BindKvCacheTransceiver.cancel_request`` reads the env var once and + caches it in ``_disagg_inflight_cancel_enabled_cache`` (see + ``tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py``). Without + this fixture an active-path test (flag=1) leaves the cache as ``True`` + and a subsequent observe-only test (flag unset) reads stale ``True``; + or vice versa. ``monkeypatch.setattr`` auto-restores the cache at + test teardown so the reset is local to each test. + """ + from tensorrt_llm._torch.pyexecutor import kv_cache_transceiver as _kct + monkeypatch.setattr(_kct, "_disagg_inflight_cancel_enabled_cache", None) + + +@pytest.fixture +def disagg_transceiver_pair(request): + """Build a single-process disagg ctx/gen transceiver pair. + + Used by the cancellation-flow regression tests. + + Yields ``(kv_cache_manager_ctx, kv_cache_manager_gen, + kv_cache_transceiver_ctx, kv_cache_transceiver_gen)``. + + Pass ``(attention_type, backend)`` as the parametrize argument; the + backend defaults to ``"DEFAULT"`` (UCX) for backend-agnostic fixes + and can be set to ``"NIXL"`` for tests that depend on NIXL-specific + paths (e.g. the recv buffer index manager). + """ + param = request.param + if isinstance(param, tuple): + attention_type = param[0] + backend = param[1] if len(param) > 1 else "DEFAULT" + else: + attention_type = param + backend = "DEFAULT" + + tensorrt_llm.logger.set_level("info") + mapping = Mapping(world_size=1, rank=0) + dist = Distributed.get(mapping) + kv_cache_manager_ctx = create_kv_cache_manager(mapping, DataType.HALF) + kv_cache_manager_gen = create_kv_cache_manager(mapping, DataType.HALF) + + cache_transceiver_config = CacheTransceiverConfig(backend=backend, + max_tokens_in_buffer=512) + + kv_cache_transceiver_ctx = create_kv_cache_transceiver( + mapping, dist, kv_cache_manager_ctx, attention_type, + cache_transceiver_config) + kv_cache_transceiver_gen = create_kv_cache_transceiver( + mapping, dist, kv_cache_manager_gen, attention_type, + cache_transceiver_config) + + fill_kv_cache_buffer(kv_cache_manager_ctx) + + yield (kv_cache_manager_ctx, kv_cache_manager_gen, kv_cache_transceiver_ctx, + kv_cache_transceiver_gen) + + @pytest.fixture(scope="function") def ctx_gen_kv_cache_dtype(request): if request.param == "ctx_fp8_gen_fp8": @@ -98,6 +245,117 @@ def test_kv_cache_transceiver_single_process(ctx_gen_kv_cache_dtype, mapping, dist, kv_cache_manager_gen, attention_type, cache_transceiver_config) + try: + fill_kv_cache_buffer(kv_cache_manager_ctx) + + # init ctx request + sampling_params = SamplingParams() + ctx_request = LlmRequest( + request_id=0, + max_new_tokens=1, + input_tokens=list(range(256)), + sampling_config=tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config()), + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) + + if transceiver_runtime == "PYTHON": + disaggregated_params = tensorrt_llm.DisaggregatedParams( + request_type="context_only", + disagg_request_id=uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF) + ctx_request.py_disaggregated_params = disaggregated_params + + kv_cache_manager_ctx.impl.add_sequence_batch( + [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], + [ctx_request]) + # add_sequence_batch must be paired with refresh_blocks so the C++ + # transceiver path observes the freshly-allocated KV-block offsets; + # without this the receiver-side pool buffers stay zero. See + # _add_sequence helper for the same pairing in cancellation tests. + kv_cache_manager_ctx.impl.refresh_blocks() + # send ctx request + kv_cache_transceiver_ctx.respond_and_send_async(ctx_request) + + # init gen request + gen_request = LlmRequest( + request_id=0, + max_new_tokens=1, + input_tokens=list(range(256)), + sampling_config=tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config()), + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + context_phase_params=ctx_request.context_phase_params) + + if transceiver_runtime == "PYTHON": + disaggregated_params = tensorrt_llm.DisaggregatedParams( + request_type="generation_only", + disagg_request_id=ctx_request.py_disaggregated_params. + disagg_request_id, + ctx_request_id=ctx_request.request_id, + ctx_dp_rank=ctx_request.context_phase_params.ctx_dp_rank, + ctx_info_endpoint=ctx_request.context_phase_params. + disagg_info_endpoint, + first_gen_tokens=ctx_request.context_phase_params. + first_gen_tokens, + draft_tokens=ctx_request.context_phase_params.draft_tokens) + + gen_request.py_disaggregated_params = disaggregated_params + + kv_cache_manager_gen.impl.add_sequence_batch( + [(gen_request.py_request_id, gen_request.prompt_len, 1)], + [gen_request]) + kv_cache_manager_gen.impl.refresh_blocks() + # send gen request + kv_cache_transceiver_gen.request_and_receive_async(gen_request) + + kv_cache_transceiver_ctx.check_context_transfer_status(1) + kv_cache_transceiver_gen.check_gen_transfer_status(1) + + assert torch.equal( + kv_cache_manager_gen.get_buffers(0), + kv_cache_manager_ctx.get_buffers(0)), "different kv-cache values" + finally: + # KvCacheTransceiverV2 (transceiver_runtime == "PYTHON") spawns + # daemon threads in its Messenger ("listener") and TransferWorker + # ("_process_task_queue") components. Those threads hold bound- + # method captures to self, creating a reference cycle that delays + # CPython's __del__ -- so pytest-threadleak fires at test teardown + # unless we call shutdown() explicitly. V1 (NIXL/UCX) cleans up + # via the C++ ~CacheTransceiver destructor and doesn't need this. + # shutdown() is idempotent via the self._shutdown guard, so safe + # to call even if the test body raised before send/receive. + if transceiver_runtime == "PYTHON": + kv_cache_transceiver_ctx.shutdown() + kv_cache_transceiver_gen.shutdown() + + +@pytest.mark.timeout(120) +@pytest.mark.parametrize("attention_type", + [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA], + ids=["mha", "mla"]) +def test_cancel_request_in_transmission(attention_type, monkeypatch): + # Verify the active cancel path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1). + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") + + # Init kv_cache manager and cache transceiver + mapping = Mapping(world_size=1, rank=0) + dist = Distributed.get(mapping) + ctx_kv_cache_dtype, gen_kv_cache_dtype = DataType.HALF, DataType.HALF + kv_cache_manager_ctx = create_kv_cache_manager(mapping, ctx_kv_cache_dtype) + kv_cache_manager_gen = create_kv_cache_manager(mapping, gen_kv_cache_dtype) + + cache_transceiver_config = CacheTransceiverConfig(backend="DEFAULT", + max_tokens_in_buffer=512) + + kv_cache_transceiver_ctx = create_kv_cache_transceiver( + mapping, dist, kv_cache_manager_ctx, attention_type, + cache_transceiver_config) + + kv_cache_transceiver_gen = create_kv_cache_transceiver( + mapping, dist, kv_cache_manager_gen, attention_type, + cache_transceiver_config) + fill_kv_cache_buffer(kv_cache_manager_ctx) # init ctx request @@ -111,17 +369,21 @@ def test_kv_cache_transceiver_single_process(ctx_gen_kv_cache_dtype, is_streaming=False, llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) - if transceiver_runtime == "PYTHON": - disaggregated_params = tensorrt_llm.DisaggregatedParams( - request_type="context_only", - disagg_request_id=uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF) - ctx_request.py_disaggregated_params = disaggregated_params - kv_cache_manager_ctx.impl.add_sequence_batch( [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], [ctx_request]) + # Pair with refresh_blocks; see _add_sequence helper docstring for + # why both calls are required for the C++ transceiver path. + kv_cache_manager_ctx.impl.refresh_blocks() # send ctx request kv_cache_transceiver_ctx.respond_and_send_async(ctx_request) + # wait for ctx request to be sent + time.sleep(2) + + # cancel ctx request + is_cancelled = kv_cache_transceiver_ctx.cancel_request(ctx_request) + assert is_cancelled + # init gen request gen_request = LlmRequest( request_id=0, @@ -133,44 +395,44 @@ def test_kv_cache_transceiver_single_process(ctx_gen_kv_cache_dtype, llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, context_phase_params=ctx_request.context_phase_params) - if transceiver_runtime == "PYTHON": - disaggregated_params = tensorrt_llm.DisaggregatedParams( - request_type="generation_only", - disagg_request_id=ctx_request.py_disaggregated_params. - disagg_request_id, - ctx_request_id=ctx_request.request_id, - ctx_dp_rank=ctx_request.context_phase_params.ctx_dp_rank, - ctx_info_endpoint=ctx_request.context_phase_params. - disagg_info_endpoint, - first_gen_tokens=ctx_request.context_phase_params.first_gen_tokens, - draft_tokens=ctx_request.context_phase_params.draft_tokens) - - gen_request.py_disaggregated_params = disaggregated_params - kv_cache_manager_gen.impl.add_sequence_batch( [(gen_request.py_request_id, gen_request.prompt_len, 1)], [gen_request]) + kv_cache_manager_gen.impl.refresh_blocks() # send gen request kv_cache_transceiver_gen.request_and_receive_async(gen_request) - kv_cache_transceiver_ctx.check_context_transfer_status(1) - kv_cache_transceiver_gen.check_gen_transfer_status(1) - - assert torch.equal( - kv_cache_manager_gen.get_buffers(0), - kv_cache_manager_ctx.get_buffers(0)), "different kv-cache values" + # Block the main thread due to the async operation + time.sleep(2) + assert gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR -@pytest.mark.timeout(120) +@pytest.mark.timeout(60) @pytest.mark.parametrize("attention_type", [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA], ids=["mha", "mla"]) -def test_cancel_request_in_transmission(attention_type): - # Init kv_cache manager and cache transceiver +def test_cancel_request_in_transmission_observe_only(attention_type, + monkeypatch): + """Pin down the flag-off contract for ``cancel_request``. + + Mirrors ``test_cancel_request_in_transmission`` but with + ``TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL`` unset (the default after + the PR #13713 atomicity refactor). The + ``BindKvCacheTransceiver.cancel_request`` shim must short-circuit + to return ``False`` without invoking the C++ cancel surface, so + the in-flight transfer is not interrupted and the ctx request + does not transition to ``DISAGG_TRANS_ERROR``. + + This regression-proofs the gate at the shim layer: any future + refactor that silently re-arms the cancel surface would surface + ``is_cancelled is True`` here and fail the assertion. + """ + # Verify the observe-only path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL unset). + monkeypatch.delenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", raising=False) + mapping = Mapping(world_size=1, rank=0) dist = Distributed.get(mapping) - ctx_kv_cache_dtype, gen_kv_cache_dtype = DataType.HALF, DataType.HALF + ctx_kv_cache_dtype = DataType.HALF kv_cache_manager_ctx = create_kv_cache_manager(mapping, ctx_kv_cache_dtype) - kv_cache_manager_gen = create_kv_cache_manager(mapping, gen_kv_cache_dtype) cache_transceiver_config = CacheTransceiverConfig(backend="DEFAULT", max_tokens_in_buffer=512) @@ -179,13 +441,8 @@ def test_cancel_request_in_transmission(attention_type): mapping, dist, kv_cache_manager_ctx, attention_type, cache_transceiver_config) - kv_cache_transceiver_gen = create_kv_cache_transceiver( - mapping, dist, kv_cache_manager_gen, attention_type, - cache_transceiver_config) - fill_kv_cache_buffer(kv_cache_manager_ctx) - # init ctx request sampling_params = SamplingParams() ctx_request = LlmRequest( request_id=0, @@ -198,44 +455,493 @@ def test_cancel_request_in_transmission(attention_type): kv_cache_manager_ctx.impl.add_sequence_batch( [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], [ctx_request]) - # send ctx request + kv_cache_manager_ctx.impl.refresh_blocks() kv_cache_transceiver_ctx.respond_and_send_async(ctx_request) - # wait for ctx request to be sent + # Mirror the active test's wait so the request is registered in + # ``mSenderFutures`` by the time we call cancel. time.sleep(2) - # cancel ctx request + # Observe-only contract: is_cancelled = kv_cache_transceiver_ctx.cancel_request(ctx_request) - assert is_cancelled + assert is_cancelled is False, ( + "Flag-off cancel_request must return False without invoking " + "the C++ cancel surface. A True result here indicates the " + "BindKvCacheTransceiver.cancel_request gate has leaked and " + "the active cancel path is firing despite the opt-in being " + "unset.") + assert ctx_request.state != LlmRequestState.DISAGG_TRANS_ERROR, ( + "Flag-off cancel_request must not transition the ctx request " + "to DISAGG_TRANS_ERROR; no cancellation actually occurred at " + "the C++ layer.") - # init gen request - gen_request = LlmRequest( - request_id=0, - max_new_tokens=1, - input_tokens=list(range(256)), - sampling_config=tensorrt_llm.bindings.SamplingConfig( - sampling_params._get_sampling_config()), - is_streaming=False, - llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, - context_phase_params=ctx_request.context_phase_params) - kv_cache_manager_gen.impl.add_sequence_batch( - [(gen_request.py_request_id, gen_request.prompt_len, 1)], [gen_request]) - # send gen request +@pytest.mark.timeout(120) +@pytest.mark.parametrize("disagg_transceiver_pair", + [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA], + ids=["mha", "mla"], + indirect=True) +def test_request_and_receive_async_state_ordering(disagg_transceiver_pair): + """Lock down the setState-after-receiveAsync ordering contract. + + ``cacheTransceiver.cpp::requestAndReceiveAsync`` was reordered so + ``llmRequest->setState(DISAGG_GENERATION_TRANS_IN_PROGRESS)`` runs + *after* ``mCacheReceiver->receiveAsync(llmRequest)`` returns (the + original code set the state first). The reorder is load-bearing: + if ``receiveAsync`` throws, the request is left out of both + ``mRequesterFutures`` and the IN_PROGRESS state set; if the spawned + worker ever read state at entry, the new ordering would race. + Today the worker does not read state at entry, but the contract is + fragile and easy to regress in a refactor. + + This test exercises a happy-path receive and asserts the observable + transition is ``DISAGG_GENERATION_INIT`` (pre-call) -> ``IN_PROGRESS`` + or already-``COMPLETE`` for a fast-completing transfer (post-call), + with the receiver-side bookkeeping reflecting the in-flight future. + """ + (kv_cache_manager_ctx, kv_cache_manager_gen, kv_cache_transceiver_ctx, + kv_cache_transceiver_gen) = disagg_transceiver_pair + + # Drive a real ctx handshake so the gen request has a live peer to + # receive from -- without one, request_and_receive_async would just + # park the future on an unresolvable peer. + ctx_request = _make_request(0, LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) + _add_sequence(kv_cache_manager_ctx, ctx_request) + kv_cache_transceiver_ctx.respond_and_send_async(ctx_request) + + gen_request = _make_request(0, + LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + ctx_request.context_phase_params) + _add_sequence(kv_cache_manager_gen, gen_request) + + # Pre-condition: gen-only requests built via _make_request land in + # DISAGG_GENERATION_INIT (set by the LlmRequest ctor for + # generation-only types). Locking this in protects the test from + # silent state-default drift. + assert gen_request.state == LlmRequestState.DISAGG_GENERATION_INIT, ( + f"gen_request.state should be DISAGG_GENERATION_INIT before " + f"request_and_receive_async, got {gen_request.state}") + assert kv_cache_transceiver_gen.check_gen_transfer_complete(), ( + "mRequesterFutures must be empty before request_and_receive_async") + kv_cache_transceiver_gen.request_and_receive_async(gen_request) - # Block the main thread due to the async operation + # Post-condition: setState and emplace happen together. The state + # must be IN_PROGRESS (or already TRANS_COMPLETE for a transfer that + # finished synchronously inside the call). It must NOT still be + # INIT -- if it is, the reordered setState was bypassed or didn't + # run. + assert gen_request.state in ( + LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS, + LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE, + ), (f"gen_request.state must be IN_PROGRESS or TRANS_COMPLETE after " + f"request_and_receive_async, got {gen_request.state}") + + # Drain the in-flight transfer so teardown is clean. + deadline = time.time() + 10 + while time.time() < deadline and ( + not kv_cache_transceiver_gen.check_gen_transfer_complete()): + kv_cache_transceiver_gen.check_gen_transfer_status(1) + kv_cache_transceiver_ctx.check_context_transfer_status(1) + time.sleep(0.05) + + kv_cache_manager_ctx.free_resources(ctx_request) + kv_cache_manager_gen.free_resources(gen_request) + + +@pytest.mark.timeout(120) +@pytest.mark.parametrize("disagg_transceiver_pair", + [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA], + ids=["mha", "mla"], + indirect=True) +def test_cancel_request_in_transmission_does_not_break_sender_future( + disagg_transceiver_pair, capfd, monkeypatch): + """Reproduce the sender-side broken-promise on cancel-after-ready. + + Pre-fix ``CacheSender::Impl::sendResponse`` erases the ready + ``(request, promise)`` pair on cancel without first fulfilling the + promise, and the destructor surfaces ``future_error: Broken promise`` + on the sender future returned by ``respond_and_send_async()``. + """ + # Verify the active cancel path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1). + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") + + (kv_cache_manager_ctx, kv_cache_manager_gen, kv_cache_transceiver_ctx, + kv_cache_transceiver_gen) = disagg_transceiver_pair + + ctx_request = _make_request(0, + LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY, + prompt_len=256) + _add_sequence(kv_cache_manager_ctx, ctx_request) + kv_cache_transceiver_ctx.respond_and_send_async(ctx_request) + + # Wait for the sender to reach the ready state, then cancel it. time.sleep(2) - assert gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR + is_cancelled = kv_cache_transceiver_ctx.cancel_request(ctx_request) + assert is_cancelled, ( + "ctx_request must be cancellable while still ready in " + "mReadyResponses") + + gen_request = _make_request(0, + LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + ctx_request.context_phase_params, + prompt_len=256) + _add_sequence(kv_cache_manager_gen, gen_request) + kv_cache_transceiver_gen.request_and_receive_async(gen_request) + + # Sender side: the cancelled request must surface as an error rather + # than being marked complete. + completed_ids, error_ids = [], [] + deadline = time.time() + 10 + while time.time() < deadline and not error_ids: + completed_ids, error_ids = ( + kv_cache_transceiver_ctx.check_context_transfer_status(1)) + if error_ids: + break + time.sleep(0.1) + + assert ctx_request.py_request_id not in completed_ids, ( + "cancelled ctx request must not appear in completed_ids") + assert ctx_request.py_request_id in error_ids, ( + "cancelled ctx request must surface in error_ids") + + # Receiver side: the matching gen request must observe the cancellation + # as a structured DISAGG_TRANS_ERROR rather than a Broken-promise + # exception. + deadline = time.time() + 10 + while time.time() < deadline and (gen_request.state + != LlmRequestState.DISAGG_TRANS_ERROR): + kv_cache_transceiver_gen.check_gen_transfer_status(1) + time.sleep(0.1) + + assert gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR, ( + "gen request mirroring the cancelled ctx request must end in " + "DISAGG_TRANS_ERROR") + + captured = capfd.readouterr() + merged = captured.out + captured.err + assert "Broken promise" not in merged, ( + "signature #1 reproduced: cancel-after-ready left the sender " + "promise unresolved and the destructor surfaced as Broken " + "promise on the sender future") + + kv_cache_manager_ctx.free_resources(ctx_request) + kv_cache_manager_gen.free_resources(gen_request) + + +_PROBE_TIMEOUT_S = 2.5 + + +@pytest.mark.timeout(120) +@pytest.mark.parametrize("disagg_transceiver_pair", + [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA], + ids=["mha", "mla"], + indirect=True) +def test_check_gen_transfer_status_at_least_one_does_not_block_on_unready_future( + disagg_transceiver_pair): + """Reproduce the gen-side blocking hang in checkGenTransferStatus(1). + + On stock ``rc11`` the polling path called from the PyExecutor disagg + loop unconditionally ``future.get()``s the first selected requester + future, even when its ``wait_for(0)`` is still ``timeout``. A single + in-flight generation request whose context-side ready signal has not + yet arrived therefore blocks the entire decoder event loop, which is + indistinguishable from a wedge. + + The test exercises the same shape as the wedge: drive one full + ctx/gen handshake to completion to capture an opaque comm/cache + state, then enqueue a generation request whose context counterpart + has not yet been ``respond_and_send_async()``-ed and call + ``check_gen_transfer_status(1)`` from a separate thread. The call + must return within a bounded probe timeout larger than the + production 1000ms sender-future wait instead of blocking indefinitely + on the unresolved future. + """ + (kv_cache_manager_ctx, kv_cache_manager_gen, kv_cache_transceiver_ctx, + kv_cache_transceiver_gen) = disagg_transceiver_pair + + # Complete one normal transfer first so we can reuse its opaque + # comm/cache state for the second (intentionally unresolved) generation + # request. + template_ctx_request = _make_request( + 100, LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY, prompt_len=256) + _add_sequence(kv_cache_manager_ctx, template_ctx_request) + kv_cache_transceiver_ctx.respond_and_send_async(template_ctx_request) + + template_gen_request = _make_request( + 100, + LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + template_ctx_request.context_phase_params, + prompt_len=256) + _add_sequence(kv_cache_manager_gen, template_gen_request) + kv_cache_transceiver_gen.request_and_receive_async(template_gen_request) + + _drive_template_handshake_to_completion(kv_cache_transceiver_ctx, + kv_cache_transceiver_gen, + template_ctx_request.py_request_id) + + opaque_state = template_ctx_request.context_phase_params.opaque_state + assert opaque_state is not None, ( + "template ctx_request must expose its opaque comm/cache state for " + "the staged blocked request below to reuse") + + kv_cache_manager_ctx.free_resources(template_ctx_request) + kv_cache_manager_gen.free_resources(template_gen_request) + + # Build a generation request for a different ctx_request_id before the + # sender has any matching ready response. This leaves a real unresolved + # future in the C++ transceiver and reproduces the blocking pattern + # behind signature #4. + blocked_request_id = 101 + blocked_ctx_request = _make_request( + blocked_request_id, + LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY, + prompt_len=256) + _add_sequence(kv_cache_manager_ctx, blocked_ctx_request) + + blocked_context_phase_params = trtllm.ContextPhaseParams( + list(template_ctx_request.context_phase_params.first_gen_tokens), + blocked_request_id, + bytes(opaque_state), + template_ctx_request.context_phase_params.draft_tokens, + template_ctx_request.context_phase_params.ctx_dp_rank, + template_ctx_request.context_phase_params.disagg_info_endpoint, + ) + blocked_gen_request = _make_request( + blocked_request_id, + LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + blocked_context_phase_params, + prompt_len=256) + _add_sequence(kv_cache_manager_gen, blocked_gen_request) + kv_cache_transceiver_gen.request_and_receive_async(blocked_gen_request) + + # Sanity: at_least_request_num=0 must not block under any circumstance. + start = time.time() + kv_cache_transceiver_gen.check_gen_transfer_status(0) + assert time.time() - start < 1.0, ( + "check_gen_transfer_status(0) must be non-blocking even when " + "futures are unresolved") + + # Real reproducer: at_least_request_num=1 must NOT hang on an + # unresolved future. Run it on a worker thread so that pre-fix this + # thread stays blocked past the bounded probe timeout (the failure + # signature), and post-fix it returns after at most the configured + # per-iteration sender-future wait because the unready future is skipped. + check_result = {"returned": False, "error": None} + + def call_blocking_check(): + try: + kv_cache_transceiver_gen.check_gen_transfer_status(1) + except BaseException as exc: # noqa: BLE001 + check_result["error"] = exc + finally: + check_result["returned"] = True + + blocked_check = threading.Thread(target=call_blocking_check, daemon=True) + blocked_check.start() + blocked_check.join(timeout=_PROBE_TIMEOUT_S) + blocked_during_probe = blocked_check.is_alive() + + # Allow the wedged context request to complete cleanly so the worker + # thread can finish (in either pre- or post-fix behaviour) and we can + # tear down the test without leaking threads. + kv_cache_transceiver_ctx.respond_and_send_async(blocked_ctx_request) + + deadline = time.time() + 10 + completed_ids, error_ids = [], [] + while time.time() < deadline and (blocked_ctx_request.py_request_id + not in completed_ids): + completed_ids, error_ids = ( + kv_cache_transceiver_ctx.check_context_transfer_status(1)) + assert blocked_ctx_request.py_request_id not in error_ids, ( + "blocked ctx request must not error during teardown polling") + if blocked_ctx_request.py_request_id in completed_ids: + break + time.sleep(0.1) + assert blocked_ctx_request.py_request_id in completed_ids, ( + "blocked ctx request must complete on the sender side once " + "respond_and_send_async unblocks the receiver") + + if blocked_during_probe: + # Pre-fix path: the thread is wedged inside + # check_gen_transfer_status(1). It will only return after the + # sender's respond_and_send_async unblocks the receive that the + # in-thread call is parked on. + blocked_check.join(timeout=10) + assert not blocked_check.is_alive(), ( + "blocked check thread must drain within 10s after the sender " + "supplies the ready signal") + else: + # Post-fix path: the in-thread call already returned because it + # skipped the unready future. Drain mRequesterFutures via + # additional polling so subsequent assertions see an empty queue. + deadline = time.time() + 10 + while time.time() < deadline and ( + not kv_cache_transceiver_gen.check_gen_transfer_complete()): + kv_cache_transceiver_gen.check_gen_transfer_status(1) + time.sleep(0.1) + + if check_result["error"] is not None: + raise check_result["error"] + assert check_result["returned"], ( + "blocked check thread must signal completion via check_result") + assert not blocked_during_probe, ( + "signature #4 reproduced: check_gen_transfer_status(1) blocked on " + "an unresolved generation future for longer than the bounded probe") + assert kv_cache_transceiver_gen.check_gen_transfer_complete(), ( + "gen-side mRequesterFutures must drain after the blocked request " + "is unblocked") + + kv_cache_manager_ctx.free_resources(blocked_ctx_request) + kv_cache_manager_gen.free_resources(blocked_gen_request) + + +@pytest.mark.timeout(120) +@pytest.mark.parametrize("disagg_transceiver_pair", + [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA], + ids=["mha", "mla"], + indirect=True) +def test_cancel_queued_gen_request_fulfills_receiver_future( + disagg_transceiver_pair, capfd, monkeypatch): + """Reproduce the receiver-side queued-cancel broken-promise. + + Mirrors the sender-side reproducer but on + ``CacheReceiver::Impl::cancelRequest()``. Pre-fix that function erases + the queued ``(request, promise)`` pair without first fulfilling the + promise, so the ``std::promise`` destructor sets + ``future_error: Broken promise`` on the future returned by + ``request_and_receive_async()``. + + The test holds the receiver worker thread busy with a first + generation request that has no matching context counterpart, + enqueues a second generation request behind it, then cancels the + queued request. Post-fix the queued promise is fulfilled with a + structured cancellation exception before the queue entry is erased. + """ + # Verify the active cancel path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1). + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") + + (kv_cache_manager_ctx, kv_cache_manager_gen, kv_cache_transceiver_ctx, + kv_cache_transceiver_gen) = disagg_transceiver_pair + + # Drive one full ctx/gen handshake to completion so we can reuse a + # real opaque comm/cache state for the orphan requests below. + template_ctx_request = _make_request( + 100, LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) + _add_sequence(kv_cache_manager_ctx, template_ctx_request) + kv_cache_transceiver_ctx.respond_and_send_async(template_ctx_request) + + template_gen_request = _make_request( + 100, LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + template_ctx_request.context_phase_params) + _add_sequence(kv_cache_manager_gen, template_gen_request) + kv_cache_transceiver_gen.request_and_receive_async(template_gen_request) + + _drive_template_handshake_to_completion(kv_cache_transceiver_ctx, + kv_cache_transceiver_gen, + template_ctx_request.py_request_id) + + opaque_state = template_ctx_request.context_phase_params.opaque_state + assert opaque_state is not None, ( + "template ctx_request must expose its opaque comm/cache state for " + "the staged orphan requests below to reuse") + + kv_cache_manager_ctx.free_resources(template_ctx_request) + kv_cache_manager_gen.free_resources(template_gen_request) + + def make_orphan_gen_request(request_id): + ctx_phase_params = trtllm.ContextPhaseParams( + list(template_ctx_request.context_phase_params.first_gen_tokens), + request_id, + bytes(opaque_state), + template_ctx_request.context_phase_params.draft_tokens, + template_ctx_request.context_phase_params.ctx_dp_rank, + template_ctx_request.context_phase_params.disagg_info_endpoint, + ) + gen_request = _make_request( + request_id, LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + ctx_phase_params) + _add_sequence(kv_cache_manager_gen, gen_request) + return gen_request + + # Submit a first orphan gen request whose context counterpart will + # never respond_and_send_async. The receiver worker dequeues it and + # parks inside requestSync() / sendRequestInfo(), tying up the worker + # thread. + blocking_gen_request = make_orphan_gen_request(101) + kv_cache_transceiver_gen.request_and_receive_async(blocking_gen_request) + + # Submit a second orphan gen request. The first one is still in + # requestSync(), so this one stays in mRequestsQueue and is the + # actual subject of the queued-cancel reproducer. + queued_gen_request = make_orphan_gen_request(102) + kv_cache_transceiver_gen.request_and_receive_async(queued_gen_request) + + # Wait briefly so the receiver worker has had time to dequeue the + # first request and block on it, leaving the second one queued. + time.sleep(1) + + # Cancel the queued request. Pre-fix this erases the (request, + # promise) pair without fulfilling the promise; post-fix it + # set_exception()s a structured kNETWORK_ERROR before erasing. + is_cancelled = kv_cache_transceiver_gen.cancel_request(queued_gen_request) + assert is_cancelled, ( + "queued_gen_request must still be in the receiver queue when we " + "call cancel_request(); if this fails, the receiver worker may " + "have dequeued faster than expected and the test setup needs to " + "be tightened") + + # Poll the gen-side polling loop and assert the cancelled request + # lands in DISAGG_TRANS_ERROR within a reasonable window. Pre-fix + # this returns via a Broken-promise exception with no useful + # diagnostic; post-fix it returns via the structured kNETWORK_ERROR + # set by the fix. + deadline = time.time() + 10 + while time.time() < deadline and (queued_gen_request.state + != LlmRequestState.DISAGG_TRANS_ERROR): + kv_cache_transceiver_gen.check_gen_transfer_status(1) + time.sleep(0.1) + + assert queued_gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR, ( + "queued gen request must land in DISAGG_TRANS_ERROR after " + "cancellation rather than surfacing a Broken-promise exception") + + # Clean up the intentionally blocked request before the transceiver + # is destroyed. Leaving it in requestSync() would make teardown race + # the worker thread that this test deliberately parked. + is_blocking_cancelled = ( + kv_cache_transceiver_gen.cancel_request(blocking_gen_request)) + assert is_blocking_cancelled, ( + "blocking_gen_request must be cancellable so teardown does not " + "race the parked receiver worker") + deadline = time.time() + 10 + while time.time() < deadline and (blocking_gen_request.state + != LlmRequestState.DISAGG_TRANS_ERROR): + kv_cache_transceiver_gen.check_gen_transfer_status(1) + time.sleep(0.1) + assert blocking_gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR, ( + "blocking_gen_request must drain to DISAGG_TRANS_ERROR before " + "transceiver destruction") + + kv_cache_manager_gen.free_resources(blocking_gen_request) + kv_cache_manager_gen.free_resources(queued_gen_request) + + captured = capfd.readouterr() + merged = captured.out + captured.err + assert "Broken promise" not in merged, ( + "signature #5 reproduced: cancelling a queued generation request " + "left its std::promise unresolved and the destructor surfaced " + "as Broken promise on the consumer side") def create_hybrid_cache_manager(mapping, dtype, mamba_conv_dtype=torch.float16, mamba_ssm_dtype=torch.float16): - """ - Create a MambaHybridCacheManager for testing hybrid models. - This manager handles both KV cache (attention layers) and Mamba cache (RNN layers). + """Create a MambaHybridCacheManager for testing hybrid models. + + This manager handles both KV cache (attention layers) and Mamba cache + (RNN layers). Args: mapping: The mapping configuration. @@ -261,6 +967,11 @@ def create_hybrid_cache_manager(mapping, mamba_layer_mask=mamba_layer_mask, mamba_cache_dtype=mamba_conv_dtype, mamba_ssm_cache_dtype=mamba_ssm_dtype, + # Disagg KV transfer is the whole point of this test file, and the + # MixedMambaHybridCacheManager (selected when is_disagg=True) is the + # only variant whose mamba state is exposed via the disagg + # transceivers exercised below. + is_disagg=True, kv_cache_config=KvCacheConfig( max_tokens=256, enable_block_reuse=False, @@ -303,8 +1014,7 @@ def fill_hybrid_cache_buffers(hybrid_cache_manager): @pytest.fixture(scope="function") def hybrid_dtypes(request): - """ - Returns (kv_dtype, mamba_conv_dtype, mamba_ssm_dtype) based on the parametrized string. + """Return dtypes for the parametrized hybrid-cache test case. KV dtype: fp8, bf16 Conv dtype: fp8, bf16, fp32 @@ -459,6 +1169,8 @@ def test_hybrid_cache_transceiver_single_process(backend, hybrid_dtypes, @pytest.mark.parametrize("backend", ["NIXL", "UCX"], ids=["NIXL", "UCX"]) def test_hybrid_cache_transceiver_cancel_request(backend, monkeypatch): monkeypatch.setenv("TRTLLM_USE_CPP_MAMBA", "1") + # Verify the active cancel path (TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1). + monkeypatch.setenv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL", "1") mapping = Mapping(world_size=1, rank=0) dtype = DataType.HALF