From f5a8d73377d0f0a4e00cba172f9fbd0d50471b07 Mon Sep 17 00:00:00 2001 From: Benjamin Chislett Date: Wed, 1 Jul 2026 11:30:24 -0400 Subject: [PATCH] [Spec Decode] DSpark (#46995) Signed-off-by: Benjamin Chislett Signed-off-by: Giancarlo Delfin Signed-off-by: mgoin Co-authored-by: Giancarlo Delfin Co-authored-by: mgoin --- tests/models/registry.py | 13 + tests/models/test_registry.py | 4 + .../test_dspark_noncausal_sparse_mla.py | 529 ++++++++++++++++++ tests/v1/e2e/spec_decode/test_spec_decode.py | 61 ++ vllm/benchmarks/datasets/datasets.py | 1 + vllm/config/speculative.py | 39 +- vllm/config/vllm.py | 29 +- vllm/model_executor/models/qwen3_dflash.py | 12 +- vllm/model_executor/models/qwen3_dspark.py | 153 +++++ vllm/model_executor/models/registry.py | 2 + vllm/models/deepseek_v4/__init__.py | 9 + vllm/models/deepseek_v4/nvidia/dspark.py | 477 ++++++++++++++++ vllm/models/deepseek_v4/nvidia/model.py | 39 +- vllm/v1/attention/backends/mla/sparse_swa.py | 162 +++++- vllm/v1/core/sched/scheduler.py | 5 + vllm/v1/worker/gpu/model_runner.py | 2 +- vllm/v1/worker/gpu/sample/gumbel.py | 5 + vllm/v1/worker/gpu/spec_decode/__init__.py | 6 + .../gpu/spec_decode/dflash/speculator.py | 147 +++-- .../worker/gpu/spec_decode/dspark/__init__.py | 2 + .../gpu/spec_decode/dspark/speculator.py | 133 +++++ .../v1/worker/gpu/spec_decode/dspark/utils.py | 64 +++ .../gpu/spec_decode/eagle/eagle3_utils.py | 9 + vllm/v1/worker/gpu/spec_decode/utils.py | 12 +- 24 files changed, 1821 insertions(+), 94 deletions(-) create mode 100644 tests/v1/attention/test_dspark_noncausal_sparse_mla.py create mode 100644 vllm/model_executor/models/qwen3_dspark.py create mode 100644 vllm/models/deepseek_v4/nvidia/dspark.py create mode 100644 vllm/v1/worker/gpu/spec_decode/dspark/__init__.py create mode 100644 vllm/v1/worker/gpu/spec_decode/dspark/speculator.py create mode 100644 vllm/v1/worker/gpu/spec_decode/dspark/utils.py diff --git a/tests/models/registry.py b/tests/models/registry.py index da2dc8bb9d..87299d90e2 100644 --- a/tests/models/registry.py +++ b/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", diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index 7e3ecb372e..46c887838a 100644 --- a/tests/models/test_registry.py +++ b/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 diff --git a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py new file mode 100644 index 0000000000..ebb29c9af0 --- /dev/null +++ b/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" + ) diff --git a/tests/v1/e2e/spec_decode/test_spec_decode.py b/tests/v1/e2e/spec_decode/test_spec_decode.py index 0b94836cb0..2eff02ea6e 100644 --- a/tests/v1/e2e/spec_decode/test_spec_decode.py +++ b/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 diff --git a/vllm/benchmarks/datasets/datasets.py b/vllm/benchmarks/datasets/datasets.py index cf7cb91821..08c1613d7e 100644 --- a/vllm/benchmarks/datasets/datasets.py +++ b/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, diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 0a3347fd23..cf7f299eef 100644 --- a/vllm/config/speculative.py +++ b/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 diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index b36c02a48e..eb13a379ff 100644 --- a/vllm/config/vllm.py +++ b/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") diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 99c16f49df..f697f065a0 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/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( diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py new file mode 100644 index 0000000000..276b90358b --- /dev/null +++ b/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() diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index fd85729ca3..92db2ad4b8 100644 --- a/vllm/model_executor/models/registry.py +++ b/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"), diff --git a/vllm/models/deepseek_v4/__init__.py b/vllm/models/deepseek_v4/__init__.py index 44e486db77..d03eba3d41 100644 --- a/vllm/models/deepseek_v4/__init__.py +++ b/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", diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py new file mode 100644 index 0000000000..ab5d1a4b47 --- /dev/null +++ b/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}" diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index fed4b5d6a1..5f4a638c24 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/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. diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index df23f34378..ac722dca9f 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/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, + ) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index ec479f0930..fe59a29e12 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/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: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 58509f72d1..5b067c87e0 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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: diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index 74ee8abb2c..190307d5e7 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/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) diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index 09153dd20f..c70f169f7b 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/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, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 7978fc5133..d5d68a0146 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/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, ) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/__init__.py b/vllm/v1/worker/gpu/spec_decode/dspark/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/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 diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py new file mode 100644 index 0000000000..90012239bc --- /dev/null +++ b/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) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py new file mode 100644 index 0000000000..acc32dafcd --- /dev/null +++ b/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 diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py index 360f64921e..66d0ba8b43 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py +++ b/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 diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index 732d893855..37ca166593 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/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." )