[kimi k3][pd disagg] support pp prefill + dcp decode with dspark (#40045)

python/sglang/srt/arg_groups/speculative_hook.py

@@ -601,8 +601,10 @@ def _handle_dspark(server_args: ServerArgs) -> None:
             )
 
     if cfg.pp_size != 1:
+        if cfg.disaggregation_mode != "prefill":
             raise ValueError(
-            "Currently DSpark speculative decoding only supports pp_size == 1."
+                "DSpark pipeline parallelism requires PD prefill; "
+                "decode and non-disaggregated serving require pp_size == 1."
             )
 
     if cfg.speculative_draft_model_path is None:

python/sglang/srt/arg_groups/validation_hook.py

@@ -57,7 +57,14 @@ def check_pipeline_parallel_compat(
     assert cfg.disable_overlap_schedule, (
         "Pipeline parallelism is not compatible with overlap schedule"
     )
-    if cfg.speculative_algorithm is not None:
+    if cfg.speculative_algorithm == "DSPARK":
+        assert cfg.disaggregation_mode == "prefill", (
+            "Pipeline parallel DSPARK requires disaggregation-mode=prefill"
+        )
+        assert not envs.SGLANG_ENABLE_PP_SPEC.get(), (
+            "SGLANG_ENABLE_PP_SPEC does not support DSPARK PD prefill"
+        )
+    elif cfg.speculative_algorithm is not None:
         assert (
             cfg.speculative_algorithm.upper() == "EAGLE"
             and not cfg.enable_multi_layer_eagle

python/sglang/srt/disaggregation/mooncake/conn.py

@@ -154,6 +154,7 @@ class KVArgsRegisterInfo:
     dst_dcp_rank: int = 0
     requires_dcp_relayout: bool = False
     dcp_token_item_lens: Optional[List[int]] = None
+    dst_kv_item_lens: List[int] = dataclasses.field(default_factory=list)
     staging_base_ptr: int = 0
     staging_total_size: int = 0
     staging: Optional[StagingRegisterInfo] = None
@@ -201,6 +202,11 @@ class KVArgsRegisterInfo:
             dst_dcp_rank=(
                 int(msg[17].decode("ascii")) if len(msg) > 17 and msg[17] != b"" else 0
             ),
+            dst_kv_item_lens=(
+                list(struct.unpack(f"{len(msg[19]) // 8}Q", msg[19]))
+                if len(msg) > 19 and msg[19]
+                else []
+            ),
             # Note: always put the staging field at the final
             staging=StagingRegisterInfo.from_zmq_fields(msg, 14, slot_ids_index=18),
         )
@@ -1090,11 +1096,16 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
         executor: concurrent.futures.ThreadPoolExecutor,
         dst_layer_ids: List[int],
         pack_buffer=None,
+        dst_kv_item_lens: Optional[List[int]] = None,
+        dst_tp_rank: int = 0,
+        dst_attn_tp_size: Optional[int] = None,
     ) -> int:
         if num_kv_tokens is None:
             raise ValueError("PD DCP transfer requires num_kv_tokens")
         physical_page_size = self.kv_args.page_size
 
+        if dst_kv_item_lens and len(dst_kv_item_lens) != len(dst_kv_ptrs):
+            raise ValueError("PD DCP destination KV lengths must match its buffers")
         src_layer_ids = self.kv_args.kv_layer_ids
         if src_layer_ids or dst_layer_ids:
             dst_indices = resolve_dcp_dst_entry_indices(
@@ -1105,11 +1116,17 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
             )
             src_kv_ptrs = self.kv_args.kv_data_ptrs
             dst_kv_ptrs = [dst_kv_ptrs[j] for j in dst_indices]
+            if dst_kv_item_lens:
+                dst_kv_item_lens = [dst_kv_item_lens[j] for j in dst_indices]
         else:
             src_kv_ptrs, dst_kv_ptrs, _ = self.get_mla_kv_ptrs_with_pp(
                 self.kv_args.kv_data_ptrs,
                 dst_kv_ptrs,
             )
+            if dst_kv_item_lens:
+                _, dst_kv_item_lens, _ = self.get_mla_kv_ptrs_with_pp(
+                    self.kv_args.kv_item_lens, dst_kv_item_lens
+                )
         num_draft = self.kv_args.num_draft_entries
         num_target = len(src_kv_ptrs) - num_draft
 
@@ -1155,20 +1172,87 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
                 )
                 for entry in range(num_target)
             ]
+        sliced_draft_params = []
         if num_draft > 0 and plan.draft_src_token_indices.size:
+            if not dst_kv_item_lens and dst_attn_tp_size not in (
+                None,
+                self.attn_tp_size,
+            ):
+                raise ValueError(
+                    "PD DCP with different draft TP sizes requires destination KV lengths"
+                )
             draft_groups = group_concurrent_contiguous(
                 plan.draft_src_token_indices,
                 plan.draft_dst_token_indices,
             )
-            layers_params += [
+            for entry in range(num_target, num_target + num_draft):
+                src_width = dcp_token_item_lens[entry]
+                dst_width = src_width
+                if dst_kv_item_lens:
+                    dst_width, remainder = divmod(
+                        dst_kv_item_lens[entry], physical_page_size * dst_dcp_size
+                    )
+                    if remainder or dst_width <= 0:
+                        raise ValueError("Invalid PD DCP draft destination token width")
+                if src_width == dst_width:
+                    layers_params.append(
                         (
-                    src_kv_ptrs[num_target + entry],
+                            src_kv_ptrs[entry],
-                    dst_kv_ptrs[num_target + entry],
+                            dst_kv_ptrs[entry],
-                    dcp_token_item_lens[num_target + entry],
+                            src_width,
                             draft_groups,
                         )
-                for entry in range(num_draft)
+                    )
-            ]
+                    continue
+                if self.is_mla_backend:
+                    raise ValueError(
+                        "PD DCP draft head slicing is unsupported for pure MLA: "
+                        "dummy prefill senders may omit draft head shards"
+                    )
+                copy_width = min(src_width, dst_width)
+                if max(src_width, dst_width) % copy_width:
+                    raise ValueError("PD DCP draft KV head shards must divide evenly")
+                if dst_attn_tp_size is None:
+                    raise ValueError(
+                        "PD DCP draft head slicing requires destination TP size"
+                    )
+                src_span = src_width * self.attn_tp_size
+                dst_span = dst_width * dst_attn_tp_size
+                src_rank = (self.kv_args.engine_rank % self.attn_tp_size) // max(
+                    1, src_span // dst_span
+                )
+                dst_rank = dst_tp_rank // max(1, dst_span // src_span)
+                src_offset = (dst_rank * dst_width) % src_width
+                dst_offset = (src_rank * src_width) % dst_width
+                sliced_draft_params.append(
+                    (
+                        src_kv_ptrs[entry] + src_offset,
+                        dst_kv_ptrs[entry] + dst_offset,
+                        src_width,
+                        dst_width,
+                        copy_width,
+                    )
+                )
+
+        def process_sliced_draft(params) -> int:
+            batch_size = self.max_transfer_batch_indices
+            if batch_size <= 0:
+                batch_size = 4096
+            for start in range(0, plan.draft_src_token_indices.size, batch_size):
+                src_indices = plan.draft_src_token_indices[start : start + batch_size]
+                dst_indices = plan.draft_dst_token_indices[start : start + batch_size]
+                blocks = []
+                for src_ptr, dst_ptr, src_width, dst_width, copy_width in params:
+                    src_addrs = src_ptr + src_indices * src_width
+                    dst_addrs = dst_ptr + dst_indices * dst_width
+                    blocks.extend(
+                        (int(src), int(dst), copy_width)
+                        for src, dst in zip(src_addrs, dst_addrs)
+                    )
+                ret = self._transfer_data(mooncake_session_id, blocks)
+                if ret != 0:
+                    return ret
+            return 0
 
         def set_transfer_blocks(
             src_ptr: int, dst_ptr: int, token_item_len: int, groups
@@ -1196,12 +1280,19 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
                 executor.submit(process_layer, *layer_params)
                 for layer_params in layers_params
             ]
+            futures.extend(
+                executor.submit(process_sliced_draft, [params])
+                for params in sliced_draft_params
+            )
             return self._await_transfer_futures(futures)
 
         transfer_blocks = []
         for layer_params in layers_params:
             transfer_blocks.extend(set_transfer_blocks(*layer_params))
-        return self._transfer_data(mooncake_session_id, transfer_blocks)
+        ret = self._transfer_data(mooncake_session_id, transfer_blocks)
+        if ret != 0 or not sliced_draft_params:
+            return ret
+        return process_sliced_draft(sliced_draft_params)
 
     def send_kvcache_slice(
         self,
@@ -2114,6 +2205,9 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
                                     target_rank_registration_info.dst_kv_layer_ids
                                 ),
                                 pack_buffer=pack_buffer,
+                                dst_kv_item_lens=target_rank_registration_info.dst_kv_item_lens,
+                                dst_tp_rank=target_rank_registration_info.dst_tp_rank,
+                                dst_attn_tp_size=target_rank_registration_info.dst_attn_tp_size,
                             )
                         elif (
                             self.is_mla_backend
@@ -2810,6 +2904,10 @@ class MooncakeKVReceiver(MooncakeFailureExceptionMixin, CommonKVReceiver):
                             dst_dcp_size,
                             dst_dcp_rank,
                             packed_staging_slot_layer_ids,
+                            struct.pack(
+                                f"{len(self.kv_mgr.kv_args.kv_item_lens)}Q",
+                                *self.kv_mgr.kv_args.kv_item_lens,
+                            ),
                         ]
                     )
             except zmq.ZMQError:

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -886,7 +886,11 @@ class SchedulerPPMixin:
         # Draft extend runs only on the last stage, but every rank needs its relayed
         # output to fill PD auxiliary buffers.
         draft_input = result.next_draft_input
-        if draft_input is not None and draft_input.topk_p is not None:
+        if (
+            draft_input is not None
+            and not batch.spec_algorithm.is_dspark()
+            and draft_input.topk_p is not None
+        ):
             tensor_dict["draft_topk_p"] = draft_input.topk_p.contiguous()
             tensor_dict["draft_topk_index"] = draft_input.topk_index.contiguous()
             tensor_dict["draft_hidden_states"] = draft_input.hidden_states.contiguous()
@@ -1138,6 +1142,16 @@ class SchedulerPPMixin:
                 dsa_topk_indices=pp_outputs.tensors.get("draft_dsa_topk_indices"),
             )
             batch.spec_info = next_draft_input
+        elif batch.spec_algorithm.is_dspark():
+            from sglang.srt.speculative.dspark_components.dspark_draft import (
+                make_next_draft_input,
+            )
+
+            next_draft_input = make_next_draft_input(
+                bonus_tokens=next_token_ids,
+                new_seq_lens=batch.seq_lens,
+            )
+            batch.spec_info = next_draft_input
 
         if self._pp_spec_relay:
             # Gated single-instance PP+spec: the sampled first token roots

python/sglang/srt/model_executor/model_runner.py

@@ -799,6 +799,13 @@ class ModelRunner:
 
             enable_batch_invariant_mode()
 
+    def get_pp_proxy_dspark_hidden_size(self) -> int:
+        return misc_utils.resolve_pp_proxy_dspark_hidden_size(
+            model=self.model,
+            pp_size=self.ps.pp_size,
+            pp_rank=self.ps.pp_rank,
+        )
+
     def get_pp_proxy_topk_size(self) -> Optional[int]:
         return misc_utils.resolve_pp_proxy_topk_size(
             model_config=self.model_config,

python/sglang/srt/model_executor/model_runner_components/misc_utils.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Optional, Protocol, runtime_checkable
 
 from sglang.srt.configs.model_config import (
     dsa_layer_skips_topk,
@@ -90,3 +90,18 @@ def resolve_pp_proxy_residual_num_blocks(
     if block_size is None:
         return None
     return (start_layer + block_size - 1) // block_size
+
+
+@runtime_checkable
+class _SupportsDSparkPPProxy(Protocol):
+    def get_pp_proxy_dspark_hidden_size(self) -> int: ...
+
+
+def resolve_pp_proxy_dspark_hidden_size(
+    *, model: Any, pp_size: int, pp_rank: int
+) -> int:
+    if pp_size <= 1 or pp_rank == 0:
+        return 0
+    if isinstance(model, _SupportsDSparkPPProxy):
+        return model.get_pp_proxy_dspark_hidden_size()
+    return 0

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

@@ -418,6 +418,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             pp_proxy_residual_num_blocks=(
                 self.model_runner.get_pp_proxy_residual_num_blocks()
             ),
+            pp_proxy_dspark_hidden_size=(
+                self.model_runner.get_pp_proxy_dspark_hidden_size()
+            ),
         )
         self.buffers.share_buffers()
         # FB-shared slot registry adopting DecodeInputBuffers storage (same

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

@@ -389,6 +389,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             pp_proxy_residual_num_blocks=(
                 self.model_runner.get_pp_proxy_residual_num_blocks()
             ),
+            pp_proxy_dspark_hidden_size=(
+                self.model_runner.get_pp_proxy_dspark_hidden_size()
+            ),
         )
         self.buffers.share_buffers()
         # Token-axis FB-shared slot registry adopting PrefillInputBuffers
@@ -599,8 +602,13 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                     f"unsupported for this model architecture."
                 ) from exc
             params = list(inspect.signature(self.layer_model.forward).parameters)
