[None][fix] Fix disagg gen-consensus deadlock under pipeline parallelism

tensorrt_llm/_torch/disaggregation/transceiver.py

@@ -314,11 +314,13 @@ def _ctx_consensus(self, local_ids: list) -> list:
         return ready_ids
 
     def _gen_consensus(self, local_ids: list) -> list:
-        sync_size = (
-            self._mapping.pp_size if self._mapping.enable_attention_dp else self._mapping.world_size
-        )
+        # adp-off: use the TP-then-PP two-stage consensus (same ordered TP/PP
+        # communicators as ctx; tp x pp == world) instead of a lone WORLD
+        # allgather that deadlocks against the PP subgroup collectives.
+        if not self._mapping.enable_attention_dp:
+            return self._ctx_consensus(local_ids)
         all_ranks = self._gen_allgather(local_ids) if self._gen_need_sync else [local_ids]
-        return _find_consensus_request_ids(all_ranks, sync_size)
+        return _find_consensus_request_ids(all_ranks, self._mapping.pp_size)
 
     @staticmethod
     def _allgather_or_passthrough(
@@ -366,6 +368,23 @@ def _consensus_outcome(
         return new_cancelled, new_failed, new_completed
 
     def _gen_consensus_outcome(self, to_process, cancelled, failed, completed):
+        # adp-off: TP-then-PP two-stage on the ordered subgroup communicators
+        # (not a WORLD allgather) to avoid the cross-communicator PP deadlock;
+        # tp x pp == world, so the outcome is identical.
+        if not self._mapping.enable_attention_dp:
+            c, f, d = self._consensus_outcome(
+                to_process,
+                cancelled,
+                failed,
+                completed,
+                self._dist.tp_allgather,
+                self._ctx_need_tp_sync,
+            )
+            if self._ctx_need_pp_sync:
+                c, f, d = self._consensus_outcome(
+                    to_process, c, f, d, self._dist.pp_allgather, True
+                )
+            return c, f, d
         return self._consensus_outcome(
             to_process, cancelled, failed, completed, self._gen_allgather, self._gen_need_sync
         )

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -1937,9 +1937,14 @@ def _pp_schedule_and_propagate(self, microbatch_id: int):
                                    and is_dp_broadcast):
             scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule(
             )
+            disagg_gen_transfer_in_progress = (
+                self.kv_cache_transceiver is not None
+                and (any(req.is_disagg_generation_transmission_in_progress
+                         for req in self.active_requests)
+                     or bool(fitting_disagg_gen_init_requests)))
             serializable_schedule = SerializableSchedulerOutput.from_scheduler_result(
                 scheduled_batch, fitting_disagg_gen_init_requests,
-                num_fitting_reqs)
+                num_fitting_reqs, disagg_gen_transfer_in_progress)
 
         # Broadcast within first tp+cp group before send/recv chain to other tp+cp groups
         if self.dist.is_first_pp_rank:
@@ -1971,6 +1976,11 @@ def _pp_schedule_and_propagate(self, microbatch_id: int):
         if scheduled_batch is None:
             scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs = serializable_schedule.to_scheduler_result(
                 self.active_requests)
+        # Adopt the scheduling rank's disagg-gen-transfer flag on every PP rank
+        # (rode this broadcast, no extra collective) so the next iteration's
+        # gen transfer-status allgather is entered/skipped by all ranks together.
+        self._pp_disagg_gen_transfer_in_progress = (
+            serializable_schedule.disagg_gen_transfer_in_progress)
         return scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs
 
     def _pp_retry_until_can_schedule(self, scheduled_batch):
@@ -4295,6 +4305,13 @@ def _check_disagg_gen_transfer_status(self):
             req.is_disagg_generation_transmission_in_progress
             for req in non_gen_first_reqs)
 
+        # PP: gate the transceiver's group allgather on the scheduling rank's
+        # propagated flag (set only on the PP loop) so all ranks enter/skip
+        # together; a divergent per-rank need_check deadlocks that allgather.
+        pp_flag = getattr(self, "_pp_disagg_gen_transfer_in_progress", None)
+        if pp_flag is not None:
+            need_check = pp_flag
+
         if need_check:
             at_least_num = 1 if need_check_one else 0
             self._check_disagg_gen_cache_transfer_status(at_least_num)

tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py

@@ -236,13 +236,18 @@ class SerializableSchedulerOutput:
         int
     ]  # request ids of fitting disaggregated generation initialization requests
     num_fitting_requests: int  # number of fitting requests
+    # Scheduling rank's "any disagg-gen KV transfer in progress" flag; rides
+    # this existing schedule broadcast (no extra collective) so all PP ranks
+    # enter/skip the gen transfer-status allgather together (else it deadlocks).
+    disagg_gen_transfer_in_progress: bool = False
 
     @classmethod
     def from_scheduler_result(
         cls,
         scheduled_requests: ScheduledRequests,
         fitting_disagg_gen_init_requests: RequestList,
         num_fitting_requests: int,
+        disagg_gen_transfer_in_progress: bool = False,
     ) -> "SerializableSchedulerOutput":
         return cls(
             encoder_requests=[req.request_id for req in scheduled_requests.encoder_requests],
@@ -258,6 +263,7 @@ def from_scheduler_result(
                 req.request_id for req in fitting_disagg_gen_init_requests
             ],
             num_fitting_requests=num_fitting_requests,
+            disagg_gen_transfer_in_progress=disagg_gen_transfer_in_progress,
         )
 
     def to_scheduler_result(