Skip to content

Deepseek v4 support dp attention is not TP size #24952

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zhangxiaolei123456 wants to merge 11 commits into
base: main
Choose a base branch
from
+14 −7
Open

Deepseek v4 support dp attention is not TP size #24952

Changes from 1 commit
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
ad78d18
Update deepseek_v4.py
zhangxiaolei123456 May 11, 2026
bc4135a
Merge branch 'main' into main_deepseek_dp_attention
zhangxiaolei123456 May 12, 2026
a9b5b05
Merge branch 'main' into main_deepseek_dp_attention
zhangxiaolei123456 May 12, 2026
3944eda
Merge branch 'main' into main_deepseek_dp_attention
zhangxiaolei123456 May 12, 2026
45c6521
Update deepseek_v4.py
zhangxiaolei123456 May 12, 2026
d4da6dc
Update deepseek_v4.py
zhangxiaolei123456 May 12, 2026
4813ac5
Merge branch 'main' into main_deepseek_dp_attention
zhangxiaolei123456 May 12, 2026
eb1d8df
Merge branch 'main' into main_deepseek_dp_attention
zhangxiaolei123456 May 14, 2026
5ec497c
Feat/dp attn 0514 (#490)
shiyu7 May 14, 2026
fe29439
Merge branch 'main' into main_deepseek_dp_attention
zhangxiaolei123456 May 15, 2026
b5f890b
Merge branch 'main' into main_deepseek_dp_attention
zhangxiaolei123456 Jun 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 12 additions & 7 deletions python/sglang/srt/models/deepseek_v4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,9 @@
from sglang.srt.layers.deepseek_v4_rope import apply_rotary_emb_triton
from sglang.srt.layers.dp_attention import (
_DpGatheredBufferWrapper,
attn_tp_all_reduce,
attn_tp_all_gather,
dp_gather_partial,
dp_gather_replicate,
dp_scatter,
get_attention_cp_rank,
get_attention_cp_size,
Expand Down Expand Up @@ -299,7 +300,7 @@ def __init__(
self.hidden_size,
bias=False,
quant_config=quant_config,
reduce_results=attn_tp_size > 1,
reduce_results=attn_tp_size == get_tensor_model_parallel_world_size() and attn_tp_size > 1,

@gemini-code-assist gemini-code-assist Bot May 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The condition and attn_tp_size > 1 is redundant here.

  • If get_tensor_model_parallel_world_size() is 1, attn_tp_size is also 1, making attn_tp_size > 1 false. The whole expression is false.
  • If get_tensor_model_parallel_world_size() is greater than 1, and attn_tp_size == get_tensor_model_parallel_world_size(), then attn_tp_size > 1 is implicitly true.

The RowParallelLinear layer already checks if tp_size > 1 before performing an all-reduce, so setting reduce_results=True when tp_size=1 has no effect. Simplifying this condition improves clarity.

Suggested change
reduce_results=attn_tp_size == get_tensor_model_parallel_world_size() and attn_tp_size > 1,
reduce_results=attn_tp_size == get_tensor_model_parallel_world_size(),

prefix=add_prefix("wo_b", prefix),
tp_rank=attn_tp_rank,
tp_size=attn_tp_size,
Expand Down Expand Up @@ -513,9 +514,6 @@ def forward(
forward_batch: ForwardBatch,
) -> torch.Tensor:
if not get_attn_tp_context().input_scattered and x.shape[0] == 0:
assert (
not self.wo_b.reduce_results
), "short-circuiting allreduce will lead to hangs"
return x
Comment on lines 821 to 822

@gemini-code-assist gemini-code-assist Bot May 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The removal of the assertion that checks self.wo_b.reduce_results could lead to hangs if not handled carefully. The original assertion assert (not self.wo_b.reduce_results) was a safeguard against short-circuiting an all-reduce operation when x.shape[0] == 0.

While the logic for input_scattered seems to cover the DP case, it's safer to retain a check to prevent potential hangs in non-DP scenarios where token distribution might be uneven across ranks, or if all ranks have zero tokens but reduce_results is true. A more robust approach would be to handle the zero-sized tensor case within the RowParallelLinear layer itself, but as a direct fix, consider reintroducing a check or ensuring that x.shape[0] == 0 implies all ranks have zero tokens when reduce_results is true.


attn_backend = forward_batch.attn_backend
Expand Down Expand Up @@ -590,6 +588,8 @@ def forward(
o = torch.einsum("tgd,grd->tgr", o, wo_a)

o, _ = self.wo_b(o.flatten(1))
if self.tp_size > 1 and self.tp_size < get_tensor_model_parallel_world_size():
o = attn_tp_all_reduce(o)

return o

Expand Down Expand Up @@ -820,7 +820,10 @@ def forward(
input_ids_global = input_ids
elif _use_tp_moe_gather:
hidden_states, local_hidden_states = get_global_dp_buffer(), hidden_states
dp_gather_partial(hidden_states, local_hidden_states, forward_batch)
# hidden_states here follow TP_ATTN_FULL semantics: they are replicated
# within an attention-TP group. Use replicate gather to avoid summing the
# same activations across attention-TP ranks before entering MLP/MoE.
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_attention_tp_size(), get_attention_tp_rank()
Expand Down Expand Up @@ -939,7 +942,9 @@ def forward(
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 here to avoid summing duplicated ids when attention_tp_size > 1.
dp_gather_replicate(input_ids_global, input_ids[:, None], forward_batch)
input_ids_global = input_ids_global.squeeze(-1)
else:
input_ids_global = input_ids
Expand Down
Loading