-            self._input_embeds_arg_idx = (
+            self._input_embeds_arg_idx = next(
-                params.index("input_embeds") if "input_embeds" in params else None
+                (
+                    params.index(name)
+                    for name in ("input_embeds", "inputs_embeds")
+                    if name in params
+                ),
+                None,
             )
 
         # --- aiter chip info pre-warming (AMD) -------------------------
@@ -1930,6 +1938,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         """A text-only batch would otherwise replay the captured input_embeds."""
         ie_idx = self._input_embeds_arg_idx
         ie = layer_kwargs.get("input_embeds")
+        if ie is None:
+            ie = layer_kwargs.get("inputs_embeds")
         if ie is None and ie_idx is not None and len(args) > ie_idx:
             ie = args[ie_idx]
         if ie is None:
@@ -1968,7 +1978,10 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             # text-only batches they are get_input_embeddings()(input_ids).
             # Copy them into the slot before replay so the graph sees the
             # current request's embeddings (mirrors main's BCG closure).
-            if self.buffer_registry.has_slot("input_embeds"):
+            if (
+                self.model_runner.pp_group.is_first_rank
+                and self.buffer_registry.has_slot("input_embeds")
+            ):
                 self._fill_input_embeds_slot(args, layer_kwargs, static_num_tokens)
             hs = self.backend.replay(shape_key, static_forward_batch, **kwargs)
             return _slice_output_rows(hs, raw_num_tokens) if full_path else hs

python/sglang/srt/model_executor/runner_utils/buffers.py

@@ -67,6 +67,7 @@ def _allocate_pp_proxy_tensors(
     hc_hidden_size: Optional[int] = None,
     pp_proxy_topk_size: Optional[int] = None,
     pp_proxy_residual_num_blocks: Optional[int] = None,
+    pp_proxy_dspark_hidden_size: int = 0,
 ) -> Dict[str, torch.Tensor]:
     """Allocate the stable buffers consumed by an incoming PP proxy."""
     is_mhc = hc_hidden_size is not None
@@ -87,6 +88,10 @@ def _allocate_pp_proxy_tensors(
         pp_proxy_tensors["topk_indices"] = torch.zeros(
             (max_num_tokens, pp_proxy_topk_size), dtype=torch.int32
         )
+    if pp_proxy_dspark_hidden_size:
+        pp_proxy_tensors["dspark_hidden_states"] = torch.zeros(
+            (max_num_tokens, pp_proxy_dspark_hidden_size), dtype=dtype
+        )
     return pp_proxy_tensors
 
 
@@ -136,6 +141,7 @@ class DecodeInputBuffers(ForwardInputBuffers):
         hc_hidden_size: Optional[int] = None,
         pp_proxy_topk_size: Optional[int] = None,
         pp_proxy_residual_num_blocks: Optional[int] = None,
+        pp_proxy_dspark_hidden_size: int = 0,
     ) -> DecodeInputBuffers:
         with torch.device(device):
             input_ids = torch.zeros((max_num_token,), dtype=torch.int64)
@@ -173,6 +179,7 @@ class DecodeInputBuffers(ForwardInputBuffers):
                     hc_hidden_size=hc_hidden_size,
                     pp_proxy_topk_size=pp_proxy_topk_size,
                     pp_proxy_residual_num_blocks=pp_proxy_residual_num_blocks,
+                    pp_proxy_dspark_hidden_size=pp_proxy_dspark_hidden_size,
                 )
                 if pp_size > 1
                 else None
@@ -275,6 +282,7 @@ class PrefillInputBuffers(ForwardInputBuffers):
         hc_hidden_size: Optional[int] = None,
         pp_proxy_topk_size: Optional[int] = None,
         pp_proxy_residual_num_blocks: Optional[int] = None,
+        pp_proxy_dspark_hidden_size: int = 0,
     ) -> PrefillInputBuffers:
         with torch.device(device):
             input_ids = torch.zeros((max_num_tokens,), dtype=torch.int64)
@@ -311,6 +319,7 @@ class PrefillInputBuffers(ForwardInputBuffers):
                     hc_hidden_size=hc_hidden_size,
                     pp_proxy_topk_size=pp_proxy_topk_size,
                     pp_proxy_residual_num_blocks=pp_proxy_residual_num_blocks,
+                    pp_proxy_dspark_hidden_size=pp_proxy_dspark_hidden_size,
                 )
                 if pp_size > 1 and not is_first_pp_rank
                 else None

python/sglang/srt/models/kimi_k3.py

@@ -3047,6 +3047,18 @@ class KimiK3LinearModel(nn.Module):
         )
         sp_sharded = False
         aux_hidden_states = []
+        if (
+            self.dspark_layers_to_capture is not None
+            and not self.pp_group.is_first_rank
+        ):
+            if "dspark_hidden_states" in pp_proxy_tensors.tensors:
+                aux_hidden_states.append(pp_proxy_tensors["dspark_hidden_states"])
+            if self.start_layer - 1 in self.dspark_layers_to_capture:
+                aux_hidden_states.append(
+                    self._dspark_capture_stream(
+                        self.start_layer - 1, hidden_states, residual, attn_res
+                    )
+                )
         for i in range(self.start_layer, self.end_layer):
             if sp_sharded and not self.layers[i]._sp_moe:
                 hidden_states = _sp_all_gather_rows(hidden_states)
