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

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)