Draft: [None][fix] Use KV cache transfer future timeout on generation worker

cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -599,6 +599,13 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus(
 void CacheTransceiver::checkGenTransferStatus(std::optional<int> const& atLeastRequestNum)
 {
     bool blockAll = !atLeastRequestNum.has_value();
+    std::optional<int> receiverFutureTimeoutMs = std::nullopt;
+    // If blockAll is true, we want to block and not use a timeout
+    if (!blockAll && mCacheTransceiverConfig.has_value())
+    {
+        receiverFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs();
+    }
+
     std::vector<LlmRequest::RequestIdType> genTransferReadyRequestIds;
     for (auto&& [request, future] : mRequesterFutures)
     {
@@ -709,41 +716,67 @@ void CacheTransceiver::checkGenTransferStatus(std::optional<int> const& atLeastR
             " checkGenTransferStatus toCompleteIdSet size: %zu, atLeastRequestNum: %d ", toCompleteIdSet.size(),
             atLeastRequestNum.value_or(0));
     }
+    auto const syncSize = (syncComm != nullptr) ? syncComm->getSize() : 1;
     for (auto it = mRequesterFutures.begin(); it != mRequesterFutures.end();)
     {
         if (blockAll || toCompleteIdSet.find(it->first->mRequestId) != toCompleteIdSet.end())
         {
             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
-                if (!common::getEnvKVCacheTimeOutputPath().empty())
+                // Wait for up to a specified timeout
+                auto status = it->second.wait_for(std::chrono::milliseconds(receiverFutureTimeoutMs.value_or(0)));
+                if (status == std::future_status::ready || !receiverFutureTimeoutMs.has_value())
                 {
-                    auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm;
-                    updateKVCacheTransferBW(syncComm, it->first);
+                    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.
+                    // Only call the timing collective when either all ranks block together (blockAll)
+                    // or the request was confirmed ready on every rank in the initial poll, to avoid
+                    // hanging in allgather when a peer timed out and skipped this request.
+                    if (!common::getEnvKVCacheTimeOutputPath().empty())
+                    {
+                        auto const freqIt = frequencyMap.find(it->first->mRequestId);
+                        if (blockAll || (freqIt != frequencyMap.end() && freqIt->second == syncSize))
+                        {
+                            updateKVCacheTransferBW(syncComm, it->first);
+                        }
+                    }
+                    if (useMPI())
+                    {
+                        TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(),
+                            "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***",
+                            it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId());
+                    }
+                    else
+                    {
+                        TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(),
+                            "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***",
+                            it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId());
+                    }
+                    it = mRequesterFutures.erase(it);
+                }
+                else if (status == std::future_status::timeout)
+                {
+                    TLLM_LOG_WARNING("Timed out waiting for generation KV cache transfer after %d milliseconds.",
+                        receiverFutureTimeoutMs.value());
+                    ++it;
+                }
+                else
+                {
+                    TLLM_LOG_ERROR(
+                        "Future returned unexpected status for request %ld. Marking as error", it->first->mRequestId);
+                    it->first->setState(LlmRequestState::kDISAGG_TRANS_ERROR);
+                    it = mRequesterFutures.erase(it);
                 }
             }
             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);
+                it = mRequesterFutures.erase(it);
             }
