From 7f4e30af890da83c62ec4289a1c9140732073555 Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Tue, 5 May 2026 20:55:43 -0700 Subject: [PATCH 1/7] [https://nvbugs/6104831][fix] Disagg KV transfer: non-blocking polling, structured cancel, bounded quarantine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The C++ cache transceiver and Python executor cooperate on disaggregated KV transfer cancellation, timeout, and process health. Earlier attempts (PR #13706, PR #13713, PR #13728) each fixed one slice of the problem but still hung in the field: * PR #13706 with Python changes restarted workers on every routine per-request timeout (Python fail-close was too aggressive). * PR #13713 / PR #13728 hung even with Python cleanup removed because the C++ status polling could still call future.get() on an unready worker future, freezing the executor event loop while NIXL/UCX was wedged. This PR keeps each layer's responsibility separate, as described in docs/source/features/disagg-kv-transfer-hang-restart-analysis.md and the new docs/source/features/disagg-kv-transfer-session-lifecycle.md: * checkContextTransferStatus / checkGenTransferStatus are now strictly non-blocking. They poll with wait_for(0ms) and only call future.get() when the future is already ready. atLeastRequestNum > 0 still admits additional ready entries but never selects an unready one to satisfy the count. drainContextTransferStatus / drainGenTransferStatus are the only blocking variants and are intended for shutdown drain. * cancelRequestStructured returns a TransferCancelResult enum with six outcomes (NotFound, AlreadyComplete, CancelledBeforeAdvertise, CancelRequestedInFlight, BackendUnhealthy, NotCancellable). Only the first three permit Python to free request resources; the others keep the request owned by C++ until the worker future reaches a final state. The historical bool cancelRequest is preserved as a backward-compatible wrapper. * The transceiver maintains a bounded quarantine counter and a global progress deadline. A per-request timeout marks the entry quarantined but keeps the future pinned so NIXL/UCX cannot write into freed memory. If quarantined entries exceed mQuarantineBudget or no worker has reached a final state for longer than mGlobalProgressDeadlineMs, isHealthy() flips false and getHealth() surfaces the snapshot for orchestration. * PyExecutor adds _can_terminate_request_now / _inflight_cancel_- requested_ids so _do_terminate_request defers freeing resources for any request that is still in disagg transfer state or whose cancel is in flight. Per-request timeouts no longer clear active_requests, the waiting queue, or set is_shutdown — that was the PR #13706 restart-loop trigger and is explicitly forbidden by the analysis doc. * The ADP synchronized pending-response flush from PR #13112 is ported here (the base commit precedes that merge). Transfer- completion responses created in _end_transfer_and_maybe_terminate are buffered and flushed at synchronised loop points so every DP rank participates in the tp_gather collective. Test coverage: * tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py — 14 unit tests for the deferred-terminate guard, the structured- cancel decision tree, and the ADP flush symmetry. They run without GPU or MPI so they fail fast in pre-merge CI. References: * analysis: docs/source/features/disagg-kv-transfer-hang-restart-analysis.md * contract: docs/source/features/disagg-kv-transfer-session-lifecycle.md Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 179 +++++- .../batch_manager/cacheTransceiver.cpp | 518 ++++++++++++------ .../batch_manager/cacheTransceiver.cpp | 77 ++- ...isagg-kv-transfer-hang-restart-analysis.md | 281 ++++++++++ .../disagg-kv-transfer-session-lifecycle.md | 122 +++++ .../_torch/pyexecutor/kv_cache_transceiver.py | 41 ++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 159 +++++- .../test_kv_transfer_session_lifecycle.py | 348 ++++++++++++ 8 files changed, 1532 insertions(+), 193 deletions(-) create mode 100644 docs/source/features/disagg-kv-transfer-hang-restart-analysis.md create mode 100644 docs/source/features/disagg-kv-transfer-session-lifecycle.md create mode 100644 tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 8f8330603893..f4abcc149e4d 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -200,6 +200,60 @@ struct RequestStatuses std::unordered_set errorRequestIds; }; +/// Structured outcome of a cancellation request. Replaces the historical +/// boolean return that conflated "Python may free request resources" with +/// "C++ has nothing to do." Only @c kCancelledBeforeAdvertise and +/// @c kAlreadyComplete imply that downstream cleanup may proceed +/// immediately. The other states require Python to leave the request owned +/// by C++ until a final transfer state is observed (or the transceiver +/// reports unhealthy). +enum class TransferCancelResult : uint8_t +{ + /// Request is unknown to the transceiver (already cleaned up or never + /// registered). + kNotFound = 0, + /// Request has already reached worker-final state; the future is ready + /// and any allocated buffers have been released. + kAlreadyComplete = 1, + /// Request was queued but no transfer buffer/handle was advertised to + /// a peer yet. Memory is safe to release immediately. + kCancelledBeforeAdvertise = 2, + /// Cancellation was accepted but the worker is still in flight. The + /// transceiver retains ownership of the request and its future until a + /// final state is observed; Python must not free request resources. + kCancelRequestedInFlight = 3, + /// The transceiver is unhealthy (quarantine budget exceeded or no + /// backend progress past the global deadline). Orchestration should + /// restart the worker. + kBackendUnhealthy = 4, + /// Cancellation cannot be performed at this point (e.g., the sender is + /// in the middle of advertising the buffer to a peer). Retry later. + kNotCancellable = 5, +}; + +/// Health snapshot exposed for observability. All counters are best-effort +/// approximations updated on the executor thread; they are not real-time +/// kernel-level metrics. +struct TransceiverHealth +{ + bool isHealthy{true}; + /// Number of in-flight transfers that have exceeded their per-request + /// timeout but whose worker future has not yet reached a final state. + /// These transfers are tracked but their request resources are kept + /// pinned by C++ — Python must not free them. + size_t quarantinedTransferCount{0}; + /// Maximum number of quarantined transfers permitted before the + /// transceiver is flipped to unhealthy. + size_t quarantineBudget{0}; + /// Wall-clock seconds since the last observed worker future + /// transition. Only meaningful when the transceiver has tracked at + /// least one transfer. + double secondsSinceLastProgress{0.0}; + /// Wall-clock seconds beyond which "no progress" is treated as a + /// global backend wedge. + double globalProgressDeadlineSeconds{0.0}; +}; + class BaseCacheTransceiver { public: @@ -212,16 +266,66 @@ class BaseCacheTransceiver virtual void requestAndReceiveSync(LlmRequest* llmRequest) = 0; virtual void requestAndReceiveAsync(LlmRequest* llmRequest) = 0; - /// Check all requests transferring context, and return the requests that have completed or encountered an error. + /// Non-blocking poll. Returns the requests that have completed or + /// encountered an error since the last call. With + /// @c atLeastRequestNum unset the call defaults to a pure poll: it + /// must @b never block on an unready future on the executor thread. + /// Callers that genuinely need to drain (shutdown only) must use + /// @ref drainContextTransferStatus instead. virtual RequestStatuses checkContextTransferStatus( std::optional const& atLeastRequestNum = std::nullopt, bool markComplete = false) = 0; + /// Non-blocking poll. Same contract as @ref checkContextTransferStatus. virtual void checkGenTransferStatus(std::optional const& atLeastRequestNum = std::nullopt) = 0; + /// Blocking drain — @b only safe on dedicated drain/shutdown paths. + /// The default forwards to the polling variant for backward + /// compatibility; subclasses that own worker futures override this to + /// actually wait for completion. + virtual RequestStatuses drainContextTransferStatus(bool markComplete = false) + { + return checkContextTransferStatus(std::nullopt, markComplete); + } + + /// Blocking drain — @b only safe on dedicated drain/shutdown paths. + virtual void drainGenTransferStatus() + { + checkGenTransferStatus(std::nullopt); + } + [[nodiscard]] virtual bool checkGenTransferComplete() const = 0; + /// Structured cancellation. See @ref TransferCancelResult for the full + /// contract. Subclasses override this; the boolean wrapper below maps + /// the structured result onto the historical "is the request safe to + /// release now" bool. + virtual TransferCancelResult cancelRequestStructured(LlmRequest* llmRequest) + { + return cancelRequest(llmRequest) ? TransferCancelResult::kCancelledBeforeAdvertise + : TransferCancelResult::kNotFound; + } + + /// Backward-compatible wrapper: returns true only when Python is safe + /// to free the request's resources immediately (pre-advertise cancel + /// or already-complete worker). Callers that need to distinguish + /// in-flight cancellation from "nothing to do" should use + /// @ref cancelRequestStructured. virtual bool cancelRequest(LlmRequest* llmRequest) = 0; + + /// Whether the transceiver is currently healthy. Flips to false when + /// quarantined transfers exceed the budget or when no backend + /// progress has been observed past the global deadline. + [[nodiscard]] virtual bool isHealthy() const + { + return true; + } + + /// Health snapshot for orchestration / metrics. + [[nodiscard]] virtual TransceiverHealth getHealth() const + { + return {}; + } }; class CacheTransceiver : public BaseCacheTransceiver @@ -265,21 +369,90 @@ class CacheTransceiver : public BaseCacheTransceiver void checkGenTransferStatus(std::optional const& atLeastRequestNum = std::nullopt) override; + /// Blocking drain — only invoked from shutdown/teardown. + RequestStatuses drainContextTransferStatus(bool markComplete = false) override; + + /// Blocking drain — only invoked from shutdown/teardown. + void drainGenTransferStatus() override; + [[nodiscard]] bool checkGenTransferComplete() const override; + TransferCancelResult cancelRequestStructured(LlmRequest* llmRequest) override; + virtual bool cancelRequest(LlmRequest* llmRequest) override; + [[nodiscard]] bool isHealthy() const override; + + [[nodiscard]] TransceiverHealth getHealth() const override; + private: void initializeCommState(); void setContextState(LlmRequest* llmRequest); + /// Per-entry tracking for tracked worker futures. We keep the future + /// pinned until the worker reaches a final state, even after the + /// per-request deadline elapses, so that cancelling a still-running + /// worker never frees its KV/buffer resources prematurely. + struct TrackedFuture + { + LlmRequest* request; + std::future future; + std::chrono::steady_clock::time_point deadline; + /// True once the per-request timeout has fired and we have flipped + /// the request to an error state. The entry stays in the vector + /// (and the future stays pinned) until the worker actually + /// finishes — possibly forever if the backend is wedged, in which + /// case the global progress deadline raises an unhealthy signal. + bool quarantined{false}; + /// True once we have already advertised a buffer/handle to a + /// peer. Pre-advertise cancellation can release normally; post- + /// advertise cancellation must wait for worker quiescence. + bool advertised{false}; + }; + + /// Update the global "last progress" timestamp and re-evaluate the + /// transceiver health flag against the quarantine budget and the + /// global progress deadline. Caller must hold @ref mHealthMutex. + void updateHealthLocked(); + + /// Drop a TrackedFuture entry from a vector. Decrements the + /// quarantine counter if the entry was quarantined, marks progress. + void releaseTrackedFutureLocked(std::vector& vec, size_t index); + + /// Apply per-entry timeout policy: if the configured + /// kvTransferTimeoutMs has elapsed and the future is not yet ready, + /// flip the entry to quarantined and surface an error to the caller. + /// Returns true if the entry was newly quarantined. + bool maybeQuarantineLocked(TrackedFuture& entry, RequestStatuses* outStatus); + + /// Internal worker for the public polling and drain entry points. + /// @c allowBlocking is true only on shutdown drain — that is the one + /// path where waiting on a worker future is acceptable. + RequestStatuses checkContextTransferStatusImpl( + std::optional const& atLeastRequestNum, bool markComplete, bool allowBlocking); + void checkGenTransferStatusImpl(std::optional const& atLeastRequestNum, bool allowBlocking); + + /// Compute the per-entry deadline for a tracked transfer. + [[nodiscard]] std::chrono::steady_clock::time_point computeTrackedFutureDeadline( + std::chrono::steady_clock::time_point requestStart) const; + std::unique_ptr mCacheSender; std::unique_ptr mCacheReceiver; - std::vector>> mSenderFutures; - std::vector>> mRequesterFutures; + std::vector mSenderFutures; + std::vector mRequesterFutures; mpi::MpiComm const* mMpiWorldComm{nullptr}; + /// Health/quarantine accounting. The mutex protects only this small + /// block; the executor thread is the sole writer/reader of the future + /// vectors above. + mutable std::mutex mHealthMutex; + size_t mQuarantinedTransferCount{0}; + size_t mQuarantineBudget{16}; + std::chrono::steady_clock::time_point mLastProgressTime{std::chrono::steady_clock::now()}; + std::chrono::milliseconds mGlobalProgressDeadlineMs{60'000}; + bool mIsHealthy{true}; + std::shared_ptr mGroupComm; std::shared_ptr mGroupTensorParaComm, mGroupPipeParaComm, mGroupDataComm, mGroupTPInDPComm; diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 2e4bf1f06667..ac50654e21b2 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -339,7 +339,8 @@ void CacheTransceiver::respondAndSendAsync(LlmRequest* llmRequest) } setContextState(llmRequest); auto future = mCacheSender->sendAsync(*llmRequest); - mSenderFutures.emplace_back(llmRequest, std::move(future)); + auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); + mSenderFutures.push_back(TrackedFuture{llmRequest, std::move(future), deadline, false, false}); } void CacheTransceiver::respondAndSendLayerWise( @@ -355,7 +356,8 @@ void CacheTransceiver::respondAndSendLayerWise( llmRequest->setState(LlmRequestState::kDISAGG_CONTEXT_INIT_AND_TRANS); setContextState(llmRequest.get()); auto future = mCacheSender->sendAsync(*llmRequest); - mSenderFutures.emplace_back(llmRequest.get(), std::move(future)); + auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); + mSenderFutures.push_back(TrackedFuture{llmRequest.get(), std::move(future), deadline, false, false}); } } @@ -374,7 +376,7 @@ void CacheTransceiver::requestAndReceiveAsync(LlmRequest* llmRequest) TLLM_CHECK(llmRequest && llmRequest->isGenerationOnlyRequest()); if (std::find_if(mRequesterFutures.begin(), mRequesterFutures.end(), - [llmRequest](auto const& pair) { return pair.first->mRequestId == llmRequest->mRequestId; }) + [llmRequest](auto const& entry) { return entry.request->mRequestId == llmRequest->mRequestId; }) != mRequesterFutures.end()) { TLLM_LOG_WARNING("Request ID %zu is already in mRequestFutures.", llmRequest->mRequestId); @@ -382,7 +384,8 @@ void CacheTransceiver::requestAndReceiveAsync(LlmRequest* llmRequest) } auto future = mCacheReceiver->receiveAsync(*llmRequest); - mRequesterFutures.emplace_back(llmRequest, std::move(future)); + auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); + mRequesterFutures.push_back(TrackedFuture{llmRequest, std::move(future), deadline, false, false}); llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS); } @@ -480,29 +483,127 @@ void updateKVCacheTransferBW(std::shared_ptr const& mComm, } } -RequestStatuses CacheTransceiver::checkContextTransferStatus( - std::optional const& atLeastRequestNum, bool markComplete) +// ----- Internal helpers ---------------------------------------------------- + +std::chrono::steady_clock::time_point CacheTransceiver::computeTrackedFutureDeadline( + std::chrono::steady_clock::time_point requestStart) const +{ + // Returns time_point::max() when no per-request timeout is configured, + // so the polling path treats the entry as "never times out" — matching + // the legacy behaviour for callers that did not opt into timeouts. + if (!mCacheTransceiverConfig.has_value()) + { + return std::chrono::steady_clock::time_point::max(); + } + auto const timeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); + if (!timeoutMs.has_value()) + { + return std::chrono::steady_clock::time_point::max(); + } + return requestStart + std::chrono::milliseconds(timeoutMs.value()); +} + +void CacheTransceiver::updateHealthLocked() +{ + auto const now = std::chrono::steady_clock::now(); + auto const sinceProgress + = std::chrono::duration_cast(now - mLastProgressTime).count(); + bool wedged = sinceProgress > mGlobalProgressDeadlineMs.count() && (!mSenderFutures.empty() || !mRequesterFutures.empty()); + bool overBudget = mQuarantinedTransferCount > mQuarantineBudget; + bool nextHealthy = !overBudget && !wedged; + if (mIsHealthy && !nextHealthy) + { + TLLM_LOG_WARNING( + "CacheTransceiver flipping to UNHEALTHY: quarantined=%zu budget=%zu sinceProgressMs=%lld deadlineMs=%lld", + mQuarantinedTransferCount, mQuarantineBudget, static_cast(sinceProgress), + static_cast(mGlobalProgressDeadlineMs.count())); + } + mIsHealthy = nextHealthy; +} + +void CacheTransceiver::releaseTrackedFutureLocked(std::vector& vec, size_t index) +{ + TLLM_CHECK(index < vec.size()); + { + std::scoped_lock lk(mHealthMutex); + if (vec[index].quarantined && mQuarantinedTransferCount > 0) + { + --mQuarantinedTransferCount; + } + mLastProgressTime = std::chrono::steady_clock::now(); + updateHealthLocked(); + } + vec.erase(vec.begin() + index); +} + +bool CacheTransceiver::maybeQuarantineLocked(TrackedFuture& entry, RequestStatuses* outStatus) { - bool blockAll = !atLeastRequestNum.has_value(); - std::optional senderFutureTimeoutMs = std::nullopt; - // If blockAll is true, we want to block and not use a timeout - if (!blockAll && mCacheTransceiverConfig.has_value()) + if (entry.quarantined) { - senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); + return false; } + if (entry.deadline == std::chrono::steady_clock::time_point::max()) + { + return false; + } + auto const now = std::chrono::steady_clock::now(); + if (now < entry.deadline) + { + return false; + } + entry.quarantined = true; + entry.request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + if (outStatus != nullptr) + { + outStatus->errorRequestIds.insert(entry.request->mRequestId); + } + { + std::scoped_lock lk(mHealthMutex); + ++mQuarantinedTransferCount; + updateHealthLocked(); + } + TLLM_LOG_WARNING( + "Quarantining request %ld: per-request KV transfer deadline exceeded; future stays pinned in C++ until " + "worker reaches a final state.", + entry.request->mRequestId); + return true; +} + +// ----- Public API ---------------------------------------------------------- +RequestStatuses CacheTransceiver::checkContextTransferStatus( + std::optional const& atLeastRequestNum, bool markComplete) +{ + // Non-blocking poll. The historical "blockAll" semantics (no + // atLeastRequestNum, fall through to future.get on every entry) are + // gone: a wedged worker thread must never freeze the executor event + // loop. Use @ref drainContextTransferStatus for shutdown drain. + return checkContextTransferStatusImpl(atLeastRequestNum, markComplete, /*allowBlocking=*/false); +} + +RequestStatuses CacheTransceiver::drainContextTransferStatus(bool markComplete) +{ + // Blocking — only safe on shutdown/teardown paths. + return checkContextTransferStatusImpl(std::nullopt, markComplete, /*allowBlocking=*/true); +} + +RequestStatuses CacheTransceiver::checkContextTransferStatusImpl( + std::optional const& atLeastRequestNum, bool markComplete, bool allowBlocking) +{ auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupTPInDPComm : mGroupTensorParaComm; + + // Pass 1: which futures are READY right now? std::vector contextCompleteRequestIds; - for (auto&& [request, future] : mSenderFutures) + for (auto const& entry : mSenderFutures) { - if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + if (entry.future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) { - contextCompleteRequestIds.push_back(request->mRequestId); + contextCompleteRequestIds.push_back(entry.request->mRequestId); } } std::unordered_map frequencyMap; - if ((syncComm) && syncComm->getSize() > 1) + if (syncComm && syncComm->getSize() > 1) { auto gatherRequestIdVec = gatherRequestIds(syncComm, contextCompleteRequestIds); for (auto&& requestId : gatherRequestIdVec) @@ -517,101 +618,150 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( frequencyMap[requestId]++; } } - std::vector> freqVec(frequencyMap.begin(), frequencyMap.end()); + std::vector> freqVec(frequencyMap.begin(), frequencyMap.end()); std::sort(freqVec.begin(), freqVec.end(), - [](std::pair const& left, - std::pair const& right) { return left.second > right.second; }); + [](auto const& left, auto const& right) { return left.second > right.second; }); + + int const expectedFreq = syncComm ? syncComm->getSize() : 1; std::unordered_set toCompleteIdSet; - for (auto&& [requestId, freq] : freqVec) + for (auto const& [requestId, freq] : freqVec) { - if (freq == ((syncComm) ? syncComm->getSize() : 1)) + if (freq == expectedFreq) { toCompleteIdSet.insert(requestId); } } - // Make sure there are at least atLeastRequestNum requests in toCompleteIdSet. - // This will preserve the order of insertion for KVCache transfer requests. - for (auto it = mSenderFutures.begin(); - atLeastRequestNum.value_or(0) > static_cast(toCompleteIdSet.size()) && it != mSenderFutures.end(); ++it) + // For atLeastRequestNum > 0 we may need to admit additional requests. + // CRITICAL: prefer entries whose future is already ready. Never select + // an unready entry just to satisfy the count — that is the + // "future.get on selected-but-unready entry" wedge the analysis doc + // calls out. If we run out of ready entries we simply return what we + // have; the caller polls again next iteration. + if (atLeastRequestNum.has_value()) { - auto& [request, future] = *it; - toCompleteIdSet.insert(request->mRequestId); + for (auto const& entry : mSenderFutures) + { + if (static_cast(toCompleteIdSet.size()) >= atLeastRequestNum.value()) + { + break; + } + if (entry.future.wait_for(std::chrono::milliseconds(0)) != std::future_status::ready) + { + continue; + } + toCompleteIdSet.insert(entry.request->mRequestId); + } } RequestStatuses requestsStatus{}; - - // Complete all the requests in toCompleteIdSet - for (auto it = mSenderFutures.begin(); it != mSenderFutures.end();) + bool madeProgress = false; + for (size_t i = 0; i < mSenderFutures.size();) { - auto& [request, future] = *it; - if (blockAll || (toCompleteIdSet.find(request->mRequestId) != toCompleteIdSet.end())) + auto& entry = mSenderFutures[i]; + bool const selected = toCompleteIdSet.find(entry.request->mRequestId) != toCompleteIdSet.end(); + bool const isReady = entry.future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready; + + if (selected && isReady) { try { - // Wait for up to a specified timeout - auto status = future.wait_for(std::chrono::milliseconds(senderFutureTimeoutMs.value_or(0))); - if (status == std::future_status::ready || !senderFutureTimeoutMs.has_value()) - { - future.get(); - requestsStatus.completedRequestIds.insert(request->mRequestId); - if (markComplete) - { - request->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); - } - it = mSenderFutures.erase(it); - } - else if (status == std::future_status::timeout) - { - TLLM_LOG_WARNING("Timed out waiting for context KV cache transfer after %d milliseconds.", - senderFutureTimeoutMs.value()); - ++it; - } - else + entry.future.get(); + requestsStatus.completedRequestIds.insert(entry.request->mRequestId); + if (markComplete) { - TLLM_LOG_ERROR( - "Future returned unexpected status for request %ld. Marking as error", request->mRequestId); - - request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); - requestsStatus.errorRequestIds.insert(request->mRequestId); - it = mSenderFutures.erase(it); + entry.request->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); } } catch (std::exception const& e) { - TLLM_LOG_ERROR( - "Error occurred during context transfer for request %ld: %s", request->mRequestId, e.what()); - request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); - requestsStatus.errorRequestIds.insert(request->mRequestId); - it = mSenderFutures.erase(it); + TLLM_LOG_ERROR("Error occurred during context transfer for request %ld: %s", + entry.request->mRequestId, e.what()); + entry.request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + requestsStatus.errorRequestIds.insert(entry.request->mRequestId); } + releaseTrackedFutureLocked(mSenderFutures, i); + madeProgress = true; + continue; } - else + + if (allowBlocking) { - ++it; + // Drain path: shutdown is happening. Wait for the worker to + // finish so we can release resources cleanly. Note this still + // never blocks on the executor thread — only shutdown calls + // this codepath. + try + { + entry.future.get(); + requestsStatus.completedRequestIds.insert(entry.request->mRequestId); + if (markComplete) + { + entry.request->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); + } + } + catch (std::exception const& e) + { + TLLM_LOG_ERROR("Error occurred during context transfer drain for request %ld: %s", + entry.request->mRequestId, e.what()); + entry.request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + requestsStatus.errorRequestIds.insert(entry.request->mRequestId); + } + releaseTrackedFutureLocked(mSenderFutures, i); + madeProgress = true; + continue; } + + // Per-entry timeout: surface error to caller, but keep the future + // pinned. Worker may still be writing; releasing the entry now + // would let Python free request resources that NIXL/UCX may still + // touch. The future stays here until the worker finishes (or the + // global progress deadline marks the transceiver unhealthy and + // orchestration restarts). + maybeQuarantineLocked(entry, &requestsStatus); + ++i; } + if (madeProgress) + { + std::scoped_lock lk(mHealthMutex); + mLastProgressTime = std::chrono::steady_clock::now(); + updateHealthLocked(); + } + else + { + std::scoped_lock lk(mHealthMutex); + updateHealthLocked(); + } return requestsStatus; } void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastRequestNum) { - bool blockAll = !atLeastRequestNum.has_value(); + checkGenTransferStatusImpl(atLeastRequestNum, /*allowBlocking=*/false); +} + +void CacheTransceiver::drainGenTransferStatus() +{ + checkGenTransferStatusImpl(std::nullopt, /*allowBlocking=*/true); +} + +void CacheTransceiver::checkGenTransferStatusImpl(std::optional const& atLeastRequestNum, bool allowBlocking) +{ + auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; + std::vector genTransferReadyRequestIds; - for (auto&& [request, future] : mRequesterFutures) + for (auto const& entry : mRequesterFutures) { - if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + if (entry.future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) { - genTransferReadyRequestIds.push_back(request->mRequestId); + genTransferReadyRequestIds.push_back(entry.request->mRequestId); } } - std::unordered_map frequencyMap; - std::vector toBlockRequestIds; - auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; - if ((syncComm) && syncComm->getSize() > 1) + std::unordered_map frequencyMap; + if (syncComm && syncComm->getSize() > 1) { auto gatherRequestIdVec = gatherRequestIds(syncComm, genTransferReadyRequestIds); for (auto&& requestId : gatherRequestIdVec) @@ -628,127 +778,97 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR } std::vector> freqVec(frequencyMap.begin(), frequencyMap.end()); - std::sort(freqVec.begin(), freqVec.end(), - [](std::pair const& left, - std::pair const& right) { return left.second > right.second; }); + [](auto const& left, auto const& right) { return left.second > right.second; }); + + int const expectedFreq = syncComm ? syncComm->getSize() : 1; std::unordered_set toCompleteIdSet; - size_t idx = 0; - while (atLeastRequestNum.value_or(0) > static_cast(toCompleteIdSet.size())) + for (auto const& [requestId, freq] : freqVec) { - if (idx >= freqVec.size()) - { - break; - } - toCompleteIdSet.insert(freqVec.at(idx).first); - if (useMPI()) - { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - " checkGenTransferStatus at least from freqVec requestId: %zu ", freqVec.at(idx).first); - } - else + if (freq == expectedFreq) { - TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), - " checkGenTransferStatus at least from freqVec requestId: %zu ", freqVec.at(idx).first); + toCompleteIdSet.insert(requestId); } - idx++; } - idx = 0; - // insert order - while (atLeastRequestNum.value_or(0) > static_cast(toCompleteIdSet.size())) + if (atLeastRequestNum.has_value()) { - if (idx >= mRequesterFutures.size()) + for (auto const& entry : mRequesterFutures) { - break; - } - if (toCompleteIdSet.find(mRequesterFutures.at(idx).first->mRequestId) == toCompleteIdSet.end()) - { - toCompleteIdSet.insert(mRequesterFutures.at(idx).first->mRequestId); - if (useMPI()) + if (static_cast(toCompleteIdSet.size()) >= atLeastRequestNum.value()) { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - " checkGenTransferStatus at least from RequesterFuture requestId: %zu atLeastRequestNum:%d", - mRequesterFutures.at(idx).first->mRequestId, atLeastRequestNum.value_or(0)); + break; } - else + if (entry.future.wait_for(std::chrono::milliseconds(0)) != std::future_status::ready) { - TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), - " checkGenTransferStatus at least from RequesterFuture requestId: %zu atLeastRequestNum:%d", - mRequesterFutures.at(idx).first->mRequestId, atLeastRequestNum.value_or(0)); + continue; } + toCompleteIdSet.insert(entry.request->mRequestId); } - idx++; } - for (auto&& [requestId, freq] : freqVec) - { - if (freq == ((syncComm != nullptr) ? syncComm->getSize() : 1)) - { - toCompleteIdSet.insert(requestId); - } - if (useMPI()) - { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), " checkGenTransferStatus freqVec requestId: %zu,freq:%d ", - requestId, freq); - } - else - { - TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), - " checkGenTransferStatus freqVec requestId: %zu,freq:%d ", requestId, freq); - } - } - if (useMPI()) - { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - " checkGenTransferStatus toCompleteIdSet size: %zu, atLeastRequestNum: %d ", toCompleteIdSet.size(), - atLeastRequestNum.value_or(0)); - } - else - { - TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), - " checkGenTransferStatus toCompleteIdSet size: %zu, atLeastRequestNum: %d ", toCompleteIdSet.size(), - atLeastRequestNum.value_or(0)); - } - for (auto it = mRequesterFutures.begin(); it != mRequesterFutures.end();) + + bool madeProgress = false; + for (size_t i = 0; i < mRequesterFutures.size();) { - if (blockAll || toCompleteIdSet.find(it->first->mRequestId) != toCompleteIdSet.end()) + auto& entry = mRequesterFutures[i]; + bool const selected = toCompleteIdSet.find(entry.request->mRequestId) != toCompleteIdSet.end(); + bool const isReady = entry.future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready; + + if (selected && isReady) { try { - it->second.get(); - it->first->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); - - // Gather the kv cache transfer time from all workers and update to leader rank + entry.future.get(); + entry.request->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); if (!common::getEnvKVCacheTimeOutputPath().empty()) { - auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; - updateKVCacheTransferBW(syncComm, it->first); + updateKVCacheTransferBW(syncComm, entry.request); } } 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); + TLLM_LOG_ERROR("Error occurred during generation transfer for request %ld: %s", + entry.request->mRequestId, e.what()); + entry.request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); } - if (useMPI()) + releaseTrackedFutureLocked(mRequesterFutures, i); + madeProgress = true; + continue; + } + + if (allowBlocking) + { + try { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***", - it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId()); + entry.future.get(); + entry.request->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); + if (!common::getEnvKVCacheTimeOutputPath().empty()) + { + updateKVCacheTransferBW(syncComm, entry.request); + } } - else + catch (std::exception const& e) { - TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), - "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***", - it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId()); + TLLM_LOG_ERROR("Error occurred during generation transfer drain for request %ld: %s", + entry.request->mRequestId, e.what()); + entry.request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); } - it = mRequesterFutures.erase(it); + releaseTrackedFutureLocked(mRequesterFutures, i); + madeProgress = true; + continue; } - else + + maybeQuarantineLocked(entry, /*outStatus=*/nullptr); + ++i; + } + + { + std::scoped_lock lk(mHealthMutex); + if (madeProgress) { - ++it; + mLastProgressTime = std::chrono::steady_clock::now(); } + updateHealthLocked(); } } @@ -757,17 +877,87 @@ bool CacheTransceiver::checkGenTransferComplete() const return mRequesterFutures.empty(); } -bool CacheTransceiver::cancelRequest(LlmRequest* llmRequest) +TransferCancelResult CacheTransceiver::cancelRequestStructured(LlmRequest* llmRequest) { + TLLM_CHECK(llmRequest != nullptr); + + if (!isHealthy()) + { + return TransferCancelResult::kBackendUnhealthy; + } + if (llmRequest->isContextOnlyRequest()) { - return mCacheSender->cancelRequest(*llmRequest); + // Pre-advertise: the sender's queue still owns the request. + if (mCacheSender->cancelRequest(*llmRequest)) + { + return TransferCancelResult::kCancelledBeforeAdvertise; + } + // Did the worker future already finish? + for (size_t i = 0; i < mSenderFutures.size(); ++i) + { + if (mSenderFutures[i].request->mRequestId != llmRequest->mRequestId) + { + continue; + } + auto status = mSenderFutures[i].future.wait_for(std::chrono::milliseconds(0)); + if (status == std::future_status::ready) + { + return TransferCancelResult::kAlreadyComplete; + } + return TransferCancelResult::kCancelRequestedInFlight; + } + return TransferCancelResult::kNotFound; } - else if (llmRequest->isGenerationOnlyRequest()) + if (llmRequest->isGenerationOnlyRequest()) { - return mCacheReceiver->cancelRequest(*llmRequest); + if (mCacheReceiver->cancelRequest(*llmRequest)) + { + return TransferCancelResult::kCancelledBeforeAdvertise; + } + for (size_t i = 0; i < mRequesterFutures.size(); ++i) + { + if (mRequesterFutures[i].request->mRequestId != llmRequest->mRequestId) + { + continue; + } + auto status = mRequesterFutures[i].future.wait_for(std::chrono::milliseconds(0)); + if (status == std::future_status::ready) + { + return TransferCancelResult::kAlreadyComplete; + } + return TransferCancelResult::kCancelRequestedInFlight; + } + return TransferCancelResult::kNotFound; } - return false; + return TransferCancelResult::kNotCancellable; +} + +bool CacheTransceiver::cancelRequest(LlmRequest* llmRequest) +{ + auto result = cancelRequestStructured(llmRequest); + return result == TransferCancelResult::kCancelledBeforeAdvertise || result == TransferCancelResult::kAlreadyComplete; +} + +bool CacheTransceiver::isHealthy() const +{ + std::scoped_lock lk(mHealthMutex); + return mIsHealthy; +} + +TransceiverHealth CacheTransceiver::getHealth() const +{ + std::scoped_lock lk(mHealthMutex); + auto const sinceProgress = std::chrono::duration_cast>( + std::chrono::steady_clock::now() - mLastProgressTime) + .count(); + return TransceiverHealth{ + mIsHealthy, + mQuarantinedTransferCount, + mQuarantineBudget, + sinceProgress, + std::chrono::duration_cast>(mGlobalProgressDeadlineMs).count(), + }; } } // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index 61cac5df4c7e..6b2db97ca7d5 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -17,8 +17,10 @@ #include "cacheTransceiver.h" #include "tensorrt_llm/batch_manager/cacheTransceiver.h" +#include "tensorrt_llm/batch_manager/contextProgress.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/rnnStateManager.h" +#include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/bindingUtils.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/nanobind/common/customCasters.h" @@ -51,6 +53,14 @@ class PyCacheTransceiver : public tb::BaseCacheTransceiver NB_OVERRIDE_PURE(respondAndSendAsync, llmRequest); } + void respondAndSendLayerWise( + tb::RequestVector const& /*requests*/, std::shared_ptr const& /*progress*/) override + { + // Layer-wise sender path is not exposed to Python overrides; this is + // a stub so the trampoline class is fully concrete. + TLLM_THROW("respondAndSendLayerWise is not overridable from Python"); + } + void requestAndReceiveSync(tb::LlmRequest* llmRequest) override { NB_OVERRIDE_PURE(requestAndReceiveSync, llmRequest); @@ -86,31 +96,64 @@ class PyCacheTransceiver : public tb::BaseCacheTransceiver void tb::CacheTransceiverBindings::initBindings(nb::module_& m) { + nb::enum_(m, "TransferCancelResult") + .value("NotFound", tb::TransferCancelResult::kNotFound) + .value("AlreadyComplete", tb::TransferCancelResult::kAlreadyComplete) + .value("CancelledBeforeAdvertise", tb::TransferCancelResult::kCancelledBeforeAdvertise) + .value("CancelRequestedInFlight", tb::TransferCancelResult::kCancelRequestedInFlight) + .value("BackendUnhealthy", tb::TransferCancelResult::kBackendUnhealthy) + .value("NotCancellable", tb::TransferCancelResult::kNotCancellable); + + nb::class_(m, "TransceiverHealth") + .def_ro("is_healthy", &tb::TransceiverHealth::isHealthy) + .def_ro("quarantined_transfer_count", &tb::TransceiverHealth::quarantinedTransferCount) + .def_ro("quarantine_budget", &tb::TransceiverHealth::quarantineBudget) + .def_ro("seconds_since_last_progress", &tb::TransceiverHealth::secondsSinceLastProgress) + .def_ro("global_progress_deadline_seconds", &tb::TransceiverHealth::globalProgressDeadlineSeconds); + + auto checkContextStatusLambda + = [](tb::BaseCacheTransceiver& self, std::optional const& atLeastRequestNum, bool markComplete) + { + RequestStatuses result; + { + nb::gil_scoped_release release; + result = self.checkContextTransferStatus(atLeastRequestNum, markComplete); + } + auto completedRequestIds + = std::vector(result.completedRequestIds.begin(), result.completedRequestIds.end()); + auto errorRequestIds = std::vector(result.errorRequestIds.begin(), result.errorRequestIds.end()); + return nb::make_tuple(completedRequestIds, errorRequestIds); + }; + + auto drainContextStatusLambda = [](tb::BaseCacheTransceiver& self, bool markComplete) + { + RequestStatuses result; + { + nb::gil_scoped_release release; + result = self.drainContextTransferStatus(markComplete); + } + auto completedRequestIds + = std::vector(result.completedRequestIds.begin(), result.completedRequestIds.end()); + auto errorRequestIds = std::vector(result.errorRequestIds.begin(), result.errorRequestIds.end()); + return nb::make_tuple(completedRequestIds, errorRequestIds); + }; + nb::class_(m, "BaseCacheTransceiver") .def("respond_and_send_async", &BaseCacheTransceiver::respondAndSendAsync) .def("request_and_receive_sync", &BaseCacheTransceiver::requestAndReceiveSync) .def("request_and_receive_async", &BaseCacheTransceiver::requestAndReceiveAsync) - .def( - "check_context_transfer_status", - [](tb::BaseCacheTransceiver& self, std::optional const& atLeastRequestNum, bool markComplete = false) - { - RequestStatuses result; - { - nb::gil_scoped_release release; - result = self.checkContextTransferStatus(atLeastRequestNum, markComplete); - } - - auto completedRequestIds - = std::vector(result.completedRequestIds.begin(), result.completedRequestIds.end()); - auto errorRequestIds - = std::vector(result.errorRequestIds.begin(), result.errorRequestIds.end()); - return nb::make_tuple(completedRequestIds, errorRequestIds); - }, + .def("check_context_transfer_status", checkContextStatusLambda, nb::arg("at_least_request_num") = std::nullopt, nb::arg("mark_complete") = false) + .def("drain_context_transfer_status", drainContextStatusLambda, nb::arg("mark_complete") = false) .def("check_gen_transfer_status", &BaseCacheTransceiver::checkGenTransferStatus, nb::call_guard()) + .def("drain_gen_transfer_status", &BaseCacheTransceiver::drainGenTransferStatus, + nb::call_guard()) .def("check_gen_transfer_complete", &BaseCacheTransceiver::checkGenTransferComplete) - .def("cancel_request", &BaseCacheTransceiver::cancelRequest); + .def("cancel_request", &BaseCacheTransceiver::cancelRequest) + .def("cancel_request_structured", &BaseCacheTransceiver::cancelRequestStructured) + .def("is_healthy", &BaseCacheTransceiver::isHealthy) + .def("get_health", &BaseCacheTransceiver::getHealth); nb::enum_(m, "AttentionType") .value("DEFAULT", executor::kv_cache::CacheState::AttentionType::kDEFAULT) diff --git a/docs/source/features/disagg-kv-transfer-hang-restart-analysis.md b/docs/source/features/disagg-kv-transfer-hang-restart-analysis.md new file mode 100644 index 000000000000..a8107d6043a7 --- /dev/null +++ b/docs/source/features/disagg-kv-transfer-hang-restart-analysis.md @@ -0,0 +1,281 @@ +# Disaggregated KV transfer restart and hang analysis + +This note summarizes the latest behavior observed across PR 13706, PR 13713, +and PR 13728, then proposes the fix direction that avoids both unsafe memory +reuse and permanent serving hangs. + +## Observed behavior + +The latest tests split the failure into two different classes: + +| Change set | `py_executor.py` changes included | Observed behavior | Immediate implication | +| --- | --- | --- | --- | +| PR 13706 | yes | frequent prefill worker restarts | The Python policy is too aggressive and turns common in-flight transfer timeouts into process-level fail-close events. | +| PR 13706 | no | no frequent restarts | The C++ side alone is not what triggers those restart loops. | +| PR 13713 | yes | hangs | The hang is not explained only by Python fail-close behavior. | +| PR 13713 | no | hangs | There is still a lower-level C++ or transport progress problem. | +| PR 13728 | yes | hangs | The added Python containment does not repair the lower-level hang. | +| PR 13728 | no | hangs | The lower-level hang persists independently of Python cleanup policy. | + +```mermaid +flowchart TD + A["PR 13706 with Python changes"] --> B["Python sees transfer timeout"] + B --> C["Python fails closed"] + C --> D["Worker restarts often"] + + E["PR 13706 without Python changes"] --> F["No Python fail-close trigger"] + F --> G["No frequent restart loop"] + + H["PR 13713 or PR 13728"] --> I["Hangs with Python changes"] + H --> J["Hangs without Python changes"] + I --> K["Hang is below Python policy"] + J --> K +``` + +The PR 13706 result tells us that per-request Python fail-close is the wrong +default policy. It can be useful as an emergency containment mechanism, but it +is too blunt when it fires for routine transfer timeouts or cancellable races. + +The PR 13713 and PR 13728 results tell us something different: removing the +Python policy is not sufficient. If those images still hang without the +`py_executor.py` changes, then at least one C++ or backend path can still stop +making progress without surfacing a final request error. + +## Relevant code paths + +The Python-side restart behavior comes from timeout and cleanup paths around: + +- `PyExecutor._check_kv_transfer_timeout()` in + `tensorrt_llm/_torch/pyexecutor/py_executor.py`: marks in-flight context or + generation transfers as timed out. +- `PyExecutor._check_disagg_ctx_cache_transfer_status()`: observes timed-out + context transfers and calls `kv_cache_transceiver.cancel_request(request)`. +- `PyExecutor._handle_responses()`: observes timed-out generation transfers and + calls `kv_cache_transceiver.cancel_request(request)`. +- The PR 13706-style Python policy then treated unknown quiescence as a fatal + executor condition, cleared active queues, and made the worker exit or become + unhealthy. + +The lower-level hang behavior must be analyzed around the C++ transfer +progress paths: + +- `CacheTransceiver::checkContextTransferStatus()` and + `CacheTransceiver::checkGenTransferStatus()` in + `cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp`. +- `CacheSender::Impl` and `CacheReceiver::Impl` in + `cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp`. +- Transfer-buffer assignment and release in + `cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp`. +- Agent/NIXL calls reached through + `cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp`. + +If a PR hangs with Python cleanup removed, the main event loop is either still +blocking on a C++ wait, or C++ has stranded an in-flight request/future without +making it observable as ready, error, or unhealthy. + +## Why the restart-only fix is not acceptable + +A pure fail-closed policy is memory safe, but it is operationally noisy: + +```mermaid +sequenceDiagram + participant Py as "Python executor" + participant Cpp as "C++ transceiver" + participant Orch as "Orchestrator" + + Py->>Py: "KV transfer timeout" + Py->>Cpp: "cancel_request(request)" + Cpp-->>Py: "cancel accepted" + Py->>Py: "quiescence unknown" + Py->>Py: "set shutdown or unhealthy" + Orch->>Py: "restart worker" +``` + +This avoids reusing memory that a transport may still write, but it restarts on +events that are not necessarily fatal. Under bursty traffic, cancellations, or +slow backend progress, that can turn a recoverable per-request timeout into a +cluster-wide restart loop. + +## Why the C++-only fix is not sufficient + +A C++-only approach can improve lifetime and avoid some leaks, but the latest +PR 13713 and PR 13728 observations show that it can still hang. That means the +implementation still lacks at least one of these properties: + +- every transfer future eventually becomes ready, error, or explicitly + abandoned with a health signal; +- no executor-facing status function blocks on a future or backend wait unless + the caller explicitly requested blocking behavior; +- cancellation has a final state that distinguishes "safe to free" from + "cancel requested but remote writes may still exist"; +- poisoned or quarantined resources are bounded and converted into an + explicit unhealthy signal before the executor silently starves. + +## Proper fix direction + +The fix should separate three concepts that the earlier PRs mixed together: + +1. Request failure: the user request can receive a timeout or network error. +2. Resource quiescence: C++ has proof that no sender, receiver, or backend + worker can still touch the request's KV blocks or transfer buffers. +3. Process health: the worker should restart only when the backend or resource + pool is no longer capable of safe progress. + +### 1. Python should not fail-close per request + +Python should keep these responsibilities: + +- mark transfer start time; +- mark a request timed out after `kv_transfer_timeout_ms`; +- call `cancel_request()` once for a timed-out transfer; +- avoid `_terminate_request()` while the request is still in a disaggregated + transfer state; +- report a request-level error only when C++ reports a final transfer state. + +Python should not set global shutdown or clear serving queues solely because a +single in-flight transfer timed out. That is the PR 13706 restart-loop trigger. + +### 2. C++ should own transfer sessions until final state + +C++ needs a transfer-session concept, explicit or equivalent, that owns all +objects needed by the in-flight transfer: + +- a `shared_ptr`; +- receive or send buffer-index leases; +- any request-owned KV-block lease needed by direct or zero-copy paths; +- the promise/future pair; +- cancellation state and timeout deadline. + +The session can release resources only on proven-safe outcomes: + +| Outcome | Safe action | +| --- | --- | +| Completed transfer | release buffers and KV-block leases | +| Peer returned explicit not-ready before data phase | release buffers | +| Cancelled before buffer advertisement | release buffers | +| Worker future ready with structured error | release or quarantine according to where the error occurred | +| Timeout after buffer advertisement and no backend completion | quarantine or poison advertised resources | +| Backend worker wedged or future never resolves | keep resources pinned or quarantined and raise health signal | + +The important rule is that `cancel_request()` must not mean "safe to free". +It should return a structured result such as: + +- `not_found`; +- `already_complete`; +- `cancelled_before_advertise`; +- `cancel_requested_in_flight`; +- `backend_unhealthy`; +- `not_cancellable`. + +Only the first three can permit immediate Python cleanup. The in-flight cases +must remain owned by C++ until a final status arrives or until the resource is +quarantined and the process is marked unhealthy. + +### 3. C++ status checks must be non-blocking by default + +The executor-facing status functions should be polling APIs: + +- `checkContextTransferStatus(0)` and `checkGenTransferStatus(0)` must not + block on an unready future. +- `atLeastNum=1` should skip unready futures rather than calling `future.get()` + on a selected-but-unready entry. +- Timeout handling should mark the transfer session cancelled or failed, but + should not erase the tracking future until the worker reaches a final state. +- Blocking behavior should exist only for explicit drain/shutdown paths. + +If a backend call can block indefinitely inside NIXL or UCX, it must not run on +the Python event-loop thread. If it runs on a dedicated worker and wedges, the +event loop must still be able to observe "no progress" and produce a health +signal. + +### 4. Quarantine should be bounded and health-driven + +Quarantine or poisoning is still needed for memory safety after a buffer has +been advertised to a peer and quiescence is unknown. But poisoning should not +immediately restart the worker for a single request unless no safe capacity +remains. + +Recommended policy: + +- pre-advertise cancellation releases normally; +- post-advertise unknown-quiescence quarantines the specific buffer or transfer + slot; +- the transceiver refuses to reuse quarantined resources; +- if the number of quarantined slots or pinned KV blocks exceeds a small budget, + or if no transfer progress occurs for a global backend deadline, mark the + transceiver unhealthy; +- health failure should be explicit and observable, so orchestration restarts + the worker instead of the server silently hanging. + +```mermaid +sequenceDiagram + participant Py as "Python executor" + participant S as "C++ transfer session" + participant B as "Transfer buffer or KV block" + participant H as "Health monitor" + + Py->>S: "timeout, cancel requested" + alt not advertised + S->>B: "release normally" + S-->>Py: "final request error" + else advertised and backend completes + S->>B: "release after completion" + S-->>Py: "final request error" + else advertised and backend does not quiesce + S->>B: "quarantine, do not reuse" + S-->>Py: "request error if safe to report" + S->>H: "increase poisoned or pinned resource count" + H-->>Py: "unhealthy only if budget or progress deadline is exceeded" + end +``` + +This policy avoids both bad extremes: + +- it does not free or reuse memory that may still be touched by the transport; +- it does not restart the worker for every ordinary timeout. + +## Expected behavior after the proper fix + +| Scenario | Expected result | +| --- | --- | +| Ordinary successful transfer | request completes and resources release normally | +| User cancellation before transfer advertisement | request fails or cancels; resources release normally | +| Timeout before transfer advertisement | request fails; resources release normally | +| Timeout after advertisement but backend later completes | request fails; resources release after worker final state | +| Timeout after advertisement and backend never completes | resource is quarantined; server keeps serving if capacity remains | +| Repeated non-quiescing transfers exhaust quarantine budget | worker becomes unhealthy and restarts with a clear reason | +| NIXL or UCX internal deadlock | no silent hang; health fails after progress deadline unless the backend root cause is fixed | + +## Recommended implementation order + +1. Remove the per-request Python fail-close behavior from PR 13706-style + branches. Keep only the "do not terminate while transfer is in progress" + guard. +2. Make C++ cancellation/status return a structured final-state enum instead + of a boolean. +3. Keep futures and transfer sessions tracked until the worker reaches ready, + error, or quarantined final state. +4. Add bounded quarantine counters for advertised transfer buffers and any + request-owned KV-block leases. +5. Add a transceiver health state that flips only when quarantine capacity is + exhausted or backend progress is absent past a global deadline. +6. Ensure ADP response enqueueing uses the synchronized pending-response flush + path from PR 13112 when testing on rc13-derived images. +7. Add regression tests for: + - PR 13706-style timeout bursts without restart loops; + - PR 13713/13728-style C++ path without Python changes, proving status APIs + return error or unhealthy instead of hanging; + - advertised-buffer timeout followed by block reuse, proving the block is + quarantined and not reused; + - ADP response flush, proving all ranks enter response collectives together. + +## Bottom line + +The Python fail-close policy explains PR 13706's frequent restarts, but it does +not explain PR 13713 or PR 13728 hanging without Python changes. The correct +solution is therefore not "more Python fail-close" and not "C++ lifetime only". + +The correct solution is a C++ transfer-session lifecycle with structured cancel +states, non-blocking status polling, bounded quarantine for unknown-quiescence +resources, and an explicit health signal only when safe serving capacity or +backend progress is genuinely lost. diff --git a/docs/source/features/disagg-kv-transfer-session-lifecycle.md b/docs/source/features/disagg-kv-transfer-session-lifecycle.md new file mode 100644 index 000000000000..f81f4438135f --- /dev/null +++ b/docs/source/features/disagg-kv-transfer-session-lifecycle.md @@ -0,0 +1,122 @@ +# Disaggregated KV transfer: session lifecycle and cancel contract + +> **Status:** beta — applies to disaggregated serving with the C++ cache +> transceiver (NIXL / UCX backends). + +This note describes the contract between the Python executor and the C++ +cache transceiver for cancellation, timeout, quarantine, and process +health. It supersedes the per-request fail-close behaviour explored in +PR 13706 and the C++-only lifetime hardening in PRs 13713 / 13728: +both versions still hung in the field because neither isolated the +"backend is wedged but the executor thread is blocked" failure mode. + +The fix direction is documented in +`disagg-kv-transfer-hang-restart-analysis.md`. This file describes the +mechanism that this codebase actually implements. + +## Three responsibilities, separated + +| Responsibility | Owner | Mechanism | +| ----------------- | -------------------- | ----------------------------------------- | +| Request failure | Python executor | Surface a request-level error to the user when C++ reports a final transfer state. | +| Resource quiescence | C++ transceiver | Worker future stays pinned until ready / error / quarantine. | +| Process health | C++ transceiver | Quarantine budget + global progress deadline → `isHealthy() == false`. | + +## Status polling is non-blocking by default + +`BaseCacheTransceiver::checkContextTransferStatus` and +`checkGenTransferStatus` poll worker futures with `wait_for(0ms)` and +**never** call `future.get()` on an unready entry. When +`atLeastRequestNum > 0` is requested, the polling code admits additional +ready futures but skips unready ones rather than blocking the executor +thread. + +For shutdown / drain there are explicit blocking variants: +`drainContextTransferStatus` and `drainGenTransferStatus`. These are the +only paths permitted to block on a worker future. + +## Structured cancellation + +`cancelRequest(LlmRequest*)` returns a backward-compatible bool; +`cancelRequestStructured(LlmRequest*)` returns a +`TransferCancelResult` enum: + +| Value | Caller may release request resources? | +| --------------------------- | -------------------------------------- | +| `NotFound` | yes — C++ already reached a final state | +| `AlreadyComplete` | yes — worker future is ready | +| `CancelledBeforeAdvertise` | yes — no buffer was advertised to the peer | +| `CancelRequestedInFlight` | **no** — worker may still touch buffers | +| `BackendUnhealthy` | **no** — orchestration must restart | +| `NotCancellable` | **no** — retry later | + +Only the first three states allow Python to free the request's KV / +transfer-buffer resources. The remaining states keep the request owned +by C++; the Python executor adds the request id to +`_inflight_cancel_requested_ids`, which `_can_terminate_request_now` +checks before any `_do_terminate_request` call. + +## Bounded quarantine, explicit health signal + +When a tracked transfer exceeds its per-request `kvTransferTimeoutMs` +deadline and the worker future is still not ready, the entry is flipped +to *quarantined*: an error is surfaced to the caller, but the future +stays in the tracking vector. A counter (`mQuarantinedTransferCount`) +increments. If the counter exceeds `mQuarantineBudget`, **or** no +worker has reached a final state for longer than +`mGlobalProgressDeadlineMs`, the transceiver flips +`isHealthy()` to false. Orchestration is expected to restart the +worker on this signal — neither C++ nor Python frees pinned memory +locally. + +`getHealth()` returns a `TransceiverHealth` snapshot exposing +`is_healthy`, `quarantined_transfer_count`, `quarantine_budget`, +`seconds_since_last_progress`, and `global_progress_deadline_seconds` +for metrics / observability. + +## Python policy + +The `PyExecutor` keeps these responsibilities only: + +* mark transfer start time (`py_kv_transfer_start_time`); +* mark a request timed out after `kv_transfer_timeout_ms` + (`py_kv_transfer_timed_out`); +* call `cancel_request_structured()` once for a timed-out transfer; +* defer `_terminate_request()` while the request is still in a + disaggregated transfer state — see `_can_terminate_request_now`; +* report a request-level error only when the structured cancel result + is one of the safe-to-release outcomes. + +The PR 13706 anti-pattern of clearing `active_requests`, +`waiting_queue`, `request_accumulated`, and setting `is_shutdown=True` +on a single in-flight transfer timeout is **not** present and must not +be added back. Per-request timeouts are recoverable; only an unhealthy +transceiver should drive a process restart, and that restart is +orchestration's responsibility. + +## ADP response flush + +When attention data parallelism (`enable_attention_dp=True`) is +enabled, `_enqueue_responses` runs a `tp_gather` collective. Calling it +inside `_end_transfer_and_maybe_terminate` deadlocks because only the +DP rank that owns the request reaches that point. The executor instead +buffers transfer-completion responses in `_pending_transfer_responses` +and flushes the buffer at synchronised loop points where every DP rank +participates in the collective: + +* before `_handle_canceled_requests` in the PP path; +* before `_handle_canceled_requests` in the standard executor loop; +* unconditionally in the overlap loop (so ranks with differing + `should_process_previous_batch` still participate). + +`_flush_pending_transfer_responses` always calls `_enqueue_responses` +when ADP is enabled, even if the local buffer is empty, so the +collective is symmetric across ranks. + +## Test coverage + +* `tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py` + — unit tests for the deferred-terminate guard, the structured-cancel + decision tree, and the ADP flush symmetry. These tests intentionally + do not require GPU or MPI; they replay the policy decisions as plain + Python so failures fail-fast in pre-merge CI. diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index e27290b472fa..a521842672a1 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -18,6 +18,12 @@ CacheTransBufferManagerCpp = tensorrt_llm.bindings.internal.batch_manager.CacheTransBufferManager BackendTypeCpp = tensorrt_llm.bindings.executor.CacheTransceiverBackendType +# Re-export the structured cancel-result enum so the rest of the Python +# executor can match on it without reaching into the bindings module +# directly. +TransferCancelResult = tensorrt_llm.bindings.internal.batch_manager.TransferCancelResult +TransceiverHealth = tensorrt_llm.bindings.internal.batch_manager.TransceiverHealth + def mapping_to_world_config(mapping: Mapping) -> WorldConfig: @@ -122,6 +128,32 @@ def check_gen_transfer_complete(self): def cancel_request(self, req: LlmRequest): raise NotImplementedError + def cancel_request_structured(self, req: LlmRequest): + """Structured cancel result. + + Default implementation maps the historical bool result onto the + enum so subclasses (e.g., the Python transceiver) keep working + without changes. The C++ binding overrides this with a real + implementation that distinguishes pre-advertise cancellation from + in-flight cancellation. + """ + cancelled = self.cancel_request(req) + return TransferCancelResult.CancelledBeforeAdvertise if cancelled else TransferCancelResult.NotFound + + def is_healthy(self) -> bool: + """Whether the transceiver is currently healthy. + + Defaults to True; subclasses with a real health signal override. + """ + return True + + def get_health(self): + """Return a TransceiverHealth snapshot. Defaults to a healthy + snapshot for transceivers that have not yet implemented health + accounting. + """ + return None + @abstractmethod def prepare_context_requests(self, requests: List[LlmRequest]): """ @@ -205,6 +237,15 @@ def check_gen_transfer_complete(self): def cancel_request(self, req: LlmRequest): return self.impl.cancel_request(req) + def cancel_request_structured(self, req: LlmRequest): + return self.impl.cancel_request_structured(req) + + def is_healthy(self) -> bool: + return self.impl.is_healthy() + + def get_health(self): + return self.impl.get_health() + def prepare_context_requests(self, requests: List[LlmRequest]): # not implemented, an empty placeholder to allow being invoked unconditionally ... diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 98cf521d1ff4..8d6eb1304334 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -55,7 +55,7 @@ from .handle_additional_outputs import HandleAdditionalOutputs from .handle_logits import HandleLogits from .hang_detector import HangDetector -from .kv_cache_transceiver import KvCacheTransceiver +from .kv_cache_transceiver import KvCacheTransceiver, TransferCancelResult from .llm_request import (ExecutorRequest, LlmRequest, LlmRequestState, LlmResponse, get_draft_token_length) from .model_engine import ModelEngine @@ -414,6 +414,21 @@ def __init__( self.max_num_active_requests = model_engine.get_max_num_sequences() self.active_requests: List[LlmRequest] = [] self.expected_num_active_requests = 0 + # Buffer for transfer-completion responses created inside + # _end_transfer_and_maybe_terminate. With ADP enabled, + # _enqueue_responses runs a tp_gather collective; calling it from + # _send_kv_async causes a collective mismatch because only the + # owning DP rank reaches that point. Buffering and flushing at a + # synchronised loop point (where every DP rank participates) + # avoids the deadlock. + self._pending_transfer_responses: List[Tuple[int, LlmResponse]] = [] + # Requests whose transfer timed out and we have already issued a + # cancel_request to C++. We must not call _terminate_request on + # these until C++ reports a final transfer state, otherwise we + # would free request resources while NIXL/UCX may still be + # writing into them. Keyed by py_request_id to survive request + # state transitions. + self._inflight_cancel_requested_ids: set[int] = set() # TODO: Remove the condition on the PP size once disagg support from KVCache reuse # path is fixed. self.async_transfer_manager = AsyncTransferManager( @@ -634,7 +649,14 @@ def _end_transfer_and_maybe_terminate(self, request: LlmRequest): response = request.create_response(False, self.dist.rank) if response: response.result.cached_tokens = request.cached_tokens - self._enqueue_responses([(request.py_request_id, response)]) + # Buffer the response instead of enqueueing immediately. + # _enqueue_responses runs a tp_gather collective when ADP + # is enabled; calling it here would deadlock because only + # the owning DP rank reaches this point. The buffer is + # flushed later at _flush_pending_transfer_responses + # where all ranks participate. + self._pending_transfer_responses.append( + (request.py_request_id, response)) if self.async_transfer_manager.end_transfer(request): self.active_requests.remove(request) self._terminate_request(request) @@ -647,6 +669,20 @@ def _end_transfer_and_maybe_terminate(self, request: LlmRequest): if not self.async_transfer_manager.should_store_blocks: self._terminate_request(request) + def _flush_pending_transfer_responses(self): + """Enqueue buffered transfer-completion responses. + + Must be called at a point where ALL DP ranks execute in lockstep so + the tp_gather inside _enqueue_responses cannot deadlock. + """ + responses = self._pending_transfer_responses + self._pending_transfer_responses = [] + # Even when this rank has nothing to enqueue, we must still call + # _enqueue_responses with an empty list so that other DP ranks' + # tp_gather can complete. The collective is symmetric. + if responses or self.enable_attention_dp: + self._enqueue_responses(responses) + # Performance metrics methods are in PerfMetricsManager (self.perf_manager) def _event_loop_wrapper(self): @@ -1700,6 +1736,7 @@ def _handle_executed_batch(self, executed_batch: Optional[BatchStatePP]): if self.kv_cache_transceiver: finished_ctx_reqs = scheduled_requests.context_requests_last_chunk self._send_kv_async(finished_ctx_reqs) + self._flush_pending_transfer_responses() self._handle_canceled_requests() finished_requests = self._handle_responses() @@ -2174,6 +2211,7 @@ def _executor_loop(self): self._update_requests(sample_state, self.resource_manager) self._send_kv_async(scheduled_batch.all_requests()) + self._flush_pending_transfer_responses() self._handle_canceled_requests() finished_requests = self._handle_responses() @@ -2421,6 +2459,11 @@ def _executor_loop_overlap(self): self._send_kv_async( self.previous_batch.scheduled_requests.all_requests()) + # Flush outside the conditional so all DP ranks participate + # in the tp_gather collective even when + # should_process_previous_batch differs between ranks. + self._flush_pending_transfer_responses() + if self.drafter is not None and self.use_spec_decode and should_process_previous_batch: # Cleanup previous draft resources used in the draft model self.drafter.cleanup_previous_draft_resources() @@ -3331,6 +3374,11 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): requests_in_transfer = self.async_transfer_manager.requests_in_transfer( ) + # Worker reached a final state — clear any pending in-flight cancel + # bookkeeping for those requests. + if completed_req_ids: + self._inflight_cancel_requested_ids -= completed_req_ids + for request_id in completed_req_ids: if request_id not in requests_in_transfer: @@ -3349,14 +3397,42 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): for request_id in list(requests_in_transfer.keys()): request = requests_in_transfer[request_id] if request.py_kv_transfer_timed_out and request_id not in completed_req_ids: - is_cancelled = self.kv_cache_transceiver.cancel_request(request) - # If cancel is successful, mark as complete so it can be cleaned up - # Otherwise, try at next iteration - if is_cancelled: + # Use structured cancel so we can distinguish + # "safe-to-free" from "cancellation accepted but worker + # still in flight". For the in-flight case we keep the + # request owned by C++ until a final transfer state is + # observed; freeing it now would let NIXL/UCX overwrite + # released memory. + cancel_result = self.kv_cache_transceiver.cancel_request_structured( + request) + if cancel_result in ( + TransferCancelResult.CancelledBeforeAdvertise, + TransferCancelResult.AlreadyComplete, + TransferCancelResult.NotFound): + # NotFound at this point means the C++ worker future + # has already reached a final state and was erased + # by check_context_transfer_status above; continue + # cleanup as if it were AlreadyComplete. request.py_kv_transfer_start_time = None request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE - + self._inflight_cancel_requested_ids.discard(request_id) self._end_transfer_and_maybe_terminate(request) + elif cancel_result == TransferCancelResult.CancelRequestedInFlight: + # Cancellation accepted but the worker is still + # touching the buffer/KV memory. Mark the request as + # in-flight-cancel so _do_terminate_request defers + # freeing resources, and re-check next iteration. + self._inflight_cancel_requested_ids.add(request_id) + elif cancel_result == TransferCancelResult.BackendUnhealthy: + # Soft signal — the transceiver is unhealthy. We do + # not touch this request; orchestration health checks + # will restart the worker once the global deadline + # elapses. + logger.warning( + f"Transceiver reports unhealthy while cancelling " + f"context request {request_id}; deferring to " + f"health-driven restart.") + self._inflight_cancel_requested_ids.add(request_id) self._check_cache_transfer_errors("context requests") @@ -3568,6 +3644,33 @@ def _handle_errors(self, for request in failed_requests: self._terminate_request(request) + def _can_terminate_request_now(self, request: LlmRequest) -> bool: + """Return False if the request still has an in-flight C++ + transfer that may write to its KV/buffer memory. + + Calling _do_terminate_request on such a request would let + free_resources() return KV blocks to the pool while NIXL/UCX is + still writing into them. We instead leave the request owned by + C++ until cancel_request_structured reports CancelledBefore- + Advertise / AlreadyComplete, or the transceiver is flipped to + unhealthy and orchestration restarts the worker. + """ + if self.kv_cache_transceiver is None: + return True + if request.is_dummy_request: + return True + # Active disagg transfer state: the C++ worker future is still + # tracked. Don't free. + if (request.state == LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + or request.state == + LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS): + return False + # We have already issued a cancel for this request and C++ has + # not yet returned a final state. Do not free. + if request.py_request_id in self._inflight_cancel_requested_ids: + return False + return True + def _terminate_request(self, request: LlmRequest): # Dummy requests don't participate in disagg KV cache transfers, # so they must bypass the PP termination handler to avoid stale @@ -3580,6 +3683,17 @@ def _terminate_request(self, request: LlmRequest): self._do_terminate_request(request) def _do_terminate_request(self, request: LlmRequest): + if not self._can_terminate_request_now(request): + # Defer: the C++ worker future is still alive. We will retry + # termination on the next iteration once the structured + # cancel result transitions to CancelledBeforeAdvertise / + # AlreadyComplete, or the request reaches a final transfer + # state via _check_disagg_ctx_cache_transfer_status / + # check_gen_transfer_status. + logger.debug( + f"Deferring _do_terminate_request for {request.py_request_id}: " + f"transfer still in flight (state={request.state})") + return self.resource_manager.free_resources(request) if self.gather_all_responses or self.dist.rank == 0: @@ -3729,11 +3843,38 @@ def _handle_responses(self): # Check if generation request needs cleanup due to KV cache transfer timeout if request.py_kv_transfer_timed_out: - is_cancelled = self.kv_cache_transceiver.cancel_request(request) - if is_cancelled: + # Use structured cancel so an "accepted but worker + # in-flight" outcome does not free request resources + # while NIXL/UCX may still write into them. + cancel_result = self.kv_cache_transceiver.cancel_request_structured( + request) + if cancel_result in ( + TransferCancelResult.CancelledBeforeAdvertise, + TransferCancelResult.AlreadyComplete, + TransferCancelResult.NotFound): + # NotFound at this stage means C++ already moved the + # request to a final state (worker future ready and + # erased by check_gen_transfer_status); treat the + # same as AlreadyComplete so we surface the timeout + # error to the user and free resources. + self._inflight_cancel_requested_ids.discard(req_id) self._handle_errors( error_msg=f"Request {request.py_request_id} timed out", requests=[request]) + elif cancel_result == TransferCancelResult.CancelRequestedInFlight: + # Worker still touching buffers — keep the request + # active so we re-poll next iteration. Once C++ + # transitions the worker future to a final state we + # will see NotFound / AlreadyComplete and clean up. + self._inflight_cancel_requested_ids.add(req_id) + new_active_requests.append(request) + elif cancel_result == TransferCancelResult.BackendUnhealthy: + logger.warning( + f"Transceiver reports unhealthy while cancelling " + f"generation request {req_id}; deferring to " + f"health-driven restart.") + self._inflight_cancel_requested_ids.add(req_id) + new_active_requests.append(request) continue if request.is_generation_only_request() and not request.is_finished: diff --git a/tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py b/tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py new file mode 100644 index 000000000000..44a4f41d571d --- /dev/null +++ b/tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py @@ -0,0 +1,348 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 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. +"""Unit tests for the disaggregated KV-transfer session lifecycle. + +These tests exercise the Python policy layer added on top of the C++ +transceiver and intentionally avoid requiring a real GPU, MPI runtime, or +the full PyExecutor stack. The focus is the contract that the analysis +doc requires us to keep: + +* per-request timeouts must NOT trigger a global executor shutdown; +* a CancelRequestedInFlight result must keep the request owned by C++ + (no _do_terminate_request, no removal from active queues); +* CancelledBeforeAdvertise / AlreadyComplete / NotFound (after a final + state has been observed) must release resources promptly; +* the ADP pending-response flush is symmetric across DP ranks. +""" + +from __future__ import annotations + +import enum +import types +from dataclasses import dataclass, field +from typing import List, Optional + +import pytest + + +class _FakeTransferCancelResult(enum.Enum): + """Mirror of the C++ enum exposed by the binding so the tests do not + pull in the heavyweight binding module on import.""" + NotFound = 0 + AlreadyComplete = 1 + CancelledBeforeAdvertise = 2 + CancelRequestedInFlight = 3 + BackendUnhealthy = 4 + NotCancellable = 5 + + +@dataclass +class _FakeRequest: + py_request_id: int + state: str = "DISAGG_GENERATION_TRANS_IN_PROGRESS" + py_kv_transfer_timed_out: bool = False + py_kv_transfer_start_time: Optional[float] = None + is_dummy_request: bool = False + + @property + def is_disagg_generation_transmission_in_progress(self) -> bool: + return self.state == "DISAGG_GENERATION_TRANS_IN_PROGRESS" + + +@dataclass +class _FakeTransceiver: + """Programmable stand-in for the C++-backed KvCacheTransceiver. The + test sets `next_cancel_results` to a list keyed by request id; each + cancel call pops the head.""" + next_cancel_results: dict = field(default_factory=dict) + cancel_calls: List[int] = field(default_factory=list) + healthy: bool = True + + def cancel_request_structured(self, req: _FakeRequest): + self.cancel_calls.append(req.py_request_id) + return self.next_cancel_results.get( + req.py_request_id, _FakeTransferCancelResult.CancelRequestedInFlight) + + def is_healthy(self) -> bool: + return self.healthy + + +def _make_executor_under_test(): + """Build a minimal object whose interface matches what the helper + methods need. We avoid constructing a real PyExecutor because that + drags in CUDA, MPI, and the full model engine.""" + obj = types.SimpleNamespace() + + obj.kv_cache_transceiver = _FakeTransceiver() + obj._inflight_cancel_requested_ids = set() + obj._pending_transfer_responses = [] + obj.enable_attention_dp = False + obj.active_requests = [] + obj.terminated_request_ids = [] + obj.enqueue_calls = [] + + # Bind the real implementations (copied from py_executor.py) by + # patching through closures so we do not have to import the whole + # module — the unit-under-test functions are pure Python with no + # CUDA/MPI dependencies. + + def _can_terminate_request_now(self, request) -> bool: + if self.kv_cache_transceiver is None: + return True + if request.is_dummy_request: + return True + if request.state in ("DISAGG_CONTEXT_TRANS_IN_PROGRESS", + "DISAGG_GENERATION_TRANS_IN_PROGRESS"): + return False + if request.py_request_id in self._inflight_cancel_requested_ids: + return False + return True + + def _do_terminate_request(self, request): + if not self._can_terminate_request_now(request): + return + self.terminated_request_ids.append(request.py_request_id) + + def _enqueue_responses(self, responses): + # Mimics the dist-side aggregator; we just record what was sent. + self.enqueue_calls.append(list(responses)) + + def _flush_pending_transfer_responses(self): + responses = self._pending_transfer_responses + self._pending_transfer_responses = [] + if responses or self.enable_attention_dp: + self._enqueue_responses(responses) + + obj._can_terminate_request_now = types.MethodType( + _can_terminate_request_now, obj) + obj._do_terminate_request = types.MethodType(_do_terminate_request, obj) + obj._enqueue_responses = types.MethodType(_enqueue_responses, obj) + obj._flush_pending_transfer_responses = types.MethodType( + _flush_pending_transfer_responses, obj) + return obj + + +# --------------------------------------------------------------------------- +# _can_terminate_request_now +# --------------------------------------------------------------------------- + + +def test_can_terminate_request_defers_during_active_transfer(): + executor = _make_executor_under_test() + req = _FakeRequest(py_request_id=42, + state="DISAGG_GENERATION_TRANS_IN_PROGRESS") + assert not executor._can_terminate_request_now(req) + + +def test_can_terminate_request_defers_when_cancel_in_flight(): + executor = _make_executor_under_test() + req = _FakeRequest(py_request_id=42, state="GENERATION_COMPLETE") + executor._inflight_cancel_requested_ids.add(42) + assert not executor._can_terminate_request_now(req) + + +def test_can_terminate_request_allows_when_safe(): + executor = _make_executor_under_test() + req = _FakeRequest(py_request_id=42, state="GENERATION_COMPLETE") + assert executor._can_terminate_request_now(req) + + +def test_dummy_request_is_always_terminable(): + executor = _make_executor_under_test() + req = _FakeRequest(py_request_id=42, + state="DISAGG_GENERATION_TRANS_IN_PROGRESS", + is_dummy_request=True) + assert executor._can_terminate_request_now(req) + + +def test_no_transceiver_means_terminable(): + executor = _make_executor_under_test() + executor.kv_cache_transceiver = None + req = _FakeRequest(py_request_id=42, + state="DISAGG_GENERATION_TRANS_IN_PROGRESS") + assert executor._can_terminate_request_now(req) + + +# --------------------------------------------------------------------------- +# _do_terminate_request honors the guard +# --------------------------------------------------------------------------- + + +def test_do_terminate_defers_when_unsafe(): + executor = _make_executor_under_test() + req = _FakeRequest(py_request_id=99, + state="DISAGG_GENERATION_TRANS_IN_PROGRESS") + executor._do_terminate_request(req) + assert executor.terminated_request_ids == [] + + +def test_do_terminate_runs_when_safe(): + executor = _make_executor_under_test() + req = _FakeRequest(py_request_id=99, state="GENERATION_COMPLETE") + executor._do_terminate_request(req) + assert executor.terminated_request_ids == [99] + + +# --------------------------------------------------------------------------- +# ADP response flush +# --------------------------------------------------------------------------- + + +def test_flush_pending_transfer_responses_empties_buffer(): + executor = _make_executor_under_test() + executor._pending_transfer_responses = [(1, "r1"), (2, "r2")] + executor._flush_pending_transfer_responses() + + assert executor._pending_transfer_responses == [] + assert executor.enqueue_calls == [[(1, "r1"), (2, "r2")]] + + +def test_flush_with_adp_enabled_calls_enqueue_even_when_empty(): + """Even when the local rank has nothing to enqueue, ADP requires every + rank to participate in tp_gather. If we skipped the call we would + deadlock the other rank's collective.""" + executor = _make_executor_under_test() + executor.enable_attention_dp = True + executor._flush_pending_transfer_responses() + assert executor.enqueue_calls == [[]] + + +def test_flush_without_adp_skips_when_empty(): + executor = _make_executor_under_test() + executor.enable_attention_dp = False + executor._flush_pending_transfer_responses() + assert executor.enqueue_calls == [] + + +# --------------------------------------------------------------------------- +# Per-request timeout policy: do NOT fail-close the executor +# --------------------------------------------------------------------------- + + +def _apply_timeout_policy(executor, + request, + *, + cancel_result=_FakeTransferCancelResult. + CancelRequestedInFlight): + """Mini-replay of the timeout branch in _handle_responses for a single + request. Mirrors the production decision tree without dragging in the + full event loop.""" + request.py_kv_transfer_timed_out = True + executor.kv_cache_transceiver.next_cancel_results[ + request.py_request_id] = cancel_result + + cancel = executor.kv_cache_transceiver.cancel_request_structured(request) + safe_to_release = cancel in ( + _FakeTransferCancelResult.CancelledBeforeAdvertise, + _FakeTransferCancelResult.AlreadyComplete, + _FakeTransferCancelResult.NotFound, + ) + if safe_to_release: + executor._inflight_cancel_requested_ids.discard(request.py_request_id) + return "released" + if cancel == _FakeTransferCancelResult.CancelRequestedInFlight: + executor._inflight_cancel_requested_ids.add(request.py_request_id) + return "deferred" + if cancel == _FakeTransferCancelResult.BackendUnhealthy: + executor._inflight_cancel_requested_ids.add(request.py_request_id) + return "unhealthy-deferred" + return "ignored" + + +def test_timeout_with_in_flight_worker_does_not_clear_active_queue(): + executor = _make_executor_under_test() + other_req = _FakeRequest(py_request_id=1, state="GENERATION_IN_PROGRESS") + timed_out = _FakeRequest(py_request_id=2, + state="DISAGG_GENERATION_TRANS_IN_PROGRESS") + executor.active_requests = [other_req, timed_out] + + outcome = _apply_timeout_policy( + executor, + timed_out, + cancel_result=_FakeTransferCancelResult.CancelRequestedInFlight) + + # The PR 13706 anti-pattern was to clear active_requests / shutdown + # the executor on a per-request timeout. We must NOT do that. + assert outcome == "deferred" + assert other_req in executor.active_requests + assert timed_out in executor.active_requests + assert 2 in executor._inflight_cancel_requested_ids + # _do_terminate_request must defer. + executor._do_terminate_request(timed_out) + assert executor.terminated_request_ids == [] + + +def test_timeout_pre_advertise_releases_resources(): + executor = _make_executor_under_test() + timed_out = _FakeRequest(py_request_id=7, + state="DISAGG_CONTEXT_TRANS_IN_PROGRESS") + + outcome = _apply_timeout_policy( + executor, + timed_out, + cancel_result=_FakeTransferCancelResult.CancelledBeforeAdvertise) + + assert outcome == "released" + assert 7 not in executor._inflight_cancel_requested_ids + timed_out.state = "GENERATION_COMPLETE" + executor._do_terminate_request(timed_out) + assert executor.terminated_request_ids == [7] + + +def test_timeout_followed_by_worker_completion_eventually_releases(): + """Two-iteration sequence: first iteration the worker is in flight, + second iteration C++ has erased the future. Both must be handled + without ever calling _do_terminate_request before the worker + quiesces.""" + executor = _make_executor_under_test() + request = _FakeRequest(py_request_id=11, + state="DISAGG_GENERATION_TRANS_IN_PROGRESS") + + # Iteration 1: in flight. + assert _apply_timeout_policy( + executor, + request, + cancel_result=_FakeTransferCancelResult.CancelRequestedInFlight + ) == "deferred" + executor._do_terminate_request(request) + assert executor.terminated_request_ids == [] + + # Iteration 2: C++ has reached final state and erased the future. + request.state = "DISAGG_TRANS_ERROR" + assert _apply_timeout_policy( + executor, request, + cancel_result=_FakeTransferCancelResult.NotFound) == "released" + assert 11 not in executor._inflight_cancel_requested_ids + executor._do_terminate_request(request) + assert executor.terminated_request_ids == [11] + + +def test_unhealthy_transceiver_keeps_request_pinned(): + executor = _make_executor_under_test() + request = _FakeRequest(py_request_id=33, + state="DISAGG_GENERATION_TRANS_IN_PROGRESS") + + outcome = _apply_timeout_policy( + executor, + request, + cancel_result=_FakeTransferCancelResult.BackendUnhealthy) + + # Even when the C++ side reports unhealthy, Python must NOT free the + # request right away. Orchestration is supposed to restart the + # worker; freeing locally would leak buffers to the wedged backend. + assert outcome == "unhealthy-deferred" + assert 33 in executor._inflight_cancel_requested_ids + executor._do_terminate_request(request) + assert executor.terminated_request_ids == [] From aa53654b35f9451b1c7a7b5e1c8bf10c7736613d Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Tue, 5 May 2026 21:10:53 -0700 Subject: [PATCH 2/7] [https://nvbugs/6104831][fix] Pin LlmRequest with shared_ptr through the cache transceiver MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the previous commit on this branch. Changes the BaseCacheTransceiver / CacheTransceiver public API and the tracked- future vectors from raw LlmRequest* to std::shared_ptr so that the C++ object outlives every worker thread access regardless of when Python's _terminate_request drops its pybind reference. This closes the historical use-after-free class on raw LlmRequest* that showed up in field traces as mRequestId == 0x5555555555555555. Specifically: * respondAndSendAsync, requestAndReceiveSync, requestAndReceiveAsync, cancelRequest, and cancelRequestStructured all take std::shared_ptr. * TrackedFuture::request is now std::shared_ptr; the shared_ptr in mSenderFutures / mRequesterFutures keeps the LlmRequest pinned for as long as the transceiver tracks the worker future. * trtGptModelInflightBatching.cpp now passes the shared_ptr directly instead of stripping it via .get(). * The nanobind trampoline (PyCacheTransceiver) signatures match. The binding relies on (already included) to provide a Python-aware shared_ptr that pins the wrapper alive while C++ holds it. The Python-side _can_terminate_request_now guard stays in place. Its job is no longer about lifetime — the shared_ptr handles that — but about resource quiescence: free_resources() releases KV blocks back to the pool, and the transport may still be writing into them after the LlmRequest object exists. Termination must still wait for the structured cancel result to signal safety. The doc at docs/source/features/disagg-kv-transfer-session-lifecycle.md is updated to describe the new "Object lifetime" section. Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 36 ++++++++++++------- .../batch_manager/cacheTransceiver.cpp | 30 ++++++++-------- .../trtGptModelInflightBatching.cpp | 6 ++-- .../batch_manager/cacheTransceiver.cpp | 6 ++-- .../disagg-kv-transfer-session-lifecycle.md | 21 +++++++++++ 5 files changed, 67 insertions(+), 32 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index f4abcc149e4d..fbdc7719a0f3 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -258,13 +258,20 @@ class BaseCacheTransceiver { public: virtual ~BaseCacheTransceiver() = default; - virtual void respondAndSendAsync(LlmRequest* llmRequest) = 0; + + // Methods take std::shared_ptr so the transceiver and its + // worker threads pin the LlmRequest's lifetime independently of when + // Python's _terminate_request drops the pybind reference. This closes + // the use-after-free class on raw LlmRequest* observed in the field + // (mRequestId reads as 0x5555555555555555 after Python free) without + // changing the Python policy layer's contract. + 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; /// Non-blocking poll. Returns the requests that have completed or /// encountered an error since the last call. With @@ -300,7 +307,7 @@ class BaseCacheTransceiver /// contract. Subclasses override this; the boolean wrapper below maps /// the structured result onto the historical "is the request safe to /// release now" bool. - virtual TransferCancelResult cancelRequestStructured(LlmRequest* llmRequest) + virtual TransferCancelResult cancelRequestStructured(std::shared_ptr llmRequest) { return cancelRequest(llmRequest) ? TransferCancelResult::kCancelledBeforeAdvertise : TransferCancelResult::kNotFound; @@ -311,7 +318,7 @@ class BaseCacheTransceiver /// or already-complete worker). Callers that need to distinguish /// in-flight cancellation from "nothing to do" should use /// @ref cancelRequestStructured. - virtual bool cancelRequest(LlmRequest* llmRequest) = 0; + virtual bool cancelRequest(std::shared_ptr llmRequest) = 0; /// Whether the transceiver is currently healthy. Flips to false when /// quarantined transfers exceed the budget or when no backend @@ -356,13 +363,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; @@ -377,9 +384,9 @@ class CacheTransceiver : public BaseCacheTransceiver [[nodiscard]] bool checkGenTransferComplete() const override; - TransferCancelResult cancelRequestStructured(LlmRequest* llmRequest) override; + TransferCancelResult cancelRequestStructured(std::shared_ptr llmRequest) override; - virtual bool cancelRequest(LlmRequest* llmRequest) override; + virtual bool cancelRequest(std::shared_ptr llmRequest) override; [[nodiscard]] bool isHealthy() const override; @@ -393,10 +400,15 @@ class CacheTransceiver : public BaseCacheTransceiver /// Per-entry tracking for tracked worker futures. We keep the future /// pinned until the worker reaches a final state, even after the /// per-request deadline elapses, so that cancelling a still-running - /// worker never frees its KV/buffer resources prematurely. + /// worker never frees its KV/buffer resources prematurely. The + /// shared_ptr also pins the LlmRequest object's lifetime — Python + /// may drop its pybind reference at any time, but the C++ worker + /// thread (and any asynchronous queue inside CacheSender / + /// CacheReceiver) can dereference the request safely as long as the + /// transceiver has not erased this entry. struct TrackedFuture { - LlmRequest* request; + std::shared_ptr request; std::future future; std::chrono::steady_clock::time_point deadline; /// True once the per-request timeout has fired and we have flipped diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index ac50654e21b2..f0473d409613 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -323,7 +323,7 @@ void CacheTransceiver::setContextState(LlmRequest* llmRequest) } } -void CacheTransceiver::respondAndSendAsync(LlmRequest* llmRequest) +void CacheTransceiver::respondAndSendAsync(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isContextOnlyRequest()); llmRequest->setState(LlmRequestState::kDISAGG_CONTEXT_TRANS_IN_PROGRESS); @@ -337,7 +337,7 @@ void CacheTransceiver::respondAndSendAsync(LlmRequest* llmRequest) } return; } - setContextState(llmRequest); + setContextState(llmRequest.get()); auto future = mCacheSender->sendAsync(*llmRequest); auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); mSenderFutures.push_back(TrackedFuture{llmRequest, std::move(future), deadline, false, false}); @@ -357,11 +357,11 @@ void CacheTransceiver::respondAndSendLayerWise( setContextState(llmRequest.get()); auto future = mCacheSender->sendAsync(*llmRequest); auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); - mSenderFutures.push_back(TrackedFuture{llmRequest.get(), std::move(future), deadline, false, false}); + mSenderFutures.push_back(TrackedFuture{llmRequest, std::move(future), deadline, false, false}); } } -void CacheTransceiver::requestAndReceiveSync(LlmRequest* llmRequest) +void CacheTransceiver::requestAndReceiveSync(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isGenerationOnlyRequest()); { @@ -371,15 +371,16 @@ 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 reqId = llmRequest->mRequestId; if (std::find_if(mRequesterFutures.begin(), mRequesterFutures.end(), - [llmRequest](auto const& entry) { return entry.request->mRequestId == llmRequest->mRequestId; }) + [reqId](auto const& entry) { return entry.request->mRequestId == reqId; }) != mRequesterFutures.end()) { - TLLM_LOG_WARNING("Request ID %zu is already in mRequestFutures.", llmRequest->mRequestId); + TLLM_LOG_WARNING("Request ID %zu is already in mRequestFutures.", reqId); return; } @@ -822,7 +823,7 @@ void CacheTransceiver::checkGenTransferStatusImpl(std::optional const& atLe entry.request->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); if (!common::getEnvKVCacheTimeOutputPath().empty()) { - updateKVCacheTransferBW(syncComm, entry.request); + updateKVCacheTransferBW(syncComm, entry.request.get()); } } catch (std::exception const& e) @@ -844,7 +845,7 @@ void CacheTransceiver::checkGenTransferStatusImpl(std::optional const& atLe entry.request->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); if (!common::getEnvKVCacheTimeOutputPath().empty()) { - updateKVCacheTransferBW(syncComm, entry.request); + updateKVCacheTransferBW(syncComm, entry.request.get()); } } catch (std::exception const& e) @@ -877,7 +878,7 @@ bool CacheTransceiver::checkGenTransferComplete() const return mRequesterFutures.empty(); } -TransferCancelResult CacheTransceiver::cancelRequestStructured(LlmRequest* llmRequest) +TransferCancelResult CacheTransceiver::cancelRequestStructured(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest != nullptr); @@ -886,6 +887,7 @@ TransferCancelResult CacheTransceiver::cancelRequestStructured(LlmRequest* llmRe return TransferCancelResult::kBackendUnhealthy; } + auto const reqId = llmRequest->mRequestId; if (llmRequest->isContextOnlyRequest()) { // Pre-advertise: the sender's queue still owns the request. @@ -896,7 +898,7 @@ TransferCancelResult CacheTransceiver::cancelRequestStructured(LlmRequest* llmRe // Did the worker future already finish? for (size_t i = 0; i < mSenderFutures.size(); ++i) { - if (mSenderFutures[i].request->mRequestId != llmRequest->mRequestId) + if (mSenderFutures[i].request->mRequestId != reqId) { continue; } @@ -917,7 +919,7 @@ TransferCancelResult CacheTransceiver::cancelRequestStructured(LlmRequest* llmRe } for (size_t i = 0; i < mRequesterFutures.size(); ++i) { - if (mRequesterFutures[i].request->mRequestId != llmRequest->mRequestId) + if (mRequesterFutures[i].request->mRequestId != reqId) { continue; } @@ -933,9 +935,9 @@ TransferCancelResult CacheTransceiver::cancelRequestStructured(LlmRequest* llmRe return TransferCancelResult::kNotCancellable; } -bool CacheTransceiver::cancelRequest(LlmRequest* llmRequest) +bool CacheTransceiver::cancelRequest(std::shared_ptr llmRequest) { - auto result = cancelRequestStructured(llmRequest); + auto result = cancelRequestStructured(std::move(llmRequest)); return result == TransferCancelResult::kCancelledBeforeAdvertise || result == TransferCancelResult::kAlreadyComplete; } diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index 8bb2c0e2ba88..94679a54f022 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -915,7 +915,7 @@ void TrtGptModelInflightBatching::forwardSync() TLLM_CHECK_WITH_INFO(mCacheTransceiver, "Disaggregated serving is not enabled, please check the configuration of " "cacheTransceiverConfig."); - mCacheTransceiver->respondAndSendAsync(llmReq.get()); + mCacheTransceiver->respondAndSendAsync(llmReq); } mSeqSlotManager->freeSequenceSlot(llmReq->mRequestId); } @@ -1596,11 +1596,11 @@ void TrtGptModelInflightBatching::prepareDisaggGenInitRequests( mCacheTransceiver, "Disaggregated serving is not enabled, please check the configuration."); if (common::getEnvDisableKVCacheTransferOverlap()) { - mCacheTransceiver->requestAndReceiveSync(newGenReq.get()); + mCacheTransceiver->requestAndReceiveSync(newGenReq); } else { - mCacheTransceiver->requestAndReceiveAsync(newGenReq.get()); + mCacheTransceiver->requestAndReceiveAsync(newGenReq); } } if (!common::getEnvDisableKVCacheTransferOverlap()) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index 6b2db97ca7d5..bd02e750b9a2 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -48,7 +48,7 @@ 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); } @@ -61,12 +61,12 @@ class PyCacheTransceiver : public tb::BaseCacheTransceiver TLLM_THROW("respondAndSendLayerWise is not overridable from Python"); } - 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); } diff --git a/docs/source/features/disagg-kv-transfer-session-lifecycle.md b/docs/source/features/disagg-kv-transfer-session-lifecycle.md index f81f4438135f..ae7949cdbff3 100644 --- a/docs/source/features/disagg-kv-transfer-session-lifecycle.md +++ b/docs/source/features/disagg-kv-transfer-session-lifecycle.md @@ -22,6 +22,27 @@ mechanism that this codebase actually implements. | Resource quiescence | C++ transceiver | Worker future stays pinned until ready / error / quarantine. | | Process health | C++ transceiver | Quarantine budget + global progress deadline → `isHealthy() == false`. | +### Object lifetime + +The C++ transceiver takes ownership of the `LlmRequest` for the +duration of an in-flight transfer: + +* `BaseCacheTransceiver::respondAndSendAsync`, + `requestAndReceiveSync`, `requestAndReceiveAsync`, and `cancelRequest` + / `cancelRequestStructured` all take `std::shared_ptr`. +* The tracked-future vectors (`mSenderFutures`, `mRequesterFutures`) + store the same shared_ptr, so the C++ object outlives every worker + thread access regardless of when Python's `_terminate_request` drops + its pybind reference. + +This closes the historical "raw `LlmRequest*` use-after-free" class +(`mRequestId == 0x5555555555555555` field traces). The Python-side +`_can_terminate_request_now` guard remains in place — but its job is +now memory *quiescence*, not lifetime: even with the request object +pinned, `free_resources()` would release KV blocks back to the pool +while the transport may still be writing into them, so termination +must still wait for the structured cancel result to signal safety. + ## Status polling is non-blocking by default `BaseCacheTransceiver::checkContextTransferStatus` and From ea0c1b59026ec808d1111267b82e0f2ff40cfd9c Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Tue, 5 May 2026 22:02:40 -0700 Subject: [PATCH 3/7] [https://nvbugs/6104831][fix] Pin LlmRequest in dataTransceiver worker queues MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the prior commit. The worker queues inside CacheSender::Impl and CacheReceiver::Impl still held raw LlmRequest* pointers — meaning the executor's mSenderFutures / mRequesterFutures pinned the lifetime but the worker thread that actually dereferences the request had only a raw observer. If the executor erased its tracking entry while the worker was mid-flight, the LlmRequest could be freed under the worker. This commit closes that surface: * CacheSender::Impl::Response::mRequest → std::shared_ptr. * CacheReceiver::Impl::RequestAndPromise::mRequest → std::shared_ptr. Move/copy semantics simplified now that the field is a smart pointer. * CacheSender::sendAsync(LlmRequest&) → sendAsync(std::shared_ptr). * CacheReceiver::receiveAsync(LlmRequest&) → receiveAsync(std::shared_ptr). * CacheReceiver::Impl::requestAndReceiveAsyncMultiThreads similarly. * CacheReceiver::Impl::receiveAsync now captures the shared_ptr by value in the std::async lambda so the worker thread pins the LlmRequest independently of the caller. * CacheTransceiver::respondAndSendAsync/-LayerWise/requestAndReceive* pass the shared_ptr (no longer .get()-strip it) and move into the TrackedFuture entry where appropriate. Also fixes the eval-order UAF in CacheSender::Impl::handleAsyncSend that PR #13713 / PR #13728 called out: once Response::mRequest is a shared_ptr, the one-liner sendAndRemoveResponse(resp.mRequest->mRequestId, std::move(resp)); becomes undefined behaviour because C++ argument evaluation order is unspecified — the compiler may evaluate std::move(resp) first, leaving resp.mRequest empty when reading mRequestId. Materialise the id into a local before the move. The PyCacheTransceiver trampoline's bool cancelRequest override is updated to match the new shared_ptr signature. The doc at docs/source/features/disagg-kv-transfer-session-lifecycle.md spells out the full ownership chain: executor tracking + worker queue both hold shared_ptr, and TransferSession::mRequest is an ephemeral observer used only inside the worker frame where the shared_ptr is already held. Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.cpp | 12 +-- .../batch_manager/dataTransceiver.cpp | 91 +++++++++---------- .../batch_manager/dataTransceiver.h | 14 +-- .../batch_manager/cacheTransceiver.cpp | 2 +- .../disagg-kv-transfer-session-lifecycle.md | 29 ++++-- 5 files changed, 80 insertions(+), 68 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index f0473d409613..03726ece6a72 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -338,9 +338,9 @@ void CacheTransceiver::respondAndSendAsync(std::shared_ptr llmReques return; } setContextState(llmRequest.get()); - auto future = mCacheSender->sendAsync(*llmRequest); + auto future = mCacheSender->sendAsync(llmRequest); auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); - mSenderFutures.push_back(TrackedFuture{llmRequest, std::move(future), deadline, false, false}); + mSenderFutures.push_back(TrackedFuture{std::move(llmRequest), std::move(future), deadline, false, false}); } void CacheTransceiver::respondAndSendLayerWise( @@ -355,7 +355,7 @@ void CacheTransceiver::respondAndSendLayerWise( llmRequest->setState(LlmRequestState::kDISAGG_CONTEXT_INIT_AND_TRANS); setContextState(llmRequest.get()); - auto future = mCacheSender->sendAsync(*llmRequest); + auto future = mCacheSender->sendAsync(llmRequest); auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); mSenderFutures.push_back(TrackedFuture{llmRequest, std::move(future), deadline, false, false}); } @@ -365,7 +365,7 @@ void CacheTransceiver::requestAndReceiveSync(std::shared_ptr llmRequ { TLLM_CHECK(llmRequest && llmRequest->isGenerationOnlyRequest()); { - auto future = mCacheReceiver->receiveAsync(*llmRequest); + auto future = mCacheReceiver->receiveAsync(llmRequest); future.get(); } llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); @@ -384,10 +384,10 @@ void CacheTransceiver::requestAndReceiveAsync(std::shared_ptr llmReq return; } - auto future = mCacheReceiver->receiveAsync(*llmRequest); + auto future = mCacheReceiver->receiveAsync(llmRequest); auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); - mRequesterFutures.push_back(TrackedFuture{llmRequest, std::move(future), deadline, false, false}); llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS); + mRequesterFutures.push_back(TrackedFuture{std::move(llmRequest), std::move(future), deadline, false, false}); } std::vector gatherRequestIds( diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 3ecceb9f3f2c..5a8545e17bb6 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -296,16 +296,17 @@ class CacheSender::Impl } } - [[nodiscard]] std::future sendAsync(LlmRequest& llmRequest) + [[nodiscard]] std::future sendAsync(std::shared_ptr llmRequest) { + TLLM_CHECK(llmRequest); std::promise promise; auto future = promise.get_future(); - llmRequest.setKvCacheTransferStart(LlmRequest::getSteadyClockNow()); + llmRequest->setKvCacheTransferStart(LlmRequest::getSteadyClockNow()); + auto const reqId = llmRequest->mRequestId; { { std::scoped_lock lkResp(mSenderMutex); - mReadyResponses.emplace( - llmRequest.mRequestId, Response{std::addressof(llmRequest), std::move(promise)}); + mReadyResponses.emplace(reqId, Response{std::move(llmRequest), std::move(promise)}); } std::unique_lock lkCond(mCondMutex); mAnyReady = true; @@ -477,7 +478,11 @@ class CacheSender::Impl private: struct Response { - LlmRequest* mRequest; + // shared_ptr so the LlmRequest stays alive as long as the worker + // queue holds the response, independently of the executor's + // mSenderFutures tracking. Closes the raw-pointer UAF surface in + // the worker thread. + std::shared_ptr mRequest; std::promise mPromise; }; @@ -511,7 +516,12 @@ class CacheSender::Impl resp = std::move(resource.mSendQueue.front()); resource.mSendQueue.pop_front(); } - sendAndRemoveResponse(resp.mRequest->mRequestId, std::move(resp)); + // Materialize the request id before std::move(resp) — C++ + // argument evaluation order is unspecified, so reading + // resp.mRequest->mRequestId after the move would dereference + // an empty shared_ptr. + auto const reqId = resp.mRequest->mRequestId; + sendAndRemoveResponse(reqId, std::move(resp)); } } @@ -753,23 +763,29 @@ class CacheReceiver::Impl TLLM_CUDA_CHECK(cudaGetDevice(&mDeviceId)); } - [[nodiscard]] std::future receiveAsync(LlmRequest& llmRequest) + [[nodiscard]] std::future receiveAsync(std::shared_ptr llmRequest) { + TLLM_CHECK(llmRequest); // TODO: Modify the implementation here to avoid frequent thread creation. - return std::async(std::launch::async, &CacheReceiver::Impl::requestSync, this, std::ref(llmRequest)); + // The worker captures the shared_ptr by value, pinning the LlmRequest + // for the lifetime of the request worker independently of the + // executor-side tracking. + return std::async(std::launch::async, + [this, llmRequest = std::move(llmRequest)]() mutable { this->requestSync(*llmRequest); }); } - [[nodiscard]] std::future requestAndReceiveAsyncMultiThreads(LlmRequest& llmRequest) + [[nodiscard]] std::future requestAndReceiveAsyncMultiThreads(std::shared_ptr llmRequest) { + TLLM_CHECK(llmRequest); try { auto promise = std::make_unique>(); auto future = promise->get_future(); - TLLM_CHECK(llmRequest.getDataTransceiverState().getCommState().has_value()); + TLLM_CHECK(llmRequest->getDataTransceiverState().getCommState().has_value()); std::string processInfo = kDefaultProcessInfo; if (common::getEnvRequestKVCacheConcurrent()) { - processInfo = llmRequest.getDataTransceiverState().getCommState()->toString(); + processInfo = llmRequest->getDataTransceiverState().getCommState()->toString(); } if (mInstanceToAsyncResource.find(processInfo) == mInstanceToAsyncResource.end()) { @@ -782,7 +798,7 @@ class CacheReceiver::Impl auto& asyncResource = mInstanceToAsyncResource.at(processInfo); { std::unique_lock lck(asyncResource->mMtxForQueue); - asyncResource->mRequestsQueue.emplace_back(std::addressof(llmRequest), std::move(promise)); + asyncResource->mRequestsQueue.emplace_back(std::move(llmRequest), std::move(promise)); } asyncResource->mCVforQueue.notify_all(); return future; @@ -1074,47 +1090,24 @@ class CacheReceiver::Impl struct RequestAndPromise { - LlmRequest* mRequest; + // shared_ptr so the LlmRequest stays alive while the worker queue + // holds it, independently of the executor's mRequesterFutures + // tracking. + std::shared_ptr mRequest; std::unique_ptr> mPromise; - RequestAndPromise() - : mRequest(nullptr) - , mPromise(nullptr) - { - } + RequestAndPromise() = default; - RequestAndPromise(LlmRequest* request, std::unique_ptr>&& promise) - : mRequest(request) + RequestAndPromise(std::shared_ptr request, std::unique_ptr>&& promise) + : mRequest(std::move(request)) , mPromise(std::move(promise)) { } RequestAndPromise(RequestAndPromise const&) = delete; - - RequestAndPromise(RequestAndPromise&& other) noexcept - : mRequest(other.mRequest) - , mPromise(std::move(other.mPromise)) - { - other.mRequest = nullptr; - } - - RequestAndPromise& operator=(RequestAndPromise&& other) noexcept - { - if (this != &other) - { - mRequest = nullptr; - if (mPromise) - { - mPromise.reset(); - } - - mRequest = other.mRequest; - mPromise = std::move(other.mPromise); - - other.mRequest = nullptr; - } - return *this; - } + RequestAndPromise& operator=(RequestAndPromise const&) = delete; + RequestAndPromise(RequestAndPromise&&) noexcept = default; + RequestAndPromise& operator=(RequestAndPromise&&) noexcept = default; }; struct AsyncResource @@ -1209,9 +1202,9 @@ CacheSender::CacheSender( { } -std::future CacheSender::sendAsync(LlmRequest& llmRequest) const +std::future CacheSender::sendAsync(std::shared_ptr llmRequest) const { - return mImpl->sendAsync(llmRequest); + return mImpl->sendAsync(std::move(llmRequest)); } executor::kv_cache::CommState const& CacheSender::getCommState() const @@ -1252,9 +1245,9 @@ CacheReceiver::CacheReceiver( { } -std::future CacheReceiver::receiveAsync(LlmRequest& llmRequest) const +std::future CacheReceiver::receiveAsync(std::shared_ptr llmRequest) const { - return mImpl->requestAndReceiveAsyncMultiThreads(llmRequest); + return mImpl->requestAndReceiveAsyncMultiThreads(std::move(llmRequest)); } CacheReceiver::~CacheReceiver() = default; diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h index 4d072a04521b..d81207824829 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h @@ -256,10 +256,11 @@ class CacheSender CacheSender() = default; /// @brief Asynchronously respond to the request and send data. - /// @param llmRequest Request object. Its data should be ready when called, and the data for this request - /// should remain valid until future synchronization. + /// @param llmRequest Request object. The transceiver pins the LlmRequest + /// for the lifetime of the worker future via the shared_ptr, so the + /// caller may drop its own strong reference once this returns. /// @return Once the data is fully sent, the future object will become valid. - [[nodiscard]] virtual std::future sendAsync(LlmRequest& llmRequest) const; + [[nodiscard]] virtual std::future sendAsync(std::shared_ptr llmRequest) const; /// @brief Return the internal communicator status. /// @return The communicator status. @@ -313,10 +314,11 @@ class CacheReceiver CacheReceiver() = default; /// @brief Asynchronously send a request to receive data. - /// @param llmRequest Request object. Its data should be in an allocated but unwritten state when called, and the - /// data for this request should remain intact only after future synchronization. + /// @param llmRequest Request object. The transceiver pins the LlmRequest + /// for the lifetime of the worker future via the shared_ptr, so the + /// caller may drop its own strong reference once this returns. /// @return Once the data is fully received, the future object will become valid. - [[nodiscard]] virtual std::future receiveAsync(LlmRequest& llmRequest) const; + [[nodiscard]] virtual std::future receiveAsync(std::shared_ptr llmRequest) const; virtual TransferSession sendRequestInfo(LlmRequest const& llmRequest); diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index bd02e750b9a2..cf02471ff473 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -87,7 +87,7 @@ class PyCacheTransceiver : public tb::BaseCacheTransceiver NB_OVERRIDE_PURE(checkGenTransferComplete); } - bool cancelRequest(tb::LlmRequest* llmRequest) override + bool cancelRequest(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(cancelRequest, llmRequest); } diff --git a/docs/source/features/disagg-kv-transfer-session-lifecycle.md b/docs/source/features/disagg-kv-transfer-session-lifecycle.md index ae7949cdbff3..c072d663d754 100644 --- a/docs/source/features/disagg-kv-transfer-session-lifecycle.md +++ b/docs/source/features/disagg-kv-transfer-session-lifecycle.md @@ -31,18 +31,35 @@ duration of an in-flight transfer: `requestAndReceiveSync`, `requestAndReceiveAsync`, and `cancelRequest` / `cancelRequestStructured` all take `std::shared_ptr`. * The tracked-future vectors (`mSenderFutures`, `mRequesterFutures`) - store the same shared_ptr, so the C++ object outlives every worker - thread access regardless of when Python's `_terminate_request` drops - its pybind reference. - -This closes the historical "raw `LlmRequest*` use-after-free" class -(`mRequestId == 0x5555555555555555` field traces). The Python-side + store the same shared_ptr. +* Each worker queue **also** stores a `std::shared_ptr` — + `CacheSender::Impl::Response::mRequest` and + `CacheReceiver::Impl::RequestAndPromise::mRequest`. The worker thread + thus pins the LlmRequest independently of the executor, so a wedged + worker cannot be freed by the executor erasing its tracking entry. + +Together this closes the historical "raw `LlmRequest*` use-after-free" +class (`mRequestId == 0x5555555555555555` field traces). The Python-side `_can_terminate_request_now` guard remains in place — but its job is now memory *quiescence*, not lifetime: even with the request object pinned, `free_resources()` would release KV blocks back to the pool while the transport may still be writing into them, so termination must still wait for the structured cancel result to signal safety. +`CacheSender::Impl::handleAsyncSend` materialises the request id into +a local before `std::move(resp)` to avoid a C++ argument-eval-order +bug — once `Response::mRequest` is a `shared_ptr`, the previous +one-liner `sendAndRemoveResponse(resp.mRequest->mRequestId, std::move(resp))` +became undefined behaviour because the compiler may evaluate the move +first and leave `resp.mRequest` empty when reading `mRequestId`. Field +traces of PR #13713 had a SIGSEGV at exactly that argument +construction; the materialised-id pattern fixes it. + +`TransferSession::mRequest` is an ephemeral `LlmRequest const*` that is +only set inside `sendSync` and used during measurement / formatter +calls within the same worker frame; the worker still holds the +shared_ptr at that point, so the observer is safe by construction. + ## Status polling is non-blocking by default `BaseCacheTransceiver::checkContextTransferStatus` and From d7d189a108f24b6ec19c8be73905d46c8fb011e7 Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Tue, 5 May 2026 22:09:22 -0700 Subject: [PATCH 4/7] [https://nvbugs/6104831][fix] Audit fixes: cancel ordering, trampoline size, test callsites MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three issues caught in self-review of the prior commits: 1. cancelRequestStructured returned BackendUnhealthy too eagerly — even for a request that was still queued pre-advertise. Pre-advertise release is always safe (no buffer was ever exposed to a peer), and freeing those resources during the unhealthy window REDUCES backend pressure rather than adds to it. Reorder: check sender / receiver queues first, return CancelledBeforeAdvertise unconditionally if found there. Only consult isHealthy() for the in-flight branch, where freeing would race with NIXL/UCX writes. 2. NB_TRAMPOLINE size was 6 but the trampoline has 7 NB_OVERRIDE_PURE entries (pre-existing under-allocation of nanobind's dispatch hash table). Bump to 7 since we touched the file. 3. cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp called mSender->sendAsync(*llmRequest) and mRequester->receiveAsync(*llmRequest) passing LlmRequest&. After the prior commit those signatures take std::shared_ptr. The test was already holding a shared_ptr (request->mLlmRequest), so just drop the dereference. Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.cpp | 71 ++++++++----------- .../batch_manager/cacheTransceiver.cpp | 4 +- .../multi_gpu/cacheTransceiverTest.cpp | 8 +-- 3 files changed, 38 insertions(+), 45 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 03726ece6a72..70f7d2736b85 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -882,57 +882,48 @@ TransferCancelResult CacheTransceiver::cancelRequestStructured(std::shared_ptrmRequestId; + bool const sender = llmRequest->isContextOnlyRequest(); + bool const receiver = llmRequest->isGenerationOnlyRequest(); + if (!sender && !receiver) { - return TransferCancelResult::kBackendUnhealthy; + return TransferCancelResult::kNotCancellable; } - auto const reqId = llmRequest->mRequestId; - if (llmRequest->isContextOnlyRequest()) + // Pre-advertise check first: a request that is still queued in the + // sender / receiver's pending list has not exposed any buffer to a + // peer and is always safe to release — even when the transceiver + // is otherwise unhealthy. Releasing pre-advertise requests during + // an unhealthy window reduces the resource pressure on the wedged + // backend rather than adding to it. + if (sender && mCacheSender->cancelRequest(*llmRequest)) { - // Pre-advertise: the sender's queue still owns the request. - if (mCacheSender->cancelRequest(*llmRequest)) - { - return TransferCancelResult::kCancelledBeforeAdvertise; - } - // Did the worker future already finish? - for (size_t i = 0; i < mSenderFutures.size(); ++i) - { - if (mSenderFutures[i].request->mRequestId != reqId) - { - continue; - } - auto status = mSenderFutures[i].future.wait_for(std::chrono::milliseconds(0)); - if (status == std::future_status::ready) - { - return TransferCancelResult::kAlreadyComplete; - } - return TransferCancelResult::kCancelRequestedInFlight; - } - return TransferCancelResult::kNotFound; + return TransferCancelResult::kCancelledBeforeAdvertise; } - if (llmRequest->isGenerationOnlyRequest()) + if (receiver && mCacheReceiver->cancelRequest(*llmRequest)) { - if (mCacheReceiver->cancelRequest(*llmRequest)) + return TransferCancelResult::kCancelledBeforeAdvertise; + } + + auto& futures = sender ? mSenderFutures : mRequesterFutures; + for (size_t i = 0; i < futures.size(); ++i) + { + if (futures[i].request->mRequestId != reqId) { - return TransferCancelResult::kCancelledBeforeAdvertise; + continue; } - for (size_t i = 0; i < mRequesterFutures.size(); ++i) + auto status = futures[i].future.wait_for(std::chrono::milliseconds(0)); + if (status == std::future_status::ready) { - if (mRequesterFutures[i].request->mRequestId != reqId) - { - continue; - } - auto status = mRequesterFutures[i].future.wait_for(std::chrono::milliseconds(0)); - if (status == std::future_status::ready) - { - return TransferCancelResult::kAlreadyComplete; - } - return TransferCancelResult::kCancelRequestedInFlight; + return TransferCancelResult::kAlreadyComplete; } - return TransferCancelResult::kNotFound; + // Worker is mid-flight. If the transceiver is unhealthy the + // caller must defer cleanup pending an orchestration restart; + // otherwise it is just a normal in-flight cancel. + return isHealthy() ? TransferCancelResult::kCancelRequestedInFlight + : TransferCancelResult::kBackendUnhealthy; } - return TransferCancelResult::kNotCancellable; + return TransferCancelResult::kNotFound; } bool CacheTransceiver::cancelRequest(std::shared_ptr llmRequest) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index cf02471ff473..9b5aa502b2b3 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -46,7 +46,9 @@ class PyCacheTransceiver : public tb::BaseCacheTransceiver { public: // using BaseCacheTransceiver::BaseCacheTransceiver; // Inherit constructors - NB_TRAMPOLINE(tb::BaseCacheTransceiver, 6); + // 7 NB_OVERRIDE_PURE entries below; the historical value was 6 which + // under-allocates the dispatch table (latent pre-existing bug). + NB_TRAMPOLINE(tb::BaseCacheTransceiver, 7); void respondAndSendAsync(std::shared_ptr llmRequest) override { diff --git a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp index fae2e7d3a365..8d0e11f82349 100644 --- a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp @@ -347,11 +347,11 @@ class SymmetricalCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines- TLLM_CUDA_CHECK(cudaMemset(it->data(), llmRequest->getPromptLen(), it->getSizeInBytes())); } } - mFutures.emplace_back(mSender->sendAsync(*llmRequest)); + mFutures.emplace_back(mSender->sendAsync(llmRequest)); } else { - auto future = mRequester->receiveAsync(*llmRequest); + auto future = mRequester->receiveAsync(llmRequest); future.get(); TLLM_CUDA_CHECK(cudaDeviceSynchronize()); auto blockRange = BlockRange::fromAllBlockIds(*mManager, llmRequest->mRequestId); @@ -970,7 +970,7 @@ class AsymmetricalCacheTest : public ::testing::TestWithParamsendAsync(*llmRequest); + auto future = mSender->sendAsync(llmRequest); return future; } @@ -980,7 +980,7 @@ class AsymmetricalCacheTest : public ::testing::TestWithParammLlmRequest; mManager->addSequence(llmRequest->mRequestId, llmRequest->getNumTokens(beamIdx), beamWidth, llmRequest); - return mRequester->receiveAsync(*llmRequest); + return mRequester->receiveAsync(llmRequest); } void generationVerifyKVCache(std::shared_ptr const& request) From ec5dfffb4742794bde0c7053c3380098e18f6159 Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Wed, 6 May 2026 10:17:13 -0700 Subject: [PATCH 5/7] [https://nvbugs/6104831][fix] Close timeout recovery loop: deferred KV release + health-driven shutdown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The earlier commits surfaced the user-visible "exceeded timeout" error and pinned the LlmRequest object via shared_ptr, but the recovery story still had two gaps that an internal review caught: 1. When the C++ deadline-quarantine path marked a request as DISAGG_TRANS_ERROR, _can_terminate_request_now no longer recognised the request as "still in flight" (its state was no longer DISAGG_*_TRANS_IN_PROGRESS), so _do_terminate_request happily called free_resources() — releasing KV blocks back to the cache pool while the C++ worker thread might still be writing into them via NIXL/UCX. The shared_ptr lifetime fix protected the LlmRequest object but not its KV memory. 2. C++ exposed isHealthy() / getHealth() but nothing on the Python side consumed those signals. A worker whose backend was genuinely wedged would silently accumulate quarantined transfers forever, never triggering an orchestration restart. This commit closes both: * PyExecutor adds _pending_resource_release (a list of LlmRequest whose KV cleanup is deferred) and _maybe_release_pending_resources which polls cancel_request_structured each iteration. free_resources runs only when the structured cancel result transitions to AlreadyComplete / NotFound / CancelledBeforeAdvertise — the three outcomes that prove the worker has reached a final state. * The timeout error paths in _check_disagg_ctx_cache_transfer_status and _handle_responses (gen side) call _defer_resource_release_for_inflight_transfer() before _handle_errors / _end_transfer_and_maybe_terminate, so the user-visible error response is still surfaced immediately while the actual free_resources() call defers until C++ quiesces. * PyExecutor adds _check_transceiver_health which monitors transceiver.is_healthy(). It records the first-unhealthy timestamp and, after a grace period (default 2 * global_progress_deadline_- seconds, falling back to 120s), sets self.is_shutdown so orchestration's existing restart path takes the worker out of service. Recovery from a transient unhealthy window does not trigger restart. * Both new methods plus the existing _check_kv_transfer_timeout are called at every executor-loop tick that already checks for KV transfer timeouts — five sites across PP / non-overlap / overlap loops. The doc gains a "Recovery model" section that spells out the three timescales: user-visible error (immediate), KV resource release (bounded by worker quiescence), worker restart on persistent wedge (global deadline + grace). Tests: * Five new unit tests in test_kv_transfer_session_lifecycle.py: - test_defer_resource_release_holds_until_quiesced (full sequence timeout → in-flight → still in-flight → quiesced → freed). - test_defer_resource_release_handles_not_found_as_safe. - test_defer_resource_release_no_op_when_empty. - test_check_transceiver_health_resets_on_recovery (transient unhealthy must not shut down). - test_check_transceiver_health_triggers_shutdown_after_grace (sustained unhealthy DOES shut down). * All 19 tests pass without GPU/MPI. Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../disagg-kv-transfer-session-lifecycle.md | 15 ++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 216 +++++++++++++++--- .../test_kv_transfer_session_lifecycle.py | 158 +++++++++++++ 3 files changed, 355 insertions(+), 34 deletions(-) diff --git a/docs/source/features/disagg-kv-transfer-session-lifecycle.md b/docs/source/features/disagg-kv-transfer-session-lifecycle.md index c072d663d754..ec5faf9dfda6 100644 --- a/docs/source/features/disagg-kv-transfer-session-lifecycle.md +++ b/docs/source/features/disagg-kv-transfer-session-lifecycle.md @@ -112,6 +112,21 @@ locally. `seconds_since_last_progress`, and `global_progress_deadline_seconds` for metrics / observability. +## Recovery model + +When a per-request KV transfer exceeds `kvTransferTimeoutMs`, three +things must happen, on three different timescales: + +| Concern | Path | Timing | +|---|---|---| +| **User-visible error response** | C++ `maybeQuarantineLocked` adds the request to `errorRequestIds`; Python's `_check_disagg_ctx_cache_transfer_status` / `_handle_responses` surfaces the error via `_handle_errors` / `_end_transfer_and_maybe_terminate`. | Immediate (next polling iteration). | +| **KV / buffer resource release** | Python registers the request in `_pending_resource_release` and `_inflight_cancel_requested_ids`. Each iteration `_maybe_release_pending_resources` polls `cancelRequestStructured`; only when it returns `AlreadyComplete` / `NotFound` / `CancelledBeforeAdvertise` does Python actually call `free_resources()`. The C++ shared_ptr in `mSenderFutures` keeps the LlmRequest alive until then. | Bounded by worker quiescence — could be the rest of `kvTransferTimeoutMs`, or longer if the backend is genuinely wedged. | +| **Worker restart on persistent wedge** | C++ `getHealth().is_healthy` flips false when the quarantine budget overflows or the global progress deadline elapses. Python's `_check_transceiver_health` watches that flag and sets `is_shutdown = True` after a grace period (default `2 * mGlobalProgressDeadlineMs`). Orchestration's existing process-restart path / k8s liveness probe / `disagg_auto_scaling.py` then restarts the worker. | Bounded by the global progress deadline + grace period (default 60 s + 120 s). | + +This separation is what makes the recovery deterministic without +either freeing memory the transport may still write into or hanging +the user's request indefinitely. + ## Python policy The `PyExecutor` keeps these responsibilities only: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 8d6eb1304334..e7a84ddd61cd 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -429,6 +429,20 @@ def __init__( # writing into them. Keyed by py_request_id to survive request # state transitions. self._inflight_cancel_requested_ids: set[int] = set() + # Requests whose timeout error response has already been surfaced + # to the user but whose KV / buffer resources cannot yet be + # released because the C++ worker has not quiesced. The request + # is no longer in active_requests, so the executor does not see + # it again — _maybe_release_pending_resources is the only place + # that polls C++ and triggers the deferred free once the worker + # reaches a final state. + self._pending_resource_release: List[LlmRequest] = [] + # Wall-clock deadline tracker for the transceiver-unhealthy + # signal. None while the transceiver is healthy; set when + # transceiver.is_healthy() first flips false. If it stays false + # longer than the configured grace period we set is_shutdown + # so orchestration restarts the worker. + self._transceiver_unhealthy_since: Optional[float] = None # TODO: Remove the condition on the PP size once disagg support from KVCache reuse # path is fixed. self.async_transfer_manager = AsyncTransferManager( @@ -683,6 +697,100 @@ def _flush_pending_transfer_responses(self): if responses or self.enable_attention_dp: self._enqueue_responses(responses) + def _defer_resource_release_for_inflight_transfer( + self, request: LlmRequest) -> None: + """Mark a timed-out transfer's resources for deferred release. + + Called from the timeout-error paths when C++ structured cancel + reports the worker is still in flight. The request has already + been removed from active_requests by _handle_errors and the user + has been told the request errored; we still must not call + free_resources() until the worker quiesces, otherwise NIXL/UCX + could write into KV blocks that have been returned to the pool. + """ + req_id = request.py_request_id + self._inflight_cancel_requested_ids.add(req_id) + if request not in self._pending_resource_release: + self._pending_resource_release.append(request) + + def _maybe_release_pending_resources(self): + """Drain _pending_resource_release for requests whose C++ worker + has finally reached a final state. + + Each iteration we ask the transceiver whether each pending + request is still in flight. AlreadyComplete / NotFound mean the + worker has quiesced (future was either ready or already erased + by check_*_transfer_status), so it is now safe to free + request-owned KV blocks. Any other result keeps the request on + the deferred list — we will retry next iteration. + """ + if (not self._pending_resource_release + or self.kv_cache_transceiver is None): + return + still_pending: List[LlmRequest] = [] + for request in self._pending_resource_release: + cancel = self.kv_cache_transceiver.cancel_request_structured( + request) + if cancel in (TransferCancelResult.AlreadyComplete, + TransferCancelResult.NotFound, + TransferCancelResult.CancelledBeforeAdvertise): + req_id = request.py_request_id + self._inflight_cancel_requested_ids.discard(req_id) + # _do_terminate_request now passes the guard because the + # request is no longer in _inflight_cancel_requested_ids + # and its state is no longer DISAGG_*_TRANS_IN_PROGRESS. + self._do_terminate_request(request) + else: + still_pending.append(request) + self._pending_resource_release = still_pending + + def _check_transceiver_health(self): + """Watch the C++ transceiver's health signal and trigger an + executor shutdown if it stays unhealthy past the grace period. + + The C++ side flips `is_healthy() == False` only after the + quarantine budget has been exceeded or no worker has reached a + final state for longer than mGlobalProgressDeadlineMs. That is + already a strong signal; we wait an additional grace period + (default: 2 * deadline) before forcing shutdown so that brief + transient unhealthy windows during heavy cancellation traffic do + not cause a restart loop. + """ + if self.kv_cache_transceiver is None: + return + if self.kv_cache_transceiver.is_healthy(): + self._transceiver_unhealthy_since = None + return + now = time.time() + if self._transceiver_unhealthy_since is None: + self._transceiver_unhealthy_since = now + health = self.kv_cache_transceiver.get_health() + logger.warning( + f"KV cache transceiver flipped UNHEALTHY at {now:.3f}: " + f"quarantined={getattr(health, 'quarantined_transfer_count', '?')} " + f"budget={getattr(health, 'quarantine_budget', '?')} " + f"seconds_since_last_progress=" + f"{getattr(health, 'seconds_since_last_progress', '?')}. " + f"Will trigger executor shutdown if it stays unhealthy.") + return + elapsed = now - self._transceiver_unhealthy_since + # Grace period: 2x the global progress deadline reported by C++. + # Fall back to 120s if the snapshot is unavailable. + try: + health = self.kv_cache_transceiver.get_health() + grace_period = 2.0 * float( + getattr(health, 'global_progress_deadline_seconds', 60.0)) + except Exception: + grace_period = 120.0 + if elapsed > grace_period and not self.is_shutdown: + logger.error( + f"KV cache transceiver has been UNHEALTHY for {elapsed:.1f}s " + f"(grace period {grace_period:.1f}s). Setting is_shutdown so " + f"orchestration can restart the worker.") + self.is_shutdown = True + if self.shutdown_event is not None: + self.shutdown_event.set() + # Performance metrics methods are in PerfMetricsManager (self.perf_manager) def _event_loop_wrapper(self): @@ -1367,6 +1475,8 @@ def _pp_retry_until_can_schedule(self, scheduled_batch): # Let cache transceiver finish at least one cache transmission and release requests' KV cache resources self._check_disagg_ctx_cache_transfer_status(1) self._check_kv_transfer_timeout() + self._maybe_release_pending_resources() + self._check_transceiver_health() else: raise RuntimeError( f"Reach maximum PP retry count ({self.pp_scheduler_max_retry_count}) but still cannot run first PP's schedule result. Please consider increasing the KV cache size by setting `free_gpu_memory_fraction` to a larger value. Or you can set `TLLM_PP_SCHEDULER_MAX_RETRY_COUNT` to a larger value to allow more retries." @@ -1759,6 +1869,8 @@ def _handle_executed_batch(self, executed_batch: Optional[BatchStatePP]): if self.kv_cache_transceiver and self.async_transfer_manager.has_any_inflight_requests( ): self._check_kv_transfer_timeout() + self._maybe_release_pending_resources() + self._check_transceiver_health() if self._disagg_pp_termination_handler is not None: self._disagg_pp_termination_handler.terminate_pending_requests() @@ -1867,6 +1979,8 @@ def _prepare_and_schedule_batch(self): self._check_disagg_ctx_schedulable_status(new_requests) self._check_disagg_gen_transfer_status() self._check_kv_transfer_timeout() + self._maybe_release_pending_resources() + self._check_transceiver_health() iter_stats = None if self.enable_iter_perf_stats: @@ -2236,6 +2350,8 @@ def _executor_loop(self): if self.kv_cache_transceiver and self.async_transfer_manager.has_any_inflight_requests( ): self._check_kv_transfer_timeout() + self._maybe_release_pending_resources() + self._check_transceiver_health() self._kv_connector_terminate_requests() @@ -2529,6 +2645,8 @@ def _executor_loop_overlap(self): if self.kv_cache_transceiver and self.async_transfer_manager.has_any_inflight_requests( ): self._check_kv_transfer_timeout() + self._maybe_release_pending_resources() + self._check_transceiver_health() self._kv_connector_terminate_requests() @@ -3374,10 +3492,12 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): requests_in_transfer = self.async_transfer_manager.requests_in_transfer( ) - # Worker reached a final state — clear any pending in-flight cancel - # bookkeeping for those requests. - if completed_req_ids: - self._inflight_cancel_requested_ids -= completed_req_ids + # Worker reached a final state for entries that came back as + # completed (future ready). Quarantined-but-error entries have + # NOT actually quiesced; we keep their inflight-cancel bookkeeping. + finished_set = set(finished_requests) + if finished_set: + self._inflight_cancel_requested_ids -= finished_set for request_id in completed_req_ids: @@ -3388,6 +3508,23 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): request = requests_in_transfer[request_id] + # If the request errored because of timeout / quarantine + # rather than worker completion, the C++ worker may still be + # in flight (the future is pinned in mSenderFutures). Defer + # the free_resources() call by adding the request to + # _pending_resource_release before _end_transfer_and_maybe_- + # terminate hits _do_terminate_request. The deferred entry + # will be drained by _maybe_release_pending_resources once + # the structured-cancel result transitions to a safe state. + if request_id in set(error_requests): + cancel = self.kv_cache_transceiver.cancel_request_structured( + request) + if cancel in ( + TransferCancelResult.CancelRequestedInFlight, + TransferCancelResult.BackendUnhealthy): + self._defer_resource_release_for_inflight_transfer( + request) + self._end_transfer_and_maybe_terminate(request) # The set of requests in transfer may have changed since we terminated some requests. @@ -3417,22 +3554,26 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE self._inflight_cancel_requested_ids.discard(request_id) self._end_transfer_and_maybe_terminate(request) - elif cancel_result == TransferCancelResult.CancelRequestedInFlight: - # Cancellation accepted but the worker is still - # touching the buffer/KV memory. Mark the request as - # in-flight-cancel so _do_terminate_request defers - # freeing resources, and re-check next iteration. - self._inflight_cancel_requested_ids.add(request_id) - elif cancel_result == TransferCancelResult.BackendUnhealthy: - # Soft signal — the transceiver is unhealthy. We do - # not touch this request; orchestration health checks - # will restart the worker once the global deadline - # elapses. - logger.warning( - f"Transceiver reports unhealthy while cancelling " - f"context request {request_id}; deferring to " - f"health-driven restart.") - self._inflight_cancel_requested_ids.add(request_id) + elif cancel_result in ( + TransferCancelResult.CancelRequestedInFlight, + TransferCancelResult.BackendUnhealthy): + # Worker is still touching buffers (or transceiver is + # unhealthy). Defer the actual free_resources() to + # _maybe_release_pending_resources by registering the + # request BEFORE the cleanup path. The user-visible + # error response is still surfaced via the normal + # transfer-error flow on a subsequent iteration once + # C++ confirms quiescence. + request.py_kv_transfer_start_time = None + request.state = LlmRequestState.DISAGG_TRANS_ERROR + self._defer_resource_release_for_inflight_transfer( + request) + if cancel_result == TransferCancelResult.BackendUnhealthy: + logger.warning( + f"Transceiver reports unhealthy while cancelling " + f"context request {request_id}; deferring resource " + f"release to health-driven restart.") + self._end_transfer_and_maybe_terminate(request) self._check_cache_transfer_errors("context requests") @@ -3861,20 +4002,27 @@ def _handle_responses(self): self._handle_errors( error_msg=f"Request {request.py_request_id} timed out", requests=[request]) - elif cancel_result == TransferCancelResult.CancelRequestedInFlight: - # Worker still touching buffers — keep the request - # active so we re-poll next iteration. Once C++ - # transitions the worker future to a final state we - # will see NotFound / AlreadyComplete and clean up. - self._inflight_cancel_requested_ids.add(req_id) - new_active_requests.append(request) - elif cancel_result == TransferCancelResult.BackendUnhealthy: - logger.warning( - f"Transceiver reports unhealthy while cancelling " - f"generation request {req_id}; deferring to " - f"health-driven restart.") - self._inflight_cancel_requested_ids.add(req_id) - new_active_requests.append(request) + elif cancel_result in ( + TransferCancelResult.CancelRequestedInFlight, + TransferCancelResult.BackendUnhealthy): + # Worker still touching buffers (or transceiver + # unhealthy). Surface the timeout error to the user + # immediately, but defer free_resources() to + # _maybe_release_pending_resources, which polls + # cancel_request_structured each iteration and frees + # only once the worker quiesces. We register the + # request BEFORE _handle_errors so _do_terminate_- + # request's guard sees req_id in the inflight set. + self._defer_resource_release_for_inflight_transfer( + request) + if cancel_result == TransferCancelResult.BackendUnhealthy: + logger.warning( + f"Transceiver reports unhealthy while cancelling " + f"generation request {req_id}; deferring resource " + f"release to health-driven restart.") + self._handle_errors( + error_msg=f"Request {request.py_request_id} timed out", + requests=[request]) continue if request.is_generation_only_request() and not request.is_finished: diff --git a/tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py b/tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py index 44a4f41d571d..3d81768d2f15 100644 --- a/tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py +++ b/tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py @@ -61,6 +61,15 @@ def is_disagg_generation_transmission_in_progress(self) -> bool: return self.state == "DISAGG_GENERATION_TRANS_IN_PROGRESS" +@dataclass +class _FakeHealth: + is_healthy: bool = True + quarantined_transfer_count: int = 0 + quarantine_budget: int = 16 + seconds_since_last_progress: float = 0.0 + global_progress_deadline_seconds: float = 60.0 + + @dataclass class _FakeTransceiver: """Programmable stand-in for the C++-backed KvCacheTransceiver. The @@ -69,6 +78,7 @@ class _FakeTransceiver: next_cancel_results: dict = field(default_factory=dict) cancel_calls: List[int] = field(default_factory=list) healthy: bool = True + health: _FakeHealth = field(default_factory=_FakeHealth) def cancel_request_structured(self, req: _FakeRequest): self.cancel_calls.append(req.py_request_id) @@ -78,6 +88,9 @@ def cancel_request_structured(self, req: _FakeRequest): def is_healthy(self) -> bool: return self.healthy + def get_health(self) -> _FakeHealth: + return self.health + def _make_executor_under_test(): """Build a minimal object whose interface matches what the helper @@ -88,6 +101,10 @@ def _make_executor_under_test(): obj.kv_cache_transceiver = _FakeTransceiver() obj._inflight_cancel_requested_ids = set() obj._pending_transfer_responses = [] + obj._pending_resource_release = [] + obj._transceiver_unhealthy_since = None + obj.is_shutdown = False + obj.shutdown_event = None obj.enable_attention_dp = False obj.active_requests = [] obj.terminated_request_ids = [] @@ -125,12 +142,58 @@ def _flush_pending_transfer_responses(self): if responses or self.enable_attention_dp: self._enqueue_responses(responses) + def _defer_resource_release_for_inflight_transfer(self, request): + self._inflight_cancel_requested_ids.add(request.py_request_id) + if request not in self._pending_resource_release: + self._pending_resource_release.append(request) + + def _maybe_release_pending_resources(self): + if not self._pending_resource_release or self.kv_cache_transceiver is None: + return + still_pending = [] + safe = ( + _FakeTransferCancelResult.AlreadyComplete, + _FakeTransferCancelResult.NotFound, + _FakeTransferCancelResult.CancelledBeforeAdvertise, + ) + for request in self._pending_resource_release: + cancel = self.kv_cache_transceiver.cancel_request_structured(request) + if cancel in safe: + self._inflight_cancel_requested_ids.discard(request.py_request_id) + self._do_terminate_request(request) + else: + still_pending.append(request) + self._pending_resource_release = still_pending + + def _check_transceiver_health(self): + if self.kv_cache_transceiver is None: + return + if self.kv_cache_transceiver.is_healthy(): + self._transceiver_unhealthy_since = None + return + import time as _time + now = _time.time() + if self._transceiver_unhealthy_since is None: + self._transceiver_unhealthy_since = now + return + elapsed = now - self._transceiver_unhealthy_since + health = self.kv_cache_transceiver.get_health() + grace = 2.0 * float(getattr(health, "global_progress_deadline_seconds", 60.0)) + if elapsed > grace and not self.is_shutdown: + self.is_shutdown = True + obj._can_terminate_request_now = types.MethodType( _can_terminate_request_now, obj) obj._do_terminate_request = types.MethodType(_do_terminate_request, obj) obj._enqueue_responses = types.MethodType(_enqueue_responses, obj) obj._flush_pending_transfer_responses = types.MethodType( _flush_pending_transfer_responses, obj) + obj._defer_resource_release_for_inflight_transfer = types.MethodType( + _defer_resource_release_for_inflight_transfer, obj) + obj._maybe_release_pending_resources = types.MethodType( + _maybe_release_pending_resources, obj) + obj._check_transceiver_health = types.MethodType( + _check_transceiver_health, obj) return obj @@ -346,3 +409,98 @@ def test_unhealthy_transceiver_keeps_request_pinned(): assert 33 in executor._inflight_cancel_requested_ids executor._do_terminate_request(request) assert executor.terminated_request_ids == [] + + +# --------------------------------------------------------------------------- +# Deferred-resource-release polling +# --------------------------------------------------------------------------- + + +def test_defer_resource_release_holds_until_quiesced(): + """Replay the full recovery sequence: timeout fires, error response + is surfaced, KV cleanup is deferred. Subsequent iterations poll + cancel_request_structured; when it transitions to AlreadyComplete / + NotFound, the deferred terminate runs.""" + executor = _make_executor_under_test() + request = _FakeRequest(py_request_id=77, state="DISAGG_TRANS_ERROR") + executor._defer_resource_release_for_inflight_transfer(request) + + # Iteration 1: worker still in flight — must NOT free. + executor.kv_cache_transceiver.next_cancel_results[77] = ( + _FakeTransferCancelResult.CancelRequestedInFlight) + executor._maybe_release_pending_resources() + assert executor.terminated_request_ids == [] + assert request in executor._pending_resource_release + assert 77 in executor._inflight_cancel_requested_ids + + # Iteration 2: still in flight. + executor._maybe_release_pending_resources() + assert executor.terminated_request_ids == [] + + # Iteration 3: worker finally quiesced. + executor.kv_cache_transceiver.next_cancel_results[77] = ( + _FakeTransferCancelResult.AlreadyComplete) + executor._maybe_release_pending_resources() + assert executor.terminated_request_ids == [77] + assert request not in executor._pending_resource_release + assert 77 not in executor._inflight_cancel_requested_ids + + +def test_defer_resource_release_handles_not_found_as_safe(): + """C++ may erase the worker entry between iterations; NotFound at + that point means the worker is gone and resources are safe to + free.""" + executor = _make_executor_under_test() + request = _FakeRequest(py_request_id=88, state="DISAGG_TRANS_ERROR") + executor._defer_resource_release_for_inflight_transfer(request) + + executor.kv_cache_transceiver.next_cancel_results[88] = ( + _FakeTransferCancelResult.NotFound) + executor._maybe_release_pending_resources() + assert executor.terminated_request_ids == [88] + + +def test_defer_resource_release_no_op_when_empty(): + executor = _make_executor_under_test() + executor._maybe_release_pending_resources() + assert executor.terminated_request_ids == [] + + +# --------------------------------------------------------------------------- +# Health-driven shutdown escape +# --------------------------------------------------------------------------- + + +def test_check_transceiver_health_resets_on_recovery(): + executor = _make_executor_under_test() + executor.kv_cache_transceiver.healthy = False + executor._check_transceiver_health() + assert executor._transceiver_unhealthy_since is not None + + # Recovery before grace period elapses — clear the timestamp, + # do not shut down. + executor.kv_cache_transceiver.healthy = True + executor._check_transceiver_health() + assert executor._transceiver_unhealthy_since is None + assert not executor.is_shutdown + + +def test_check_transceiver_health_triggers_shutdown_after_grace(): + executor = _make_executor_under_test() + # Set an extremely short deadline so the grace period (2x) is + # easy to exceed in a unit test. + executor.kv_cache_transceiver.health = _FakeHealth( + is_healthy=False, + global_progress_deadline_seconds=0.01, + ) + executor.kv_cache_transceiver.healthy = False + + # First call records the timestamp, does not shut down yet. + executor._check_transceiver_health() + assert not executor.is_shutdown + assert executor._transceiver_unhealthy_since is not None + + # Move the recorded timestamp into the past beyond the grace period. + executor._transceiver_unhealthy_since = executor._transceiver_unhealthy_since - 1.0 + executor._check_transceiver_health() + assert executor.is_shutdown From 05c8074de10b2ab5deb1e4e70fe62e9bdce11f82 Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Wed, 6 May 2026 14:07:50 -0700 Subject: [PATCH 6/7] [https://nvbugs/6104831][feat] [wedge-trace] heartbeats at the four disagg KV transfer wedge points MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Builds on PR #13796. Adds periodic WARNING-level [wedge-trace] log lines at the four C++ polling loops where a disaggregated KV transfer can stall, so an operator can identify the wedge location from log output alone — without attaching gdb to the worker. Sites instrumented: * NixlTransferStatus::wait Sender-side getXferStatus poll loop. Heartbeat fires when the same nixlXferReqH has been NIXL_IN_PROG for longer than the heartbeat interval. Logs the handle pointer + elapsed ms + the requested timeout_ms. * AgentConnectionManager::waitForNotification Receiver-side getNotifs poll loop. Heartbeat fires when no matching NotificationSyncInfo / ReadySignalInfo has arrived for longer than the heartbeat interval. Logs the remote agent name + expected tag + notification kind + elapsed ms. * BaseTransBufferManager::assignBufferIndex Slot-acquisition CV wait. Converts the previously unbounded resource.mBuffersCV.wait() into wait_for(interval) with a heartbeat log on each timeout, so the downstream effect of a stuck slot (next request blocked, slot held by a wedged worker thread) is visible. * CacheTransceiver::requestAndReceiveSync Synchronous gen-side path used when TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1. Replaces future.get() with a wait_for(interval) loop + heartbeat. Pairs with one of the worker-thread heartbeats to identify the actual wedge location. All four sites share the same env knob: TRTLLM_DISAGG_WEDGE_TRACE_INTERVAL_MS=5000 (default) TRTLLM_DISAGG_WEDGE_TRACE_INTERVAL_MS=0 (disable) Default 5 s is short enough to surface a wedge promptly without flooding logs on the happy path (transfers that complete within 5 s do not log anything). Cost on the happy path is one chrono::steady_clock::now() per polling iteration. Doc: docs/source/features/disagg-kv-transfer-wedge-trace-logs.md spells out each log line, sample output during a wedge, and the log-line → recovery-action mapping. Stacked on PR #13796 (fix/disagg-kv-transfer-session). Same base commit b9ce4b69d12fe5ba65d13893111b1a2ea29413ee. Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../batch_manager/baseTransBuffer.cpp | 38 ++++++++- .../batch_manager/cacheTransceiver.cpp | 26 ++++++ cpp/tensorrt_llm/common/envUtils.cpp | 7 ++ cpp/tensorrt_llm/common/envUtils.h | 7 ++ .../agent_utils/connection.cpp | 31 +++++++ .../nixl_utils/transferAgent.cpp | 27 +++++- .../disagg-kv-transfer-wedge-trace-logs.md | 82 +++++++++++++++++++ 7 files changed, 215 insertions(+), 3 deletions(-) create mode 100644 docs/source/features/disagg-kv-transfer-wedge-trace-logs.md diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp index 58092897ebbe..48e5d7149ca8 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp @@ -21,6 +21,7 @@ #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/opUtils.h" +#include #include namespace tensorrt_llm::batch_manager @@ -233,8 +234,41 @@ std::optional BaseTransBufferManager::assignBufferIndex( return std::nullopt; } std::unique_lock lk(resource.mBuffersMutex); - resource.mBuffersCV.wait( - lk, [&resource, bufferCount]() { return static_cast(resource.mConcurrence) < bufferCount; }); + + // [wedge-trace] convert the unbounded CV wait into wait_for(interval) + // so we can periodically log when a thread has been blocked waiting + // for a slot. Slot exhaustion at this layer is the downstream + // symptom of a stuck transfer (slot held by a wedged worker thread) + // — see docs/source/features/disagg-kv-transfer-debug-stuck-slot.md. + auto const heartbeatIntervalMs = common::getEnvDisaggWedgeTraceIntervalMs(); + auto const haveSlot + = [&resource, bufferCount]() { return static_cast(resource.mConcurrence) < bufferCount; }; + if (heartbeatIntervalMs > 0) + { + auto const startTime = std::chrono::steady_clock::now(); + auto lastHeartbeat = startTime; + while (!haveSlot()) + { + if (resource.mBuffersCV.wait_for(lk, std::chrono::milliseconds(heartbeatIntervalMs), haveSlot)) + { + break; + } + auto const now = std::chrono::steady_clock::now(); + auto const elapsedMs + = std::chrono::duration_cast(now - startTime).count(); + TLLM_LOG_WARNING( + "[wedge-trace] BaseTransBufferManager::assignBufferIndex blocked waiting for a slot " + "for %lld ms (concurrence=%d budget=%zu). The slot is most likely held by a wedged " + "worker thread; check sender/receiver heartbeats and run the stuck-slot diagnosis.", + static_cast(elapsedMs), static_cast(resource.mConcurrence), bufferCount); + lastHeartbeat = now; + (void) lastHeartbeat; + } + } + else + { + resource.mBuffersCV.wait(lk, haveSlot); + } int bufferId = -1; for (size_t i = 0; i < bufferCount; i++) { diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 70f7d2736b85..ab78e444484d 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -364,8 +364,34 @@ void CacheTransceiver::respondAndSendLayerWise( void CacheTransceiver::requestAndReceiveSync(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isGenerationOnlyRequest()); + auto const reqId = llmRequest->mRequestId; { auto future = mCacheReceiver->receiveAsync(llmRequest); + + // [wedge-trace] turn the previously unbounded future.get() into a + // wait_for(interval) loop so the executor thread can log a + // heartbeat when this synchronous gen-side path + // (TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1) is blocked. The + // worker thread's own heartbeats will tell us whether the wedge + // is in NIXL/UCX or somewhere downstream. See + // docs/source/features/disagg-kv-transfer-wedge-trace-logs.md. + auto const heartbeatIntervalMs = common::getEnvDisaggWedgeTraceIntervalMs(); + if (heartbeatIntervalMs > 0) + { + auto const startTime = std::chrono::steady_clock::now(); + while (future.wait_for(std::chrono::milliseconds(heartbeatIntervalMs)) != std::future_status::ready) + { + auto const elapsedMs = std::chrono::duration_cast( + std::chrono::steady_clock::now() - startTime) + .count(); + TLLM_LOG_WARNING( + "[wedge-trace] CacheTransceiver::requestAndReceiveSync executor thread blocked on " + "future.get for request %ld for %lld ms (TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP path). " + "Worker thread should be emitting its own [wedge-trace] heartbeat for the actual " + "wedge location.", + reqId, static_cast(elapsedMs)); + } + } future.get(); } llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); diff --git a/cpp/tensorrt_llm/common/envUtils.cpp b/cpp/tensorrt_llm/common/envUtils.cpp index bf3142a160dd..9a0117879ac7 100644 --- a/cpp/tensorrt_llm/common/envUtils.cpp +++ b/cpp/tensorrt_llm/common/envUtils.cpp @@ -454,6 +454,13 @@ std::string const& getEnvKVCacheTimeOutputPath() return outputPath; } +size_t getEnvDisaggWedgeTraceIntervalMs() +{ + // Default: 5 s. Set TRTLLM_DISAGG_WEDGE_TRACE_INTERVAL_MS=0 to disable. + static size_t const intervalMs = getUInt64Env("TRTLLM_DISAGG_WEDGE_TRACE_INTERVAL_MS").value_or(5000); + return intervalMs; +} + bool getEnvKVCacheTransferUseAsyncBuffer() { diff --git a/cpp/tensorrt_llm/common/envUtils.h b/cpp/tensorrt_llm/common/envUtils.h index 82c56f300362..546bd9e363f6 100644 --- a/cpp/tensorrt_llm/common/envUtils.h +++ b/cpp/tensorrt_llm/common/envUtils.h @@ -115,6 +115,13 @@ std::string const& getEnvKVCacheTimeOutputPath(); bool getEnvTryZCopyForKVCacheTransfer(); +/// Periodic [wedge-trace] heartbeat interval (ms) for the disaggregated +/// KV transfer wedge points (NixlTransferStatus::wait, +/// waitForNotification, assignBufferIndex CV wait, +/// requestAndReceiveSync future.get). When set to 0 the heartbeats are +/// disabled. Default: 5000 ms. +size_t getEnvDisaggWedgeTraceIntervalMs(); + // Force deterministic behavior for all kernels. bool getEnvForceDeterministic(); diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp index d46defdf50ad..e1403dac4f72 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp @@ -18,6 +18,7 @@ #include "connection.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h" +#include #include #include #include @@ -585,6 +586,12 @@ template void AgentConnectionManager::waitForNotification( std::string const& remoteAgentName, NotificationType& expectedInfo, std::atomic const& terminateFlag) { + auto const startTime = std::chrono::steady_clock::now(); + auto lastHeartbeat = startTime; + auto const heartbeatIntervalMs = common::getEnvDisaggWedgeTraceIntervalMs(); + char const* notificationKindName + = std::is_same_v ? "NotificationSyncInfo" : "ReadySignalInfo"; + while (!terminateFlag.load()) { @@ -593,6 +600,30 @@ void AgentConnectionManager::waitForNotification( return; } updateUnhandledNotifications(); + + // [wedge-trace] periodic heartbeat for the receiver-side + // getNotifs poll loop. Fires when no matching notification has + // arrived for more than the configured interval. See + // docs/source/features/disagg-kv-transfer-wedge-trace-logs.md. + if (heartbeatIntervalMs > 0) + { + auto const now = std::chrono::steady_clock::now(); + auto const sinceHeartbeatMs + = std::chrono::duration_cast(now - lastHeartbeat).count(); + if (sinceHeartbeatMs >= static_cast(heartbeatIntervalMs)) + { + auto const elapsedMs + = std::chrono::duration_cast(now - startTime).count(); + TLLM_LOG_WARNING( + "[wedge-trace] AgentConnectionManager::waitForNotification still pending after %lld ms " + "(remoteAgent='%s' expectedTag=%d kind=%s). Receiver-side getNotifs poll is not making " + "progress; sender may have crashed or its notification was lost.", + static_cast(elapsedMs), remoteAgentName.c_str(), + static_cast(expectedInfo.mContext.getTag()), notificationKindName); + lastHeartbeat = now; + } + } + std::scoped_lock lock(mNotificationMutex); auto it = mUnhandledNotifications.begin(); while (it != mUnhandledNotifications.end()) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp index bad3e184f983..99fdf55d38d9 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp @@ -484,7 +484,9 @@ NixlTransferStatus::NixlTransferStatus(nixlAgent* agent, nixlXferReqH* handle) TransferState NixlTransferStatus::wait(int64_t timeout_ms) const { - auto startTime = std::chrono::steady_clock::now(); + auto const startTime = std::chrono::steady_clock::now(); + auto lastHeartbeat = startTime; + auto const heartbeatIntervalMs = common::getEnvDisaggWedgeTraceIntervalMs(); while (true) { @@ -498,6 +500,29 @@ TransferState NixlTransferStatus::wait(int64_t timeout_ms) const return TransferState::kFAILURE; } + // [wedge-trace] periodic heartbeat for the sender-side + // getXferStatus poll loop. Fires when the same nixlXferReqH has + // been NIXL_IN_PROG for more than the configured interval. See + // docs/source/features/disagg-kv-transfer-wedge-trace-logs.md. + if (heartbeatIntervalMs > 0) + { + auto const now = std::chrono::steady_clock::now(); + auto const sinceHeartbeatMs + = std::chrono::duration_cast(now - lastHeartbeat).count(); + if (sinceHeartbeatMs >= static_cast(heartbeatIntervalMs)) + { + auto const elapsedMs + = std::chrono::duration_cast(now - startTime).count(); + TLLM_LOG_WARNING( + "[wedge-trace] NixlTransferStatus::wait still NIXL_IN_PROG after %lld ms " + "(handle=%p, requested timeout_ms=%lld). Sender-side getXferStatus poll " + "is not making progress; peer or UCX endpoint may be wedged.", + static_cast(elapsedMs), static_cast(mHandle), + static_cast(timeout_ms)); + lastHeartbeat = now; + } + } + // If timeout_ms < 0, wait indefinitely until status is not NIXL_IN_PROG if (timeout_ms < 0) { diff --git a/docs/source/features/disagg-kv-transfer-wedge-trace-logs.md b/docs/source/features/disagg-kv-transfer-wedge-trace-logs.md new file mode 100644 index 000000000000..31536c3de98a --- /dev/null +++ b/docs/source/features/disagg-kv-transfer-wedge-trace-logs.md @@ -0,0 +1,82 @@ +# `[wedge-trace]` log lines — quick reference + +Companion doc to the verify-wedge runbook (in TRT-LLM PR #13796). When +the disagg KV transfer wedges, the four polling loops where the +process can stall now emit a periodic `[wedge-trace]` log at +`WARNING` level, so an operator can identify the wedge location from +log output alone — no `gdb` attach required. + +All four log sites share the same heartbeat interval, controlled by +**`TRTLLM_DISAGG_WEDGE_TRACE_INTERVAL_MS`** (default `5000`). Set to +`0` to disable. + +## The four heartbeats + +| Log prefix | Source | Wedge it surfaces | +|---|---|---| +| `[wedge-trace] NixlTransferStatus::wait still NIXL_IN_PROG …` | `cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp` (inside `NixlTransferStatus::wait`) | **Sender** worker stuck polling `getXferStatus` — the in-flight `ucp_request` is not completing. UCX endpoint or peer wedged. | +| `[wedge-trace] AgentConnectionManager::waitForNotification still pending …` | `cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp` (inside `waitForNotification`) | **Receiver** worker stuck polling `getNotifs` — the post-write notification from the sender never arrived. Sender crashed or notification dropped. | +| `[wedge-trace] BaseTransBufferManager::assignBufferIndex blocked …` | `cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp` (inside `assignBufferIndex`) | A new request can't acquire a slot — slot is held by a wedged worker. **Downstream effect** of one of the above wedges. | +| `[wedge-trace] CacheTransceiver::requestAndReceiveSync executor thread blocked …` | `cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp` (inside `requestAndReceiveSync`) | Executor thread blocked on `future.get`. **Only on the** `TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1` **path.** Pairs with one of the worker-thread heartbeats above. | + +## Sample log output during a wedge + +``` +[2026-05-06 10:42:18.131] WARN [wedge-trace] NixlTransferStatus::wait still NIXL_IN_PROG after 5023 ms (handle=0x7f1c4c01a3b0, requested timeout_ms=-1). Sender-side getXferStatus poll is not making progress; peer or UCX endpoint may be wedged. +[2026-05-06 10:42:18.155] WARN [wedge-trace] BaseTransBufferManager::assignBufferIndex blocked waiting for a slot for 5009 ms (concurrence=1 budget=1). The slot is most likely held by a wedged worker thread; check sender/receiver heartbeats and run the stuck-slot diagnosis. +[2026-05-06 10:42:23.131] WARN [wedge-trace] NixlTransferStatus::wait still NIXL_IN_PROG after 10025 ms (handle=0x7f1c4c01a3b0, requested timeout_ms=-1). … +[2026-05-06 10:42:23.155] WARN [wedge-trace] BaseTransBufferManager::assignBufferIndex blocked waiting for a slot for 10010 ms (concurrence=1 budget=1). … +``` + +Reading this: + +* The handle pointer `0x7f1c4c01a3b0` identifies which NIXL transfer + is wedged. Cross-reference with earlier `start recv bufferIdx` + debug logs (see + [`disagg-kv-transfer-debug-stuck-slot.md`](disagg-kv-transfer-debug-stuck-slot.md)) + to map the handle to a request id. +* `concurrence=1 budget=1` means the recv-buffer pool has the default + one slot and it is fully consumed; this is the classic "single + wedged transfer pins the rank" pattern documented in + [`disagg-kv-transfer-session-lifecycle.md`](disagg-kv-transfer-session-lifecycle.md). + +## Mapping log line → recovery action + +| Log line you see | Most likely cause | Action | +|---|---|---| +| Only `NixlTransferStatus::wait` heartbeats fire (sender side) | Sender's `ucp_request` not completing — peer down or UCX endpoint wedge | Health-driven worker restart via PR #13796's `_check_transceiver_health` (default ~3 min total). Ask NIXL Q1.2 — should `getXferStatus` surface UCX endpoint-error as a final state? | +| Only `waitForNotification` heartbeats fire (receiver side) | Sender notification dropped or sender crashed before sending it | Same health-driven restart path. Ask NIXL Q4.1 — are notifications reliable? | +| Both sides heartbeat | Both peers wedged; usually one started it (often a downstream pod restart) | Restart the worker that's hardest to recover. | +| Only `assignBufferIndex` heartbeats fire (no NIXL-level heartbeat) | The thread holding the slot is wedged in something OTHER than NIXL — most likely CUDA, MPI, or a TRT-LLM mutex | `gdb -p` the process to find the slot owner; this is NOT a NIXL wedge. | +| Only `requestAndReceiveSync` heartbeat fires (no worker heartbeat) | Possible. Means the receive worker died without setting the promise. Or `TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1` is on AND the worker is healthy but slow | Check for crashed worker thread via `cat /proc/$PID/task/*/comm`. | + +## Disabling the heartbeats + +```bash +TRTLLM_DISAGG_WEDGE_TRACE_INTERVAL_MS=0 +``` + +…or for finer control: + +```bash +# Only fire after 30 s of no progress (suppresses occasional slow transfers) +TRTLLM_DISAGG_WEDGE_TRACE_INTERVAL_MS=30000 + +# Tight loop, useful in test environments +TRTLLM_DISAGG_WEDGE_TRACE_INTERVAL_MS=1000 +``` + +The heartbeats themselves are cheap (a single `chrono::steady_clock::now()` +plus a `TLLM_LOG_WARNING` per interval per wedged thread). They do +not affect throughput on the happy path. + +## Cross-references + +* [`disagg-kv-transfer-debug-verify-wedge.md`](disagg-kv-transfer-debug-verify-wedge.md) + — the gdb / py-spy / NIXL-telemetry / UCX-stats verification recipes. +* [`disagg-kv-transfer-debug-stuck-slot.md`](disagg-kv-transfer-debug-stuck-slot.md) + — the start-recv-without-finish-recv slot-leak diagnosis. +* [`disagg-kv-transfer-session-lifecycle.md`](disagg-kv-transfer-session-lifecycle.md) + — recovery model and global progress deadline. +* [`disagg-nixl-open-questions.md`](disagg-nixl-open-questions.md) — + upstream NIXL questions whose answers would shorten the recovery path. From bf5e7188d00b03b9b962334145efb01182ab0888 Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Wed, 6 May 2026 19:34:42 -0700 Subject: [PATCH 7/7] [https://nvbugs/6104831][fix] Anchor KV transfer deadline to actual-transfer-start, not admit-time MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stacked on PR #13818. PR #13796 captured the per-request KV transfer deadline at the moment the executor admitted the request to its tracking (in respondAndSendAsync / requestAndReceiveAsync). The actual transfer doesn't begin until the worker thread dequeues the request, calls assignBufferIndexFor{Send,Recv}, and starts session.{recv,send}. Under burst load with the default small staged-buffer pools (recvBufferCount=1, TRTLLM_KVCACHE_SEND_MAX_CONCURRENCY_NUM=1) and moderate transport throughput, the deadline burns down while the request is queued, not while it's transferring. Empirically — see the d224/d227 testing in dynamo-gb200-test/docs/disagg-wedge-pr13818-findings-2026-05-06.md — a 60s deadline on TCP at ~3 Gbps with a conc=128 burst quarantined ~80% of requests for queue starvation rather than for any actual NIXL/UCX wedge. This commit anchors the deadline to actual-transfer-start instead. Changes: * New LlmRequest field mKvCacheActualTransferStart with setKvCacheActualTransferStart / getKvCacheActualTransferStart / hasKvCacheActualTransferStart accessors. Setter is idempotent so hybrid (KV + indexer-K, KV + RNN-state) models capture the first slot acquisition. * The formatters set the field once a transfer-buffer slot has been acquired: - cacheFormatter.cpp:497 (sender, after assignBufferIndexForSend) - cacheFormatter.cpp:822 (receiver, after assignBufferIndexForRecv or pre-assigned id) - mlaCacheFormatter.cpp:256, 496 (MLA indexer-K send/recv) - rnnCacheFormatter.cpp:151, 324 (RNN-state send/recv) * CacheTransceiver::maybeQuarantineLocked consults the new field. Returns early when hasKvCacheActualTransferStart() is false — requests still in the worker queue or waiting for a slot are not eligible for quarantine. When the field is set, deadline = transferStart + kvTransferTimeoutMs. * TrackedFuture::deadline (admit-time captured) and CacheTransceiver::computeTrackedFutureDeadline are removed. TrackedFuture initialization sites in respondAndSendAsync, respondAndSendLayerWise, and requestAndReceiveAsync drop the deadline argument. The existing perf metric kv_cache_transfer_time_ms is unchanged — mPerfMetrics.timingMetrics.kvCacheTransferStart still records the admit / worker-dequeue time and feeds the existing perf metric. The new field is a separate "when did the actual transfer start" clock used only by the quarantine logic. Doc: docs/source/features/disagg-kv-transfer-deadline-anchor.md The 19 existing Python unit tests in tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py continue to pass without modification — they exercise the Python policy layer which this commit does not touch. Same base commit as PR #13796 / #13818 (b9ce4b69d12fe5ba65d13893111b1a2ea29413ee). Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 13 +- .../tensorrt_llm/batch_manager/llmRequest.h | 51 +++++ .../batch_manager/cacheFormatter.cpp | 12 ++ .../batch_manager/cacheTransceiver.cpp | 63 +++--- .../batch_manager/mlaCacheFormatter.cpp | 8 + .../batch_manager/rnnCacheFormatter.cpp | 9 + .../disagg-kv-transfer-deadline-anchor.md | 204 ++++++++++++++++++ 7 files changed, 326 insertions(+), 34 deletions(-) create mode 100644 docs/source/features/disagg-kv-transfer-deadline-anchor.md diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index fbdc7719a0f3..2b8c797565ac 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -406,11 +406,18 @@ class CacheTransceiver : public BaseCacheTransceiver /// thread (and any asynchronous queue inside CacheSender / /// CacheReceiver) can dereference the request safely as long as the /// transceiver has not erased this entry. + /// + /// The per-request KV transfer timeout is anchored to + /// @c LlmRequest::getKvCacheActualTransferStart(), recorded by the + /// formatters once a transfer-buffer slot has been acquired. There + /// is intentionally no captured deadline on this struct: a + /// deeply-queued request without an actual-transfer-start time + /// cannot be quarantined, which prevents queue-starvation false + /// positives. struct TrackedFuture { std::shared_ptr request; std::future future; - std::chrono::steady_clock::time_point deadline; /// True once the per-request timeout has fired and we have flipped /// the request to an error state. The entry stays in the vector /// (and the future stays pinned) until the worker actually @@ -445,10 +452,6 @@ class CacheTransceiver : public BaseCacheTransceiver std::optional const& atLeastRequestNum, bool markComplete, bool allowBlocking); void checkGenTransferStatusImpl(std::optional const& atLeastRequestNum, bool allowBlocking); - /// Compute the per-entry deadline for a tracked transfer. - [[nodiscard]] std::chrono::steady_clock::time_point computeTrackedFutureDeadline( - std::chrono::steady_clock::time_point requestStart) const; - std::unique_ptr mCacheSender; std::unique_ptr mCacheReceiver; std::vector mSenderFutures; diff --git a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h index c1ecffff026b..3be16c9e608f 100644 --- a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h +++ b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h @@ -1797,6 +1797,46 @@ class GenericLlmRequest .count()); } + /// @brief Mark when the actual KV transfer (post slot acquisition, + /// post pre-data handshake) began for this request. + /// + /// This is distinct from @ref setKvCacheTransferStart: that one is + /// called at admit / worker-dequeue time and feeds the + /// `kv_cache_transfer_time_ms` perf metric (which has historically + /// included queue + slot wait). The "actual transfer start" is + /// recorded by the formatters after `assignBufferIndexFor*` returns, + /// when the worker is about to begin the data movement, and is the + /// correct anchor for the `kvTransferTimeoutMs` deadline used by + /// CacheTransceiver's quarantine logic. + /// + /// Idempotent: subsequent calls within the same transfer (e.g., for + /// MLA / hybrid models that acquire multiple pool slots) are no-ops + /// once the field is set. + void setKvCacheActualTransferStart(TimePoint time) const + { + if (mKvCacheActualTransferStart == TimePoint{}) + { + mKvCacheActualTransferStart = maybeToGlobalSteadyClock(time); + } + } + + [[nodiscard]] TimePoint getKvCacheActualTransferStart() const noexcept + { + return mKvCacheActualTransferStart; + } + + [[nodiscard]] bool hasKvCacheActualTransferStart() const noexcept + { + return mKvCacheActualTransferStart != TimePoint{}; + } + + /// Reset before re-issuing a transfer (e.g., on retry). Not used in + /// PR #13796's flow but kept symmetric with the perf-metric API. + void clearKvCacheActualTransferStart() const noexcept + { + mKvCacheActualTransferStart = TimePoint{}; + } + void updateKvCacheSize(size_t targetBufferSize) const { mPerfMetrics.timingMetrics.kvCacheSize += targetBufferSize; @@ -2118,6 +2158,17 @@ class GenericLlmRequest bool mReturnPerfMetrics{false}; mutable executor::RequestPerfMetrics mPerfMetrics; + /// Time at which the disagg KV cache transfer actually started for + /// this request — recorded by the formatters once a transfer-buffer + /// slot has been acquired (i.e., once the worker is about to begin + /// data movement). Distinct from + /// @c mPerfMetrics.timingMetrics.kvCacheTransferStart, which is set + /// at admit / worker-dequeue time and thus includes queue / slot + /// wait. Used by @c CacheTransceiver::maybeQuarantineLocked as the + /// anchor for @c kvTransferTimeoutMs so that a deeply-queued + /// request is not quarantined for queue starvation. + mutable TimePoint mKvCacheActualTransferStart{}; + // Guided decoding params. std::optional mGuidedDecodingParams{std::nullopt}; diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index e05e8d6f76fc..6af1242c8743 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -495,6 +495,12 @@ 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(); + // The slot is now reserved and the worker is about to begin + // actual data movement. Anchor the per-request KV transfer + // timeout from this point so queue / slot wait does not + // consume the deadline budget. See + // CacheTransceiver::maybeQuarantineLocked. + session.getLlmRequest().setKvCacheActualTransferStart(LlmRequest::getSteadyClockNow()); int peerDuplicateHeadFactor = targetInfo.mPeerDupHeadFactor; auto bufferTargetNum = targetNum / peerDuplicateHeadFactor; auto ppRank = selfIdx @@ -815,6 +821,12 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess { cacheBufferId = mCacheTransBufferManager->assignBufferIndexForRecv(); } + // Slot reserved (dynamically or pre-assigned). Anchor + // the per-request KV transfer timeout from this point — + // before this, the worker was still queued or waiting + // for a slot, neither of which is "in transfer." See + // CacheTransceiver::maybeQuarantineLocked. + llmRequest.setKvCacheActualTransferStart(LlmRequest::getSteadyClockNow()); auto [recvSplitCachestmp, bufferCoverTargetNumtmp, onlyUseDynamicBuffer] = mCacheTransBufferManager->getOrAllocateRecvBuffers( cacheBufferId, static_cast(targetNum), bufferEleSizes, bufferManager); diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index ab78e444484d..999757df6ccb 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -339,8 +339,11 @@ void CacheTransceiver::respondAndSendAsync(std::shared_ptr llmReques } setContextState(llmRequest.get()); auto future = mCacheSender->sendAsync(llmRequest); - auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); - mSenderFutures.push_back(TrackedFuture{std::move(llmRequest), std::move(future), deadline, false, false}); + // Note: no admit-time deadline is captured. The per-request KV + // transfer timeout (kvTransferTimeoutMs) is anchored to + // LlmRequest::getKvCacheActualTransferStart(), set by the + // formatters once the worker has acquired a transfer-buffer slot. + mSenderFutures.push_back(TrackedFuture{std::move(llmRequest), std::move(future), false, false}); } void CacheTransceiver::respondAndSendLayerWise( @@ -356,8 +359,7 @@ void CacheTransceiver::respondAndSendLayerWise( llmRequest->setState(LlmRequestState::kDISAGG_CONTEXT_INIT_AND_TRANS); setContextState(llmRequest.get()); auto future = mCacheSender->sendAsync(llmRequest); - auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); - mSenderFutures.push_back(TrackedFuture{llmRequest, std::move(future), deadline, false, false}); + mSenderFutures.push_back(TrackedFuture{llmRequest, std::move(future), false, false}); } } @@ -411,9 +413,9 @@ void CacheTransceiver::requestAndReceiveAsync(std::shared_ptr llmReq } auto future = mCacheReceiver->receiveAsync(llmRequest); - auto const deadline = computeTrackedFutureDeadline(/*requestStart=*/std::chrono::steady_clock::now()); llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS); - mRequesterFutures.push_back(TrackedFuture{std::move(llmRequest), std::move(future), deadline, false, false}); + // Note: no admit-time deadline; see respondAndSendAsync. + mRequesterFutures.push_back(TrackedFuture{std::move(llmRequest), std::move(future), false, false}); } std::vector gatherRequestIds( @@ -512,24 +514,6 @@ void updateKVCacheTransferBW(std::shared_ptr const& mComm, // ----- Internal helpers ---------------------------------------------------- -std::chrono::steady_clock::time_point CacheTransceiver::computeTrackedFutureDeadline( - std::chrono::steady_clock::time_point requestStart) const -{ - // Returns time_point::max() when no per-request timeout is configured, - // so the polling path treats the entry as "never times out" — matching - // the legacy behaviour for callers that did not opt into timeouts. - if (!mCacheTransceiverConfig.has_value()) - { - return std::chrono::steady_clock::time_point::max(); - } - auto const timeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); - if (!timeoutMs.has_value()) - { - return std::chrono::steady_clock::time_point::max(); - } - return requestStart + std::chrono::milliseconds(timeoutMs.value()); -} - void CacheTransceiver::updateHealthLocked() { auto const now = std::chrono::steady_clock::now(); @@ -569,12 +553,33 @@ bool CacheTransceiver::maybeQuarantineLocked(TrackedFuture& entry, RequestStatus { return false; } - if (entry.deadline == std::chrono::steady_clock::time_point::max()) + if (!mCacheTransceiverConfig.has_value()) + { + return false; + } + auto const timeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); + if (!timeoutMs.has_value()) + { + return false; + } + // Defer the deadline check until the worker has actually started + // the transfer (post slot acquisition). Requests still sitting in + // the worker queue or waiting on a slot are not "stuck" in any + // transport-relevant sense; quarantining them for queue starvation + // turned every burst into a false-positive cascade — see + // claude-swe/disagg-deadline-at-transfer-start-fix.md and the + // d224/d227 findings (TCP at ~3 Gbps, 1-slot serialization, conc + // 128) where this mis-accounting caused ~80% of requests to be + // quarantined for queue wait rather than for any actual NIXL/UCX + // wedge. + if (!entry.request->hasKvCacheActualTransferStart()) { return false; } + auto const transferStart = entry.request->getKvCacheActualTransferStart(); + auto const deadline = transferStart + std::chrono::milliseconds(timeoutMs.value()); auto const now = std::chrono::steady_clock::now(); - if (now < entry.deadline) + if (now < deadline) { return false; } @@ -590,9 +595,9 @@ bool CacheTransceiver::maybeQuarantineLocked(TrackedFuture& entry, RequestStatus updateHealthLocked(); } TLLM_LOG_WARNING( - "Quarantining request %ld: per-request KV transfer deadline exceeded; future stays pinned in C++ until " - "worker reaches a final state.", - entry.request->mRequestId); + "Quarantining request %ld: per-request KV transfer deadline exceeded after slot acquisition " + "(actual transfer time > %lld ms). Future stays pinned in C++ until worker reaches a final state.", + entry.request->mRequestId, static_cast(timeoutMs.value())); return true; } diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp index c72090867f29..429def3f9dcf 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp @@ -254,6 +254,11 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses }; auto bufferEleSizes = getBufferSizeForTarget(); auto cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForSend(); + // MLA indexer-K-cache slot reserved. Idempotent — if the + // primary KV slot was already acquired earlier in + // CacheFormatter::format, this call is a no-op. See + // CacheTransceiver::maybeQuarantineLocked. + llmRequest.setKvCacheActualTransferStart(LlmRequest::getSteadyClockNow()); auto result = mCacheTransBufferManagers[transferIndexerKCache]->getOrAllocateSendBuffers( cacheBufferId, static_cast(pPDomainSize * cPDomainSize), bufferEleSizes, bufferManager); auto& outputSplitCaches = std::get<0>(result); @@ -495,6 +500,9 @@ void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& s { cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForRecv(); } + // MLA indexer-K-cache slot reserved on the receive side. + // Idempotent (see CacheFormatter::unformat). + llmRequest.setKvCacheActualTransferStart(LlmRequest::getSteadyClockNow()); auto targetNum = pickUpConnections.size(); diff --git a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp index 1fd1cbdc253f..54a80832ae08 100644 --- a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp @@ -149,6 +149,12 @@ void RnnCacheFormatter::format(TransferSession& session) } auto cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForSend(); + // RNN-state slot reserved on the send side. Anchor the per-request + // KV transfer timeout from this point. Idempotent — for hybrid + // models that also acquire a KV slot in CacheFormatter::format, + // whichever fires first wins. See + // CacheTransceiver::maybeQuarantineLocked. + llmRequest.setKvCacheActualTransferStart(LlmRequest::getSteadyClockNow()); auto allocationResult = mRnnCacheTransBufferManager->getOrAllocateSendBuffers( cacheBufferId, static_cast(bufferTargetNum), bufferSizesPerTarget, bufferManager); auto& outputBuffers = std::get<0>(allocationResult); @@ -323,6 +329,9 @@ void RnnCacheFormatter::unformat(TransferSession& session) { cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForRecv(); } + // RNN-state slot reserved on the recv side. Idempotent (see send + // side / CacheFormatter::unformat). + llmRequest.setKvCacheActualTransferStart(LlmRequest::getSteadyClockNow()); auto allocationResult = mRnnCacheTransBufferManager->getOrAllocateRecvBuffers( cacheBufferId, static_cast(sourceNum), bufferSizesPerSource, bufferManager); diff --git a/docs/source/features/disagg-kv-transfer-deadline-anchor.md b/docs/source/features/disagg-kv-transfer-deadline-anchor.md new file mode 100644 index 000000000000..96cb94a58bc0 --- /dev/null +++ b/docs/source/features/disagg-kv-transfer-deadline-anchor.md @@ -0,0 +1,204 @@ +# Anchoring `kv_transfer_timeout_ms` to actual transfer start + +> Companion to +> [`disagg-kv-transfer-session-lifecycle.md`](disagg-kv-transfer-session-lifecycle.md). +> Documents the bookkeeping change introduced in PR #13830 (this PR), +> stacked on PRs #13796 and #13818. + +## Background — the bug we're fixing + +PR #13796 captured the per-request deadline at admit time: + +```cpp +// cacheTransceiver.cpp::requestAndReceiveAsync (pre-fix) +auto future = mCacheReceiver->receiveAsync(llmRequest); +auto const deadline = computeTrackedFutureDeadline( + /*requestStart=*/std::chrono::steady_clock::now()); +mRequesterFutures.push_back(TrackedFuture{..., deadline, ...}); +``` + +The actual transfer doesn't start until much later inside the worker +thread — it has to wait for the worker to dequeue the request, then +for `assignBufferIndexFor{Send,Recv}` to acquire a slot, then for +`session.{recv,send}` to exchange data with the peer. + +Under burst load with a small staged-buffer pool (default +`recvBufferCount=1` and `TRTLLM_KVCACHE_SEND_MAX_CONCURRENCY_NUM=1`), +the deadline burns down while the request is queued, not while it's +transferring. Empirically — see the d224/d227 testing under +`dynamo-gb200-test/docs/disagg-wedge-pr13818-findings-2026-05-06.md` +— a 60 s deadline on TCP at ~3 Gbps with a conc=128 burst caused +~80 % of requests to be quarantined for **queue starvation** rather +than for any actual NIXL/UCX wedge. They never even started +transferring; they just sat waiting for their turn. + +## What this PR changes + +`kv_transfer_timeout_ms` is now anchored to **actual transfer +start** — the moment the worker has acquired a slot and is about to +begin data movement. + +### New `LlmRequest` field + +```cpp +mutable TimePoint mKvCacheActualTransferStart{}; + +void setKvCacheActualTransferStart(TimePoint time) const; // idempotent +TimePoint getKvCacheActualTransferStart() const noexcept; +bool hasKvCacheActualTransferStart() const noexcept; +``` + +This is **distinct from** the existing +`mPerfMetrics.timingMetrics.kvCacheTransferStart` (set at +admit / worker-dequeue time and exposed as +`req.kv_cache_transfer_time_ms`). The perf metric semantics are +preserved — downstream consumers / dashboards continue to see the +same timing model they did before. + +### Formatters set the field after slot acquisition + +| Path | File:line | Slot | +|---|---|---| +| Sender, primary KV | [`cacheFormatter.cpp:497`](https://github.com/NVIDIA/TensorRT-LLM/blob/fix/disagg-kv-transfer-deadline-fix/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp#L497) | `assignBufferIndexForSend` | +| Receiver, primary KV | [`cacheFormatter.cpp:822`](https://github.com/NVIDIA/TensorRT-LLM/blob/fix/disagg-kv-transfer-deadline-fix/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp#L822) | `assignBufferIndexForRecv` (or pre-assigned id) | +| MLA indexer-K, sender | [`mlaCacheFormatter.cpp:256`](https://github.com/NVIDIA/TensorRT-LLM/blob/fix/disagg-kv-transfer-deadline-fix/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp#L256) | indexer-K send slot | +| MLA indexer-K, receiver | [`mlaCacheFormatter.cpp:496`](https://github.com/NVIDIA/TensorRT-LLM/blob/fix/disagg-kv-transfer-deadline-fix/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp#L496) | indexer-K recv slot | +| RNN/Mamba, sender | [`rnnCacheFormatter.cpp:151`](https://github.com/NVIDIA/TensorRT-LLM/blob/fix/disagg-kv-transfer-deadline-fix/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp#L151) | RNN-state send slot | +| RNN/Mamba, receiver | [`rnnCacheFormatter.cpp:324`](https://github.com/NVIDIA/TensorRT-LLM/blob/fix/disagg-kv-transfer-deadline-fix/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp#L324) | RNN-state recv slot | + +The setter is idempotent — for hybrid (KV + indexer-K, KV + RNN-state) +models that acquire multiple pool slots in sequence, the field +captures the **first** slot acquisition only. That's the point at +which the worker can be considered "in transfer." + +### `maybeQuarantineLocked` consults the new field + +```cpp +// cacheTransceiver.cpp::maybeQuarantineLocked (post-fix) +if (!entry.request->hasKvCacheActualTransferStart()) +{ + return false; // worker still queued / waiting for slot — no deadline +} +auto const transferStart = entry.request->getKvCacheActualTransferStart(); +auto const deadline = transferStart + std::chrono::milliseconds(timeoutMs.value()); +if (now < deadline) return false; +// quarantine +``` + +Requests that haven't yet reached the formatter's slot-acquire path +are not eligible for quarantine — they're not "stuck" in any +transport-relevant sense. + +### Dropped: admit-time deadline plumbing + +`TrackedFuture::deadline` is removed. `CacheTransceiver::computeTrackedFutureDeadline` +is removed. The three call sites in `respondAndSendAsync`, +`respondAndSendLayerWise`, and `requestAndReceiveAsync` no longer +capture admit-time deadlines. + +## Behavioural impact + +### Intended + +| Scenario | Before | After | +|---|---|---| +| TCP, conc=128, 60 s deadline, 1-slot serialization | ~80 % of burst quarantined for queue wait | only transfers whose actual transfer time exceeds 60 s are quarantined (rare on functional TCP) | +| Healthy transfer in queue for 30 s, transfers in 5 s | quarantined at T+60 s mid-transfer | not quarantined; user sees normal completion | +| Genuinely stuck mid-transfer (peer crash after slot acquired, NIXL hung) | quarantined at T + 60 s after admit | quarantined at T + 60 s after slot acquisition | +| Request never gets a slot (queue depth ≫ capacity, pool wedged elsewhere) | quarantined at admit + 60 s | not quarantined for queue wait. Orchestrator's `req_timeout_secs=180` is the user-visible backstop. | + +### Side effects + +- **Perf metric `kv_cache_transfer_time_ms`** is unchanged. Still + measured from admit to `kvCacheTransferEnd` via the existing + `mPerfMetrics.timingMetrics.kvCacheTransferStart` — this PR adds + a separate field, doesn't repurpose the existing one. +- **`_check_transceiver_health` (Python side)** is unchanged. It + flips `is_shutdown` on `is_healthy=false` after the grace period. + The trigger for `is_healthy=false` (no progress for > + `mGlobalProgressDeadlineMs` OR quarantine count > budget) is + unchanged at the C++ level. With this fix in place, false-positive + quarantines from queue starvation no longer drive the budget side + of that condition. +- **Wedge-trace heartbeats from PR #13818** are unchanged — they + fire inside the worker thread regardless of where the deadline + anchor is. +- **Structured cancel result enum / deferred resource release / + `_inflight_cancel_requested_ids`** are unchanged. + +### Risks and mitigations + +1. **Unbounded queue wait.** A request in the queue without a slot + never gets its deadline to start ticking. Mitigations: + - The disagg orchestrator's HTTP-level `req_timeout_secs=180` + fires a client-visible 504/408 at 180 s. + - The transceiver's no-progress detector (`mLastProgressTime` + vs `mGlobalProgressDeadlineMs`) flips `is_healthy=false` if + the queue head stops draining for > 60 s with at least one + tracked entry, triggering health-driven shutdown via the + PR #13796 mechanism. + - For deployments that want a hard ceiling on queue wait, a + follow-up could add a separate `kv_queue_wait_timeout_ms` + knob that fires at admit + N seconds if `hasKvCacheActualTransferStart()` + is still false. Not needed for the d224/d227 scenario. + +2. **Pre-existing behavior change for users who set + `kv_transfer_timeout_ms` thinking "end-to-end including queue."** + The new semantics are "actual transfer time only." For most + deployments this is the more useful interpretation. Users who + wanted the old admit-to-completion bound should consider using + the orchestrator's `req_timeout_secs` as a wrapper budget. + +3. **MLA / RNN / hybrid-pool ordering.** If a hybrid model acquires + the indexer-K slot before the primary KV slot, the deadline + anchors at the earlier acquisition. That's correct — it's the + first moment the worker is doing real transfer work. + +## Empirical validation plan + +Re-run the d224 / d227 burst suite: + +| Repro | Pre-fix | Post-fix expectation | +|---|---|---| +| d224 r1 (TCP 3 Gbps, conc=128, no canary) | ~80 % quarantined; `is_shutdown` at T+4.4 min | quarantines drop to near zero (only true mid-transfer stalls); `is_shutdown` does not fire; bursts drain over time | +| d227 r4 (Pattern C, canary on, peer crash) | both sides quarantine + canary-driven restart | mid-transfer wedges still get quarantined → still trigger restart path; verdict unchanged | +| Single-request smoke | unchanged | unchanged | + +If d224 r1 still reaches `is_shutdown` after this fix, look at the +no-progress side of `updateHealthLocked` — that's the next layer +of the diagnosis (covered separately in +`disagg-deadline-vs-quarantine-proposal.md`). + +## Test plan + +The existing 19 Python unit tests in +`tests/unittest/disaggregated/test_kv_transfer_session_lifecycle.py` +continue to pass without modification — they exercise the Python +policy layer, which is unchanged by this PR. + +C++-side coverage is via the existing multi-GPU integration tests in +`cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp`. No +adjustments needed; that suite uses the public API and the change +is internal. + +For empirical validation of the deadline-anchor behaviour itself, +the d224/d227 reruns above are the canonical signal. + +## Sequencing relative to the broader recovery work + +This PR is **independent of** the deadline-vs-quarantine proposal +in `claude-swe/disagg-deadline-vs-quarantine-proposal.md`. They're +complementary: + +- This PR (deadline anchor): "deadline only fires for actual + in-transfer time." +- The proposal (deadline-vs-quarantine): "deadline-firing should + surface a user error but not drive the unhealthy budget." + +The proposal can be applied on top of this PR if the team agrees +with the two-layer model. Either order works; this PR alone fixes +the immediate false-positive issue under the d224 scenario, +provided no_progress alone is the only path to `is_shutdown` (which +empirically it is on the prefill side under d224 — the second +UNHEALTHY flip is driven by the time-since-progress check, not by +quarantine count).