diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index cef6ef31e91f..897599b0e2f2 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3325,6 +3325,27 @@ def _dist_size(dist, name: str) -> int: except (AttributeError, TypeError, ValueError): return 1 + def _allgather_model_parallel_status( + self, local_status: Tuple[int, bool]) -> List[Tuple[int, bool]]: + """Gather a status over the TP+CP scheduling group. + + Args: + local_status: Caller-defined ``(state, flag)`` pair from this rank. + The fill gate uses ``(ready, synchronous_progress)`` and the + fail-fast path uses ``(all_fetched, terminal_no_fit)``. + + Returns: + One status pair per TP+CP rank in the current pipeline-parallel + slice. A singleton group returns ``[local_status]``. + """ + # CP may coexist with TP; tp_cp_allgather covers both CP-only and + # TP+CP configurations. + if self._dist_size(self.dist, "cp_size") > 1: + return self.dist.tp_cp_allgather(local_status) + if self._dist_size(self.dist, "tp_size") > 1: + 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) @@ -3383,6 +3404,49 @@ def _check_disagg_transfer_progress_when_idle( # blocking on un-finished ones. 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], + wait_for_disagg_gen_transfer_progress: bool) -> bool: + """Return whether benchmark fill has terminal KV exhaustion. + + Model-parallel ranks can make different local scheduling decisions. + Every rank must therefore vote before entering the collective error- + handling path. One terminal rank prevents the global benchmark fill + gate from opening. The vote is fill-only to avoid adding a collective + to every decode iteration after the gate opens. + + Args: + scheduler_fitting_disagg_gen_init_requests: Generation INIT + requests that fit KV capacity before transfer admission. A + nonempty list means KV capacity exists even if transfer + admission temporarily defers every request. + wait_for_disagg_gen_transfer_progress: Whether active generation + transfers are consuming the admission budget and transfer + progress can unblock a deferred request. + + Returns: + True when every TP+CP rank has fetched its full benchmark queue and + at least one rank has an INIT request that cannot fit KV capacity + and has no transfer progress that can unblock it; otherwise False. + """ + if (self.benchmark_req_queues_size <= 0 or self.is_warmup + or not self._benchmark_fill_phase_active): + return False + + local_has_stuck = any(req.is_disagg_generation_init_state + for req in self.active_requests) + local_all_fetched = (self.num_fetch_requests + >= self.benchmark_req_queues_size) + local_terminal_no_fit = (local_has_stuck and + not scheduler_fitting_disagg_gen_init_requests + and not wait_for_disagg_gen_transfer_progress) + local_status = (local_all_fetched, local_terminal_no_fit) + + all_rank_status = self._allgather_model_parallel_status(local_status) + all_ranks_fetched = all(status[0] for status in all_rank_status) + any_rank_terminal_no_fit = any(status[1] for status in all_rank_status) + return all_ranks_fetched and any_rank_terminal_no_fit + def _prepare_and_schedule_batch(self): self._sync_disagg_transfer_made_progress = False new_requests = self._fetch_and_activate_new_requests() @@ -3448,7 +3512,7 @@ def _prepare_and_schedule_batch(self): # that speculation is about to happen. self._prepare_draft_requests() - scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule( + scheduled_batch, scheduler_fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule( ) if self.drafter is not None and not self.use_spec_decode: @@ -3457,18 +3521,19 @@ def _prepare_and_schedule_batch(self): if self.kv_cache_transceiver: wait_for_disagg_gen_transfer_progress = False - fitting_disagg_gen_init_requests, wait_for_disagg_gen_transfer_progress = ( + admitted_disagg_gen_init_requests, wait_for_disagg_gen_transfer_progress = ( self._apply_disagg_transfer_admission( - fitting_disagg_gen_init_requests)) - # For requests that are fitting disagg gen init, also prepare resources for KV cache manager - self._prepare_disagg_gen_init(fitting_disagg_gen_init_requests) + scheduler_fitting_disagg_gen_init_requests)) + # Prepare KV cache manager resources only for requests admitted + # 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, fitting_disagg_gen_init_requests, + num_fitting_reqs, admitted_disagg_gen_init_requests, wait_for_disagg_gen_transfer_progress, all_gen_first) # In gen-only benchmark mode, all requests must fit in KV cache @@ -3476,30 +3541,25 @@ def _prepare_and_schedule_batch(self): # scheduler could not allocate KV for any of them, the benchmark # will hang forever because in-progress generation requests won't # release their KV cache. - if (self.benchmark_req_queues_size > 0 and not self.is_warmup - and not fitting_disagg_gen_init_requests): - stuck_init_requests = [ - req for req in self.active_requests - if req.is_disagg_generation_init_state - ] - # Only fail once all benchmark requests have been fetched - # so that _handle_errors covers every request and every - # client receives an error response. - if (stuck_init_requests and self.num_fetch_requests - >= self.benchmark_req_queues_size): - error_msg = ( - f"Insufficient KV cache for gen-only benchmark mode: " - f"{len(stuck_init_requests)} request(s) are waiting for " - f"KV cache allocation but the scheduler could not fit " - f"any of them. Increase free_gpu_memory_fraction or " - f"reduce TLLM_BENCHMARK_REQ_QUEUES_SIZE (currently " - f"{self.benchmark_req_queues_size}).") - logger.error(error_msg) - # Fail all active and waiting requests so every - # client receives an error instead of hanging. - self._handle_errors(error_msg, - requests=self.active_requests) - return None, None + # Check the scheduler result from before transfer admission. An + # empty admitted list can mean that active transfers are + # temporarily consuming the transfer budget. + has_insufficient_kv = self._sync_gen_only_benchmark_has_insufficient_kv( + scheduler_fitting_disagg_gen_init_requests, + wait_for_disagg_gen_transfer_progress) + if has_insufficient_kv: + error_msg = ( + f"Insufficient KV cache for gen-only benchmark mode: " + f"one or more requests are waiting for KV cache allocation " + f"on a model-parallel rank whose scheduler could not fit " + f"any of them. Increase free_gpu_memory_fraction or reduce " + f"TLLM_BENCHMARK_REQ_QUEUES_SIZE (currently " + f"{self.benchmark_req_queues_size}).") + logger.error(error_msg) + # Fail all active and waiting requests on every rank so every + # client receives an error instead of hanging. + self._handle_errors(error_msg, requests=self.active_requests) + return None, None self.num_scheduled_requests = scheduled_batch.batch_size logger.debug( @@ -3542,9 +3602,9 @@ def _is_benchmark_disagg_fill_complete( KV-transfer phase (not in INIT, TRANS_IN_PROGRESS, or ERROR). (C) The KV cache transceiver has no pending receive sessions. - For ADP, the conditions and synchronous-progress signal are gathered - together across TP ranks so every rank makes the same gate and sleep - decision. + The conditions and synchronous-progress signal are gathered across the + TP+CP scheduling group so every model-parallel rank makes the same gate + and sleep decision. This method must only be called when ``is_benchmark_disagg`` is True. @@ -3564,8 +3624,8 @@ def _is_benchmark_disagg_fill_complete( "outside benchmark disagg mode.") # (A) All benchmark requests have been fetched from the queue. Keep - # going to the shared allgather even when this rank is not done so TP - # ranks cannot diverge in collective order. + # going to the shared allgather even when this rank is not done so + # model-parallel ranks cannot diverge in collective order. local_all_fetched = (self.num_fetch_requests >= self.benchmark_req_queues_size) if not local_all_fetched: @@ -3592,10 +3652,7 @@ def _is_benchmark_disagg_fill_complete( and local_no_inflight) local_status = (local_ok, bool(local_sync_progress)) - if self.enable_attention_dp: - all_rank_status = self.dist.tp_allgather(local_status) - else: - all_rank_status = [local_status] + all_rank_status = self._allgather_model_parallel_status(local_status) all_ranks_ok = [status[0] for status in all_rank_status] global_ok = min(all_ranks_ok) == 1 self._benchmark_sync_progress_global = any( diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py index b1f94f130bd6..dabb2b910d55 100644 --- a/tests/unittest/_torch/executor/test_benchmark_disagg.py +++ b/tests/unittest/_torch/executor/test_benchmark_disagg.py @@ -22,6 +22,7 @@ - ADP dummy suppression during fill vs taper-down - ADP router imbalance regression (nvbug 6071070) - Non-blocking behaviour of `_prepare_and_schedule_batch` +- Insufficient-KV fail-fast vs transfer-admission backpressure """ from unittest.mock import Mock, patch @@ -100,6 +101,8 @@ def __init__( from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + _dist_size = staticmethod(PyExecutor._dist_size) + _allgather_model_parallel_status = PyExecutor._allgather_model_parallel_status _is_benchmark_disagg_fill_complete = PyExecutor._is_benchmark_disagg_fill_complete _check_benchmark_disagg_gate = PyExecutor._check_benchmark_disagg_gate @@ -227,7 +230,8 @@ def test_allgather_sends_local_ok_int(self): ex._is_benchmark_disagg_fill_complete(ScheduledRequests()) ex.dist.tp_allgather.assert_called_once_with((1, False)) - def test_no_allgather_without_adp(self): + def test_single_rank_skips_model_parallel_allgather(self): + """A singleton group returns local status without a collective.""" reqs = [_make_active_request() for _ in range(4)] ex = MockBenchmarkExecutor( benchmark_req_queues_size=4, @@ -236,8 +240,12 @@ def test_no_allgather_without_adp(self): num_fetch_requests=4, active_requests=reqs, ) + ex.dist.tp_size = 1 + ex.dist.cp_size = 1 + ex.dist.world_size = 1 ex._is_benchmark_disagg_fill_complete(ScheduledRequests()) ex.dist.tp_allgather.assert_not_called() + ex.dist.tp_cp_allgather.assert_not_called() class TestFillCompleteADPRouterImbalance: @@ -391,6 +399,52 @@ def test_gate_skips_sleep_on_all_adp_ranks_when_peer_makes_progress(self, mock_t ex.dist.tp_allgather.assert_called_once_with((0, False)) mock_time.sleep.assert_not_called() + @pytest.mark.parametrize( + "enable_attention_dp, tp_size, cp_size, gather_name", + [ + pytest.param(False, 2, 1, "tp_allgather", id="tensor_parallel"), + pytest.param(False, 1, 2, "tp_cp_allgather", id="context_parallel"), + pytest.param(True, 2, 2, "tp_cp_allgather", id="attention_dp_with_cp"), + ], + ) + def test_gate_waits_for_blocked_model_parallel_peer( + self, enable_attention_dp, tp_size, cp_size, gather_name + ): + """A ready local slice cannot open the gate ahead of a peer. + + Args: + enable_attention_dp: Whether to simulate attention data parallelism. + tp_size: Tensor-parallel group size. + cp_size: Context-parallel group size. + gather_name: Expected model-parallel allgather method. + """ + reqs = [_make_active_request() for _ in range(4)] + ex = MockBenchmarkExecutor( + benchmark_req_queues_size=4, + kv_cache_transceiver=_make_transceiver(transfer_complete=True), + enable_attention_dp=enable_attention_dp, + tp_size=tp_size, + num_fetch_requests=4, + active_requests=reqs, + ) + ex.dist.cp_size = cp_size + ex.dist.world_size = tp_size * cp_size + all_rank_status = [(1, False)] * (tp_size * cp_size) + all_rank_status[-1] = (0, False) + gather = getattr(ex.dist, gather_name) + gather.return_value = all_rank_status + + with patch("tensorrt_llm._torch.pyexecutor.py_executor.time") as mock_time: + can_forward, should_retry = ex._check_benchmark_disagg_gate(ScheduledRequests(), False) + + assert can_forward is False + assert should_retry is True + assert ex._benchmark_fill_phase_active is True + gather.assert_called_once_with((1, False)) + if gather_name == "tp_cp_allgather": + ex.dist.tp_allgather.assert_not_called() + mock_time.sleep.assert_called_once_with(0.1) + @pytest.mark.parametrize( "is_warmup, can_forward_in", [ @@ -957,7 +1011,7 @@ def _make_executor( ex._fill_admit_cap = 0 ex.enable_attention_dp = False ex.num_fetch_requests = num_fetch_requests - ex.dist = Mock(rank=0, tp_size=1) + ex.dist = Mock(rank=0, tp_size=1, cp_size=1, world_size=1) ex.dist.allreduce.return_value = 0 ex.is_shutdown = False ex._is_warmup = False @@ -976,6 +1030,7 @@ def _make_executor( ex._fetch_and_activate_new_requests = Mock(return_value=[]) ex._check_disagg_ctx_schedulable_status = Mock() ex._check_disagg_gen_transfer_status = Mock() + ex._check_disagg_gen_cache_transfer_status = Mock() ex._check_kv_transfer_timeout = Mock() ex._check_disagg_ctx_cache_transfer_status = Mock() ex._pad_attention_dp_dummy_request = Mock() @@ -1014,6 +1069,137 @@ def test_healthy_fill_phase_does_not_kill(self): ) ex._handle_errors.assert_not_called() + def test_partial_transfer_admission_uses_only_admitted_requests(self): + """The admitted subset is prepared and passed to the idle check.""" + admitted_req = _make_active_request(in_init=True) + deferred_req = _make_active_request(in_init=True) + candidates = [admitted_req, deferred_req] + ex = self._make_executor(fill_phase_active=True, fitting_init_requests=candidates) + ex._apply_disagg_transfer_admission = Mock(return_value=([admitted_req], False)) + ex._check_disagg_transfer_progress_when_idle = Mock() + + result, _ = ex._prepare_and_schedule_batch() + + 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._handle_errors.assert_not_called() + + def test_fill_with_no_init_requests_does_not_kill(self): + """The final fill iteration is ready for the gate, not terminal.""" + ex = self._make_executor(fill_phase_active=True, num_init_requests=0) + + result, _ = ex._prepare_and_schedule_batch() + + assert result is not None + ex._handle_errors.assert_not_called() + + def test_transfer_admission_backpressure_does_not_kill(self, monkeypatch): + """NVBug 6438658: admission backpressure is not KV exhaustion. + + Args: + monkeypatch: Pytest fixture used to select asynchronous transfer + behavior. + """ + monkeypatch.delenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY", raising=False) + monkeypatch.delenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", raising=False) + fitting_req = _make_active_request(in_init=True) + ex = self._make_executor(fill_phase_active=True, fitting_init_requests=[fitting_req]) + ex._apply_disagg_transfer_admission = Mock(return_value=([], True)) + + result, _ = ex._prepare_and_schedule_batch() + + assert result is not None, ( + "Fail-fast should NOT fire when the scheduler fit an INIT request " + "that transfer admission temporarily deferred" + ) + 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._handle_errors.assert_not_called() + + @pytest.mark.parametrize( + "enable_attention_dp, tp_size, cp_size, gather_name", + [ + pytest.param(False, 2, 1, "tp_allgather", id="tensor_parallel"), + pytest.param(False, 1, 2, "tp_cp_allgather", id="context_parallel"), + pytest.param(True, 2, 2, "tp_cp_allgather", id="attention_dp_with_cp"), + ], + ) + def test_model_parallel_peer_terminal_no_fit_kills_all_ranks( + self, enable_attention_dp, tp_size, cp_size, gather_name + ): + """A terminal peer makes every model-parallel rank fail together. + + Args: + enable_attention_dp: Whether to simulate attention data parallelism. + tp_size: Tensor-parallel group size. + cp_size: Context-parallel group size. + gather_name: Expected model-parallel allgather method. + """ + fitting_req = _make_active_request(in_init=True) + ex = self._make_executor(fill_phase_active=True, fitting_init_requests=[fitting_req]) + ex.enable_attention_dp = enable_attention_dp + ex.dist.tp_size = tp_size + ex.dist.cp_size = cp_size + ex.dist.world_size = tp_size * cp_size + all_rank_status = [(True, False)] * (tp_size * cp_size) + all_rank_status[-1] = (True, True) + gather = getattr(ex.dist, gather_name) + gather.return_value = all_rank_status + ex._apply_disagg_transfer_admission = Mock(return_value=([], True)) + ex._check_disagg_transfer_progress_when_idle = Mock() + + result, _ = ex._prepare_and_schedule_batch() + + assert result is None + gather.assert_called_once_with((True, False)) + if gather_name == "tp_cp_allgather": + ex.dist.tp_allgather.assert_not_called() + ex._handle_errors.assert_called_once() + assert "one or more requests" in ex._handle_errors.call_args.args[0] + + def test_attention_dp_backpressure_without_terminal_peer_does_not_kill(self): + """Admission backpressure stays non-terminal on every rank.""" + fitting_req = _make_active_request(in_init=True) + ex = self._make_executor(fill_phase_active=True, fitting_init_requests=[fitting_req]) + ex.enable_attention_dp = True + ex.dist.tp_size = 2 + ex.dist.world_size = 2 + ex.dist.tp_allgather.return_value = [ + (True, False), + (True, False), + ] + ex._apply_disagg_transfer_admission = Mock(return_value=([], True)) + ex._check_disagg_transfer_progress_when_idle = Mock() + + result, _ = ex._prepare_and_schedule_batch() + + assert result is not None + ex.dist.tp_allgather.assert_called_once_with((True, False)) + ex._handle_errors.assert_not_called() + + def test_model_parallel_waits_until_all_ranks_have_fetched(self): + """A terminal rank cannot fail peers that are still fetching.""" + ex = self._make_executor(fill_phase_active=True) + ex.dist.tp_size = 2 + ex.dist.world_size = 2 + ex.dist.tp_allgather.return_value = [ + (True, True), + (False, False), + ] + ex._check_disagg_transfer_progress_when_idle = Mock() + + result, _ = ex._prepare_and_schedule_batch() + + assert result is not None + ex.dist.tp_allgather.assert_called_once_with((True, True)) + ex._handle_errors.assert_not_called() + def test_mid_fetch_does_not_kill(self): """Before all benchmark requests are fetched, keep filling.""" ex = self._make_executor(fill_phase_active=True, num_fetch_requests=4) @@ -1025,23 +1211,25 @@ def test_mid_fetch_does_not_kill(self): ) ex._handle_errors.assert_not_called() - def test_kills_after_fill_phase(self): - """After fill phase completes, stuck INIT requests trigger fail-fast.""" + def test_post_fill_skips_fail_fast_vote(self): + """Decode iterations must not pay for the fill-only collective.""" ex = self._make_executor(fill_phase_active=False) + ex.enable_attention_dp = True + ex.dist.tp_size = 2 + ex.dist.world_size = 2 + ex._check_disagg_transfer_progress_when_idle = Mock() result, _ = ex._prepare_and_schedule_batch() - assert result is None, ( - "Fail-fast SHOULD fire after fill phase — " - "stuck INIT requests indicate genuine KV insufficiency" - ) - ex._handle_errors.assert_called_once() + assert result is not None + ex.dist.tp_allgather.assert_not_called() + ex._handle_errors.assert_not_called() @pytest.mark.parametrize( "fill_active, is_warmup, expected_alive", [ pytest.param(True, False, False, id="stalled_fill_kills"), - pytest.param(False, False, False, id="post_fill_kills"), + pytest.param(False, False, True, id="post_fill_suppresses"), pytest.param(False, True, True, id="warmup_suppresses"), pytest.param(True, True, True, id="both_suppress"), ], @@ -1080,7 +1268,7 @@ class TestFillPhaseEndToEnd: 3. Scheduler can't fit INIT requests 4. Verify: fail-fast does NOT fire while scheduler fits INIT requests 5. Transfers complete, gate opens, fill phase clears - 6. Verify: fail-fast DOES fire if INIT requests remain after fill + 6. Verify: the fill-only fail-fast vote stops after the gate opens This test catches all three bugs we found iteratively: - Bug 1: Count-based gate unsatisfiable under ADP router skew @@ -1172,6 +1360,7 @@ def test_full_lifecycle(self): # Phase 2b: Healthy fill keeps making progress, so fail-fast must not # fire even though some active requests remain in INIT. ex._schedule = Mock(return_value=(ScheduledRequests(), [init_reqs[0]], 0)) + ex.dist.tp_allgather = Mock(return_value=[(True, False), (True, False)]) result, _ = ex._prepare_and_schedule_batch() assert result is not None, ( "Fail-fast must not kill requests while the scheduler can still fit INIT requests" @@ -1190,16 +1379,13 @@ def test_full_lifecycle(self): # Phase 4: Gate opens, fill phase clears ex._benchmark_fill_phase_active = False - # Phase 5: After fill, if new INIT requests appear and the scheduler - # cannot fit any of them, fail-fast fires. Reset _schedule from - # Phase 2b's healthy-fill mock so it once again returns no fitting - # INIT requests, mirroring genuine insufficient-KV conditions. + # Phase 5: Decode iterations do not re-enter the fill-only vote. ex._schedule = Mock(return_value=(ScheduledRequests(), [], 0)) - stuck_req = _make_active_request(in_init=True) - ex.active_requests = [stuck_req] + ready_reqs + ex.active_requests = ready_reqs + ex.dist.tp_allgather = Mock() + ex._check_disagg_transfer_progress_when_idle = Mock() result, _ = ex._prepare_and_schedule_batch() - assert result is None, ( - "Fail-fast SHOULD fire after fill phase completes — " - "stuck INIT requests now indicate genuine KV insufficiency" - ) + assert result is not None + ex.dist.tp_allgather.assert_not_called() + ex._handle_errors.assert_not_called()