From 4c8a022f38e386731d2e7026618a491e7fa6f0c2 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Sat, 6 Jun 2026 18:24:33 -0700 Subject: [PATCH] Fix DeepSeek V4 DP reduce scatter when use attention DP + MoE TP (#27191) --- python/sglang/srt/models/deepseek_v4.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index c5fc558c9..7786c9edb 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -59,6 +59,7 @@ from sglang.srt.layers.dp_attention import ( get_attention_dp_size, get_attention_tp_rank, get_attention_tp_size, + get_dp_global_num_tokens, get_global_dp_buffer, get_local_dp_buffer, is_dp_attention_enabled, @@ -67,7 +68,7 @@ from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.mhc import mhc_fused_post_pre -from sglang.srt.layers.moe import get_moe_a2a_backend +from sglang.srt.layers.moe import get_moe_a2a_backend, should_use_dp_reduce_scatterv from sglang.srt.layers.moe.fused_moe_triton import FusedMoE from sglang.srt.layers.quantization.fp8_kernel import sglang_per_token_group_quant_fp8 from sglang.srt.layers.rotary_embedding import get_rope_wrapper @@ -1431,7 +1432,14 @@ class DeepseekV4DecoderLayer(nn.Module): get_local_dp_buffer(get_tp_group()), hidden_states, ) - dp_scatter(hidden_states, global_hidden_states, forward_batch) + if should_use_dp_reduce_scatterv(): + get_tp_group().reduce_scatterv( + global_hidden_states, + output=hidden_states, + sizes=get_dp_global_num_tokens(), + ) + else: + dp_scatter(hidden_states, global_hidden_states, forward_batch) if _use_tp_attn_a2a_scatter: assert _a2a_scatter_chunks is not None gathered = [torch.empty_like(t) for t in _a2a_scatter_chunks]