diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 2fa5a5d0607b..ed3fb6520a7a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2568,8 +2568,8 @@ def _executor_loop_pp(self): # 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) + 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. self._pp_retry_until_can_schedule(scheduled_batch) @@ -2593,14 +2593,10 @@ 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) + 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 @@ -3514,63 +3510,27 @@ 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 _check_disagg_transfer_progress_when_idle(self) -> None: + """Reap completed context KV transfers so their blocks can be freed. - 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) -> None: - local_need_check = (num_fitting_reqs == 0 - and not fitting_disagg_gen_init_requests) + 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. + # still be receiving while another is idle, so entering the context + # progress collective here is unsafe. if not self._uses_async_disagg_gen_transfer(): return - 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) - return - - any_need_check = self._sync_disagg_ctx_status_entry(local_need_check) - if any_need_check > 0: - if local_need_check and not all_gen_first: - logger.warning( - "num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache" - ) - # 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) def _sync_gen_only_benchmark_has_insufficient_kv( self, scheduler_fitting_disagg_gen_init_requests: List[LlmRequest], @@ -3715,13 +3675,10 @@ 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) + 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 # 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 b8cb399b1203..b68d78e420c6 100644 --- a/tests/unittest/_torch/executor/test_benchmark_disagg.py +++ b/tests/unittest/_torch/executor/test_benchmark_disagg.py @@ -1168,9 +1168,7 @@ def test_partial_transfer_admission_uses_only_admitted_requests(self): 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 - ) + 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): @@ -1203,8 +1201,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 c9c05b39137f..be7c73a80910 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -20,7 +20,6 @@ import pytest -from tensorrt_llm._torch.distributed.communicator import ReduceOp from tensorrt_llm._torch.pyexecutor.executor_request_queue import ( SHUTDOWN_REQUEST_ID, RequestQueueItem, @@ -636,57 +635,51 @@ 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=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) + + 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_skips_idle_polls(self, monkeypatch): + monkeypatch.setenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY", "1") 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=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_ctx_cache_transfer_status.assert_not_called() executor._check_disagg_gen_cache_transfer_status.assert_not_called() def test_sync_benchmark_skips_idle_transfer_collectives(self, monkeypatch): @@ -697,13 +690,7 @@ def test_sync_benchmark_skips_idle_transfer_collectives(self, monkeypatch): 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.dist.allreduce.assert_not_called() executor.dist.tp_allreduce.assert_not_called() @@ -718,13 +705,7 @@ def test_sync_non_benchmark_skips_idle_transfer_collectives(self, monkeypatch): 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.dist.allreduce.assert_not_called() executor.dist.tp_allreduce.assert_not_called() @@ -813,26 +794,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: