From 6f6a7b7451c524ae116351d8b6a0a4e3f03b4fb2 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 29 May 2026 00:59:11 -0700 Subject: [PATCH] [https://nvbugs/6104831][fix] disagg buffer-pool RAII and NIXL agent lifetime fixes Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/baseTransBuffer.cpp | 186 +++++++++++++++++- .../batch_manager/baseTransBuffer.h | 175 +++++++++++++++- .../batch_manager/cacheFormatter.cpp | 6 +- .../batch_manager/dataTransceiver.cpp | 68 ++++++- .../batch_manager/mlaCacheFormatter.cpp | 5 +- .../nixl_utils/agentBindings.cpp | 5 +- 6 files changed, 419 insertions(+), 26 deletions(-) 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 17dd557be1a3..092f9fe31cfe 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -523,6 +523,8 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio // 5. send the buffer to the corresponding target. Ideally, we send only once (one buffer) for each target. auto cacheBufferId = mCacheTransBufferManager->assignBufferIndexForSend(); + // RAII wrapper releases the slot on any exception unwind. + BufferIndexHolder sendHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/false); int peerDuplicateHeadFactor = targetInfo.mPeerDupHeadFactor; auto bufferTargetNum = targetNum / peerDuplicateHeadFactor; auto ppRank = selfIdx @@ -604,9 +606,9 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio sendAllBuffers(session, deviceId, outputSplitCaches, bufferCoverTargetNum, preAllocSendBuffer, bufferManager, targetInfo, pickUpConnections); - session.setTime(TransferSession::kTimeTransmissions); + sendHolder.release(); - mCacheTransBufferManager->freeBufferIndexForSend(cacheBufferId); + session.setTime(TransferSession::kTimeTransmissions); session.setTime(TransferSession::kTimePostprocess); } TLLM_LOG_DEBUG( diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 13e95dd86e48..0da619c23cd6 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -819,7 +819,25 @@ class CacheReceiver::Impl } } + // Backward-compat overload: acquires buffer indices internally with no RAII. + // requestSync uses the 2-arg overload below to wrap the indices in + // BufferIndexHolder for exception-safe release. TransferSession sendRequestInfo(LlmRequest const& llmRequest) + { + std::vector> cacheBufferIds; + auto* agentConnectionManager = dynamic_cast(mManager); + if (agentConnectionManager) + { + for (auto& cacheTransBufferManager : agentConnectionManager->getCacheTransBufferManagers()) + { + cacheBufferIds.push_back(cacheTransBufferManager->assignBufferIndexForRecv()); + } + TLLM_CHECK(!cacheBufferIds.empty()); + } + return sendRequestInfo(llmRequest, std::move(cacheBufferIds)); + } + + TransferSession sendRequestInfo(LlmRequest const& llmRequest, std::vector> cacheBufferIds) { uint64_t requestId = llmRequest.getContextPhaseParams().value().getReqId(); auto const& contextState = llmRequest.getDataTransceiverState(); @@ -855,15 +873,6 @@ class CacheReceiver::Impl } auto* agentConnectionManager = dynamic_cast(mManager); - std::vector> cacheBufferIds; - if (agentConnectionManager) - { - for (auto& cacheTransBufferManager : agentConnectionManager->getCacheTransBufferManagers()) - { - cacheBufferIds.push_back(cacheTransBufferManager->assignBufferIndexForRecv()); - } - TLLM_CHECK(!cacheBufferIds.empty()); - } auto allCounterparts = mCacheTransferLayer.computeCounterparts(mSelfState.getCommState().value().getSelfIdx(), contextState); @@ -1063,12 +1072,42 @@ class CacheReceiver::Impl llmRequest.getContextPhaseParams().value().getReqId()); llmRequest.setKvCacheTransferStart(std::chrono::steady_clock::now()); TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); - auto session = sendRequestInfo(llmRequest); + + // Own the receive-side buffer indices with BufferIndexHolder across + // the full receive flow so any throw or early-return releases the + // slot via destructor. detach() is invoked after a successful + // receiveSync because the formatter already freed each slot. + std::vector recvHolders; + std::vector> cacheBufferIds; + auto const reqIdForLog = std::make_optional(static_cast(llmRequest.mRequestId)); + if (auto* agentConnectionManagerForAcq = dynamic_cast(mManager)) + { + auto const& managers = agentConnectionManagerForAcq->getCacheTransBufferManagers(); + recvHolders.reserve(managers.size()); + cacheBufferIds.reserve(managers.size()); + for (auto* cacheTransBufferManager : managers) + { + auto rawIdx = cacheTransBufferManager->assignBufferIndexForRecv(); + 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); + } + } + TLLM_CHECK(!cacheBufferIds.empty()); + } + + auto session = sendRequestInfo(llmRequest, std::move(cacheBufferIds)); session.setTime(TransferSession::kTimeRequestInfo); bool isReady = receiveReadySignal(session); if (!isReady) { // Reuse the error state for the cancelled request. + // recvHolders' destructors release the slots back to the pool. llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); return; @@ -1076,6 +1115,15 @@ class CacheReceiver::Impl receiveSync(session); llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now()); + // Happy path: the formatter invoked inside receiveSync already + // released each buffer index via freeBufferIndexForRecv. Detach + // each holder so its destructor does not double-release on + // stack unwind. + for (auto& h : recvHolders) + { + (void) h.detach(); + } + TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "End calling requestSync for request ID: %zu, context request ID: %zu.", llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId()); diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp index c72090867f29..e607a9d7bdf0 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp @@ -254,6 +254,9 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses }; auto bufferEleSizes = getBufferSizeForTarget(); auto cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForSend(); + // RAII wrapper releases the slot on any exception unwind. + BufferIndexHolder sendHolder( + *mCacheTransBufferManagers[transferIndexerKCache], cacheBufferId, /*isRecv=*/false); auto result = mCacheTransBufferManagers[transferIndexerKCache]->getOrAllocateSendBuffers( cacheBufferId, static_cast(pPDomainSize * cPDomainSize), bufferEleSizes, bufferManager); auto& outputSplitCaches = std::get<0>(result); @@ -380,7 +383,7 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses { sendBufferFun(deviceId, pickUpConnections[0]); } - mCacheTransBufferManagers[transferIndexerKCache]->freeBufferIndexForSend(cacheBufferId); + sendHolder.release(); } session.setTime(TransferSession::kTimeTransmissions); session.setTime(TransferSession::kTimePostprocess); 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..220868893f08 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp @@ -230,7 +230,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) @@ -263,7 +263,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)