From 268200e31fed8d57a19046f5b4e3a173f0174d3f Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 21 May 2026 15:50:02 -0700 Subject: [PATCH 1/9] [NVBUGS-6104831][fix] keep NIXL agent alive while its TransferStatus exists Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../executor/cache_transmission/nixl_utils/agentBindings.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp index 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) From 0a6f7b54ae810cd89d22512d0dc1681db42e3a81 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Sat, 23 May 2026 02:18:20 -0700 Subject: [PATCH 2/9] [https://nvbugs/6104831][fix] remove rank-asymmetric gates on disagg ctx-side checkContextTransferStatus Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 66 ++++++++++++------- 1 file changed, 42 insertions(+), 24 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index a893264d1ef1..6243b11754eb 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1884,17 +1884,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 @@ -2429,19 +2446,20 @@ def _prepare_and_schedule_batch(self): req.py_disaggregated_params and req.py_disaggregated_params. schedule_style == DisaggScheduleStyle.GENERATION_FIRST for req in self.active_requests) - if num_fitting_reqs == 0 and not fitting_disagg_gen_init_requests: - if not all_gen_first: - logger.warning( - "num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache" - ) - self._check_disagg_ctx_cache_transfer_status(1) - elif self.async_transfer_manager.has_any_inflight_requests(): - # Non-blocking cleanup of completed/timed-out transfers - # to free KV blocks. We avoid the blocking check because - # gen-first requests may be waiting for peer info (which - # would block indefinitely), but completed transfers must - # still be reaped so that KV cache can be reclaimed. - self._check_disagg_ctx_cache_transfer_status(0) + # Same rank-symmetric pattern as _executor_loop_pp above; see + # rationale comment there. The C++ gatherRequestIds collective + # inside checkContextTransferStatus must be entered by every + # ctx-side rank together. at_least_num is rank-local and safe + # to diverge (it only controls per-request loop behavior inside + # the C++ function, not collective entry). + at_least_num = 0 + if (num_fitting_reqs == 0 and not fitting_disagg_gen_init_requests + and not all_gen_first): + logger.warning( + "num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache" + ) + at_least_num = 1 + self._check_disagg_ctx_cache_transfer_status(at_least_num) # In gen-only benchmark mode, all requests must fit in KV cache # simultaneously because generation requests do not release their From c06e70ce672b7df1342215958dd9f823828ffd7a Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 29 May 2026 17:33:06 -0700 Subject: [PATCH 3/9] [https://nvbugs/6104831][fix] Enforce request and buffer index lifecycle integrity. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 43 +++++-- .../batch_manager/baseTransBuffer.cpp | 30 +++++ .../batch_manager/baseTransBuffer.h | 111 ++++++++++++++++++ .../batch_manager/cacheTransceiver.cpp | 105 ++++++++++++++--- .../trtGptModelInflightBatching.cpp | 6 +- .../batch_manager/cacheTransceiver.cpp | 8 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 33 ++++-- 7 files changed, 289 insertions(+), 47 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 452b90462448..869aa38d9c2a 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -35,6 +35,7 @@ #include #include #include +#include #include using SizeType32 = tensorrt_llm::runtime::SizeType32; @@ -204,13 +205,17 @@ class BaseCacheTransceiver { public: virtual ~BaseCacheTransceiver() = default; - virtual void respondAndSendAsync(LlmRequest* llmRequest) = 0; + // These methods take std::shared_ptr so the transceiver and + // its async workers can hold a strong reference for the duration of the + // transfer. See the comment on CacheTransceiver::mSenderFutures for the + // lifetime invariant (kept in one place to avoid drift). + virtual void respondAndSendAsync(std::shared_ptr llmRequest) = 0; virtual void respondAndSendLayerWise( RequestVector const& requests, std::shared_ptr const& progress) = 0; - virtual void requestAndReceiveSync(LlmRequest* llmRequest) = 0; - virtual void requestAndReceiveAsync(LlmRequest* llmRequest) = 0; + virtual void requestAndReceiveSync(std::shared_ptr llmRequest) = 0; + virtual void requestAndReceiveAsync(std::shared_ptr llmRequest) = 0; /// Check all requests transferring context, and return the requests that have completed or encountered an error. virtual RequestStatuses checkContextTransferStatus( @@ -221,7 +226,7 @@ class BaseCacheTransceiver [[nodiscard]] virtual bool checkGenTransferComplete() const = 0; - virtual bool cancelRequest(LlmRequest* llmRequest) = 0; + virtual bool cancelRequest(std::shared_ptr llmRequest) = 0; }; class CacheTransceiver : public BaseCacheTransceiver @@ -252,13 +257,13 @@ class CacheTransceiver : public BaseCacheTransceiver virtual ~CacheTransceiver(); - void respondAndSendAsync(LlmRequest* llmRequest) override; + void respondAndSendAsync(std::shared_ptr llmRequest) override; void respondAndSendLayerWise( RequestVector const& requests, std::shared_ptr const& progress) override; - void requestAndReceiveSync(LlmRequest* llmRequest) override; - void requestAndReceiveAsync(LlmRequest* llmRequest) override; + void requestAndReceiveSync(std::shared_ptr llmRequest) override; + void requestAndReceiveAsync(std::shared_ptr llmRequest) override; RequestStatuses checkContextTransferStatus( std::optional const& atLeastRequestNum = std::nullopt, bool markComplete = false) override; @@ -267,7 +272,7 @@ class CacheTransceiver : public BaseCacheTransceiver [[nodiscard]] bool checkGenTransferComplete() const override; - virtual bool cancelRequest(LlmRequest* llmRequest) override; + virtual bool cancelRequest(std::shared_ptr llmRequest) override; private: void initializeCommState(); @@ -276,8 +281,26 @@ class CacheTransceiver : public BaseCacheTransceiver std::unique_ptr mCacheSender; std::unique_ptr mCacheReceiver; - std::vector>> mSenderFutures; - std::vector>> mRequesterFutures; + // Store shared_ptr rather than raw LlmRequest* so the futures map holds a + // strong reference for the duration of the transfer. Otherwise Python's + // _terminate_request can drop its pybind shared_ptr while the C++ side's + // raw pointer is still dereferenced by checkGenTransferStatus / + // checkContextTransferStatus (the UAF forensically confirmed via + // MALLOC_PERTURB_=85 producing mRequestId=0x5555555555555555). Entries are + // erased only in the normal eviction paths (completion, deadline, + // exception), ensuring the request outlives every C++ access but doesn't + // leak past the transfer lifecycle. + std::vector, std::future>> mSenderFutures; + std::vector, std::future>> mRequesterFutures; + // Dedup sets for observe-only timeout WARN logs. Each requestId is logged + // at most once per transfer so a stuck transfer doesn't produce a WARN + // flood as the polling loop revisits the entry each tick. Indexed by role: + // mTimedOutSenderIds tracks context-side (mSenderFutures) and + // mTimedOutRequesterIds tracks generation-side (mRequesterFutures). + // Entries are erased when the corresponding future is removed from the + // futures vector (completion, error, or status transition). + std::unordered_set mTimedOutSenderIds; + std::unordered_set mTimedOutRequesterIds; mpi::MpiComm const* mMpiWorldComm{nullptr}; std::shared_ptr mGroupComm; diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp index 58092897ebbe..5bd2d6938fcf 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp @@ -26,6 +26,36 @@ namespace tensorrt_llm::batch_manager { +void BufferIndexHolder::release() noexcept +{ + // Happy-path release: frees the slot and disarms the holder in one + // noexcept call. Used in place of an older detach() + explicit + // freeBufferIndex*() sequence so a throw between the two calls cannot + // leave the holder in a partially-released state. + if (!mHeld || mMgr == nullptr) + { + return; + } + try + { + if (mIsRecv) + { + mMgr->freeBufferIndexForRecv(mIndex); + } + else + { + mMgr->freeBufferIndexForSend(mIndex); + } + } + catch (...) + { + // Swallow; the destructor must be noexcept and any exit path that + // failed to release explicitly relies on this fallback to free the + // slot. + } + mHeld = false; +} + BaseTransBufferManager::BaseTransBufferManager( size_t transferBufferSize, nvinfer1::DataType dataType, std::optional maxNumTokens) : mDataType{dataType} diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h index 1efeb89ccc04..26bf9be8bd53 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h @@ -46,6 +46,117 @@ enum class BufferKind : uint8_t kRNN = 2 }; +class BaseTransBufferManager; + +/// @brief RAII scoped holder for a buffer index acquired from +/// BaseTransBufferManager::assignBufferIndexForRecv / +/// assignBufferIndexForSend. Releases the index on destruction, +/// including stack unwind from exceptions. +/// +/// Motivation: CacheReceiver::Impl::requestSync has at least six exit +/// paths (normal, early-cancel, not-ready, cancel-after-ready, +/// receiveReadySignal cancelled, exception from requestSync). Pre-fix, +/// the buffer-index release lived inside receiveSync's formatter, so any +/// exit path that skipped receiveSync leaked one index. Under saturation +/// even a single leaked index permanently wedged the (default size-1) +/// pool, so every subsequent request waited forever for an index that +/// would never be released. +/// +/// This holder closes that class of bug rather than patching the one +/// observed branch. Move-only so ownership is unambiguous; `detach()` +/// hands off ownership when the formatter inside receiveSync takes the +/// buffer's release responsibility on the happy path. +class BufferIndexHolder +{ +public: + BufferIndexHolder() = default; + + BufferIndexHolder(BaseTransBufferManager& mgr, std::optional index, bool isRecv, + std::optional requestIdForLog = std::nullopt) noexcept + : mMgr(&mgr) + , mIndex(index) + , mHeld(index.has_value()) + , mIsRecv(isRecv) + , mRequestIdForLog(requestIdForLog) + { + } + + // The destructor is inline but `release()` is out-of-line in + // baseTransBuffer.cpp because it dereferences BaseTransBufferManager and + // needs the full definition; defining it inline would create an include + // cycle between this header and the manager definition. + ~BufferIndexHolder() + { + release(); + } + + BufferIndexHolder(BufferIndexHolder const&) = delete; + BufferIndexHolder& operator=(BufferIndexHolder const&) = delete; + + BufferIndexHolder(BufferIndexHolder&& other) noexcept + : mMgr(other.mMgr) + , mIndex(other.mIndex) + , mHeld(other.mHeld) + , mIsRecv(other.mIsRecv) + , mRequestIdForLog(other.mRequestIdForLog) + { + other.mHeld = false; + } + + BufferIndexHolder& operator=(BufferIndexHolder&& other) noexcept + { + if (this != &other) + { + release(); + mMgr = other.mMgr; + mIndex = other.mIndex; + mHeld = other.mHeld; + mIsRecv = other.mIsRecv; + mRequestIdForLog = other.mRequestIdForLog; + other.mHeld = false; + } + return *this; + } + + [[nodiscard]] std::optional index() const noexcept + { + return mIndex; + } + + [[nodiscard]] bool held() const noexcept + { + return mHeld; + } + + /// @brief Relinquish ownership without releasing. Use when a downstream + /// owner (e.g. the formatter inside receiveSync) takes over the + /// release responsibility on the happy path. + std::optional detach() noexcept + { + mHeld = false; + return mIndex; + } + + /// @brief Happy-path release. Frees the slot immediately and disarms the + /// destructor. Use this on any path where the caller has confirmed + /// the slot is no longer needed and the release is the expected + /// outcome (e.g. the sender formatter after sendAllBuffers + /// returns). After this call, the holder owns nothing; subsequent + /// destructor or move-assignment is a no-op. + /// + /// If the holder goes out of scope with mHeld still true (exception + /// or early return that forgot to call release/detach), the + /// destructor calls release() to free the slot. + void release() noexcept; + +private: + BaseTransBufferManager* mMgr{nullptr}; + std::optional mIndex{}; + bool mHeld{false}; + bool mIsRecv{true}; + std::optional mRequestIdForLog{}; +}; + /// @brief Base class for cache transfer buffer management. /// Handles buffer pool allocation, index assignment, and slicing. /// Derived classes provide cache-specific size calculations. diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 42fa3d9eac7c..d75936939ab0 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -379,7 +379,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 +393,9 @@ void CacheTransceiver::respondAndSendAsync(LlmRequest* llmRequest) } return; } - setContextState(llmRequest); + setContextState(llmRequest.get()); auto future = mCacheSender->sendAsync(*llmRequest); - mSenderFutures.emplace_back(llmRequest, std::move(future)); + mSenderFutures.emplace_back(std::move(llmRequest), std::move(future)); } void CacheTransceiver::respondAndSendLayerWise( @@ -411,11 +411,11 @@ 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)); + mSenderFutures.emplace_back(llmRequest, std::move(future)); } } -void CacheTransceiver::requestAndReceiveSync(LlmRequest* llmRequest) +void CacheTransceiver::requestAndReceiveSync(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isGenerationOnlyRequest()); { @@ -425,21 +425,27 @@ void CacheTransceiver::requestAndReceiveSync(LlmRequest* llmRequest) llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); } -void CacheTransceiver::requestAndReceiveAsync(LlmRequest* llmRequest) +void CacheTransceiver::requestAndReceiveAsync(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isGenerationOnlyRequest()); + auto const requestId = llmRequest->mRequestId; if (std::find_if(mRequesterFutures.begin(), mRequesterFutures.end(), - [llmRequest](auto const& pair) { return pair.first->mRequestId == llmRequest->mRequestId; }) + [requestId](auto const& pair) { return pair.first->mRequestId == requestId; }) != mRequesterFutures.end()) { - TLLM_LOG_WARNING("Request ID %zu is already in mRequestFutures.", llmRequest->mRequestId); + TLLM_LOG_WARNING("Request ID %zu is already in mRequestFutures.", requestId); return; } auto future = mCacheReceiver->receiveAsync(*llmRequest); - mRequesterFutures.emplace_back(llmRequest, std::move(future)); - llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS); + // Cache the raw pointer before moving the shared_ptr into the futures vector + // so we can still call setState() to preserve upstream's + // receiveAsync -> emplace -> setState ordering (A5 reorder explicitly NOT + // included in this Tier-1 baseline). + auto* requestPtr = llmRequest.get(); + mRequesterFutures.emplace_back(std::move(llmRequest), std::move(future)); + requestPtr->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS); } std::vector gatherRequestIds( @@ -546,6 +552,17 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( { senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); } + // Observe-only overall-deadline check. The configured kvTransferTimeoutMs + // is the wall-clock budget for an entire transfer; if exceeded we emit a + // single WARN (dedup'd via mTimedOutSenderIds) but do NOT take action — + // this Tier-1 baseline intentionally leaves cancel / state transitions / + // eviction to higher layers so we can isolate the source of the + // cross-rank wedge. + std::optional kvTransferTimeoutMs = std::nullopt; + if (mCacheTransceiverConfig.has_value()) + { + kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); + } auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupTPInDPComm : mGroupTensorParaComm; std::vector contextCompleteRequestIds; @@ -602,6 +619,26 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( for (auto it = mSenderFutures.begin(); it != mSenderFutures.end();) { auto& [request, future] = *it; + // Observe-only overall-deadline WARN. Emitted at most once per + // requestId (dedup'd via mTimedOutSenderIds) when the wall-clock + // elapsed time since transfer start exceeds the configured budget. + // No action is taken — this is purely diagnostic for the Tier-1 + // baseline. If the request later completes or errors, the entry is + // removed from mTimedOutSenderIds in the normal eviction paths + // below so a re-enqueued requestId can be warned about again. + if (kvTransferTimeoutMs.has_value()) + { + auto elapsed = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()); + auto elapsedMs = static_cast(elapsed.count()); + if (elapsedMs > kvTransferTimeoutMs.value() && mTimedOutSenderIds.insert(request->mRequestId).second) + { + TLLM_LOG_WARNING( + "Context KV cache transfer for request %ld exceeded configured timeout: " + "elapsed %ld ms > limit %d ms (observe-only).", + request->mRequestId, elapsedMs, kvTransferTimeoutMs.value()); + } + } if (blockAll || (toCompleteIdSet.find(request->mRequestId) != toCompleteIdSet.end())) { try @@ -616,6 +653,7 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( { request->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); } + mTimedOutSenderIds.erase(request->mRequestId); it = mSenderFutures.erase(it); } else if (status == std::future_status::timeout) @@ -631,6 +669,7 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); requestsStatus.errorRequestIds.insert(request->mRequestId); + mTimedOutSenderIds.erase(request->mRequestId); it = mSenderFutures.erase(it); } } @@ -640,6 +679,7 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( "Error occurred during context transfer for request %ld: %s", request->mRequestId, e.what()); request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); requestsStatus.errorRequestIds.insert(request->mRequestId); + mTimedOutSenderIds.erase(request->mRequestId); it = mSenderFutures.erase(it); } } @@ -765,40 +805,67 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR " checkGenTransferStatus toCompleteIdSet size: %zu, atLeastRequestNum: %d ", toCompleteIdSet.size(), atLeastRequestNum.value_or(0)); } + // Observe-only overall-deadline check for the generation side. Mirrors + // the context-side logic in checkContextTransferStatus: emit a WARN at + // most once per requestId via the mTimedOutRequesterIds dedup set when + // the wall-clock elapsed time since transfer start exceeds the + // configured budget. No action is taken — eviction / state transitions + // / cancellation are intentionally left to higher layers in this + // Tier-1 baseline. + std::optional kvTransferTimeoutMs = std::nullopt; + if (mCacheTransceiverConfig.has_value()) + { + kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); + } for (auto it = mRequesterFutures.begin(); it != mRequesterFutures.end();) { - if (blockAll || toCompleteIdSet.find(it->first->mRequestId) != toCompleteIdSet.end()) + auto& request = it->first; + if (kvTransferTimeoutMs.has_value()) + { + auto elapsed = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()); + auto elapsedMs = static_cast(elapsed.count()); + if (elapsedMs > kvTransferTimeoutMs.value() && mTimedOutRequesterIds.insert(request->mRequestId).second) + { + TLLM_LOG_WARNING( + "Generation KV cache transfer for request %ld exceeded configured timeout: " + "elapsed %ld ms > limit %d ms (observe-only).", + request->mRequestId, elapsedMs, kvTransferTimeoutMs.value()); + } + } + if (blockAll || toCompleteIdSet.find(request->mRequestId) != toCompleteIdSet.end()) { try { it->second.get(); - it->first->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); + request->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); // Gather the kv cache transfer time from all workers and update to leader rank if (!common::getEnvKVCacheTimeOutputPath().empty()) { - auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; - updateKVCacheTransferBW(syncComm, it->first); + auto bwSyncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; + updateKVCacheTransferBW(bwSyncComm, request.get()); } } 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); + "Error occurred during generation transfer for request %ld: %s", request->mRequestId, e.what()); + request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); } if (useMPI()) { TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***", - it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId()); + request->mRequestId, request->getContextPhaseParams().value().getReqId()); } else { TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***", - it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId()); + request->mRequestId, request->getContextPhaseParams().value().getReqId()); } + mTimedOutRequesterIds.erase(request->mRequestId); it = mRequesterFutures.erase(it); } else @@ -813,7 +880,7 @@ bool CacheTransceiver::checkGenTransferComplete() const return mRequesterFutures.empty(); } -bool CacheTransceiver::cancelRequest(LlmRequest* llmRequest) +bool CacheTransceiver::cancelRequest(std::shared_ptr llmRequest) { if (llmRequest->isContextOnlyRequest()) { diff --git a/cpp/tensorrt_llm/batch_manager/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/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index 8adddf1c86c0..5257c7adb1c9 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -46,17 +46,17 @@ class PyCacheTransceiver : public tb::BaseCacheTransceiver // using BaseCacheTransceiver::BaseCacheTransceiver; // Inherit constructors NB_TRAMPOLINE(tb::BaseCacheTransceiver, 6); - void respondAndSendAsync(tb::LlmRequest* llmRequest) override + void respondAndSendAsync(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(respondAndSendAsync, llmRequest); } - void requestAndReceiveSync(tb::LlmRequest* llmRequest) override + void requestAndReceiveSync(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(requestAndReceiveSync, llmRequest); } - void requestAndReceiveAsync(tb::LlmRequest* llmRequest) override + void requestAndReceiveAsync(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(requestAndReceiveAsync, llmRequest); } @@ -77,7 +77,7 @@ class PyCacheTransceiver : public tb::BaseCacheTransceiver NB_OVERRIDE_PURE(checkGenTransferComplete); } - bool cancelRequest(tb::LlmRequest* llmRequest) override + bool cancelRequest(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(cancelRequest, llmRequest); } diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 6243b11754eb..4badd200f56e 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3700,11 +3700,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. @@ -3713,10 +3718,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 @@ -3868,8 +3871,16 @@ def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests): 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) + # 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): From a60ca3baadf639b10424103cb9c4fa10739abd59 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 29 May 2026 18:20:07 -0700 Subject: [PATCH 4/9] [https://nvbugs/6104831][fix] Wire BufferIndexHolder at formatter sites Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 26 ++------- .../batch_manager/baseTransBuffer.cpp | 8 +-- .../batch_manager/baseTransBuffer.h | 38 ++----------- .../batch_manager/cacheFormatter.cpp | 10 ++-- .../batch_manager/cacheTransceiver.cpp | 27 +-------- .../batch_manager/mlaCacheFormatter.cpp | 12 ++-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 56 ++++--------------- 7 files changed, 40 insertions(+), 137 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 869aa38d9c2a..ea061af74db7 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -205,10 +205,8 @@ class BaseCacheTransceiver { public: virtual ~BaseCacheTransceiver() = default; - // These methods take std::shared_ptr so the transceiver and - // its async workers can hold a strong reference for the duration of the - // transfer. See the comment on CacheTransceiver::mSenderFutures for the - // lifetime invariant (kept in one place to avoid drift). + // Async entry points take shared_ptr so the async worker holds a strong + // reference for the transfer lifetime; see CacheTransceiver::mSenderFutures. virtual void respondAndSendAsync(std::shared_ptr llmRequest) = 0; virtual void respondAndSendLayerWise( RequestVector const& requests, std::shared_ptr const& progress) @@ -281,24 +279,12 @@ class CacheTransceiver : public BaseCacheTransceiver std::unique_ptr mCacheSender; std::unique_ptr mCacheReceiver; - // Store shared_ptr rather than raw LlmRequest* so the futures map holds a - // strong reference for the duration of the transfer. Otherwise Python's - // _terminate_request can drop its pybind shared_ptr while the C++ side's - // raw pointer is still dereferenced by checkGenTransferStatus / - // checkContextTransferStatus (the UAF forensically confirmed via - // MALLOC_PERTURB_=85 producing mRequestId=0x5555555555555555). Entries are - // erased only in the normal eviction paths (completion, deadline, - // exception), ensuring the request outlives every C++ access but doesn't - // leak past the transfer lifecycle. + // shared_ptr (not raw LlmRequest*) so the futures hold a strong reference for + // the transfer lifetime; otherwise Python's _terminate_request can drop the + // request while a C++ status check still dereferences it. std::vector, std::future>> mSenderFutures; std::vector, std::future>> mRequesterFutures; - // Dedup sets for observe-only timeout WARN logs. Each requestId is logged - // at most once per transfer so a stuck transfer doesn't produce a WARN - // flood as the polling loop revisits the entry each tick. Indexed by role: - // mTimedOutSenderIds tracks context-side (mSenderFutures) and - // mTimedOutRequesterIds tracks generation-side (mRequesterFutures). - // Entries are erased when the corresponding future is removed from the - // futures vector (completion, error, or status transition). + // Dedup sets so observe-only timeout WARN logs fire at most once per stuck request. std::unordered_set mTimedOutSenderIds; std::unordered_set mTimedOutRequesterIds; mpi::MpiComm const* mMpiWorldComm{nullptr}; diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp index 5bd2d6938fcf..3be8ee22bc1f 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp @@ -28,10 +28,6 @@ namespace tensorrt_llm::batch_manager void BufferIndexHolder::release() noexcept { - // Happy-path release: frees the slot and disarms the holder in one - // noexcept call. Used in place of an older detach() + explicit - // freeBufferIndex*() sequence so a throw between the two calls cannot - // leave the holder in a partially-released state. if (!mHeld || mMgr == nullptr) { return; @@ -49,9 +45,7 @@ void BufferIndexHolder::release() noexcept } catch (...) { - // Swallow; the destructor must be noexcept and any exit path that - // failed to release explicitly relies on this fallback to free the - // slot. + // noexcept: swallow so the destructor can never throw. } mHeld = false; } diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h index 26bf9be8bd53..bce94702e121 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h @@ -48,24 +48,9 @@ enum class BufferKind : uint8_t 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. +/// @brief RAII holder for an index from BaseTransBufferManager::assignBufferIndexFor{Send,Recv}. +/// Releases on destruction (incl. exception unwind). Move-only; call release() on +/// the happy path or detach() when ownership is handed off downstream. class BufferIndexHolder { public: @@ -81,10 +66,8 @@ class BufferIndexHolder { } - // 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. + // release() is out-of-line in baseTransBuffer.cpp to avoid the include cycle + // with BaseTransBufferManager. ~BufferIndexHolder() { release(); @@ -137,16 +120,7 @@ class BufferIndexHolder 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. + /// @brief Release the slot now and disarm the destructor. Safe to call multiple times. void release() noexcept; private: diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index 17dd557be1a3..68a398f1b52d 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -523,6 +523,7 @@ 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(); + BufferIndexHolder sendHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/false); int peerDuplicateHeadFactor = targetInfo.mPeerDupHeadFactor; auto bufferTargetNum = targetNum / peerDuplicateHeadFactor; auto ppRank = selfIdx @@ -606,7 +607,7 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio session.setTime(TransferSession::kTimeTransmissions); - mCacheTransBufferManager->freeBufferIndexForSend(cacheBufferId); + sendHolder.release(); session.setTime(TransferSession::kTimePostprocess); } TLLM_LOG_DEBUG( @@ -849,6 +850,7 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess size_t remainNoCoverTargetNum = 0; size_t bufferCoverTargetNum = 0; std::optional cacheBufferId = std::nullopt; + BufferIndexHolder recvHolder; { NVTX3_SCOPED_RANGE(formatInputAllocBuffer); @@ -864,6 +866,7 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess { cacheBufferId = mCacheTransBufferManager->assignBufferIndexForRecv(); } + recvHolder = BufferIndexHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/true); auto [recvSplitCachestmp, bufferCoverTargetNumtmp, onlyUseDynamicBuffer] = mCacheTransBufferManager->getOrAllocateRecvBuffers( cacheBufferId, static_cast(targetNum), bufferEleSizes, bufferManager); @@ -997,10 +1000,7 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess recvSplitCaches, outputBuffersPerWindow, destConfig, selfConfig, selfIdx, bufferManager); bufferManager.getStream().synchronize(); - if (cacheBufferId.has_value()) - { - mCacheTransBufferManager->freeBufferIndexForRecv(cacheBufferId); - } + recvHolder.release(); } session.setTime(TransferSession::kTimePostprocess); } diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index d75936939ab0..de146a7652de 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -439,10 +439,6 @@ void CacheTransceiver::requestAndReceiveAsync(std::shared_ptr llmReq } auto future = mCacheReceiver->receiveAsync(*llmRequest); - // Cache the raw pointer before moving the shared_ptr into the futures vector - // so we can still call setState() to preserve upstream's - // receiveAsync -> emplace -> setState ordering (A5 reorder explicitly NOT - // included in this Tier-1 baseline). auto* requestPtr = llmRequest.get(); mRequesterFutures.emplace_back(std::move(llmRequest), std::move(future)); requestPtr->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS); @@ -552,12 +548,8 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( { senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); } - // Observe-only overall-deadline check. The configured kvTransferTimeoutMs - // is the wall-clock budget for an entire transfer; if exceeded we emit a - // single WARN (dedup'd via mTimedOutSenderIds) but do NOT take action — - // this Tier-1 baseline intentionally leaves cancel / state transitions / - // eviction to higher layers so we can isolate the source of the - // cross-rank wedge. + // Observe-only: WARN per-request when the wall-clock transfer time exceeds + // kvTransferTimeoutMs. No cancellation, eviction, or state transition. std::optional kvTransferTimeoutMs = std::nullopt; if (mCacheTransceiverConfig.has_value()) { @@ -619,13 +611,6 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( for (auto it = mSenderFutures.begin(); it != mSenderFutures.end();) { auto& [request, future] = *it; - // Observe-only overall-deadline WARN. Emitted at most once per - // requestId (dedup'd via mTimedOutSenderIds) when the wall-clock - // elapsed time since transfer start exceeds the configured budget. - // No action is taken — this is purely diagnostic for the Tier-1 - // baseline. If the request later completes or errors, the entry is - // removed from mTimedOutSenderIds in the normal eviction paths - // below so a re-enqueued requestId can be warned about again. if (kvTransferTimeoutMs.has_value()) { auto elapsed = std::chrono::duration_cast( @@ -805,13 +790,7 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR " checkGenTransferStatus toCompleteIdSet size: %zu, atLeastRequestNum: %d ", toCompleteIdSet.size(), atLeastRequestNum.value_or(0)); } - // Observe-only overall-deadline check for the generation side. Mirrors - // the context-side logic in checkContextTransferStatus: emit a WARN at - // most once per requestId via the mTimedOutRequesterIds dedup set when - // the wall-clock elapsed time since transfer start exceeds the - // configured budget. No action is taken — eviction / state transitions - // / cancellation are intentionally left to higher layers in this - // Tier-1 baseline. + // Observe-only: gen-side mirror of the context-side timeout WARN. std::optional kvTransferTimeoutMs = std::nullopt; if (mCacheTransceiverConfig.has_value()) { diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp index c72090867f29..812323e92b81 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp @@ -254,6 +254,8 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses }; auto bufferEleSizes = getBufferSizeForTarget(); auto cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForSend(); + 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 +382,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); @@ -459,6 +461,7 @@ void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& s int deviceId = bufferManager.getStream().getDevice(); std::optional cacheBufferId = std::nullopt; + BufferIndexHolder recvHolder; if (common::getEnvTryZCopyForKVCacheTransfer() && destConfig.getParallelConfig().mPipelineParallelism @@ -495,6 +498,8 @@ void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& s { cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForRecv(); } + recvHolder + = BufferIndexHolder(*mCacheTransBufferManagers[transferIndexerKCache], cacheBufferId, /*isRecv=*/true); auto targetNum = pickUpConnections.size(); @@ -642,10 +647,7 @@ void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& s bufferManager.getStream().synchronize(); } - if (cacheBufferId.has_value()) - { - mCacheTransBufferManagers[transferIndexerKCache]->freeBufferIndexForRecv(cacheBufferId); - } + recvHolder.release(); } session.setTime(TransferSession::kTimePostprocess); diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 4badd200f56e..2b171d01d1ab 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1884,25 +1884,11 @@ def _executor_loop_pp(self): and req.py_disaggregated_params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST for req in self.active_requests) - # 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]. + # Rank-symmetric entry into _check_disagg_ctx_cache_transfer_status: + # every ctx rank must call this so the C++ gatherRequestIds Allgather + # inside it is reached symmetrically, else ABBA against the next + # unconditional collective. at_least_num is rank-local; only the + # entry must be symmetric. at_least_num = 0 if (num_fitting_reqs == 0 and not fitting_disagg_gen_init_requests @@ -2446,12 +2432,7 @@ 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) - # 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). + # Rank-symmetric entry; see _executor_loop_pp for rationale. at_least_num = 0 if (num_fitting_reqs == 0 and not fitting_disagg_gen_init_requests and not all_gen_first): @@ -3700,16 +3681,9 @@ def _schedule(self): @nvtx_range("_check_disagg_gen_transfer_status") def _check_disagg_gen_transfer_status(self): - # 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. + # Entered unconditionally so the C++ gatherRequestIds Allgather inside + # checkGenTransferStatus is reached by every gen rank. at_least_num is + # rank-local; only the entry must be symmetric. non_gen_first_reqs = [ req for req in self.active_requests if req.py_disaggregated_params and req.py_disaggregated_params. @@ -3871,15 +3845,9 @@ def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests): resource_mgr_type].prepare_resources( disagg_gen_init_to_prepare) - # 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. + # Unconditional call: every gen rank must reach _recv_disagg_gen_cache + # so the downstream gatherRequestIds Allgather is entered symmetrically. + # Empty-input case is a no-op for the receive work. self._recv_disagg_gen_cache(fitting_disagg_gen_init_requests) @nvtx_range("_prepare_disagg_gen_transmission_complete") From 5dfa4e9b3eaea00200b91f4e07999c041036d71e Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 29 May 2026 19:30:57 -0700 Subject: [PATCH 5/9] [https://nvbugs/6104831][fix] Revert A9 + A10 rank-symmetric collective entry Restores the pre-A9/A10 behavior in _check_disagg_gen_transfer_status, _prepare_disagg_gen_init, _executor_loop_pp, and _executor_loop to match upstream/main. Reduces this PR's delta vs main and isolates the rank-symmetric collective-entry changes for separate verification. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 60 +++++++++---------- 1 file changed, 28 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 2b171d01d1ab..9829e9027eee 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1884,20 +1884,15 @@ def _executor_loop_pp(self): and req.py_disaggregated_params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST for req in self.active_requests) - # Rank-symmetric entry into _check_disagg_ctx_cache_transfer_status: - # every ctx rank must call this so the C++ gatherRequestIds Allgather - # inside it is reached symmetrically, else ABBA against the next - # unconditional collective. at_least_num is rank-local; only the - # entry must be symmetric. - 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) + 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( + ): + self._check_disagg_ctx_cache_transfer_status(0) self.num_scheduled_requests = scheduled_batch.batch_size @@ -2432,15 +2427,14 @@ 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) - # Rank-symmetric entry; see _executor_loop_pp for rationale. - 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) + 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(): + self._check_disagg_ctx_cache_transfer_status(0) # In gen-only benchmark mode, all requests must fit in KV cache # simultaneously because generation requests do not release their @@ -3681,9 +3675,11 @@ def _schedule(self): @nvtx_range("_check_disagg_gen_transfer_status") def _check_disagg_gen_transfer_status(self): - # Entered unconditionally so the C++ gatherRequestIds Allgather inside - # checkGenTransferStatus is reached by every gen rank. at_least_num is - # rank-local; only the entry must be symmetric. + + need_check = any([ + req.is_disagg_generation_transmission_in_progress + for req in self.active_requests + ]) non_gen_first_reqs = [ req for req in self.active_requests if req.py_disaggregated_params and req.py_disaggregated_params. @@ -3692,8 +3688,10 @@ 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) - at_least_num = 1 if need_check_one else 0 - self._check_disagg_gen_cache_transfer_status(at_least_num) + + if need_check: + at_least_num = 1 if need_check_one else 0 + self._check_disagg_gen_cache_transfer_status(at_least_num) return @@ -3845,10 +3843,8 @@ def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests): resource_mgr_type].prepare_resources( disagg_gen_init_to_prepare) - # Unconditional call: every gen rank must reach _recv_disagg_gen_cache - # so the downstream gatherRequestIds Allgather is entered symmetrically. - # Empty-input case is a no-op for the receive work. - self._recv_disagg_gen_cache(fitting_disagg_gen_init_requests) + # Trigger KV cache exchange for new disagg_gen_init_requests + 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): From fa5991793dd4fb66813970d422ac126d9f0dabef Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Sat, 30 May 2026 21:54:13 -0700 Subject: [PATCH 6/9] [https://nvbugs/6104831][test] Add unit test coverage Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../unit_tests/batch_manager/CMakeLists.txt | 1 + .../batch_manager/bufferIndexHolderTest.cpp | 291 ++++++++++++++++++ 2 files changed, 292 insertions(+) create mode 100644 cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index 02e293b8d102..2ee994f76226 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -17,6 +17,7 @@ add_gtest(radixTreeTest radixTreeTest.cpp) add_gtest(blockKeyTest blockKeyTest.cpp) add_gtest(radixBlockTreeTest radixBlockTreeTest.cpp) add_gtest(cacheTransBufferTest cacheTransBufferTest.cpp) +add_gtest(bufferIndexHolderTest bufferIndexHolderTest.cpp) add_gtest(capacitySchedulerTest capacitySchedulerTest.cpp) add_gtest(contextProgressTest contextProgressTest.cu) add_gtest(evictionPolicyTest evictionPolicyTest.cpp) diff --git a/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp b/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp new file mode 100644 index 000000000000..d4d561f55d07 --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp @@ -0,0 +1,291 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * 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. + */ + +#include "tensorrt_llm/batch_manager/baseTransBuffer.h" +#include "tensorrt_llm/batch_manager/cacheTransBuffer.h" +#include "tensorrt_llm/batch_manager/kvCacheManager.h" +#include +#include +#include +#include +#include +#include + +using namespace tensorrt_llm::batch_manager; +using namespace tensorrt_llm::batch_manager::kv_cache_manager; +using namespace tensorrt_llm::runtime; + +namespace +{ + +// Subclass that exposes concurrence counters so tests can assert how many +// slots are currently held without relying on indirect observable effects. +class ObservableTransBufferManager : public CacheTransBufferManager +{ +public: + using CacheTransBufferManager::CacheTransBufferManager; + + [[nodiscard]] int sendInUse() const + { + return mConcurrenceSendResource.mConcurrence.load(); + } + + [[nodiscard]] int recvInUse() const + { + return mConcurrenceRecvResource.mConcurrence.load(); + } +}; + +} // namespace + +enum class Side +{ + Send, + Recv +}; + +struct HolderCase +{ + std::string name; + Side side; +}; + +class BufferIndexHolderLifecycleTest : public ::testing::TestWithParam +{ +protected: + void SetUp() override + { + setenv("TRTLLM_USE_UCX_KVCACHE", "1", 1); + + int constexpr numLayers = 2; + int constexpr numHeads = 2; + int constexpr sizePerHead = 8; + int constexpr tokensPerBlock = 4; + SizeType32 constexpr maxBlocksPerSeq = 4; + SizeType32 constexpr maxBeamWidth = 1; + SizeType32 constexpr maxNumSequences = 4; + SizeType32 constexpr sinkTokenLength = 0; + auto stream = std::make_shared(); + auto const kvMaxNumTokens = tokensPerBlock * maxBlocksPerSeq; + auto const totalNumBlocks = maxNumSequences * maxBlocksPerSeq; + using BlocksPerWindow = std::map>; + auto const blocksPerWindow = BlocksPerWindow{{kvMaxNumTokens, {totalNumBlocks, 0}}}; + + mKv = std::make_unique(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, + maxNumSequences, maxBeamWidth, std::vector{kvMaxNumTokens}, + nvinfer1::DataType::kFLOAT, sinkTokenLength, stream, kvMaxNumTokens, kvMaxNumTokens, + /*enableBlockReuse=*/true, CacheType::kSELF, std::nullopt, nullptr, true); + mKv->allocatePools(false); + mTrans = std::make_unique(mKv.get(), std::optional{kvMaxNumTokens}); + } + + void TearDown() override + { + mTrans.reset(); + mKv.reset(); + } + + [[nodiscard]] std::optional acquire() const + { + return GetParam().side == Side::Send ? mTrans->assignBufferIndexForSend() : mTrans->assignBufferIndexForRecv(); + } + + [[nodiscard]] int inUse() const + { + return GetParam().side == Side::Send ? mTrans->sendInUse() : mTrans->recvInUse(); + } + + [[nodiscard]] bool isRecv() const + { + return GetParam().side == Side::Recv; + } + + [[nodiscard]] ObservableTransBufferManager& mgr() const + { + return *mTrans; + } + + std::unique_ptr mKv; + std::unique_ptr mTrans; +}; + +// Default-constructed holder owns nothing; destruction is a no-op. +TEST_P(BufferIndexHolderLifecycleTest, DefaultConstructedHolderReleasesNothing) +{ + int const before = inUse(); + { + BufferIndexHolder holder; + EXPECT_FALSE(holder.held()); + EXPECT_FALSE(holder.index().has_value()); + } + EXPECT_EQ(inUse(), before); +} + +// Explicit release() on a default-constructed holder (no manager bound) +// must not dereference the null manager pointer. +TEST_P(BufferIndexHolderLifecycleTest, DefaultConstructedExplicitReleaseIsNoOp) +{ + int const before = inUse(); + BufferIndexHolder holder; + holder.release(); + EXPECT_FALSE(holder.held()); + EXPECT_EQ(inUse(), before); +} + +// Holder built from a nullopt index is disarmed (mHeld == false). +TEST_P(BufferIndexHolderLifecycleTest, NulloptIndexReleasesNothing) +{ + int const before = inUse(); + { + BufferIndexHolder holder{mgr(), std::nullopt, isRecv()}; + EXPECT_FALSE(holder.held()); + } + EXPECT_EQ(inUse(), before); +} + +// RAII: a held slot is released when the holder goes out of scope. +TEST_P(BufferIndexHolderLifecycleTest, ValidIndexReleasedOnDestruction) +{ + int const before = inUse(); + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + EXPECT_EQ(inUse(), before + 1); + { + BufferIndexHolder holder{mgr(), idx, isRecv()}; + EXPECT_TRUE(holder.held()); + EXPECT_EQ(holder.index(), idx); + } + EXPECT_EQ(inUse(), before); +} + +// release() frees the slot immediately and disarms the destructor. +TEST_P(BufferIndexHolderLifecycleTest, ExplicitReleaseFreesSlotEagerly) +{ + int const before = inUse(); + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + BufferIndexHolder holder{mgr(), idx, isRecv()}; + EXPECT_EQ(inUse(), before + 1); + holder.release(); + EXPECT_EQ(inUse(), before); + EXPECT_FALSE(holder.held()); +} + +// release() is idempotent: a second call is a no-op. +TEST_P(BufferIndexHolderLifecycleTest, DoubleReleaseIsSafe) +{ + int const before = inUse(); + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + BufferIndexHolder holder{mgr(), idx, isRecv()}; + holder.release(); + EXPECT_EQ(inUse(), before); + holder.release(); + EXPECT_EQ(inUse(), before); +} + +// detach() returns the index and disarms the destructor. The caller is +// responsible for freeing the slot — we free it manually to keep the pool +// balanced for subsequent tests. +TEST_P(BufferIndexHolderLifecycleTest, DetachDisarmsDestructor) +{ + int const before = inUse(); + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + { + BufferIndexHolder holder{mgr(), idx, isRecv()}; + auto detached = holder.detach(); + EXPECT_EQ(detached, idx); + EXPECT_FALSE(holder.held()); + } + // Destructor was disarmed: slot still in use. + EXPECT_EQ(inUse(), before + 1); + + // Manually free so the harness is balanced. + if (isRecv()) + { + mgr().freeBufferIndexForRecv(idx); + } + else + { + mgr().freeBufferIndexForSend(idx); + } + EXPECT_EQ(inUse(), before); +} + +// Move construction transfers ownership; the moved-from holder is disarmed. +TEST_P(BufferIndexHolderLifecycleTest, MoveConstructTransfersOwnership) +{ + int const before = inUse(); + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + { + BufferIndexHolder source{mgr(), idx, isRecv()}; + BufferIndexHolder sink{std::move(source)}; + EXPECT_FALSE(source.held()); + EXPECT_TRUE(sink.held()); + EXPECT_EQ(sink.index(), idx); + EXPECT_EQ(inUse(), before + 1); + } + // Sink released on scope exit; source was disarmed. + EXPECT_EQ(inUse(), before); +} + +// Move assignment releases the prior holding before taking new ownership. +TEST_P(BufferIndexHolderLifecycleTest, MoveAssignReleasesPriorThenTransfers) +{ + int const before = inUse(); + auto firstIdx = acquire(); + auto secondIdx = acquire(); + ASSERT_TRUE(firstIdx.has_value()); + ASSERT_TRUE(secondIdx.has_value()); + EXPECT_EQ(inUse(), before + 2); + { + BufferIndexHolder sink{mgr(), firstIdx, isRecv()}; + BufferIndexHolder source{mgr(), secondIdx, isRecv()}; + sink = std::move(source); + // Sink now owns secondIdx; firstIdx was released by the move-assign. + EXPECT_EQ(inUse(), before + 1); + EXPECT_EQ(sink.index(), secondIdx); + EXPECT_FALSE(source.held()); + } + EXPECT_EQ(inUse(), before); +} + +// Exception unwind through a scope containing a held holder still releases +// the slot. +TEST_P(BufferIndexHolderLifecycleTest, ExceptionUnwindStillReleases) +{ + int const before = inUse(); + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + try + { + BufferIndexHolder holder{mgr(), idx, isRecv()}; + EXPECT_EQ(inUse(), before + 1); + throw std::runtime_error("forced unwind"); + } + catch (std::runtime_error const&) + { + // Holder destructor ran during unwind. + } + EXPECT_EQ(inUse(), before); +} + +INSTANTIATE_TEST_SUITE_P(SideVariants, BufferIndexHolderLifecycleTest, + ::testing::Values(HolderCase{"send", Side::Send}, HolderCase{"recv", Side::Recv}), + [](::testing::TestParamInfo const& info) { return info.param.name; }); From 3e879f5c3a4c1517626f756de3f08e2e5ec613f6 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Sat, 30 May 2026 22:28:19 -0700 Subject: [PATCH 7/9] [https://nvbugs/6104831][test] Add unit test coverage Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../others/test_kv_cache_transceiver.py | 83 ++++++++++++++++++- 1 file changed, 81 insertions(+), 2 deletions(-) diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index cf9d544a9ce0..9576f4076406 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -1,5 +1,7 @@ +import gc import time import uuid +import weakref import pytest import torch @@ -229,12 +231,89 @@ def test_cancel_request_in_transmission(attention_type): assert gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR +@pytest.mark.timeout(120) +@pytest.mark.parametrize("side", ["sender", "receiver"], + ids=["sender", "receiver"]) +def test_async_transfer_keeps_llm_request_alive(side): + """Async transfer entry points keep LlmRequest alive while in flight. + + The futures vector on each side holds a strong shared_ptr; + if a future regression switched back to storing a raw pointer the + Python wrapper would be GC'd and the weakref below would expire + prematurely. + """ + 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="DEFAULT", + max_tokens_in_buffer=512) + transceiver_ctx = create_kv_cache_transceiver(mapping, dist, + kv_cache_manager_ctx, + AttentionTypeCpp.DEFAULT, + cache_transceiver_config) + transceiver_gen = create_kv_cache_transceiver(mapping, dist, + kv_cache_manager_gen, + AttentionTypeCpp.DEFAULT, + cache_transceiver_config) + + fill_kv_cache_buffer(kv_cache_manager_ctx) + + sampling_params = SamplingParams() + ctx_request = LlmRequest( + request_id=0, + max_new_tokens=1, + input_tokens=list(range(256)), + sampling_config=tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config()), + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) + kv_cache_manager_ctx.impl.add_sequence_batch( + [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], [ctx_request]) + + 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]) + + ctx_ref = weakref.ref(ctx_request) + gen_ref = weakref.ref(gen_request) + + transceiver_ctx.respond_and_send_async(ctx_request) + transceiver_gen.request_and_receive_async(gen_request) + + # Drop the external Python references. After this, the only owners + # of the underlying C++ objects are the shared_ptr entries inside + # mSenderFutures / mRequesterFutures. + del ctx_request + del gen_request + gc.collect() + + target_ref = ctx_ref if side == "sender" else gen_ref + assert target_ref() is not None, ( + f"{side}-side LlmRequest was destroyed prematurely; the async " + f"transfer entry point must hold a strong shared_ptr while the " + f"future is in flight") + + # Drive the transfer to completion so the harness tears down cleanly. + transceiver_ctx.check_context_transfer_status(1) + transceiver_gen.check_gen_transfer_status(1) + + def create_hybrid_cache_manager(mapping, dtype, mamba_conv_dtype=torch.float16, mamba_ssm_dtype=torch.float16): - """ - Create a MambaHybridCacheManager for testing hybrid models. + """Create a MambaHybridCacheManager for testing hybrid models. + This manager handles both KV cache (attention layers) and Mamba cache (RNN layers). Args: From 149da450b0580f2fa0e8984bd2ff24e521bbeee2 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Sun, 31 May 2026 21:29:34 -0700 Subject: [PATCH 8/9] [https://nvbugs/6104831][fix] Add test coverage and cleanup. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/baseTransBuffer.h | 12 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 7 + .../test_lists/test-db/l0_sanity_check.yml | 3 + .../others/test_kv_cache_transceiver.py | 155 ++++++++++++++++-- 4 files changed, 152 insertions(+), 25 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h index bce94702e121..2cbf9f514bd7 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h @@ -56,18 +56,17 @@ class BufferIndexHolder public: BufferIndexHolder() = default; - BufferIndexHolder(BaseTransBufferManager& mgr, std::optional index, bool isRecv, - std::optional requestIdForLog = std::nullopt) noexcept + BufferIndexHolder(BaseTransBufferManager& mgr, std::optional index, bool isRecv) noexcept : mMgr(&mgr) , mIndex(index) , mHeld(index.has_value()) , mIsRecv(isRecv) - , mRequestIdForLog(requestIdForLog) { } - // release() is out-of-line in baseTransBuffer.cpp to avoid the include cycle - // with BaseTransBufferManager. + // Defined out-of-line in baseTransBuffer.cpp: release() calls + // BaseTransBufferManager methods, whose full definition appears later + // in this header. ~BufferIndexHolder() { release(); @@ -81,7 +80,6 @@ class BufferIndexHolder , mIndex(other.mIndex) , mHeld(other.mHeld) , mIsRecv(other.mIsRecv) - , mRequestIdForLog(other.mRequestIdForLog) { other.mHeld = false; } @@ -95,7 +93,6 @@ class BufferIndexHolder mIndex = other.mIndex; mHeld = other.mHeld; mIsRecv = other.mIsRecv; - mRequestIdForLog = other.mRequestIdForLog; other.mHeld = false; } return *this; @@ -128,7 +125,6 @@ class BufferIndexHolder std::optional mIndex{}; bool mHeld{false}; bool mIsRecv{true}; - std::optional mRequestIdForLog{}; }; /// @brief Base class for cache transfer buffer management. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 9829e9027eee..a893264d1ef1 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1892,6 +1892,8 @@ def _executor_loop_pp(self): 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) self.num_scheduled_requests = scheduled_batch.batch_size @@ -2434,6 +2436,11 @@ def _prepare_and_schedule_batch(self): ) 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) # In gen-only benchmark mode, all requests must fit in KV cache diff --git a/tests/integration/test_lists/test-db/l0_sanity_check.yml b/tests/integration/test_lists/test-db/l0_sanity_check.yml index 231e545f50d4..945a7d9460b9 100644 --- a/tests/integration/test_lists/test-db/l0_sanity_check.yml +++ b/tests/integration/test_lists/test-db/l0_sanity_check.yml @@ -34,6 +34,9 @@ l0_sanity_check: - unittest/others/test_kv_cache_transceiver.py::test_kv_cache_transceiver_single_process[UCX-mha-ctx_fp16_gen_fp16] - unittest/others/test_kv_cache_transceiver.py::test_cancel_request_in_transmission[mha] - unittest/others/test_kv_cache_transceiver.py::test_cancel_request_in_transmission[mla] + - unittest/others/test_kv_cache_transceiver.py::test_async_transfer_keeps_llm_request_alive + - unittest/others/test_kv_cache_transceiver.py::test_kv_transfer_timeout_warns_once_per_request + - unittest/others/test_kv_cache_transceiver.py::test_kv_transfer_timeout_silent_when_unset - condition: ranges: system_gpu_count: diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index 9576f4076406..908a7d503235 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -1,4 +1,5 @@ import gc +import sys import time import uuid import weakref @@ -232,15 +233,12 @@ def test_cancel_request_in_transmission(attention_type): @pytest.mark.timeout(120) -@pytest.mark.parametrize("side", ["sender", "receiver"], - ids=["sender", "receiver"]) -def test_async_transfer_keeps_llm_request_alive(side): - """Async transfer entry points keep LlmRequest alive while in flight. - - The futures vector on each side holds a strong shared_ptr; - if a future regression switched back to storing a raw pointer the - Python wrapper would be GC'd and the weakref below would expire - prematurely. +def test_async_transfer_keeps_llm_request_alive(): + """Async entry points must hold a strong shared_ptr to the LlmRequest. + + A regression to raw-pointer storage in mSenderFutures / + mRequesterFutures lets Python GC the wrapper while a C++ status check + still dereferences it. """ mapping = Mapping(world_size=1, rank=0) dist = Distributed.get(mapping) @@ -284,30 +282,153 @@ def test_async_transfer_keeps_llm_request_alive(side): kv_cache_manager_gen.impl.add_sequence_batch( [(gen_request.py_request_id, gen_request.prompt_len, 1)], [gen_request]) + # Pin the reference graph *before* submission. add_sequence_batch's + # nanobind binding takes std::reference_wrapper, not a + # shared_ptr, so it does not retain a Python-side ref. Snapshotting + # the baseline here lets us assert exactly +1 after each async + # submission and catches regressions that don't capture a shared_ptr. ctx_ref = weakref.ref(ctx_request) gen_ref = weakref.ref(gen_request) + baseline_ctx_refcount = sys.getrefcount(ctx_request) + baseline_gen_refcount = sys.getrefcount(gen_request) transceiver_ctx.respond_and_send_async(ctx_request) transceiver_gen.request_and_receive_async(gen_request) - # Drop the external Python references. After this, the only owners - # of the underlying C++ objects are the shared_ptr entries inside - # mSenderFutures / mRequesterFutures. + # The nanobind binding for respond_and_send_async / request_and_receive_async + # takes std::shared_ptr; nanobind binds Python refcount into the + # shared_ptr, so post-submission each request has +1 ref held by + # mSenderFutures / mRequesterFutures. A regression to raw-pointer storage + # would not increment the refcount. + assert sys.getrefcount(ctx_request) == baseline_ctx_refcount + 1, ( + f"respond_and_send_async did not capture a shared_ptr; " + f"refcount={sys.getrefcount(ctx_request)} " + f"(expected baseline {baseline_ctx_refcount} + 1)") + assert sys.getrefcount(gen_request) == baseline_gen_refcount + 1, ( + f"request_and_receive_async did not capture a shared_ptr; " + f"refcount={sys.getrefcount(gen_request)} " + f"(expected baseline {baseline_gen_refcount} + 1)") + + # Drop the external Python refs. After this, the only owners are the + # nanobind-managed shared_ptr entries inside mSenderFutures / + # mRequesterFutures. del ctx_request del gen_request gc.collect() - target_ref = ctx_ref if side == "sender" else gen_ref - assert target_ref() is not None, ( - f"{side}-side LlmRequest was destroyed prematurely; the async " - f"transfer entry point must hold a strong shared_ptr while the " - f"future is in flight") + assert ctx_ref() is not None, ( + "Sender-side LlmRequest was destroyed prematurely; " + "respond_and_send_async must hold a strong shared_ptr while the " + "sender future is in flight") + assert gen_ref() is not None, ( + "Receiver-side LlmRequest was destroyed prematurely; " + "request_and_receive_async must hold a strong shared_ptr while " + "the requester future is in flight") # Drive the transfer to completion so the harness tears down cleanly. transceiver_ctx.check_context_transfer_status(1) transceiver_gen.check_gen_transfer_status(1) +def _build_ctx_request_for_timeout_test(request_id: int) -> LlmRequest: + sampling_params = SamplingParams() + return LlmRequest( + request_id=request_id, + 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) + + +@pytest.mark.timeout(60) +def test_kv_transfer_timeout_warns_once_per_request(capfd): + """Observe-only timeout WARN must fire exactly once per stuck request. + + checkContextTransferStatus emits a TLLM_LOG_WARNING when elapsed time + exceeds kv_transfer_timeout_ms; mTimedOutSenderIds dedup suppresses + repeat emissions across subsequent polls of the same in-flight future. + """ + mapping = Mapping(world_size=1, rank=0) + dist = Distributed.get(mapping) + kv_cache_manager_ctx = create_kv_cache_manager(mapping, DataType.HALF) + + cache_transceiver_config = CacheTransceiverConfig( + backend="DEFAULT", max_tokens_in_buffer=512, kv_transfer_timeout_ms=100) + transceiver_ctx = create_kv_cache_transceiver(mapping, dist, + kv_cache_manager_ctx, + AttentionTypeCpp.DEFAULT, + cache_transceiver_config) + + fill_kv_cache_buffer(kv_cache_manager_ctx) + + ctx_request = _build_ctx_request_for_timeout_test(request_id=42) + kv_cache_manager_ctx.impl.add_sequence_batch( + [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], [ctx_request]) + + capfd.readouterr() # drain prior noise + + transceiver_ctx.respond_and_send_async(ctx_request) + time.sleep(0.3) # > kv_transfer_timeout_ms + + transceiver_ctx.check_context_transfer_status(0) + first = capfd.readouterr() + marker = "Context KV cache transfer for request 42 exceeded configured timeout" + assert marker in first.err, ( + f"Expected observe-only WARN after timeout; stderr:\n{first.err}") + assert first.err.count(marker) == 1, ( + f"WARN should fire exactly once per poll for a single stuck request; " + f"got {first.err.count(marker)} emissions:\n{first.err}") + + transceiver_ctx.check_context_transfer_status(0) + second = capfd.readouterr() + assert marker not in second.err, ( + f"WARN re-emitted on a subsequent poll; dedup broken. " + f"stderr:\n{second.err}") + + transceiver_ctx.cancel_request(ctx_request) + + +@pytest.mark.timeout(60) +def test_kv_transfer_timeout_silent_when_unset(capfd): + """Without kv_transfer_timeout_ms the observe-only WARN must stay silent. + + The elapsed-time check is gated by the optional config field; absence + of the field must short-circuit the WARN path even on a long-running + transfer. + """ + mapping = Mapping(world_size=1, rank=0) + dist = Distributed.get(mapping) + kv_cache_manager_ctx = create_kv_cache_manager(mapping, DataType.HALF) + + cache_transceiver_config = CacheTransceiverConfig(backend="DEFAULT", + max_tokens_in_buffer=512) + transceiver_ctx = create_kv_cache_transceiver(mapping, dist, + kv_cache_manager_ctx, + AttentionTypeCpp.DEFAULT, + cache_transceiver_config) + + fill_kv_cache_buffer(kv_cache_manager_ctx) + + ctx_request = _build_ctx_request_for_timeout_test(request_id=99) + kv_cache_manager_ctx.impl.add_sequence_batch( + [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], [ctx_request]) + + capfd.readouterr() + + transceiver_ctx.respond_and_send_async(ctx_request) + time.sleep(0.3) + transceiver_ctx.check_context_transfer_status(0) + + out = capfd.readouterr() + assert "exceeded configured timeout" not in out.err, ( + f"Observe-only WARN must not fire when kv_transfer_timeout_ms is " + f"unset; stderr:\n{out.err}") + + transceiver_ctx.cancel_request(ctx_request) + + def create_hybrid_cache_manager(mapping, dtype, mamba_conv_dtype=torch.float16, From 21e4890277c8da8f5c541f208a7c5e3886eb9bb5 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 1 Jun 2026 09:46:13 -0700 Subject: [PATCH 9/9] [https://nvbugs/6104831][fix] Fix test order and stream-check bugs in new unit tests Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../others/test_kv_cache_transceiver.py | 51 +++++++++---------- 1 file changed, 25 insertions(+), 26 deletions(-) diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index 908a7d503235..9eed2dea5e6f 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -270,6 +270,21 @@ def test_async_transfer_keeps_llm_request_alive(): kv_cache_manager_ctx.impl.add_sequence_batch( [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], [ctx_request]) + # Snapshot ctx refcount *before* submission. add_sequence_batch takes + # std::reference_wrapper (no Python ref retained), so the + # baseline here is clean and the +1 below isolates exactly the + # shared_ptr captured by respond_and_send_async. + ctx_ref = weakref.ref(ctx_request) + baseline_ctx_refcount = sys.getrefcount(ctx_request) + + # respond_and_send_async also populates ctx_request.context_phase_params + # as a side effect — gen_request below requires this to be non-empty. + transceiver_ctx.respond_and_send_async(ctx_request) + assert sys.getrefcount(ctx_request) == baseline_ctx_refcount + 1, ( + f"respond_and_send_async did not capture a shared_ptr; " + f"refcount={sys.getrefcount(ctx_request)} " + f"(expected baseline {baseline_ctx_refcount} + 1)") + gen_request = LlmRequest( request_id=0, max_new_tokens=1, @@ -282,28 +297,10 @@ def test_async_transfer_keeps_llm_request_alive(): kv_cache_manager_gen.impl.add_sequence_batch( [(gen_request.py_request_id, gen_request.prompt_len, 1)], [gen_request]) - # Pin the reference graph *before* submission. add_sequence_batch's - # nanobind binding takes std::reference_wrapper, not a - # shared_ptr, so it does not retain a Python-side ref. Snapshotting - # the baseline here lets us assert exactly +1 after each async - # submission and catches regressions that don't capture a shared_ptr. - ctx_ref = weakref.ref(ctx_request) gen_ref = weakref.ref(gen_request) - baseline_ctx_refcount = sys.getrefcount(ctx_request) baseline_gen_refcount = sys.getrefcount(gen_request) - transceiver_ctx.respond_and_send_async(ctx_request) transceiver_gen.request_and_receive_async(gen_request) - - # The nanobind binding for respond_and_send_async / request_and_receive_async - # takes std::shared_ptr; nanobind binds Python refcount into the - # shared_ptr, so post-submission each request has +1 ref held by - # mSenderFutures / mRequesterFutures. A regression to raw-pointer storage - # would not increment the refcount. - assert sys.getrefcount(ctx_request) == baseline_ctx_refcount + 1, ( - f"respond_and_send_async did not capture a shared_ptr; " - f"refcount={sys.getrefcount(ctx_request)} " - f"(expected baseline {baseline_ctx_refcount} + 1)") assert sys.getrefcount(gen_request) == baseline_gen_refcount + 1, ( f"request_and_receive_async did not capture a shared_ptr; " f"refcount={sys.getrefcount(gen_request)} " @@ -375,17 +372,18 @@ def test_kv_transfer_timeout_warns_once_per_request(capfd): transceiver_ctx.check_context_transfer_status(0) first = capfd.readouterr() marker = "Context KV cache transfer for request 42 exceeded configured timeout" - assert marker in first.err, ( - f"Expected observe-only WARN after timeout; stderr:\n{first.err}") - assert first.err.count(marker) == 1, ( + # TLLM_LOG_WARNING writes to stdout; check first.out (not first.err). + assert marker in first.out, ( + f"Expected observe-only WARN after timeout; stdout:\n{first.out}") + assert first.out.count(marker) == 1, ( f"WARN should fire exactly once per poll for a single stuck request; " - f"got {first.err.count(marker)} emissions:\n{first.err}") + f"got {first.out.count(marker)} emissions:\n{first.out}") transceiver_ctx.check_context_transfer_status(0) second = capfd.readouterr() - assert marker not in second.err, ( + assert marker not in second.out, ( f"WARN re-emitted on a subsequent poll; dedup broken. " - f"stderr:\n{second.err}") + f"stdout:\n{second.out}") transceiver_ctx.cancel_request(ctx_request) @@ -422,9 +420,10 @@ def test_kv_transfer_timeout_silent_when_unset(capfd): transceiver_ctx.check_context_transfer_status(0) out = capfd.readouterr() - assert "exceeded configured timeout" not in out.err, ( + # TLLM_LOG_WARNING writes to stdout; check out.out (not out.err). + assert "exceeded configured timeout" not in out.out, ( f"Observe-only WARN must not fire when kv_transfer_timeout_ms is " - f"unset; stderr:\n{out.err}") + f"unset; stdout:\n{out.out}") transceiver_ctx.cancel_request(ctx_request)