[Feature] DSpark Support Prefill TP+CP

python/sglang/srt/arg_groups/speculative_hook.py

@@ -275,23 +275,48 @@ def _target_checkpoint_bundles_dspark_draft(server_args: ServerArgs) -> bool:
     return checkpoint_bundles_dspark_draft(server_args.get_model_config().hf_config)
 
 
+def _is_supported_dspark_pd_prefill_cp(server_args: ServerArgs) -> bool:
+    model_arch = server_args.get_model_config().hf_config.architectures[0]
+    attn_tp_size = (
+        server_args.tp_size // server_args.dp_size // server_args.attn_cp_size
+    )
+    return (
+        server_args.disaggregation_mode == "prefill"
+        and server_args.disaggregation_transfer_backend == "mooncake"
+        and server_args.pp_size == 1
+        and server_args.attn_cp_size > 1
+        and attn_tp_size == 1
+        and server_args.enable_prefill_cp
+        and server_args.cp_strategy == "interleave"
+        and model_arch == "DeepseekV4ForCausalLM"
+    )
+
+
 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:
+    pd_prefill_cp = _is_supported_dspark_pd_prefill_cp(server_args)
+    if server_args.attn_cp_size > 1 and not pd_prefill_cp:
+        raise ValueError(
+            "DSpark context parallel is only supported for DeepSeek-V4 PD prefill "
+            "with Mooncake, pp_size == 1, attn_tp_size == 1, and interleave "
+            f"(round-robin-split) CP; got disaggregation_mode="
+            f"{server_args.disaggregation_mode!r}, pp_size={server_args.pp_size}, "
+            f"attn_cp_size={server_args.attn_cp_size}, attn_tp_size="
+            f"{server_args.tp_size // server_args.dp_size // server_args.attn_cp_size}, "
+            f"cp_strategy={server_args.cp_strategy!r}, transfer_backend="
+            f"{server_args.disaggregation_transfer_backend!r}."
+        )
+
+    if server_args.enable_dp_attention and not pd_prefill_cp:
         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":
             raise ValueError(
                 "DSpark with dp attention only supports the built-in TP MoE "
                 f"(moe_a2a_backend='none'), got {server_args.moe_a2a_backend!r}."
             )
-        if server_args.attn_cp_size > 1:
-            raise ValueError(
-                "DSpark with dp attention does not support context parallel "
-                f"(attn_cp_size={server_args.attn_cp_size})."
-            )
         if (
             server_args.speculative_moe_a2a_backend is not None
             and server_args.speculative_moe_a2a_backend != server_args.moe_a2a_backend

python/sglang/srt/models/deepseek_v4.py

@@ -53,6 +53,7 @@
 )
 from sglang.srt.layers.attention.dsv4.compressor import Compressor
 from sglang.srt.layers.attention.dsv4.indexer import C4Indexer
