fix(spec decode): harden DSpark and DFlash edge paths

tests/v1/attention/test_deepseek_v4_dspark_metadata.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import torch
+
+from vllm.v1.attention.backends.mla.sparse_swa import (
+    DeepseekSparseSWAMetadataBuilder,
+)
+from vllm.v1.kv_cache_interface import MLAAttentionSpec
+
+
+def test_dspark_swa_decode_threshold_matches_target_verification() -> None:
+    """DSpark verifies 1 + K target tokens, not the generic 1 + 2K."""
+    speculative_config = SimpleNamespace(
+        num_speculative_tokens=5,
+        parallel_drafting=True,
+        use_dspark=lambda: True,
+    )
+    hf_config = SimpleNamespace(sliding_window=128, compress_ratios=[1, 4, 128])
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(max_model_len=4096, hf_config=hf_config),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=16),
+        speculative_config=speculative_config,
+        parallel_config=SimpleNamespace(
+            decode_context_parallel_size=1,
+            prefill_context_parallel_size=1,
+            cp_kv_cache_interleave_size=1,
+        ),
+    )
+    kv_cache_spec = MLAAttentionSpec(
+        block_size=256,
+        num_kv_heads=1,
+        head_size=512,
+        dtype=torch.bfloat16,
+    )
+
+    builder = DeepseekSparseSWAMetadataBuilder(
+        kv_cache_spec,
+        ["placeholder"],
+        vllm_config,
+        torch.device("cpu"),
+    )
+
+    assert builder.decode_threshold == 6

tests/v1/spec_decode/test_acceptance_length_controller.py

@@ -268,6 +268,9 @@ def test_synthetic_scheduler_output_uses_default_speculative_depth():
     output.num_spec_tokens_to_schedule = 2
     assert output.resolve_num_spec_tokens_to_schedule(default=5) == 2
 
+    output.num_spec_tokens_to_schedule = 0
+    assert output.resolve_num_spec_tokens_to_schedule(default=5) == 0
+
 
 def test_runner_v2_autoregressive_drafter_stops_at_adaptive_depth(monkeypatch):
     monkeypatch.setattr(AutoRegressiveSpeculator, "__abstractmethods__", frozenset())

tests/v1/spec_decode/test_dflash_cudagraph_lifetime.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import torch
+
+from vllm.config.compilation import CUDAGraphMode
+from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
+
+
+def _make_speculator() -> SimpleNamespace:
+    hidden_states = torch.randn(2, 8)
+    return SimpleNamespace(
+        _run_model=Mock(return_value=hidden_states),
+        _captured_backbone_outputs=[],
+        num_speculative_steps=2,
+        sample_indices=torch.tensor([0, 1]),
+        sample_pos=torch.tensor([1, 2]),
+        sample_idx_mapping=torch.tensor([0, 0]),
+        temperature=torch.ones(1),
+        seeds=torch.zeros(1, dtype=torch.int64),
+        sample_col=torch.tensor([0, 1]),
+        draft_logits=None,
+        sample_draft=Mock(return_value=torch.tensor([11, 12])),
+        draft_tokens=torch.zeros(1, 2, dtype=torch.int64),
+    )
+
+
+def test_dflash_retains_backbone_output_during_cudagraph_capture(monkeypatch):
+    speculator = _make_speculator()
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
+
+    DFlashSpeculator._generate_draft(
+        speculator,
+        num_reqs=1,
+        num_tokens_padded=2,
+        attn_metadata=None,
+        slot_mappings=None,
+        num_tokens_across_dp=None,
+        cudagraph_runtime_mode=CUDAGraphMode.NONE,
+    )
+
+    assert len(speculator._captured_backbone_outputs) == 1
+    assert (
+        speculator._captured_backbone_outputs[0] is speculator._run_model.return_value
+    )
+
+
+def test_dflash_does_not_retain_eager_backbone_output(monkeypatch):
+    speculator = _make_speculator()
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+
+    DFlashSpeculator._generate_draft(
+        speculator,
+        num_reqs=1,
+        num_tokens_padded=2,
+        attn_metadata=None,
+        slot_mappings=None,
+        num_tokens_across_dp=None,
+        cudagraph_runtime_mode=CUDAGraphMode.NONE,
+    )
+
+    assert speculator._captured_backbone_outputs == []

tests/v1/spec_decode/test_dflash_prefix_cache_masking.py

