From 32c7843f0d1c45fa70e7493ffd311d9478924743 Mon Sep 17 00:00:00 2001
From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
Date: Wed, 5 Aug 2026 16:01:02 -0700
Subject: [PATCH 1/7] [None][fix] Simplify idle disagg KV transfer progress
 check

`_check_disagg_transfer_progress_when_idle` gated its work behind two
rank-collectives (`_sync_disagg_gen_status_entry` /
`_sync_disagg_ctx_status_entry`) and then issued a blocking
`atLeastNum=1` wait on whichever direction won the vote. The vote input
was derived from purely local scheduler state (`num_fitting_reqs`,
`fitting_disagg_gen_init_requests`, `wait_for_disagg_gen_transfer_progress`,
`all_gen_first`), so every disagg iteration paid for an extra allreduce or
allgather just to decide whether to poll, and the winning branch could
block the executor loop on an unfinished transfer.

Both `_check_disagg_ctx_cache_transfer_status` and
`_check_disagg_gen_cache_transfer_status` already perform their own
internal cross-rank consensus and are safe to enter unconditionally with
`atLeastNum=0`. Entering both non-blocking polls on every iteration keeps
all ranks symmetric without the extra collective, and reaps completed
transfers so their KV blocks are freed just the same. Ranks with nothing
in flight simply reap nothing.

The synchronous-transfer early return is preserved: a synchronous GEN
receive is rank-local and blocking, so one rank can still be receiving
while another is idle, which makes entering either progress collective
unsafe.

Removes the now-unused `_sync_disagg_gen_status_entry` and
`_sync_disagg_ctx_status_entry` helpers and drops the per-iteration
`all_gen_first` scan over `active_requests` at both call sites.

Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py |  98 +++-------------
 .../_torch/executor/test_benchmark_disagg.py  |   8 +-
 .../_torch/executor/test_py_executor.py       | 109 +++---------------
 3 files changed, 37 insertions(+), 178 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 7838491f1e8a..9658484e2d32 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -2788,9 +2788,8 @@ def _executor_loop_pp(self):
                 self._pad_attention_dp_dummy_request()
 
                 # Stage 0: first PP rank schedules requests and propagates the result to all other PP ranks.
-                (scheduled_batch, fitting_disagg_gen_init_requests,
-                 num_fitting_reqs, wait_for_disagg_gen_transfer_progress
-                 ) = self._pp_schedule_and_propagate(microbatch_id)
+                (scheduled_batch, fitting_disagg_gen_init_requests, _,
+                 _) = self._pp_schedule_and_propagate(microbatch_id)
                 if self.dist.rank != 0:
                     # Retry until current rank can run first PP's schedule result.
                     self._pp_retry_until_can_schedule(scheduled_batch)
@@ -2822,17 +2821,7 @@ def _executor_loop_pp(self):
                     self._prepare_disagg_gen_init(
                         fitting_disagg_gen_init_requests)
 
-                    all_gen_first = self.active_requests and all(
-                        req.py_disaggregated_params
-                        and req.py_disaggregated_params.schedule_style ==
-                        DisaggScheduleStyle.GENERATION_FIRST
-                        for req in self.active_requests)
-                    self._check_disagg_transfer_progress_when_idle(
-                        num_fitting_reqs,
-                        fitting_disagg_gen_init_requests,
-                        wait_for_disagg_gen_transfer_progress,
-                        all_gen_first,
-                        is_idle=scheduled_batch.batch_size == 0)
+                    self._check_disagg_transfer_progress_when_idle()
 
                 self.num_scheduled_requests = scheduled_batch.batch_size
 
@@ -3780,73 +3769,25 @@ def _allgather_model_parallel_status(
             return self.dist.tp_allgather(local_status)
         return [local_status]
 
-    def _sync_disagg_gen_status_entry(self, local_need_check: bool) -> int:
-        if self._dist_size(self.dist, "world_size") > 1:
-            return self.dist.allreduce(int(local_need_check), op=ReduceOp.MAX)
-        return int(local_need_check)
-
-    def _sync_disagg_ctx_status_entry(self, local_need_check: bool) -> int:
-        if self._dist_size(self.dist, "cp_size") > 1:
-            return int(any(self.dist.tp_cp_allgather(int(local_need_check))))
-        if self._dist_size(self.dist, "tp_size") > 1:
-            return self.dist.tp_allreduce(int(local_need_check),
-                                          op=ReduceOp.MAX)
-        return int(local_need_check)
-
-    def _check_disagg_transfer_progress_when_idle(
-            self,
-            num_fitting_reqs: int,
-            fitting_disagg_gen_init_requests: List[LlmRequest],
-            wait_for_disagg_gen_transfer_progress: bool,
-            all_gen_first: bool,
-            is_idle: bool = False) -> None:
-        local_needs_progress = (num_fitting_reqs == 0
-                                and not fitting_disagg_gen_init_requests)
-
-        uses_async_gen_transfer = self._uses_async_disagg_gen_transfer()
+    def _check_disagg_transfer_progress_when_idle(self) -> None:
+        """Reap completed KV transfers so their blocks can be freed.
 
+        Both polls are non-blocking and rank-symmetric: every rank enters them
+        unconditionally on every disagg iteration, so the consensus performed
+        inside the status calls stays aligned without an extra collective here.
+        Ranks with nothing in flight simply reap nothing.
+        """
         # A synchronous GEN receive is rank-local and blocking. One rank can
         # still be receiving while another is idle, so entering either the
         # generation or context progress collective here is unsafe. The
         # gen-only-no-context benchmark skips KV transfer entirely, so its
         # ranks remain aligned and may safely poll context progress.
-        if (not uses_async_gen_transfer
+        if (not self._uses_async_disagg_gen_transfer()
                 and not self._is_disagg_gen_only_no_context_benchmark()):
             return
 
-        local_need_gen_check = (uses_async_gen_transfer and local_needs_progress
-                                and wait_for_disagg_gen_transfer_progress)
-
-        any_need_gen_check = self._sync_disagg_gen_status_entry(
-            local_need_gen_check)
-        if any_need_gen_check > 0:
-            if local_need_gen_check:
-                logger.debug(
-                    "Waiting for generation KV cache transfer progress to "
-                    "free disagg admission budget")
-            self._check_disagg_gen_cache_transfer_status(1)
-            return
-
-        local_need_ctx_check = is_idle or (uses_async_gen_transfer
-                                           and local_needs_progress)
-        any_need_check = self._sync_disagg_ctx_status_entry(
-            local_need_ctx_check)
-        if any_need_check > 0:
-            if local_need_ctx_check and not all_gen_first:
-                logger.warning(
-                    "Executor is idle or no disaggregated generation request "
-                    "fits; waiting for context KV cache transfer progress")
-                # Local conditions warrant a blocking wait for at least one
-                # in-flight transfer to complete so KV blocks can be freed.
-                self._check_disagg_ctx_cache_transfer_status(1)
-            else:
-                # Either (a) a peer rank needed the call but we didn't, or
-                # (b) all active requests are gen-first so we don't
-                # actively block. In both cases the non-blocking variant
-                # still runs the internal allgather (keeping all ranks in
-                # sync) and reaps any already-completed transfers without
-                # blocking on un-finished ones.
-                self._check_disagg_ctx_cache_transfer_status(0)
+        self._check_disagg_ctx_cache_transfer_status(0)
+        self._check_disagg_gen_cache_transfer_status(0)
 
     def _sync_gen_only_benchmark_has_insufficient_kv(
             self, scheduler_fitting_disagg_gen_init_requests: List[LlmRequest],
@@ -3984,7 +3925,7 @@ def _prepare_and_schedule_batch(self):
                     request.py_draft_tokens = [0] * self.max_total_draft_tokens
                 request.draft_tokens = [0] * self.max_total_draft_tokens
 
-        scheduled_batch, scheduler_fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule(
+        scheduled_batch, scheduler_fitting_disagg_gen_init_requests, _ = self._schedule(
         )
 
         # Must run after _schedule(): the empty scheduled batch it repairs does
@@ -4004,16 +3945,7 @@ def _prepare_and_schedule_batch(self):
             # into the transfer window this iteration.
             self._prepare_disagg_gen_init(admitted_disagg_gen_init_requests)
 
-            all_gen_first = self.active_requests and all(
-                req.py_disaggregated_params and req.py_disaggregated_params.
-                schedule_style == DisaggScheduleStyle.GENERATION_FIRST
-                for req in self.active_requests)
-            self._check_disagg_transfer_progress_when_idle(
-                num_fitting_reqs,
-                admitted_disagg_gen_init_requests,
-                wait_for_disagg_gen_transfer_progress,
-                all_gen_first,
-                is_idle=scheduled_batch.batch_size == 0)
+            self._check_disagg_transfer_progress_when_idle()
 
             # In gen-only benchmark mode, all requests must fit in KV cache
             # simultaneously. If some requests are stuck in INIT state and the
diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py
index 5aaadbb2e35a..673527105c27 100644
--- a/tests/unittest/_torch/executor/test_benchmark_disagg.py
+++ b/tests/unittest/_torch/executor/test_benchmark_disagg.py
@@ -1196,9 +1196,7 @@ def test_partial_transfer_admission_uses_only_admitted_requests(self) -> None:
         assert result is not None
         ex._apply_disagg_transfer_admission.assert_called_once_with(candidates)
         ex._prepare_disagg_gen_init.assert_called_once_with([admitted_req])
-        ex._check_disagg_transfer_progress_when_idle.assert_called_once_with(
-            0, [admitted_req], False, False, is_idle=True
-        )
+        ex._check_disagg_transfer_progress_when_idle.assert_called_once_with()
         ex._handle_errors.assert_not_called()
 
     def test_fill_with_no_init_requests_does_not_kill(self):
@@ -1231,8 +1229,8 @@ def test_transfer_admission_backpressure_does_not_kill(self, monkeypatch):
         )
         ex._apply_disagg_transfer_admission.assert_called_once_with([fitting_req])
         ex._prepare_disagg_gen_init.assert_called_once_with([])
-        ex._check_disagg_gen_cache_transfer_status.assert_called_once_with(1)
-        ex._check_disagg_ctx_cache_transfer_status.assert_not_called()
+        ex._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0)
+        ex._check_disagg_gen_cache_transfer_status.assert_called_once_with(0)
         ex._handle_errors.assert_not_called()
 
     @pytest.mark.parametrize(
diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py
index 70a04185aa3c..e30d0702b70e 100644
--- a/tests/unittest/_torch/executor/test_py_executor.py
+++ b/tests/unittest/_torch/executor/test_py_executor.py
@@ -87,14 +87,7 @@ def _run_sync_idle_progress_rank(rank: int, world_size: int, rendezvous_file: st
         executor.dist = _TorchCollectiveDist()
         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,
-        )
+        PyExecutor._check_disagg_transfer_progress_when_idle(executor)
     finally:
         torch_dist.destroy_process_group()
 
@@ -969,58 +962,31 @@ def test_gen_transfer_status_skips_sync_mode(self, monkeypatch):
 
         executor._check_disagg_gen_cache_transfer_status.assert_not_called()
 
-    def test_polls_generation_transfer_when_admission_blocked(self):
+    def test_polls_both_transfer_directions_without_blocking(self):
         executor = object.__new__(PyExecutor)
         executor.dist = Mock(tp_size=1)
         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,
-        )
-
-        executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(1)
-        executor._check_disagg_ctx_cache_transfer_status.assert_not_called()
-
-    def test_peer_rank_enters_bounded_progress_poll(self):
-        executor = object.__new__(PyExecutor)
-        executor.dist = Mock(tp_size=1, cp_size=4, world_size=4)
-        executor.dist.allreduce.return_value = 1
-        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=1,
-            fitting_disagg_gen_init_requests=[],
-            wait_for_disagg_gen_transfer_progress=True,
-            all_gen_first=False,
-        )
+        PyExecutor._check_disagg_transfer_progress_when_idle(executor)
 
-        executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(1)
-        executor._check_disagg_ctx_cache_transfer_status.assert_not_called()
-        executor.dist.allreduce.assert_called_once_with(0, op=ReduceOp.MAX)
+        executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0)
+        executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(0)
 
-    def test_falls_back_to_context_transfer_when_not_generation_blocked(self):
+    def test_idle_poll_enters_no_extra_collective(self):
+        """Both polls are rank-symmetric, so no gating collective is needed."""
         executor = object.__new__(PyExecutor)
-        executor.dist = Mock(tp_size=1)
+        executor.dist = Mock(tp_size=4, cp_size=4, world_size=16)
         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=False,
-            all_gen_first=False,
-        )
+        PyExecutor._check_disagg_transfer_progress_when_idle(executor)
 
-        executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(1)
-        executor._check_disagg_gen_cache_transfer_status.assert_not_called()
+        executor.dist.allreduce.assert_not_called()
+        executor.dist.tp_allreduce.assert_not_called()
+        executor.dist.tp_cp_allgather.assert_not_called()
+        executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0)
+        executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(0)
 
     def test_gen_only_no_context_benchmark_polls_context_when_idle(
         self, monkeypatch: pytest.MonkeyPatch
@@ -1028,24 +994,14 @@ def test_gen_only_no_context_benchmark_polls_context_when_idle(
         monkeypatch.setenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY", "1")
         executor = object.__new__(PyExecutor)
         executor.dist = Mock(tp_size=4, cp_size=1, world_size=4)
-        executor.dist.allreduce.return_value = 0
-        executor.dist.tp_allreduce.return_value = 1
         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,
-        )
+        PyExecutor._check_disagg_transfer_progress_when_idle(executor)
 
-        executor.dist.allreduce.assert_called_once_with(0, op=ReduceOp.MAX)
-        executor.dist.tp_allreduce.assert_called_once_with(1, op=ReduceOp.MAX)
-        executor._check_disagg_gen_cache_transfer_status.assert_not_called()
-        executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(1)
+        executor.dist.allreduce.assert_not_called()
+        executor.dist.tp_allreduce.assert_not_called()
+        executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0)
 
     def test_sync_transfer_skips_idle_progress_collectives(
         self, monkeypatch: pytest.MonkeyPatch
@@ -1056,14 +1012,7 @@ def test_sync_transfer_skips_idle_progress_collectives(
         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,
-        )
+        PyExecutor._check_disagg_transfer_progress_when_idle(executor)
 
         executor.dist.allreduce.assert_not_called()
         executor.dist.tp_allreduce.assert_not_called()
@@ -1165,26 +1114,6 @@ def complete_or_error(req):
             charge_budget=False,
         )
 
