[https://nvbugs/6527301][fix] Do not block ranks with scheduled forward work on disagg gen KV-transfer progress

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -155,6 +155,11 @@ class _SleepWakeupAction(StrEnum):
 _SLEEP_WAKEUP_ACK_TIMEOUT_S = 30.0
 _SLEEP_WAKEUP_ACK_POLL_INTERVAL_S = 0.01
 
+# Per-rank status bits combined (OR across ranks) when deciding how to poll
+# disaggregated generation KV-transfer progress before a forward step.
+_DISAGG_GEN_NEED_TRANSFER_PROGRESS = 1
+_DISAGG_GEN_HAS_FORWARD_WORK = 2
+
 
 def _sleep_wakeup_ack_ready(comm, source: int, tag: _SleepWakeupTag) -> bool:
     """Return whether an ACK is ready without blocking on recv."""
@@ -2595,7 +2600,8 @@ def _executor_loop_pp(self):
                         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)
+                        wait_for_disagg_gen_transfer_progress, all_gen_first,
+                        scheduled_batch.batch_size > 0)
 
                 self.num_scheduled_requests = scheduled_batch.batch_size
 
@@ -3493,10 +3499,15 @@ 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:
+    def _sync_disagg_gen_progress_state(self, local_state: int) -> int:
+        """OR the per-rank disagg gen progress status bits across all ranks.
+
+        BOR is supported by both backends here: MPI object-mode reduction
+        applies Python ``|``, and torch routes CPU scalars to gloo.
+        """
         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)
+            return self.dist.allreduce(local_state, op=ReduceOp.BOR)
+        return local_state
 
     def _sync_disagg_ctx_status_entry(self, local_need_check: bool) -> int:
         if self._dist_size(self.dist, "cp_size") > 1:
@@ -3509,8 +3520,22 @@ def _sync_disagg_ctx_status_entry(self, local_need_check: bool) -> int:
     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) -> None:
+            wait_for_disagg_gen_transfer_progress: bool, all_gen_first: bool,
+            has_forward_work: bool) -> None:
+        """Poll disagg KV-transfer progress between scheduling and forward.
+
+        All ranks branch on the same OR-combined global state, so every rank
+        passes the same ``at_least_request_num`` to the consensus collectives
+        inside ``check_gen_transfer_status``.
+
+        ``has_forward_work`` reflects the final scheduled batch and is the
+        authoritative work signal; ``num_fitting_reqs`` carries
+        scheduler-specific semantics and post-scheduling adjustments can make
+        the two diverge on one rank (e.g. the MixedMamba disagg WAR in
+        ``_schedule``). Whenever any rank can run forward, progress is reaped
+        non-blockingly; the bounded blocking poll is reserved for the globally
+        idle case.
+        """
         local_need_check = (num_fitting_reqs == 0
                             and not fitting_disagg_gen_init_requests)
 
@@ -3523,14 +3548,24 @@ def _check_disagg_transfer_progress_when_idle(
         local_need_gen_check = (local_need_check
                                 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)
+        local_state = (
+            _DISAGG_GEN_NEED_TRANSFER_PROGRESS if local_need_gen_check else
+            0) | (_DISAGG_GEN_HAS_FORWARD_WORK if has_forward_work else 0)
+        global_state = self._sync_disagg_gen_progress_state(local_state)
+        if global_state & _DISAGG_GEN_NEED_TRANSFER_PROGRESS:
+            if global_state & _DISAGG_GEN_HAS_FORWARD_WORK:
+                # Some rank has a scheduled batch. Blocking cannot add the
+                # awaited transfers to the already-fixed batch and would stall
+                # decode on every lockstepped rank; reap without blocking.
+                self._check_disagg_gen_cache_transfer_status(0)
+            else:
+                # Globally idle: keep the bounded blocking poll so the loop
+                # does not busy-spin while admission stays blocked.
+                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
 
         any_need_check = self._sync_disagg_ctx_status_entry(local_need_check)
@@ -3700,7 +3735,8 @@ def _prepare_and_schedule_batch(self):
                 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)
+                wait_for_disagg_gen_transfer_progress, all_gen_first,
+                scheduled_batch.batch_size > 0)
 
             # In gen-only benchmark mode, all requests must fit in KV cache
             # simultaneously. If some requests are stuck in INIT state and the

tests/unittest/_torch/executor/test_benchmark_disagg.py

@@ -1166,7 +1166,7 @@ def test_partial_transfer_admission_uses_only_admitted_requests(self):
         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
+            0, [admitted_req], False, False, False
         )
         ex._handle_errors.assert_not_called()
 

tests/unittest/_torch/executor/test_py_executor.py