@@ -0,0 +1,170 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DFlash/DSpark draft context masking under prefix caching.
+
+Cache-restored tokens never flow through the target forward, so the draft's
+context KV is never written for them. shift_draft_block_tables hides those
+slots from the draft's attention by left-shifting each request's block-table
+row by the restored whole blocks (seq_lens is shortened to match by
+_prepare_dflash_inputs_kernel).
+"""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+from vllm.v1.worker.gpu.spec_decode.dflash.speculator import (
+    DFlashSpeculator,
+    shift_draft_block_tables,
+)
+
+pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
+
+DEVICE = "cuda"
+BLOCK_SIZE = 16
+MAX_BLOCKS = 64
+MAX_NUM_REQS = 8
+
+
+def test_unaligned_cached_prefix_detection():
+    speculator = SimpleNamespace(
+        num_cached_tokens_np=np.array([32, 35, 64], dtype=np.int32),
+        block_tables=SimpleNamespace(kernel_block_sizes=[16]),
+    )
+
+    aligned = SimpleNamespace(
+        idx_mapping_np=np.array([0, 2], dtype=np.int32),
+        num_reqs=2,
+    )
+    unaligned = SimpleNamespace(
+        idx_mapping_np=np.array([0, 1], dtype=np.int32),
+        num_reqs=2,
+    )
+
+    assert not DFlashSpeculator._has_unaligned_cached_prefix(speculator, aligned)
+    assert DFlashSpeculator._has_unaligned_cached_prefix(speculator, unaligned)
+
+
+def _make_block_table(num_reqs: int) -> torch.Tensor:
+    # Distinct block ids per (request, slot) so shifts are detectable.
+    table = torch.arange(
+        MAX_NUM_REQS * MAX_BLOCKS, dtype=torch.int32, device=DEVICE
+    ).view(MAX_NUM_REQS, MAX_BLOCKS)
+    return table[:num_reqs].contiguous()
+
+
+@pytest.mark.parametrize(
+    "num_cached,expected_shift",
+    [
+        (0, 0),  # no cache hit: no-op
+        (BLOCK_SIZE * 3, 3),  # block-aligned hit (the common APC case)
+        (BLOCK_SIZE * 3 + 5, 3),  # unaligned: floor to whole blocks
+        (BLOCK_SIZE - 1, 0),  # less than one block: no-op
+    ],
+)
+def test_shift_single_request(num_cached: int, expected_shift: int):
+    block_table = _make_block_table(1)
+    original = block_table.clone()
+    idx_mapping = torch.zeros(1, dtype=torch.int32, device=DEVICE)
+    num_cached_tokens = torch.full(
+        (MAX_NUM_REQS,), num_cached, dtype=torch.int32, device=DEVICE
+    )
+
+    seq_lens = torch.full(
+        (idx_mapping.shape[0],),
+        MAX_BLOCKS * BLOCK_SIZE,
+        dtype=torch.int32,
+        device=DEVICE,
+    )
+    shift_draft_block_tables(
+        block_table, idx_mapping, num_cached_tokens, seq_lens, BLOCK_SIZE
+    )
+
+    kept = MAX_BLOCKS - expected_shift
+    torch.testing.assert_close(block_table[0, :kept], original[0, expected_shift:])
+
+
+def test_shift_per_request_and_idx_mapping():
+    # Requests in batch order 0..3 map to request-state slots 3..0, with a
+    # different cached count per slot. Each row must shift by its own count.
+    num_reqs = 4
+    block_table = _make_block_table(num_reqs)
+    original = block_table.clone()
+    idx_mapping = torch.tensor([3, 2, 1, 0], dtype=torch.int32, device=DEVICE)
+    # Slot i has i whole cached blocks.
+    num_cached_tokens = torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=DEVICE)
+    num_cached_tokens[:4] = (
+        torch.arange(4, dtype=torch.int32, device=DEVICE) * BLOCK_SIZE
+    )
+
+    seq_lens = torch.full(
+        (idx_mapping.shape[0],),
+        MAX_BLOCKS * BLOCK_SIZE,
+        dtype=torch.int32,
+        device=DEVICE,
+    )
+    shift_draft_block_tables(
+        block_table, idx_mapping, num_cached_tokens, seq_lens, BLOCK_SIZE
+    )
+
+    for batch_idx in range(num_reqs):
+        shift = int(idx_mapping[batch_idx])  # slot id == cached blocks
+        kept = MAX_BLOCKS - shift
+        torch.testing.assert_close(
+            block_table[batch_idx, :kept],
+            original[batch_idx, shift:],
+            msg=f"batch row {batch_idx} (slot {shift})",
+        )
+
+
+def test_shift_large_row_in_place_overlap():
+    # Shift smaller than the copy chunk (1024) exercises the overlapping
+    # in-place load-before-store path on a long row.
+    max_blocks = 4096
+    block_table = (
+        torch.arange(max_blocks, dtype=torch.int32, device=DEVICE)
+        .unsqueeze(0)
+        .contiguous()
+    )
+    original = block_table.clone()
+    idx_mapping = torch.zeros(1, dtype=torch.int32, device=DEVICE)
+    num_cached_tokens = torch.full(
+        (1,), 7 * BLOCK_SIZE, dtype=torch.int32, device=DEVICE
+    )
+
+    seq_lens = torch.full(
+        (idx_mapping.shape[0],),
+        max_blocks * BLOCK_SIZE,
+        dtype=torch.int32,
+        device=DEVICE,
+    )
+    shift_draft_block_tables(
+        block_table, idx_mapping, num_cached_tokens, seq_lens, BLOCK_SIZE
+    )
+
+    torch.testing.assert_close(block_table[0, : max_blocks - 7], original[0, 7:])
+
+
+def test_shift_copy_bounded_by_seq_len():
+    # Only the blocks referenced by the shifted sequence move; the tail of the
+    # row must stay untouched (perf guard for long-context block tables).
+    block_table = _make_block_table(1)
+    original = block_table.clone()
+    idx_mapping = torch.zeros(1, dtype=torch.int32, device=DEVICE)
+    num_cached_tokens = torch.full(
+        (MAX_NUM_REQS,), 4 * BLOCK_SIZE, dtype=torch.int32, device=DEVICE
+    )
+    # Shifted draft length of 3.5 blocks -> exactly 4 blocks copied.
+    seq_lens = torch.full(
+        (1,), 3 * BLOCK_SIZE + BLOCK_SIZE // 2, dtype=torch.int32, device=DEVICE
+    )
+
+    shift_draft_block_tables(
+        block_table, idx_mapping, num_cached_tokens, seq_lens, BLOCK_SIZE
+    )
+
+    torch.testing.assert_close(block_table[0, :4], original[0, 4:8])
+    torch.testing.assert_close(block_table[0, 4:], original[0, 4:])

tests/v1/spec_decode/test_dynamic_sd_cug.py

@@ -62,7 +62,10 @@ def _create_vllm_config_for_dsd(
     vllm_config.num_speculative_tokens = max_spec_tokens
 
     speculative_config = MagicMock()
-    speculative_config.uses_dynamic_speculative_decoding.return_value = use_dynamic_sd
+    speculative_config.uses_batch_size_dynamic_speculative_decoding.return_value = (
+        use_dynamic_sd
+    )
+    speculative_config.uses_acceptance_length_adaptation.return_value = False
     if use_dynamic_sd:
         # DSD reads the per-batch-size schedule; a schedule entry with K
         # speculative tokens maps to decode query length K + 1. By default
@@ -326,3 +329,51 @@ def test_dynamic_sd_only_captures_scheduled_query_lengths(monkeypatch):
                 assert desc.num_tokens == num_tokens
                 assert desc.num_reqs is None
             assert desc.num_active_loras == 0
+
+
+def test_dynamic_sd_skips_zero_draft_tokens_in_cudagraph_schedule(monkeypatch):
+    """K=0 in the DSD schedule must not produce decode_query_len=0.
+
+    DSpark (anchor-as-first) passes ``num_query_per_req == num_speculative_tokens``
+    to the draft CudaGraphManager, so ``num_new_sampled_tokens_per_step`` recovers
+    as 0. A schedule entry with K=0 would otherwise crash during candidate init.
+    """
+
+    max_num_seqs = 128
+    max_spec_tokens = 5
+
+    monkeypatch.setattr(
+        gpu_cudagraph_utils,
+        "get_pp_group",
+        lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
+    )
+
+    vllm_config = _create_vllm_config_for_dsd(
+        max_num_seqs=max_num_seqs,
+        max_spec_tokens=max_spec_tokens,
+        cudagraph_mode="FULL_AND_PIECEWISE",
+        use_dynamic_sd=True,
+        num_spec_per_batch_size=[
+            (1, 32, 5),
+            (33, 64, 3),
+            (65, 96, 1),
+            (97, 128, 0),
+        ],
+    )
+    draft_decode_query_len = max_spec_tokens
+
+    manager = gpu_cudagraph_utils.CudaGraphManager(
+        vllm_config=vllm_config,
+        device=torch.device("cpu"),
+        cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,
+        decode_query_len=draft_decode_query_len,
+    )
+
+    scheduled_query_lens = {5, 3, 1}
+    captured_query_lens = {
+        desc.uniform_token_count
+        for descs in manager._candidates.values()
+        for desc in descs
+        if desc.cg_mode == CUDAGraphMode.FULL and desc.uniform_token_count is not None
+    }
+    assert captured_query_lens == scheduled_query_lens

tests/v1/worker/test_gpu_sampling_states_seed.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import torch
+
+from vllm.sampling_params import SamplingParams
+from vllm.v1.worker.gpu.sample import states
+
+
+class _HostBackedTensor:
+    def __init__(self, size: int, dtype: torch.dtype):
+        self.cpu = torch.zeros(size, dtype=dtype)
+        self.np = self.cpu.numpy()
+        self.gpu = self.cpu
+
+    def copy_to_uva(self, n: int | None = None) -> torch.Tensor:
+        return self.gpu[:n] if n is not None else self.gpu
+
+
+def test_fallback_seeds_do_not_depend_on_global_numpy_rng(monkeypatch) -> None:
+    monkeypatch.setattr(states, "UvaBackedTensor", _HostBackedTensor)
+    rank0 = states.SamplingStates(4, 128, seed=17)
+
+    np.random.seed(1234)
+    np.random.random(1000)
+    rank1 = states.SamplingStates(4, 128, seed=17)
+
+    params = SamplingParams(seed=None)
+    for req_idx in range(4):
+        rank0.add_request(req_idx, params)
+        np.random.random(req_idx + 1)
+        rank1.add_request(req_idx, params)
+
+    np.testing.assert_array_equal(rank0.seeds.np, rank1.seeds.np)