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

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":

python/sglang/srt/model_executor/runner/eager_runner.py

@@ -371,13 +371,26 @@ class EagerRunner(BaseRunner):
                 else hidden_states
             )
 
-        hidden_states = cp_gather_after_forward(
+        stream = torch.cuda.current_stream()
-            hidden_states, forward_batch, 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
+            # 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,

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.

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
+        # Follow lm_head's group choice; attn_tp_group degenerates to size 1
-        if attn_tp_size != tp_size:
+        # under prefill CP while lm_head still shards over the full TP group.
-            raise ValueError(
+        parallel = get_parallel()
-                "DSpark markov_w2 TP-shard needs the attn-TP group (used for the per-step "
+        shard_group = (
-                f"all-gather) to equal the lm_head shard group, got attn_tp_size="
+            parallel.attn_tp_group
-                f"{attn_tp_size} vs lm_head tp_size={tp_size}. This config (e.g. DP "
+            if getattr(lm_head, "use_attn_tp_group", False)
-                "attention without --enable-dp-lm-head, where lm_head shards over the "
+            else parallel.tp_group
-                "global TP group) is unsupported; disable "
-                "SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD."
         )
+        shard_group_size = shard_group.world_size
+        if shard_group_size != tp_size:
+            raise ValueError(
+                "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]

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()