[Bugfix][Spec Decode] DFlash2: accept unquantized linear LM heads in the candidate selector

tests/models/registry.py

@@ -1486,6 +1486,13 @@ 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",
+        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

@@ -127,6 +127,31 @@ def test_rocm_defaults_deepseek_v4_to_mrv1(monkeypatch):
         default_v2_model_runner_architectures.cache_clear()
 
 
+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,91 @@
+# 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 test_selector_always_keeps_proposal_logits(monkeypatch):
+    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 = None
+
+    monkeypatch.setattr(DFlashSpeculator, "__init__", init_base)
+    speculator = DFlash2Speculator(None, torch.device("cpu"))
+
+    assert speculator.draft_logits.shape == (2, 4, 17)
+    assert speculator.draft_logits.dtype == torch.float32
+    assert torch.isneginf(speculator.draft_logits).all()

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

@@ -636,6 +636,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
 
@@ -659,6 +665,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/models/qwen3_dflash.py

@@ -56,10 +56,13 @@
 
 
 def _dflash_layer_causal(config: Qwen3Config, layer_idx: int) -> bool:
-    """``dflash_config.causal`` overrides all layers; else only SWA layers causal."""
+    """Resolve explicit causality before falling back to legacy layer defaults."""
+    is_causal = getattr(config, "is_causal", None)
+    if is_causal is not None:
+        return bool(is_causal)
     override = (getattr(config, "dflash_config", None) or {}).get("causal")
     if override is not None:
-        return override
+        return bool(override)
     layer_types = getattr(config, "layer_types", None)
     return bool(layer_types) and layer_types[layer_idx] == _SLIDING_ATTENTION
 
@@ -374,6 +377,8 @@ def forward(
 
 @support_torch_compile
 class DFlashQwen3Model(nn.Module):
+    decoder_layer_cls = DFlashQwen3DecoderLayer
+
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_substr={
             "midlayer.": "layers.0.",
@@ -434,7 +439,7 @@ def __init__(
 
         self.layers = nn.ModuleList(
             [
-                DFlashQwen3DecoderLayer(
+                self.decoder_layer_cls(
                     current_vllm_config,
                     config=self.config,
                     layer_idx=layer_idx,
@@ -699,6 +704,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
 
 
 class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):
+    model_cls = DFlashQwen3Model
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         nn.Module.__init__(self)
         self.draft_model_config = vllm_config.speculative_config.draft_model_config
@@ -708,7 +715,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         target_layer_num = vllm_config.model_config.get_num_layers(
             vllm_config.parallel_config
         )
-        self.model = DFlashQwen3Model(
+        self.model = self.model_cls(
             vllm_config=vllm_config,
             prefix=maybe_prefix(prefix, "model"),
             start_layer_id=target_layer_num,