-            if (useMPI())
-            {
-                TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(),
-                    "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***",
-                    it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId());
-            }
-            else
-            {
-                TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(),
-                    "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***",
-                    it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId());
-            }
-            it = mRequesterFutures.erase(it);
         }
         else
         {

tensorrt_llm/llmapi/llm_args.py

@@ -2563,7 +2563,7 @@ class CacheTransceiverConfig(StrictBaseModel, PybindMirror):
     kv_transfer_sender_future_timeout_ms: Optional[PositiveInt] = Field(
         default=1000,
         description=
-        "Timeout in milliseconds to wait for the sender future to be ready when scheduled batch size is 0. This allows the request to be eventually cancelled by the user or because of kv_transfer_timeout_ms"
+        "Timeout in milliseconds to wait for the sender/receiver future to be ready. This allows the request to be eventually cancelled by the user or because of kv_transfer_timeout_ms"
     )
 
     def _to_pybind(self):

tests/unittest/llmapi/test_llm_pytorch.py

@@ -1651,6 +1651,105 @@ def test_llm_context_only_timed_out_kv_cache_exhausted(sender_future_timeout_ms,
     assert final_used_num_blocks == 0
 
 
+@pytest.mark.threadleak(enabled=False)
+@pytest.mark.part0
+@skip_ray
+def test_llm_gen_only_transfer_timeout():
+    """Gen-only request flagged for KV transfer that never arrives is eventually cancelled.
+
+    A live ctx endpoint is used so the gen worker enters
+    kDISAGG_GENERATION_TRANS_IN_PROGRESS.  The ctx worker receives the request
+    but can never match the bogus request-ID to a response, so it never sends
+    the ready signal and the gen-side future never resolves.
+
+    After kv_transfer_timeout_ms the Python layer marks the request as timed
+    out.  Shutting down the ctx LLM closes the connection, which unblocks the
+    gen-side background thread; the request is then terminated and KV blocks
+    are freed.
+    """
+    kv_transfer_timeout_ms = 2000
+    kv_transfer_sender_future_timeout_ms = 100
+
+    llm_args_extra = dict(
+        enable_iter_perf_stats=True,
+        disable_overlap_scheduler=True,
+    )
+
+    llm_ctx = LLM(model=llama_model_path,
+                  kv_cache_config=global_kvcache_config_no_reuse,
+                  tensor_parallel_size=1,
+                  cache_transceiver_config=CacheTransceiverConfig(
+                      backend="UCX",
+                      kv_transfer_timeout_ms=kv_transfer_timeout_ms),
+                  **llm_args_extra)
+
+    llm_gen = LLM(model=llama_model_path,
+                  kv_cache_config=global_kvcache_config_no_reuse,
+                  tensor_parallel_size=1,
+                  cache_transceiver_config=CacheTransceiverConfig(
+                      backend="UCX",
+                      kv_transfer_timeout_ms=kv_transfer_timeout_ms,
+                      kv_transfer_sender_future_timeout_ms=
+                      kv_transfer_sender_future_timeout_ms),
+                  **llm_args_extra)
+
+    ctx_shutdown = False
+    try:
+        prompt = "What is your name?"
+
+        # Run a real context-only request to obtain a live ctx_info_endpoint.
+        ctx_outputs = list(
+            llm_ctx.generate([prompt],
+                             sampling_params=SamplingParams(max_tokens=1),
+                             disaggregated_params=DisaggregatedParams(
+                                 request_type="context_only")))
+        assert len(ctx_outputs) == 1
+        ctx_disagg = ctx_outputs[0].disaggregated_params
+
+        # Build gen-only params that reuse the valid endpoint but carry a bogus
+        # ctx_request_id.  The ctx worker will connect and then block waiting
+        # for a response with that nonexistent ID, so it never sends the ready
+        # signal and the gen-side future never resolves.
+        gen_disagg_params = DisaggregatedParams(
+            request_type="generation_only",
+            first_gen_tokens=ctx_disagg.first_gen_tokens,
+            ctx_request_id=ctx_disagg.ctx_request_id + 99999,
+            opaque_state=ctx_disagg.opaque_state,
+            ctx_info_endpoint=ctx_disagg.ctx_info_endpoint,
+        )
+
+        # Submit the stuck gen request without waiting for it.
+        llm_gen.generate_async(prompt,
+                               sampling_params=SamplingParams(max_tokens=10),
+                               disaggregated_params=gen_disagg_params)
+
+        # Wait for kv_transfer_timeout_ms to fire on the gen side.
+        time.sleep((kv_transfer_timeout_ms / 1000) * 3)
+
+        # Shutting down the ctx LLM closes its network connection to the gen
+        # worker, which unblocks the gen-side background thread and causes it
+        # to complete with a connection error.  The gen executor then frees
+        # the KV blocks for the timed-out request.
+        llm_ctx.shutdown()
+        ctx_shutdown = True
+
+        # Verify that all KV blocks on the gen side are freed.
+        max_retries = 15
+        for _ in range(max_retries):
+            results = llm_gen.get_stats(2)
+            if results and results[-1]["kvCacheStats"]["usedNumBlocks"] == 0:
+                break
+            time.sleep(1)
+        else:
+            pytest.fail(
+                "KV blocks were not freed after gen-side KV transfer timeout")
+
+    finally:
+        if not ctx_shutdown:
+            llm_ctx.shutdown()
+        llm_gen.shutdown()
+
+
 @pytest.mark.threadleak(enabled=False)
 @pytest.mark.part0
 @skip_ray