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

Fix DSpark + DeepSeek V4 prefill CP compatibility (#33865)

Browse Source
This commit is contained in:
ybyang
2026-08-10 23:26:29 -07:00
committed by GitHub
parent afa2d5570b
commit 9d4be40124
5 changed files with 91 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/speculative_hook.py
+2 -1
View File
@@ -279,7 +279,8 @@ def _handle_dspark(server_args: ServerArgs) -> None:
if not server_args.device.startswith("cuda"):
raise ValueError("DSpark speculative decoding only supports CUDA device.")
if server_args.enable_dp_attention:
# dp_size==1 with dp_attention is a degenerate flag under DSV4 CP; skip DP-only checks.
if server_args.enable_dp_attention and server_args.dp_size > 1:
if not server_args.enable_dp_lm_head:
raise ValueError("DSpark with dp attention requires --enable-dp-lm-head.")
if server_args.moe_a2a_backend != "none":
python/sglang/srt/model_executor/runner/eager_runner.py
+17 -4
View File
@@ -371,14 +371,27 @@ class EagerRunner(BaseRunner):
else hidden_states
)
hidden_states = cp_gather_after_forward(
hidden_states, forward_batch, torch.cuda.current_stream()
)
stream = torch.cuda.current_stream()
hidden_states = cp_gather_after_forward(hidden_states, forward_batch, stream)
# DSpark aux tensors ride the same CP token split; gather them the same way.
if aux_hidden_states is not None:
if isinstance(aux_hidden_states, torch.Tensor):
aux_hidden_states = cp_gather_after_forward(
aux_hidden_states, forward_batch, stream
)
else:
aux_hidden_states = [
cp_gather_after_forward(aux, forward_batch, stream)
for aux in aux_hidden_states
]
logits_kwargs = {}
# DSV4 returns (hidden_states, hidden_states_before_norm) from its model body.
if isinstance(hidden_states, tuple):
hidden_states, hidden_states_before_norm = hidden_states
logits_kwargs["hidden_states_before_norm"] = hidden_states_before_norm
# Mirror DeepseekV4ForCausalLM.forward: drop pre_hc_head when
# DSpark aux capture is on, else it overrides the packed aux.
if aux_hidden_states is None:
logits_kwargs["hidden_states_before_norm"] = hidden_states_before_norm
return model.logits_processor(
forward_batch.input_ids,
hidden_states,
python/sglang/srt/models/deepseek_v4.py
+10 -7
View File
@@ -2492,12 +2492,6 @@ class DeepseekV4Model(nn.Module):
if hasattr(forward_batch, _attr):
delattr(forward_batch, _attr)
capture_dspark = self.dspark_layers_to_capture is not None
if capture_dspark and use_prefill_cp:
raise NotImplementedError(
"DSpark aux hidden-state capture is not supported together with "
"DeepSeek-V4 prefill context parallelism (attn_cp_size > 1). Disable one "
"of them: DSpark static-verify is CP-off for v1."
)
dspark_aux_hidden_states: List[torch.Tensor] = []
# DSpark aux capture needs the per-layer eager loop (TBO's overlapped
# execution cannot expose per-layer completed hidden states), so skip
@@ -2548,12 +2542,21 @@ class DeepseekV4Model(nn.Module):
# CP all-gather only on the last PP rank; PP IPC carries CP-split tensors.
if self.pp_group.is_last_rank and use_prefill_cp and not cp_v2_active:
stream = torch.cuda.current_stream()
hidden_states = cp_all_gather_rerange_output(
hidden_states,
self.cp_size,
forward_batch,
torch.cuda.current_stream(),
stream,
)
# Gather DSpark aux tensors on the same CP token split.
if capture_dspark:
dspark_aux_hidden_states = [
cp_all_gather_rerange_output(
aux, self.cp_size, forward_batch, stream
)
for aux in dspark_aux_hidden_states
]
if not self.pp_group.is_last_rank:
# Flatten 3D mHC tensor for PP IPC.
python/sglang/srt/models/deepseek_v4_dspark.py
+18 -9
View File
@@ -319,6 +319,7 @@ class DSparkV4MarkovHead(nn.Module):
self.markov_rank, self.vocab_size, bias=False, dtype=markov_w2_dtype
)
self._tp_shard: Optional[MarkovW2ShardGeometry] = None
self._shard_group = None
def configure_tp_shard(self, *, lm_head: nn.Module) -> None:
if not self._opt_markov_w2_tp_shard:
@@ -338,16 +339,23 @@ class DSparkV4MarkovHead(nn.Module):
f"num_embeddings_per_partition({per_partition}) * tp_size({tp_size}) != "
f"num_embeddings_padded({num_padded})."
)
attn_tp_size = get_parallel().attn_tp_group.world_size
if attn_tp_size != tp_size:
# Follow lm_head's group choice; attn_tp_group degenerates to size 1
# under prefill CP while lm_head still shards over the full TP group.
parallel = get_parallel()
shard_group = (
parallel.attn_tp_group
if getattr(lm_head, "use_attn_tp_group", False)
else parallel.tp_group
)
shard_group_size = shard_group.world_size
if shard_group_size != tp_size:
raise ValueError(
"DSpark markov_w2 TP-shard needs the attn-TP group (used for the per-step "
f"all-gather) to equal the lm_head shard group, got attn_tp_size="
f"{attn_tp_size} vs lm_head tp_size={tp_size}. This config (e.g. DP "
"attention without --enable-dp-lm-head, where lm_head shards over the "
"global TP group) is unsupported; disable "
"SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD."
"DSpark markov_w2 TP-shard needs the per-step all-gather group to "
f"equal the lm_head shard group, got shard_group_size="
f"{shard_group_size} vs lm_head tp_size={tp_size}. "
"Disable SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD."
)
self._shard_group = shard_group
self._tp_shard = MarkovW2ShardGeometry(
tp_size=tp_size,
org_vocab_start=int(lm_head.shard_indices.org_vocab_start_index),
@@ -400,7 +408,8 @@ class DSparkV4MarkovHead(nn.Module):
bias = F.linear(latent.float(), weight_local)
step_local = BuildStepLocal.execute(bias=bias, base_local=base_local)
if shard.tp_size > 1:
full = get_parallel().attn_tp_group.all_gather(step_local, dim=-1)
assert self._shard_group is not None
full = self._shard_group.all_gather(step_local, dim=-1)
else:
full = step_local
return full[..., : self.vocab_size]