[Spec Decode] DFlash2: local convolution + candidate selector

tests/models/registry.py

@@ -1501,6 +1501,14 @@ def check_available_online(
         max_model_len=8192,  # Reduce max len to ensure test runs in low-VRAM CI env
         max_num_seqs=32,
     ),
+    "DFlash2DraftModel": _HfExamplesInfo(
+        "Qwen/Qwen3.8-27B",
+        speculative_model="z-lab/Qwen3.8-27B-DFlash2",
+        is_available_online=False,
+        use_original_num_layers=True,
+        max_model_len=8192,
+        max_num_seqs=32,
+    ),
     "DFlashLagunaForCausalLM": _HfExamplesInfo(
         "poolside/Laguna-XS-2.1-NVFP4",
         speculative_model="poolside/Laguna-XS-2.1-DFlash-NVFP4",

tests/test_config.py

@@ -185,6 +185,7 @@ def test_dsa_models_default_to_mrv2_and_breakable_cudagraph(
         ),
     )
     config._dflash_needs_multi_kv_group = lambda: False
+    config._is_dflash2_draft = lambda: False
     config._is_default_v2_model_runner_model = lambda: (
         VllmConfig._is_default_v2_model_runner_model(config)
     )
@@ -278,6 +279,29 @@ def test_v2_model_runner_supports_extract_hidden_states():
     assert config._get_v2_model_runner_unsupported_features() == []
 
 
+def test_dflash2_draft_forces_v2_model_runner():
+    """A DFlash2 draft must reach the V2 speculator, the only one that runs its
+    candidate selector; on V1 it would draft as DFlash1 without raising."""
+
+    def config(method, architectures):
+        return SimpleNamespace(
+            speculative_config=SimpleNamespace(
+                method=method,
+                draft_model_config=SimpleNamespace(architectures=architectures),
+            )
+        )
+
+    assert VllmConfig._is_dflash2_draft(config("dflash", ["DFlash2DraftModel"]))
+    assert not VllmConfig._is_dflash2_draft(config("dflash", ["DFlashDraftModel"]))
+    assert not VllmConfig._is_dflash2_draft(config("eagle", ["DFlash2DraftModel"]))
+    assert not VllmConfig._is_dflash2_draft(SimpleNamespace(speculative_config=None))
+    assert not VllmConfig._is_dflash2_draft(
+        SimpleNamespace(
+            speculative_config=SimpleNamespace(method="dflash", draft_model_config=None)
+        )
+    )
+
+
 @pytest.mark.parametrize(
     ("use_v2_model_runner", "expected_capture_sizes"),
     [

tests/v1/spec_decode/test_dflash2.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.model_executor.models.qwen3_dflash2 import _grouped_conv, _score_edges
+from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
+from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import DFlash2Speculator
+
+
+@pytest.mark.parametrize("block_size", [5, 8])
+def test_grouped_conv_matches_reference(block_size: int):
+    torch.manual_seed(0)
+    batch, taps, num_groups, group_size = 3, 3, 4, 2
+    hidden = torch.randn(batch * block_size, num_groups * group_size)
+    delta = torch.randn(batch * block_size, taps, num_groups)
+    base = torch.randn(taps, num_groups * group_size)
+
+    actual = _grouped_conv(
+        hidden, delta, base, block_size, num_groups, group_size, taps
+    )
+    hidden_blocks = hidden.view(batch, block_size, num_groups, group_size)
+    expected = torch.zeros_like(hidden_blocks)
+    base = base.view(taps, num_groups, group_size)
+    delta = delta.view(batch, block_size, taps, num_groups)
+    for position in range(block_size):
+        for tap in range(min(taps, position + 1)):
+            expected[:, position] += (
+                base[tap] + delta[:, position, tap, :, None]
+            ) * hidden_blocks[:, position - tap]
+
+    torch.testing.assert_close(actual, expected.flatten(0, 1).flatten(-2))
+
+
+def test_selector_edges_match_sequential_reference():
+    torch.manual_seed(1)
+    batch, steps, top_k, rank = 2, 4, 3, 5
+    vocab = 17
+    predecessors = torch.randn(vocab, rank)
+    successors = torch.randn(vocab, rank)
+    candidate_ids = torch.randint(vocab, (batch, steps, top_k))
+    unary = torch.randn(batch, steps, top_k)
+    hidden = torch.randn(batch, steps, rank)
+    anchors = torch.randint(vocab, (batch,))
+
+    actual = _score_edges(
+        predecessors,
+        successors,
+        candidate_ids,
+        unary,
+        hidden,
+        anchors,
+        top_k,
+    )
+    expected = torch.empty_like(actual)
+    for step in range(steps):
+        pred = (
+            anchors[:, None].expand(-1, top_k)
+            if step == 0
+            else candidate_ids[:, step - 1]
+        )
+        expected[:, step] = unary[:, step, None] + torch.einsum(
+            "bpr,bcr->bpc",
+            predecessors[pred] * hidden[:, step, None],
+            successors[candidate_ids[:, step]],
+        )
+
+    torch.testing.assert_close(actual, expected)
+
+
+def _stub_base(monkeypatch, draft_logits):
+    """A DFlashSpeculator.__init__ that allocates only what the base class would.
+
+    The real base class fills draft_logits from draft_logits_spec, so callers
+    pass a tensor already in that state.
+    """
+
+    def init_base(self, _vllm_config, device):
+        self.draft_model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(dflash_config={"selector_top_k": 3})
+        )
+        self.max_num_reqs = 2
+        self.num_query_per_req = 5
+        self.num_speculative_steps = 4
+        self.vocab_size = 17
+        self.draft_tokens = torch.empty((2, 4), dtype=torch.int64, device=device)
+        self.draft_logits = draft_logits
+
+    monkeypatch.setattr(DFlashSpeculator, "__init__", init_base)
+
+
+def test_selector_leaves_greedy_drafting_without_proposal_logits(monkeypatch):
+    """Greedy is the default, and it caches no proposal distribution.
+
+    The base class allocates draft_logits only for "probabilistic"; verification
+    reads `draft_logits is None` to decide whether a distribution is on offer, so
+    allocating one here would claim a proposal the walk never sampled from.
+    """
+    _stub_base(monkeypatch, None)
+    speculator = DFlash2Speculator(None, torch.device("cpu"))
+
+    assert speculator.draft_logits is None
+
+
+def test_selector_asks_for_fp32_proposal_logits():
+    """The spec the base class allocates from: fp32, filled -inf.
+
+    Not the head dtype -- rounding selector scores to bf16 moves the argmax of a
+    candidate row often enough that the walk and the rejection sampler checking it
+    would no longer read the same distribution.
+    """
+    dtype, fill = DFlash2Speculator.draft_logits_spec(None, None)
+
+    assert dtype is torch.float32
+    assert fill == float("-inf")

tests/v1/spec_decode/test_dflash_causality.py

@@ -21,12 +21,13 @@
 )
 
 
