From b5d46d1f0230358125610e26e606c5a1d100421b Mon Sep 17 00:00:00 2001 From: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> Date: Wed, 5 Aug 2026 15:12:30 +0800 Subject: [PATCH 1/2] [https://nvbugs/6527301][fix] Do not block ranks with scheduled work on disagg gen transfer progress Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 66 +++++++++++++---- .../_torch/executor/test_benchmark_disagg.py | 2 +- .../_torch/executor/test_py_executor.py | 73 +++++++++++++++++-- 3 files changed, 119 insertions(+), 22 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 6e3252b48688..4bf40c7e7e87 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/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 diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py index d7992e1b445c..146ed3e2eaa0 100644 --- a/tests/unittest/_torch/executor/test_benchmark_disagg.py +++ b/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() diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index c9c05b39137f..185c27decb46 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/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) From 4452576feb8326b54a33d484ac646b91bf0f22de Mon Sep 17 00:00:00 2001 From: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> Date: Wed, 5 Aug 2026 16:20:54 +0800 Subject: [PATCH 2/2] [https://nvbugs/6527301][chore] Prefix internal disagg gen progress bit constants with underscore Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 14 +++++++------- tests/unittest/_torch/executor/test_py_executor.py | 12 ++++++------ 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 4bf40c7e7e87..a4cb3c73d7ee 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -157,8 +157,8 @@ class _SleepWakeupAction(StrEnum): # 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 +_DISAGG_GEN_NEED_TRANSFER_PROGRESS = 1 +_DISAGG_GEN_HAS_FORWARD_WORK = 2 def _sleep_wakeup_ack_ready(comm, source: int, tag: _SleepWakeupTag) -> bool: @@ -3548,12 +3548,12 @@ def _check_disagg_transfer_progress_when_idle( local_need_gen_check = (local_need_check and wait_for_disagg_gen_transfer_progress) - 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) + 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: + 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. diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 185c27decb46..07ca23acbe6f 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -27,8 +27,8 @@ ) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, SamplingConfig from tensorrt_llm._torch.pyexecutor.py_executor import ( - DISAGG_GEN_HAS_FORWARD_WORK, - DISAGG_GEN_NEED_TRANSFER_PROGRESS, + _DISAGG_GEN_HAS_FORWARD_WORK, + _DISAGG_GEN_NEED_TRANSFER_PROGRESS, DisaggTransferAdmissionController, PyExecutor, ) @@ -688,7 +688,7 @@ def test_peer_backpressure_with_forward_work_reaps_nonblocking(self): 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 | DISAGG_GEN_HAS_FORWARD_WORK + _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() @@ -705,14 +705,14 @@ def test_peer_backpressure_with_forward_work_reaps_nonblocking(self): 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 + _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.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() @@ -876,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 = DISAGG_GEN_HAS_FORWARD_WORK + 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()