@@ -26,7 +26,12 @@
     RequestQueueItem,
 )
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, SamplingConfig
-from tensorrt_llm._torch.pyexecutor.py_executor import DisaggTransferAdmissionController, PyExecutor
+from tensorrt_llm._torch.pyexecutor.py_executor import (
+    _DISAGG_GEN_HAS_FORWARD_WORK,
+    _DISAGG_GEN_NEED_TRANSFER_PROGRESS,
+    DisaggTransferAdmissionController,
+    PyExecutor,
+)
 from tensorrt_llm._torch.pyexecutor.resource_manager import NoFreeSlotsError, ResourceManagerType
 from tensorrt_llm._torch.pyexecutor.scheduler import (
     FCFSWaitingQueue,
@@ -636,7 +641,8 @@ 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_generation_transfer_when_admission_blocked_and_idle(self):
+        """Single rank, admission blocked, nothing to forward: bounded poll."""
         executor = object.__new__(PyExecutor)
         executor.dist = Mock(tp_size=1)
         executor._check_disagg_gen_cache_transfer_status = Mock()
@@ -648,15 +654,42 @@ def test_polls_generation_transfer_when_admission_blocked(self):
             fitting_disagg_gen_init_requests=[],
             wait_for_disagg_gen_transfer_progress=True,
             all_gen_first=False,
+            has_forward_work=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):
+    def test_same_rank_progress_need_with_forward_work_reaps_nonblocking(self):
+        """NEED and WORK can coexist on one rank when a post-scheduling
+        adjustment shrinks num_fitting_reqs below the final batch (e.g. the
+        MixedMamba disagg WAR in _schedule filters context requests but keeps
+        generation requests). The rank can forward, so it must not block."""
+        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,
+            has_forward_work=True,
+        )
+
+        executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(0)
+        executor._check_disagg_ctx_cache_transfer_status.assert_not_called()
+
+    def test_peer_backpressure_with_forward_work_reaps_nonblocking(self):
+        """A rank with a scheduled batch must not be dragged into the bounded
+        poll by a peer rank's admission backpressure."""
         executor = object.__new__(PyExecutor)
         executor.dist = Mock(tp_size=1, cp_size=4, world_size=4)
-        executor.dist.allreduce.return_value = 1
+        executor.dist.allreduce.return_value = (
+            _DISAGG_GEN_NEED_TRANSFER_PROGRESS | _DISAGG_GEN_HAS_FORWARD_WORK
+        )
         executor._check_disagg_gen_cache_transfer_status = Mock()
         executor._check_disagg_ctx_cache_transfer_status = Mock()
 
@@ -666,11 +699,35 @@ def test_peer_rank_enters_bounded_progress_poll(self):
             fitting_disagg_gen_init_requests=[],
             wait_for_disagg_gen_transfer_progress=True,
             all_gen_first=False,
+            has_forward_work=True,
+        )
+
+        executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(0)
+        executor._check_disagg_ctx_cache_transfer_status.assert_not_called()
+        executor.dist.allreduce.assert_called_once_with(
+            _DISAGG_GEN_HAS_FORWARD_WORK, op=ReduceOp.BOR
+        )
+
+    def test_peer_backpressure_with_all_ranks_idle_enters_bounded_poll(self):
+        """When no rank can run forward, the bounded blocking poll survives."""
+        executor = object.__new__(PyExecutor)
+        executor.dist = Mock(tp_size=1, cp_size=4, world_size=4)
+        executor.dist.allreduce.return_value = _DISAGG_GEN_NEED_TRANSFER_PROGRESS
+        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,
+            has_forward_work=False,
         )
 
         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.dist.allreduce.assert_called_once_with(0, op=ReduceOp.BOR)
 
     def test_falls_back_to_context_transfer_when_not_generation_blocked(self):
         executor = object.__new__(PyExecutor)
@@ -684,6 +741,7 @@ def test_falls_back_to_context_transfer_when_not_generation_blocked(self):
             fitting_disagg_gen_init_requests=[],
             wait_for_disagg_gen_transfer_progress=False,
             all_gen_first=False,
+            has_forward_work=False,
         )
 
         executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(1)
@@ -703,6 +761,7 @@ def test_sync_benchmark_skips_idle_transfer_collectives(self, monkeypatch):
             fitting_disagg_gen_init_requests=[],
             wait_for_disagg_gen_transfer_progress=True,
             all_gen_first=False,
+            has_forward_work=False,
         )
 
         executor.dist.allreduce.assert_not_called()
@@ -724,6 +783,7 @@ def test_sync_non_benchmark_skips_idle_transfer_collectives(self, monkeypatch):
             fitting_disagg_gen_init_requests=[],
             wait_for_disagg_gen_transfer_progress=True,
             all_gen_first=False,
+            has_forward_work=False,
         )
 
         executor.dist.allreduce.assert_not_called()
@@ -816,7 +876,7 @@ def complete_or_error(req):
     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.allreduce.return_value = _DISAGG_GEN_HAS_FORWARD_WORK
         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()
@@ -827,6 +887,7 @@ def test_peer_cp_rank_enters_context_progress_poll(self):
             fitting_disagg_gen_init_requests=[],
             wait_for_disagg_gen_transfer_progress=False,
             all_gen_first=False,
+            has_forward_work=True,
         )
 
         executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0)