-def _config(num_hidden_layers, layer_types=None, causal_override=None):
+def _config(num_hidden_layers, layer_types=None, causal_override=None, is_causal=None):
     dflash_config = None if causal_override is None else {"causal": causal_override}
     return SimpleNamespace(
         num_hidden_layers=num_hidden_layers,
         layer_types=layer_types,
         dflash_config=dflash_config,
+        is_causal=is_causal,
     )
 
 
@@ -40,6 +41,19 @@ def _config(num_hidden_layers, layer_types=None, causal_override=None):
             _config(2, layer_types=["sliding_attention"] * 2, causal_override=False),
             True,
         ),
+        # DFlash2 stores the explicit attention semantics at the top level.
+        (
+            _config(
+                2,
+                layer_types=["sliding_attention"] * 2,
+                is_causal=False,
+            ),
+            True,
+        ),
+        (
+            _config(2, layer_types=["full_attention"] * 2, is_causal=True),
+            False,
+        ),
         # SWA-derived: full-attention layers are non-causal.
         (_config(2, layer_types=["sliding_attention", "full_attention"]), True),
         # SWA-derived: all-sliding is fully causal.
@@ -59,6 +73,16 @@ def test_dflash_layer_causal_is_per_layer():
     assert _dflash_layer_causal(config, 1) is False
 
 