-    def test_peer_cp_rank_enters_context_progress_poll(self):
-        executor = object.__new__(PyExecutor)
-        executor.dist = Mock(tp_size=1, cp_size=4, world_size=4)
-        executor.dist.allreduce.return_value = 0
-        executor.dist.tp_cp_allgather.return_value = [0, 1, 0, 0]
-        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=1,
-            fitting_disagg_gen_init_requests=[],
-            wait_for_disagg_gen_transfer_progress=False,
-            all_gen_first=False,
-        )
-
-        executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0)
-        executor._check_disagg_gen_cache_transfer_status.assert_not_called()
-        executor.dist.tp_cp_allgather.assert_called_once_with(0)
-
 
 @pytest.mark.usefixtures("_clear_disagg_transfer_mode_env")
 class TestDisaggTransferAdmissionPP:

From 348e737795760c161251488af364b62afeb2641c Mon Sep 17 00:00:00 2001
From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
Date: Thu, 6 Aug 2026 20:32:10 -0700
Subject: [PATCH 2/7] [None][fix] Drop redundant GEN transfer poll from the
 idle progress check

`_check_disagg_transfer_progress_when_idle` polled both directions, but the
GEN poll was always a repeat of one that already ran earlier in the same
iteration:

- The loop head (`_executor_loop_pp` / `_prepare_and_schedule_batch`) calls
  `_check_disagg_gen_transfer_status`, which enters
  `_check_disagg_gen_cache_transfer_status(0)` unconditionally.
- If scheduling started new receives, `_prepare_disagg_gen_init` ->
  `_recv_disagg_gen_cache` already polls GEN status right after issuing them.

So in both cases the second call re-ran the GEN status query and its internal
cross-rank consensus for nothing. Keep only the CTX poll here.

