@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f5a8d73377

[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>

Benjamin Chislett committed Jul 1, 2026 at 11:30 UTC f5a8d73377d0f0a4e00cba172f9fbd0d50471b07
24 files changed +1821 -94
tests/models/registry.py
+13
@@ -1407,6 +1407,19 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1407 max_num_seqs=32,
1408 min_transformers_version="4.56.3", # Required for Qwen3Next
1409 ),
1410 + # [DSpark]
1411 + "DSparkDraftModel": _HfExamplesInfo(
1412 + "deepseek-ai/DeepSeek-V4-Pro-DSpark",
1413 + speculative_model="deepseek-ai/DeepSeek-V4-Pro-DSpark", # draft in mtp.*
1414 + is_available_online=False,
1415 + use_original_num_layers=True, # DSpark has >1 draft block
1416 + ),
1417 + "Qwen3DSparkModel": _HfExamplesInfo(
1418 + "Qwen/Qwen3-8B",
1419 + speculative_model="deepseek-ai/dspark_qwen3_8b_block7",
1420 + is_available_online=False,
1421 + use_original_num_layers=True, # DSpark backbone requires all layers
1422 + ),
1423 # [Eagle]
1424 "EagleCohereForCausalLM": _HfExamplesInfo(
1425 "/host/engines/cohere-moe",
tests/models/test_registry.py
+4
@@ -48,6 +48,10 @@ def test_registry_imports(model_arch):
48 "(see #41376)"
49 )
50
51 + # DSpark draft model is NVIDIA-only; class is stubbed to None on ROCm/XPU.
52 + if model_arch == "DSparkDraftModel" and not current_platform.is_cuda():
53 + pytest.skip("DSparkDraftModel is only supported on CUDA")
54 +
55 # Ensure all model classes can be imported successfully
56 model_cls = ModelRegistry._try_load_model_cls(model_arch)
57 assert model_cls is not None
tests/v1/attention/test_dspark_noncausal_sparse_mla.py new
+529
@@ -0,0 +1,529 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Correctness tests for DSpark non-causal sliding-window MLA via sparse indices.
4 +
5 +DSpark drafts a block of N tokens whose attention is NON-CAUSAL within the block:
6 +every block token attends to the sliding window of context AND to all block
7 +tokens (including ones at later positions than itself).
8 +
9 +We can implement this using the existing sparse-MLA pathway by expanding the window size
10 +to include the rest of the block tokens: instead of setting topk indices to the 127
11 +previous tokens, we expand it to the next power of 2 (256) and include up to
12 +swa_size + block_size - 1 topk indices, so that each query attends to the rest. The
13 +remaining slots are filled with padding.
14 +
15 +The sparse-MLA decode kernels (FlashMLA on SM90/SM100, FlashInfer TRTLLM on
16 +SM100/SM120) are index-driven: each query attends over exactly the slots in its
17 +index list, with no causal mask (see ``flash_mla_with_kvcache(..., indices=...)``
18 +and ``_forward_decode``'s "attend only by generated indices"). The existing
19 +``test_sparse_mla_backends`` suite already validates arbitrary index lists, but
20 +only ones whose entries are <= the query's own position. This test suite specifically
21 +ensures correctness of the non-causal attention case.
22 +
23 +This reuses the harness/helpers of ``test_sparse_mla_backends.py`` (same model
24 +shapes, fp8_ds_mla round-trip, mock indexer, MockSparseMLAAttentionLayer); only
25 +the index construction differs.
26 +"""
27 +
28 +import math
29 +from types import MethodType, SimpleNamespace
30 +
31 +import pytest
32 +import torch
33 +
34 +from tests.v1.attention.test_mla_backends import (
35 + BatchSpec,
36 + MockSparseMLAAttentionLayer,
37 + create_and_prepopulate_kv_cache,
38 +)
39 +from tests.v1.attention.test_sparse_mla_backends import (
40 + _quantize_dequantize_fp8_ds_mla,
41 +)
42 +from tests.v1.attention.utils import (
43 + create_common_attn_metadata,
44 + create_standard_kv_cache_spec,
45 + create_vllm_config,
46 +)
47 +from vllm.config import set_current_vllm_config
48 +from vllm.model_executor.layers.linear import ColumnParallelLinear
49 +from vllm.platforms import current_platform
50 +
51 +if not current_platform.is_cuda():
52 + pytest.skip(
53 + "DSpark non-causal sparse MLA tests currently only support CUDA.",
54 + allow_module_level=True,
55 + )
56 +
57 +from vllm.utils.math_utils import cdiv
58 +from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
59 + FlashInferMLASparseTRTLLMBackend,
60 +)
61 +from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseBackend
62 +from vllm.v1.attention.ops import flashmla
63 +
64 +DEVICE_TYPE = current_platform.device_type
65 +
66 +# (window, block_size, topk_width). topk_width must be a multiple of the kernel's
67 +# B_TOPK (= padded query-head count, 64 or 128); we use 128-multiples to cover
68 +# both. The "wide" case needs window + block > 128 -> width must grow past 128.
69 +_DSPARK_CONFIGS = {
70 + "small_block": (8, 4, 128),
71 + "full_window_block": (128, 5, 256),
72 +}
73 +
74 +
75 +def _build_dspark_noncausal_indices(
76 + seq_lens: list[int],
77 + query_lens: list[int],
78 + window: int,
79 + topk_width: int,
80 + device: torch.device,
81 +) -> torch.Tensor:
82 + """Per-token sparse indices for the DSpark non-causal block.
83 +
84 + For a request with context length ``ctx`` and a query block of ``q_len``
85 + tokens (block positions ``ctx .. ctx+q_len-1``), EVERY block query attends to
86 + the same set: the trailing ``window`` context positions plus all block
87 + positions, i.e. the contiguous range ``[max(ctx-window,0) .. ctx+q_len-1]``.
88 + This is non-causal: an early block query's list contains later block tokens
89 + (future-pointing). The list is padded to ``topk_width`` with ``-1``.
90 + """
91 + total_query_tokens = sum(query_lens)
92 + sparse_indices = torch.full(
93 + (total_query_tokens, topk_width), -1, dtype=torch.int32, device=device
94 + )
95 + gt = 0
96 + for s_len, q_len in zip(seq_lens, query_lens):
97 + ctx_len = s_len - q_len
98 + lo = max(ctx_len - window, 0)
99 + hi = ctx_len + q_len # exclusive: window context + the full block
100 + idx_list = torch.arange(lo, hi, dtype=torch.int32, device=device)
101 + n = idx_list.numel()
102 + assert n <= topk_width, (
103 + f"index list ({n}) exceeds aligned topk width ({topk_width})"
104 + )
105 + for _ in range(q_len):
106 + sparse_indices[gt, :n] = idx_list
107 + gt += 1
108 + return sparse_indices
109 +
110 +
111 +def _run_sparse_backend_vs_sdpa(
112 + backend_cls,
113 + seq_lens: list[int],
114 + query_lens: list[int],
115 + sparse_indices: torch.Tensor,
116 + kv_cache_dtype: str,
117 + block_size: int,
118 + num_heads: int,
119 + device: torch.device,
120 + force_future_dominance: bool = False,
121 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
122 + """Run a sparse-MLA backend with the given per-token indices and compute a
123 + dense per-token SDPA reference over the SAME indices.
124 +
125 + Mirrors ``test_sparse_mla_backends.test_sparse_backend_decode_correctness``
126 + but with externally-supplied (non-causal) ``sparse_indices``.
127 +
128 + ``num_heads`` selects the kernel's B_TOPK (= padded q-head count): 128 -> 128,
129 + 64 -> 64. The aligned widths (128/256) are multiples of both, so num_heads=64
130 + exercises the head64 decode path that the SM100 alignment assert guards.
131 +
132 + ``force_future_dominance`` scales the LAST block token's latent KV so it
133 + dominates the softmax for every query that attends to it. With random data the
134 + few future block tokens carry negligible attention mass (especially with a wide
135 + window), so causal and non-causal outputs coincide; this knob makes the
136 + future-token contribution provably large for the differentiation test. It is
137 + OFF for the correctness test (which needs sensitivity to all tokens).
138 +
139 + Returns (backend_output, noncausal_reference, causal_reference). The causal
140 + reference restricts each query to indices <= its own absolute position.
141 + """
142 + batch_spec = BatchSpec(seq_lens=seq_lens, query_lens=query_lens)
143 + topk_tokens = sparse_indices.shape[1]
144 + dtype = torch.bfloat16
145 + use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla"
146 +
147 + kv_lora_rank = 512
148 + qk_nope_head_dim = 128
149 + qk_rope_head_dim = 64
150 + v_head_dim = 128
151 + head_size = kv_lora_rank + qk_rope_head_dim
152 +
153 + max_seqlen = max(seq_lens)
154 + total_cache_tokens = sum(seq_lens)
155 +
156 + vllm_config = create_vllm_config(
157 + model_name="deepseek-ai/DeepSeek-V2-Lite-Chat",
158 + tensor_parallel_size=1,
159 + max_model_len=max_seqlen,
160 + num_gpu_blocks=max(2048, cdiv(total_cache_tokens, block_size) + 1),
161 + block_size=block_size,
162 + hf_config_override={
163 + "index_topk": topk_tokens,
164 + "attn_module_list_cfg": [{"topk_tokens": topk_tokens}],
165 + },
166 + )
167 + model_config = vllm_config.model_config
168 + model_config.hf_text_config = SimpleNamespace(
169 + q_lora_rank=None,
170 + kv_lora_rank=kv_lora_rank,
171 + qk_nope_head_dim=qk_nope_head_dim,
172 + qk_rope_head_dim=qk_rope_head_dim,
173 + v_head_dim=v_head_dim,
174 + model_type="deepseek_v2",
175 + )
176 + model_config.dtype = dtype
177 + model_config.get_num_attention_heads = MethodType(
178 + lambda self, parallel_config: num_heads, model_config
179 + )
180 + model_config.get_num_kv_heads = MethodType(
181 + lambda self, parallel_config: 1, model_config
182 + )
183 + model_config.get_head_size = MethodType(lambda self: head_size, model_config)
184 + model_config.get_sliding_window = MethodType(lambda self: None, model_config)
185 +
186 + kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
187 +
188 + torch.manual_seed(0)
189 + scale = 1.0 / math.sqrt(head_size)
190 +
191 + # Shared MLA projection weights, used by both reference and backend.
192 + W_UK = torch.rand(
193 + kv_lora_rank, num_heads, qk_nope_head_dim, dtype=dtype, device=device
194 + )
195 + W_UV = torch.rand(kv_lora_rank, num_heads, v_head_dim, dtype=dtype, device=device)
196 +
197 + all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], []
198 + kv_c_contexts, k_pe_contexts = [], []
199 + reference_outputs = []
200 + # Causal counterpart of the reference: same index lists, but each query is
201 + # restricted to indices <= its own absolute position (drops future-pointing
202 + # block tokens). Used to prove the non-causal result is genuinely different.
203 + causal_reference_outputs = []
204 +
205 + kv_cache_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
206 + global_token_idx = 0
207 +
208 + for s_len, q_len in zip(seq_lens, query_lens):
209 + ctx_len = s_len - q_len
210 +
211 + q_c = torch.rand(
212 + q_len,
213 + num_heads,
214 + qk_nope_head_dim + qk_rope_head_dim,
215 + dtype=dtype,
216 + device=device,
217 + )
218 + kv_c_full = torch.rand(s_len, kv_lora_rank, dtype=dtype, device=device)
219 + k_pe_full = torch.rand(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device)
220 +
221 + if force_future_dominance:
222 + # Scale the last block token's latent KV so its key/value dominate the
223 + # softmax for any query attending to it. 4x in the latent dot makes its
224 + # pre-softmax score exceed the others by a wide margin, so non-causal
225 + # queries (which include it) diverge sharply from causal ones (which,
226 + # for all but the last query, exclude it). Applied before quantization
227 + # so cache and reference stay consistent.
228 + kv_c_full[s_len - 1] = kv_c_full[s_len - 1] * 4.0 + 2.0
229 +
230 + if use_fp8_ds_mla_quantization:
231 + is_sm100 = torch.cuda.get_device_capability()[0] >= 10
232 + kv_c_full, k_pe_squeezed = _quantize_dequantize_fp8_ds_mla(
233 + kv_c_full,
234 + k_pe_full.squeeze(1),
235 + block_size=block_size,
236 + scale=kv_cache_scale,
237 + simulate_sm100_e8m0_scales=is_sm100,
238 + )
239 + k_pe_full = k_pe_squeezed.unsqueeze(1)
240 +
241 + q_nope, q_pe = q_c.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)
242 + ql_nope = torch.einsum("qnh,lnh->qnl", q_nope, W_UK)
243 + q_mqa = torch.cat([ql_nope, q_pe], dim=-1)
244 +
245 + k_mqa = torch.cat([kv_c_full, k_pe_full.squeeze(1)], dim=-1)
246 + v_mqa = kv_c_full
247 +
248 + # Per-token sparse SDPA reference over the supplied (non-causal) indices.
249 + def _sparse_sdpa(idx_tensor, q_tok, k_mqa=k_mqa, v_mqa=v_mqa):
250 + k_sparse = k_mqa[idx_tensor].unsqueeze(1).expand(-1, num_heads, -1)
251 + v_sparse = v_mqa[idx_tensor].unsqueeze(1).expand(-1, num_heads, -1)
252 + out = torch.nn.functional.scaled_dot_product_attention(
253 + q_tok.unsqueeze(0).transpose(1, 2),
254 + k_sparse.unsqueeze(0).transpose(1, 2),
255 + v_sparse.unsqueeze(0).transpose(1, 2),
256 + scale=scale,
257 + )
258 + out = out.transpose(1, 2).squeeze(0)
259 + out = torch.einsum("qnl,lnv->qnv", out, W_UV)
260 + return out.flatten(start_dim=-2)
261 +
262 + for q_idx in range(q_len):
263 + tok_sparse_idx = sparse_indices[global_token_idx]
264 + valid_indices = tok_sparse_idx[tok_sparse_idx >= 0].long()
265 +
266 + q_tok = q_mqa[q_idx : q_idx + 1]
267 + reference_outputs.append(_sparse_sdpa(valid_indices, q_tok))
268 +
269 + # Causal: drop indices pointing past this query's own position.
270 + abs_pos = ctx_len + q_idx
271 + causal_indices = valid_indices[valid_indices <= abs_pos]
272 + causal_reference_outputs.append(_sparse_sdpa(causal_indices, q_tok))
273 + global_token_idx += 1
274 +
275 + all_q_vllm.append(q_c)
276 + all_kv_c_vllm.append(kv_c_full[ctx_len:])
277 + all_k_pe_vllm.append(k_pe_full[ctx_len:])
278 + kv_c_contexts.append(kv_c_full[: ctx_len + 1])
279 + k_pe_contexts.append(k_pe_full[: ctx_len + 1])
280 +
281 + query_vllm = torch.cat(all_q_vllm, dim=0)
282 + kv_c_vllm = torch.cat(all_kv_c_vllm, dim=0)
283 + k_pe_vllm = torch.cat(all_k_pe_vllm, dim=0)
284 + sdpa_reference = torch.cat(reference_outputs, dim=0)
285 + causal_reference = torch.cat(causal_reference_outputs, dim=0)
286 +
287 + vllm_config.cache_config.cache_dtype = kv_cache_dtype
288 + vllm_config.model_config.hf_config.index_topk = topk_tokens
289 +
290 + common_attn_metadata = create_common_attn_metadata(
291 + batch_spec, block_size, device, arange_block_indices=True
292 + )
293 + kv_cache = create_and_prepopulate_kv_cache(
294 + kv_c_contexts=kv_c_contexts,
295 + k_pe_contexts=k_pe_contexts,
296 + block_size=block_size,
297 + head_size=head_size,
298 + dtype=dtype,
299 + device=device,
300 + num_blocks=vllm_config.cache_config.num_gpu_blocks,
301 + common_attn_metadata=common_attn_metadata,
302 + randomize_blocks=False,
303 + kv_cache_dtype=kv_cache_dtype,
304 + scale=kv_cache_scale,
305 + )
306 +
307 + builder = backend_cls.get_builder_cls()(
308 + kv_cache_spec, ["placeholder"], vllm_config, device
309 + )
310 + metadata = builder.build(
311 + common_prefix_len=0, common_attn_metadata=common_attn_metadata
312 + )
313 +
314 + mock_indexer = SimpleNamespace(topk_indices_buffer=sparse_indices)
315 +
316 + kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1).view(
317 + kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim)
318 + )
319 + mock_kv_b_proj = ColumnParallelLinear(
320 + input_size=kv_lora_rank,
321 + output_size=num_heads * (qk_nope_head_dim + v_head_dim),
322 + bias=False,
323 + ).to(device=device, dtype=dtype)
324 + mock_kv_b_proj.weight = torch.nn.Parameter(kv_b_proj_weight.T.contiguous())
325 +
326 + with set_current_vllm_config(vllm_config):
327 + impl = backend_cls.get_impl_cls()(
328 + num_heads=num_heads,
329 + head_size=head_size,
330 + scale=scale,
331 + num_kv_heads=1,
332 + alibi_slopes=None,
333 + sliding_window=None,
334 + kv_cache_dtype=vllm_config.cache_config.cache_dtype,
335 + logits_soft_cap=None,
336 + attn_type="decoder",
337 + kv_sharing_target_layer_name=None,
338 + q_lora_rank=None,
339 + kv_lora_rank=kv_lora_rank,
340 + qk_nope_head_dim=qk_nope_head_dim,
341 + qk_rope_head_dim=qk_rope_head_dim,
342 + qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,
343 + v_head_dim=v_head_dim,
344 + kv_b_proj=mock_kv_b_proj,
345 + indexer=mock_indexer,
346 + )
347 + impl.process_weights_after_loading(dtype)
348 + mock_layer = MockSparseMLAAttentionLayer(
349 + impl=impl,
350 + num_heads=num_heads,
351 + qk_nope_head_dim=qk_nope_head_dim,
352 + qk_rope_head_dim=qk_rope_head_dim,
353 + v_head_dim=v_head_dim,
354 + kv_lora_rank=kv_lora_rank,
355 + device=device,
356 + W_UK=W_UK,
357 + W_UV=W_UV,
358 + q_scale=1.0,
359 + k_scale=1.0,
360 + )
361 +
362 + out_buffer = torch.empty(
363 + metadata.num_actual_tokens, num_heads * v_head_dim, dtype=dtype, device=device
364 + )
365 + with torch.inference_mode():
366 + backend_output = mock_layer.forward_impl(
367 + query_vllm, kv_c_vllm, k_pe_vllm, kv_cache, metadata, out_buffer
368 + )
369 + return backend_output, sdpa_reference, causal_reference
370 +
371 +
372 +def _skip_if_backend_unavailable(backend_cls, kv_cache_dtype: str, block_size: int):
373 + if kv_cache_dtype not in backend_cls.supported_kv_cache_dtypes:
374 + pytest.skip(f"{backend_cls.get_name()} does not support {kv_cache_dtype}")
375 + if (
376 + backend_cls is FlashMLASparseBackend
377 + and kv_cache_dtype.startswith("fp8")
378 + and kv_cache_dtype != "fp8_ds_mla"
379 + ):
380 + pytest.skip("FlashMLA Sparse fp8 only supports fp8_ds_mla kv-cache dtype")
381 + if block_size not in backend_cls.get_supported_kernel_block_sizes():
382 + pytest.skip(
383 + f"{backend_cls.get_name()} does not support block_size={block_size}"
384 + )
385 + if backend_cls is FlashMLASparseBackend:
386 + ok, reason = flashmla.is_flashmla_sparse_supported()
387 + if not ok:
388 + pytest.skip(reason)
389 + elif backend_cls is FlashInferMLASparseTRTLLMBackend:
390 + cap = current_platform.get_device_capability()
391 + if cap is None or not backend_cls.supports_compute_capability(cap):
392 + pytest.skip("FlashInferMLASparseTRTLLMBackend requires SM 10.x capability")
393 +
394 +
395 +@pytest.mark.parametrize(
396 + "backend_cls",
397 + [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend],
398 + ids=["FlashMLA", "FlashInferTRTLLM"],
399 +)
400 +@pytest.mark.parametrize("config_name", list(_DSPARK_CONFIGS.keys()))
401 +# Per backend, the skip logic routes fp8 to the supported flavor: FlashMLA tests
402 +# auto + fp8_ds_mla (and skips per-tensor "fp8", which it aliases to ds_mla);
403 +# FlashInfer TRTLLM tests auto + per-tensor "fp8" (and skips fp8_ds_mla, which it
404 +# does not implement). So both backends get a bf16 case and an fp8 case.
405 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_ds_mla", "fp8"])
406 +@pytest.mark.parametrize("block_size", [64])
407 +# h_q=128 -> B_TOPK=128; h_q=64 -> B_TOPK=64 (covers the head64 decode path the
408 +# SM100 alignment assert specifically guards). Aligned widths (128/256) satisfy both.
409 +@pytest.mark.parametrize("num_heads", [128, 64], ids=["h128", "h64"])
410 +def test_dspark_noncausal_sparse_mla_matches_sdpa(
411 + default_vllm_config,
412 + dist_init,
413 + workspace_init,
414 + backend_cls,
415 + config_name,
416 + kv_cache_dtype,
417 + block_size,
418 + num_heads,
419 +):
420 + """Non-causal (window ∪ block, future-pointing) per-token indices must match
421 + a dense SDPA reference over the same indices, for both sparse-MLA backends."""
422 + _skip_if_backend_unavailable(backend_cls, kv_cache_dtype, block_size)
423 +
424 + window, block, topk_width = _DSPARK_CONFIGS[config_name]
425 + device = torch.device(DEVICE_TYPE)
426 +
427 + # Decode-style batch: each request has `block` query tokens and enough
428 + # context for a full sliding window.
429 + seq_lens = [window + block + 123, window + block + 50]
430 + query_lens = [block, block]
431 +
432 + sparse_indices = _build_dspark_noncausal_indices(
433 + seq_lens, query_lens, window, topk_width, device
434 + )
435 +
436 + # Sanity: the construction must actually be non-causal (an early block query
437 + # must reference a later block position than itself).
438 + ctx0 = seq_lens[0] - query_lens[0]
439 + first_query_valid = sparse_indices[0][sparse_indices[0] >= 0]
440 + assert int(first_query_valid.max()) >= ctx0 + query_lens[0] - 1, (
441 + "expected the first block query to attend to a future block token"
442 + )
443 +
444 + backend_output, sdpa_reference, _ = _run_sparse_backend_vs_sdpa(
445 + backend_cls,
446 + seq_lens,
447 + query_lens,
448 + sparse_indices,
449 + kv_cache_dtype,
450 + block_size,
451 + num_heads,
452 + device,
453 + )
454 +
455 + assert backend_output.shape == sdpa_reference.shape
456 + assert backend_output.dtype == sdpa_reference.dtype
457 + assert torch.isfinite(backend_output).all()
458 + if kv_cache_dtype.startswith("fp8"):
459 + rtol, atol = 0.065, 0.05
460 + else:
461 + rtol, atol = 0.01, 0.01
462 + torch.testing.assert_close(backend_output, sdpa_reference, rtol=rtol, atol=atol)
463 +
464 +
465 +@pytest.mark.parametrize(
466 + "backend_cls",
467 + [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend],
468 + ids=["FlashMLA", "FlashInferTRTLLM"],
469 +)
470 +@pytest.mark.parametrize("config_name", list(_DSPARK_CONFIGS.keys()))
471 +@pytest.mark.parametrize("block_size", [64])
472 +def test_dspark_noncausal_differs_from_causal(
473 + default_vllm_config,
474 + dist_init,
475 + workspace_init,
476 + backend_cls,
477 + config_name,
478 + block_size,
479 +):
480 + """Differentiation guard: prove the backend genuinely attends to the
481 + future-pointing indices (not silently applying a causal mask, and not merely
482 + coinciding with a causal result because future tokens carry little weight).
483 +
484 + With random data the few future block tokens are a negligible fraction of the
485 + attended set (especially with a wide window), so causal and non-causal outputs
486 + are numerically indistinguishable -- that is correct physics, not a backend
487 + bug. To make the check meaningful we use ``force_future_dominance`` so the last
488 + block token dominates the softmax: the backend must then match the non-causal
489 + reference and diverge sharply from the causal one. bf16 (``auto``) suffices;
490 + the property is dtype-independent and fp8 correctness is covered above.
491 + """
492 + _skip_if_backend_unavailable(backend_cls, "auto", block_size)
493 +
494 + window, block, topk_width = _DSPARK_CONFIGS[config_name]
495 + device = torch.device(DEVICE_TYPE)
496 + seq_lens = [window + block + 123, window + block + 50]
497 + query_lens = [block, block]
498 +
499 + sparse_indices = _build_dspark_noncausal_indices(
500 + seq_lens, query_lens, window, topk_width, device
501 + )
502 +
503 + backend_output, sdpa_reference, causal_reference = _run_sparse_backend_vs_sdpa(
504 + backend_cls,
505 + seq_lens,
506 + query_lens,
507 + sparse_indices,
508 + "auto",
509 + block_size,
510 + 128,
511 + device,
512 + force_future_dominance=True,
513 + )
514 +
515 + # The two references must be clearly distinguishable for the check to mean
516 + # anything (dominance guarantees this).
517 + ref_gap = (sdpa_reference - causal_reference).abs().max().item()
518 + assert ref_gap > 0.1, (
519 + f"non-causal and causal references are too close (gap={ref_gap}); "
520 + "force_future_dominance did not create a separable scenario"
521 + )
522 +
523 + # Backend must track the NON-causal reference, not the causal one.
524 + torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.01, atol=0.01)
525 + causal_err = (backend_output - causal_reference).abs().max().item()
526 + assert causal_err > 0.1, (
527 + f"non-causal backend output matches the causal reference "
528 + f"(max abs diff={causal_err}); future-pointing indices are not attended to"
529 + )
tests/v1/e2e/spec_decode/test_spec_decode.py
+61
@@ -1412,6 +1412,67 @@ def test_dflash_acceptance_rates(
1412 cleanup_dist_env_and_memory()
1413
1414
1415 +@pytest.fixture
1416 +def dspark_config():
1417 + target_model = "Qwen/Qwen3-4B-FP8"
1418 + draft_model = "deepseek-ai/dspark_qwen3_4b_block7"
1419 +
1420 + return dict(
1421 + model=target_model,
1422 + trust_remote_code=True,
1423 + speculative_config={
1424 + "method": "dspark",
1425 + "model": draft_model,
1426 + "num_speculative_tokens": 7,
1427 + "attention_backend": "FLASH_ATTN",
1428 + "draft_sample_method": "probabilistic",
1429 + },
1430 + max_model_len=4096,
1431 + disable_log_stats=False,
1432 + )
1433 +
1434 +
1435 +@single_gpu_only
1436 +@large_gpu_mark(min_gb=24)
1437 +def test_dspark_correctness_and_acceptance_rate(dspark_config):
1438 + """
1439 + E2E test for DSpark speculative decoding: acceptance rate/length
1440 + regression coverage plus GSM8K correctness, at temperature=1.0 to
1441 + exercise the probabilistic draft-sampling/rejection-sampling path
1442 + (not just greedy).
1443 +
1444 + Uses Qwen/Qwen3-4B-FP8 as target with the dspark_qwen3_4b_block7 draft
1445 + model. Reference: measured over 12 runs of the full GSM8K set at
1446 + temperature=1.0 (prefix caching disabled to avoid cross-run reuse):
1447 + accuracy: min=0.782 max=0.814 mean=0.801
1448 + acceptance_rate: min=0.418 max=0.434 mean=0.428
1449 + acceptance_len: min=3.928 max=4.037 mean=3.994
1450 + Thresholds set conservatively to 10% to avoid flaking due to unlucky sampling
1451 + """
1452 + spec_llm = LLM(**dspark_config)
1453 +
1454 + results = evaluate_gsm8k_offline(spec_llm, temperature=1.0)
1455 + gsm8k_accuracy = results["accuracy"]
1456 +
1457 + metrics = spec_llm.get_metrics()
1458 + acceptance_rate = compute_acceptance_rate(metrics)
1459 + acceptance_len = compute_acceptance_len(metrics)
1460 +
1461 + print(
1462 + f"DSpark acceptance_rate={acceptance_rate:.2f}, "
1463 + f"acceptance_len={acceptance_len:.2f}, "
1464 + f"gsm8k_accuracy={gsm8k_accuracy:.3f}"
1465 + )
1466 +
1467 + assert acceptance_rate >= 0.428 * 0.9
1468 + assert acceptance_len >= 3.994 * 0.9
1469 + assert gsm8k_accuracy >= 0.801 * 0.9
1470 +
1471 + del spec_llm
1472 + torch.accelerator.empty_cache()
1473 + cleanup_dist_env_and_memory()
1474 +
1475 +
1476 @single_gpu_only
1477 def test_synthetic_acceptance_rate():
1478 """Verify that synthetic rejection sampling produces an acceptance
vllm/benchmarks/datasets/datasets.py
+1
@@ -2452,6 +2452,7 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
2452 num_requests=args.num_prompts,
2453 tokenizer=tokenizer,
2454 output_len=args.speed_bench_output_len,
2455 + skip_chat_template=args.skip_chat_template,
2456 chat_template_kwargs=getattr(args, "chat_template_kwargs", None),
2457 enable_multimodal_chat=args.enable_multimodal_chat,
2458 request_id_prefix=args.request_id_prefix,
vllm/config/speculative.py
+36 -3
@@ -54,6 +54,7 @@ MTPModelTypes = Literal[
54 ]
55 NgramGPUTypes = Literal["ngram_gpu"]
56 DFlashModelTypes = Literal["dflash"]
57 +DSparkModelTypes = Literal["dspark"]
58 EagleModelTypes = Literal[
59 "eagle", "eagle3", "extract_hidden_states", MTPModelTypes, DFlashModelTypes
60 ]
@@ -66,6 +67,7 @@ SpeculativeMethod = Literal[
67 "custom_class",
68 EagleModelTypes,
69 NgramGPUTypes,
70 + DSparkModelTypes,
71 ]
72 RejectionSampleMethod = Literal["standard", "synthetic", "block"]
73 DraftSampleMethod = Literal["greedy", "probabilistic"]
@@ -291,6 +293,7 @@ class SpeculativeConfig:
293 "eagle3",
294 "extract_hidden_states",
295 "dflash",
296 + "dspark",
297 )
298 factors.append(uses_aux_hidden_states)
299
@@ -608,6 +611,13 @@ class SpeculativeConfig:
611 # --quantization fp8 with a bf16 checkpoint.
612 if not self.quantization:
613 self.quantization = self.target_model_config.quantization
614 + elif self.method == "dspark":
615 + # DeepSeek DSpark can ship the weights inside the target checkpoint
616 + if self.target_model_config is None:
617 + raise ValueError("target_model_config must be present for dspark")
618 + self.model = self.target_model_config.model
619 + if not self.quantization:
620 + self.quantization = self.target_model_config.quantization
621 elif self.method in ("ngram", "[ngram]"):
622 self.model = "ngram"
623 elif self.method == "ngram_gpu":
@@ -755,18 +765,24 @@ class SpeculativeConfig:
765 draft_hf.truncated_vocab_size = target_vocab
766
767 # Automatically detect the method
758 - if self.method in ("eagle", "eagle3", "dflash"):
768 + if self.method in ("eagle", "eagle3", "dflash", "dspark"):
769 pass
770 # examples:
771 # yuhuili/EAGLE-LLaMA3-Instruct-8B
772 # yuhuili/EAGLE3-LLaMA3.1-Instruct-8B
773 # AngelSlim/Qwen3-8B_eagle3
774 + # deepseek-ai/dspark_qwen3_8b_block7
775 elif "eagle-" in self.draft_model_config.model.lower():
776 self.method = "eagle"
777 elif "eagle3" in self.draft_model_config.model.lower():
778 self.method = "eagle3"
779 elif "dflash" in self.draft_model_config.model.lower():
780 self.method = "dflash"
781 + elif (
782 + "dspark" in self.draft_model_config.model.lower()
783 + or "Qwen3DSparkModel" in self.draft_model_config.architectures
784 + ):
785 + self.method = "dspark"
786 elif self.draft_model_config.hf_config.model_type == "medusa":
787 self.method = "medusa"
788 elif self.draft_model_config.hf_config.model_type == "mlp_speculator":
@@ -813,7 +829,18 @@ class SpeculativeConfig:
829 self.draft_model_config.hf_config = eagle_config
830 self.update_arch_()
831
816 - if self.method == "dflash":
832 + if self.method == "dspark" and (
833 + "Qwen3DSparkModel" not in self.draft_model_config.architectures
834 + ):
835 + # DeepSeek-V4 DSpark reuses the full DeepSeek-V4 config
836 + # and its weights ship in the target checkpoint.
837 + self.draft_model_config.hf_config.model_type = "deepseek_v4"
838 + self.draft_model_config.hf_config.architectures = [
839 + "DSparkDraftModel"
840 + ]
841 + self.update_arch_()
842 +
843 + if self.method in ("dflash", "dspark"):
844 self.parallel_drafting = True
845
846 if self.num_speculative_tokens is not None and hasattr(
@@ -1129,11 +1156,17 @@ class SpeculativeConfig:
1156 )
1157
1158 def use_eagle(self) -> bool:
1132 - return self.method in ("eagle", "eagle3", "mtp", "dflash")
1159 + # NOTE: This method is usually a stand-in for "speculative decoding using
1160 + # target model hidden states"
1161 + # TODO(ben): Refactor this so the naming is clearer
1162 + return self.method in ("eagle", "eagle3", "mtp", "dflash", "dspark")
1163
1164 def use_dflash(self) -> bool:
1165 return self.method == "dflash"
1166
1167 + def use_dspark(self) -> bool:
1168 + return self.method == "dspark"
1169 +
1170 def uses_dynamic_speculative_decoding(self) -> bool:
1171 return self.num_speculative_tokens_per_batch_size is not None
1172
vllm/config/vllm.py
+24 -5
@@ -524,6 +524,16 @@ class VllmConfig:
524 if use_v2_model_runner is not None:
525 return use_v2_model_runner
526
527 + # DSpark is implemented only by the V2 GPU model runner, and DeepSeek-V4
528 + # is not otherwise a default-V2 architecture, so force V2 for it. If V2
529 + # is unsupported for the rest of the config, _validate_v2_model_runner
530 + # raises rather than silently falling back to V1 (which can't run dspark).
531 + if (
532 + self.speculative_config is not None
533 + and self.speculative_config.method == "dspark"
534 + ):
535 + return True
536 +
537 if self.model_config is not None and self.model_config.is_diffusion:
538 return True
539
@@ -958,10 +968,11 @@ class VllmConfig:
968 self.speculative_config.method not in get_args(EagleModelTypes)
969 and self.speculative_config.method not in get_args(NgramGPUTypes)
970 and self.speculative_config.method != "draft_model"
971 + and self.speculative_config.method != "dspark"
972 ):
973 raise ValueError(
974 "Currently, async scheduling is only supported "
964 - "with EAGLE/MTP/Draft Model/NGram GPU kind of "
975 + "with EAGLE/MTP/Draft Model/NGram GPU/DSpark kind of "
976 "speculative decoding"
977 )
978 if self.speculative_config.disable_padded_drafter_batch:
@@ -989,6 +1000,7 @@ class VllmConfig:
1000 self.speculative_config is not None
1001 and self.speculative_config.method not in get_args(EagleModelTypes)
1002 and self.speculative_config.method not in get_args(NgramGPUTypes)
1003 + and self.speculative_config.method != "dspark"
1004 ):
1005 logger.warning_once(
1006 "Async scheduling not supported with %s-based "
@@ -2038,17 +2050,24 @@ class VllmConfig:
2050 # TODO: ngram / ngram_gpu are not supported by the v2 model runner yet
2051 if speculative_config.method in ("ngram", "ngram_gpu"):
2052 unsupported.append("ngram/ngram_gpu speculative decoding")
2041 - elif speculative_config.method not in ("eagle", "eagle3", "mtp", "dflash"):
2053 + elif speculative_config.method not in (
2054 + "eagle",
2055 + "eagle3",
2056 + "mtp",
2057 + "dflash",
2058 + "dspark",
2059 + ):
2060 unsupported.append(f"speculative method '{speculative_config.method}'")
2061
2062 if speculative_config.uses_dynamic_speculative_decoding():
2063 unsupported.append("dynamic speculative decoding")
2064
2047 - # V2 EagleSpeculator does not support parallel_drafting (for P-Eagle)
2048 - # DFlash uses parallel drafting natively in V2 via DFlashSpeculator.
2065 + # V2 EagleSpeculator does not support parallel_drafting (for P-Eagle).
2066 + # DFlash and DSpark use parallel drafting natively in V2 via their
2067 + # own speculators.
2068 if (
2069 speculative_config.parallel_drafting
2051 - and speculative_config.method != "dflash"
2070 + and speculative_config.method not in ("dflash", "dspark")
2071 ):
2072 unsupported.append("parallel drafting for EAGLE speculative decoding")
2073
vllm/model_executor/models/qwen3_dflash.py
+9 -3
@@ -472,7 +472,7 @@ class DFlashQwen3Model(nn.Module):
472 self,
473 context_states: torch.Tensor,
474 context_positions: torch.Tensor,
475 - context_slot_mapping: torch.Tensor | None = None,
475 + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
476 ) -> None:
477 """Precompute K/V for context states write them into each layer's KV cache.
478
@@ -551,7 +551,13 @@ class DFlashQwen3Model(nn.Module):
551
552 # --- Per-layer cache insert ---
553 all_k_final = all_k_flat.view(L, num_ctx, nkv, hd)
554 + per_layer = isinstance(context_slot_mapping, (list, tuple))
555 for i in range(L):
556 + slot_mapping = (
557 + context_slot_mapping[i] if per_layer else context_slot_mapping
558 + )
559 + if slot_mapping is None:
560 + continue # dummy run: skip cache ops
561 attn = self._attn_layers[i]
562 kv_cache = attn.kv_cache
563 attn.impl.do_kv_cache_update(
@@ -559,7 +565,7 @@ class DFlashQwen3Model(nn.Module):
565 all_k_final[i],
566 all_v[i],
567 kv_cache,
562 - context_slot_mapping,
568 + slot_mapping,
569 )
570
571 def forward(
@@ -701,7 +707,7 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):
707 self,
708 context_states: torch.Tensor,
709 context_positions: torch.Tensor,
704 - context_slot_mapping: torch.Tensor | None = None,
710 + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
711 ) -> None:
712 """Precompute projected + RoPE'd K/V and write to cache."""
713 self.model.precompute_and_store_context_kv(
vllm/model_executor/models/qwen3_dspark.py new
+153
@@ -0,0 +1,153 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Qwen3 DSpark draft model for semi-autoregressive drafting.
4 +
5 +DSpark drafts a whole block in one parallel pass (DFlash-style: context-KV
6 +precompute + a non-causal query-block forward) and then injects intra-block
7 +dependency with a lightweight sequential Markov head.
8 +
9 +The parallel backbone is a standard Qwen3 decoder stack reused from the
10 +DFlash Qwen3 draft (see qwen3_dflash.py). DSpark adds:
11 + * ``markov_head``: low-rank V x r / r x V transition bias added to the base
12 + logits, sampled left-to-right by the speculator (the sequential stage).
13 +
14 +DSparkMarkovHead is shared with the DSV4-style DSpark model.
15 +"""
16 +
17 +from collections.abc import Iterable
18 +
19 +import torch
20 +import torch.nn as nn
21 +
22 +from vllm.config import VllmConfig
23 +from vllm.logger import init_logger
24 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
25 +from vllm.model_executor.layers.vocab_parallel_embedding import (
26 + ParallelLMHead,
27 + VocabParallelEmbedding,
28 +)
29 +
30 +from .qwen3_dflash import DFlashQwen3ForCausalLM, DFlashQwen3Model
31 +from .utils import AutoWeightsLoader, maybe_prefix, process_eagle_weight
32 +
33 +logger = init_logger(__name__)
34 +
35 +
36 +class DSparkMarkovHead(nn.Module):
37 + """Sequential transition-bias head (low-rank V x r, r x V).
38 +
39 + ``markov_w1[token]`` is an r-dim embedding of the previously sampled token;
40 + ``markov_w2`` projects it back to a vocab-size bias added to the base logits.
41 + """
42 +
43 + def __init__(self, vocab_size: int, markov_rank: int, prefix: str) -> None:
44 + super().__init__()
45 + # TODO(ben): profile for which (if any) it makes sense to replicate or TP-shard
46 + self.markov_w1 = VocabParallelEmbedding(
47 + vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w1")
48 + )
49 + self.markov_w2 = ParallelLMHead(
50 + vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w2")
51 + )
52 +
53 + def embed(self, token_ids: torch.Tensor) -> torch.Tensor:
54 + """r-dim Markov embedding of ``token_ids`` ([B] -> [B, r])."""
55 + return self.markov_w1(token_ids)
56 +
57 + def bias(self, markov_embed: torch.Tensor, logits_processor) -> torch.Tensor:
58 + """Vocab-size transition bias from a Markov embedding ([B, r] -> [B, V])."""
59 + return logits_processor(self.markov_w2, markov_embed)
60 +
61 +
62 +class Qwen3DSparkModel(DFlashQwen3Model):
63 + """DFlash Qwen3 backbone + DSpark Markov head."""
64 +
65 + def __init__(
66 + self,
67 + *,
68 + vllm_config: VllmConfig,
69 + start_layer_id: int = 0,
70 + prefix: str = "",
71 + ) -> None:
72 + super().__init__(
73 + vllm_config=vllm_config, start_layer_id=start_layer_id, prefix=prefix
74 + )
75 + config = self.config
76 + self.markov_head = DSparkMarkovHead(
77 + config.vocab_size,
78 + config.markov_rank,
79 + prefix=maybe_prefix(prefix, "markov_head"),
80 + )
81 +
82 +
83 +class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
84 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
85 + nn.Module.__init__(self)
86 + self.draft_model_config = vllm_config.speculative_config.draft_model_config
87 + self.config = self.draft_model_config.hf_config
88 + if getattr(self.config, "draft_vocab_size", None) is None:
89 + self.config.draft_vocab_size = getattr(self.config, "vocab_size", None)
90 + target_layer_num = vllm_config.model_config.get_num_layers(
91 + vllm_config.parallel_config
92 + )
93 + self.model = Qwen3DSparkModel(
94 + vllm_config=vllm_config,
95 + prefix=maybe_prefix(prefix, "model"),
96 + start_layer_id=target_layer_num,
97 + )
98 +
99 + logit_scale = getattr(self.config, "logit_scale", 1.0)
100 + self.lm_head = ParallelLMHead(
101 + self.config.draft_vocab_size,
102 + self.config.hidden_size,
103 + prefix=maybe_prefix(prefix, "lm_head"),
104 + )
105 + self.logits_processor = LogitsProcessor(
106 + self.config.draft_vocab_size, scale=logit_scale
107 + )
108 + target_vocab_size = vllm_config.model_config.get_vocab_size()
109 + if self.config.draft_vocab_size != target_vocab_size:
110 + self.draft_id_to_target_id = nn.Parameter(
111 + torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
112 + requires_grad=False,
113 + )
114 + else:
115 + self.draft_id_to_target_id = None
116 +
117 + def get_draft_kv_cache_layer_names(self) -> list[str]:
118 + return [layer.self_attn.attn.layer_name for layer in self.model.layers]
119 +
120 + def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor:
121 + return self.model.markov_head.embed(token_ids)
122 +
123 + def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
124 + return self.model.markov_head.bias(markov_embed, self.logits_processor)
125 +
126 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
127 + model_weights = {}
128 + includes_embed_tokens = False
129 + includes_lm_head = False
130 + for name, loaded_weight in weights:
131 + if "lm_head" not in name:
132 + name = "model." + name
133 + if "embed_tokens" in name:
134 + includes_embed_tokens = True
135 + if "lm_head" in name:
136 + includes_lm_head = True
137 + model_weights[name] = loaded_weight
138 + # Sets has_own_embed_tokens / has_own_lm_head so load_dspark_model
139 + # knows whether to keep these or alias the target's.
140 + process_eagle_weight(self, name)
141 +
142 + # mask_embedding is an unused placeholder param; DSpark masks via the vocab row.
143 + # confidence_head is not wired into inference yet; skip its weights.
144 + # embed_tokens / lm_head are optional; when omitted they are shared from
145 + # the target by load_dspark_model, so skip the unloaded params here.
146 + skip_substrs = ["mask_embedding", "confidence_head"]
147 + if not includes_embed_tokens:
148 + skip_substrs.append("embed_tokens")
149 + if not includes_lm_head:
150 + skip_substrs.append("lm_head")
151 + loader = AutoWeightsLoader(self, skip_substrs=skip_substrs)
152 + loader.load_weights(model_weights.items())
153 + self.model._build_fused_kv_buffers()
vllm/model_executor/models/registry.py
+2
@@ -585,6 +585,8 @@ _SPECULATIVE_DECODING_MODELS = {
585 "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),
586 "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
587 "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
588 + "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"),
589 + "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
590 "PEagleDraftModel": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
591 "PeagleLlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
592 "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
vllm/models/deepseek_v4/__init__.py
+9
@@ -17,14 +17,23 @@ from .quant_config import DeepseekV4FP8Config
17 if current_platform.is_rocm():
18 from .amd.model import DeepseekV4ForCausalLM
19 from .amd.mtp import DeepSeekV4MTP
20 +
21 + # DSpark is NVIDIA-only for now.
22 + DSparkDeepseekV4ForCausalLM = None # type: ignore[assignment]
23 elif current_platform.is_xpu():
24 from .xpu.model import DeepseekV4ForCausalLM # type: ignore[assignment]
25 from .xpu.mtp import DeepSeekV4MTP # type: ignore[assignment]
26 +
27 + DSparkDeepseekV4ForCausalLM = None # type: ignore[assignment]
28 else:
29 + from .nvidia.dspark import ( # type: ignore[assignment]
30 + DSparkDeepseekV4ForCausalLM,
31 + )
32 from .nvidia.model import DeepseekV4ForCausalLM # type: ignore[assignment]
33 from .nvidia.mtp import DeepSeekV4MTP # type: ignore[assignment]
34
35 __all__ = [
36 + "DSparkDeepseekV4ForCausalLM",
37 "DeepSeekV4MTP",
38 "DeepseekV4FP8Config",
39 "DeepseekV4ForCausalLM",
vllm/models/deepseek_v4/nvidia/dspark.py new
+477
@@ -0,0 +1,477 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DSpark draft model for DeepSeek-V4 (semi-autoregressive speculative decoding).
4 +
5 +See: qwen3_dspark.py for base architecture. This one is specialized to the DSV4 DSpark,
6 +which reuses the target model's architecture similarly to MTP.
7 +
8 +To implement non-causal attention, we leverage the sparse attention implementation to
9 +include the future query tokens in the top-k indices for each query token.
10 +"""
11 +
12 +from collections.abc import Iterable
13 +
14 +import regex as re
15 +import torch
16 +import torch.nn as nn
17 +
18 +from vllm.config import VllmConfig, get_current_vllm_config
19 +from vllm.distributed import (
20 + get_tensor_model_parallel_rank,
21 + get_tensor_model_parallel_world_size,
22 +)
23 +from vllm.logger import init_logger
24 +from vllm.model_executor.kernels.mhc.tilelang import (
25 + hc_head_fused_kernel_tilelang,
26 + mhc_post_tilelang,
27 +)
28 +from vllm.model_executor.layers.fused_moe import (
29 + fused_moe_make_expert_params_mapping,
30 +)
31 +from vllm.model_executor.layers.layernorm import RMSNorm
32 +from vllm.model_executor.layers.linear import ReplicatedLinear
33 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
34 +from vllm.model_executor.layers.vocab_parallel_embedding import (
35 + ParallelLMHead,
36 + VocabParallelEmbedding,
37 +)
38 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
39 +from vllm.model_executor.models.qwen3_dspark import (
40 + DSparkMarkovHead,
41 +)
42 +from vllm.model_executor.models.utils import maybe_prefix
43 +
44 +from .model import (
45 + DeepseekV4DecoderLayer,
46 + make_deepseek_v4_expert_params_mapping,
47 +)
48 +
49 +logger = init_logger(__name__)
50 +
51 +# MoE expert scale suffix differs by expert dtype (mirrors deepseek_v4 loaders):
52 +# fp4 experts register ``.weight_scale``; block-fp8 experts ``.weight_scale_inv``.
53 +_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")
54 +
55 +
56 +class DSparkDeepseekV4Model(nn.Module):
57 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
58 + super().__init__()
59 + assert vllm_config.speculative_config is not None
60 + config = vllm_config.speculative_config.draft_model_config.hf_config
61 + self.config = config
62 + self.hidden_size = config.hidden_size
63 + self.hc_mult = config.hc_mult
64 + self.hc_eps = config.hc_eps
65 + self.rms_norm_eps = config.rms_norm_eps
66 + self.num_hidden_layers = config.num_hidden_layers
67 + self.target_layer_ids = tuple(config.dspark_target_layer_ids)
68 +
69 + self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3
70 +
71 + # Shared with the target (aliased by the speculator's loading utility).
72 + self.embed_tokens = VocabParallelEmbedding(
73 + config.vocab_size,
74 + config.hidden_size,
75 + prefix=maybe_prefix(prefix, "embed_tokens"),
76 + )
77 +
78 + self.main_proj = ReplicatedLinear(
79 + config.hidden_size * len(self.target_layer_ids),
80 + config.hidden_size,
81 + bias=False,
82 + return_bias=False,
83 + quant_config=vllm_config.quant_config,
84 + prefix=maybe_prefix(prefix, "main_proj"),
85 + )
86 + self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
87 +
88 + current_vllm_config = get_current_vllm_config()
89 + self.layers = nn.ModuleList(
90 + [
91 + DeepseekV4DecoderLayer(
92 + current_vllm_config,
93 + prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"),
94 + )
95 + for i in range(self.num_dspark_layers)
96 + ]
97 + )
98 +
99 + # Heads: final norm + hc_head, and the Markov head
100 + # Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint
101 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
102 + hc_dim = self.hc_mult * config.hidden_size
103 + self.hc_head_fn = nn.Parameter(
104 + torch.empty(self.hc_mult, hc_dim, dtype=torch.float32),
105 + requires_grad=False,
106 + )
107 + self.hc_head_base = nn.Parameter(
108 + torch.empty(self.hc_mult, dtype=torch.float32), requires_grad=False
109 + )
110 + self.hc_head_scale = nn.Parameter(
111 + torch.empty(1, dtype=torch.float32), requires_grad=False
112 + )
113 + self.markov_head = DSparkMarkovHead(
114 + config.vocab_size,
115 + config.dspark_markov_rank,
116 + prefix=maybe_prefix(prefix, "markov_head"),
117 + )
118 +
119 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
120 + return self.embed_tokens(input_ids)
121 +
122 + def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor:
123 + """main_x = main_norm(main_proj(concat of target aux hidden states)).
124 +
125 + ``aux_hidden_states`` is [T, hidden_size * len(target_layer_ids)].
126 + """
127 + return self.main_norm(self.main_proj(aux_hidden_states))
128 +
129 + @torch.inference_mode()
130 + def precompute_and_store_context_kv(
131 + self,
132 + main_x: torch.Tensor,
133 + context_positions: torch.Tensor,
134 + context_slot_mappings: list[torch.Tensor | None] | None = None,
135 + ) -> None:
136 + """Insert the sliding-window context KV for every draft layer.
137 +
138 + Mirrors the reference DSparkAttention: each layer derives its context KV
139 + from the SAME projected target hidden ``main_x``, via that layer's own
140 + ``wkv`` + ``kv_norm`` + RoPE + quant, then writes it at the
141 + layer's context slots.
142 +
143 + ``context_slot_mappings`` is a per-layer list (each entry is the context
144 + slot mapping for that layer's kv-cache group, since the hybrid manager may
145 + place draft layers in different groups). ``None`` (or a ``None`` entry)
146 + runs the projection to reserve workspace but writes nothing (profiling).
147 + """
148 + for i, layer in enumerate(self.layers):
149 + slot_mapping = (
150 + None if context_slot_mappings is None else context_slot_mappings[i]
151 + )
152 + attn = layer.attn
153 + # Optimized DSV4 MLA path: wkv part of the fused wq_a|wkv projection
154 + # (q_lora part discarded), then RoPE/quant/insert via the fused op.
155 + qr_kv, _ = attn.fused_wqa_wkv(main_x)
156 + kv = qr_kv[..., attn.q_lora_rank :]
157 + kv = attn.kv_norm(kv)
158 + if slot_mapping is None:
159 + continue
160 + _insert_context_kv(attn, kv, context_positions, slot_mapping)
161 +
162 + def forward(
163 + self,
164 + input_ids: torch.Tensor,
165 + positions: torch.Tensor,
166 + inputs_embeds: torch.Tensor | None = None,
167 + ) -> torch.Tensor:
168 + if inputs_embeds is None:
169 + inputs_embeds = self.embed_input_ids(input_ids)
170 + # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]).
171 + hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1)
172 +
173 + residual = post_mix = res_mix = None
174 + for layer in self.layers:
175 + hidden_states, residual, post_mix, res_mix = layer(
176 + hidden_states,
177 + positions,
178 + input_ids,
179 + post_mix,
180 + res_mix,
181 + residual,
182 + )
183 + hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix)
184 + # hc_head reduces the hc copies; return the PRE-norm head hidden
185 + hidden_states = hc_head_fused_kernel_tilelang(
186 + hidden_states,
187 + self.hc_head_fn,
188 + self.hc_head_scale,
189 + self.hc_head_base,
190 + self.rms_norm_eps,
191 + self.hc_eps,
192 + )
193 + return hidden_states
194 +
195 +
196 +def _insert_context_kv(
197 + attn: nn.Module,
198 + kv: torch.Tensor,
199 + positions: torch.Tensor,
200 + slot_mapping: torch.Tensor,
201 +) -> None:
202 + """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV.
203 +
204 + Reuses the DSV4 fused insert ops (which also process a query; we pass a dummy
205 + query and discard it, since context tokens have no query). Mirrors
206 + ``DeepseekV4Attention._fused_qnorm_rope_kv_insert``.
207 + """
208 + swa_cache = attn.swa_cache_layer.kv_cache
209 + block_size = attn.swa_cache_layer.block_size
210 + cos_sin_cache = attn.rotary_emb.cos_sin_cache
211 + cache_dtype = swa_cache.dtype
212 + n_ctx = kv.shape[0]
213 + dummy_q = torch.zeros(
214 + (n_ctx, attn.n_local_heads, attn.head_dim),
215 + dtype=kv.dtype,
216 + device=kv.device,
217 + )
218 + if cache_dtype == torch.uint8:
219 + # fp8_ds_mla UE8M0 paged layout
220 + swa_2d = swa_cache.view(swa_cache.shape[0], -1)
221 + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
222 + dummy_q,
223 + kv,
224 + swa_2d,
225 + slot_mapping,
226 + positions,
227 + cos_sin_cache,
228 + attn.padded_heads,
229 + attn.eps,
230 + block_size,
231 + )
232 + elif cache_dtype == torch.bfloat16:
233 + swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
234 + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
235 + dummy_q,
236 + kv,
237 + swa_3d,
238 + slot_mapping,
239 + positions,
240 + cos_sin_cache,
241 + attn.eps,
242 + block_size,
243 + )
244 + else: # per-tensor fp8 (torch.float8_e4m3fn)
245 + # TODO(ben): double-check if this is being dispatched correctly for FI backend
246 + swa_3d = swa_cache.view(-1, block_size, attn.head_dim)
247 + dummy_q_fp8 = torch.zeros_like(dummy_q, dtype=torch.float8_e4m3fn)
248 + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
249 + dummy_q,
250 + kv,
251 + dummy_q_fp8,
252 + swa_3d,
253 + slot_mapping,
254 + positions,
255 + cos_sin_cache,
256 + attn._flashinfer_fp8_kv_scale,
257 + attn._flashinfer_fp8_q_scale_inv,
258 + attn.eps,
259 + block_size,
260 + )
261 +
262 +
263 +class DSparkDeepseekV4ForCausalLM(nn.Module):
264 + # Draft weights ship in the target checkpoint (mtp.*) without embed/head, so
265 + # load_dspark_model always aliases the target's.
266 + has_own_embed_tokens = False
267 + has_own_lm_head = False
268 +
269 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
270 + super().__init__()
271 + assert vllm_config.speculative_config is not None
272 + self.draft_model_config = vllm_config.speculative_config.draft_model_config
273 + self.config = self.draft_model_config.hf_config
274 + self.model = DSparkDeepseekV4Model(
275 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
276 + )
277 + # Shared with the target (aliased by the speculator's load utility).
278 + self.lm_head = ParallelLMHead(
279 + self.config.vocab_size,
280 + self.config.hidden_size,
281 + prefix=maybe_prefix(prefix, "lm_head"),
282 + )
283 + self.logits_processor = LogitsProcessor(self.config.vocab_size)
284 +
285 + # --- Hooks used by the speculator -------------------------------------
286 +
287 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
288 + return self.model.embed_input_ids(input_ids)
289 +
290 + def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor:
291 + return self.model.combine_hidden_states(aux_hidden_states)
292 +
293 + def get_draft_kv_cache_layer_names(self) -> list[str]:
294 + # DSV4 MLA path: each draft layer's sliding-window cache is a separate
295 + # layer, named by its prefix.
296 + return [layer.attn.swa_cache_layer.prefix for layer in self.model.layers]
297 +
298 + def precompute_and_store_context_kv(
299 + self,
300 + context_states: torch.Tensor,
301 + context_positions: torch.Tensor,
302 + context_slot_mappings: list[torch.Tensor | None] | None = None,
303 + ) -> None:
304 + self.model.precompute_and_store_context_kv(
305 + context_states, context_positions, context_slot_mappings
306 + )
307 +
308 + def forward(
309 + self,
310 + input_ids: torch.Tensor,
311 + positions: torch.Tensor,
312 + inputs_embeds: torch.Tensor | None = None,
313 + ) -> torch.Tensor:
314 + # Returns the pre-norm hc_head hidden ([T, hidden_size]).
315 + return self.model(input_ids, positions, inputs_embeds)
316 +
317 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
318 + """Base logits U_k = lm_head(norm(head_hidden))."""
319 + return self.logits_processor(self.lm_head, self.model.norm(hidden_states))
320 +
321 + def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor:
322 + return self.model.markov_head.embed(token_ids)
323 +
324 + def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
325 + return self.model.markov_head.bias(markov_embed, self.logits_processor)
326 +
327 + # --- Weight loading ----------------------------------------------------
328 +
329 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
330 + """Load the ``mtp.{0,1,2}.*`` draft weights from the target checkpoint.
331 +
332 + Non-mtp weights (embed/head/main layers) belong to the target model and
333 + are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target.
334 + """
335 + first_layer = self.model.layers[0]
336 + use_mega_moe = first_layer.ffn.use_mega_moe
337 + if use_mega_moe:
338 + expert_mapping = make_deepseek_v4_expert_params_mapping(
339 + self.config.n_routed_experts
340 + )
341 + else:
342 + expert_mapping = fused_moe_make_expert_params_mapping(
343 + self,
344 + ckpt_gate_proj_name="w1",
345 + ckpt_down_proj_name="w2",
346 + ckpt_up_proj_name="w3",
347 + num_experts=self.config.n_routed_experts,
348 + )
349 + expert_scale_suffix = (
350 + ".weight_scale"
351 + if getattr(self.config, "expert_dtype", "fp4") == "fp4"
352 + else ".weight_scale_inv"
353 + )
354 +
355 + # (param_name, ckpt_shard_name, shard_id) for non-expert stacked params.
356 + stacked_params_mapping = [
357 + ("gate_up_proj", "w1", 0),
358 + ("gate_up_proj", "w3", 1),
359 + ("attn.fused_wqa_wkv", "attn.wq_a", 0),
360 + ("attn.fused_wqa_wkv", "attn.wkv", 1),
361 + ]
362 +
363 + params_dict = dict(self.named_parameters())
364 + loaded_params: set[str] = set()
365 +
366 + tp_size = get_tensor_model_parallel_world_size()
367 + tp_rank = get_tensor_model_parallel_rank()
368 + n_local_head = self.config.num_attention_heads // tp_size
369 + head_start = n_local_head * tp_rank
370 + head_end = n_local_head * (tp_rank + 1)
371 +
372 + for name, loaded_weight in weights:
373 + mapped = self._remap_dspark_name(name)
374 + if mapped is None:
375 + continue
376 + name = mapped
377 +
378 + # ``.scale`` -> per-method scale suffix.
379 + if name.endswith(".scale"):
380 + suffix = (
381 + expert_scale_suffix
382 + if _EXPERT_SCALE_RE.search(name)
383 + else ".weight_scale_inv"
384 + )
385 + name = name.removesuffix(".scale") + suffix
386 +
387 + # E8M0 expert scales: keep raw exponent bytes.
388 + if ".experts." in name:
389 + if (
390 + "weight_scale" in name
391 + and loaded_weight.dtype == torch.float8_e8m0fnu
392 + ):
393 + loaded_weight = loaded_weight.view(torch.uint8)
394 + for param_name, weight_name, expert_id, shard_id in expert_mapping:
395 + if weight_name not in name:
396 + continue
397 + name_mapped = name.replace(weight_name, param_name)
398 + param = params_dict[name_mapped]
399 + success = param.weight_loader(
400 + param,
401 + loaded_weight,
402 + name_mapped,
403 + shard_id=shard_id,
404 + expert_id=expert_id,
405 + return_success=True,
406 + )
407 + if success:
408 + loaded_params.add(name_mapped)
409 + break
410 + continue
411 +
412 + # Stacked rules only apply to decoder-layer weights. Head-stack params
413 + # (main_proj/norm/hc_head/markov_head) load directly — otherwise e.g.
414 + # "markov_w1" would collide with the "w1" shard rule.
415 + is_layer_param = name.startswith("model.layers.")
416 + for param_name, weight_name, stacked_shard_id in stacked_params_mapping:
417 + if not is_layer_param or weight_name not in name:
418 + continue
419 + name = name.replace(weight_name, param_name)
420 + param = params_dict[name]
421 + param.weight_loader(param, loaded_weight, stacked_shard_id)
422 + loaded_params.add(name)
423 + break
424 + else:
425 + if "attn_sink" in name:
426 + narrow = loaded_weight[head_start:head_end]
427 + params_dict[name][: narrow.shape[0]].copy_(narrow)
428 + loaded_params.add(name)
429 + continue
430 + if ".shared_experts.w2" in name:
431 + name = name.replace(
432 + ".shared_experts.w2", ".shared_experts.down_proj"
433 + )
434 + if name.endswith(".ffn.gate.bias"):
435 + name = name.replace(
436 + ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias"
437 + )
438 + param = params_dict[name]
439 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
440 + weight_loader(param, loaded_weight)
441 + loaded_params.add(name)
442 +
443 + self._finalize_moe()
444 + logger.info_once("DSpark draft model loaded: %d params", len(loaded_params))
445 + return loaded_params
446 +
447 + def _finalize_moe(self) -> None:
448 + for layer in self.model.layers:
449 + layer.ffn.finalize_mega_moe_weights()
450 +
451 + def _remap_dspark_name(self, name: str) -> str | None:
452 + """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path.
453 +
454 + Returns None for non-mtp weights (owned by the target model).
455 + """
456 + m = re.match(r"mtp\.(\d+)\.(.*)", name)
457 + if m is None:
458 + return None
459 + stage = int(m.group(1))
460 + rest = m.group(2)
461 + # The confidence head is not wired into inference yet; drop its weights.
462 + if rest.startswith("confidence_head."):
463 + return None
464 + # Head-stack params live at model level (mtp.last), context combiner at
465 + # model level (mtp.0); everything else is a per-layer decoder block.
466 + head_prefixes = (
467 + "norm.",
468 + "hc_head_fn",
469 + "hc_head_base",
470 + "hc_head_scale",
471 + "markov_head.",
472 + )
473 + if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith(
474 + head_prefixes
475 + ):
476 + return f"model.{rest}"
477 + return f"model.layers.{stage}.{rest}"
vllm/models/deepseek_v4/nvidia/model.py
+32 -7
@@ -48,7 +48,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
48 VocabParallelEmbedding,
49 )
50 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
51 -from vllm.model_executor.models.interfaces import MixtureOfExperts, SupportsPP
51 +from vllm.model_executor.models.interfaces import (
52 + EagleModelMixin,
53 + MixtureOfExperts,
54 + SupportsEagle3,
55 + SupportsPP,
56 +)
57 from vllm.model_executor.models.utils import (
58 AutoWeightsLoader,
59 PPMissingLayer,
@@ -933,7 +938,7 @@ class DeepseekV4DecoderLayer(nn.Module):
938 return x, residual, post_mix, res_mix
939
940
936 -class DeepseekV4Model(nn.Module):
941 +class DeepseekV4Model(nn.Module, EagleModelMixin):
942 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
943 super().__init__()
944
@@ -1074,7 +1079,12 @@ class DeepseekV4Model(nn.Module):
1079 input_ids = input_ids.to(torch.int64)
1080
1081 residual, post_mix, res_mix = None, None, None
1077 - for layer in islice(self.layers, self.start_layer, self.end_layer):
1082 + aux_hidden_states: list[torch.Tensor] = []
1083 + final_aux_recon: torch.Tensor | None = None # avoid duplicate mhc_post call
1084 + for idx, layer in enumerate(
1085 + islice(self.layers, self.start_layer, self.end_layer),
1086 + start=self.start_layer,
1087 + ):
1088 hidden_states, residual, post_mix, res_mix = layer(
1089 hidden_states,
1090 positions,
@@ -1083,10 +1093,21 @@ class DeepseekV4Model(nn.Module):
1093 res_mix,
1094 residual,
1095 )
1096 + if idx + 1 in self.aux_hidden_state_layers:
1097 + # Reconstruct the aux hidden state for draft models
1098 + aux_recon = mhc_post_tilelang(
1099 + hidden_states, residual, post_mix, res_mix
1100 + )
1101 + aux_hidden_states.append(aux_recon.mean(dim=1))
1102 + final_aux_recon = aux_recon
1103 if layer is not None:
1087 - hidden_states = mhc_post_tilelang(
1088 - hidden_states, residual, post_mix, res_mix
1089 - )
1104 + # Reuse if the last layer was captured as an aux hidden state
1105 + if self.end_layer in self.aux_hidden_state_layers:
1106 + hidden_states = final_aux_recon
1107 + else:
1108 + hidden_states = mhc_post_tilelang(
1109 + hidden_states, residual, post_mix, res_mix
1110 + )
1111
1112 if not get_pp_group().is_last_rank:
1113 return IntermediateTensors({"hidden_states": hidden_states})
@@ -1104,6 +1125,8 @@ class DeepseekV4Model(nn.Module):
1125 self.hc_eps,
1126 )
1127 hidden_states = self.norm(hidden_states)
1128 + if len(aux_hidden_states) > 0:
1129 + return hidden_states, aux_hidden_states
1130 return hidden_states
1131
1132 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -1330,7 +1353,9 @@ class DeepseekV4MixtureOfExperts(MixtureOfExperts):
1353 moe.experts.update_expert_map()
1354
1355
1333 -class DeepseekV4ForCausalLM(nn.Module, SupportsPP, DeepseekV4MixtureOfExperts):
1356 +class DeepseekV4ForCausalLM(
1357 + nn.Module, SupportsPP, SupportsEagle3, DeepseekV4MixtureOfExperts
1358 +):
1359 model_cls = DeepseekV4Model
1360
1361 # Default mapper assumes the original FP4-expert checkpoint layout.
vllm/v1/attention/backends/mla/sparse_swa.py
+139 -23
@@ -306,16 +306,19 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
306 )
307
308 # Handle MTP: adjust decode_threshold like the indexer does
309 + spec_config = self.vllm_config.speculative_config
310 self.num_speculative_tokens = (
310 - self.vllm_config.speculative_config.num_speculative_tokens
311 - if self.vllm_config.speculative_config
312 - else 0
311 + spec_config.num_speculative_tokens if spec_config else 0
312 + )
313 + # Decode can have query_len up to
314 + # 1 + (2 if parallel drafting else 1) * num_speculative_tokens.
315 + # This MUST match the flashmla_sparse / indexer threshold so that
316 + # all backends agree on the decode/prefill split.
317 + spec_mult = (
318 + 2 if (spec_config is not None and spec_config.parallel_drafting) else 1
319 )
314 - # With MTP, decode can have query_len up to 1 + num_speculative_tokens.
315 - # Must match the threshold used by the indexer and flashmla_sparse so
316 - # that all backends agree on the decode/prefill split.
320 self.decode_threshold = (
318 - self.reorder_batch_threshold + self.num_speculative_tokens
321 + self.reorder_batch_threshold + spec_mult * self.num_speculative_tokens
322 )
323
324 hf_config = self.vllm_config.model_config.hf_config
@@ -368,6 +371,20 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
371 device=self.device,
372 )
373
374 + # DSpark draft: the block is non-causal (every query attends to the
375 + # trailing window of context PLUS all query tokens, including future ones),
376 + # so its per-token index list is wider than `window_size`. The kernel pads
377 + # the q-head count to B_TOPK (64/128), which requires the index width to be
378 + # a multiple of 128.
379 + self.is_dspark = spec_config is not None and spec_config.use_dspark()
380 + self.noncausal_index_width = (
381 + cdiv(self.window_size + self.num_speculative_tokens, 128) * 128
382 + if self.is_dspark
383 + else 0
384 + )
385 + self.decode_swa_indices_noncausal: torch.Tensor | None = None
386 + self._max_tokens = max_tokens
387 +
388 def build(
389 self,
390 common_prefix_len: int,
@@ -407,23 +424,56 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
424 is_valid_token = self.is_valid_token[: slot_mapping.shape[0]]
425 is_valid_token.copy_(slot_mapping >= 0)
426
427 + non_causal = not common_attn_metadata.causal
428 + decode_swa_indices = self.decode_swa_indices
429 if num_decode_tokens > 0:
430 self.decode_swa_lens[num_decode_tokens:] = 0
412 - _compute_swa_indices_and_lens_kernel[(num_decode_tokens,)](
413 - self.decode_swa_indices,
414 - self.decode_swa_indices.stride(0),
415 - self.decode_swa_lens,
416 - self.window_size,
417 - query_start_loc,
418 - seq_lens,
419 - token_to_req_indices,
420 - is_valid_token,
421 - block_table,
422 - block_table.stride(0),
423 - self.block_size,
424 - token_offset=0,
425 - TRITON_BLOCK_SIZE=1024,
426 - )
431 + if non_causal:
432 + assert self.is_dspark, (
433 + "Non-causal DeepseekV4 SWA is only supported for the DSpark "
434 + "speculation mode, but causal=False was set without DSpark."
435 + )
436 + if self.decode_swa_indices_noncausal is None:
437 + self.decode_swa_indices_noncausal = torch.zeros(
438 + self._max_tokens,
439 + 1,
440 + self.noncausal_index_width,
441 + dtype=torch.int32,
442 + device=self.device,
443 + )
444 + decode_swa_indices = self.decode_swa_indices_noncausal
445 + _compute_dspark_noncausal_swa_indices_kernel[(num_decode_tokens,)](
446 + decode_swa_indices,
447 + decode_swa_indices.stride(0),
448 + self.decode_swa_lens,
449 + self.window_size,
450 + self.noncausal_index_width,
451 + query_start_loc,
452 + seq_lens,
453 + token_to_req_indices,
454 + is_valid_token,
455 + block_table,
456 + block_table.stride(0),
457 + self.block_size,
458 + token_offset=0,
459 + TRITON_BLOCK_SIZE=1024,
460 + )
461 + else:
462 + _compute_swa_indices_and_lens_kernel[(num_decode_tokens,)](
463 + decode_swa_indices,
464 + decode_swa_indices.stride(0),
465 + self.decode_swa_lens,
466 + self.window_size,
467 + query_start_loc,
468 + seq_lens,
469 + token_to_req_indices,
470 + is_valid_token,
471 + block_table,
472 + block_table.stride(0),
473 + self.block_size,
474 + token_offset=0,
475 + TRITON_BLOCK_SIZE=1024,
476 + )
477
478 # Prefill SWA indices live in paged coordinates. `token_offset` lets
479 # the kernel read is_valid_token / token_to_req_indices at absolute
@@ -471,7 +521,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
521 slot_mapping=slot_mapping,
522 is_valid_token=is_valid_token,
523 token_to_req_indices=token_to_req_indices,
474 - decode_swa_indices=self.decode_swa_indices[:num_decode_tokens],
524 + decode_swa_indices=decode_swa_indices[:num_decode_tokens],
525 decode_swa_lens=self.decode_swa_lens[:num_decode_tokens],
526 prefill_swa_indices=(
527 self.prefill_swa_indices[:num_prefill_tokens]
@@ -665,3 +715,69 @@ def _compute_swa_indices_and_lens_kernel(
715 slot_ids,
716 mask=offset < window_size,
717 )
718 +
719 +
720 +# TODO(ben): unify this kernel to reduce duplication
721 +@triton.jit(do_not_specialize=["token_offset"])
722 +def _compute_dspark_noncausal_swa_indices_kernel(
723 + swa_indices_ptr,
724 + swa_indices_stride,
725 + swa_lens_ptr,
726 + window_size,
727 + index_width,
728 + query_start_loc_ptr,
729 + seq_lens_ptr,
730 + token_to_req_indices_ptr,
731 + is_valid_token_ptr,
732 + block_table_ptr,
733 + block_table_stride,
734 + block_size,
735 + token_offset,
736 + TRITON_BLOCK_SIZE: tl.constexpr,
737 +):
738 + """Non-causal per-token indices for the DSpark draft block.
739 +
740 + Here, we populate the topk indices with the trailing window of context tokens,
741 + plus all query tokens (including future ones).
742 + """
743 + pid = tl.program_id(0)
744 + token_idx = pid + token_offset
745 + is_valid = tl.load(is_valid_token_ptr + token_idx)
746 + if not is_valid:
747 + tl.store(swa_lens_ptr + pid, 0)
748 + return
749 +
750 + req_idx = tl.load(token_to_req_indices_ptr + token_idx)
751 +
752 + query_start = tl.load(query_start_loc_ptr + req_idx)
753 + query_end = tl.load(query_start_loc_ptr + req_idx + 1)
754 + query_len = query_end - query_start
755 +
756 + seq_len = tl.load(seq_lens_ptr + req_idx)
757 + prefix_len = seq_len - query_len
758 +
759 + # Block-anchored window (shared by every token in the block) + full block.
760 + start_pos = tl.maximum(prefix_len - window_size, 0)
761 + end_pos = seq_len
762 +
763 + swa_len = end_pos - start_pos
764 + tl.store(swa_lens_ptr + pid, swa_len)
765 +
766 + for i in range(0, index_width, TRITON_BLOCK_SIZE):
767 + offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
768 +
769 + pos_offset = start_pos + offset
770 + block_indices = pos_offset // block_size
771 + block_numbers = tl.load(
772 + block_table_ptr + req_idx * block_table_stride + block_indices,
773 + mask=pos_offset < end_pos,
774 + )
775 + block_offsets = pos_offset % block_size
776 + slot_ids = block_numbers * block_size + block_offsets
777 +
778 + slot_ids = tl.where(offset < swa_len, slot_ids, -1)
779 + tl.store(
780 + swa_indices_ptr + pid * swa_indices_stride + offset,
781 + slot_ids,
782 + mask=offset < index_width,
783 + )
vllm/v1/core/sched/scheduler.py
+5
@@ -247,6 +247,11 @@ class Scheduler(SchedulerInterface):
247 # decoding instead of standard next-token sampling, so it has a query
248 # for the last sampled token plus queries for each draft token.
249 self.num_lookahead_tokens = self.num_spec_tokens + 1
250 + if speculative_config.use_dspark():
251 + # DSpark drafts a block of num_spec_tokens query tokens in which the
252 + # anchor itself is the first prediction position (no separate bonus
253 + # query), so it needs exactly num_spec_tokens lookahead slots.
254 + self.num_lookahead_tokens = self.num_spec_tokens
255
256 # Create the KV cache manager.
257 if hash_block_size is None:
vllm/v1/worker/gpu/model_runner.py
+1 -1
@@ -193,7 +193,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
193 if self.is_last_pp_rank:
194 self.speculator = init_speculator(self.vllm_config, self.device)
195
196 - if self.speculative_config.method in ("eagle3", "dflash"):
196 + if self.speculative_config.method in ("eagle3", "dflash", "dspark"):
197 # Drafting may require auxiliary hidden states from target model outputs
198 self.use_aux_hidden_state_outputs = True
199 if self.use_pp:
vllm/v1/worker/gpu/sample/gumbel.py
+5
@@ -223,6 +223,11 @@ def gumbel_sample(
223 output_processed_logits_col: torch.Tensor | None = None,
224 use_fp64: bool = False,
225 ) -> torch.Tensor:
226 + # Enforce contiguity on non-strided input tensors
227 + expanded_idx_mapping = expanded_idx_mapping.contiguous()
228 + pos = pos.contiguous()
229 + if output_processed_logits_col is not None:
230 + output_processed_logits_col = output_processed_logits_col.contiguous()
231 num_tokens, vocab_size = logits.shape
232 BLOCK_SIZE = 1024
233 num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
vllm/v1/worker/gpu/spec_decode/__init__.py
+6
@@ -14,6 +14,12 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device):
14 )
15
16 return DFlashSpeculator(vllm_config, device)
17 + elif speculative_config.method == "dspark":
18 + from vllm.v1.worker.gpu.spec_decode.dspark.speculator import (
19 + DSparkSpeculator,
20 + )
21 +
22 + return DSparkSpeculator(vllm_config, device)
23 elif speculative_config.use_gemma4_mtp():
24 from vllm.v1.worker.gpu.spec_decode.gemma4.speculator import (
25 Gemma4Speculator,
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+98 -49
@@ -29,6 +29,8 @@ logger = init_logger(__name__)
29
30
31 class DFlashSpeculator(DraftModelSpeculator):
32 + _speculator_name = "DFlash" # For logging, so we can share methods with subclasses
33 +
34 def __init__(self, vllm_config: VllmConfig, device: torch.device):
35 super().__init__(vllm_config, device)
36
@@ -48,15 +50,17 @@ class DFlashSpeculator(DraftModelSpeculator):
50
51 self.dflash_causal = get_dflash_causal(self.draft_model_config)
52
51 - # Buffers for context K/V precomputation. Populated by prepare_dflash_inputs,
52 - # and processed by the model's precompute_and_store_context_kv method.
53 - # NOT captured by CUDA graphs.
53 + # Whether the anchor query position is itself a prediction. DFlash default uses
54 + # the anchor as the bonus token (only mask tokens predict); DSpark samples from
55 + # the anchor and the N-1 mask token positions. See _prepare_dflash_inputs_kernel
56 + self.sample_from_anchor = False
57 +
58 + # Context positions for the K/V precompute. Populated by
59 + # prepare_dflash_inputs, and processed by the model's
60 + # precompute_and_store_context_kv method. NOT captured by CUDA graphs.
61 self.context_positions = torch.zeros(
62 self.max_num_tokens, dtype=torch.int64, device=device
63 )
57 - self.context_slot_mapping = torch.zeros(
58 - self.max_num_tokens, dtype=torch.int64, device=device
59 - )
64
65 # Per-mask-token sampling buffers. Flattened from (num_reqs, num_spec_tokens).
66 max_num_sampled_tokens = self.max_num_reqs * self.num_speculative_steps
@@ -94,7 +98,7 @@ class DFlashSpeculator(DraftModelSpeculator):
98 )
99
100 def capture(self, attn_states: dict | None = None) -> None:
97 - logger.info("Capturing model for DFlash speculator...")
101 + logger.info("Capturing model for %s speculator...", self._speculator_name)
102 # Reset sampling indices to zero to prevent stale values from prior
103 # dummy runs from being baked into the captured graph.
104 self.sample_indices.zero_()
@@ -108,7 +112,7 @@ class DFlashSpeculator(DraftModelSpeculator):
112 self.attn_groups,
113 self.kv_cache_config,
114 self.max_model_len,
111 - progress_bar_desc="Capturing dflash CUDA graphs",
115 + progress_bar_desc=f"Capturing {self._speculator_name.lower()} CUDA graphs",
116 )
117
118 def load_draft_model(
@@ -126,18 +130,39 @@ class DFlashSpeculator(DraftModelSpeculator):
130 ) -> None:
131 super().set_attn(model_state, kv_cache_config, block_tables)
132
129 - # DFlash precomputes context K/V with a single block_size; mixing
130 - # kv-cache groups would silently corrupt the cache for the non-matching group.
131 - draft_groups = [gid for gid, g in enumerate(self.attn_groups) if g]
132 - assert len(draft_groups) == 1, (
133 - "DFlash currently requires all draft attention layers to share "
134 - "a single kv-cache group."
135 - )
136 - self.draft_kv_cache_group_id = draft_groups[0]
133 + self.draft_kv_cache_group_ids = [
134 + gid for gid, g in enumerate(self.attn_groups) if g
135 + ]
136 + assert self.draft_kv_cache_group_ids, "No draft attention groups found."
137 + self.draft_kv_cache_group_id = self.draft_kv_cache_group_ids[0]
138 self.draft_block_size = self.block_tables.block_sizes[
139 self.draft_kv_cache_group_id
140 ]
141
142 + # Per-group context slot buffers for the precompute (one row per group).
143 + self._context_slot_mappings = torch.zeros(
144 + len(self.draft_kv_cache_group_ids),
145 + self.max_num_tokens,
146 + dtype=torch.int64,
147 + device=self.device,
148 + )
149 +
150 + # Map each draft decoder layer to the index (within draft_kv_cache_group_ids)
151 + # of the kv-cache group its cache belongs to. Models that share a single group
152 + # leave this as None and share one context slot mapping.
153 + self._layer_group_idx: list[int] | None = None
154 + if hasattr(self.model, "get_draft_kv_cache_layer_names"):
155 + name_to_gid = {
156 + ln: gid
157 + for gid, group in enumerate(kv_cache_config.kv_cache_groups)
158 + for ln in group.layer_names
159 + }
160 + gid_to_idx = {gid: i for i, gid in enumerate(self.draft_kv_cache_group_ids)}
161 + self._layer_group_idx = [
162 + gid_to_idx[name_to_gid[name]]
163 + for name in self.model.get_draft_kv_cache_layer_names()
164 + ]
165 +
166 @torch.inference_mode()
167 def _run_model(
168 self,
@@ -183,9 +208,11 @@ class DFlashSpeculator(DraftModelSpeculator):
208
209 num_sample = num_reqs * self.num_speculative_steps
210 sample_hidden_states = last_hidden_states[self.sample_indices[:num_sample]]
211 + # sample_pos is the predicted token's position Q; verification keys
212 + # Gumbel by the predecessor (Q-1). sample_draft adds +1, so pass Q-2.
213 draft_tokens = self.sample_draft(
214 sample_hidden_states,
188 - self.sample_pos[:num_sample],
215 + self.sample_pos[:num_sample] - 2,
216 self.sample_idx_mapping[:num_sample],
217 self.temperature,
218 self.seeds,
@@ -294,40 +321,49 @@ class DFlashSpeculator(DraftModelSpeculator):
321 # The query slot mapping is written into the shared BlockTables slot_mappings.
322 # That buffer's address is what the captured CUDA graph reads from at replay.
323 assert self.draft_kv_cache_group_id >= 0
297 - query_slot_mapping = self.block_tables.slot_mappings[
298 - self.draft_kv_cache_group_id
299 - ]
300 - prepare_dflash_inputs(
301 - self.input_buffers,
302 - query_slot_mapping,
303 - self.context_positions,
304 - self.context_slot_mapping,
305 - self.sample_indices,
306 - self.sample_pos,
307 - self.sample_idx_mapping,
308 - input_batch,
309 - num_sampled,
310 - num_rejected,
311 - last_sampled,
312 - next_prefill_tokens,
313 - self.block_tables.input_block_tables[self.draft_kv_cache_group_id],
314 - self.draft_block_size,
315 - self.parallel_drafting_token_id,
316 - self.num_query_per_req,
317 - self.num_speculative_steps,
318 - self.max_num_reqs,
319 - self.max_num_tokens,
320 - )
324 + # Support multiple draft KV cache groups by preparing inputs once for each
325 + for i, gid in enumerate(self.draft_kv_cache_group_ids):
326 + prepare_dflash_inputs(
327 + self.input_buffers,
328 + self.block_tables.slot_mappings[gid],
329 + self.context_positions,
330 + self._context_slot_mappings[i],
331 + self.sample_indices,
332 + self.sample_pos,
333 + self.sample_idx_mapping,
334 + input_batch,
335 + num_sampled,
336 + num_rejected,
337 + last_sampled,
338 + next_prefill_tokens,
339 + self.block_tables.input_block_tables[gid],
340 + self.block_tables.block_sizes[gid],
341 + self.parallel_drafting_token_id,
342 + self.num_query_per_req,
343 + self.num_speculative_steps,
344 + self.max_num_reqs,
345 + self.max_num_tokens,
346 + self.max_model_len,
347 + self.sample_from_anchor,
348 + )
349
350 # Pre-insert context K/V into the cache. Runs eagerly outside the captured graph
351 # because the context shape varies per step. During dummy runs the block tables
352 # are placeholders, so we skip the cache write to avoid clobbering real entries.
353 + # Each layer uses the context slots of its own kv-cache group.
354 + if dummy_run:
355 + context_slots: torch.Tensor | list[torch.Tensor | None] | None = None
356 + elif self._layer_group_idx is not None:
357 + context_slots = [
358 + self._context_slot_mappings[gidx][:num_target_tokens]
359 + for gidx in self._layer_group_idx
360 + ]
361 + else:
362 + context_slots = self._context_slot_mappings[0][:num_target_tokens]
363 self.model.precompute_and_store_context_kv(
364 self.hidden_states[:num_target_tokens],
365 self.context_positions[:num_target_tokens],
328 - context_slot_mapping=(
329 - None if dummy_run else self.context_slot_mapping[:num_target_tokens]
330 - ),
366 + context_slots,
367 )
368
369 # 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(
444 num_speculative_steps,
445 max_num_reqs,
446 max_num_tokens,
447 + max_model_len,
448 + SAMPLE_FROM_ANCHOR: tl.constexpr,
449 PAD_SLOT_ID: tl.constexpr,
450 BLOCK_SIZE: tl.constexpr,
451 ):
@@ -468,14 +506,21 @@ def _prepare_dflash_inputs_kernel(
506 q_slot = q_block_id * block_size + (query_pos % block_size)
507
508 tl.store(out_input_ids_ptr + query_idx, input_id, mask=is_query)
471 - tl.store(out_query_positions_ptr + query_idx, query_pos, mask=is_query)
509 + clamped_query_pos = tl.minimum(query_pos, max_model_len - 1)
510 + tl.store(out_query_positions_ptr + query_idx, clamped_query_pos, mask=is_query)
511 tl.store(out_query_slot_mapping_ptr + query_idx, q_slot, mask=is_query)
512
474 - # --- Sample indices / positions / idx_mapping (mask tokens only) ---
475 - is_sample = is_query & (query_off > 0)
476 - sample_idx = req_idx * num_speculative_steps + (query_off - 1)
513 + # --- Sample indices / positions / idx_mapping ---
514 + # When SAMPLE_FROM_ANCHOR (DSpark), so we sample at EVERY query position
515 + # and each position k predicts the NEXT token (sampled position = query_pos + 1).
516 + # Otherwise (DFlash default) the anchor is the bonus token and only the mask tokens
517 + # at offsets > 0 are sampled from, each AT its own position.
518 + sample_off = 0 if SAMPLE_FROM_ANCHOR else 1
519 + is_sample = is_query & (query_off >= sample_off)
520 + sample_idx = req_idx * num_speculative_steps + (query_off - sample_off)
521 + sample_pos = query_pos + 1 if SAMPLE_FROM_ANCHOR else query_pos
522 tl.store(out_sample_indices_ptr + sample_idx, query_idx, mask=is_sample)
478 - tl.store(out_sample_pos_ptr + sample_idx, query_pos, mask=is_sample)
523 + tl.store(out_sample_pos_ptr + sample_idx, sample_pos, mask=is_sample)
524 tl.store(out_sample_idx_mapping_ptr + sample_idx, req_state_idx, mask=is_sample)
525
526 if block_idx == 0:
@@ -542,6 +587,8 @@ def prepare_dflash_inputs(
587 num_speculative_steps: int,
588 max_num_reqs: int,
589 max_num_tokens: int,
590 + max_model_len: int,
591 + sample_from_anchor: bool = False,
592 ) -> None:
593 num_reqs = input_batch.num_reqs
594 assert num_reqs > 0
@@ -577,6 +624,8 @@ def prepare_dflash_inputs(
624 num_speculative_steps,
625 max_num_reqs,
626 max_num_tokens,
627 + max_model_len,
628 + SAMPLE_FROM_ANCHOR=sample_from_anchor,
629 PAD_SLOT_ID=PAD_SLOT_ID,
630 BLOCK_SIZE=BLOCK_SIZE,
631 )
vllm/v1/worker/gpu/spec_decode/dspark/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/v1/worker/gpu/spec_decode/dspark/speculator.py new
+133
@@ -0,0 +1,133 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DSpark speculator: semi-autoregressive parallel drafting.
4 +
5 +DSpark drafts a block of ``num_speculative_tokens`` tokens in one parallel pass
6 +(reusing the DFlash machinery: context-KV precompute + a query-block forward),
7 +then injects intra-block dependency with a lightweight sequential Markov head.
8 +
9 +Differences from DFlash:
10 + * Anchor-as-first-prediction: each request emits exactly ``N =
11 + num_speculative_tokens`` query tokens (anchor + N-1 noise), NOT ``1 + N``.
12 + Every query position is a prediction (the anchor predicts the first draft
13 + token), so we sample at all N positions and ``sample_pos = query_pos + 1``
14 + (standard next-token), whereas DFlash's masks sit AT the predicted position.
15 + This is the ``sample_from_anchor`` path in the shared prepare-inputs kernel.
16 + * Sequential Markov sampling: instead of DFlash's single parallel sample, we
17 + sample left-to-right, adding a prefix-dependent Markov bias derived from the
18 + previously sampled token at each step.
19 +
20 +CUDA graphs (FULL, mirroring DFlash) cover the whole draft step: the parallel
21 +backbone forward AND the sequential Markov sampling.
22 +"""
23 +
24 +from typing import Any
25 +
26 +import torch
27 +
28 +from vllm.config import VllmConfig
29 +from vllm.config.compilation import CUDAGraphMode
30 +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
31 +from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
32 +from vllm.v1.worker.gpu.spec_decode.dspark.utils import load_dspark_model
33 +
34 +
35 +class DSparkSpeculator(DFlashSpeculator):
36 + _speculator_name = "DSpark"
37 +
38 + def __init__(self, vllm_config: VllmConfig, device: torch.device):
39 + super().__init__(vllm_config, device)
40 +
41 + # Anchor-first: N query tokens per request (anchor + N-1 noise), not 1+N.
42 + self.num_query_per_req = self.num_speculative_steps
43 +
44 + # DSpark consumes mean-pooled target aux hidden states at the target
45 + # layers, combined to hidden_size via main_proj. Store that combined
46 + # main_x (hidden_size wide). DSpark does not use the same pre-allocated buffer
47 + # that DeepSeek-V4's MTP uses.
48 + draft_hidden = self.draft_model_config.get_hidden_size()
49 + self.hidden_states = torch.zeros(
50 + self.max_num_tokens, draft_hidden, dtype=self.dtype, device=device
51 + )
52 +
53 + self.dflash_causal = False
54 +
55 + # The anchor query position is itself a prediction (see module docstring).
56 + self.sample_from_anchor = True
57 +
58 + self._step_cols = torch.arange(
59 + self.num_speculative_steps, dtype=torch.int32, device=device
60 + )
61 +
62 + self._anchor_idx = (
63 + torch.arange(self.max_num_reqs, dtype=torch.int64, device=device)
64 + * self.num_query_per_req
65 + )
66 +
67 + def load_draft_model(
68 + self,
69 + target_model: torch.nn.Module,
70 + target_attn_layer_names: set[str],
71 + ) -> torch.nn.Module:
72 + return load_dspark_model(target_model, self.vllm_config)
73 +
74 + def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None:
75 + # Sequential Markov sampling over the backbone's output hidden states.
76 + n_spec = self.num_speculative_steps
77 + num_sample = num_reqs * n_spec
78 + # Per-(req, position) head hidden, ordered (req, step).
79 + sample_hidden = head_hidden[self.sample_indices[:num_sample]]
80 + base_logits = self.model.compute_logits(sample_hidden)
81 + vocab_size = base_logits.shape[-1]
82 + base_logits = base_logits.view(num_reqs, n_spec, vocab_size)
83 +
84 + idx_map = self.sample_idx_mapping[:num_sample].view(num_reqs, n_spec)
85 + sample_pos = self.sample_pos[:num_sample].view(num_reqs, n_spec)
86 +
87 + # Anchor (bonus) token per request = the input id at query offset 0,
88 + # read via the precomputed persistent index (fixed buffer for capture).
89 + prev = self.input_buffers.input_ids[self._anchor_idx[:num_reqs]]
90 +
91 + for i in range(n_spec):
92 + # Sequential stage: Markov bias from the previously sampled token.
93 + markov_embed = self.model.markov_embed(prev)
94 + bias = self.model.markov_bias(markov_embed)
95 + logits_i = base_logits[:, i] + bias
96 + if self.draft_logits is not None:
97 + # sample_pos is the predicted token's position Q; the target
98 + # verifies it with the predecessor's Gumbel key (Q-1). Pass Q-1.
99 + draft_i = gumbel_sample(
100 + logits_i,
101 + idx_map[:, i],
102 + self.temperature,
103 + self.seeds,
104 + sample_pos[:, i] - 1,
105 + apply_temperature=True,
106 + output_processed_logits=self.draft_logits,
107 + output_processed_logits_col=self._step_cols[i],
108 + use_fp64=self.use_fp64_gumbel,
109 + )
110 + else:
111 + draft_i = logits_i.argmax(dim=-1)
112 + self.draft_tokens[:num_reqs, i] = draft_i
113 + prev = draft_i
114 +
115 + def _generate_draft(
116 + self,
117 + num_reqs: int,
118 + num_tokens_padded: int,
119 + attn_metadata: dict[str, Any] | None,
120 + slot_mappings: dict[str, torch.Tensor] | None,
121 + num_tokens_across_dp: torch.Tensor | None,
122 + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
123 + ) -> None:
124 + # Full draft step (captured under CUDA graph): parallel backbone forward
125 + # then sequential Markov sampling over its hidden state outputs.
126 + head_hidden = self._run_model(
127 + num_tokens_padded,
128 + attn_metadata,
129 + slot_mappings,
130 + num_tokens_across_dp,
131 + cudagraph_runtime_mode,
132 + )
133 + self._sample_sequential(num_reqs, head_hidden)
vllm/v1/worker/gpu/spec_decode/dspark/utils.py new
+64
@@ -0,0 +1,64 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch.nn as nn
5 +
6 +from vllm.config import VllmConfig, replace
7 +from vllm.distributed.parallel_state import get_pp_group
8 +from vllm.model_executor.model_loader import get_model
9 +from vllm.v1.worker.gpu.spec_decode.eagle.utils import _should_share
10 +
11 +
12 +def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module:
13 + speculative_config = vllm_config.speculative_config
14 + assert speculative_config is not None
15 + draft_model_config = speculative_config.draft_model_config
16 +
17 + from vllm.compilation.backends import set_model_tag
18 +
19 + # DSpark uses non-causal attention.
20 + causal = False
21 + draft_vllm_config = replace(
22 + vllm_config,
23 + attention_config=replace(
24 + vllm_config.attention_config,
25 + use_non_causal=not causal,
26 + backend=speculative_config.attention_backend,
27 + ),
28 + )
29 +
30 + with set_model_tag("dspark_head"):
31 + draft_model = get_model(
32 + vllm_config=draft_vllm_config, model_config=draft_model_config
33 + )
34 +
35 + if get_pp_group().world_size != 1:
36 + raise NotImplementedError("DSpark does not support pipeline parallelism.")
37 +
38 + target_language_model = (
39 + target_model.get_language_model()
40 + if hasattr(target_model, "get_language_model")
41 + else target_model
42 + )
43 + target_inner = target_language_model.model
44 + draft_inner = draft_model.model
45 +
46 + target_embed = getattr(target_inner, "embed_tokens", None)
47 + draft_embed = getattr(draft_inner, "embed_tokens", None)
48 + if target_embed is not None and _should_share(
49 + draft_model, "has_own_embed_tokens", draft_embed, target_embed
50 + ):
51 + if draft_embed is not None:
52 + del draft_inner.embed_tokens
53 + draft_inner.embed_tokens = target_embed
54 +
55 + target_lm_head = getattr(target_model, "lm_head", None)
56 + draft_lm_head = getattr(draft_model, "lm_head", None)
57 + if target_lm_head is not None and _should_share(
58 + draft_model, "has_own_lm_head", draft_lm_head, target_lm_head
59 + ):
60 + if draft_lm_head is not None:
61 + del draft_model.lm_head
62 + draft_model.lm_head = target_lm_head
63 +
64 + return draft_model
vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py
+9
@@ -44,6 +44,15 @@ def get_eagle3_aux_layers_from_config(
44 if dflash_config and isinstance(dflash_config, dict):
45 # Add 1 to convert DFlash's aux layer id semantics
46 layer_ids = [i + 1 for i in (dflash_config.get("target_layer_ids") or [])]
47 + if not layer_ids:
48 + dspark_layer_ids = getattr(hf_config, "dspark_target_layer_ids", None)
49 + if dspark_layer_ids:
50 + layer_ids = [i + 1 for i in dspark_layer_ids]
51 + if not layer_ids:
52 + # Dense DSpark (e.g. Qwen3) also uses different aux layer semantics.
53 + target_layer_ids = getattr(hf_config, "target_layer_ids", None)
54 + if target_layer_ids:
55 + layer_ids = [i + 1 for i in target_layer_ids]
56 if layer_ids and isinstance(layer_ids, (list, tuple)):
57 return tuple(layer_ids)
58 return None
vllm/v1/worker/gpu/spec_decode/utils.py
+9 -3
@@ -54,17 +54,23 @@ class DraftTokensHandler:
54 def get_parallel_drafting_token_id(hf_config) -> int:
55 """Resolve the mask token id used for parallel drafting slots.
56
57 - Checks (in order): `dflash_config.mask_token_id`, `pard_token`,
58 - `ptd_token_id`. Raises ValueError if none are present.
57 + Checks (in order): `dflash_config.mask_token_id`, top-level `mask_token_id`,
58 + `dspark_noise_token_id`, `pard_token`, `ptd_token_id`. Raises ValueError if
59 + none are present.
60 """
61 dflash_config = getattr(hf_config, "dflash_config", None) or {}
62 if "mask_token_id" in dflash_config:
63 return int(dflash_config["mask_token_id"])
64 + if getattr(hf_config, "mask_token_id", None) is not None:
65 + return int(hf_config.mask_token_id)
66 + if hasattr(hf_config, "dspark_noise_token_id"):
67 + return int(hf_config.dspark_noise_token_id)
68 if hasattr(hf_config, "pard_token"):
69 return int(hf_config.pard_token)
70 if hasattr(hf_config, "ptd_token_id"):
71 return int(hf_config.ptd_token_id)
72 raise ValueError(
73 "Model config must specify `dflash_config.mask_token_id`,"
69 - " `pard_token`, or `ptd_token_id` for parallel drafting."
74 + " `mask_token_id`, `dspark_noise_token_id`, `pard_token`, or"
75 + " `ptd_token_id` for parallel drafting."
76 )