+def test_dflash_layer_causal_honors_top_level_override():
+    config = _config(
+        2,
+        layer_types=["sliding_attention", "full_attention"],
+        is_causal=False,
+    )
+    assert _dflash_layer_causal(config, 0) is False
+    assert _dflash_layer_causal(config, 1) is False
+
+
 def _vllm_config(**draft_config):
     config = SimpleNamespace(**draft_config)
     return SimpleNamespace(

vllm/config/vllm.py

@@ -670,6 +670,12 @@ def use_v2_model_runner(self) -> bool:
         if self._dflash_needs_multi_kv_group():
             return True
 
+        # The DFlash2 candidate selector exists only in the V2 speculator. On V1
+        # the same checkpoint drafts through DFlashProposer, which never calls
+        # it, so the draft degrades to DFlash1 silently. Force V2 as for dspark.
+        if self._is_dflash2_draft():
+            return True
+
         if self.model_config is not None and self.model_config.is_diffusion:
             return True
 
@@ -693,6 +699,17 @@ def use_v2_model_runner(self) -> bool:
 
         return True
 
+    def _is_dflash2_draft(self) -> bool:
+        """Whether the DFlash draft is a DFlash2 one, by the architecture the
+        speculator selects on (v1/worker/gpu/spec_decode/__init__.py)."""
+        spec = self.speculative_config
+        if spec is None or spec.method != "dflash":
+            return False
+        draft_config = getattr(spec, "draft_model_config", None)
+        if draft_config is None:
+            return False
+        return "DFlash2DraftModel" in (draft_config.architectures or [])
+
     def _dflash_needs_multi_kv_group(self) -> bool:
         """Whether a DFlash draft mixes sliding-window and full attention."""
         spec = self.speculative_config

vllm/model_executor/layers/logits_processor.py

@@ -2,6 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """A layer that compute logits from hidden_stats."""
 
+from collections.abc import Callable
+from functools import cache
+
 import torch
 import torch.nn.functional as F
 
@@ -10,12 +13,43 @@
     tensor_model_parallel_all_gather,
     tensor_model_parallel_gather,
 )
+from vllm.logger import init_logger
 from vllm.model_executor.custom_op import PluggableLayer
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     UnquantizedEmbeddingMethod,
     VocabParallelEmbedding,
 )
 from vllm.platforms import current_platform
+from vllm.utils.flashinfer import has_flashinfer
+
+logger = init_logger(__name__)
+
+
+@cache
+def _flashinfer_topk() -> Callable[..., tuple[torch.Tensor, torch.Tensor]] | None:
+    """FlashInfer's radix top-k, or None for torch.topk.
+
+    The top-k spans the vocabulary, where the radix kernel is about twice
+    torch.topk.
+    """
+    if not current_platform.is_cuda():
+        return None
+    if not has_flashinfer():
+        logger.info_once(
+            "flashinfer is unavailable; vocab-parallel top-k uses torch.topk, "
+            "at roughly half the speed."
+        )
+        return None
+    from flashinfer import top_k
+
+    return top_k
+
+
+def _topk(scores: torch.Tensor, k: int) -> tuple[torch.Tensor, torch.Tensor]:
+    impl = _flashinfer_topk()
+    if impl is None or not scores.is_cuda:
+        return torch.topk(scores, k, dim=-1)
+    return impl(scores, k, sorted=True, deterministic=True)
 
 
 # --8<-- [start:logits_processor]
@@ -204,6 +238,53 @@ def get_top_tokens(
         top_tokens = gathered[:, :, 1].gather(dim=-1, index=max_rank_idx)
         return top_tokens.squeeze(-1).to(torch.int64)
 
+    def get_top_k_tokens(
+        self,
+        lm_head: VocabParallelEmbedding,
+        hidden_states: torch.Tensor,
+        k: int,
+        embedding_bias: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Vocab-parallel top-k without all-gathering full logits.
+
+        The `get_top_tokens` reduction widened from one token to k, returning
+        the values as well as the global ids. Communication is
+        O(batch * 2k * tp_size) rather than O(batch * vocab_size).
+
+        Scale and soft cap are applied to the k selected values rather than
+        the whole vocabulary; both are monotonic, so the selection is the same
+        and only k entries are touched.
+        """
+        if self.scale <= 0.0 and self.scale != 1.0:
+            raise ValueError(
+                "The local top-k reduction optimization is not supported for "
+                "non-positive logit scaling factors."
+            )
+
+        logits = self._apply_head(lm_head, hidden_states, embedding_bias)
+
+        # Mask out padding entries beyond org_vocab_size on this shard.
+        num_pad = lm_head.shard_indices.num_org_vocab_padding
+        if num_pad > 0:
+            logits[..., -num_pad:] = -float("inf")
+
+        values, ids = _topk(logits, k)
+        # Convert shard-local indices to global vocab indices.
+        ids = ids.to(torch.int64) + lm_head.shard_indices.org_vocab_start_index
+
+        if lm_head.tp_size > 1:
+            values = tensor_model_parallel_all_gather(values, dim=-1)
+            ids = tensor_model_parallel_all_gather(ids, dim=-1)
+            values, selected = _topk(values, k)
+            ids = ids.gather(-1, selected)
+
+        values = values.float()
+        if self.scale != 1.0:
+            values = values * self.scale
+        if self.soft_cap is not None:
+            values = torch.tanh(values / self.soft_cap) * self.soft_cap
+        return ids, values
+
     def extra_repr(self) -> str:
         s = f"vocab_size={self.vocab_size}"
         s += f", org_vocab_size={self.org_vocab_size}"