diff --git a/python/sglang/srt/managers/data_parallel_controller.py b/python/sglang/srt/managers/data_parallel_controller.py index 3256e07f0..0c7921e92 100644 --- a/python/sglang/srt/managers/data_parallel_controller.py +++ b/python/sglang/srt/managers/data_parallel_controller.py @@ -163,7 +163,13 @@ class DataParallelController: if server_args.enable_dp_attention: self.launch_dp_attention_schedulers(server_args, port_args) - self.control_message_step = server_args.tp_size + # When local control broadcast is enabled, send control messages to + # every DP group leader (attn_tp_rank=0) so each leader broadcasts + # within its own attn_tp_group instead of the full tp_group. + # Otherwise fall back to the original behaviour: send to only the + # first leader, which then broadcasts over the full tp_group. + local_ctrl = server_args.enable_dp_attention_local_control_broadcast + self.control_message_step = 1 if local_ctrl else server_args.tp_size else: self.launch_dp_schedulers(server_args, port_args) self.control_message_step = 1 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 16193d132..7bb566aec 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1575,7 +1575,28 @@ class Scheduler( src=self.attn_cp_group.ranks[0], ) - if self.tp_size != 1: + # When dp_attention_local_control_broadcast is enabled, each DP + # group leader already receives control messages from the DP + # controller, so we broadcast within attn_tp_group + attn_cp_group + # instead of the full tp_group. This avoids an expensive + # all-ranks gloo sync. + _local_ctrl = self.server_args.enable_dp_attention_local_control_broadcast + if _local_ctrl: + if self.attn_tp_size != 1: + control_reqs = broadcast_pyobj( + control_reqs, + self.attn_tp_group.rank, + self.attn_tp_cpu_group, + src=self.attn_tp_group.ranks[0], + ) + if self.attn_cp_size != 1: + control_reqs = broadcast_pyobj( + control_reqs, + self.attn_cp_group.rank, + self.attn_cp_cpu_group, + src=self.attn_cp_group.ranks[0], + ) + elif self.tp_size != 1: control_reqs = broadcast_pyobj( control_reqs, self.tp_group.rank, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 2390b3414..7afa57281 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -639,6 +639,7 @@ class ServerArgs: disable_overlap_schedule: bool = False enable_mixed_chunk: bool = False enable_dp_attention: bool = False + enable_dp_attention_local_control_broadcast: bool = False enable_dp_lm_head: bool = False enable_two_batch_overlap: bool = False enable_single_batch_overlap: bool = False @@ -5830,6 +5831,13 @@ class ServerArgs: action="store_true", help="Enabling data parallelism for attention and tensor parallelism for FFN. The dp size should be equal to the tp size. Currently DeepSeek-V2 and Qwen 2/3 MoE models are supported.", ) + parser.add_argument( + "--enable-dp-attention-local-control-broadcast", + action="store_true", + help="With DP-attention, send control messages to every DP group leader " + "and broadcast within attn_tp_group instead of the full tp_group. " + "Eliminates a costly all-ranks gloo sync on every scheduler iteration.", + ) parser.add_argument( "--enable-dp-lm-head", action="store_true",