From d8a8ed1e1443dbe560a4608c30bbca3790fbf563 Mon Sep 17 00:00:00 2001
From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Date: Thu, 13 Aug 2026 12:32:39 -0700
Subject: [PATCH] [https://nvbugs/6590664][fix] Reap idle single-rank CTX
 transfers

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py |  6 +++++
 .../_torch/executor/test_py_executor.py       | 26 +++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 7838491f1e8a..733c6bc7b62e 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -3812,6 +3812,12 @@ def _check_disagg_transfer_progress_when_idle(
         # ranks remain aligned and may safely poll context progress.
         if (not uses_async_gen_transfer
                 and not self._is_disagg_gen_only_no_context_benchmark()):
+            # A single-rank CTX worker cannot diverge on a collective. Reap
+            # completed sends while it is idle so their pinned KV blocks can
+            # be reused by the next context requests.
+            if (is_idle and self._dist_size(self.dist, "world_size") == 1 and
+                    self.async_transfer_manager.has_any_inflight_requests()):
+                self._check_disagg_ctx_cache_transfer_status(0)
             return
 
         local_need_gen_check = (uses_async_gen_transfer and local_needs_progress
diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py
index 70a04185aa3c..6160a102d472 100644
--- a/tests/unittest/_torch/executor/test_py_executor.py
+++ b/tests/unittest/_torch/executor/test_py_executor.py
@@ -1071,6 +1071,32 @@ def test_sync_transfer_skips_idle_progress_collectives(
         executor._check_disagg_gen_cache_transfer_status.assert_not_called()
         executor._check_disagg_ctx_cache_transfer_status.assert_not_called()
 
+    def test_sync_single_rank_ctx_reaps_idle_transfer(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.setenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", "1")
+        executor = object.__new__(PyExecutor)
+        executor.dist = Mock(tp_size=1, cp_size=1, world_size=1)
+        executor.async_transfer_manager = Mock()
+        executor.async_transfer_manager.has_any_inflight_requests.return_value = True
+        executor._check_disagg_gen_cache_transfer_status = Mock()
+        executor._check_disagg_ctx_cache_transfer_status = Mock()
+
+        PyExecutor._check_disagg_transfer_progress_when_idle(
+            executor,
+            num_fitting_reqs=0,
+            fitting_disagg_gen_init_requests=[],
+            wait_for_disagg_gen_transfer_progress=True,
+            all_gen_first=False,
+            is_idle=True,
+        )
+
+        executor.dist.allreduce.assert_not_called()
+        executor.dist.tp_allreduce.assert_not_called()
+        executor.dist.tp_cp_allgather.assert_not_called()
+        executor._check_disagg_gen_cache_transfer_status.assert_not_called()
+        executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0)
+
     def test_sync_multi_rank_does_not_wait_for_blocked_peer(
         self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
     ) -> None: