diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index bcd1363138d7..4633677e3fba 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3687,6 +3687,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/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 32ca8df0c60d..e3b404de9f9c 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -331,7 +331,6 @@ perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4 perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb256_mtp3_ccb-NIXL] SKIP (https://nvbugs/6601537) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_8k1k_con1_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6601537) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_8k1k_con4096_ctx1_dep4_gen1_dep32_eplb256_mtp0_ccb-NIXL] SKIP (https://nvbugs/6601537) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6581075) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL] SKIP (https://nvbugs/6581075) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6601537) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_deepseek-r1-fp4_128k8k_con256_ctx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL] SKIP (https://nvbugs/6581075) diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index d6592dfedd5e..bc9687112a40 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1065,6 +1065,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: