minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Fix DeepSeek V4 DP reduce scatter when use attention DP + MoE TP (#27191)

Browse Source
This commit is contained in:
Yueming Yuan
2026-06-06 18:24:33 -07:00
committed by GitHub
parent 5160f7914e
commit 4c8a022f38
1 changed files with 10 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_v4.py
+9 -1
View File
@@ -59,6 +59,7 @@ from sglang.srt.layers.dp_attention import (
get_attention_dp_size, get_attention_dp_size,
get_attention_tp_rank, get_attention_tp_rank,
get_attention_tp_size, get_attention_tp_size,
get_dp_global_num_tokens,
get_global_dp_buffer, get_global_dp_buffer,
get_local_dp_buffer, get_local_dp_buffer,
is_dp_attention_enabled, 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.linear import ColumnParallelLinear, RowParallelLinear
from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.mhc import mhc_fused_post_pre 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.moe.fused_moe_triton import FusedMoE
from sglang.srt.layers.quantization.fp8_kernel import sglang_per_token_group_quant_fp8 from sglang.srt.layers.quantization.fp8_kernel import sglang_per_token_group_quant_fp8
from sglang.srt.layers.rotary_embedding import get_rope_wrapper from sglang.srt.layers.rotary_embedding import get_rope_wrapper
@@ -1431,6 +1432,13 @@ class DeepseekV4DecoderLayer(nn.Module):
get_local_dp_buffer(get_tp_group()), get_local_dp_buffer(get_tp_group()),
hidden_states, hidden_states,
) )
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) dp_scatter(hidden_states, global_hidden_states, forward_batch)
if _use_tp_attn_a2a_scatter: if _use_tp_attn_a2a_scatter:
assert _a2a_scatter_chunks is not None assert _a2a_scatter_chunks is not None