From 7d121448b21ca05a2700e8e74f3fced179061af0 Mon Sep 17 00:00:00 2001 From: popsiclexu Date: Fri, 31 Oct 2025 05:22:47 +0800 Subject: [PATCH] [Bug fix][PD Dissaggregation] fix prefill hanging issue with PP and DP Attention, (#12368) --- python/sglang/srt/disaggregation/prefill.py | 9 ++++++--- python/sglang/srt/managers/scheduler_pp_mixin.py | 9 +++++++-- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 8fad0c0dc..5b2e5cb41 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -588,7 +588,7 @@ class SchedulerDisaggregationPrefillMixin: """ polls = poll_and_all_reduce( [req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue], - self.tp_worker.get_tp_group().cpu_group, + self.tp_worker.get_attention_tp_cpu_group(), ) transferred_rids: List[str] = [] @@ -722,8 +722,11 @@ class SchedulerDisaggregationPrefillMixin: else: data = None - if self.tp_size != 1: + if self.attn_tp_size != 1: data = broadcast_pyobj( - data, self.tp_group.rank, self.tp_cpu_group, src=self.tp_group.ranks[0] + data, + self.attn_tp_group.rank, + self.attn_tp_cpu_group, + src=self.attn_tp_group.ranks[0], ) return data diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index e177d3b56..c49ae937f 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -4,7 +4,7 @@ from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.model_executor.forward_batch_info import PPProxyTensors -from sglang.srt.utils import DynamicGradMode, point_to_point_pyobj +from sglang.srt.utils import DynamicGradMode, point_to_point_pyobj, require_mlp_sync class SchedulerPPMixin: @@ -236,7 +236,12 @@ class SchedulerPPMixin: tmbs[mb_id] = transferred_rids self.process_prefill_chunk() - mbs[mb_id] = self.get_new_batch_prefill() + + batch = self.get_new_batch_prefill() + if require_mlp_sync(self.server_args): + batch = self.prepare_mlp_sync_batch(batch) + mbs[mb_id] = batch + self.running_mbs[mb_id] = self.running_batch self.cur_batch = mbs[mb_id]