diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 78869315038a..2354917daf3e 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/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 ) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 009043a97847..c7774a28bbae 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/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) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index ace36551f5c0..483cca2f183d 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -236,6 +236,10 @@ 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( @@ -243,6 +247,7 @@ def from_scheduler_result( 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(