[https://nvbugs/6104831][fix] Anchor KV transfer deadline to actual-transfer-start, not admit-time

cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h

@@ -200,28 +200,139 @@ struct RequestStatuses
     std::unordered_set<LlmRequest::RequestIdType> 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:
     virtual ~BaseCacheTransceiver() = default;
-    virtual void respondAndSendAsync(LlmRequest* llmRequest) = 0;
+
+    // Methods take std::shared_ptr<LlmRequest> 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> llmRequest) = 0;
     virtual void respondAndSendLayerWise(
         RequestVector const& requests, std::shared_ptr<ContextProgress> const& progress)
         = 0;
 
-    virtual void requestAndReceiveSync(LlmRequest* llmRequest) = 0;
-    virtual void requestAndReceiveAsync(LlmRequest* llmRequest) = 0;
+    virtual void requestAndReceiveSync(std::shared_ptr<LlmRequest> llmRequest) = 0;
+    virtual void requestAndReceiveAsync(std::shared_ptr<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<int> const& atLeastRequestNum = std::nullopt, bool markComplete = false)
         = 0;
 
+    /// Non-blocking poll. Same contract as @ref checkContextTransferStatus.
     virtual void checkGenTransferStatus(std::optional<int> 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;
 
-    virtual bool cancelRequest(LlmRequest* llmRequest) = 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(std::shared_ptr<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(std::shared_ptr<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
@@ -252,34 +363,111 @@ class CacheTransceiver : public BaseCacheTransceiver
 
     virtual ~CacheTransceiver();
 
-    void respondAndSendAsync(LlmRequest* llmRequest) override;
+    void respondAndSendAsync(std::shared_ptr<LlmRequest> llmRequest) override;
 
     void respondAndSendLayerWise(
         RequestVector const& requests, std::shared_ptr<ContextProgress> const& progress) override;
 
-    void requestAndReceiveSync(LlmRequest* llmRequest) override;
-    void requestAndReceiveAsync(LlmRequest* llmRequest) override;
+    void requestAndReceiveSync(std::shared_ptr<LlmRequest> llmRequest) override;
+    void requestAndReceiveAsync(std::shared_ptr<LlmRequest> llmRequest) override;
 
     RequestStatuses checkContextTransferStatus(
         std::optional<int> const& atLeastRequestNum = std::nullopt, bool markComplete = false) override;
 
     void checkGenTransferStatus(std::optional<int> 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;
 
-    virtual bool cancelRequest(LlmRequest* llmRequest) override;
+    TransferCancelResult cancelRequestStructured(std::shared_ptr<LlmRequest> llmRequest) override;
+
+    virtual bool cancelRequest(std::shared_ptr<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. 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.
+    ///
+    /// 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<LlmRequest> request;
+        std::future<void> future;
+        /// 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<TrackedFuture>& 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<int> const& atLeastRequestNum, bool markComplete, bool allowBlocking);
+    void checkGenTransferStatusImpl(std::optional<int> const& atLeastRequestNum, bool allowBlocking);
+
     std::unique_ptr<CacheSender> mCacheSender;
     std::unique_ptr<CacheReceiver> mCacheReceiver;
-    std::vector<std::pair<LlmRequest*, std::future<void>>> mSenderFutures;
-    std::vector<std::pair<LlmRequest*, std::future<void>>> mRequesterFutures;
+    std::vector<TrackedFuture> mSenderFutures;
+    std::vector<TrackedFuture> 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<CacheTransceiverComm> mGroupComm;
     std::shared_ptr<CacheTransceiverComm> mGroupTensorParaComm, mGroupPipeParaComm, mGroupDataComm, mGroupTPInDPComm;
 

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<executor::GuidedDecodingParams> mGuidedDecodingParams{std::nullopt};
 

cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp

@@ -21,6 +21,7 @@
 #include "tensorrt_llm/common/logger.h"
 #include "tensorrt_llm/common/opUtils.h"
 
+#include <chrono>
 #include <mutex>
 
 namespace tensorrt_llm::batch_manager
@@ -233,8 +234,41 @@ std::optional<int> BaseTransBufferManager::assignBufferIndex(
         return std::nullopt;
     }
     std::unique_lock lk(resource.mBuffersMutex);
-    resource.mBuffersCV.wait(
-        lk, [&resource, bufferCount]() { return static_cast<size_t>(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<size_t>(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<std::chrono::milliseconds>(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<long long>(elapsedMs), static_cast<int>(resource.mConcurrence), bufferCount);
+            lastHeartbeat = now;
+            (void) lastHeartbeat;
+        }
+    }
+    else
+    {
+        resource.mBuffersCV.wait(lk, haveSlot);
+    }
     int bufferId = -1;
     for (size_t i = 0; i < bufferCount; i++)
     {