From 09e2571e2e427d3fa1dd4fbb676a285f9e6a6d03 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sun, 11 Jan 2026 13:04:43 +0800 Subject: [PATCH] Clarify the meaning of `cpu_group` / `entry_rank` when dp + tp is enabled. (#16876) --- python/sglang/srt/disaggregation/prefill.py | 2 +- python/sglang/srt/managers/scheduler.py | 49 +++++++++++-------- .../sglang/srt/managers/scheduler_pp_mixin.py | 2 +- .../srt/managers/scheduler_profiler_mixin.py | 8 +-- python/sglang/srt/managers/tp_worker.py | 9 ---- 5 files changed, 34 insertions(+), 36 deletions(-) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index cbd18af03..39a824c3a 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -603,7 +603,7 @@ class SchedulerDisaggregationPrefillMixin: """ polls = poll_and_all_reduce( [req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue], - self.tp_worker.get_attention_tp_cpu_group(), + self.attn_tp_cpu_group, ) transferred_rids: List[str] = [] diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index f06a9a8e8..a3765d261 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -60,10 +60,14 @@ from sglang.srt.disaggregation.utils import ( prepare_abort, ) from sglang.srt.distributed import get_pp_group, get_world_group +from sglang.srt.distributed.parallel_state import get_tp_group from sglang.srt.dllm.config import DllmConfig from sglang.srt.environ import envs from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder -from sglang.srt.layers.dp_attention import compute_dp_attention_world_info +from sglang.srt.layers.dp_attention import ( + compute_dp_attention_world_info, + get_attention_tp_group, +) from sglang.srt.layers.moe import initialize_moe_config from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config from sglang.srt.managers.io_struct import ( @@ -546,25 +550,24 @@ class Scheduler( self.max_running_requests // self.pp_size, 1 ) - self.tp_group = self.tp_worker.get_tp_group() + self.tp_group = get_tp_group() self.tp_cpu_group = self.tp_group.cpu_group - self.attn_tp_group = self.tp_worker.get_attention_tp_group() - self.attn_tp_cpu_group = self.tp_worker.get_attention_tp_cpu_group() + self.attn_tp_group = get_attention_tp_group() + self.attn_tp_cpu_group = self.attn_tp_group.cpu_group self.pp_group = get_pp_group() self.world_group = get_world_group() - # With DP attention enabled, the entry rank is attn_tp_rank==0; - # otherwise the entry rank is TP group local rank 0. - # For #11910, use the CPU communication group to broadcast VLM Python objects, - # avoiding any coupling with CUDA streams/devices. - if self.server_args.enable_dp_attention: - self.cpu_group = self.attn_tp_cpu_group - self.entry_rank = self.attn_tp_group.first_rank - self.is_entry_rank = self.attn_tp_rank == 0 - else: - self.cpu_group = self.tp_cpu_group - self.entry_rank = self.tp_group.first_rank - self.is_entry_rank = self.tp_group.rank_in_group == 0 + # NOTE: dp_tp_* are request/data-plane coordination groups (not tensor collectives). + # When DP attention is enabled, scope to the attention-TP group; otherwise use + # the base TP group. Entry rank is the local rank 0 in that group. + # Use the CPU (gloo) group to broadcast VLM Python objects and avoid CUDA + # stream/device coupling (#11910). + self.dp_tp_group = ( + self.attn_tp_group + if self.server_args.enable_dp_attention + else self.tp_group + ) + self.dp_tp_cpu_group = self.dp_tp_group.cpu_group self.pad_input_ids_func = self.tp_worker.get_pad_input_ids_func() set_random_seed(self.random_seed) @@ -1358,10 +1361,10 @@ class Scheduler( if ( torch.distributed.is_available() and torch.distributed.is_initialized() - and self.cpu_group is not None + and self.dp_tp_cpu_group is not None ): group_world_size = torch.distributed.get_world_size( - group=self.cpu_group + group=self.dp_tp_cpu_group ) except Exception as e: logger.warning( @@ -1374,14 +1377,16 @@ class Scheduler( # Since the Scheduler is single-threaded, any large CPU cost will impact # handling of other messages. For example, CPU hits 99.9% can significantly # increase the CUDA kernel launch time. - if self.is_entry_rank: + if self.dp_tp_group.rank_in_group == 0: # Only the entry rank materializes once from dict. image_inputs = MultimodalInputs.from_dict(raw_mm_inputs) # Broadcast to other TP ranks (use src=0 within the group). if group_world_size > 1: obj_list = [image_inputs] torch.distributed.broadcast_object_list( - obj_list, src=self.entry_rank, group=self.cpu_group + obj_list, + src=self.dp_tp_group.first_rank, + group=self.dp_tp_cpu_group, ) image_inputs = obj_list[0] else: @@ -1389,7 +1394,9 @@ class Scheduler( if group_world_size > 1: obj_list = [None] torch.distributed.broadcast_object_list( - obj_list, src=self.entry_rank, group=self.cpu_group + obj_list, + src=self.dp_tp_group.first_rank, + group=self.dp_tp_cpu_group, ) image_inputs = obj_list[0] else: diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 6a3069868..b785f8073 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -1079,7 +1079,7 @@ class SchedulerPPMixin: """ polls = poll_and_all_reduce( [req.disagg_kv_sender if is_send else req.kv_receiver for req in req_queue], - self.tp_worker.get_attention_tp_cpu_group(), + self.attn_tp_cpu_group, ) rids: List = [] for poll_statuses in poll_statuses_group: diff --git a/python/sglang/srt/managers/scheduler_profiler_mixin.py b/python/sglang/srt/managers/scheduler_profiler_mixin.py index 559997280..6b30b790f 100644 --- a/python/sglang/srt/managers/scheduler_profiler_mixin.py +++ b/python/sglang/srt/managers/scheduler_profiler_mixin.py @@ -39,7 +39,7 @@ class SchedulerProfilerMixin: if envs.SGLANG_PROFILE_V2.get(): self._profile_manager = ProfileManager( tp_rank=self.tp_rank, - cpu_group=self.cpu_group, + cpu_group=self.dp_tp_cpu_group, gpu_id=self.gpu_id, ) return @@ -180,7 +180,7 @@ class SchedulerProfilerMixin: schema.writeSchema(connection) connection.commit() del connection - torch.distributed.barrier(self.cpu_group) + torch.distributed.barrier(self.dp_tp_cpu_group) self.rpd_profiler = rpdTracerControl() self.rpd_profiler.setPythonTrace(True) @@ -291,14 +291,14 @@ class SchedulerProfilerMixin: self.torch_profiler.export_chrome_trace( os.path.join(self.torch_profiler_output_dir, filename) ) - torch.distributed.barrier(self.cpu_group) + torch.distributed.barrier(self.dp_tp_cpu_group) if self.rpd_profiler is not None: self.rpd_profiler.rangePop() self.rpd_profiler.stop() self.rpd_profiler.flush() - torch.distributed.barrier(self.cpu_group) + torch.distributed.barrier(self.dp_tp_cpu_group) if self.tp_rank == 0: from sglang.srt.utils.rpd_utils import rpd_to_chrome_trace diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 88076228a..af4000729 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -83,15 +83,6 @@ class BaseTpWorker(ABC): def get_pad_input_ids_func(self): return getattr(self.model_runner.model, "pad_input_ids", None) - def get_tp_group(self): - return self.model_runner.tp_group - - def get_attention_tp_group(self): - return self.model_runner.attention_tp_group - - def get_attention_tp_cpu_group(self): - return getattr(self.model_runner.attention_tp_group, "cpu_group", None) - def get_memory_pool(self): return ( self.model_runner.req_to_token_pool,