The synchronous-transfer early return is unchanged: a synchronous GEN receive
is rank-local and blocking, so one rank can still be receiving while another is
idle, which makes entering the context progress collective unsafe.

Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 21 ++++++++++++-------
 .../_torch/executor/test_benchmark_disagg.py  |  1 -
 .../_torch/executor/test_py_executor.py       | 17 +++++++++++----
 3 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 9658484e2d32..336294ec5296 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -3770,24 +3770,29 @@ def _allgather_model_parallel_status(
         return [local_status]
 
     def _check_disagg_transfer_progress_when_idle(self) -> None:
-        """Reap completed KV transfers so their blocks can be freed.
+        """Reap completed context KV transfers so their blocks can be freed.
 
-        Both polls are non-blocking and rank-symmetric: every rank enters them
+        The poll is non-blocking and rank-symmetric: every rank enters it
         unconditionally on every disagg iteration, so the consensus performed
-        inside the status calls stays aligned without an extra collective here.
+        inside the status call stays aligned without an extra collective here.
         Ranks with nothing in flight simply reap nothing.
+
+        Generation transfers are deliberately not polled here: the loop head
+        already ran `_check_disagg_gen_transfer_status` this iteration, and any
+        receive started since then by `_prepare_disagg_gen_init` is polled by
+        `_recv_disagg_gen_cache` right after it is issued. A poll here would
+        only repeat the GEN status call and its consensus.
         """
         # A synchronous GEN receive is rank-local and blocking. One rank can
-        # still be receiving while another is idle, so entering either the
-        # generation or context progress collective here is unsafe. The
-        # gen-only-no-context benchmark skips KV transfer entirely, so its
-        # ranks remain aligned and may safely poll context progress.
+        # still be receiving while another is idle, so entering the context
+        # progress collective here is unsafe. The gen-only-no-context
+        # benchmark skips KV transfer entirely, so its ranks remain aligned
+        # and may safely poll context progress.
         if (not self._uses_async_disagg_gen_transfer()
                 and not self._is_disagg_gen_only_no_context_benchmark()):
             return
 
         self._check_disagg_ctx_cache_transfer_status(0)
-        self._check_disagg_gen_cache_transfer_status(0)
 
     def _sync_gen_only_benchmark_has_insufficient_kv(
             self, scheduler_fitting_disagg_gen_init_requests: List[LlmRequest],
diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py
index 673527105c27..217b354223dc 100644
--- a/tests/unittest/_torch/executor/test_benchmark_disagg.py
+++ b/tests/unittest/_torch/executor/test_benchmark_disagg.py
@@ -1230,7 +1230,6 @@ def test_transfer_admission_backpressure_does_not_kill(self, monkeypatch):
         ex._apply_disagg_transfer_admission.assert_called_once_with([fitting_req])
         ex._prepare_disagg_gen_init.assert_called_once_with([])
         ex._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0)
-        ex._check_disagg_gen_cache_transfer_status.assert_called_once_with(0)
         ex._handle_errors.assert_not_called()
 
     @pytest.mark.parametrize(
diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py
index e30d0702b70e..93fab06ec096 100644
--- a/tests/unittest/_torch/executor/test_py_executor.py
+++ b/tests/unittest/_torch/executor/test_py_executor.py
@@ -962,7 +962,7 @@ def test_gen_transfer_status_skips_sync_mode(self, monkeypatch):
 
         executor._check_disagg_gen_cache_transfer_status.assert_not_called()
 
-    def test_polls_both_transfer_directions_without_blocking(self):
+    def test_polls_context_transfers_without_blocking(self):
         executor = object.__new__(PyExecutor)
         executor.dist = Mock(tp_size=1)
         executor._check_disagg_gen_cache_transfer_status = Mock()
@@ -971,10 +971,20 @@ def test_polls_both_transfer_directions_without_blocking(self):
         PyExecutor._check_disagg_transfer_progress_when_idle(executor)
 
         executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0)
-        executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(0)
+
+    def test_does_not_repeat_gen_status_polled_by_loop_head(self):
+        """The loop head already polls GEN status every iteration."""
+        executor = object.__new__(PyExecutor)
+        executor.dist = Mock(tp_size=1)
+        executor._check_disagg_gen_cache_transfer_status = Mock()
+        executor._check_disagg_ctx_cache_transfer_status = Mock()
+
+        PyExecutor._check_disagg_transfer_progress_when_idle(executor)
+
+        executor._check_disagg_gen_cache_transfer_status.assert_not_called()
 
     def test_idle_poll_enters_no_extra_collective(self):
-        """Both polls are rank-symmetric, so no gating collective is needed."""
+        """The context poll is rank-symmetric, so no gating collective is needed."""
         executor = object.__new__(PyExecutor)
         executor.dist = Mock(tp_size=4, cp_size=4, world_size=16)
         executor._check_disagg_gen_cache_transfer_status = Mock()
@@ -986,7 +996,6 @@ def test_idle_poll_enters_no_extra_collective(self):
         executor.dist.tp_allreduce.assert_not_called()
         executor.dist.tp_cp_allgather.assert_not_called()
         executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0)
-        executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(0)
 
     def test_gen_only_no_context_benchmark_polls_context_when_idle(
         self, monkeypatch: pytest.MonkeyPatch

From a028c572212c94e317f8ac80f8e8730706380cfb Mon Sep 17 00:00:00 2001
From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
Date: Fri, 7 Aug 2026 18:20:22 -0700
Subject: [PATCH 3/7] [None][fix] Warn when the scheduler fits nothing on the
 idle disagg path

The idle progress check no longer blocks on a transfer, and the request queue
does not block either while INIT/TRANS requests are active, so nothing named
the KV-starved state after the blocking branch was removed. Log it again.

Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 336294ec5296..f3e242b95d7f 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -2788,7 +2788,8 @@ def _executor_loop_pp(self):
                 self._pad_attention_dp_dummy_request()
 
                 # Stage 0: first PP rank schedules requests and propagates the result to all other PP ranks.
-                (scheduled_batch, fitting_disagg_gen_init_requests, _,
+                (scheduled_batch, fitting_disagg_gen_init_requests,
+                 num_fitting_reqs,
                  _) = self._pp_schedule_and_propagate(microbatch_id)
                 if self.dist.rank != 0:
                     # Retry until current rank can run first PP's schedule result.
@@ -2821,6 +2822,9 @@ def _executor_loop_pp(self):
                     self._prepare_disagg_gen_init(
                         fitting_disagg_gen_init_requests)
 
+                    if num_fitting_reqs == 0:
+                        logger.warning(
+                            "num_fitting_reqs=0, may not have enough kvCache")
                     self._check_disagg_transfer_progress_when_idle()
 
                 self.num_scheduled_requests = scheduled_batch.batch_size
@@ -3930,7 +3934,7 @@ def _prepare_and_schedule_batch(self):
                     request.py_draft_tokens = [0] * self.max_total_draft_tokens
                 request.draft_tokens = [0] * self.max_total_draft_tokens
 
-        scheduled_batch, scheduler_fitting_disagg_gen_init_requests, _ = self._schedule(
+        scheduled_batch, scheduler_fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule(
         )
 
         # Must run after _schedule(): the empty scheduled batch it repairs does
@@ -3950,6 +3954,9 @@ def _prepare_and_schedule_batch(self):
             # into the transfer window this iteration.
             self._prepare_disagg_gen_init(admitted_disagg_gen_init_requests)
 
+            if num_fitting_reqs == 0:
+                logger.warning(
+                    "num_fitting_reqs=0, may not have enough kvCache")
             self._check_disagg_transfer_progress_when_idle()
 
             # In gen-only benchmark mode, all requests must fit in KV cache

From feb2e1184554d8966f59909c86d404787bbf92b2 Mon Sep 17 00:00:00 2001
From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
Date: Thu, 20 Aug 2026 15:47:53 -0700
Subject: [PATCH 4/7] [None][fix] Pace the executor loop when the forward pass
 is skipped

Removing the blocking `atLeastNum=1` wait also removed the only pacing for
the idle-blocked case. `_fetch_and_enqueue_requests` uses a zero timeout
while any request is active, so once nothing fits, the loop re-ran the ADP
allgather, both transfer-status gathers and a full schedule pass at full
speed until a transfer completed, burning a core and contending with the
transceiver's own progress threads.

Sleep 1ms on iterations where `_can_queue` came back False. `can_queue` is
already a fleet-wide consensus, so every rank pauses on the same iterations
without adding a collective, and the loop head still runs each iteration so
cancellations, control requests, shutdown and generation-transfer reaping
keep being serviced.

Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index f3e242b95d7f..86a3c4c6f124 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -4348,6 +4348,11 @@ def _executor_loop(self):
                 if not can_queue and scheduled_batch.encoder_requests:
                     self._run_encoder_step(scheduled_batch.encoder_requests)
 
+                if not can_queue:
+                    # Nothing runs this iteration; only a KV transfer completing
+                    # can unblock it, so pace the loop instead of spinning.
+                    time.sleep(0.001)
+
                 if can_queue:
                     # init_disagg_gen_requests must be before drafter loop, otherwise draft requests do not have initialized matchers.
                     # init_disagg_gen_requests must be before engine forward, where the prev_seq_slot is updated.
@@ -5179,6 +5184,11 @@ def _executor_loop_overlap(self):
                 if not can_queue and scheduled_batch.encoder_requests:
                     self._run_encoder_step(scheduled_batch.encoder_requests)
 
+                if not can_queue:
+                    # Nothing runs this iteration; only a KV transfer completing
+                    # can unblock it, so pace the loop instead of spinning.
+                    time.sleep(0.001)
+
                 # If the batch is not empty on this rank, but empty on other ranks,
                 # we need to delay the update of the previous batch's sample state,
                 # and let the later iteration to update it.

From cd39ea83bb01658df9e629c47fa345fae9094a9a Mon Sep 17 00:00:00 2001
From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
Date: Thu, 20 Aug 2026 15:53:00 -0700
Subject: [PATCH 5/7] [None][fix] Drop the num_fitting_reqs=0 warning on the
 idle disagg path

The restored warning is broader than the one it replaced, which also
required `not fitting_disagg_gen_init_requests` and suppressed the
generation-first case. `num_fitting_reqs` does not count fitting gen-init
requests -- the capacity scheduler returns those separately -- so the
condition holds during normal generation ramp-up and while a context server
drains in-flight transfers, making it fire every iteration in states that
are not actually degraded.

Drop it rather than re-adding the guards: the caller already has the
information, and a log that cannot distinguish KV exhaustion from routine
ramp-up is not worth the rate-limiting it would need.

`num_fitting_reqs` becomes unused at both scheduling sites, so those revert
to `_`. `_pp_schedule_and_propagate` still returns it as part of its
serialized contract.

Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 11 ++---------
 1 file changed, 2 insertions(+), 9 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 86a3c4c6f124..eef1483934d5 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -2788,8 +2788,7 @@ def _executor_loop_pp(self):
                 self._pad_attention_dp_dummy_request()
 
                 # Stage 0: first PP rank schedules requests and propagates the result to all other PP ranks.
-                (scheduled_batch, fitting_disagg_gen_init_requests,
-                 num_fitting_reqs,
+                (scheduled_batch, fitting_disagg_gen_init_requests, _,
                  _) = self._pp_schedule_and_propagate(microbatch_id)
                 if self.dist.rank != 0:
                     # Retry until current rank can run first PP's schedule result.
@@ -2822,9 +2821,6 @@ def _executor_loop_pp(self):
                     self._prepare_disagg_gen_init(
                         fitting_disagg_gen_init_requests)
 
-                    if num_fitting_reqs == 0:
-                        logger.warning(
-                            "num_fitting_reqs=0, may not have enough kvCache")
                     self._check_disagg_transfer_progress_when_idle()
 
                 self.num_scheduled_requests = scheduled_batch.batch_size
@@ -3934,7 +3930,7 @@ def _prepare_and_schedule_batch(self):
                     request.py_draft_tokens = [0] * self.max_total_draft_tokens
                 request.draft_tokens = [0] * self.max_total_draft_tokens
 
-        scheduled_batch, scheduler_fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule(
+        scheduled_batch, scheduler_fitting_disagg_gen_init_requests, _ = self._schedule(
         )
 
         # Must run after _schedule(): the empty scheduled batch it repairs does
@@ -3954,9 +3950,6 @@ def _prepare_and_schedule_batch(self):
             # into the transfer window this iteration.
             self._prepare_disagg_gen_init(admitted_disagg_gen_init_requests)
 
-            if num_fitting_reqs == 0:
-                logger.warning(
-                    "num_fitting_reqs=0, may not have enough kvCache")
             self._check_disagg_transfer_progress_when_idle()
 
             # In gen-only benchmark mode, all requests must fit in KV cache

From efa5e80b4d01c15a212768f8f200d1bf5e791f0e Mon Sep 17 00:00:00 2001
From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
Date: Mon, 24 Aug 2026 14:52:48 -0700
Subject: [PATCH 6/7] [None][fix] Pace the idle disagg loop after the pass
 drains its ready work

The 1ms sleep sat between `_can_queue` and the rest of the iteration, so it
delayed work that was already finished and waiting to be handed off: on a
`not can_queue` pass the overlap loop still runs `_update_requests`,
`_send_kv_async` and the first-token responses for the previous batch, and
both loops still flush pending transfer responses and enter the synced
transfer-timeout collective. None of that needed to wait a millisecond.

Move the sleep to the end of the iteration, once that work has drained, and
put it behind `_pace_idle_disagg_loop`. Sleeping there also means the
pending-transfer check reads the state the pass left behind rather than the
state it started with, so a transfer that landed during the pass no longer
costs a sleep.

Only pace when a KV transfer is what the loop is waiting on. Context sends
are tracked by the transfer manager and generation receives by request
state, so both directions are covered. Outside disaggregated serving, and on
a rank with nothing in flight, the sleep is skipped entirely -- it never
gates a collective, so ranks taking it on different iterations is safe.

Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 47 +++++++++++++++----
 .../_torch/executor/test_py_executor.py       | 47 +++++++++++++++++++
 2 files changed, 84 insertions(+), 10 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index eef1483934d5..2cafad649c2b 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -3794,6 +3794,37 @@ def _check_disagg_transfer_progress_when_idle(self) -> None:
 
         self._check_disagg_ctx_cache_transfer_status(0)
 
+    def _pace_idle_disagg_loop(self) -> None:
+        """Sleep briefly when only a KV transfer completing can make progress.
+
+        Dropping the blocking `atLeastNum=1` wait also dropped the only pacing
+        for the idle-blocked case: `_fetch_and_enqueue_requests` uses a zero
+        timeout while any request is active, so the loop would otherwise re-run
+        the schedule pass and its collectives at full speed until a transfer
+        lands.
+
+        Call this at the end of an iteration that queued nothing, once the pass
+        has drained its ready work. Request updates, KV sends and responses
+        must not be held behind the sleep, and running them first means the
+        pending-transfer check below sees the state they left behind rather
+        than a stale one.
+
+        That check is rank-local. The sleep only paces and never gates a
+        collective, so ranks taking it on different iterations is safe.
+        """
+        if self.kv_cache_transceiver is None:
+            return
+
+        # Context sends are tracked by the transfer manager; generation
+        # receives live in the request state, so both directions are covered.
+        waiting_on_transfer = (
+            self.async_transfer_manager.has_any_inflight_requests()
+            or any(req.is_disagg_generation_init_state
+                   or req.is_disagg_generation_transmission_in_progress
+                   for req in self.active_requests))
+        if waiting_on_transfer:
+            time.sleep(0.001)
+
     def _sync_gen_only_benchmark_has_insufficient_kv(
             self, scheduler_fitting_disagg_gen_init_requests: List[LlmRequest],
             wait_for_disagg_gen_transfer_progress: bool) -> bool:
@@ -4341,11 +4372,6 @@ def _executor_loop(self):
                 if not can_queue and scheduled_batch.encoder_requests:
                     self._run_encoder_step(scheduled_batch.encoder_requests)
 
-                if not can_queue:
-                    # Nothing runs this iteration; only a KV transfer completing
-                    # can unblock it, so pace the loop instead of spinning.
-                    time.sleep(0.001)
-
                 if can_queue:
                     # init_disagg_gen_requests must be before drafter loop, otherwise draft requests do not have initialized matchers.
                     # init_disagg_gen_requests must be before engine forward, where the prev_seq_slot is updated.
@@ -4508,6 +4534,9 @@ def _executor_loop(self):
                 # TLLM_METRICS_ALL_RANKS=0.
                 self._flush_iter_stats_synced()
 
+                if not can_queue:
+                    self._pace_idle_disagg_loop()
+
                 self.iter_counter += 1
 
     def _prepare_draft_requests(self):
@@ -5177,11 +5206,6 @@ def _executor_loop_overlap(self):
                 if not can_queue and scheduled_batch.encoder_requests:
                     self._run_encoder_step(scheduled_batch.encoder_requests)
 
-                if not can_queue:
-                    # Nothing runs this iteration; only a KV transfer completing
-                    # can unblock it, so pace the loop instead of spinning.
-                    time.sleep(0.001)
-
                 # If the batch is not empty on this rank, but empty on other ranks,
                 # we need to delay the update of the previous batch's sample state,
                 # and let the later iteration to update it.
@@ -5388,6 +5412,9 @@ def _executor_loop_overlap(self):
 
                 self._kv_connector_terminate_requests()
 
+                if not can_queue:
+                    self._pace_idle_disagg_loop()
+
                 self.iter_counter += 1
 
     @nvtx_range("_accept_draft_tokens")
diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py
index 93fab06ec096..bf40b080d26d 100644
--- a/tests/unittest/_torch/executor/test_py_executor.py
+++ b/tests/unittest/_torch/executor/test_py_executor.py
@@ -1124,6 +1124,53 @@ def complete_or_error(req):
         )
 
 
+class TestIdleDisaggLoopPacing:
+    """The idle poll no longer blocks, so the executor loops pace themselves.
+
+    Pacing must cost nothing when a transfer is not what is holding the loop
+    back, and the PP loop must not pace while the ring still has work.
+    """
+
+    @staticmethod
+    def _make_request(*, init_state: bool = False, transfer_in_progress: bool = False) -> Mock:
+        req = Mock()
+        req.is_disagg_generation_init_state = init_state
+        req.is_disagg_generation_transmission_in_progress = transfer_in_progress
+        return req
+
+    @pytest.mark.parametrize(
+        "has_transceiver, ctx_inflight, request_kwargs, expect_sleep",
+        [
+            pytest.param(False, True, {"init_state": True}, False, id="not_disagg"),
+            pytest.param(True, False, {}, False, id="nothing_pending"),
+            pytest.param(True, True, {}, True, id="context_send_inflight"),
+            pytest.param(True, False, {"init_state": True}, True, id="gen_awaiting_transfer"),
+            pytest.param(
+                True, False, {"transfer_in_progress": True}, True, id="gen_receive_inflight"
+            ),
+        ],
+    )
+    def test_paces_only_when_a_transfer_can_unblock_the_loop(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        has_transceiver: bool,
+        ctx_inflight: bool,
+        request_kwargs: dict,
+        expect_sleep: bool,
+    ) -> None:
+        sleep = Mock()
+        monkeypatch.setattr("tensorrt_llm._torch.pyexecutor.py_executor.time.sleep", sleep)
+        executor = object.__new__(PyExecutor)
+        executor.kv_cache_transceiver = Mock() if has_transceiver else None
+        executor.async_transfer_manager = Mock()
+        executor.async_transfer_manager.has_any_inflight_requests.return_value = ctx_inflight
+        executor.active_requests = [self._make_request(**request_kwargs)]
+
+        PyExecutor._pace_idle_disagg_loop(executor)
+
+        assert sleep.called is expect_sleep
+
+
 @pytest.mark.usefixtures("_clear_disagg_transfer_mode_env")
 class TestDisaggTransferAdmissionPP:
     def test_pp_schedule_applies_gate_before_serializing(self):

From 4a177b024891f26709b1a831f3bbb68f2e640960 Mon Sep 17 00:00:00 2001
From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
Date: Mon, 24 Aug 2026 14:53:22 -0700
Subject: [PATCH 7/7] [None][fix] Pace the PP executor loop when the ring is
 idle

The pacing added for the other two loops skipped `_executor_loop_pp`, which
can reach the same state: a microbatch that cannot be queued leaves the loop
re-running the schedule pass and its ring collectives at full speed.

It is only the same state once the ring is empty. While microbatches are
still in flight, `fetch_executed_batches` sets `must_get = not can_queue` and
blocks on the response queue until one finishes, so the loop is already
paced -- and sleeping there would delay the relay of batches that are
finished and waiting to be handled.

Pace only when nothing is queued and nothing is outstanding, and do it after
stage 3.4 so a rebalance that the drained ring just unblocked still completes
in the same iteration.

Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py   | 13 +++++++++++++
 .../_torch/executor/test_py_executor.py         | 17 +++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 2cafad649c2b..e2a034f818c3 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -3081,6 +3081,9 @@ def handle_executed_batches(executed_batch_num: int):
                 if self._uses_kv_manager_v2():
                     self._maybe_finish_pp_rebalance()
 
+                if not can_queue and self._pp_ring_is_drained():
+                    self._pace_idle_disagg_loop()
+
                 # Stage 4: March forward in microbatch slots
                 microbatch_id = (microbatch_id + 1) % self.num_micro_batches
                 self.iter_counter += 1
@@ -3825,6 +3828,16 @@ def _pace_idle_disagg_loop(self) -> None:
         if waiting_on_transfer:
             time.sleep(0.001)
 
+    def _pp_ring_is_drained(self) -> bool:
+        """Return whether no microbatch is queued or awaiting handling.
+
+        While microbatches are still in flight `fetch_executed_batches` blocks
+        on the response queue, which paces the loop on its own; sleeping on top
+        of that would only delay their relay.
+        """
+        return (self.unhandled_batch_counter == 0
+                and all(batch is None for batch in self.micro_batches))
+
     def _sync_gen_only_benchmark_has_insufficient_kv(
             self, scheduler_fitting_disagg_gen_init_requests: List[LlmRequest],
             wait_for_disagg_gen_transfer_progress: bool) -> bool:
diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py
index bf40b080d26d..5254945b446d 100644
--- a/tests/unittest/_torch/executor/test_py_executor.py
+++ b/tests/unittest/_torch/executor/test_py_executor.py
@@ -1170,6 +1170,23 @@ def test_paces_only_when_a_transfer_can_unblock_the_loop(
 
         assert sleep.called is expect_sleep
 
+    @pytest.mark.parametrize(
+        "unhandled_batches, micro_batches, expected",
+        [
+            pytest.param(0, [None, None], True, id="ring_empty"),
+            pytest.param(1, [None, None], False, id="batch_awaiting_handling"),
+            pytest.param(0, [None, "batch"], False, id="batch_still_queued"),
+        ],
+    )
+    def test_pp_ring_drained_only_when_no_microbatch_is_outstanding(
+        self, unhandled_batches: int, micro_batches: list, expected: bool
+    ) -> None:
+        executor = object.__new__(PyExecutor)
+        executor.unhandled_batch_counter = unhandled_batches
+        executor.micro_batches = micro_batches
+
+        assert PyExecutor._pp_ring_is_drained(executor) is expected
+
 
 @pytest.mark.usefixtures("_clear_disagg_transfer_mode_env")
 class TestDisaggTransferAdmissionPP: