[https://nvbugs/6438586][fix] Make gen-only benchmark insufficient-KV fail-fast ADP-safe

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")
 

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: