diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index cef6ef31e91f..e21128c6783f 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3476,17 +3476,34 @@ 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): + if self.benchmark_req_queues_size > 0 and not self.is_warmup: 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): + # client receives an error response. Requests deferred by + # the transfer admission controller + # (wait_for_disagg_gen_transfer_progress) are waiting on + # in-flight transfers, not on KV capacity, so they are + # not stuck. + local_stuck = bool(stuck_init_requests + and not fitting_disagg_gen_init_requests + and not wait_for_disagg_gen_transfer_progress + and self.num_fetch_requests + >= self.benchmark_req_queues_size) + # All DP ranks must agree before failing: entering + # _handle_errors on a subset of ranks desyncs its response + # gather against the fill-gate allgather running on the + # other ranks (https://nvbugs/6438586). The consensus + # allgather runs on every iteration regardless of local + # state so all ranks stay collective-aligned. + if self.enable_attention_dp and self.dist.world_size != 1: + should_fail = any(self.dist.tp_allgather(local_stuck)) + else: + should_fail = local_stuck + if should_fail: error_msg = ( f"Insufficient KV cache for gen-only benchmark mode: " f"{len(stuck_init_requests)} request(s) are waiting for " @@ -5992,8 +6009,14 @@ def _handle_errors(self, client_id=getattr(item.request, 'client_id', None)))) - if waiting_responses: + # Under ADP, waiting_responses is only populated on rank 0 + # (or with gather_all_responses), but _enqueue_responses runs + # a collective gather — every rank must enter it, even with an + # empty list, to stay in lockstep. + if waiting_responses or (self.enable_attention_dp + and self.dist.world_size != 1): self._enqueue_responses(waiting_responses) + if waiting_responses: logger.info(f"Drained {len(waiting_responses)} queued requests " "on fatal error") diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py index b1f94f130bd6..45f436b3f1fb 100644 --- a/tests/unittest/_torch/executor/test_benchmark_disagg.py +++ b/tests/unittest/_torch/executor/test_benchmark_disagg.py @@ -1065,6 +1065,80 @@ def test_suppression_matrix(self, fill_active, is_warmup, expected_alive): ) ex._handle_errors.assert_called_once() + def test_admission_deferral_does_not_kill(self): + """Requests deferred by the transfer admission controller are not stuck. + + When the scheduler fits INIT requests but the admission controller + defers all of them behind in-flight transfers + (wait_for_disagg_gen_transfer_progress=True), the fill is making + progress and the fail-fast must not fire (nvbug 6438586). + """ + ex = self._make_executor(fill_phase_active=True) + 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 while the admission controller is " + "deferring INIT requests behind active KV transfers" + ) + ex._handle_errors.assert_not_called() + + def _make_adp_executor(self, **kwargs): + """ADP variant of the executor stub (2 DP ranks).""" + ex = self._make_executor(**kwargs) + ex.enable_attention_dp = True + ex.dist = Mock(rank=0, tp_size=2, world_size=2) + ex.dist.allreduce.return_value = 0 + ex.dist.tp_allreduce.return_value = 0 + return ex + + def test_adp_consensus_runs_allgather_even_when_healthy(self): + """Under ADP the stuck-consensus allgather must run every iteration. + + A rank that skips the collective while a peer enters it desyncs all + subsequent collectives (nvbug 6438586: peers crashed with + `TypeError: '<' not supported between instances of 'list' and 'int'` + in the fill-gate allgather). + """ + fitting_req = _make_active_request(in_init=True) + ex = self._make_adp_executor(fill_phase_active=True, fitting_init_requests=[fitting_req]) + ex.dist.tp_allgather.return_value = [False, False] + + result, _ = ex._prepare_and_schedule_batch() + + assert result is not None + ex.dist.tp_allgather.assert_called_once_with(False) + ex._handle_errors.assert_not_called() + + def test_adp_consensus_kills_all_ranks_when_peer_is_stuck(self): + """A healthy rank must fail together with a stuck peer rank.""" + fitting_req = _make_active_request(in_init=True) + ex = self._make_adp_executor(fill_phase_active=True, fitting_init_requests=[fitting_req]) + # This rank is healthy (local flag False) but a peer reports stuck. + ex.dist.tp_allgather.return_value = [False, True] + + result, _ = ex._prepare_and_schedule_batch() + + assert result is None, ( + "All DP ranks must enter _handle_errors together when any rank " + "is stuck; failing on a subset desyncs the response gather " + "against the fill-gate allgather on the healthy ranks" + ) + ex.dist.tp_allgather.assert_called_once_with(False) + ex._handle_errors.assert_called_once() + + def test_adp_consensus_local_stuck_reported(self): + """A locally stuck rank contributes True to the consensus.""" + ex = self._make_adp_executor(fill_phase_active=True) + ex.dist.tp_allgather.return_value = [True, False] + + result, _ = ex._prepare_and_schedule_batch() + + assert result is None + ex.dist.tp_allgather.assert_called_once_with(True) + ex._handle_errors.assert_called_once() + # --------------------------------------------------------------------------- # End-to-end fill phase reproducer @@ -1172,10 +1246,15 @@ 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)) + # The stuck-consensus allgather (nvbug 6438586) shares + # dist.tp_allgather with the fill gate; during a healthy fill no + # rank reports stuck. + ex.dist.tp_allgather = Mock(return_value=[False, 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" ) + ex.dist.tp_allgather.assert_called_once_with(False) # Phase 3: All transfers complete, gate opens for req in init_reqs: