[Spec Decode] DSpark (#46995)

Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: Giancarlo Delfin <gdelfin@inferact.ai>
Co-authored-by: mgoin <mgoin64@gmail.com>

tests/models/registry.py

@@ -1407,6 +1407,19 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
         max_num_seqs=32,
         min_transformers_version="4.56.3",  # Required for Qwen3Next
     ),
+    # [DSpark]
+    "DSparkDraftModel": _HfExamplesInfo(
+        "deepseek-ai/DeepSeek-V4-Pro-DSpark",
+        speculative_model="deepseek-ai/DeepSeek-V4-Pro-DSpark",  # draft in mtp.*
+        is_available_online=False,
+        use_original_num_layers=True,  # DSpark has >1 draft block
+    ),
+    "Qwen3DSparkModel": _HfExamplesInfo(
+        "Qwen/Qwen3-8B",
+        speculative_model="deepseek-ai/dspark_qwen3_8b_block7",
+        is_available_online=False,
+        use_original_num_layers=True,  # DSpark backbone requires all layers
+    ),
     # [Eagle]
     "EagleCohereForCausalLM": _HfExamplesInfo(
         "/host/engines/cohere-moe",

tests/models/test_registry.py

@@ -48,6 +48,10 @@ def test_registry_imports(model_arch):
                 "(see #41376)"
             )
 
+    # DSpark draft model is NVIDIA-only; class is stubbed to None on ROCm/XPU.
+    if model_arch == "DSparkDraftModel" and not current_platform.is_cuda():
+        pytest.skip("DSparkDraftModel is only supported on CUDA")
+
     # Ensure all model classes can be imported successfully
     model_cls = ModelRegistry._try_load_model_cls(model_arch)
     assert model_cls is not None

tests/v1/attention/test_dspark_noncausal_sparse_mla.py

