[DSpark] Support pipeline-parallel targets in disaggregated and aggregated serving (+ padded graph batch safety)

tests/config/test_dspark_prefill_only.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm.config import KVTransferConfig, ParallelConfig, SpeculativeConfig
+from vllm.config.kv_transfer import KVRole
+from vllm.config.speculative import SpeculativeMethod
+
+
+def _spec_config(
+    *, pp: int, role: KVRole, method: SpeculativeMethod = "dspark"
+) -> SpeculativeConfig:
+    config = object.__new__(SpeculativeConfig)
+    config.method = method
+    config.target_parallel_config = ParallelConfig(pipeline_parallel_size=pp)
+    config.target_kv_transfer_config = KVTransferConfig(
+        kv_connector="NixlConnector",
+        kv_role=role,
+    )
+    return config
+
+
+@pytest.mark.parametrize(
+    ("pp", "role", "method", "expected"),
+    [
+        (2, "kv_producer", "dspark", True),
+        (4, "kv_producer", "dspark", True),
+        (1, "kv_producer", "dspark", False),
+        (2, "kv_consumer", "dspark", False),
+        (2, "kv_both", "dspark", False),
+        (2, "kv_producer", "dflash", False),
+    ],
+)
+def test_dspark_prefill_only_role_detection(pp, role, method, expected):
+    assert _spec_config(pp=pp, role=role, method=method).is_dspark_prefill_only() is (
+        expected
+    )
+
+
+def test_dspark_prefill_materializer_uses_pp1_draft_config():
+    target = ParallelConfig(pipeline_parallel_size=4, tensor_parallel_size=1)
+
+    draft = SpeculativeConfig.create_draft_parallel_config(
+        target,
+        speculative_draft_tensor_parallel_size=1,
+    )
+
+    assert draft.pipeline_parallel_size == 1
+    assert draft.tensor_parallel_size == 1
+
+
+def _spec_config_no_kv(
+    *, pp: int, method: SpeculativeMethod = "dspark"
+) -> SpeculativeConfig:
+    config = object.__new__(SpeculativeConfig)
+    config.method = method
+    config.target_parallel_config = ParallelConfig(pipeline_parallel_size=pp)
+    config.target_kv_transfer_config = None  # type: ignore[assignment]
+    return config
+
+
+@pytest.mark.parametrize(
+    ("pp", "method", "expected"),
+    [
+        (2, "dspark", True),
+        (4, "dspark", True),
+        (1, "dspark", False),
+        (2, "dflash", False),
+        (2, "mtp", False),
+    ],
+)
+def test_dspark_last_stage_drafter_aggregated(pp, method, expected):
+    # Aggregated (IFB) serving: no KV transfer config at all.
+    assert _spec_config_no_kv(pp=pp, method=method).use_dspark_last_stage_drafter() is (
+        expected
+    )
+
+
+def test_dspark_last_stage_drafter_covers_pd_roles():
+    # kv_producer (prefill-only) and kv_consumer/aggregated all take the
+    # last-stage drafter path once the target is pipeline-parallel.
+    assert (
+        _spec_config(pp=2, role="kv_producer").use_dspark_last_stage_drafter() is True
+    )
+    assert (
+        _spec_config(pp=2, role="kv_consumer").use_dspark_last_stage_drafter() is True
+    )

tests/kernels/moe/test_topk_softplus_sqrt.py

@@ -280,6 +280,136 @@ def test_dsv4_fast_topk(
     )
 
 