+from sglang.srt.layers.aux_hidden_states import AuxHiddenStatePacker
 from sglang.srt.layers.communicator import get_attn_tp_context
 from sglang.srt.layers.communicator_dsa_cp import (
     dsa_cp_gather_hidden_states,
@@ -2389,13 +2390,11 @@ def forward(
             if hasattr(forward_batch, _attr):
                 delattr(forward_batch, _attr)
         capture_dspark = self.dspark_layers_to_capture is not None
-        if capture_dspark and dsa_use_prefill_cp(forward_batch):
-            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_hidden_states = (
+            AuxHiddenStatePacker(len(self.dspark_layers_to_capture))
+            if capture_dspark
+            else None
+        )
         # DSpark aux capture needs the per-layer eager loop (TBO's overlapped
         # execution cannot expose per-layer completed hidden states), so skip
         # TBO when capturing -- a perf-only downgrade, not a correctness one.
@@ -2437,6 +2436,7 @@ def forward(
                         )
                     else:
                         completed = hidden_states
+                    assert dspark_aux_hidden_states is not None
                     dspark_aux_hidden_states.append(completed.mean(dim=1))
             if use_fused and last_layer is not None:
                 hidden_states = last_layer.hc_post(
@@ -2451,6 +2451,19 @@ def forward(
                 forward_batch,
                 torch.cuda.current_stream(),
             )
+            if dspark_aux_hidden_states is not None:
+                dspark_aux_hidden = cp_all_gather_rerange_output(
+                    dspark_aux_hidden_states.finalize(),
+                    self.cp_size,
+                    forward_batch,
+                    torch.cuda.current_stream(),
+                )
+            else:
+                dspark_aux_hidden = None
+        elif dspark_aux_hidden_states is not None:
+            dspark_aux_hidden = dspark_aux_hidden_states.finalize()
+        else:
+            dspark_aux_hidden = None
 
         if not self.pp_group.is_last_rank:
             # Flatten 3D mHC tensor for PP IPC.
@@ -2464,7 +2477,8 @@ def forward(
         hidden_states = self.norm(hidden_states)
 
         if capture_dspark:
-            return (hidden_states, pre_hc_head), dspark_aux_hidden_states
+            assert dspark_aux_hidden is not None
+            return (hidden_states, pre_hc_head), dspark_aux_hidden
 
         return hidden_states, pre_hc_head
 

python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py

@@ -46,10 +46,22 @@ def inject_target_hidden(
         n_real = positions.shape[0]
         if target_hidden.shape[0] > n_real:
             target_hidden = target_hidden[:n_real]
+        if target_hidden.shape[0] != n_real or cache_loc.shape[0] != n_real:
+            raise ValueError(
+                "DSpark target hidden injection requires one hidden row and cache "
+                f"location per position; got hidden_rows={target_hidden.shape[0]}, "
+                f"cache_locs={cache_loc.shape[0]}, positions={n_real}."
+            )
         if cache_loc_2d is not None:
             cache_loc_2d = cache_loc_2d.to(
                 device=device, dtype=torch.int64, non_blocking=True
             )
+            if cache_loc_2d.numel() != n_real:
+                raise ValueError(
+                    "DSpark target hidden injection requires cache_loc_2d to cover "
+                    f"every position; got cache_loc_2d={cache_loc_2d.numel()}, "
+                    f"positions={n_real}."
+                )
         if commit_lens is not None:
             commit_lens = commit_lens.to(
                 device=device, dtype=torch.int32, non_blocking=True

python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py

@@ -162,16 +162,7 @@ def __init__(
             draft_token_num=int(self.gamma), device=self.device
         )
 
-        target_model = self.target_worker.model_runner.model
-        lm_head = getattr(target_model, "lm_head", None)
-        if lm_head is None or not hasattr(lm_head, "weight"):
-            raise RuntimeError(
-                "DSpark requires the target model to expose `lm_head` with `weight`."
-            )
-        self.draft_model.attach_shared_modules(
-            embed_tokens=self._resolve_target_embed_tokens(target_model),
-            lm_head=lm_head,
-        )
+        self._attach_shared_modules()
 
         self._verify_planner = DSparkVerifyPlanner(
             draft_model=self.draft_model,
@@ -278,6 +269,21 @@ def __init__(
         if self._is_pd_prefill and not self._draft_is_moe:
             self.draft_model.prune_to_ctx_kv_injection()
 
+    def _attach_shared_modules(self) -> None:
+        if self._is_pd_prefill:
+            return
+
+        target_model = self.target_worker.model_runner.model
+        lm_head = getattr(target_model, "lm_head", None)
+        if lm_head is None or not hasattr(lm_head, "weight"):
+            raise RuntimeError(
+                "DSpark requires the target model to expose `lm_head` with `weight`."
+            )
+        self.draft_model.attach_shared_modules(
+            embed_tokens=self._resolve_target_embed_tokens(target_model),
+            lm_head=lm_head,
+        )
+
     def _resolve_target_embed_tokens(self, target_model):
         if hasattr(target_model, "get_input_embeddings"):
             return target_model.get_input_embeddings()

test/registered/spec/dspark/test_dspark_draft_path_default.py

@@ -42,6 +42,18 @@ def _make_dspark_server_args(
     return server_args
 
 
+def _enable_pd_prefill_cp(server_args: ServerArgs) -> None:
+    server_args.disaggregation_mode = "prefill"
+    server_args.tp_size = 8
+    server_args.dp_size = 1
+    server_args.pp_size = 1
+    server_args.attn_cp_size = 8
+    server_args.enable_prefill_cp = True
+    server_args.cp_strategy = "interleave"
+    server_args.enable_dp_attention = True
+    server_args.enable_dp_lm_head = False
+
+
 class TestTargetCheckpointBundlesDsparkDraft(CustomTestCase):
     def test_bundled_dsv4_config_is_detected(self):
         server_args = _make_dspark_server_args(
@@ -83,6 +95,30 @@ def test_explicit_draft_path_is_not_overwritten(self):
             "deepseek-ai/some-other-dspark-draft",
         )
 
+    def test_pd_prefill_cp_does_not_require_dp_lm_head(self):
+        """Prefill CP sets the DP-attention flag but does not use DP LM head."""
+        server_args = _make_dspark_server_args(
+            model_path=_BUNDLED_MODEL_PATH, hf_config=_bundled_hf_config()
+        )
+        _enable_pd_prefill_cp(server_args)
+
+        _handle_dspark(server_args)
+
+        self.assertEqual(server_args.speculative_draft_model_path, _BUNDLED_MODEL_PATH)
+
+    def test_context_parallel_remains_rejected_outside_pd_prefill(self):
+        """Decode CP must not enter the prefill-only hidden-gather path."""
+        server_args = _make_dspark_server_args(
+            model_path=_BUNDLED_MODEL_PATH, hf_config=_bundled_hf_config()
+        )
+        _enable_pd_prefill_cp(server_args)
+        server_args.disaggregation_mode = "decode"
+
+        with self.assertRaisesRegex(
+            ValueError, "only supported for DeepSeek-V4 PD prefill"
+        ):
+            _handle_dspark(server_args)
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/spec/dspark/test_dspark_kv_inject.py

@@ -0,0 +1,44 @@
+import unittest
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import torch
+
+from sglang.srt.speculative.dspark_components.dspark_kv_inject import (
+    TargetHiddenKvInjector,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class TestTargetHiddenKvInjector(CustomTestCase):
+    def _make_injector(self) -> TargetHiddenKvInjector:
+        draft_model = SimpleNamespace(write_target_hidden_kv=Mock())
+        draft_model_runner = SimpleNamespace(token_to_kv_pool=SimpleNamespace())
+        return TargetHiddenKvInjector(
+            draft_model=draft_model,
+            draft_model_runner=draft_model_runner,
+            model_runner=SimpleNamespace(device="cpu"),
+            device="cpu",
+            verify_num_draft_tokens=2,
+            block_pos_offsets=torch.arange(2),
+        )
+
+    def test_rejects_cp_local_hidden_with_global_indices(self):
+        """A CP-local hidden shard must not be written with global cache indices."""
+        injector = self._make_injector()
+
+        with self.assertRaisesRegex(
+            ValueError, "one hidden row and cache location per position"
+        ):
+            injector.inject_target_hidden(
+                target_hidden=torch.zeros((2, 4)),
+                cache_loc=torch.arange(4),
+                positions=torch.arange(4),
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/spec/dspark/test_dspark_worker_pd_prefill.py

@@ -0,0 +1,26 @@
+import unittest
+from unittest.mock import Mock
+
+from sglang.srt.speculative.dspark_components.dspark_worker_v2 import (
+    DSparkWorkerV2,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class TestDsparkWorkerPdPrefill(CustomTestCase):
+    def test_shared_decode_modules_are_not_attached(self):
+        """PD prefill must not configure the unused Markov/LM-head TP shard."""
+        worker = object.__new__(DSparkWorkerV2)
+        worker._is_pd_prefill = True
+        worker.draft_model = Mock()
+
+        worker._attach_shared_modules()
+
+        worker.draft_model.attach_shared_modules.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()