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

[Spec Decode] DFlash2: local convolution + candidate selector (#52816)

Zihan Zhang committed Aug 20, 2026 at 22:27 UTC b389ac29465b33f9e9c534df221ea3c129e9793f
14 files changed +866 -44
tests/models/registry.py
+8
@@ -1501,6 +1501,14 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1501 max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env
1502 max_num_seqs=32,
1503 ),
1504 + "DFlash2DraftModel": _HfExamplesInfo(
1505 + "Qwen/Qwen3.8-27B",
1506 + speculative_model="z-lab/Qwen3.8-27B-DFlash2",
1507 + is_available_online=False,
1508 + use_original_num_layers=True,
1509 + max_model_len=8192,
1510 + max_num_seqs=32,
1511 + ),
1512 "DFlashLagunaForCausalLM": _HfExamplesInfo(
1513 "poolside/Laguna-XS-2.1-NVFP4",
1514 speculative_model="poolside/Laguna-XS-2.1-DFlash-NVFP4",
tests/test_config.py
+24
@@ -199,6 +199,7 @@ def test_dsa_models_default_to_mrv2_and_breakable_cudagraph(
199 ),
200 )
201 config._dflash_needs_multi_kv_group = lambda: False
202 + config._is_dflash2_draft = lambda: False
203 config._is_default_v2_model_runner_model = lambda: (
204 VllmConfig._is_default_v2_model_runner_model(config)
205 )
@@ -292,6 +293,29 @@ def test_v2_model_runner_supports_extract_hidden_states():
293 assert config._get_v2_model_runner_unsupported_features() == []
294
295
296 +def test_dflash2_draft_forces_v2_model_runner():
297 + """A DFlash2 draft must reach the V2 speculator, the only one that runs its
298 + candidate selector; on V1 it would draft as DFlash1 without raising."""
299 +
300 + def config(method, architectures):
301 + return SimpleNamespace(
302 + speculative_config=SimpleNamespace(
303 + method=method,
304 + draft_model_config=SimpleNamespace(architectures=architectures),
305 + )
306 + )
307 +
308 + assert VllmConfig._is_dflash2_draft(config("dflash", ["DFlash2DraftModel"]))
309 + assert not VllmConfig._is_dflash2_draft(config("dflash", ["DFlashDraftModel"]))
310 + assert not VllmConfig._is_dflash2_draft(config("eagle", ["DFlash2DraftModel"]))
311 + assert not VllmConfig._is_dflash2_draft(SimpleNamespace(speculative_config=None))
312 + assert not VllmConfig._is_dflash2_draft(
313 + SimpleNamespace(
314 + speculative_config=SimpleNamespace(method="dflash", draft_model_config=None)
315 + )
316 + )
317 +
318 +
319 @pytest.mark.parametrize(
320 ("use_v2_model_runner", "expected_capture_sizes"),
321 [
tests/v1/spec_decode/test_dflash2.py new
+118
@@ -0,0 +1,118 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from types import SimpleNamespace
5 +
6 +import pytest
7 +import torch
8 +
9 +from vllm.model_executor.models.qwen3_dflash2 import _grouped_conv, _score_edges
10 +from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
11 +from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import DFlash2Speculator
12 +
13 +
14 +@pytest.mark.parametrize("block_size", [5, 8])
15 +def test_grouped_conv_matches_reference(block_size: int):
16 + torch.manual_seed(0)
17 + batch, taps, num_groups, group_size = 3, 3, 4, 2
18 + hidden = torch.randn(batch * block_size, num_groups * group_size)
19 + delta = torch.randn(batch * block_size, taps, num_groups)
20 + base = torch.randn(taps, num_groups * group_size)
21 +
22 + actual = _grouped_conv(
23 + hidden, delta, base, block_size, num_groups, group_size, taps
24 + )
25 + hidden_blocks = hidden.view(batch, block_size, num_groups, group_size)
26 + expected = torch.zeros_like(hidden_blocks)
27 + base = base.view(taps, num_groups, group_size)
28 + delta = delta.view(batch, block_size, taps, num_groups)
29 + for position in range(block_size):
30 + for tap in range(min(taps, position + 1)):
31 + expected[:, position] += (
32 + base[tap] + delta[:, position, tap, :, None]
33 + ) * hidden_blocks[:, position - tap]
34 +
35 + torch.testing.assert_close(actual, expected.flatten(0, 1).flatten(-2))
36 +
37 +
38 +def test_selector_edges_match_sequential_reference():
39 + torch.manual_seed(1)
40 + batch, steps, top_k, rank = 2, 4, 3, 5
41 + vocab = 17
42 + predecessors = torch.randn(vocab, rank)
43 + successors = torch.randn(vocab, rank)
44 + candidate_ids = torch.randint(vocab, (batch, steps, top_k))
45 + unary = torch.randn(batch, steps, top_k)
46 + hidden = torch.randn(batch, steps, rank)
47 + anchors = torch.randint(vocab, (batch,))
48 +
49 + actual = _score_edges(
50 + predecessors,
51 + successors,
52 + candidate_ids,
53 + unary,
54 + hidden,
55 + anchors,
56 + top_k,
57 + )
58 + expected = torch.empty_like(actual)
59 + for step in range(steps):
60 + pred = (
61 + anchors[:, None].expand(-1, top_k)
62 + if step == 0
63 + else candidate_ids[:, step - 1]
64 + )
65 + expected[:, step] = unary[:, step, None] + torch.einsum(
66 + "bpr,bcr->bpc",
67 + predecessors[pred] * hidden[:, step, None],
68 + successors[candidate_ids[:, step]],
69 + )
70 +
71 + torch.testing.assert_close(actual, expected)
72 +
73 +
74 +def _stub_base(monkeypatch, draft_logits):
75 + """A DFlashSpeculator.__init__ that allocates only what the base class would.
76 +
77 + The real base class fills draft_logits from draft_logits_spec, so callers
78 + pass a tensor already in that state.
79 + """
80 +
81 + def init_base(self, _vllm_config, device):
82 + self.draft_model_config = SimpleNamespace(
83 + hf_config=SimpleNamespace(dflash_config={"selector_top_k": 3})
84 + )
85 + self.max_num_reqs = 2
86 + self.num_query_per_req = 5
87 + self.num_speculative_steps = 4
88 + self.vocab_size = 17
89 + self.draft_tokens = torch.empty((2, 4), dtype=torch.int64, device=device)
90 + self.draft_logits = draft_logits
91 +
92 + monkeypatch.setattr(DFlashSpeculator, "__init__", init_base)
93 +
94 +
95 +def test_selector_leaves_greedy_drafting_without_proposal_logits(monkeypatch):
96 + """Greedy is the default, and it caches no proposal distribution.
97 +
98 + The base class allocates draft_logits only for "probabilistic"; verification
99 + reads `draft_logits is None` to decide whether a distribution is on offer, so
100 + allocating one here would claim a proposal the walk never sampled from.
101 + """
102 + _stub_base(monkeypatch, None)
103 + speculator = DFlash2Speculator(None, torch.device("cpu"))
104 +
105 + assert speculator.draft_logits is None
106 +
107 +
108 +def test_selector_asks_for_fp32_proposal_logits():
109 + """The spec the base class allocates from: fp32, filled -inf.
110 +
111 + Not the head dtype -- rounding selector scores to bf16 moves the argmax of a
112 + candidate row often enough that the walk and the rejection sampler checking it
113 + would no longer read the same distribution.
114 + """
115 + dtype, fill = DFlash2Speculator.draft_logits_spec(None, None)
116 +
117 + assert dtype is torch.float32
118 + assert fill == float("-inf")
tests/v1/spec_decode/test_dflash_causality.py
+25 -1
@@ -21,12 +21,13 @@ from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
21 )
22
23
24 -def _config(num_hidden_layers, layer_types=None, causal_override=None):
24 +def _config(num_hidden_layers, layer_types=None, causal_override=None, is_causal=None):
25 dflash_config = None if causal_override is None else {"causal": causal_override}
26 return SimpleNamespace(
27 num_hidden_layers=num_hidden_layers,
28 layer_types=layer_types,
29 dflash_config=dflash_config,
30 + is_causal=is_causal,
31 )
32
33
@@ -40,6 +41,19 @@ def _config(num_hidden_layers, layer_types=None, causal_override=None):
41 _config(2, layer_types=["sliding_attention"] * 2, causal_override=False),
42 True,
43 ),
44 + # DFlash2 stores the explicit attention semantics at the top level.
45 + (
46 + _config(
47 + 2,
48 + layer_types=["sliding_attention"] * 2,
49 + is_causal=False,
50 + ),
51 + True,
52 + ),
53 + (
54 + _config(2, layer_types=["full_attention"] * 2, is_causal=True),
55 + False,
56 + ),
57 # SWA-derived: full-attention layers are non-causal.
58 (_config(2, layer_types=["sliding_attention", "full_attention"]), True),
59 # SWA-derived: all-sliding is fully causal.
@@ -59,6 +73,16 @@ def test_dflash_layer_causal_is_per_layer():
73 assert _dflash_layer_causal(config, 1) is False
74
75
76 +def test_dflash_layer_causal_honors_top_level_override():
77 + config = _config(
78 + 2,
79 + layer_types=["sliding_attention", "full_attention"],
80 + is_causal=False,
81 + )
82 + assert _dflash_layer_causal(config, 0) is False
83 + assert _dflash_layer_causal(config, 1) is False
84 +
85 +
86 def _vllm_config(**draft_config):
87 config = SimpleNamespace(**draft_config)
88 return SimpleNamespace(
vllm/config/vllm.py
+17
@@ -670,6 +670,12 @@ class VllmConfig:
670 if self._dflash_needs_multi_kv_group():
671 return True
672
673 + # The DFlash2 candidate selector exists only in the V2 speculator. On V1
674 + # the same checkpoint drafts through DFlashProposer, which never calls
675 + # it, so the draft degrades to DFlash1 silently. Force V2 as for dspark.
676 + if self._is_dflash2_draft():
677 + return True
678 +
679 if self.model_config is not None and self.model_config.is_diffusion:
680 return True
681
@@ -693,6 +699,17 @@ class VllmConfig:
699
700 return True
701
702 + def _is_dflash2_draft(self) -> bool:
703 + """Whether the DFlash draft is a DFlash2 one, by the architecture the
704 + speculator selects on (v1/worker/gpu/spec_decode/__init__.py)."""
705 + spec = self.speculative_config
706 + if spec is None or spec.method != "dflash":
707 + return False
708 + draft_config = getattr(spec, "draft_model_config", None)
709 + if draft_config is None:
710 + return False
711 + return "DFlash2DraftModel" in (draft_config.architectures or [])
712 +
713 def _dflash_needs_multi_kv_group(self) -> bool:
714 """Whether a DFlash draft mixes sliding-window and full attention."""
715 spec = self.speculative_config
vllm/model_executor/layers/logits_processor.py
+81
@@ -2,6 +2,9 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """A layer that compute logits from hidden_stats."""
4
5 +from collections.abc import Callable
6 +from functools import cache
7 +
8 import torch
9 import torch.nn.functional as F
10
@@ -10,12 +13,43 @@ from vllm.distributed import (
13 tensor_model_parallel_all_gather,
14 tensor_model_parallel_gather,
15 )
16 +from vllm.logger import init_logger
17 from vllm.model_executor.custom_op import PluggableLayer
18 from vllm.model_executor.layers.vocab_parallel_embedding import (
19 UnquantizedEmbeddingMethod,
20 VocabParallelEmbedding,
21 )
22 from vllm.platforms import current_platform
23 +from vllm.utils.flashinfer import has_flashinfer
24 +
25 +logger = init_logger(__name__)
26 +
27 +
28 +@cache
29 +def _flashinfer_topk() -> Callable[..., tuple[torch.Tensor, torch.Tensor]] | None:
30 + """FlashInfer's radix top-k, or None for torch.topk.
31 +
32 + The top-k spans the vocabulary, where the radix kernel is about twice
33 + torch.topk.
34 + """
35 + if not current_platform.is_cuda():
36 + return None
37 + if not has_flashinfer():
38 + logger.info_once(
39 + "flashinfer is unavailable; vocab-parallel top-k uses torch.topk, "
40 + "at roughly half the speed."
41 + )
42 + return None
43 + from flashinfer import top_k
44 +
45 + return top_k
46 +
47 +
48 +def _topk(scores: torch.Tensor, k: int) -> tuple[torch.Tensor, torch.Tensor]:
49 + impl = _flashinfer_topk()
50 + if impl is None or not scores.is_cuda:
51 + return torch.topk(scores, k, dim=-1)
52 + return impl(scores, k, sorted=True, deterministic=True)
53
54
55 # --8<-- [start:logits_processor]
@@ -204,6 +238,53 @@ class LogitsProcessor(PluggableLayer):
238 top_tokens = gathered[:, :, 1].gather(dim=-1, index=max_rank_idx)
239 return top_tokens.squeeze(-1).to(torch.int64)
240
241 + def get_top_k_tokens(
242 + self,
243 + lm_head: VocabParallelEmbedding,
244 + hidden_states: torch.Tensor,
245 + k: int,
246 + embedding_bias: torch.Tensor | None = None,
247 + ) -> tuple[torch.Tensor, torch.Tensor]:
248 + """Vocab-parallel top-k without all-gathering full logits.
249 +
250 + The `get_top_tokens` reduction widened from one token to k, returning
251 + the values as well as the global ids. Communication is
252 + O(batch * 2k * tp_size) rather than O(batch * vocab_size).
253 +
254 + Scale and soft cap are applied to the k selected values rather than
255 + the whole vocabulary; both are monotonic, so the selection is the same
256 + and only k entries are touched.
257 + """
258 + if self.scale <= 0.0 and self.scale != 1.0:
259 + raise ValueError(
260 + "The local top-k reduction optimization is not supported for "
261 + "non-positive logit scaling factors."
262 + )
263 +
264 + logits = self._apply_head(lm_head, hidden_states, embedding_bias)
265 +
266 + # Mask out padding entries beyond org_vocab_size on this shard.
267 + num_pad = lm_head.shard_indices.num_org_vocab_padding
268 + if num_pad > 0:
269 + logits[..., -num_pad:] = -float("inf")
270 +
271 + values, ids = _topk(logits, k)
272 + # Convert shard-local indices to global vocab indices.
273 + ids = ids.to(torch.int64) + lm_head.shard_indices.org_vocab_start_index
274 +
275 + if lm_head.tp_size > 1:
276 + values = tensor_model_parallel_all_gather(values, dim=-1)
277 + ids = tensor_model_parallel_all_gather(ids, dim=-1)
278 + values, selected = _topk(values, k)
279 + ids = ids.gather(-1, selected)
280 +
281 + values = values.float()
282 + if self.scale != 1.0:
283 + values = values * self.scale
284 + if self.soft_cap is not None:
285 + values = torch.tanh(values / self.soft_cap) * self.soft_cap
286 + return ids, values
287 +
288 def extra_repr(self) -> str:
289 s = f"vocab_size={self.vocab_size}"
290 s += f", org_vocab_size={self.org_vocab_size}"
vllm/model_executor/models/qwen3_dflash.py
+11 -4
@@ -56,10 +56,13 @@ _SLIDING_ATTENTION = "sliding_attention"
56
57
58 def _dflash_layer_causal(config: Qwen3Config, layer_idx: int) -> bool:
59 - """``dflash_config.causal`` overrides all layers; else only SWA layers causal."""
59 + """Resolve explicit causality before falling back to legacy layer defaults."""
60 + is_causal = getattr(config, "is_causal", None)
61 + if is_causal is not None:
62 + return bool(is_causal)
63 override = (getattr(config, "dflash_config", None) or {}).get("causal")
64 if override is not None:
62 - return override
65 + return bool(override)
66 layer_types = getattr(config, "layer_types", None)
67 return bool(layer_types) and layer_types[layer_idx] == _SLIDING_ATTENTION
68
@@ -374,6 +377,8 @@ class DFlashQwen3DecoderLayer(nn.Module):
377
378 @support_torch_compile
379 class DFlashQwen3Model(nn.Module):
380 + decoder_layer_cls = DFlashQwen3DecoderLayer
381 +
382 hf_to_vllm_mapper = WeightsMapper(
383 orig_to_new_substr={
384 "midlayer.": "layers.0.",
@@ -434,7 +439,7 @@ class DFlashQwen3Model(nn.Module):
439
440 self.layers = nn.ModuleList(
441 [
437 - DFlashQwen3DecoderLayer(
442 + self.decoder_layer_cls(
443 current_vllm_config,
444 config=self.config,
445 layer_idx=layer_idx,
@@ -699,6 +704,8 @@ class DFlashQwen3Model(nn.Module):
704
705
706 class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):
707 + model_cls = DFlashQwen3Model
708 +
709 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
710 nn.Module.__init__(self)
711 self.draft_model_config = vllm_config.speculative_config.draft_model_config
@@ -708,7 +715,7 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):
715 target_layer_num = vllm_config.model_config.get_num_layers(
716 vllm_config.parallel_config
717 )
711 - self.model = DFlashQwen3Model(
718 + self.model = self.model_cls(
719 vllm_config=vllm_config,
720 prefix=maybe_prefix(prefix, "model"),
721 start_layer_id=target_layer_num,
vllm/model_executor/models/qwen3_dflash2.py new
+290
@@ -0,0 +1,290 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +import torch.nn.functional as F
6 +from torch import nn
7 +
8 +from vllm.compilation.backends import set_model_tag
9 +from vllm.compilation.decorators import support_torch_compile
10 +from vllm.config import CacheConfig, VllmConfig
11 +from vllm.model_executor.layers.linear import ReplicatedLinear
12 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
13 +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
14 +
15 +from .qwen3_dflash import (
16 + DFlashQwen3DecoderLayer,
17 + DFlashQwen3ForCausalLM,
18 + DFlashQwen3Model,
19 +)
20 +from .utils import maybe_prefix
21 +
22 +
23 +def _grouped_conv(
24 + hidden_states: torch.Tensor,
25 + delta: torch.Tensor,
26 + base: torch.Tensor,
27 + block_size: int,
28 + num_groups: int,
29 + group_size: int,
30 + taps: int,
31 +) -> torch.Tensor:
32 + blocks = hidden_states.unflatten(-1, (num_groups, group_size))
33 + coefficients = base.view(1, taps, num_groups, group_size) + delta.unsqueeze(-1)
34 + output = coefficients[:, 0] * blocks
35 + position = torch.arange(hidden_states.shape[0], device=hidden_states.device)
36 + if block_size & (block_size - 1) == 0:
37 + position = position & (block_size - 1)
38 + else:
39 + position = position % block_size
40 + for tap in range(1, taps):
41 + shifted = F.pad(blocks[:-tap], (0, 0, 0, 0, tap, 0))
42 + output += coefficients[:, tap] * shifted * (position >= tap).view(-1, 1, 1)
43 + return output.flatten(-2)
44 +
45 +
46 +class DFlashGroupedConv(nn.Module):
47 + def __init__(
48 + self,
49 + hidden_size: int,
50 + taps: int,
51 + group_size: int,
52 + block_size: int,
53 + params_dtype: torch.dtype,
54 + prefix: str,
55 + ) -> None:
56 + super().__init__()
57 + if hidden_size % group_size:
58 + raise ValueError(
59 + f"conv_group_size={group_size} must divide hidden_size={hidden_size}."
60 + )
61 + self.block_size = block_size
62 + self.taps = taps
63 + self.group_size = group_size
64 + self.num_groups = hidden_size // group_size
65 + self.base_kernel = nn.Parameter(
66 + torch.empty(2, taps, hidden_size, dtype=params_dtype),
67 + requires_grad=False,
68 + )
69 + self.kernel_projection = ReplicatedLinear(
70 + hidden_size,
71 + 2 * taps * self.num_groups,
72 + bias=False,
73 + params_dtype=params_dtype,
74 + quant_config=None,
75 + prefix=maybe_prefix(prefix, "kernel_projection"),
76 + return_bias=False,
77 + )
78 +
79 + def _convolve(
80 + self, hidden_states: torch.Tensor, delta: torch.Tensor, side: int
81 + ) -> torch.Tensor:
82 + return _grouped_conv(
83 + hidden_states,
84 + delta,
85 + self.base_kernel[side],
86 + self.block_size,
87 + self.num_groups,
88 + self.group_size,
89 + self.taps,
90 + )
91 +
92 + def prepare(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
93 + coefficients = self.kernel_projection(hidden_states).reshape(
94 + hidden_states.shape[0], 2, self.taps, self.num_groups
95 + )
96 + return self._convolve(hidden_states, coefficients[:, 0], 0), coefficients[:, 1]
97 +
98 + def finish(
99 + self, hidden_states: torch.Tensor, coefficients: torch.Tensor
100 + ) -> torch.Tensor:
101 + return self._convolve(hidden_states, coefficients, 1)
102 +
103 +
104 +class DFlash2Qwen3DecoderLayer(DFlashQwen3DecoderLayer):
105 + def __init__(
106 + self,
107 + vllm_config: VllmConfig,
108 + *,
109 + config,
110 + layer_idx: int,
111 + cache_config: CacheConfig | None = None,
112 + quant_config: QuantizationConfig | None = None,
113 + prefix: str = "",
114 + ) -> None:
115 + super().__init__(
116 + vllm_config,
117 + config=config,
118 + layer_idx=layer_idx,
119 + cache_config=cache_config,
120 + quant_config=quant_config,
121 + prefix=prefix,
122 + )
123 + draft_config = config.dflash_config
124 + speculative_config = vllm_config.speculative_config
125 + assert speculative_config is not None
126 + conv_args = dict(
127 + hidden_size=config.hidden_size,
128 + taps=int(draft_config["conv_kernel_size"]),
129 + group_size=int(draft_config["conv_group_size"]),
130 + # Query tokens per request: the bonus token plus the mask tokens.
131 + block_size=1 + speculative_config.num_speculative_tokens,
132 + params_dtype=vllm_config.model_config.dtype,
133 + )
134 + self.attention_conv = DFlashGroupedConv(
135 + **conv_args, prefix=maybe_prefix(prefix, "attention_conv")
136 + )
137 + self.mlp_conv = DFlashGroupedConv(
138 + **conv_args, prefix=maybe_prefix(prefix, "mlp_conv")
139 + )
140 +
141 + def forward(
142 + self,
143 + positions: torch.Tensor,
144 + hidden_states: torch.Tensor,
145 + residual: torch.Tensor | None,
146 + ) -> tuple[torch.Tensor, torch.Tensor]:
147 + if residual is None:
148 + residual = hidden_states
149 + hidden_states = self.input_layernorm(hidden_states)
150 + else:
151 + hidden_states, residual = self.input_layernorm(hidden_states, residual)
152 +
153 + hidden_states, coefficients = self.attention_conv.prepare(hidden_states)
154 + hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
155 + hidden_states = self.attention_conv.finish(hidden_states, coefficients)
156 +
157 + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
158 + hidden_states, coefficients = self.mlp_conv.prepare(hidden_states)
159 + hidden_states = self.mlp(hidden_states)
160 + hidden_states = self.mlp_conv.finish(hidden_states, coefficients)
161 + return hidden_states, residual
162 +
163 +
164 +def _score_edges(
165 + predecessor_table: torch.Tensor,
166 + successor_table: torch.Tensor,
167 + candidate_ids: torch.Tensor,
168 + unary_logits: torch.Tensor,
169 + hidden: torch.Tensor,
170 + anchor_token_ids: torch.Tensor,
171 + top_k: int,
172 +) -> torch.Tensor:
173 + successors = successor_table[candidate_ids]
174 + predecessor_ids = torch.cat(
175 + (
176 + anchor_token_ids[:, None, None].expand(-1, 1, top_k),
177 + candidate_ids[:, :-1],
178 + ),
179 + dim=1,
180 + )
181 + predecessors = predecessor_table[predecessor_ids]
182 + return unary_logits[:, :, None] + torch.einsum(
183 + "blpr,blcr->blpc", predecessors * hidden[:, :, None], successors
184 + )
185 +
186 +
187 +@support_torch_compile
188 +class CandidateSelector(nn.Module):
189 + def __init__(
190 + self,
191 + hidden_size: int,
192 + vocab_size: int,
193 + rank: int,
194 + top_k: int,
195 + params_dtype: torch.dtype,
196 + prefix: str,
197 + ) -> None:
198 + super().__init__()
199 + self.top_k = top_k
200 + self.predecessor_codebook = nn.Parameter(
201 + torch.empty(vocab_size, rank, dtype=params_dtype), requires_grad=False
202 + )
203 + self.successor_codebook = nn.Parameter(
204 + torch.empty(vocab_size, rank, dtype=params_dtype), requires_grad=False
205 + )
206 + self.hidden_projection = ReplicatedLinear(
207 + hidden_size,
208 + rank,
209 + bias=False,
210 + params_dtype=params_dtype,
211 + quant_config=None,
212 + prefix=maybe_prefix(prefix, "hidden_projection"),
213 + return_bias=False,
214 + )
215 +
216 + def forward(
217 + self,
218 + candidate_ids: torch.Tensor,
219 + unary_logits: torch.Tensor,
220 + hidden_states: torch.Tensor,
221 + anchor_token_ids: torch.Tensor,
222 + ) -> torch.Tensor:
223 + hidden = self.hidden_projection(hidden_states)
224 + return _score_edges(
225 + self.predecessor_codebook,
226 + self.successor_codebook,
227 + candidate_ids,
228 + unary_logits,
229 + hidden,
230 + anchor_token_ids,
231 + self.top_k,
232 + )
233 +
234 +
235 +class DFlash2Qwen3Model(DFlashQwen3Model):
236 + decoder_layer_cls = DFlash2Qwen3DecoderLayer
237 +
238 + def __init__(
239 + self,
240 + *,
241 + vllm_config: VllmConfig,
242 + start_layer_id: int = 0,
243 + prefix: str = "",
244 + ) -> None:
245 + super().__init__(
246 + vllm_config=vllm_config,
247 + start_layer_id=start_layer_id,
248 + prefix=prefix,
249 + )
250 + draft_config = self.config.dflash_config
251 + self.input_embedding_scale = float(
252 + draft_config.get("input_embedding_scale", 1.0)
253 + )
254 + # Without its own tag the selector shares the draft head's compile cache.
255 + with set_model_tag("dflash2_candidate_selector"):
256 + self.candidate_selector = CandidateSelector(
257 + hidden_size=self.config.hidden_size,
258 + vocab_size=self.config.vocab_size,
259 + rank=int(draft_config["selector_rank"]),
260 + top_k=int(draft_config["selector_top_k"]),
261 + params_dtype=vllm_config.model_config.dtype,
262 + prefix=maybe_prefix(prefix, "candidate_selector"),
263 + )
264 +
265 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
266 + return super().embed_input_ids(input_ids) * self.input_embedding_scale
267 +
268 +
269 +class DFlash2Qwen3ForCausalLM(DFlashQwen3ForCausalLM):
270 + model_cls = DFlash2Qwen3Model
271 +
272 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
273 + super().__init__(vllm_config=vllm_config, prefix=prefix)
274 + draft_config = self.config.dflash_config
275 + softcap = float(draft_config.get("final_logit_softcapping") or 0.0)
276 + self.candidate_logits_processor = LogitsProcessor(
277 + vllm_config.model_config.get_vocab_size(),
278 + scale=float(draft_config.get("output_multiplier", 1.0)),
279 + soft_cap=softcap if softcap > 0 else None,
280 + )
281 +
282 + def compute_candidates(
283 + self, hidden_states: torch.Tensor
284 + ) -> tuple[torch.Tensor, torch.Tensor]:
285 + return self.candidate_logits_processor.get_top_k_tokens(
286 + self.lm_head, hidden_states, self.model.candidate_selector.top_k
287 + )
288 +
289 +
290 +EntryClass = DFlash2Qwen3ForCausalLM
vllm/model_executor/models/registry.py
+1
@@ -632,6 +632,7 @@ _SPECULATIVE_DECODING_MODELS = {
632 "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),
633 "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
634 "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
635 + "DFlash2DraftModel": ("qwen3_dflash2", "DFlash2Qwen3ForCausalLM"),
636 # Muse Glimmer's DFlash draft head, reusing the generic qwen3_dflash
637 # implementation. EAGLEConfig rewrites a dflash draft's architecture to
638 # DFlash{arch} unless it already starts or ends with "DFlash" (see
vllm/v1/worker/gpu/sample/gumbel.py
+52 -34
@@ -81,6 +81,46 @@ def tl_rand32(seed, offset, includes_zero: tl.constexpr):
81 return u
82
83
84 +@triton.jit
85 +def gumbel_noised_argmax(
86 + logits,
87 + keys,
88 + mask,
89 + seed,
90 + pos,
91 + temp,
92 + USE_FP64: tl.constexpr,
93 + APPLY_TEMPERATURE: tl.constexpr = True,
94 +):
95 + """Argmax of logits under Gumbel-max sampling, or plain argmax at temp 0.
96 +
97 + `keys` indexes the noise, so the same token draws the same noise wherever it
98 + appears; `pos` and `seed` place the draw in the request's stream, which is
99 + what lets a draft and its verification agree.
100 + """
101 + if temp != 0.0 and APPLY_TEMPERATURE:
102 + # Match the behavior of _temperature_kernel: if that kernel uses
103 + # tl.div_rn, this must too.
104 + logits = logits / temp
105 +
106 + # fp32 is the default reduction dtype; fp64 is ~1/32-1/64x the throughput
107 + # on H100/Ada/Blackwell and empirically indistinguishable for Gumbel-max.
108 + if USE_FP64:
109 + logits = logits.to(tl.float64)
110 + if temp != 0.0:
111 + gumbel_seed = tl.randint(seed, pos)
112 + if USE_FP64:
113 + u = tl_rand64(gumbel_seed, keys, includes_zero=False)
114 + gumbel_noise = -tl.log(-tl.log(u))
115 + else:
116 + u = tl_rand32(gumbel_seed, keys, includes_zero=False)
117 + # log1p keeps the winning tail at u -> 0, where fp32 resolves it.
118 + gumbel_noise = -tl.log(-tldevice.log1p(-u))
119 + logits = tl.where(mask, logits + gumbel_noise, float("-inf"))
120 +
121 + return tl.max(logits, axis=0, return_indices=True)
122 +
123 +
124 @triton.jit
125 def gumbel_block_argmax(
126 logits,
@@ -124,40 +164,18 @@ def gumbel_block_argmax(
164 mask=mask & is_valid_req,
165 )
166
127 - if temp != 0.0 and APPLY_TEMPERATURE:
128 - # Apply temperature.
129 - # NOTE(woosuk): Match the behavior of _temperature_kernel.
130 - # E.g., if the kernel uses tl.div_rn, we should use tl.div_rn here too.
131 - logits = logits / temp
132 -
133 - # fp32 is the default reduction dtype; fp64 is ~1/32–1/64x the throughput
134 - # on H100/Ada/Blackwell and empirically indistinguishable for Gumbel-max.
135 - if USE_FP64:
136 - logits = logits.to(tl.float64)
137 - if temp != 0.0:
138 - # Calculate the seed for gumbel noise.
139 - seed = tl.load(seeds_ptr + req_state_idx, mask=is_valid_req, other=0)
140 - pos = tl.load(pos_ptr + token_idx)
141 - gumbel_seed = tl.randint(seed, pos)
142 -
143 - if USE_FP64:
144 - u = tl_rand64(gumbel_seed, block, includes_zero=False)
145 - gumbel_noise = -tl.log(-tl.log(u))
146 - else:
147 - u = tl_rand32(gumbel_seed, block, includes_zero=False)
148 - # Draw the large-noise tail (which decides the argmax winner) from u -> 0,
149 - # where fp32 has fine resolution, instead of u -> 1, where fp32 spacing is
150 - # ~2**-24. The naive `-log(-log(u))` puts the winning tail at u -> 1,
151 - # hard-capping the noise at ~16.6 and coarsely quantizing it; using
152 - # `log1p(-u)` == `log(1 - u)` keeps the tail in the well-resolved region.
153 - # Note `1 - u` would lose precision for small u, so `log1p` is required.
154 - gumbel_noise = -tl.log(-tldevice.log1p(-u))
155 -
156 - # Apply gumbel noise.
157 - logits = tl.where(mask, logits + gumbel_noise, float("-inf"))
158 -
159 - value, idx = tl.max(logits, axis=0, return_indices=True)
160 - return value, idx
167 + seed = tl.load(seeds_ptr + req_state_idx, mask=is_valid_req, other=0)
168 + pos = tl.load(pos_ptr + token_idx)
169 + return gumbel_noised_argmax(
170 + logits,
171 + block,
172 + mask,
173 + seed,
174 + pos,
175 + temp,
176 + USE_FP64=USE_FP64,
177 + APPLY_TEMPERATURE=APPLY_TEMPERATURE,
178 + )
179
180
181 @triton.jit
vllm/v1/worker/gpu/spec_decode/__init__.py
+6
@@ -15,6 +15,12 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device):
15
16 return ExtractHiddenStatesSpeculator(vllm_config, device)
17 elif speculative_config.method == "dflash":
18 + if "DFlash2DraftModel" in speculative_config.draft_model_config.architectures:
19 + from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import (
20 + DFlash2Speculator,
21 + )
22 +
23 + return DFlash2Speculator(vllm_config, device)
24 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import (
25 DFlashSpeculator,
26 )
vllm/v1/worker/gpu/spec_decode/dflash2/__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/dflash2/speculator.py new
+215
@@ -0,0 +1,215 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import Any
5 +
6 +import torch
7 +
8 +from vllm.config import VllmConfig
9 +from vllm.config.compilation import CUDAGraphMode
10 +from vllm.triton_utils import tl, triton
11 +from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax
12 +from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
13 +
14 +
15 +@triton.jit
16 +def _selector_walk_kernel(
17 + scores_ptr,
18 + candidate_ptr,
19 + sample_pos_ptr,
20 + req_state_ptr,
21 + temperature_ptr,
22 + seeds_ptr,
23 + tokens_ptr,
24 + realized_scores_ptr,
25 + num_steps: tl.constexpr,
26 + top_k: tl.constexpr,
27 + BLOCK_K: tl.constexpr,
28 + SAMPLE_PROBABILISTIC: tl.constexpr,
29 + USE_FP64: tl.constexpr,
30 +):
31 + row = tl.program_id(0)
32 + offsets = tl.arange(0, BLOCK_K)
33 + mask = offsets < top_k
34 + req_state = tl.load(req_state_ptr + row * num_steps)
35 + valid = req_state >= 0
36 + temperature = tl.load(temperature_ptr + req_state, mask=valid, other=0.0)
37 + seed = tl.load(seeds_ptr + req_state, mask=valid, other=0)
38 + previous = 0
39 + for step in range(num_steps):
40 + flat = row * num_steps + step
41 + score_base = (flat * top_k + previous) * top_k
42 + scores = tl.load(
43 + scores_ptr + score_base + offsets,
44 + mask=mask & valid,
45 + other=float("-inf"),
46 + ).to(tl.float64 if USE_FP64 else tl.float32)
47 + candidate_base = flat * top_k
48 + candidates = tl.load(
49 + candidate_ptr + candidate_base + offsets,
50 + mask=mask & valid,
51 + other=0,
52 + )
53 +
54 + # Candidate ids key the noise, matching the target's own sampling.
55 + position = tl.load(sample_pos_ptr + flat) - 1
56 + _, index = gumbel_noised_argmax(
57 + scores,
58 + candidates,
59 + mask & valid,
60 + seed,
61 + position,
62 + temperature if SAMPLE_PROBABILISTIC else 0.0,
63 + USE_FP64=USE_FP64,
64 + )
65 +
66 + tl.store(
67 + realized_scores_ptr + candidate_base + offsets,
68 + scores,
69 + mask=mask & valid,
70 + )
71 + token = tl.load(candidate_ptr + candidate_base + index, mask=valid, other=0)
72 + tl.store(tokens_ptr + flat, token, mask=valid)
73 + previous = index
74 +
75 +
76 +@triton.jit
77 +def _cache_draft_logits_kernel(
78 + draft_logits_ptr,
79 + cached_candidate_ptr,
80 + candidate_ptr,
81 + scores_ptr,
82 + req_state_ptr,
83 + draft_logits_stride_0,
84 + draft_logits_stride_1,
85 + num_steps: tl.constexpr,
86 + top_k: tl.constexpr,
87 + BLOCK_K: tl.constexpr,
88 +):
89 + flat = tl.program_id(0)
90 + req_state = tl.load(req_state_ptr + flat)
91 + step = flat % num_steps
92 + offsets = tl.arange(0, BLOCK_K)
93 + mask = (req_state >= 0) & (offsets < top_k)
94 + candidate_base = flat * top_k
95 + cache_base = (req_state * num_steps + step) * top_k
96 + old_token_ids = tl.load(cached_candidate_ptr + cache_base + offsets, mask=mask)
97 + logits_base = (
98 + draft_logits_ptr
99 + + req_state * draft_logits_stride_0
100 + + step * draft_logits_stride_1
101 + )
102 + tl.store(logits_base + old_token_ids, -float("inf"), mask=mask)
103 + token_ids = tl.load(candidate_ptr + candidate_base + offsets, mask=mask)
104 + scores = tl.load(scores_ptr + candidate_base + offsets, mask=mask)
105 + tl.store(logits_base + token_ids, scores, mask=mask)
106 + tl.store(cached_candidate_ptr + cache_base + offsets, token_ids, mask=mask)
107 +
108 +
109 +class DFlash2Speculator(DFlashSpeculator):
110 + _speculator_name = "DFlash2"
111 +
112 + def __init__(self, vllm_config: VllmConfig, device: torch.device):
113 + super().__init__(vllm_config, device)
114 + draft_config = self.draft_model_config.hf_config.dflash_config
115 + self.selector_top_k = int(draft_config["selector_top_k"])
116 + self._anchor_indices = (
117 + torch.arange(self.max_num_reqs, dtype=torch.int64, device=device)
118 + * self.num_query_per_req
119 + )
120 + self._selector_scores = torch.empty(
121 + self.max_num_reqs,
122 + self.num_speculative_steps,
123 + self.selector_top_k,
124 + dtype=torch.float32,
125 + device=device,
126 + )
127 + self._cached_candidate_ids = torch.zeros(
128 + self._selector_scores.shape, dtype=torch.int64, device=device
129 + )
130 +
131 + def draft_logits_spec(self, vllm_config: VllmConfig) -> tuple[torch.dtype, float]:
132 + # fp32 so the walk and the rejection that checks it read the same
133 + # distribution; -inf because the cache kernel writes only the K
134 + # candidates.
135 + return torch.float32, -float("inf")
136 +
137 + def _sample_path(
138 + self,
139 + candidate_ids: torch.Tensor,
140 + scores: torch.Tensor,
141 + num_reqs: int,
142 + ) -> None:
143 + block_k = triton.next_power_of_2(self.selector_top_k)
144 + _selector_walk_kernel[(num_reqs,)](
145 + scores.contiguous(),
146 + candidate_ids.contiguous(),
147 + self.sample_pos,
148 + self.sample_idx_mapping,
149 + self.temperature,
150 + self.seeds,
151 + self.draft_tokens,
152 + self._selector_scores,
153 + num_steps=self.num_speculative_steps,
154 + top_k=self.selector_top_k,
155 + BLOCK_K=block_k,
156 + SAMPLE_PROBABILISTIC=self.draft_logits is not None,
157 + USE_FP64=self.use_fp64_gumbel,
158 + num_warps=1,
159 + )
160 +
161 + def _cache_draft_logits(self, candidate_ids: torch.Tensor, num_sample: int) -> None:
162 + draft_logits = self.draft_logits
163 + assert draft_logits is not None
164 + block_k = triton.next_power_of_2(self.selector_top_k)
165 + _cache_draft_logits_kernel[(num_sample,)](
166 + draft_logits,
167 + self._cached_candidate_ids,
168 + candidate_ids,
169 + self._selector_scores,
170 + self.sample_idx_mapping,
171 + draft_logits.stride(0),
172 + draft_logits.stride(1),
173 + num_steps=self.num_speculative_steps,
174 + top_k=self.selector_top_k,
175 + BLOCK_K=block_k,
176 + num_warps=1,
177 + )
178 +
179 + def _generate_draft(
180 + self,
181 + num_reqs: int,
182 + num_tokens_padded: int,
183 + attn_metadata: dict[str, Any] | None,
184 + slot_mappings: dict[str, torch.Tensor] | None,
185 + num_tokens_across_dp: torch.Tensor | None,
186 + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
187 + ) -> None:
188 + last_hidden_states = self._run_model(
189 + num_tokens_padded,
190 + attn_metadata,
191 + slot_mappings,
192 + num_tokens_across_dp,
193 + cudagraph_runtime_mode,
194 + )
195 + num_sample = num_reqs * self.num_speculative_steps
196 + hidden_states = last_hidden_states[self.sample_indices[:num_sample]].view(
197 + num_reqs, self.num_speculative_steps, -1
198 + )
199 + candidate_ids, unary_logits = self.model.compute_candidates(
200 + hidden_states.flatten(0, 1)
201 + )
202 + candidate_ids = candidate_ids.view(
203 + num_reqs, self.num_speculative_steps, self.selector_top_k
204 + )
205 + unary_logits = unary_logits.view_as(candidate_ids)
206 + anchor_token_ids = self.input_buffers.input_ids[self._anchor_indices[:num_reqs]]
207 + scores = self.model.model.candidate_selector(
208 + candidate_ids,
209 + unary_logits,
210 + hidden_states,
211 + anchor_token_ids,
212 + )
213 + self._sample_path(candidate_ids, scores, num_reqs)
214 + if self.draft_logits is not None:
215 + self._cache_draft_logits(candidate_ids, num_sample)
vllm/v1/worker/gpu/spec_decode/speculator.py
+16 -5
@@ -132,11 +132,15 @@ class DraftModelSpeculator(BaseSpeculator):
132 self.draft_logits: torch.Tensor | None = None
133 if self.speculative_config.draft_sample_method == "probabilistic":
134 # Pre-temperature logits, cached from the previous decode step.
135 - self.draft_logits = torch.zeros(
136 - self.max_num_reqs,
137 - self.num_speculative_steps,
138 - self.vocab_size,
139 - dtype=vllm_config.model_config.head_dtype,
135 + dtype, fill = self.draft_logits_spec(vllm_config)
136 + self.draft_logits = torch.full(
137 + (
138 + self.max_num_reqs,
139 + self.num_speculative_steps,
140 + self.vocab_size,
141 + ),
142 + fill,
143 + dtype=dtype,
144 device=device,
145 )
146
@@ -293,6 +297,13 @@ class DraftModelSpeculator(BaseSpeculator):
297 )
298 return attn_metadata
299
300 + def draft_logits_spec(self, vllm_config: VllmConfig) -> tuple[torch.dtype, float]:
301 + """Dtype and fill for the cached proposal distribution.
302 +
303 + Speculators that write only a subset of columns each step override this.
304 + """
305 + return vllm_config.model_config.head_dtype, 0.0
306 +
307 def _validate_local_argmax_reduction(self) -> None:
308 if not self.use_local_argmax_reduction:
309 return