[DSpark] Support pipeline-parallel targets in aggregated serving

tests/kernels/moe/test_topk_softplus_sqrt.py

@@ -339,6 +339,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/v1/spec_decode/test_dflash_prepare_inputs.py

@@ -1,14 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from contextlib import contextmanager
 from types import SimpleNamespace
 
 import numpy as np
 import pytest
 import torch
 
 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
+from vllm.v1.worker.gpu.spec_decode.dflash import speculator as dflash_speculator
 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import (
+    DFlashSpeculator,
     prepare_dflash_inputs,
 )
 
@@ -33,6 +36,7 @@ def _run_prepare(
     input_buffers = SimpleNamespace(
         input_ids=torch.full((max_num_tokens,), -1, dtype=torch.int32, device=device),
         positions=torch.full((max_num_tokens,), -1, dtype=torch.int64, device=device),
+        is_padding=torch.zeros(max_num_tokens, dtype=torch.bool, device=device),
         query_start_loc=torch.full(
             (max_num_reqs + 1,), -1, dtype=torch.int32, device=device
         ),
@@ -140,6 +144,114 @@ def test_prepare_dflash_inputs_excludes_rejected_context_suffix():
     assert out.temperature[2].item() == 1.0
     assert out.seeds[2].item() == 17
 
+    assert not out.input_buffers.is_padding[:3].any()
+    assert out.input_buffers.is_padding[3:].all()
+    assert out.input_buffers.input_ids[3:].cpu().tolist() == [0] * 13
+    assert out.input_buffers.positions[3:].cpu().tolist() == [0] * 13
+
+
+def test_prepare_dflash_inputs_compacts_noncontiguous_request_slots():
+    device = torch.device("cuda")
+    max_num_reqs = 4
+    max_num_tokens = 16
+    num_speculative_steps = 3
+    input_buffers = SimpleNamespace(
+        input_ids=torch.full((max_num_tokens,), -1, dtype=torch.int32, device=device),
+        positions=torch.full((max_num_tokens,), -1, dtype=torch.int64, device=device),
+        is_padding=torch.zeros(max_num_tokens, dtype=torch.bool, device=device),
+        query_start_loc=torch.full(
+            (max_num_reqs + 1,), -1, dtype=torch.int32, device=device
+        ),
+        seq_lens=torch.full((max_num_reqs,), -1, dtype=torch.int32, device=device),
+    )
+    input_batch = SimpleNamespace(
+        num_reqs=2,
+        num_scheduled_tokens=np.array([4, 4], dtype=np.int32),
+        positions=torch.tensor(
+            [10, 11, 12, 13, 20, 21, 22, 23],
+            dtype=torch.int64,
+            device=device,
+        ),
+        query_start_loc=torch.tensor([0, 4, 8], dtype=torch.int32, device=device),
+        # Active batch rows are compact, while request state remains in slots 3 and 1.
+        idx_mapping=torch.tensor([3, 1], dtype=torch.int32, device=device),
+    )
+    query_slot_mapping = torch.full(
+        (max_num_tokens,), -2, dtype=torch.int64, device=device
+    )
+    context_positions = torch.full(
+        (max_num_tokens,), -1, dtype=torch.int64, device=device
+    )
+    context_slot_mapping = torch.full(
+        (max_num_tokens,), -2, dtype=torch.int64, device=device
+    )
+    sample_indices = torch.full(
+        (max_num_reqs * num_speculative_steps,),
+        -1,
+        dtype=torch.int64,
+        device=device,
+    )
+    sample_pos = torch.full_like(sample_indices, -1)
+    sample_idx_mapping = torch.full(
+        sample_indices.shape, -1, dtype=torch.int32, device=device
+    )
+    temperature = torch.zeros(max_num_reqs, dtype=torch.float32, device=device)
+    seeds = torch.zeros(max_num_reqs, dtype=torch.int64, device=device)
+    input_temperature = torch.tensor(
+        [0.0, 0.5, 0.0, 1.0], dtype=torch.float32, device=device
+    )
+    input_seeds = torch.tensor([0, 11, 0, 33], dtype=torch.int64, device=device)
+    last_sampled = torch.tensor([0, 77, 0, 99], dtype=torch.int64, device=device)
+    next_prefill_tokens = torch.zeros_like(last_sampled)
+    block_table = torch.tensor(
+        [[0, 0, 7, 8, 9, 10, 11, 12], [0, 0, 13, 14, 15, 16, 17, 18]],
+        dtype=torch.int32,
+        device=device,
+    )
+
+    prepare_dflash_inputs(
+        input_buffers,
+        query_slot_mapping,
+        context_positions,
+        context_slot_mapping,
+        sample_indices,
+        sample_pos,
+        sample_idx_mapping,
+        temperature,
+        seeds,
+        input_batch,
+        torch.tensor([1, 1], dtype=torch.int32, device=device),
+        torch.tensor([2, 1], dtype=torch.int32, device=device),
+        last_sampled,
+        next_prefill_tokens,
+        input_temperature,
+        input_seeds,
+        block_table,
+        4,
+        0,
+        1,
+        1,
+        123,
+        num_speculative_steps,
+        num_speculative_steps,
+        max_num_reqs,
+        max_num_tokens,
+        128,
+        sample_from_anchor=True,
+    )
+    torch.accelerator.synchronize()
+
+    # Query rows follow compact batch order, but every persistent state lookup
+    # follows idx_mapping instead of accidentally using the compact row index.
+    assert input_buffers.input_ids[:6].cpu().tolist() == [99, 123, 123, 77, 123, 123]
+    assert input_buffers.positions[:6].cpu().tolist() == [12, 13, 14, 23, 24, 25]
+    assert sample_indices[:6].cpu().tolist() == [0, 1, 2, 3, 4, 5]
+    assert sample_idx_mapping[:6].cpu().tolist() == [3, 3, 3, 1, 1, 1]
+    assert temperature.cpu().tolist() == [0.0, 0.5, 0.0, 1.0]
+    assert seeds.cpu().tolist() == [0, 11, 0, 33]
+    assert not input_buffers.is_padding[:6].any()
+    assert input_buffers.is_padding[6:].all()
+
 
 def test_prepare_dflash_inputs_excludes_rejected_context_suffix_with_dcp():
     out = _run_prepare(
@@ -174,3 +286,40 @@ def test_prepare_dflash_inputs_never_writes_the_null_block():
         PAD_SLOT_ID,
         PAD_SLOT_ID,
     ]
+
+
+def test_dflash_forward_context_receives_draft_padding_mask(monkeypatch):
+    device = torch.device("cuda")
+    input_buffers = SimpleNamespace(
+        input_ids=torch.tensor([11, 12, 0, 0], dtype=torch.int32, device=device),
+        positions=torch.tensor([7, 8, 0, 0], dtype=torch.int64, device=device),
+        is_padding=torch.tensor([False, False, True, True], device=device),
+    )
+    observed = None
+
+    @contextmanager
+    def fake_set_forward_context(*args, **kwargs):
+        nonlocal observed
+        observed = kwargs["is_padding"].clone()
+        yield
+
+    monkeypatch.setattr(
+        dflash_speculator, "set_forward_context", fake_set_forward_context
+    )
+    speculator = SimpleNamespace(
+        input_buffers=input_buffers,
+        vllm_config=SimpleNamespace(),
+        model=lambda **kwargs: kwargs["input_ids"],
+    )
+
+    result = DFlashSpeculator._run_model(
+        speculator,
+        num_tokens=4,
+        attn_metadata=None,
+        slot_mappings=None,
+        num_tokens_across_dp=None,
+    )
+
+    assert result.tolist() == [11, 12, 0, 0]
+    assert observed is not None
+    assert observed.tolist() == [False, False, True, True]