diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 7838491f1e8a..e2a034f818c3 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 @@ -3092,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 @@ -3780,73 +3772,71 @@ 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) + def _check_disagg_transfer_progress_when_idle(self) -> None: + """Reap completed context KV transfers so their blocks can be freed. - uses_async_gen_transfer = self._uses_async_disagg_gen_transfer() + The poll is non-blocking and rank-symmetric: every rank enters it + unconditionally on every disagg iteration, so the consensus performed + 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. - if (not uses_async_gen_transfer + # 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 - local_need_gen_check = (uses_async_gen_transfer and local_needs_progress - and wait_for_disagg_gen_transfer_progress) + self._check_disagg_ctx_cache_transfer_status(0) - 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) + 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 - 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) + # 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 _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], @@ -3984,7 +3974,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 +3994,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 @@ -4566,6 +4547,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): @@ -5441,6 +5425,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_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py index 5aaadbb2e35a..217b354223dc 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,7 @@ 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._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..5254945b446d 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,40 @@ 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_context_transfers_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, - ) + 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._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0) - def test_peer_rank_enters_bounded_progress_poll(self): + 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, cp_size=4, world_size=4) - executor.dist.allreduce.return_value = 1 + 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=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_gen_cache_transfer_status.assert_not_called() - def test_falls_back_to_context_transfer_when_not_generation_blocked(self): + def test_idle_poll_enters_no_extra_collective(self): + """The context poll is 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) def test_gen_only_no_context_benchmark_polls_context_when_idle( self, monkeypatch: pytest.MonkeyPatch @@ -1028,24 +1003,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 +1021,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,25 +1123,69 @@ 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, - ) +class TestIdleDisaggLoopPacing: + """The idle poll no longer blocks, so the executor loops pace themselves. - 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) + 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.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")