+@pytest.mark.skipif(
+    not current_platform.is_cuda(),
+    reason="The DeepSeek V4 fast path is CUDA-only.",
+)
+def test_dsv4_fast_topk_padding_uint32_falls_back(monkeypatch: pytest.MonkeyPatch):
+    """Padded rows need the -1 sentinel, which uint32 cannot represent: the
+    router must skip the dsv4 fast path and still route the real rows."""
+    torch.manual_seed(0)
+    num_tokens = 17
+    num_experts = 256
+    hidden_states = torch.randn((num_tokens, 64), dtype=torch.float32, device="cuda")
+    gating_output = torch.randn(
+        (num_tokens, num_experts), dtype=torch.float32, device="cuda"
+    )
+    correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda")
+    is_padding = torch.zeros(num_tokens, dtype=torch.bool, device="cuda")
+    is_padding[1::2] = True
+    gating_output[is_padding] = float("nan")
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.fused_moe.router."
+        "fused_topk_bias_router._get_padding_mask",
+        lambda _: is_padding,
+    )
+    # uint32 + padding would trip dsv4_topk's signed-indices assertion; the
+    # generic path must take over instead.
+    topk_weights, topk_ids = fused_topk_bias(
+        hidden_states=hidden_states,
+        gating_output=gating_output,
+        scoring_func="sqrtsoftplus",
+        e_score_correction_bias=correction_bias,
+        topk=6,
+        renormalize=True,
+        indices_type=torch.uint32,
+        routed_scaling_factor=1.5,
+    )
+
+    assert topk_ids.dtype == torch.uint32
+    topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
+        gating_output=gating_output[~is_padding],
+        topk=6,
+        renormalize=True,
+        routed_scaling_factor=1.5,
+        e_score_correction_bias=correction_bias,
+    )
+    # uint32 CUDA tensors do not support boolean-mask indexing; widen first.
+    torch.testing.assert_close(
+        topk_ids.to(torch.int64)[~is_padding],
+        topk_ids_ref.to(torch.int64),
+        atol=0,
+        rtol=0,
+    )
+    torch.testing.assert_close(
+        topk_weights[~is_padding], topk_weights_ref, atol=2e-5, rtol=2e-5
+    )
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda(),
+    reason="The DeepSeek V4 fast path is CUDA-only.",
+)
+def test_dsv4_fast_topk_padding(monkeypatch: pytest.MonkeyPatch):
+    """Verify the DSV4 fast path removes graph-padding rows from routing."""
+    torch.manual_seed(0)
+    num_tokens = 17
+    num_experts = 256
+    hidden_states = torch.randn((num_tokens, 64), dtype=torch.float32, device="cuda")
+    gating_output = torch.randn(
+        (num_tokens, num_experts), dtype=torch.float32, device="cuda"
+    )
+    correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda")
+    is_padding = torch.zeros(num_tokens, dtype=torch.bool, device="cuda")
+    is_padding[1::2] = True
+    gating_output[is_padding] = float("nan")
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.fused_moe.router."
+        "fused_topk_bias_router._get_padding_mask",
+        lambda _: is_padding,
+    )
+    topk_weights, topk_ids = fused_topk_bias(
+        hidden_states=hidden_states,
+        gating_output=gating_output,
+        scoring_func="sqrtsoftplus",
+        e_score_correction_bias=correction_bias,
+        topk=6,
+        renormalize=True,
+        routed_scaling_factor=1.5,
+    )
+
+    assert torch.equal(topk_ids[is_padding], torch.full_like(topk_ids[is_padding], -1))
+    assert torch.equal(
+        topk_weights[is_padding], torch.zeros_like(topk_weights[is_padding])
+    )
+
+    topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
+        gating_output=gating_output[~is_padding],
+        topk=6,
+        renormalize=True,
+        routed_scaling_factor=1.5,
+        e_score_correction_bias=correction_bias,
+    )
+    torch.testing.assert_close(topk_ids[~is_padding], topk_ids_ref, atol=0, rtol=0)
+    torch.testing.assert_close(
+        topk_weights[~is_padding], topk_weights_ref, atol=2e-5, rtol=2e-5
+    )
+
+    # The mask buffer is persistent under CUDA graph replay, but its contents
+    # change with every batch. Verify that the kernel reads those contents at
+    # runtime rather than specializing on the mask captured above.
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        graph_weights, graph_ids = dsv4_topk(
+            gating_output,
+            correction_bias,
+            torch.int32,
+            1.5,
+            is_padding=is_padding,
+        )
+
+    is_padding.logical_not_()
+    graph.replay()
+    assert torch.equal(
+        graph_ids[is_padding], torch.full_like(graph_ids[is_padding], -1)
+    )
+    assert torch.equal(
+        graph_weights[is_padding], torch.zeros_like(graph_weights[is_padding])
+    )
+
+
 @pytest.mark.skipif(
     not current_platform.is_cuda_alike(),
     reason="This test is skipped on non-CUDA platform.",

tests/models/kimi_k3/test_dspark_mla.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""K3 DSpark draft weight-mapping tests."""
+
+from vllm.models.kimi_k3.nvidia.dspark_mla import _build_weights_mapper
+
+
+def test_mapper_drops_embed_for_target_aliasing():
+    mapper = _build_weights_mapper(drop_embed=True)
+    assert mapper.apply_list(["embed_tokens.weight"]) == []
+    assert mapper.apply_list(["lm_head.weight"]) == []
+    # Draft-owned weights still map into the model namespace.
+    assert mapper.apply_list(["layers.0.mlp.gate_proj.weight"]) == [
+        "model.layers.0.mlp.gate_up_proj.weight"
+    ]
+
+
+def test_mapper_keeps_embed_under_pp():
+    # Under PP the drafter cannot alias the target's first-stage table, so the
+    # checkpoint's own embed_tokens.weight must flow through.
+    mapper = _build_weights_mapper(drop_embed=False)
+    assert mapper.apply_list(["embed_tokens.weight"]) == ["model.embed_tokens.weight"]
+    assert mapper.apply_list(["lm_head.weight"]) == []

tests/models/kimi_k3/test_eagle3.py

@@ -200,3 +200,111 @@ def test_attn_res_stream_capture_receives_the_layer_outputs_in_order(monkeypatch
     assert got_pending is layer_hidden_states
     assert got_residual is block_residual
     torch.testing.assert_close(aux_hidden_states[0], captured)
+
+
+def _make_stage(
+    *,
+    start_layer: int,
+    taps: tuple[int, ...],
+    layer_outputs: list[tuple[torch.Tensor, None, torch.Tensor]],
+) -> KimiLinearModel:
+    model = _make_kimi_linear_model()
+    end_layer = start_layer + len(layer_outputs)
+    object.__setattr__(model, "start_layer", start_layer)
+    object.__setattr__(model, "end_layer", end_layer)
+    # The real model keeps the global layer list and slices [start:end].
+    layers = [Mock() for _ in range(end_layer)]
+    for i, out in enumerate(layer_outputs):
+        layers[start_layer + i] = Mock(return_value=out)
+    object.__setattr__(model, "layers", layers)
+    object.__setattr__(model, "aux_hidden_state_layers", taps)
+    object.__setattr__(model, "config", SimpleNamespace(hidden_size=2))
+    return model
+
+
+def test_kimi_linear_aux_hidden_states_flow_across_pp_stages(monkeypatch):
+    """A tap owned by an earlier PP stage must reach the last stage intact.
+
+    The drafter's taps can reference layers outside the last stage (K3 taps
+    [24, 48, 72, 88, 92]); each stage packs the taps it owns under global
+    per-tap keys (EagleModelMixin.pack_local_aux_hidden_states) and the last
+    stage prepends the collected remote taps to its own.
+    """
+    stage0_hidden = torch.tensor([[1.0, 2.0]])
+    stage0_residual = torch.tensor([[3.0, 4.0]])
+    stage1_hidden = torch.tensor([[5.0, 6.0]])
+    stage1_residual = torch.tensor([[7.0, 8.0]])
+
+    stage0 = _make_stage(
+        start_layer=0,
+        taps=(1, 2),
+        layer_outputs=[(stage0_hidden, None, stage0_residual)],
+    )
+    stage1 = _make_stage(
+        start_layer=1,
+        taps=(1, 2),
+        layer_outputs=[(stage1_hidden, None, stage1_residual)],
+    )
+    # EagleModelMixin caches the PP aux layout in _set_aux_hidden_state_layers;
+    # the stubs set layers directly, so prime the caches by hand.
+    object.__setattr__(stage0, "_aux_slot_base_cached", 0)
+    object.__setattr__(stage1, "_aux_slot_base_cached", 1)
+    object.__setattr__(stage1, "_aux_upstream_total_cached", 1)
+
+    monkeypatch.setattr(
+        kimi_model,
+        "get_pp_group",
+        lambda: SimpleNamespace(is_first_rank=True, is_last_rank=False),
+    )
+    stage0_out = stage0.forward(
+        input_ids=None,
+        positions=torch.tensor([0]),
+        intermediate_tensors=None,
+        inputs_embeds=torch.zeros(1, 2),
+    )
+
+    # Stage 0 owns the post-layer-1 tap; it rides the wire under its global
+    # slot key.
+    stage0_aux = stage0_hidden + stage0_residual
+    torch.testing.assert_close(stage0_out.tensors["aux_hidden_states_0"], stage0_aux)
+
+    monkeypatch.setattr(
+        kimi_model,
+        "get_pp_group",
+        lambda: SimpleNamespace(is_first_rank=False, is_last_rank=True),
+    )
+    output, aux_hidden_states = stage1.forward(
+        input_ids=None,
+        positions=torch.tensor([0]),
+        intermediate_tensors=stage0_out,
+    )
+
+    # The boundary tap (position 1 == stage1's start_layer) must not be
+    # duplicated by the stage-entry capture: two taps, in ascending order.
+    assert len(aux_hidden_states) == 2
+    torch.testing.assert_close(aux_hidden_states[0], stage0_aux)
+    torch.testing.assert_close(aux_hidden_states[1], stage1_hidden + stage1_residual)
+    torch.testing.assert_close(output, stage1_hidden + stage1_residual)
+
+
+def test_kimi_linear_first_stage_without_taps_sends_no_aux_buffer(monkeypatch):
+    """No taps captured on the stage -> no aux keys on the wire."""
+    stage0 = _make_stage(
+        start_layer=0,
+        taps=(2,),
+        layer_outputs=[(torch.ones(1, 2), None, torch.zeros(1, 2))],
+    )
+    object.__setattr__(stage0, "_aux_slot_base_cached", 0)
+
+    monkeypatch.setattr(
+        kimi_model,
+        "get_pp_group",
+        lambda: SimpleNamespace(is_first_rank=True, is_last_rank=False),
+    )
+    out = stage0.forward(
+        input_ids=None,
+        positions=torch.tensor([0]),
+        intermediate_tensors=None,
+        inputs_embeds=torch.zeros(1, 2),
+    )
+    assert not any(key.startswith("aux_hidden_states_") for key in out.tensors)

tests/models/test_deepseek_v4_mega_moe.py

@@ -743,12 +743,27 @@ def test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe():
     mtp = SimpleNamespace(finalize_mega_moe_weights=lambda: calls.append("mtp"))
     DeepSeekV4MTP.process_weights_after_loading(mtp)
 
-    dspark = SimpleNamespace(_finalize_moe=lambda: calls.append("dspark"))
+    dspark = SimpleNamespace(
+        model=SimpleNamespace(context_kv_only=False),
+        _finalize_moe=lambda: calls.append("dspark"),
+    )
     DSparkDeepseekV4ForCausalLM.process_weights_after_loading(dspark)
 
     assert calls == ["mtp", "dspark"]
 
 
+def test_dspark_context_materializer_skips_absent_confidence_head():
+    """The context-only P model omits decode-only heads entirely."""
+    dspark = object.__new__(DSparkDeepseekV4ForCausalLM)
+    dspark.model = SimpleNamespace()
+
+    assert dspark._remap_dspark_name("mtp.2.confidence_head.weight") is None
+
+    dspark.model.context_kv_only = True
+    dspark._finalize_moe = lambda: pytest.fail("context-only model has no MoE")
+    dspark.process_weights_after_loading()
+
+
 @pytest.mark.skipif(
     not torch.cuda.is_available(),
     reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",