@@ -3065,6 +3077,7 @@ class KimiK3LinearModel(nn.Module):
             if (
                 self.dspark_layers_to_capture is not None
                 and i in self.dspark_layers_to_capture
+                and (i + 1 < self.end_layer or self.pp_group.is_last_rank)
             ):
                 aux_hidden_states.append(
                     self._dspark_capture_stream(i, hidden_states, residual, attn_res)
@@ -3078,9 +3091,12 @@ class KimiK3LinearModel(nn.Module):
                     # full stream head (bit-identical to the fused fold).
                     hidden_states = residual + hidden_states
                 residual = attn_res.block_residual  # raw bank across ranks
-            return PPProxyTensors(
+            proxy_tensors = {"hidden_states": hidden_states, "residual": residual}
-                {"hidden_states": hidden_states, "residual": residual}
+            if aux_hidden_states:
+                proxy_tensors["dspark_hidden_states"] = torch.cat(
+                    aux_hidden_states, dim=-1
                 )
+            return PPProxyTensors(proxy_tensors)
 
         if hidden_states.shape[0] != 0:
             if attn_res is not None:
@@ -3204,13 +3220,13 @@ class KimiK3LinearForCausalLM(nn.Module):
     def get_input_embeddings(self):
         return self.model.embed_tokens
 
+    def get_pp_proxy_dspark_hidden_size(self) -> int:
+        layers = self.model.dspark_layers_to_capture or []
+        return self.config.hidden_size * sum(
+            layer < self.model.start_layer - 1 for layer in layers
+        )
+
     def set_dspark_layers_to_capture(self, layer_ids: list[int]) -> None:
-        if self.pp_group.world_size > 1:
-            # Capture layers living on non-last PP ranks would be silently
-            # skipped (the flag is only set on the last rank).
-            raise NotImplementedError("DSPARK aux hidden capture requires PP=1.")
-        if not self.pp_group.is_last_rank:
-            return
         if layer_ids is None:
             raise ValueError(
                 "DSPARK requires explicit layer_ids for aux hidden capture."
@@ -3667,6 +3683,11 @@ class KimiK3ForConditionalGeneration(nn.Module):
             raise AttributeError("lm_head is not available in encoder-only mode")
         return self.language_model.lm_head
 
+    def get_pp_proxy_dspark_hidden_size(self) -> int:
+        if self.language_model is None:
+            return 0
+        return self.language_model.get_pp_proxy_dspark_hidden_size()
+
     def set_dspark_layers_to_capture(self, layer_ids: list[int]) -> None:
         if self.language_model is None:
             raise AttributeError(

python/sglang/srt/models/kimi_linear.py

@@ -717,6 +717,12 @@ class KimiLinearModel(nn.Module):
             device=device,
         )
         aux_hidden_states = []
+        if (
+            self.dspark_layers_to_capture is not None
+            and not self.pp_group.is_first_rank
+            and "dspark_hidden_states" in pp_proxy_tensors.tensors
+        ):
+            aux_hidden_states.append(pp_proxy_tensors["dspark_hidden_states"])
         for i in range(self.start_layer, self.end_layer):
             ctx = get_global_expert_distribution_recorder().with_current_layer(i)
             with ctx:
@@ -737,12 +743,12 @@ class KimiLinearModel(nn.Module):
                 )
 
         if not self.pp_group.is_last_rank:
-            return PPProxyTensors(
+            proxy_tensors = {"hidden_states": hidden_states, "residual": residual}
-                {
+            if aux_hidden_states:
-                    "hidden_states": hidden_states,
+                proxy_tensors["dspark_hidden_states"] = torch.cat(
-                    "residual": residual,
+                    aux_hidden_states, dim=-1
-                }
                 )
+            return PPProxyTensors(proxy_tensors)
         else:
             if hidden_states.shape[0] != 0:
                 if residual is None:
@@ -787,11 +793,13 @@ class KimiLinearForCausalLM(nn.Module):
     def get_input_embeddings(self):
         return self.model.embed_tokens
 
+    def get_pp_proxy_dspark_hidden_size(self) -> int:
+        layers = self.model.dspark_layers_to_capture or []
+        return self.config.hidden_size * sum(
+            layer < self.model.start_layer for layer in layers
+        )
+
     def set_dspark_layers_to_capture(self, layer_ids: list[int]) -> None:
-        if self.pp_group.world_size > 1:
-            raise NotImplementedError("DSPARK aux hidden capture requires PP=1.")
-        if not self.pp_group.is_last_rank:
-            return
         if layer_ids is None:
             raise ValueError(
                 "DSPARK requires explicit layer_ids for aux hidden capture."

python/sglang/srt/speculative/draft_worker_common.py

@@ -69,6 +69,7 @@ def build_draft_tp_worker(
     algo_label: str,
     attention_backend_override: Optional[str] = None,
     draft_worker_cls: type[TpModelWorker] = TpModelWorker,
+    random_seed: Optional[int] = None,
 ) -> DraftWorkerBundle:
     # An override names a draft-specific backend the caller has already
     # validated (e.g. a self-drafting architecture); it skips the generic
@@ -90,6 +91,7 @@ def build_draft_tp_worker(
             ps=ps,
             nccl_port=nccl_port,
             is_draft_worker=True,
+            random_seed=random_seed,
             # The draft runs at absolute target positions.
             context_length=target_model_config.context_len,
             draft_attention_backend=draft_backend,

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

@@ -9,6 +9,7 @@ from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
     is_unified_kv_triton,
 )
 from sglang.srt.configs.hybrid_arch import mambaish_config
+from sglang.srt.distributed import get_pp_group
 from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.environ import envs
 from sglang.srt.layers.logprob_processor import compute_spec_logprobs
@@ -20,6 +21,7 @@ from sglang.srt.model_executor.cuda_graph_config import Backend
 from sglang.srt.model_executor.forward_batch_info import (
     CaptureHiddenMode,
     ForwardMode,
+    PPProxyTensors,
     compute_position,
 )
 from sglang.srt.runtime_context import (
@@ -126,6 +128,8 @@ def _configure_target_hidden_projection(
 
 
 class DSparkWorkerV2(BaseSpecWorker):
+    """Non-last PP stages run only the target; draft state belongs to the last stage."""
+
     def __init__(
         self,
         server_args: ServerArgs,
@@ -145,6 +149,10 @@ class DSparkWorkerV2(BaseSpecWorker):
         self.model_runner = target_worker.model_runner
         self.page_size = get_schedule().page_size
         self.device = target_worker.device
+        self._draft_worker = None
+        self._hosts_draft = get_pp_group().is_last_rank
+        if not self._hosts_draft:
+            return
 
         self._draft_is_moe = draft_is_deepseek_v4()
         self._draft_dp_context_enabled = (
@@ -178,6 +186,7 @@ class DSparkWorkerV2(BaseSpecWorker):
                     DSV4_DRAFT_ATTENTION_BACKEND if self._draft_is_moe else None
                 ),
                 draft_worker_cls=draft_worker_cls,
+                random_seed=target_worker.random_seed,
             )
         self._draft_worker = bundle.draft_worker
         self.draft_model_runner = bundle.draft_model_runner
@@ -397,10 +406,12 @@ class DSparkWorkerV2(BaseSpecWorker):
 
     @property
     def carries_confidence(self) -> bool:
-        return self._verify_planner.carries_confidence
+        return self._hosts_draft and self._verify_planner.carries_confidence
 
     @property
     def spec_v2_attn_backends(self) -> tuple:
+        if not self._hosts_draft:
+            return super().spec_v2_attn_backends
         return (
             self._target_worker.model_runner.attn_backend,
             self.draft_model_runner.attn_backend,
@@ -422,6 +433,8 @@ class DSparkWorkerV2(BaseSpecWorker):
         req_to_token_pool=None,
         token_to_kv_pool_allocator=None,
     ):
+        if not self._hosts_draft:
+            return
         self._draft_worker.alloc_memory_pool(
             memory_pool_config=memory_pool_config,
             req_to_token_pool=req_to_token_pool,
@@ -429,6 +442,8 @@ class DSparkWorkerV2(BaseSpecWorker):
         )
 
     def init_attention_backends(self):
+        if not self._hosts_draft:
+            return
         with draft_pp_context(), self._draft_context():
             self._draft_worker.init_attention_backends()
         self._target_hidden_projection_enabled = _configure_target_hidden_projection(
@@ -449,6 +464,8 @@ class DSparkWorkerV2(BaseSpecWorker):
         )
 
     def init_cuda_graphs(self):
+        if not self._hosts_draft:
+            return
         capture_decode_cuda_graph = self._decode_graph_allowed
         available_mem = self._tp_sync.available_memory_gb(
             SpecTpSyncSite.DSPARK_MEM,
@@ -510,19 +527,29 @@ class DSparkWorkerV2(BaseSpecWorker):
         pass
 
     def set_dspark_forced_budget_frac(self, frac: Optional[float]) -> None:
+        if not self._hosts_draft:
+            return
         self._forced_budget_frac = frac
         self._verify_planner.set_forced_budget_frac(frac)
 
     def dump_info_records(self) -> Optional[dict]:
+        if not self._hosts_draft:
+            return None
         return self._observers.dump_info_records()
 
     def clear_info_records(self) -> None:
+        if not self._hosts_draft:
+            return
         self._observers.clear_info_records()
 
     def block_accept_estimate_log_suffix(self) -> Optional[str]:
+        if not self._hosts_draft:
+            return None
         return self._observers.block_accept_estimate_log_suffix()
 
     def note_request_finished(self, *, rid: str, natural_stop: bool) -> None:
+        if not self._hosts_draft:
+            return
         self._observers.note_request_finished(rid=rid, natural_stop=natural_stop)
 
     def forward_batch_generation(
@@ -531,19 +558,30 @@ class DSparkWorkerV2(BaseSpecWorker):
         on_publish=None,
         grammar_barrier=None,
         *,
-        pp_proxy_tensors=None,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
     ) -> GenerationBatchResult:
-        # The non-overlap scheduler passes this keyword even when PP=1.
+        if not self._hosts_draft:
-        assert pp_proxy_tensors is None, "DSpark does not support pipeline parallelism"
+            batch_output = self.target_worker.forward_batch_generation(
+                batch,
+                pp_proxy_tensors=pp_proxy_tensors,
+                capture_hidden_mode=CaptureHiddenMode.FULL,
+            )
+            batch_output.new_seq_lens = batch.seq_lens
+            if on_publish is not None:
+                on_publish(batch_output.new_seq_lens)
+            return batch_output
         if batch.forward_mode.is_extend() or batch.is_extend_in_batch:
             self._verify_planner.note_non_decode_step()
             self._observers.note_prefill_step()
-            return self._forward_prefill(batch, on_publish)
+            return self._forward_prefill(batch, on_publish, pp_proxy_tensors)
 
         return self._forward_decode(batch, on_publish, grammar_barrier)
 
     def _forward_prefill(
-        self, batch: ScheduleBatch, on_publish
+        self,
+        batch: ScheduleBatch,
+        on_publish,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
     ) -> GenerationBatchResult:
         if batch.forward_mode.is_idle():
             if get_parallel().enable_dp_attention:
@@ -553,7 +591,9 @@ class DSparkWorkerV2(BaseSpecWorker):
             return self._decode_idle_result(on_publish=on_publish)
 
         batch_output = self.target_worker.forward_batch_generation(
-            batch, capture_hidden_mode=CaptureHiddenMode.FULL
+            batch,
+            pp_proxy_tensors=pp_proxy_tensors,
+            capture_hidden_mode=CaptureHiddenMode.FULL,
         )
         # BCG replay skips model-side Python, so re-evaluate the same pure predicate.
         target_hidden_is_projected = (
@@ -1010,4 +1050,6 @@ class DSparkWorkerV2(BaseSpecWorker):
         )
 
     def get_confidence_budget_prepare(self):
+        if not self._hosts_draft:
+            return None
         return self._verify_planner.confidence_budget_prepare()

test/registered/unit/disaggregation/test_disaggregation_wire.py

@@ -96,6 +96,9 @@ class TestDisaggregationWire(unittest.TestCase):
         self.assertEqual(info.staging_total_size, 4096)
         self.assertEqual(info.dst_dcp_size, 4)
         self.assertEqual(info.dst_dcp_rank, 2)
+        self.assertEqual(info.dst_kv_item_lens, [])
+        info = KVArgsRegisterInfo.from_zmq(msg + [b"", struct.pack("Q", 128)])
+        self.assertEqual(info.dst_kv_item_lens, [128])
 
     def test_int_lists_roundtrip(self):
         cases = [

test/registered/unit/disaggregation/test_mooncake_transfer_batching.py

@@ -130,5 +130,161 @@ class TestMooncakeTransferBatching(unittest.TestCase):
         )
 
 
+class TestDcpDraftHeadTransfer(unittest.TestCase):
+    def test_transfers_draft_heads_to_logical_destination_rows(self):
+        for src_tp, dst_tp in ((4, 8), (8, 4), (8, 8), (4, 32), (32, 4)):
+            for custom_pool in (False, True):
+                for batch_size in (0, 37):
+                    with self.subTest(
+                        src_tp=src_tp,
+                        dst_tp=dst_tp,
+                        custom_pool=custom_pool,
+                        batch_size=batch_size,
+                    ):
+                        self._check_transfer(src_tp, dst_tp, custom_pool, batch_size)
+
+    def test_rejects_pure_mla_with_unequal_draft_head_widths(self):
+        for src_tp, dst_tp in ((4, 8), (8, 4)):
+            with self.subTest(src_tp=src_tp, dst_tp=dst_tp):
+                with self.assertRaisesRegex(ValueError, "dummy prefill senders"):
+                    self._check_transfer(src_tp, dst_tp, False, 37, pure_mla=True)
+
+    def test_sliced_draft_stops_after_failed_batch(self):
+        self._check_transfer(4, 8, False, 37, fail_draft=True)
+
+    def _check_transfer(
+        self, src_tp, dst_tp, custom_pool, batch_size, fail_draft=False, pure_mla=False
+    ):
+        page_size, tokens, heads, head_bytes = 64, 249, 16, 4
+        src_width, dst_width = (
+            max(1, heads // src_tp) * head_bytes,
+            max(1, heads // dst_tp) * head_bytes,
+        )
+        src_pages = np.array([1, 3, 4, 7], dtype=np.int32)
+        logical = np.arange(tokens)
+        src_rows = src_pages[logical // page_size] * page_size + logical % page_size
+        expected = (
+            np.arange(tokens * heads * head_bytes, dtype=np.int64)
+            .reshape(tokens, heads, head_bytes)
+            .astype(np.uint8)
+        )
+        for dst_rank in range(dst_tp):
+            dst_buffers = {
+                base: np.zeros(16384 * max(8, dst_width), dtype=np.uint8)
+                for base in (1000000, 2000000, 3000000, 4000000)
+            }
+            source_ranks = (
+                range(dst_rank * src_tp // dst_tp, (dst_rank + 1) * src_tp // dst_tp)
+                if src_tp >= dst_tp
+                else [dst_rank * src_tp // dst_tp]
+            )
+            for src_rank in source_ranks:
+                src_head_start = (src_rank // max(1, src_tp // heads)) * max(
+                    1, heads // src_tp
+                )
+                source = np.zeros(1024 * src_width, dtype=np.uint8)
+                source.reshape(-1, src_width)[src_rows] = expected[
+                    :, src_head_start : src_head_start + max(1, heads // src_tp)
+                ].reshape(tokens, src_width)
+                target = np.zeros(1024 * 8, dtype=np.uint8)
+                target.reshape(-1, 8)[src_rows] = (
+                    np.arange(tokens * 8).reshape(tokens, 8).astype(np.uint8)
+                )
+                src_buffers = {10000: target, 100000: source, 200000: source}
+
+                failed_batches = []
+
+                def transfer(
+                    session, blocks, src_buffers=src_buffers, dst_buffers=dst_buffers
+                ):
+                    draft_blocks = [block for block in blocks if block[1] >= 3000000]
+                    if fail_draft and draft_blocks:
+                        failed_batches.append(draft_blocks)
+                        return 17
+                    if batch_size and src_width != dst_width:
+                        self.assertLessEqual(
+                            len(draft_blocks), batch_size * (1 if custom_pool else 2)
+                        )
+                    for src, dst, size in blocks:
+                        src_base = max(base for base in src_buffers if base <= src)
+                        dst_base = max(base for base in dst_buffers if base <= dst)
+                        dst_buffers[dst_base][
+                            dst - dst_base : dst - dst_base + size
+                        ] = src_buffers[src_base][
+                            src - src_base : src - src_base + size
+                        ]
+                    return 0
+
+                manager = SimpleNamespace(
+                    is_mla_backend=pure_mla,
+                    kv_args=SimpleNamespace(
+                        page_size=page_size,
+                        kv_layer_ids=[47, 93, 93],
+                        kv_data_ptrs=[10000, 100000, 200000],
+                        num_draft_entries=2,
+                        engine_rank=src_rank + 2 * src_tp,
+                    ),
+                    attn_tp_size=src_tp,
+                    max_transfer_batch_indices=batch_size,
+                    enable_custom_mem_pool=custom_pool,
+                    _transfer_data=transfer,
+                    _await_transfer_futures=lambda futures: max(
+                        f.result() for f in futures
+                    ),
+                )
+                with concurrent.futures.ThreadPoolExecutor() as executor:
+                    result = MooncakeKVManager.send_kvcache_dcp(
+                        manager,
+                        "session",
+                        src_pages,
+                        [1000000, 2000000, 3000000, 4000000],
+                        np.array([2], dtype=np.int32),
+                        dcp_token_item_lens=[8, src_width, src_width],
+                        dst_dcp_size=dst_tp,
+                        dst_dcp_rank=dst_rank,
+                        src_page_offset=0,
+                        decode_prefix_len=0,
+                        num_kv_tokens=tokens,
+                        executor=executor,
+                        dst_layer_ids=[3, 47, 93, 93],
+                        dst_kv_item_lens=[
+                            page_size * 8,
+                            page_size * 8,
+                            page_size * dst_tp * dst_width,
+                            page_size * dst_tp * dst_width,
+                        ],
+                        dst_tp_rank=dst_rank,
+                        dst_attn_tp_size=dst_tp,
+                    )
+                if fail_draft:
+                    self.assertEqual(result, 17)
+                    self.assertEqual(len(failed_batches), 1)
+                    return
+                self.assertEqual(result, 0)
+            dst_head_start = (dst_rank // max(1, dst_tp // heads)) * max(
+                1, heads // dst_tp
+            )
+            for base in (3000000, 4000000):
+                actual = dst_buffers[base].reshape(-1, dst_width)[
+                    2 * page_size * dst_tp + logical
+                ]
+                np.testing.assert_array_equal(
+                    actual,
+                    expected[
+                        :,
+                        dst_head_start : dst_head_start + max(1, heads // dst_tp),
+                    ].reshape(tokens, dst_width),
+                )
+            owned = np.arange(dst_rank, tokens, dst_tp)
+            actual_target = dst_buffers[2000000].reshape(-1, 8)[
+                2 * page_size + owned // dst_tp
+            ]
+            np.testing.assert_array_equal(
+                actual_target,
+                np.arange(tokens * 8).reshape(tokens, 8).astype(np.uint8)[owned],
+            )
+            self.assertFalse(dst_buffers[1000000].any())
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/managers/test_auxiliary_output.py

@@ -660,6 +660,7 @@ def test_pipeline_parallel_auxiliary_output_round_trip():
         next_token_ids=torch.tensor([7]),
     )
     batch = SimpleNamespace(
+        spec_algorithm=SpeculativeAlgorithm.NONE,
         return_logprob=False,
         req_pool_indices=torch.tensor([3]),
         input_ids=torch.tensor([5]),
@@ -705,6 +706,7 @@ def test_pipeline_parallel_dsa_seed_round_trip(dsa_topk_indices):
         next_draft_input=draft_input,
     )
     batch = SimpleNamespace(
+        spec_algorithm=SpeculativeAlgorithm.EAGLE3,
         return_logprob=False,
         req_pool_indices=torch.tensor([3]),
         input_ids=torch.tensor([5]),
@@ -745,6 +747,7 @@ def test_pipeline_parallel_auxiliary_output_stays_packed_before_first_rank():
         next_token_ids=torch.tensor([7]),
     )
     batch = SimpleNamespace(
+        spec_algorithm=SpeculativeAlgorithm.NONE,
         return_logprob=False,
         req_pool_indices=torch.tensor([3]),
         input_ids=torch.tensor([5]),

test/registered/unit/managers/test_pp_cp_rank_offsets.py

@@ -2,8 +2,11 @@ import unittest
 from types import SimpleNamespace
 from unittest.mock import patch
 
+import torch
+
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import (
+    CustomTestCase,
     enter_scope,
     maybe_stub_sgl_kernel,
     published_topology,
@@ -239,5 +242,56 @@ class TestPPCPRankOffsets(unittest.TestCase):
         )
 
 
+class TestDSparkPPOutput(CustomTestCase):
+    def test_output_ring_rebinds_dspark_state_on_each_stage(self):
+        from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
+        from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
+        from sglang.srt.speculative.dspark_components.dspark_draft import (
+            make_next_draft_input,
+        )
+        from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+
+        payloads = []
+        scheduler = SimpleNamespace(
+            _pp_spec_relay=False,
+            pp_group=SimpleNamespace(is_first_rank=False),
+            future_map=SimpleNamespace(
+                stash=lambda indices, value: payloads.append(value)
+            ),
+        )
+        tokens = torch.tensor([13, 29])
+        batch = SimpleNamespace(
+            return_logprob=False,
+            req_pool_indices=torch.tensor([0, 1]),
+            seq_lens=torch.tensor([8, 15]),
+            spec_algorithm=SpeculativeAlgorithm.DSPARK,
+            spec_info=object(),
+        )
+        wire = SchedulerPPMixin._pp_prepare_tensor_dict(
+            scheduler,
+            SimpleNamespace(
+                next_token_ids=tokens,
+                next_draft_input=make_next_draft_input(
+                    bonus_tokens=tokens, new_seq_lens=batch.seq_lens
+                ),
+                logits_output=None,
+            ),
+            batch,
+        )
+        self.assertNotIn("draft_topk_p", wire)
+        result = SchedulerPPMixin._pp_prep_batch_result(
+            scheduler,
+            batch,
+            SimpleNamespace(can_run_cuda_graph=False),
+            PPProxyTensors(wire),
+        )
+        self.assertIsInstance(result.next_draft_input, DFlashDraftInputV2)
+        self.assertIs(batch.spec_info, result.next_draft_input)
+        torch.testing.assert_close(batch.spec_info.bonus_tokens, tokens)
+        torch.testing.assert_close(batch.spec_info.new_seq_lens, batch.seq_lens)
+        torch.testing.assert_close(payloads[0].bonus_tokens, tokens)
+        self.assertEqual(payloads[0].hidden_states.numel(), 0)
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py

@@ -11,6 +11,9 @@ from sglang.srt.model_executor.cuda_graph_buffer_registry import (
     build_prefill_registry,
 )
 from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
+from sglang.srt.model_executor.model_runner_components.misc_utils import (
+    resolve_pp_proxy_dspark_hidden_size,
+)
 from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import (
     PrefillCudaGraphRunner,
     _build_layer_model_forward_kwargs,
@@ -52,6 +55,25 @@ def _make_pp_buffers_and_registry():
 
 
 class TestPrefillCudaGraphRunnerHelpers(CustomTestCase):
+    def test_dspark_proxy_width_requires_receiving_stage_and_model_support(self):
+        class Model:
+            def get_pp_proxy_dspark_hidden_size(self):
+                return 16
+
+        for model, pp_size, pp_rank, expected in (
+            (Model(), 1, 0, 0),
+            (Model(), 2, 0, 0),
+            (Model(), 2, 1, 16),
+            (object(), 2, 1, 0),
+        ):
+            with self.subTest(pp_size=pp_size, pp_rank=pp_rank, expected=expected):
+                self.assertEqual(
+                    resolve_pp_proxy_dspark_hidden_size(
+                        model=model, pp_size=pp_size, pp_rank=pp_rank
+                    ),
+                    expected,
+                )
+
     def test_pp_proxy_stable_buffers_accept_full_and_hidden_only_contracts(self):
         buffers, registry = _make_pp_buffers_and_registry()
         full_proxy = PPProxyTensors(
@@ -170,6 +192,7 @@ class TestPrefillCudaGraphRunnerHelpers(CustomTestCase):
             pp_size=2,
             is_first_pp_rank=False,
             pp_proxy_residual_num_blocks=3,
+            pp_proxy_dspark_hidden_size=16,
         )
 
         self.assertEqual(
@@ -177,9 +200,151 @@ class TestPrefillCudaGraphRunnerHelpers(CustomTestCase):
                 key: tuple(value.shape)
                 for key, value in buffers.pp_proxy_tensors.items()
             },
-            {"hidden_states": (16, 8), "residual": (16, 3, 8)},
+            {
+                "hidden_states": (16, 8),
+                "residual": (16, 3, 8),
+                "dspark_hidden_states": (16, 16),
+            },
         )
 
+    def test_dspark_proxy_width_respects_deferred_k3_boundary_capture(self):
+        from sglang.srt.models.kimi_k3 import (
+            KimiK3ForConditionalGeneration,
+            KimiK3LinearForCausalLM,
+        )
+        from sglang.srt.models.kimi_linear import KimiLinearForCausalLM
+
+        for start, k3_count, linear_count in [
+            (0, 0, 0),
+            (8, 0, 1),
+            (24, 1, 2),
+            (52, 2, 3),
+        ]:
+            model = SimpleNamespace(
+                config=SimpleNamespace(hidden_size=8),
+                model=SimpleNamespace(
+                    start_layer=start, dspark_layers_to_capture=[7, 23, 51]
+                ),
+            )
+            self.assertEqual(
+                KimiK3LinearForCausalLM.get_pp_proxy_dspark_hidden_size(model),
+                k3_count * 8,
+            )
+            self.assertEqual(
+                KimiLinearForCausalLM.get_pp_proxy_dspark_hidden_size(model),
+                linear_count * 8,
+            )
+            model.get_pp_proxy_dspark_hidden_size = lambda: (
+                KimiK3LinearForCausalLM.get_pp_proxy_dspark_hidden_size(model)
+            )
+            self.assertEqual(
+                KimiK3ForConditionalGeneration.get_pp_proxy_dspark_hidden_size(
+                    SimpleNamespace(language_model=model)
+                ),
+                k3_count * 8,
+            )
+
+    def test_body_replay_uses_plural_embeds_and_skips_embedding_on_later_pp_stage(self):
+        for first_rank in (True, False):
+            for positional in (True, False):
+                with self.subTest(first_rank=first_rank, positional=positional):
+                    runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
+                    runner._is_full_backend = False
+                    runner._input_embeds_arg_idx = 3
+                    backing = torch.zeros(4, 8)
+                    supplied = torch.full((4, 8), 7.0) if first_rank else None
+                    runner.buffer_registry = SimpleNamespace(
+                        has_slot=lambda name: name == "input_embeds",
+                        get_slot=lambda _: SimpleNamespace(
+                            slice_for=lambda *args: backing
+                        ),
+                    )
+                    layer = SimpleNamespace(forward=lambda *args, **kwargs: None)
+                    original = layer.forward
+                    runner.layer_model = layer
+                    sentinel = object()
+                    runner.backend = SimpleNamespace(
+                        replay=lambda *args, **kwargs: sentinel
+                    )
+                    runner._prefill_forward_context = lambda *args, **kwargs: (
+                        nullcontext()
+                    )
+
+                    def outer_forward(ids, positions, batch):
+                        if positional:
+                            return layer.forward(None, positions, batch, supplied)
+                        return layer.forward(
+                            None, positions, batch, inputs_embeds=supplied
+                        )
+
+                    runner.model_runner = SimpleNamespace(
+                        pp_group=SimpleNamespace(is_first_rank=first_rank),
+                        model=SimpleNamespace(forward=outer_forward),
+                    )
+                    batch = SimpleNamespace(
+                        input_ids=torch.arange(4),
+                        positions=torch.arange(4),
+                        mm_input_embeds=None,
+                    )
+                    result = runner._execute_body_capture(batch, batch, 4, 4, None)
+                    self.assertIs(result, sentinel)
+                    self.assertIs(layer.forward, original)
+                    if first_rank:
+                        torch.testing.assert_close(backing, supplied)
+                    else:
+                        self.assertEqual(torch.count_nonzero(backing).item(), 0)
+
+    def test_dspark_proxy_replay_updates_features_and_clears_padding(self):
+        buffers = PrefillInputBuffers.create(
+            device=torch.device("cpu"),
+            max_bs=1,
+            max_num_tokens=8,
+            cache_loc_dtype=torch.int64,
+            is_multimodal=False,
+            hidden_size=4,
+            dtype=torch.float32,
+            enable_mamba_track=False,
+            pp_size=2,
+            pp_proxy_dspark_hidden_size=8,
+        )
+        registry = build_prefill_registry(
+            device=torch.device("cpu"),
+            max_bs=1,
+            max_num_token=8,
+            cache_loc_dtype=torch.int64,
+            share_pool=False,
+            source=buffers,
+        )
+        runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
+        runner.buffers = buffers
+        runner.model_runner = SimpleNamespace(
+            pp_group=SimpleNamespace(is_first_rank=False)
+        )
+        captured = runner._capture_pp_proxy_tensors(8)["dspark_hidden_states"]
+        ptr = captured.data_ptr()
+        for count, value in [(7, 2.0), (3, 5.0)]:
+            tokens = torch.arange(count)
+            proxy = PPProxyTensors(
+                {
+                    "hidden_states": torch.zeros(count, 4),
+                    "residual": torch.zeros(count, 4),
+                    "dspark_hidden_states": torch.full((count, 8), value),
+                }
+            )
+            registry.fill_from(
+                SimpleNamespace(
+                    input_ids=tokens, positions=tokens, out_cache_loc=tokens
+                ),
+                raw_bs=1,
+                padded_bs=1,
+                raw_num_tokens=count,
+                padded_num_tokens=8,
+                pp_proxy_tensors=proxy,
+            )
+            self.assertEqual(captured.data_ptr(), ptr)
+            torch.testing.assert_close(captured[:count], proxy["dspark_hidden_states"])
+            self.assertEqual(torch.count_nonzero(captured[count:]).item(), 0)
+
     def test_pipeline_proxy_output_is_supported(self):
         runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
         runner.raw_num_tokens = 3

test/registered/unit/server_args/test_server_args.py

@@ -2501,6 +2501,28 @@ class TestPipelineParallelCompat(CustomTestCase):
     def test_no_speculative_decoding_is_fine(self):
         check_pipeline_parallel_compat(self._cfg())
 
+    def test_dspark_pd_prefill_does_not_require_eagle_architecture(self):
+        check_pipeline_parallel_compat(self._cfg(speculative_algorithm="DSPARK"))
+
+    def test_dspark_is_rejected_outside_pd_prefill(self):
+        for mode in ("decode", "null"):
+            with self.subTest(mode=mode):
+                with self.assertRaisesRegex(AssertionError, "DSPARK.*prefill"):
+                    check_pipeline_parallel_compat(
+                        self._cfg(
+                            speculative_algorithm="DSPARK", disaggregation_mode=mode
+                        )
+                    )
+
+    def test_dspark_rejects_eagle_pp_relay(self):
+        with patch.object(
+            validation_hook.envs.SGLANG_ENABLE_PP_SPEC, "get", return_value=True
+        ):
+            with self.assertRaisesRegex(AssertionError, "SGLANG_ENABLE_PP_SPEC"):
+                check_pipeline_parallel_compat(
+                    self._cfg(speculative_algorithm="DSPARK")
+                )
+
     def test_eagle_is_allowed_on_prefill(self):
         check_pipeline_parallel_compat(
             self._cfg(speculative_algorithm="EAGLE"),

test/registered/unit/spec/test_dspark_target_hidden_projection.py

@@ -9,6 +9,7 @@ from sglang.srt.models.dspark import DSparkDraftMixin
 from sglang.srt.speculative.dspark_components.dspark_kv_inject import (
     TargetHiddenKvInjector,
 )
+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
 
@@ -37,6 +38,40 @@ class _Attention:
 
 
 class DSparkTargetHiddenProjectionTest(CustomTestCase):
+    def test_nonfinal_prefill_stage_only_forwards_target_proxies(self) -> None:
+        proxy = object()
+        target = SimpleNamespace(
+            model_runner=SimpleNamespace(attn_backend=object(), spec_algorithm=None),
+            device="cpu",
+            forward_batch_generation=lambda batch, *, pp_proxy_tensors, capture_hidden_mode: (
+                SimpleNamespace(pp_hidden_states_proxy_tensors=pp_proxy_tensors)
+            ),
+        )
+        with (
+            mock.patch(
+                "sglang.srt.speculative.dspark_components.dspark_worker_v2.get_pp_group",
+                return_value=SimpleNamespace(is_last_rank=False),
+            ),
+            mock.patch(
+                "sglang.srt.speculative.dspark_components.dspark_worker_v2.get_schedule",
+                return_value=SimpleNamespace(page_size=1),
+            ),
+        ):
+            worker = DSparkWorkerV2(None, 0, None, 0, target)
+        worker.alloc_memory_pool()
+        worker.init_attention_backends()
+        worker.init_cuda_graphs()
+        batch = SimpleNamespace(seq_lens=torch.tensor([8]))
+        result = worker.forward_batch_generation(batch, pp_proxy_tensors=proxy)
+        self.assertIs(result.pp_hidden_states_proxy_tensors, proxy)
+        self.assertIs(result.new_seq_lens, batch.seq_lens)
+        self.assertIsNone(worker.get_confidence_budget_prepare())
+        self.assertIsNone(worker.primary_draft_kv_pool)
+        self.assertEqual(worker.preloaded_weights_bytes, 0)
+        self.assertEqual(
+            worker.spec_v2_attn_backends, (target.model_runner.attn_backend,)
+        )
+
     def test_single_aux_hidden_state_is_returned_without_copy(self) -> None:
         hidden_states = torch.empty(2, 3)