From 7c7326ccb328abaf85425e766b016e58bdf6f9cd Mon Sep 17 00:00:00 2001 From: Michael Gschwind Date: Mon, 10 Aug 2026 17:03:19 -0700 Subject: [PATCH] Fix DeepSeek-V4/DeepSeek-V4-Pro DP-attention gather semantics (#31700) --- python/sglang/srt/models/deepseek_v4.py | 9 +++++++-- python/sglang/srt/models/deepseek_v4_nextn.py | 9 +++++++-- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index c6a425f31..812bb68e2 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -1942,7 +1942,9 @@ class DeepseekV4DecoderLayer(nn.Module): ) if _do_shared_local and local_hidden_states.shape[0] > 0: _shared_local = self.mlp._forward_shared_experts(local_hidden_states) - dp_gather_partial(hidden_states, local_hidden_states, forward_batch) + # self_attn has already reduced across attention TP, so these hidden + # states are replicated and must not be summed by a partial gather. + dp_gather_replicate(hidden_states, local_hidden_states, forward_batch) _a2a_scatter_chunks: Optional[List[torch.Tensor]] = None if _use_tp_attn_a2a_scatter: s, r = get_parallel().attn_tp_size, get_parallel().attn_tp_rank @@ -2465,7 +2467,10 @@ class DeepseekV4Model(nn.Module): ) # Token ids are replicated within an attention-TP group. Use replicate # gather here to avoid summing duplicated ids when attention_tp_size > 1. - dp_gather_replicate(input_ids_global, input_ids[:, None], forward_batch) + # Clone because the MAX_LEN gather may zero its local input in place. + dp_gather_replicate( + input_ids_global, input_ids[:, None].clone(), forward_batch + ) input_ids_global = input_ids_global.squeeze(-1) else: input_ids_global = input_ids diff --git a/python/sglang/srt/models/deepseek_v4_nextn.py b/python/sglang/srt/models/deepseek_v4_nextn.py index 94abf6e5d..39c064d9c 100644 --- a/python/sglang/srt/models/deepseek_v4_nextn.py +++ b/python/sglang/srt/models/deepseek_v4_nextn.py @@ -18,7 +18,7 @@ from sglang.srt.layers.cp.utils import ( is_cp_v2_active, ) from sglang.srt.layers.dp_attention import ( - dp_gather_partial, + dp_gather_replicate, get_global_dp_buffer_len, is_dp_attention_enabled, ) @@ -171,7 +171,12 @@ class DeepseekV4ModelNextN(nn.Module): dtype=input_ids.dtype, device=input_ids.device, ) - dp_gather_partial(input_ids_global, input_ids[:, None], forward_batch) + # Token IDs are replicated within an attention-TP group. Use replicate + # gather to avoid summing duplicated IDs when attention_tp_size > 1. + # Clone because the MAX_LEN gather may zero its local input in place. + dp_gather_replicate( + input_ids_global, input_ids[:, None].clone(), forward_batch + ) input_ids_global = input_ids_global.squeeze(-1) else: input_ids_global = input_ids