@@ -0,0 +1,529 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Correctness tests for DSpark non-causal sliding-window MLA via sparse indices.
+
+DSpark drafts a block of N tokens whose attention is NON-CAUSAL within the block:
+every block token attends to the sliding window of context AND to all block
+tokens (including ones at later positions than itself).
+
+We can implement this using the existing sparse-MLA pathway by expanding the window size
+to include the rest of the block tokens: instead of setting topk indices to the 127
+previous tokens, we expand it to the next power of 2 (256) and include up to
+swa_size + block_size - 1 topk indices, so that each query attends to the rest. The
+remaining slots are filled with padding.
+
+The sparse-MLA decode kernels (FlashMLA on SM90/SM100, FlashInfer TRTLLM on
+SM100/SM120) are index-driven: each query attends over exactly the slots in its
+index list, with no causal mask (see ``flash_mla_with_kvcache(..., indices=...)``
+and ``_forward_decode``'s "attend only by generated indices"). The existing
+``test_sparse_mla_backends`` suite already validates arbitrary index lists, but
+only ones whose entries are <= the query's own position. This test suite specifically
+ensures correctness of the non-causal attention case.
+
+This reuses the harness/helpers of ``test_sparse_mla_backends.py`` (same model
+shapes, fp8_ds_mla round-trip, mock indexer, MockSparseMLAAttentionLayer); only
+the index construction differs.
+"""
+
+import math
+from types import MethodType, SimpleNamespace
+
+import pytest
+import torch
+
+from tests.v1.attention.test_mla_backends import (
+    BatchSpec,
+    MockSparseMLAAttentionLayer,
+    create_and_prepopulate_kv_cache,
+)
+from tests.v1.attention.test_sparse_mla_backends import (
+    _quantize_dequantize_fp8_ds_mla,
+)
+from tests.v1.attention.utils import (
+    create_common_attn_metadata,
+    create_standard_kv_cache_spec,
+    create_vllm_config,
+)
+from vllm.config import set_current_vllm_config
+from vllm.model_executor.layers.linear import ColumnParallelLinear
+from vllm.platforms import current_platform
+
+if not current_platform.is_cuda():
+    pytest.skip(
+        "DSpark non-causal sparse MLA tests currently only support CUDA.",
+        allow_module_level=True,
+    )
+
+from vllm.utils.math_utils import cdiv
+from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
+    FlashInferMLASparseTRTLLMBackend,
+)
+from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseBackend
+from vllm.v1.attention.ops import flashmla
+
+DEVICE_TYPE = current_platform.device_type
+
+# (window, block_size, topk_width). topk_width must be a multiple of the kernel's
+# B_TOPK (= padded query-head count, 64 or 128); we use 128-multiples to cover
+# both. The "wide" case needs window + block > 128 -> width must grow past 128.
+_DSPARK_CONFIGS = {
+    "small_block": (8, 4, 128),
+    "full_window_block": (128, 5, 256),
+}
+
+
+def _build_dspark_noncausal_indices(
+    seq_lens: list[int],
+    query_lens: list[int],
+    window: int,
+    topk_width: int,
+    device: torch.device,
+) -> torch.Tensor:
+    """Per-token sparse indices for the DSpark non-causal block.
+
+    For a request with context length ``ctx`` and a query block of ``q_len``
+    tokens (block positions ``ctx .. ctx+q_len-1``), EVERY block query attends to
+    the same set: the trailing ``window`` context positions plus all block
+    positions, i.e. the contiguous range ``[max(ctx-window,0) .. ctx+q_len-1]``.
+    This is non-causal: an early block query's list contains later block tokens
+    (future-pointing). The list is padded to ``topk_width`` with ``-1``.
+    """
+    total_query_tokens = sum(query_lens)
+    sparse_indices = torch.full(
+        (total_query_tokens, topk_width), -1, dtype=torch.int32, device=device
+    )
+    gt = 0
+    for s_len, q_len in zip(seq_lens, query_lens):
+        ctx_len = s_len - q_len
+        lo = max(ctx_len - window, 0)
+        hi = ctx_len + q_len  # exclusive: window context + the full block
+        idx_list = torch.arange(lo, hi, dtype=torch.int32, device=device)
+        n = idx_list.numel()
+        assert n <= topk_width, (
+            f"index list ({n}) exceeds aligned topk width ({topk_width})"
+        )
+        for _ in range(q_len):
+            sparse_indices[gt, :n] = idx_list
+            gt += 1
+    return sparse_indices
+
+
+def _run_sparse_backend_vs_sdpa(
+    backend_cls,
+    seq_lens: list[int],
+    query_lens: list[int],
+    sparse_indices: torch.Tensor,
+    kv_cache_dtype: str,
+    block_size: int,
+    num_heads: int,
+    device: torch.device,
+    force_future_dominance: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Run a sparse-MLA backend with the given per-token indices and compute a
+    dense per-token SDPA reference over the SAME indices.
+
+    Mirrors ``test_sparse_mla_backends.test_sparse_backend_decode_correctness``
+    but with externally-supplied (non-causal) ``sparse_indices``.
+
+    ``num_heads`` selects the kernel's B_TOPK (= padded q-head count): 128 -> 128,
+    64 -> 64. The aligned widths (128/256) are multiples of both, so num_heads=64
+    exercises the head64 decode path that the SM100 alignment assert guards.
+
+    ``force_future_dominance`` scales the LAST block token's latent KV so it
+    dominates the softmax for every query that attends to it. With random data the
+    few future block tokens carry negligible attention mass (especially with a wide
+    window), so causal and non-causal outputs coincide; this knob makes the
+    future-token contribution provably large for the differentiation test. It is
+    OFF for the correctness test (which needs sensitivity to all tokens).
+
+    Returns (backend_output, noncausal_reference, causal_reference). The causal
+    reference restricts each query to indices <= its own absolute position.
+    """
+    batch_spec = BatchSpec(seq_lens=seq_lens, query_lens=query_lens)
+    topk_tokens = sparse_indices.shape[1]
+    dtype = torch.bfloat16
+    use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla"
+
+    kv_lora_rank = 512
+    qk_nope_head_dim = 128
+    qk_rope_head_dim = 64
+    v_head_dim = 128
+    head_size = kv_lora_rank + qk_rope_head_dim
+
+    max_seqlen = max(seq_lens)
+    total_cache_tokens = sum(seq_lens)
+
+    vllm_config = create_vllm_config(
+        model_name="deepseek-ai/DeepSeek-V2-Lite-Chat",
+        tensor_parallel_size=1,
+        max_model_len=max_seqlen,
+        num_gpu_blocks=max(2048, cdiv(total_cache_tokens, block_size) + 1),
+        block_size=block_size,
+        hf_config_override={
+            "index_topk": topk_tokens,
+            "attn_module_list_cfg": [{"topk_tokens": topk_tokens}],
+        },
+    )
+    model_config = vllm_config.model_config
+    model_config.hf_text_config = SimpleNamespace(
+        q_lora_rank=None,
+        kv_lora_rank=kv_lora_rank,
+        qk_nope_head_dim=qk_nope_head_dim,
+        qk_rope_head_dim=qk_rope_head_dim,
+        v_head_dim=v_head_dim,
+        model_type="deepseek_v2",
+    )
+    model_config.dtype = dtype
+    model_config.get_num_attention_heads = MethodType(
+        lambda self, parallel_config: num_heads, model_config
+    )
+    model_config.get_num_kv_heads = MethodType(
+        lambda self, parallel_config: 1, model_config
+    )
+    model_config.get_head_size = MethodType(lambda self: head_size, model_config)
+    model_config.get_sliding_window = MethodType(lambda self: None, model_config)
+
+    kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
+
+    torch.manual_seed(0)
+    scale = 1.0 / math.sqrt(head_size)
+
+    # Shared MLA projection weights, used by both reference and backend.
+    W_UK = torch.rand(
+        kv_lora_rank, num_heads, qk_nope_head_dim, dtype=dtype, device=device
+    )
+    W_UV = torch.rand(kv_lora_rank, num_heads, v_head_dim, dtype=dtype, device=device)
+
+    all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], []
+    kv_c_contexts, k_pe_contexts = [], []
+    reference_outputs = []
+    # Causal counterpart of the reference: same index lists, but each query is
+    # restricted to indices <= its own absolute position (drops future-pointing
+    # block tokens). Used to prove the non-causal result is genuinely different.
+    causal_reference_outputs = []
+
+    kv_cache_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
+    global_token_idx = 0
+
+    for s_len, q_len in zip(seq_lens, query_lens):
+        ctx_len = s_len - q_len
+
+        q_c = torch.rand(
+            q_len,
+            num_heads,
+            qk_nope_head_dim + qk_rope_head_dim,
+            dtype=dtype,
+            device=device,
+        )
+        kv_c_full = torch.rand(s_len, kv_lora_rank, dtype=dtype, device=device)
+        k_pe_full = torch.rand(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device)
+
+        if force_future_dominance:
+            # Scale the last block token's latent KV so its key/value dominate the
+            # softmax for any query attending to it. 4x in the latent dot makes its
+            # pre-softmax score exceed the others by a wide margin, so non-causal
+            # queries (which include it) diverge sharply from causal ones (which,
+            # for all but the last query, exclude it). Applied before quantization
+            # so cache and reference stay consistent.
+            kv_c_full[s_len - 1] = kv_c_full[s_len - 1] * 4.0 + 2.0
+
+        if use_fp8_ds_mla_quantization:
+            is_sm100 = torch.cuda.get_device_capability()[0] >= 10
+            kv_c_full, k_pe_squeezed = _quantize_dequantize_fp8_ds_mla(
+                kv_c_full,
+                k_pe_full.squeeze(1),
+                block_size=block_size,
+                scale=kv_cache_scale,
+                simulate_sm100_e8m0_scales=is_sm100,
+            )
+            k_pe_full = k_pe_squeezed.unsqueeze(1)
+
+        q_nope, q_pe = q_c.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)
+        ql_nope = torch.einsum("qnh,lnh->qnl", q_nope, W_UK)
+        q_mqa = torch.cat([ql_nope, q_pe], dim=-1)
+
+        k_mqa = torch.cat([kv_c_full, k_pe_full.squeeze(1)], dim=-1)
+        v_mqa = kv_c_full
+
+        # Per-token sparse SDPA reference over the supplied (non-causal) indices.
+        def _sparse_sdpa(idx_tensor, q_tok, k_mqa=k_mqa, v_mqa=v_mqa):
+            k_sparse = k_mqa[idx_tensor].unsqueeze(1).expand(-1, num_heads, -1)
+            v_sparse = v_mqa[idx_tensor].unsqueeze(1).expand(-1, num_heads, -1)
+            out = torch.nn.functional.scaled_dot_product_attention(
+                q_tok.unsqueeze(0).transpose(1, 2),
+                k_sparse.unsqueeze(0).transpose(1, 2),
+                v_sparse.unsqueeze(0).transpose(1, 2),
+                scale=scale,
+            )
+            out = out.transpose(1, 2).squeeze(0)
+            out = torch.einsum("qnl,lnv->qnv", out, W_UV)
+            return out.flatten(start_dim=-2)
+
+        for q_idx in range(q_len):
+            tok_sparse_idx = sparse_indices[global_token_idx]
+            valid_indices = tok_sparse_idx[tok_sparse_idx >= 0].long()
+
+            q_tok = q_mqa[q_idx : q_idx + 1]
+            reference_outputs.append(_sparse_sdpa(valid_indices, q_tok))
+
+            # Causal: drop indices pointing past this query's own position.
+            abs_pos = ctx_len + q_idx
+            causal_indices = valid_indices[valid_indices <= abs_pos]
+            causal_reference_outputs.append(_sparse_sdpa(causal_indices, q_tok))
+            global_token_idx += 1
+
+        all_q_vllm.append(q_c)
+        all_kv_c_vllm.append(kv_c_full[ctx_len:])
+        all_k_pe_vllm.append(k_pe_full[ctx_len:])
+        kv_c_contexts.append(kv_c_full[: ctx_len + 1])
+        k_pe_contexts.append(k_pe_full[: ctx_len + 1])
+
+    query_vllm = torch.cat(all_q_vllm, dim=0)
+    kv_c_vllm = torch.cat(all_kv_c_vllm, dim=0)
+    k_pe_vllm = torch.cat(all_k_pe_vllm, dim=0)
+    sdpa_reference = torch.cat(reference_outputs, dim=0)
+    causal_reference = torch.cat(causal_reference_outputs, dim=0)
+
+    vllm_config.cache_config.cache_dtype = kv_cache_dtype
+    vllm_config.model_config.hf_config.index_topk = topk_tokens
+
+    common_attn_metadata = create_common_attn_metadata(
+        batch_spec, block_size, device, arange_block_indices=True
+    )
+    kv_cache = create_and_prepopulate_kv_cache(
+        kv_c_contexts=kv_c_contexts,
+        k_pe_contexts=k_pe_contexts,
+        block_size=block_size,
+        head_size=head_size,
+        dtype=dtype,
+        device=device,
+        num_blocks=vllm_config.cache_config.num_gpu_blocks,
+        common_attn_metadata=common_attn_metadata,
+        randomize_blocks=False,
+        kv_cache_dtype=kv_cache_dtype,
+        scale=kv_cache_scale,
+    )
+
+    builder = backend_cls.get_builder_cls()(
+        kv_cache_spec, ["placeholder"], vllm_config, device
+    )
+    metadata = builder.build(
+        common_prefix_len=0, common_attn_metadata=common_attn_metadata
+    )
+
+    mock_indexer = SimpleNamespace(topk_indices_buffer=sparse_indices)
+
+    kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1).view(
+        kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim)
+    )
+    mock_kv_b_proj = ColumnParallelLinear(
+        input_size=kv_lora_rank,
+        output_size=num_heads * (qk_nope_head_dim + v_head_dim),
+        bias=False,
+    ).to(device=device, dtype=dtype)
+    mock_kv_b_proj.weight = torch.nn.Parameter(kv_b_proj_weight.T.contiguous())
+
+    with set_current_vllm_config(vllm_config):
+        impl = backend_cls.get_impl_cls()(
+            num_heads=num_heads,
+            head_size=head_size,
+            scale=scale,
+            num_kv_heads=1,
+            alibi_slopes=None,
+            sliding_window=None,
+            kv_cache_dtype=vllm_config.cache_config.cache_dtype,
+            logits_soft_cap=None,
+            attn_type="decoder",
+            kv_sharing_target_layer_name=None,
+            q_lora_rank=None,
+            kv_lora_rank=kv_lora_rank,
+            qk_nope_head_dim=qk_nope_head_dim,
+            qk_rope_head_dim=qk_rope_head_dim,
+            qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,
+            v_head_dim=v_head_dim,
+            kv_b_proj=mock_kv_b_proj,
+            indexer=mock_indexer,
+        )
+        impl.process_weights_after_loading(dtype)
+        mock_layer = MockSparseMLAAttentionLayer(
+            impl=impl,
+            num_heads=num_heads,
+            qk_nope_head_dim=qk_nope_head_dim,
+            qk_rope_head_dim=qk_rope_head_dim,
+            v_head_dim=v_head_dim,
+            kv_lora_rank=kv_lora_rank,
+            device=device,
+            W_UK=W_UK,
+            W_UV=W_UV,
+            q_scale=1.0,
+            k_scale=1.0,
+        )
+
+    out_buffer = torch.empty(
+        metadata.num_actual_tokens, num_heads * v_head_dim, dtype=dtype, device=device
+    )
+    with torch.inference_mode():
+        backend_output = mock_layer.forward_impl(
+            query_vllm, kv_c_vllm, k_pe_vllm, kv_cache, metadata, out_buffer
+        )
+    return backend_output, sdpa_reference, causal_reference
+
+
+def _skip_if_backend_unavailable(backend_cls, kv_cache_dtype: str, block_size: int):
+    if kv_cache_dtype not in backend_cls.supported_kv_cache_dtypes:
+        pytest.skip(f"{backend_cls.get_name()} does not support {kv_cache_dtype}")
+    if (
+        backend_cls is FlashMLASparseBackend
+        and kv_cache_dtype.startswith("fp8")
+        and kv_cache_dtype != "fp8_ds_mla"
+    ):
+        pytest.skip("FlashMLA Sparse fp8 only supports fp8_ds_mla kv-cache dtype")
+    if block_size not in backend_cls.get_supported_kernel_block_sizes():
+        pytest.skip(
+            f"{backend_cls.get_name()} does not support block_size={block_size}"
+        )
+    if backend_cls is FlashMLASparseBackend:
+        ok, reason = flashmla.is_flashmla_sparse_supported()
+        if not ok:
+            pytest.skip(reason)
+    elif backend_cls is FlashInferMLASparseTRTLLMBackend:
+        cap = current_platform.get_device_capability()
+        if cap is None or not backend_cls.supports_compute_capability(cap):
+            pytest.skip("FlashInferMLASparseTRTLLMBackend requires SM 10.x capability")
+
+
+@pytest.mark.parametrize(
+    "backend_cls",
+    [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend],
+    ids=["FlashMLA", "FlashInferTRTLLM"],
+)
+@pytest.mark.parametrize("config_name", list(_DSPARK_CONFIGS.keys()))
+# Per backend, the skip logic routes fp8 to the supported flavor: FlashMLA tests
+# auto + fp8_ds_mla (and skips per-tensor "fp8", which it aliases to ds_mla);
+# FlashInfer TRTLLM tests auto + per-tensor "fp8" (and skips fp8_ds_mla, which it
+# does not implement). So both backends get a bf16 case and an fp8 case.
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_ds_mla", "fp8"])
+@pytest.mark.parametrize("block_size", [64])
+# h_q=128 -> B_TOPK=128; h_q=64 -> B_TOPK=64 (covers the head64 decode path the
+# SM100 alignment assert specifically guards). Aligned widths (128/256) satisfy both.
+@pytest.mark.parametrize("num_heads", [128, 64], ids=["h128", "h64"])
+def test_dspark_noncausal_sparse_mla_matches_sdpa(
+    default_vllm_config,
+    dist_init,
+    workspace_init,
+    backend_cls,
+    config_name,
+    kv_cache_dtype,
+    block_size,
+    num_heads,
+):
+    """Non-causal (window ∪ block, future-pointing) per-token indices must match
+    a dense SDPA reference over the same indices, for both sparse-MLA backends."""
+    _skip_if_backend_unavailable(backend_cls, kv_cache_dtype, block_size)
+
+    window, block, topk_width = _DSPARK_CONFIGS[config_name]
+    device = torch.device(DEVICE_TYPE)
+
+    # Decode-style batch: each request has `block` query tokens and enough
+    # context for a full sliding window.
+    seq_lens = [window + block + 123, window + block + 50]
+    query_lens = [block, block]
+
+    sparse_indices = _build_dspark_noncausal_indices(
+        seq_lens, query_lens, window, topk_width, device
+    )
+
+    # Sanity: the construction must actually be non-causal (an early block query
+    # must reference a later block position than itself).
+    ctx0 = seq_lens[0] - query_lens[0]
+    first_query_valid = sparse_indices[0][sparse_indices[0] >= 0]
+    assert int(first_query_valid.max()) >= ctx0 + query_lens[0] - 1, (
+        "expected the first block query to attend to a future block token"
+    )
+
+    backend_output, sdpa_reference, _ = _run_sparse_backend_vs_sdpa(
+        backend_cls,
+        seq_lens,
+        query_lens,
+        sparse_indices,
+        kv_cache_dtype,
+        block_size,
+        num_heads,
+        device,
+    )
+
+    assert backend_output.shape == sdpa_reference.shape
+    assert backend_output.dtype == sdpa_reference.dtype
+    assert torch.isfinite(backend_output).all()
+    if kv_cache_dtype.startswith("fp8"):
+        rtol, atol = 0.065, 0.05
+    else:
+        rtol, atol = 0.01, 0.01
+    torch.testing.assert_close(backend_output, sdpa_reference, rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize(
+    "backend_cls",
+    [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend],
+    ids=["FlashMLA", "FlashInferTRTLLM"],
+)
+@pytest.mark.parametrize("config_name", list(_DSPARK_CONFIGS.keys()))
+@pytest.mark.parametrize("block_size", [64])
+def test_dspark_noncausal_differs_from_causal(
+    default_vllm_config,
+    dist_init,
+    workspace_init,
+    backend_cls,
+    config_name,
+    block_size,
+):
+    """Differentiation guard: prove the backend genuinely attends to the
+    future-pointing indices (not silently applying a causal mask, and not merely
+    coinciding with a causal result because future tokens carry little weight).
+
+    With random data the few future block tokens are a negligible fraction of the
+    attended set (especially with a wide window), so causal and non-causal outputs
+    are numerically indistinguishable -- that is correct physics, not a backend
+    bug. To make the check meaningful we use ``force_future_dominance`` so the last
+    block token dominates the softmax: the backend must then match the non-causal
+    reference and diverge sharply from the causal one. bf16 (``auto``) suffices;
+    the property is dtype-independent and fp8 correctness is covered above.
+    """
+    _skip_if_backend_unavailable(backend_cls, "auto", block_size)
+
+    window, block, topk_width = _DSPARK_CONFIGS[config_name]
+    device = torch.device(DEVICE_TYPE)
+    seq_lens = [window + block + 123, window + block + 50]
+    query_lens = [block, block]
+
+    sparse_indices = _build_dspark_noncausal_indices(
+        seq_lens, query_lens, window, topk_width, device
+    )
+
+    backend_output, sdpa_reference, causal_reference = _run_sparse_backend_vs_sdpa(
+        backend_cls,
+        seq_lens,
+        query_lens,
+        sparse_indices,
+        "auto",
+        block_size,
+        128,
+        device,
+        force_future_dominance=True,
+    )
+
+    # The two references must be clearly distinguishable for the check to mean
+    # anything (dominance guarantees this).
+    ref_gap = (sdpa_reference - causal_reference).abs().max().item()
+    assert ref_gap > 0.1, (
+        f"non-causal and causal references are too close (gap={ref_gap}); "
+        "force_future_dominance did not create a separable scenario"
+    )
+
+    # Backend must track the NON-causal reference, not the causal one.
+    torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.01, atol=0.01)
+    causal_err = (backend_output - causal_reference).abs().max().item()
+    assert causal_err > 0.1, (
+        f"non-causal backend output matches the causal reference "
+        f"(max abs diff={causal_err}); future-pointing indices are not attended to"
+    )

tests/v1/e2e/spec_decode/test_spec_decode.py

@@ -1412,6 +1412,67 @@ def test_dflash_acceptance_rates(
     cleanup_dist_env_and_memory()
 
 
+@pytest.fixture
+def dspark_config():
+    target_model = "Qwen/Qwen3-4B-FP8"
+    draft_model = "deepseek-ai/dspark_qwen3_4b_block7"
+
+    return dict(
+        model=target_model,
+        trust_remote_code=True,
+        speculative_config={
+            "method": "dspark",
+            "model": draft_model,
+            "num_speculative_tokens": 7,
+            "attention_backend": "FLASH_ATTN",
+            "draft_sample_method": "probabilistic",
+        },
+        max_model_len=4096,
+        disable_log_stats=False,
+    )
+
+
+@single_gpu_only
+@large_gpu_mark(min_gb=24)
+def test_dspark_correctness_and_acceptance_rate(dspark_config):
+    """
+    E2E test for DSpark speculative decoding: acceptance rate/length
+    regression coverage plus GSM8K correctness, at temperature=1.0 to
+    exercise the probabilistic draft-sampling/rejection-sampling path
+    (not just greedy).
+
+    Uses Qwen/Qwen3-4B-FP8 as target with the dspark_qwen3_4b_block7 draft
+    model. Reference: measured over 12 runs of the full GSM8K set at
+    temperature=1.0 (prefix caching disabled to avoid cross-run reuse):
+      accuracy:        min=0.782 max=0.814 mean=0.801
+      acceptance_rate: min=0.418 max=0.434 mean=0.428
+      acceptance_len:  min=3.928 max=4.037 mean=3.994
+    Thresholds set conservatively to 10% to avoid flaking due to unlucky sampling
+    """
+    spec_llm = LLM(**dspark_config)
+
+    results = evaluate_gsm8k_offline(spec_llm, temperature=1.0)
+    gsm8k_accuracy = results["accuracy"]
+
+    metrics = spec_llm.get_metrics()
+    acceptance_rate = compute_acceptance_rate(metrics)
+    acceptance_len = compute_acceptance_len(metrics)
+
+    print(
+        f"DSpark acceptance_rate={acceptance_rate:.2f}, "
+        f"acceptance_len={acceptance_len:.2f}, "
+        f"gsm8k_accuracy={gsm8k_accuracy:.3f}"
+    )
+
+    assert acceptance_rate >= 0.428 * 0.9
+    assert acceptance_len >= 3.994 * 0.9
+    assert gsm8k_accuracy >= 0.801 * 0.9
+
+    del spec_llm
+    torch.accelerator.empty_cache()
+    cleanup_dist_env_and_memory()
+
+
 @single_gpu_only
 def test_synthetic_acceptance_rate():
     """Verify that synthetic rejection sampling produces an acceptance

vllm/benchmarks/datasets/datasets.py

@@ -2452,6 +2452,7 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
                 num_requests=args.num_prompts,
                 tokenizer=tokenizer,
                 output_len=args.speed_bench_output_len,
+                skip_chat_template=args.skip_chat_template,
                 chat_template_kwargs=getattr(args, "chat_template_kwargs", None),
                 enable_multimodal_chat=args.enable_multimodal_chat,
                 request_id_prefix=args.request_id_prefix,

vllm/config/speculative.py

@@ -54,6 +54,7 @@ MTPModelTypes = Literal[
 ]
 NgramGPUTypes = Literal["ngram_gpu"]
 DFlashModelTypes = Literal["dflash"]
+DSparkModelTypes = Literal["dspark"]
 EagleModelTypes = Literal[
     "eagle", "eagle3", "extract_hidden_states", MTPModelTypes, DFlashModelTypes
 ]
@@ -66,6 +67,7 @@ SpeculativeMethod = Literal[
     "custom_class",
     EagleModelTypes,
     NgramGPUTypes,
+    DSparkModelTypes,
 ]
 RejectionSampleMethod = Literal["standard", "synthetic", "block"]
 DraftSampleMethod = Literal["greedy", "probabilistic"]
@@ -291,6 +293,7 @@ class SpeculativeConfig:
             "eagle3",
             "extract_hidden_states",
             "dflash",
+            "dspark",
         )
         factors.append(uses_aux_hidden_states)
 
@@ -608,6 +611,13 @@ class SpeculativeConfig:
                 # --quantization fp8 with a bf16 checkpoint.
                 if not self.quantization:
                     self.quantization = self.target_model_config.quantization
+            elif self.method == "dspark":
+                # DeepSeek DSpark can ship the weights inside the target checkpoint
+                if self.target_model_config is None:
+                    raise ValueError("target_model_config must be present for dspark")
+                self.model = self.target_model_config.model
+                if not self.quantization:
+                    self.quantization = self.target_model_config.quantization
             elif self.method in ("ngram", "[ngram]"):
                 self.model = "ngram"
             elif self.method == "ngram_gpu":
@@ -755,18 +765,24 @@ class SpeculativeConfig:
                         draft_hf.truncated_vocab_size = target_vocab
 
                 # Automatically detect the method
-                if self.method in ("eagle", "eagle3", "dflash"):
+                if self.method in ("eagle", "eagle3", "dflash", "dspark"):
                     pass
                 # examples:
                 # yuhuili/EAGLE-LLaMA3-Instruct-8B
                 # yuhuili/EAGLE3-LLaMA3.1-Instruct-8B
                 # AngelSlim/Qwen3-8B_eagle3
+                # deepseek-ai/dspark_qwen3_8b_block7
                 elif "eagle-" in self.draft_model_config.model.lower():
                     self.method = "eagle"
                 elif "eagle3" in self.draft_model_config.model.lower():
                     self.method = "eagle3"
                 elif "dflash" in self.draft_model_config.model.lower():
                     self.method = "dflash"
+                elif (
+                    "dspark" in self.draft_model_config.model.lower()
+                    or "Qwen3DSparkModel" in self.draft_model_config.architectures
+                ):
+                    self.method = "dspark"
                 elif self.draft_model_config.hf_config.model_type == "medusa":
                     self.method = "medusa"
                 elif self.draft_model_config.hf_config.model_type == "mlp_speculator":
@@ -813,7 +829,18 @@ class SpeculativeConfig:
                         self.draft_model_config.hf_config = eagle_config
                         self.update_arch_()
 
-                if self.method == "dflash":
+                if self.method == "dspark" and (
+                    "Qwen3DSparkModel" not in self.draft_model_config.architectures
+                ):
+                    # DeepSeek-V4 DSpark reuses the full DeepSeek-V4 config
+                    # and its weights ship in the target checkpoint.
+                    self.draft_model_config.hf_config.model_type = "deepseek_v4"
+                    self.draft_model_config.hf_config.architectures = [
+                        "DSparkDraftModel"
+                    ]
+                    self.update_arch_()
+
+                if self.method in ("dflash", "dspark"):
                     self.parallel_drafting = True
 
                 if self.num_speculative_tokens is not None and hasattr(
@@ -1129,11 +1156,17 @@ class SpeculativeConfig:
         )
 
     def use_eagle(self) -> bool:
-        return self.method in ("eagle", "eagle3", "mtp", "dflash")
+        # NOTE: This method is usually a stand-in for "speculative decoding using
+        # target model hidden states"
+        # TODO(ben): Refactor this so the naming is clearer
+        return self.method in ("eagle", "eagle3", "mtp", "dflash", "dspark")
 
     def use_dflash(self) -> bool:
         return self.method == "dflash"
 
+    def use_dspark(self) -> bool:
+        return self.method == "dspark"
+
     def uses_dynamic_speculative_decoding(self) -> bool:
         return self.num_speculative_tokens_per_batch_size is not None
 

vllm/config/vllm.py

@@ -524,6 +524,16 @@ class VllmConfig:
         if use_v2_model_runner is not None:
             return use_v2_model_runner
 
+        # DSpark is implemented only by the V2 GPU model runner, and DeepSeek-V4
+        # is not otherwise a default-V2 architecture, so force V2 for it. If V2
+        # is unsupported for the rest of the config, _validate_v2_model_runner
+        # raises rather than silently falling back to V1 (which can't run dspark).
+        if (
+            self.speculative_config is not None
+            and self.speculative_config.method == "dspark"
+        ):
+            return True
+
         if self.model_config is not None and self.model_config.is_diffusion:
             return True
 
@@ -958,10 +968,11 @@ class VllmConfig:
                     self.speculative_config.method not in get_args(EagleModelTypes)
                     and self.speculative_config.method not in get_args(NgramGPUTypes)
                     and self.speculative_config.method != "draft_model"
+                    and self.speculative_config.method != "dspark"
                 ):
                     raise ValueError(
                         "Currently, async scheduling is only supported "
-                        "with EAGLE/MTP/Draft Model/NGram GPU kind of "
+                        "with EAGLE/MTP/Draft Model/NGram GPU/DSpark kind of "
                         "speculative decoding"
                     )
                 if self.speculative_config.disable_padded_drafter_batch:
@@ -989,6 +1000,7 @@ class VllmConfig:
                 self.speculative_config is not None
                 and self.speculative_config.method not in get_args(EagleModelTypes)
                 and self.speculative_config.method not in get_args(NgramGPUTypes)
+                and self.speculative_config.method != "dspark"
             ):
                 logger.warning_once(
                     "Async scheduling not supported with %s-based "
@@ -2038,17 +2050,24 @@ class VllmConfig:
             # TODO: ngram / ngram_gpu are not supported by the v2 model runner yet
             if speculative_config.method in ("ngram", "ngram_gpu"):
                 unsupported.append("ngram/ngram_gpu speculative decoding")
-            elif speculative_config.method not in ("eagle", "eagle3", "mtp", "dflash"):
+            elif speculative_config.method not in (
+                "eagle",
+                "eagle3",
+                "mtp",
+                "dflash",
+                "dspark",
+            ):
                 unsupported.append(f"speculative method '{speculative_config.method}'")
 
             if speculative_config.uses_dynamic_speculative_decoding():
                 unsupported.append("dynamic speculative decoding")
 
-            # V2 EagleSpeculator does not support parallel_drafting (for P-Eagle)
-            # DFlash uses parallel drafting natively in V2 via DFlashSpeculator.
+            # V2 EagleSpeculator does not support parallel_drafting (for P-Eagle).
+            # DFlash and DSpark use parallel drafting natively in V2 via their
+            # own speculators.
             if (
                 speculative_config.parallel_drafting
-                and speculative_config.method != "dflash"
+                and speculative_config.method not in ("dflash", "dspark")
             ):
                 unsupported.append("parallel drafting for EAGLE speculative decoding")
 

vllm/model_executor/models/qwen3_dflash.py

@@ -472,7 +472,7 @@ class DFlashQwen3Model(nn.Module):
         self,
         context_states: torch.Tensor,
         context_positions: torch.Tensor,
-        context_slot_mapping: torch.Tensor | None = None,
+        context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
     ) -> None:
         """Precompute K/V for context states write them into each layer's KV cache.
 
@@ -551,7 +551,13 @@ class DFlashQwen3Model(nn.Module):
 
         # --- Per-layer cache insert ---
         all_k_final = all_k_flat.view(L, num_ctx, nkv, hd)
+        per_layer = isinstance(context_slot_mapping, (list, tuple))
         for i in range(L):
+            slot_mapping = (
+                context_slot_mapping[i] if per_layer else context_slot_mapping
+            )
+            if slot_mapping is None:
+                continue  # dummy run: skip cache ops
             attn = self._attn_layers[i]
             kv_cache = attn.kv_cache
             attn.impl.do_kv_cache_update(
@@ -559,7 +565,7 @@ class DFlashQwen3Model(nn.Module):
                 all_k_final[i],
                 all_v[i],
                 kv_cache,
-                context_slot_mapping,
+                slot_mapping,
             )
 
     def forward(
@@ -701,7 +707,7 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):
         self,
         context_states: torch.Tensor,
         context_positions: torch.Tensor,
-        context_slot_mapping: torch.Tensor | None = None,
+        context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
     ) -> None:
         """Precompute projected + RoPE'd K/V and write to cache."""
         self.model.precompute_and_store_context_kv(

vllm/model_executor/models/qwen3_dspark.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Qwen3 DSpark draft model for semi-autoregressive drafting.
+
+DSpark drafts a whole block in one parallel pass (DFlash-style: context-KV
+precompute + a non-causal query-block forward) and then injects intra-block
+dependency with a lightweight sequential Markov head.
+
+The parallel backbone is a standard Qwen3 decoder stack reused from the
+DFlash Qwen3 draft (see qwen3_dflash.py). DSpark adds:
+  * ``markov_head``: low-rank V x r / r x V transition bias added to the base
+    logits, sampled left-to-right by the speculator (the sequential stage).
+
+DSparkMarkovHead is shared with the DSV4-style DSpark model.
+"""
+
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+
+from .qwen3_dflash import DFlashQwen3ForCausalLM, DFlashQwen3Model
+from .utils import AutoWeightsLoader, maybe_prefix, process_eagle_weight
+
+logger = init_logger(__name__)
+
+
+class DSparkMarkovHead(nn.Module):
+    """Sequential transition-bias head (low-rank V x r, r x V).
+
+    ``markov_w1[token]`` is an r-dim embedding of the previously sampled token;
+    ``markov_w2`` projects it back to a vocab-size bias added to the base logits.
+    """
+
+    def __init__(self, vocab_size: int, markov_rank: int, prefix: str) -> None:
+        super().__init__()
+        # TODO(ben): profile for which (if any) it makes sense to replicate or TP-shard
+        self.markov_w1 = VocabParallelEmbedding(
+            vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w1")
+        )
+        self.markov_w2 = ParallelLMHead(
+            vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w2")
+        )
+
+    def embed(self, token_ids: torch.Tensor) -> torch.Tensor:
+        """r-dim Markov embedding of ``token_ids`` ([B] -> [B, r])."""
+        return self.markov_w1(token_ids)
+
+    def bias(self, markov_embed: torch.Tensor, logits_processor) -> torch.Tensor:
+        """Vocab-size transition bias from a Markov embedding ([B, r] -> [B, V])."""
+        return logits_processor(self.markov_w2, markov_embed)
+
+
+class Qwen3DSparkModel(DFlashQwen3Model):
+    """DFlash Qwen3 backbone + DSpark Markov head."""
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        start_layer_id: int = 0,
+        prefix: str = "",
+    ) -> None:
+        super().__init__(
+            vllm_config=vllm_config, start_layer_id=start_layer_id, prefix=prefix
+        )
+        config = self.config
+        self.markov_head = DSparkMarkovHead(
+            config.vocab_size,
+            config.markov_rank,
+            prefix=maybe_prefix(prefix, "markov_head"),
+        )
+
+
+class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        nn.Module.__init__(self)
+        self.draft_model_config = vllm_config.speculative_config.draft_model_config
+        self.config = self.draft_model_config.hf_config
+        if getattr(self.config, "draft_vocab_size", None) is None:
+            self.config.draft_vocab_size = getattr(self.config, "vocab_size", None)
+        target_layer_num = vllm_config.model_config.get_num_layers(
+            vllm_config.parallel_config
+        )
+        self.model = Qwen3DSparkModel(
+            vllm_config=vllm_config,
+            prefix=maybe_prefix(prefix, "model"),
+            start_layer_id=target_layer_num,
+        )
+
+        logit_scale = getattr(self.config, "logit_scale", 1.0)
+        self.lm_head = ParallelLMHead(
+            self.config.draft_vocab_size,
+            self.config.hidden_size,
+            prefix=maybe_prefix(prefix, "lm_head"),
+        )
+        self.logits_processor = LogitsProcessor(
+            self.config.draft_vocab_size, scale=logit_scale
+        )
+        target_vocab_size = vllm_config.model_config.get_vocab_size()
+        if self.config.draft_vocab_size != target_vocab_size:
+            self.draft_id_to_target_id = nn.Parameter(
+                torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
+                requires_grad=False,
+            )
+        else:
+            self.draft_id_to_target_id = None
+
+    def get_draft_kv_cache_layer_names(self) -> list[str]:
+        return [layer.self_attn.attn.layer_name for layer in self.model.layers]
+
+    def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.markov_head.embed(token_ids)
+
+    def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
+        return self.model.markov_head.bias(markov_embed, self.logits_processor)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+        model_weights = {}
+        includes_embed_tokens = False
+        includes_lm_head = False
+        for name, loaded_weight in weights:
+            if "lm_head" not in name:
+                name = "model." + name
+            if "embed_tokens" in name:
+                includes_embed_tokens = True
+            if "lm_head" in name:
+                includes_lm_head = True
+            model_weights[name] = loaded_weight
+            # Sets has_own_embed_tokens / has_own_lm_head so load_dspark_model
+            # knows whether to keep these or alias the target's.
+            process_eagle_weight(self, name)
+
+        # mask_embedding is an unused placeholder param; DSpark masks via the vocab row.
+        # confidence_head is not wired into inference yet; skip its weights.
+        # embed_tokens / lm_head are optional; when omitted they are shared from
+        # the target by load_dspark_model, so skip the unloaded params here.
+        skip_substrs = ["mask_embedding", "confidence_head"]
+        if not includes_embed_tokens:
+            skip_substrs.append("embed_tokens")
+        if not includes_lm_head:
+            skip_substrs.append("lm_head")
+        loader = AutoWeightsLoader(self, skip_substrs=skip_substrs)
+        loader.load_weights(model_weights.items())
+        self.model._build_fused_kv_buffers()

vllm/model_executor/models/registry.py

@@ -585,6 +585,8 @@ _SPECULATIVE_DECODING_MODELS = {
     "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),
     "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
     "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
+    "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"),
+    "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
     "PEagleDraftModel": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
     "PeagleLlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
     "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),

vllm/models/deepseek_v4/__init__.py

@@ -17,14 +17,23 @@ from .quant_config import DeepseekV4FP8Config
 if current_platform.is_rocm():
     from .amd.model import DeepseekV4ForCausalLM
     from .amd.mtp import DeepSeekV4MTP
+
+    # DSpark is NVIDIA-only for now.
+    DSparkDeepseekV4ForCausalLM = None  # type: ignore[assignment]
 elif current_platform.is_xpu():
     from .xpu.model import DeepseekV4ForCausalLM  # type: ignore[assignment]
     from .xpu.mtp import DeepSeekV4MTP  # type: ignore[assignment]
+
+    DSparkDeepseekV4ForCausalLM = None  # type: ignore[assignment]
 else:
+    from .nvidia.dspark import (  # type: ignore[assignment]
+        DSparkDeepseekV4ForCausalLM,
+    )
     from .nvidia.model import DeepseekV4ForCausalLM  # type: ignore[assignment]
     from .nvidia.mtp import DeepSeekV4MTP  # type: ignore[assignment]
 
 __all__ = [
+    "DSparkDeepseekV4ForCausalLM",
     "DeepSeekV4MTP",
     "DeepseekV4FP8Config",
     "DeepseekV4ForCausalLM",

vllm/models/deepseek_v4/nvidia/dspark.py

@@ -0,0 +1,477 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DSpark draft model for DeepSeek-V4 (semi-autoregressive speculative decoding).
+
+See: qwen3_dspark.py for base architecture. This one is specialized to the DSV4 DSpark,
+which reuses the target model's architecture similarly to MTP.
+
+To implement non-causal attention, we leverage the sparse attention implementation to
+include the future query tokens in the top-k indices for each query token.
+"""
+
+from collections.abc import Iterable
+
+import regex as re
+import torch
+import torch.nn as nn
+
+from vllm.config import VllmConfig, get_current_vllm_config
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.kernels.mhc.tilelang import (
+    hc_head_fused_kernel_tilelang,
+    mhc_post_tilelang,
+)
+from vllm.model_executor.layers.fused_moe import (
+    fused_moe_make_expert_params_mapping,
+)
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.qwen3_dspark import (
+    DSparkMarkovHead,
+)
+from vllm.model_executor.models.utils import maybe_prefix
+
+from .model import (
+    DeepseekV4DecoderLayer,
+    make_deepseek_v4_expert_params_mapping,
+)
+
+logger = init_logger(__name__)
+
+# MoE expert scale suffix differs by expert dtype (mirrors deepseek_v4 loaders):
+# fp4 experts register ``.weight_scale``; block-fp8 experts ``.weight_scale_inv``.
+_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")
+
+
+class DSparkDeepseekV4Model(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        assert vllm_config.speculative_config is not None
+        config = vllm_config.speculative_config.draft_model_config.hf_config
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.hc_mult = config.hc_mult
+        self.hc_eps = config.hc_eps
+        self.rms_norm_eps = config.rms_norm_eps
+        self.num_hidden_layers = config.num_hidden_layers
+        self.target_layer_ids = tuple(config.dspark_target_layer_ids)
+
+        self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3
+
+        # Shared with the target (aliased by the speculator's loading utility).
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+            prefix=maybe_prefix(prefix, "embed_tokens"),
+        )
+
+        self.main_proj = ReplicatedLinear(
+            config.hidden_size * len(self.target_layer_ids),
+            config.hidden_size,
+            bias=False,
+            return_bias=False,
+            quant_config=vllm_config.quant_config,
+            prefix=maybe_prefix(prefix, "main_proj"),
+        )
+        self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        current_vllm_config = get_current_vllm_config()
+        self.layers = nn.ModuleList(
+            [
+                DeepseekV4DecoderLayer(
+                    current_vllm_config,
+                    prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"),
+                )
+                for i in range(self.num_dspark_layers)
+            ]
+        )
+
+        # Heads: final norm + hc_head, and the Markov head
+        # Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        hc_dim = self.hc_mult * config.hidden_size
+        self.hc_head_fn = nn.Parameter(
+            torch.empty(self.hc_mult, hc_dim, dtype=torch.float32),
+            requires_grad=False,
+        )
+        self.hc_head_base = nn.Parameter(
+            torch.empty(self.hc_mult, dtype=torch.float32), requires_grad=False
+        )
+        self.hc_head_scale = nn.Parameter(
+            torch.empty(1, dtype=torch.float32), requires_grad=False
+        )
+        self.markov_head = DSparkMarkovHead(
+            config.vocab_size,
+            config.dspark_markov_rank,
+            prefix=maybe_prefix(prefix, "markov_head"),
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor:
+        """main_x = main_norm(main_proj(concat of target aux hidden states)).
+
+        ``aux_hidden_states`` is [T, hidden_size * len(target_layer_ids)].
+        """
+        return self.main_norm(self.main_proj(aux_hidden_states))
+
+    @torch.inference_mode()
+    def precompute_and_store_context_kv(
+        self,
+        main_x: torch.Tensor,
+        context_positions: torch.Tensor,
+        context_slot_mappings: list[torch.Tensor | None] | None = None,
+    ) -> None:
+        """Insert the sliding-window context KV for every draft layer.
+
+        Mirrors the reference DSparkAttention: each layer derives its context KV
+        from the SAME projected target hidden ``main_x``, via that layer's own
+        ``wkv`` + ``kv_norm`` + RoPE + quant, then writes it at the
+        layer's context slots.
+
+        ``context_slot_mappings`` is a per-layer list (each entry is the context
+        slot mapping for that layer's kv-cache group, since the hybrid manager may
+        place draft layers in different groups). ``None`` (or a ``None`` entry)
+        runs the projection to reserve workspace but writes nothing (profiling).
+        """
+        for i, layer in enumerate(self.layers):
+            slot_mapping = (
+                None if context_slot_mappings is None else context_slot_mappings[i]
+            )
+            attn = layer.attn
+            # Optimized DSV4 MLA path: wkv part of the fused wq_a|wkv projection
+            # (q_lora part discarded), then RoPE/quant/insert via the fused op.
+            qr_kv, _ = attn.fused_wqa_wkv(main_x)
+            kv = qr_kv[..., attn.q_lora_rank :]
+            kv = attn.kv_norm(kv)
+            if slot_mapping is None:
+                continue
+            _insert_context_kv(attn, kv, context_positions, slot_mapping)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_input_ids(input_ids)
+        # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]).
+        hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1)
+
+        residual = post_mix = res_mix = None
+        for layer in self.layers:
+            hidden_states, residual, post_mix, res_mix = layer(
+                hidden_states,
+                positions,
+                input_ids,
+                post_mix,
+                res_mix,
+                residual,
+            )
+        hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix)
+        # hc_head reduces the hc copies; return the PRE-norm head hidden
+        hidden_states = hc_head_fused_kernel_tilelang(
+            hidden_states,
+            self.hc_head_fn,
+            self.hc_head_scale,
+            self.hc_head_base,
+            self.rms_norm_eps,
+            self.hc_eps,
+        )
+        return hidden_states
+
+
+def _insert_context_kv(
+    attn: nn.Module,
+    kv: torch.Tensor,
+    positions: torch.Tensor,
+    slot_mapping: torch.Tensor,
+) -> None:
+    """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV.
+
+    Reuses the DSV4 fused insert ops (which also process a query; we pass a dummy
+    query and discard it, since context tokens have no query). Mirrors
+    ``DeepseekV4Attention._fused_qnorm_rope_kv_insert``.
+    """
+    swa_cache = attn.swa_cache_layer.kv_cache
+    block_size = attn.swa_cache_layer.block_size
+    cos_sin_cache = attn.rotary_emb.cos_sin_cache
+    cache_dtype = swa_cache.dtype
+    n_ctx = kv.shape[0]
+    dummy_q = torch.zeros(
+        (n_ctx, attn.n_local_heads, attn.head_dim),
+        dtype=kv.dtype,
+        device=kv.device,
+    )
+    if cache_dtype == torch.uint8:
+        # fp8_ds_mla UE8M0 paged layout
+        swa_2d = swa_cache.view(swa_cache.shape[0], -1)
+        torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+            dummy_q,
+            kv,
+            swa_2d,
+            slot_mapping,
+            positions,
+            cos_sin_cache,
+            attn.padded_heads,
+            attn.eps,
+            block_size,
+        )
+    elif cache_dtype == torch.bfloat16:
+        swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
+        torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
+            dummy_q,
+            kv,
+            swa_3d,
+            slot_mapping,
+            positions,
+            cos_sin_cache,
+            attn.eps,
+            block_size,
+        )
+    else:  # per-tensor fp8 (torch.float8_e4m3fn)
+        # TODO(ben): double-check if this is being dispatched correctly for FI backend
+        swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
+        dummy_q_fp8 = torch.zeros_like(dummy_q, dtype=torch.float8_e4m3fn)
+        torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
+            dummy_q,
+            kv,
+            dummy_q_fp8,
+            swa_3d,
+            slot_mapping,
+            positions,
+            cos_sin_cache,
+            attn._flashinfer_fp8_kv_scale,
+            attn._flashinfer_fp8_q_scale_inv,
+            attn.eps,
+            block_size,
+        )
+
+
+class DSparkDeepseekV4ForCausalLM(nn.Module):
+    # Draft weights ship in the target checkpoint (mtp.*) without embed/head, so
+    # load_dspark_model always aliases the target's.
+    has_own_embed_tokens = False
+    has_own_lm_head = False
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        assert vllm_config.speculative_config is not None
+        self.draft_model_config = vllm_config.speculative_config.draft_model_config
+        self.config = self.draft_model_config.hf_config
+        self.model = DSparkDeepseekV4Model(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+        # Shared with the target (aliased by the speculator's load utility).
+        self.lm_head = ParallelLMHead(
+            self.config.vocab_size,
+            self.config.hidden_size,
+            prefix=maybe_prefix(prefix, "lm_head"),
+        )
+        self.logits_processor = LogitsProcessor(self.config.vocab_size)
+
+    # --- Hooks used by the speculator -------------------------------------
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.model.combine_hidden_states(aux_hidden_states)
+
+    def get_draft_kv_cache_layer_names(self) -> list[str]:
+        # DSV4 MLA path: each draft layer's sliding-window cache is a separate
+        # layer, named by its prefix.
+        return [layer.attn.swa_cache_layer.prefix for layer in self.model.layers]
+
+    def precompute_and_store_context_kv(
+        self,
+        context_states: torch.Tensor,
+        context_positions: torch.Tensor,
+        context_slot_mappings: list[torch.Tensor | None] | None = None,
+    ) -> None:
+        self.model.precompute_and_store_context_kv(
+            context_states, context_positions, context_slot_mappings
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # Returns the pre-norm hc_head hidden ([T, hidden_size]).
+        return self.model(input_ids, positions, inputs_embeds)
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Base logits U_k = lm_head(norm(head_hidden))."""
+        return self.logits_processor(self.lm_head, self.model.norm(hidden_states))
+
+    def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.markov_head.embed(token_ids)
+
+    def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
+        return self.model.markov_head.bias(markov_embed, self.logits_processor)
+
+    # --- Weight loading ----------------------------------------------------
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load the ``mtp.{0,1,2}.*`` draft weights from the target checkpoint.
+
+        Non-mtp weights (embed/head/main layers) belong to the target model and
+        are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target.
+        """
+        first_layer = self.model.layers[0]
+        use_mega_moe = first_layer.ffn.use_mega_moe
+        if use_mega_moe:
+            expert_mapping = make_deepseek_v4_expert_params_mapping(
+                self.config.n_routed_experts
+            )
+        else:
+            expert_mapping = fused_moe_make_expert_params_mapping(
+                self,
+                ckpt_gate_proj_name="w1",
+                ckpt_down_proj_name="w2",
+                ckpt_up_proj_name="w3",
+                num_experts=self.config.n_routed_experts,
+            )
+        expert_scale_suffix = (
+            ".weight_scale"
+            if getattr(self.config, "expert_dtype", "fp4") == "fp4"
+            else ".weight_scale_inv"
+        )
+
+        # (param_name, ckpt_shard_name, shard_id) for non-expert stacked params.
+        stacked_params_mapping = [
+            ("gate_up_proj", "w1", 0),
+            ("gate_up_proj", "w3", 1),
+            ("attn.fused_wqa_wkv", "attn.wq_a", 0),
+            ("attn.fused_wqa_wkv", "attn.wkv", 1),
+        ]
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        tp_size = get_tensor_model_parallel_world_size()
+        tp_rank = get_tensor_model_parallel_rank()
+        n_local_head = self.config.num_attention_heads // tp_size
+        head_start = n_local_head * tp_rank
+        head_end = n_local_head * (tp_rank + 1)
+
+        for name, loaded_weight in weights:
+            mapped = self._remap_dspark_name(name)
+            if mapped is None:
+                continue
+            name = mapped
+
+            # ``.scale`` -> per-method scale suffix.
+            if name.endswith(".scale"):
+                suffix = (
+                    expert_scale_suffix
+                    if _EXPERT_SCALE_RE.search(name)
+                    else ".weight_scale_inv"
+                )
+                name = name.removesuffix(".scale") + suffix
+
+            # E8M0 expert scales: keep raw exponent bytes.
+            if ".experts." in name:
+                if (
+                    "weight_scale" in name
+                    and loaded_weight.dtype == torch.float8_e8m0fnu
+                ):
+                    loaded_weight = loaded_weight.view(torch.uint8)
+                for param_name, weight_name, expert_id, shard_id in expert_mapping:
+                    if weight_name not in name:
+                        continue
+                    name_mapped = name.replace(weight_name, param_name)
+                    param = params_dict[name_mapped]
+                    success = param.weight_loader(
+                        param,
+                        loaded_weight,
+                        name_mapped,
+                        shard_id=shard_id,
+                        expert_id=expert_id,
+                        return_success=True,
+                    )
+                    if success:
+                        loaded_params.add(name_mapped)
+                        break
+                continue
+
+            # Stacked rules only apply to decoder-layer weights. Head-stack params
+            # (main_proj/norm/hc_head/markov_head) load directly — otherwise e.g.
+            # "markov_w1" would collide with the "w1" shard rule.
+            is_layer_param = name.startswith("model.layers.")
+            for param_name, weight_name, stacked_shard_id in stacked_params_mapping:
+                if not is_layer_param or weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                param = params_dict[name]
+                param.weight_loader(param, loaded_weight, stacked_shard_id)
+                loaded_params.add(name)
+                break
+            else:
+                if "attn_sink" in name:
+                    narrow = loaded_weight[head_start:head_end]
+                    params_dict[name][: narrow.shape[0]].copy_(narrow)
+                    loaded_params.add(name)
+                    continue
+                if ".shared_experts.w2" in name:
+                    name = name.replace(
+                        ".shared_experts.w2", ".shared_experts.down_proj"
+                    )
+                if name.endswith(".ffn.gate.bias"):
+                    name = name.replace(
+                        ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias"
+                    )
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+                loaded_params.add(name)
+
+        self._finalize_moe()
+        logger.info_once("DSpark draft model loaded: %d params", len(loaded_params))
+        return loaded_params
+
+    def _finalize_moe(self) -> None:
+        for layer in self.model.layers:
+            layer.ffn.finalize_mega_moe_weights()
+
+    def _remap_dspark_name(self, name: str) -> str | None:
+        """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path.
+
+        Returns None for non-mtp weights (owned by the target model).
+        """
+        m = re.match(r"mtp\.(\d+)\.(.*)", name)
+        if m is None:
+            return None
+        stage = int(m.group(1))
+        rest = m.group(2)
+        # The confidence head is not wired into inference yet; drop its weights.
+        if rest.startswith("confidence_head."):
+            return None
+        # Head-stack params live at model level (mtp.last), context combiner at
+        # model level (mtp.0); everything else is a per-layer decoder block.
+        head_prefixes = (
+            "norm.",
+            "hc_head_fn",
+            "hc_head_base",
+            "hc_head_scale",
+            "markov_head.",
+        )
+        if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith(
+            head_prefixes
+        ):
+            return f"model.{rest}"
+        return f"model.layers.{stage}.{rest}"

vllm/models/deepseek_v4/nvidia/model.py

@@ -48,7 +48,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
-from vllm.model_executor.models.interfaces import MixtureOfExperts, SupportsPP
+from vllm.model_executor.models.interfaces import (
+    EagleModelMixin,
+    MixtureOfExperts,
+    SupportsEagle3,
+    SupportsPP,
+)
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     PPMissingLayer,
@@ -933,7 +938,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         return x, residual, post_mix, res_mix
 
 
-class DeepseekV4Model(nn.Module):
+class DeepseekV4Model(nn.Module, EagleModelMixin):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
 
@@ -1074,7 +1079,12 @@ class DeepseekV4Model(nn.Module):
             input_ids = input_ids.to(torch.int64)
 
         residual, post_mix, res_mix = None, None, None
-        for layer in islice(self.layers, self.start_layer, self.end_layer):
+        aux_hidden_states: list[torch.Tensor] = []
+        final_aux_recon: torch.Tensor | None = None  # avoid duplicate mhc_post call
+        for idx, layer in enumerate(
+            islice(self.layers, self.start_layer, self.end_layer),
+            start=self.start_layer,
+        ):
             hidden_states, residual, post_mix, res_mix = layer(
                 hidden_states,
                 positions,
@@ -1083,10 +1093,21 @@ class DeepseekV4Model(nn.Module):
                 res_mix,
                 residual,
             )
+            if idx + 1 in self.aux_hidden_state_layers:
+                # Reconstruct the aux hidden state for draft models
+                aux_recon = mhc_post_tilelang(
+                    hidden_states, residual, post_mix, res_mix
+                )
+                aux_hidden_states.append(aux_recon.mean(dim=1))
+                final_aux_recon = aux_recon
         if layer is not None:
-            hidden_states = mhc_post_tilelang(
-                hidden_states, residual, post_mix, res_mix
-            )
+            # Reuse if the last layer was captured as an aux hidden state
+            if self.end_layer in self.aux_hidden_state_layers:
+                hidden_states = final_aux_recon
+            else:
+                hidden_states = mhc_post_tilelang(
+                    hidden_states, residual, post_mix, res_mix
+                )
 
         if not get_pp_group().is_last_rank:
             return IntermediateTensors({"hidden_states": hidden_states})
@@ -1104,6 +1125,8 @@ class DeepseekV4Model(nn.Module):
             self.hc_eps,
         )
         hidden_states = self.norm(hidden_states)
+        if len(aux_hidden_states) > 0:
+            return hidden_states, aux_hidden_states
         return hidden_states
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -1330,7 +1353,9 @@ class DeepseekV4MixtureOfExperts(MixtureOfExperts):
             moe.experts.update_expert_map()
 
 
-class DeepseekV4ForCausalLM(nn.Module, SupportsPP, DeepseekV4MixtureOfExperts):
+class DeepseekV4ForCausalLM(
+    nn.Module, SupportsPP, SupportsEagle3, DeepseekV4MixtureOfExperts
+):
     model_cls = DeepseekV4Model
 
     # Default mapper assumes the original FP4-expert checkpoint layout.

vllm/v1/attention/backends/mla/sparse_swa.py

@@ -306,16 +306,19 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
         )
 
         # Handle MTP: adjust decode_threshold like the indexer does
+        spec_config = self.vllm_config.speculative_config
         self.num_speculative_tokens = (
-            self.vllm_config.speculative_config.num_speculative_tokens
-            if self.vllm_config.speculative_config
-            else 0
+            spec_config.num_speculative_tokens if spec_config else 0
+        )
+        # Decode can have query_len up to
+        #   1 + (2 if parallel drafting else 1) * num_speculative_tokens.
+        # This MUST match the flashmla_sparse / indexer threshold so that
+        # all backends agree on the decode/prefill split.
+        spec_mult = (
+            2 if (spec_config is not None and spec_config.parallel_drafting) else 1
         )
-        # With MTP, decode can have query_len up to 1 + num_speculative_tokens.
-        # Must match the threshold used by the indexer and flashmla_sparse so
-        # that all backends agree on the decode/prefill split.
         self.decode_threshold = (
-            self.reorder_batch_threshold + self.num_speculative_tokens
+            self.reorder_batch_threshold + spec_mult * self.num_speculative_tokens
         )
 
         hf_config = self.vllm_config.model_config.hf_config
@@ -368,6 +371,20 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
             device=self.device,
         )
 
+        # DSpark draft: the block is non-causal (every query attends to the
+        # trailing window of context PLUS all query tokens, including future ones),
+        # so its per-token index list is wider than `window_size`. The kernel pads
+        # the q-head count to B_TOPK (64/128), which requires the index width to be
+        # a multiple of 128.
+        self.is_dspark = spec_config is not None and spec_config.use_dspark()
+        self.noncausal_index_width = (
+            cdiv(self.window_size + self.num_speculative_tokens, 128) * 128
+            if self.is_dspark
+            else 0
+        )
+        self.decode_swa_indices_noncausal: torch.Tensor | None = None
+        self._max_tokens = max_tokens
+
     def build(
         self,
         common_prefix_len: int,
@@ -407,23 +424,56 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
         is_valid_token = self.is_valid_token[: slot_mapping.shape[0]]
         is_valid_token.copy_(slot_mapping >= 0)
 
+        non_causal = not common_attn_metadata.causal
+        decode_swa_indices = self.decode_swa_indices
         if num_decode_tokens > 0:
             self.decode_swa_lens[num_decode_tokens:] = 0
-            _compute_swa_indices_and_lens_kernel[(num_decode_tokens,)](
-                self.decode_swa_indices,
-                self.decode_swa_indices.stride(0),
-                self.decode_swa_lens,
-                self.window_size,
-                query_start_loc,
-                seq_lens,
-                token_to_req_indices,
-                is_valid_token,
-                block_table,
-                block_table.stride(0),
-                self.block_size,
-                token_offset=0,
-                TRITON_BLOCK_SIZE=1024,
-            )
+            if non_causal:
+                assert self.is_dspark, (
+                    "Non-causal DeepseekV4 SWA is only supported for the DSpark "
+                    "speculation mode, but causal=False was set without DSpark."
+                )
+                if self.decode_swa_indices_noncausal is None:
+                    self.decode_swa_indices_noncausal = torch.zeros(
+                        self._max_tokens,
+                        1,
+                        self.noncausal_index_width,
+                        dtype=torch.int32,
+                        device=self.device,
+                    )
+                decode_swa_indices = self.decode_swa_indices_noncausal
+                _compute_dspark_noncausal_swa_indices_kernel[(num_decode_tokens,)](
+                    decode_swa_indices,
+                    decode_swa_indices.stride(0),
+                    self.decode_swa_lens,
+                    self.window_size,
+                    self.noncausal_index_width,
+                    query_start_loc,
+                    seq_lens,
+                    token_to_req_indices,
+                    is_valid_token,
+                    block_table,
+                    block_table.stride(0),
+                    self.block_size,
+                    token_offset=0,
+                    TRITON_BLOCK_SIZE=1024,
+                )
+            else:
+                _compute_swa_indices_and_lens_kernel[(num_decode_tokens,)](
+                    decode_swa_indices,
+                    decode_swa_indices.stride(0),
+                    self.decode_swa_lens,
+                    self.window_size,
+                    query_start_loc,
+                    seq_lens,
+                    token_to_req_indices,
+                    is_valid_token,
+                    block_table,
+                    block_table.stride(0),
+                    self.block_size,
+                    token_offset=0,
+                    TRITON_BLOCK_SIZE=1024,
+                )
 
         # Prefill SWA indices live in paged coordinates. `token_offset` lets
         # the kernel read is_valid_token / token_to_req_indices at absolute
@@ -471,7 +521,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
             slot_mapping=slot_mapping,
             is_valid_token=is_valid_token,
             token_to_req_indices=token_to_req_indices,
-            decode_swa_indices=self.decode_swa_indices[:num_decode_tokens],
+            decode_swa_indices=decode_swa_indices[:num_decode_tokens],
             decode_swa_lens=self.decode_swa_lens[:num_decode_tokens],
             prefill_swa_indices=(
                 self.prefill_swa_indices[:num_prefill_tokens]
@@ -665,3 +715,69 @@ def _compute_swa_indices_and_lens_kernel(
             slot_ids,
             mask=offset < window_size,
         )
+
+
+# TODO(ben): unify this kernel to reduce duplication
+@triton.jit(do_not_specialize=["token_offset"])
+def _compute_dspark_noncausal_swa_indices_kernel(
+    swa_indices_ptr,
+    swa_indices_stride,
+    swa_lens_ptr,
+    window_size,
+    index_width,
+    query_start_loc_ptr,
+    seq_lens_ptr,
+    token_to_req_indices_ptr,
+    is_valid_token_ptr,
+    block_table_ptr,
+    block_table_stride,
+    block_size,
+    token_offset,
+    TRITON_BLOCK_SIZE: tl.constexpr,
+):
+    """Non-causal per-token indices for the DSpark draft block.
+
+    Here, we populate the topk indices with the trailing window of context tokens,
+    plus all query tokens (including future ones).
+    """
+    pid = tl.program_id(0)
+    token_idx = pid + token_offset
+    is_valid = tl.load(is_valid_token_ptr + token_idx)
+    if not is_valid:
+        tl.store(swa_lens_ptr + pid, 0)
+        return
+
+    req_idx = tl.load(token_to_req_indices_ptr + token_idx)
+
+    query_start = tl.load(query_start_loc_ptr + req_idx)
+    query_end = tl.load(query_start_loc_ptr + req_idx + 1)
+    query_len = query_end - query_start
+
+    seq_len = tl.load(seq_lens_ptr + req_idx)
+    prefix_len = seq_len - query_len
+
+    # Block-anchored window (shared by every token in the block) + full block.
+    start_pos = tl.maximum(prefix_len - window_size, 0)
+    end_pos = seq_len
+
+    swa_len = end_pos - start_pos
+    tl.store(swa_lens_ptr + pid, swa_len)
+
+    for i in range(0, index_width, TRITON_BLOCK_SIZE):
+        offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
+
+        pos_offset = start_pos + offset
+        block_indices = pos_offset // block_size
+        block_numbers = tl.load(
+            block_table_ptr + req_idx * block_table_stride + block_indices,
+            mask=pos_offset < end_pos,
+        )
+        block_offsets = pos_offset % block_size
+        slot_ids = block_numbers * block_size + block_offsets
+
+        slot_ids = tl.where(offset < swa_len, slot_ids, -1)
+        tl.store(
+            swa_indices_ptr + pid * swa_indices_stride + offset,
+            slot_ids,
+            mask=offset < index_width,
+        )

vllm/v1/core/sched/scheduler.py

@@ -247,6 +247,11 @@ class Scheduler(SchedulerInterface):
                 # decoding instead of standard next-token sampling, so it has a query
                 # for the last sampled token plus queries for each draft token.
                 self.num_lookahead_tokens = self.num_spec_tokens + 1
+            if speculative_config.use_dspark():
+                # DSpark drafts a block of num_spec_tokens query tokens in which the
+                # anchor itself is the first prediction position (no separate bonus
+                # query), so it needs exactly num_spec_tokens lookahead slots.
+                self.num_lookahead_tokens = self.num_spec_tokens
 
         # Create the KV cache manager.
         if hash_block_size is None:

vllm/v1/worker/gpu/model_runner.py

@@ -193,7 +193,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             if self.is_last_pp_rank:
                 self.speculator = init_speculator(self.vllm_config, self.device)
 
-            if self.speculative_config.method in ("eagle3", "dflash"):
+            if self.speculative_config.method in ("eagle3", "dflash", "dspark"):
                 # Drafting may require auxiliary hidden states from target model outputs
                 self.use_aux_hidden_state_outputs = True
                 if self.use_pp:

vllm/v1/worker/gpu/sample/gumbel.py

@@ -223,6 +223,11 @@ def gumbel_sample(
     output_processed_logits_col: torch.Tensor | None = None,
     use_fp64: bool = False,
 ) -> torch.Tensor:
+    # Enforce contiguity on non-strided input tensors
+    expanded_idx_mapping = expanded_idx_mapping.contiguous()
+    pos = pos.contiguous()
+    if output_processed_logits_col is not None:
+        output_processed_logits_col = output_processed_logits_col.contiguous()
     num_tokens, vocab_size = logits.shape
     BLOCK_SIZE = 1024
     num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)

vllm/v1/worker/gpu/spec_decode/__init__.py

@@ -14,6 +14,12 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device):
         )
 
         return DFlashSpeculator(vllm_config, device)
+    elif speculative_config.method == "dspark":
+        from vllm.v1.worker.gpu.spec_decode.dspark.speculator import (
+            DSparkSpeculator,
+        )
+
+        return DSparkSpeculator(vllm_config, device)
     elif speculative_config.use_gemma4_mtp():
         from vllm.v1.worker.gpu.spec_decode.gemma4.speculator import (
             Gemma4Speculator,

vllm/v1/worker/gpu/spec_decode/dflash/speculator.py

@@ -29,6 +29,8 @@ logger = init_logger(__name__)
 
 
 class DFlashSpeculator(DraftModelSpeculator):
+    _speculator_name = "DFlash"  # For logging, so we can share methods with subclasses
+
     def __init__(self, vllm_config: VllmConfig, device: torch.device):
         super().__init__(vllm_config, device)
 
@@ -48,15 +50,17 @@ class DFlashSpeculator(DraftModelSpeculator):
 
         self.dflash_causal = get_dflash_causal(self.draft_model_config)
 
-        # Buffers for context K/V precomputation. Populated by prepare_dflash_inputs,
-        # and processed by the model's precompute_and_store_context_kv method.
-        # NOT captured by CUDA graphs.
+        # Whether the anchor query position is itself a prediction. DFlash default uses
+        # the anchor as the bonus token (only mask tokens predict); DSpark samples from
+        # the anchor and the N-1 mask token positions. See _prepare_dflash_inputs_kernel
+        self.sample_from_anchor = False
+
+        # Context positions for the K/V precompute. Populated by
+        # prepare_dflash_inputs, and processed by the model's
+        # precompute_and_store_context_kv method. NOT captured by CUDA graphs.
         self.context_positions = torch.zeros(
             self.max_num_tokens, dtype=torch.int64, device=device
         )
-        self.context_slot_mapping = torch.zeros(
-            self.max_num_tokens, dtype=torch.int64, device=device
-        )
 
         # Per-mask-token sampling buffers. Flattened from (num_reqs, num_spec_tokens).
         max_num_sampled_tokens = self.max_num_reqs * self.num_speculative_steps
@@ -94,7 +98,7 @@ class DFlashSpeculator(DraftModelSpeculator):
         )
 
     def capture(self, attn_states: dict | None = None) -> None:
-        logger.info("Capturing model for DFlash speculator...")
+        logger.info("Capturing model for %s speculator...", self._speculator_name)
         # Reset sampling indices to zero to prevent stale values from prior
         # dummy runs from being baked into the captured graph.
         self.sample_indices.zero_()
@@ -108,7 +112,7 @@ class DFlashSpeculator(DraftModelSpeculator):
             self.attn_groups,
             self.kv_cache_config,
             self.max_model_len,
-            progress_bar_desc="Capturing dflash CUDA graphs",
+            progress_bar_desc=f"Capturing {self._speculator_name.lower()} CUDA graphs",
         )
 
     def load_draft_model(
@@ -126,18 +130,39 @@ class DFlashSpeculator(DraftModelSpeculator):
     ) -> None:
         super().set_attn(model_state, kv_cache_config, block_tables)
 
-        # DFlash precomputes context K/V with a single block_size; mixing
-        # kv-cache groups would silently corrupt the cache for the non-matching group.
-        draft_groups = [gid for gid, g in enumerate(self.attn_groups) if g]
-        assert len(draft_groups) == 1, (
-            "DFlash currently requires all draft attention layers to share "
-            "a single kv-cache group."
-        )
-        self.draft_kv_cache_group_id = draft_groups[0]
+        self.draft_kv_cache_group_ids = [
+            gid for gid, g in enumerate(self.attn_groups) if g
+        ]
+        assert self.draft_kv_cache_group_ids, "No draft attention groups found."
+        self.draft_kv_cache_group_id = self.draft_kv_cache_group_ids[0]
         self.draft_block_size = self.block_tables.block_sizes[
             self.draft_kv_cache_group_id
         ]
 
+        # Per-group context slot buffers for the precompute (one row per group).
+        self._context_slot_mappings = torch.zeros(
+            len(self.draft_kv_cache_group_ids),
+            self.max_num_tokens,
+            dtype=torch.int64,
+            device=self.device,
+        )
+
+        # Map each draft decoder layer to the index (within draft_kv_cache_group_ids)
+        # of the kv-cache group its cache belongs to. Models that share a single group
+        # leave this as None and share one context slot mapping.
+        self._layer_group_idx: list[int] | None = None
+        if hasattr(self.model, "get_draft_kv_cache_layer_names"):
+            name_to_gid = {
+                ln: gid
+                for gid, group in enumerate(kv_cache_config.kv_cache_groups)
+                for ln in group.layer_names
+            }
+            gid_to_idx = {gid: i for i, gid in enumerate(self.draft_kv_cache_group_ids)}
+            self._layer_group_idx = [
+                gid_to_idx[name_to_gid[name]]
+                for name in self.model.get_draft_kv_cache_layer_names()
+            ]
+
     @torch.inference_mode()
     def _run_model(
         self,
@@ -183,9 +208,11 @@ class DFlashSpeculator(DraftModelSpeculator):
 
         num_sample = num_reqs * self.num_speculative_steps
         sample_hidden_states = last_hidden_states[self.sample_indices[:num_sample]]
+        # sample_pos is the predicted token's position Q; verification keys
+        # Gumbel by the predecessor (Q-1). sample_draft adds +1, so pass Q-2.
         draft_tokens = self.sample_draft(
             sample_hidden_states,
-            self.sample_pos[:num_sample],
+            self.sample_pos[:num_sample] - 2,
             self.sample_idx_mapping[:num_sample],
             self.temperature,
             self.seeds,
@@ -294,40 +321,49 @@ class DFlashSpeculator(DraftModelSpeculator):
         # The query slot mapping is written into the shared BlockTables slot_mappings.
         # That buffer's address is what the captured CUDA graph reads from at replay.
         assert self.draft_kv_cache_group_id >= 0
-        query_slot_mapping = self.block_tables.slot_mappings[
-            self.draft_kv_cache_group_id
-        ]
-        prepare_dflash_inputs(
-            self.input_buffers,
-            query_slot_mapping,
-            self.context_positions,
-            self.context_slot_mapping,
-            self.sample_indices,
-            self.sample_pos,
-            self.sample_idx_mapping,
-            input_batch,
-            num_sampled,
-            num_rejected,
-            last_sampled,
-            next_prefill_tokens,
-            self.block_tables.input_block_tables[self.draft_kv_cache_group_id],
-            self.draft_block_size,
-            self.parallel_drafting_token_id,
-            self.num_query_per_req,
-            self.num_speculative_steps,
-            self.max_num_reqs,
-            self.max_num_tokens,
-        )
+        # Support multiple draft KV cache groups by preparing inputs once for each
+        for i, gid in enumerate(self.draft_kv_cache_group_ids):
+            prepare_dflash_inputs(
+                self.input_buffers,
+                self.block_tables.slot_mappings[gid],
+                self.context_positions,
+                self._context_slot_mappings[i],
+                self.sample_indices,
+                self.sample_pos,
+                self.sample_idx_mapping,
+                input_batch,
+                num_sampled,
+                num_rejected,
+                last_sampled,
+                next_prefill_tokens,
+                self.block_tables.input_block_tables[gid],
+                self.block_tables.block_sizes[gid],
+                self.parallel_drafting_token_id,
+                self.num_query_per_req,
+                self.num_speculative_steps,
+                self.max_num_reqs,
+                self.max_num_tokens,
+                self.max_model_len,
+                self.sample_from_anchor,
+            )
 
         # Pre-insert context K/V into the cache. Runs eagerly outside the captured graph
         # because the context shape varies per step. During dummy runs the block tables
         # are placeholders, so we skip the cache write to avoid clobbering real entries.
+        # Each layer uses the context slots of its own kv-cache group.
+        if dummy_run:
+            context_slots: torch.Tensor | list[torch.Tensor | None] | None = None
+        elif self._layer_group_idx is not None:
+            context_slots = [
+                self._context_slot_mappings[gidx][:num_target_tokens]
+                for gidx in self._layer_group_idx
+            ]
+        else:
+            context_slots = self._context_slot_mappings[0][:num_target_tokens]
         self.model.precompute_and_store_context_kv(
             self.hidden_states[:num_target_tokens],
             self.context_positions[:num_target_tokens],
-            context_slot_mapping=(
-                None if dummy_run else self.context_slot_mapping[:num_target_tokens]
-            ),
+            context_slots,
         )
 
         # Every DFlash step has exactly num_query_per_req tokens, so we can use FULL CGs
@@ -408,6 +444,8 @@ def _prepare_dflash_inputs_kernel(
     num_speculative_steps,
     max_num_reqs,
     max_num_tokens,
+    max_model_len,
+    SAMPLE_FROM_ANCHOR: tl.constexpr,
     PAD_SLOT_ID: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
@@ -468,14 +506,21 @@ def _prepare_dflash_inputs_kernel(
     q_slot = q_block_id * block_size + (query_pos % block_size)
 
     tl.store(out_input_ids_ptr + query_idx, input_id, mask=is_query)
-    tl.store(out_query_positions_ptr + query_idx, query_pos, mask=is_query)
+    clamped_query_pos = tl.minimum(query_pos, max_model_len - 1)
+    tl.store(out_query_positions_ptr + query_idx, clamped_query_pos, mask=is_query)
     tl.store(out_query_slot_mapping_ptr + query_idx, q_slot, mask=is_query)
 
-    # --- Sample indices / positions / idx_mapping (mask tokens only) ---
-    is_sample = is_query & (query_off > 0)
-    sample_idx = req_idx * num_speculative_steps + (query_off - 1)
+    # --- Sample indices / positions / idx_mapping ---
+    # When SAMPLE_FROM_ANCHOR (DSpark), so we sample at EVERY query position
+    # and each position k predicts the NEXT token (sampled position = query_pos + 1).
+    # Otherwise (DFlash default) the anchor is the bonus token and only the mask tokens
+    # at offsets > 0 are sampled from, each AT its own position.
+    sample_off = 0 if SAMPLE_FROM_ANCHOR else 1
+    is_sample = is_query & (query_off >= sample_off)
+    sample_idx = req_idx * num_speculative_steps + (query_off - sample_off)
+    sample_pos = query_pos + 1 if SAMPLE_FROM_ANCHOR else query_pos
     tl.store(out_sample_indices_ptr + sample_idx, query_idx, mask=is_sample)
-    tl.store(out_sample_pos_ptr + sample_idx, query_pos, mask=is_sample)
+    tl.store(out_sample_pos_ptr + sample_idx, sample_pos, mask=is_sample)
     tl.store(out_sample_idx_mapping_ptr + sample_idx, req_state_idx, mask=is_sample)
 
     if block_idx == 0:
@@ -542,6 +587,8 @@ def prepare_dflash_inputs(
     num_speculative_steps: int,
     max_num_reqs: int,
     max_num_tokens: int,
+    max_model_len: int,
+    sample_from_anchor: bool = False,
 ) -> None:
     num_reqs = input_batch.num_reqs
     assert num_reqs > 0
@@ -577,6 +624,8 @@ def prepare_dflash_inputs(
         num_speculative_steps,
         max_num_reqs,
         max_num_tokens,
+        max_model_len,
+        SAMPLE_FROM_ANCHOR=sample_from_anchor,
         PAD_SLOT_ID=PAD_SLOT_ID,
         BLOCK_SIZE=BLOCK_SIZE,
     )

vllm/v1/worker/gpu/spec_decode/dspark/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/v1/worker/gpu/spec_decode/dspark/speculator.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DSpark speculator: semi-autoregressive parallel drafting.
+
+DSpark drafts a block of ``num_speculative_tokens`` tokens in one parallel pass
+(reusing the DFlash machinery: context-KV precompute + a query-block forward),
+then injects intra-block dependency with a lightweight sequential Markov head.
+
+Differences from DFlash:
+  * Anchor-as-first-prediction: each request emits exactly ``N =
+    num_speculative_tokens`` query tokens (anchor + N-1 noise), NOT ``1 + N``.
+    Every query position is a prediction (the anchor predicts the first draft
+    token), so we sample at all N positions and ``sample_pos = query_pos + 1``
+    (standard next-token), whereas DFlash's masks sit AT the predicted position.
+    This is the ``sample_from_anchor`` path in the shared prepare-inputs kernel.
+  * Sequential Markov sampling: instead of DFlash's single parallel sample, we
+    sample left-to-right, adding a prefix-dependent Markov bias derived from the
+    previously sampled token at each step.
+
+CUDA graphs (FULL, mirroring DFlash) cover the whole draft step: the parallel
+backbone forward AND the sequential Markov sampling.
+"""
+
+from typing import Any
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.config.compilation import CUDAGraphMode
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
+from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
+from vllm.v1.worker.gpu.spec_decode.dspark.utils import load_dspark_model
+
+
+class DSparkSpeculator(DFlashSpeculator):
+    _speculator_name = "DSpark"
+
+    def __init__(self, vllm_config: VllmConfig, device: torch.device):
+        super().__init__(vllm_config, device)
+
+        # Anchor-first: N query tokens per request (anchor + N-1 noise), not 1+N.
+        self.num_query_per_req = self.num_speculative_steps
+
+        # DSpark consumes mean-pooled target aux hidden states at the target
+        # layers, combined to hidden_size via main_proj. Store that combined
+        # main_x (hidden_size wide). DSpark does not use the same pre-allocated buffer
+        # that DeepSeek-V4's MTP uses.
+        draft_hidden = self.draft_model_config.get_hidden_size()
+        self.hidden_states = torch.zeros(
+            self.max_num_tokens, draft_hidden, dtype=self.dtype, device=device
+        )
+
+        self.dflash_causal = False
+
+        # The anchor query position is itself a prediction (see module docstring).
+        self.sample_from_anchor = True
+
+        self._step_cols = torch.arange(
+            self.num_speculative_steps, dtype=torch.int32, device=device
+        )
+
+        self._anchor_idx = (
+            torch.arange(self.max_num_reqs, dtype=torch.int64, device=device)
+            * self.num_query_per_req
+        )
+
+    def load_draft_model(
+        self,
+        target_model: torch.nn.Module,
+        target_attn_layer_names: set[str],
+    ) -> torch.nn.Module:
+        return load_dspark_model(target_model, self.vllm_config)
+
+    def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None:
+        # Sequential Markov sampling over the backbone's output hidden states.
+        n_spec = self.num_speculative_steps
+        num_sample = num_reqs * n_spec
+        # Per-(req, position) head hidden, ordered (req, step).
+        sample_hidden = head_hidden[self.sample_indices[:num_sample]]
+        base_logits = self.model.compute_logits(sample_hidden)
+        vocab_size = base_logits.shape[-1]
+        base_logits = base_logits.view(num_reqs, n_spec, vocab_size)
+
+        idx_map = self.sample_idx_mapping[:num_sample].view(num_reqs, n_spec)
+        sample_pos = self.sample_pos[:num_sample].view(num_reqs, n_spec)
+
+        # Anchor (bonus) token per request = the input id at query offset 0,
+        # read via the precomputed persistent index (fixed buffer for capture).
+        prev = self.input_buffers.input_ids[self._anchor_idx[:num_reqs]]
+
+        for i in range(n_spec):
+            # Sequential stage: Markov bias from the previously sampled token.
+            markov_embed = self.model.markov_embed(prev)
+            bias = self.model.markov_bias(markov_embed)
+            logits_i = base_logits[:, i] + bias
+            if self.draft_logits is not None:
+                # sample_pos is the predicted token's position Q; the target
+                # verifies it with the predecessor's Gumbel key (Q-1). Pass Q-1.
+                draft_i = gumbel_sample(
+                    logits_i,
+                    idx_map[:, i],
+                    self.temperature,
+                    self.seeds,
+                    sample_pos[:, i] - 1,
+                    apply_temperature=True,
+                    output_processed_logits=self.draft_logits,
+                    output_processed_logits_col=self._step_cols[i],
+                    use_fp64=self.use_fp64_gumbel,
+                )
+            else:
+                draft_i = logits_i.argmax(dim=-1)
+            self.draft_tokens[:num_reqs, i] = draft_i
+            prev = draft_i
+
+    def _generate_draft(
+        self,
+        num_reqs: int,
+        num_tokens_padded: int,
+        attn_metadata: dict[str, Any] | None,
+        slot_mappings: dict[str, torch.Tensor] | None,
+        num_tokens_across_dp: torch.Tensor | None,
+        cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
+    ) -> None:
+        # Full draft step (captured under CUDA graph): parallel backbone forward
+        # then sequential Markov sampling over its hidden state outputs.
+        head_hidden = self._run_model(
+            num_tokens_padded,
+            attn_metadata,
+            slot_mappings,
+            num_tokens_across_dp,
+            cudagraph_runtime_mode,
+        )
+        self._sample_sequential(num_reqs, head_hidden)

vllm/v1/worker/gpu/spec_decode/dspark/utils.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch.nn as nn
+
+from vllm.config import VllmConfig, replace
+from vllm.distributed.parallel_state import get_pp_group
+from vllm.model_executor.model_loader import get_model
+from vllm.v1.worker.gpu.spec_decode.eagle.utils import _should_share
+
+
+def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module:
+    speculative_config = vllm_config.speculative_config
+    assert speculative_config is not None
+    draft_model_config = speculative_config.draft_model_config
+
+    from vllm.compilation.backends import set_model_tag
+
+    # DSpark uses non-causal attention.
+    causal = False
+    draft_vllm_config = replace(
+        vllm_config,
+        attention_config=replace(
+            vllm_config.attention_config,
+            use_non_causal=not causal,
+            backend=speculative_config.attention_backend,
+        ),
+    )
+
+    with set_model_tag("dspark_head"):
+        draft_model = get_model(
+            vllm_config=draft_vllm_config, model_config=draft_model_config
+        )
+
+    if get_pp_group().world_size != 1:
+        raise NotImplementedError("DSpark does not support pipeline parallelism.")
+
+    target_language_model = (
+        target_model.get_language_model()
+        if hasattr(target_model, "get_language_model")
+        else target_model
+    )
+    target_inner = target_language_model.model
+    draft_inner = draft_model.model
+
+    target_embed = getattr(target_inner, "embed_tokens", None)
+    draft_embed = getattr(draft_inner, "embed_tokens", None)
+    if target_embed is not None and _should_share(
+        draft_model, "has_own_embed_tokens", draft_embed, target_embed
+    ):
+        if draft_embed is not None:
+            del draft_inner.embed_tokens
+        draft_inner.embed_tokens = target_embed
+
+    target_lm_head = getattr(target_model, "lm_head", None)
+    draft_lm_head = getattr(draft_model, "lm_head", None)
+    if target_lm_head is not None and _should_share(
+        draft_model, "has_own_lm_head", draft_lm_head, target_lm_head
+    ):
+        if draft_lm_head is not None:
+            del draft_model.lm_head
+        draft_model.lm_head = target_lm_head
+
+    return draft_model

vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py

@@ -44,6 +44,15 @@ def get_eagle3_aux_layers_from_config(
         if dflash_config and isinstance(dflash_config, dict):
             # Add 1 to convert DFlash's aux layer id semantics
             layer_ids = [i + 1 for i in (dflash_config.get("target_layer_ids") or [])]
+    if not layer_ids:
+        dspark_layer_ids = getattr(hf_config, "dspark_target_layer_ids", None)
+        if dspark_layer_ids:
+            layer_ids = [i + 1 for i in dspark_layer_ids]
+    if not layer_ids:
+        # Dense DSpark (e.g. Qwen3) also uses different aux layer semantics.
+        target_layer_ids = getattr(hf_config, "target_layer_ids", None)
+        if target_layer_ids:
+            layer_ids = [i + 1 for i in target_layer_ids]
     if layer_ids and isinstance(layer_ids, (list, tuple)):
         return tuple(layer_ids)
     return None

vllm/v1/worker/gpu/spec_decode/utils.py

@@ -54,17 +54,23 @@ class DraftTokensHandler:
 def get_parallel_drafting_token_id(hf_config) -> int:
     """Resolve the mask token id used for parallel drafting slots.
 
-    Checks (in order): `dflash_config.mask_token_id`, `pard_token`,
-    `ptd_token_id`. Raises ValueError if none are present.
+    Checks (in order): `dflash_config.mask_token_id`, top-level `mask_token_id`,
+    `dspark_noise_token_id`, `pard_token`, `ptd_token_id`. Raises ValueError if
+    none are present.
     """
     dflash_config = getattr(hf_config, "dflash_config", None) or {}
     if "mask_token_id" in dflash_config:
         return int(dflash_config["mask_token_id"])
+    if getattr(hf_config, "mask_token_id", None) is not None:
+        return int(hf_config.mask_token_id)
+    if hasattr(hf_config, "dspark_noise_token_id"):
+        return int(hf_config.dspark_noise_token_id)
     if hasattr(hf_config, "pard_token"):
         return int(hf_config.pard_token)
     if hasattr(hf_config, "ptd_token_id"):
         return int(hf_config.ptd_token_id)
     raise ValueError(
         "Model config must specify `dflash_config.mask_token_id`,"
-        " `pard_token`, or `ptd_token_id` for parallel drafting."
+        " `mask_token_id`, `dspark_noise_token_id`, `pard_token`, or"
+        " `ptd_token_id` for parallel drafting."
     )