From 9d4be40124cc4ef4ac3cc3a5be086f78a030a32c Mon Sep 17 00:00:00 2001 From: ybyang <10629930+whybeyoung@users.noreply.github.com> Date: Tue, 11 Aug 2026 14:26:29 +0800 Subject: [PATCH] Fix DSpark + DeepSeek V4 prefill CP compatibility (#33865) --- .../sglang/srt/arg_groups/speculative_hook.py | 3 +- .../srt/model_executor/runner/eager_runner.py | 21 +++++++-- python/sglang/srt/models/deepseek_v4.py | 17 ++++--- .../sglang/srt/models/deepseek_v4_dspark.py | 27 ++++++++---- .../cp/test_deepseek_v4_flash_fp4_b200_cp.py | 44 +++++++++++++++++++ 5 files changed, 91 insertions(+), 21 deletions(-) diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 427acdd38..1b1f65e3c 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -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": diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index e3e72d4ae..1ceb52f91 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -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, diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 812bb68e2..0c45035b6 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -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. diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index d7ca4e373..c93c84fad 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -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] diff --git a/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py b/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py index e4cd6bbe6..2975c9d40 100644 --- a/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py +++ b/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py @@ -198,5 +198,49 @@ class TestDSV4FlashFP4B200Balanced_CP_NonDeepEP( kill_process_tree(cls.process.pid) +# DSPARK draft is bundled with the -DSpark checkpoint. +DSPARK_MODEL = "deepseek-ai/DeepSeek-V4-Flash-DSpark" + + +class TestDSV4FlashFP4B200_CP_DSpark( + BasicDecodeCorrectnessMixin, + GSM8KMixin, + CustomTestCase, +): + """DSPARK speculation + prefill CP (interleave, CP_V2, attn_cp=tp).""" + + gsm8k_accuracy_thres = 0.90 + + @classmethod + def setUpClass(cls): + cls.model = try_cached_model(DSPARK_MODEL) + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=[ + "--trust-remote-code", + "--tp", + "4", + "--attn-cp-size", + "4", + "--speculative-algorithm", + "DSPARK", + "--enable-prefill-cp", + "--cp-strategy", + "interleave", + "--moe-runner-backend", # for fp4 checkpoint + "flashinfer_mxfp4", + ], + env={"SGLANG_ENABLE_CP_V2": "1"}, + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + if __name__ == "__main__": unittest.main()