fix: Domino inference patches for speculators-format checkpoints

docs/features/speculative_decoding/README.md

@@ -11,6 +11,7 @@ vLLM supports a variety of methods of speculative decoding. Model-based methods
 - [EAGLE](eagle.md)
 - [Multi-Token Prediction (MTP)](mtp.md)
 - [Draft Model](draft_model.md)
+- [DFlash / Domino](dflash.md)
 - [Parallel Draft Model (PARD)](parallel_draft_model.md)
 - [Multi-Layer Perceptron](mlp.md)
 - [N-Gram](n_gram.md)
@@ -29,6 +30,7 @@ depend on your model family, traffic pattern, hardware, and sampling settings.
 | EAGLE | High gain | Medium to high gain | Strong general-purpose model-based method. |
 | MTP | High gain | Medium to high gain | Best when the target model has native MTP support. |
 | Draft model | High gain | Medium gain | Needs a separate draft model. |
+| DFlash / Domino | — | — | Parallel drafter with optional causal correction (Domino). Gains depend on model, hardware, and workload. |
 | Parallel Draft Model | High gain | Medium to high gain | Low draft model latency. |
 | MLP speculator | Medium to high gain | Medium gain | Good when compatible MLP speculators are available. |
 | N-gram | Low to medium gain | Medium gain | Lightweight and easy to enable. |
@@ -78,7 +80,7 @@ only apply to model-based methods such as `draft_model`, `mtp`, `eagle3`, and
 
 | Key | Type | Default | Allowed values / meaning |
 | --- | --- | --- | --- |
-| `method` | `string` | `None` | Speculation method. Common values include `draft_model`, `ngram`, `suffix`, `mtp`, `eagle3`, and `dflash`. If omitted, vLLM infers the method from the provided configuration when possible. |
+| `method` | `string` | `None` | Speculation method. Common values include `draft_model`, `ngram`, `suffix`, `mtp`, `eagle3`, and `dflash`. If omitted, vLLM infers the method from the provided configuration when possible. For Domino (causal correction head), use `dflash` with `projector_type="domino"` in `dflash_config`. |
 | `model` | `string` | `None` | Draft model, EAGLE head, or auxiliary model identifier. For `ngram`, `ngram_gpu`, `suffix`, and `mtp`, this can often be omitted. |
 | `num_speculative_tokens` | `integer > 0` | `None` | Number of speculative tokens to propose per step. Required for methods that do not infer it from model metadata. |
 | `draft_tensor_parallel_size` | `integer >= 1` | `None` | Tensor parallel size for the draft model. |
@@ -143,6 +145,31 @@ vllm serve <target-model> \
   }'
 ```
 
+#### DFlash / Domino
+
+DFlash is a parallel drafter that produces all K draft tokens in a single forward
+pass. Domino extends DFlash with a lightweight causal correction head (GRU +
+low-rank MLP) that refines the parallel base logits using causal state from
+previous draft tokens. Domino is enabled by setting `projector_type="domino"` in
+the checkpoint's `dflash_config`.
+
+| Key | Type | Default | Meaning |
+| --- | --- | --- | --- |
+| `mask_token_id` | `int` | Required | Token ID used for masked hidden states in the DFlash draft model. Must be set inside `dflash_config`. |
+
+Domino-specific `dflash_config` fields (set in the checkpoint):
+
+| Key | Type | Default | Meaning |
+| --- | --- | --- | --- |
+| `projector_type` | `string` | `"dflash"` | Set to `"domino"` to enable the Domino causal correction head. |
+| `gru_hidden_dim` | `int` | `1024` | GRU hidden dimension for the Domino correction head. |
+| `emb_dim` | `int` | `256` | Bottleneck dimension for the low-rank correction MLP. |
+| `pure_draft_prefix_len` | `int` | `1` | Number of prefix positions sampled from base logits without Domino correction. |
+| `shift_label` | `bool` | `true` | Whether to shift labels for Domino training. |
+
+Domino draft models can be trained using the
+[vllm-project/speculators](speculators.md) library.
+
 #### Cross-Vocabulary Draft Models (TLI)
 
   By default, vLLM requires the draft and target models to share the same

docs/features/speculative_decoding/dflash.md

@@ -0,0 +1,55 @@
+# DFlash / Domino
+
+DFlash is a parallel drafter for speculative decoding: it produces all K draft
+tokens in a single forward pass of the draft model, avoiding the sequential
+overhead of autoregressive drafters.
+
+[Domino](https://arxiv.org/abs/2605.29707) extends DFlash with a lightweight
+causal correction head (a GRU encoder + low-rank MLP) that refines the parallel
+base logits using causal state from previously drafted tokens. The correction
+operates in logit space, so no additional forward passes through the draft model
+or LM head are required.
+
+## Usage
+
+Domino is configured as a `projector_type` sub-mode of DFlash. Use a
+Domino-trained checkpoint with `method="dflash"`:
+
+```python
+from vllm import LLM, SamplingParams
+
+llm = LLM(
+    model="Qwen/Qwen3-8B",
+    speculative_config={
+        "method": "dflash",
+        "model": "your-username/Qwen3-8B-Domino-b16",
+        "num_speculative_tokens": 16,
+    },
+)
+```
+
+```bash
+vllm serve Qwen/Qwen3-8B \
+    --speculative-config '{
+        "method": "dflash",
+        "model": "your-username/Qwen3-8B-Domino-b16",
+        "num_speculative_tokens": 16
+    }'
+```
+
+When the checkpoint's `dflash_config.projector_type` is `"domino"`, vLLM
+automatically loads the Domino correction head weights and uses them during
+draft generation.
+
+## Training Domino draft models
+
+Domino draft models are trained using the
+[vllm-project/speculators](https://github.com/vllm-project/speculators) library.
+See the [speculators guide](speculators.md) for details.
+
+## Pre-trained models
+
+- See the [vllm-project/speculators](https://github.com/vllm-project/speculators)
+  repository for available Domino checkpoints.
+- Public DFlash checkpoints (without Domino head) are available on Hugging Face,
+  e.g. `z-lab/Qwen3-8B-DFlash-b16`.

tests/v1/spec_decode/test_domino_vocab.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for Domino pruned-vocab logit pipeline.
+
+Verifies that base logits + Domino correction happen in draft space
+before scattering to target space, so pruned-vocab Domino checkpoints
+work correctly (draft_vocab_size != target_vocab_size).
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+DRAFT_VOCAB = 32
+TARGET_VOCAB = 64
+HIDDEN_SIZE = 16
+GRU_HIDDEN = 8
+EMB_DIM = 12
+
+
+class _ScatterModel(nn.Module):
+    """Minimal nn.Module with scatter_logits_to_target from DFlashQwen3ForCausalLM."""
+
+    def __init__(self, draft_vocab, target_vocab):
+        super().__init__()
+        from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM
+
+        self.config = SimpleNamespace(
+            draft_vocab_size=draft_vocab,
+            vocab_size=target_vocab,
+        )
+        # d2t stores offsets: target_id = draft_id + d2t[draft_id].
+        # Identity mapping (offset=0) places draft tokens at positions 0..N-1.
+        d2t = torch.zeros(draft_vocab, dtype=torch.long)
+        if draft_vocab == target_vocab:
+            self.draft_id_to_target_id = None
+        else:
+            self.draft_id_to_target_id = nn.Parameter(d2t, requires_grad=False)
+
+        self.scatter_logits_to_target = (
+            DFlashQwen3ForCausalLM.scatter_logits_to_target.__get__(
+                self, type(self)
+            )
+        )
+
+
+def _make_scatter_model(draft_vocab=DRAFT_VOCAB, target_vocab=TARGET_VOCAB):
+    return _ScatterModel(draft_vocab, target_vocab)
+
+
+class TestScatterLogitsToTarget:
+    def test_shape_with_d2t_mapping(self):
+        model = _make_scatter_model()
+        logits = torch.randn(4, DRAFT_VOCAB)
+        result = model.scatter_logits_to_target(logits)
+        assert result.shape == (4, TARGET_VOCAB)
+
+    def test_values_scattered_correctly(self):
+        model = _make_scatter_model()
+        logits = torch.ones(1, DRAFT_VOCAB) * 42.0
+        result = model.scatter_logits_to_target(logits)
+        # Identity offset (d2t=0): draft[i] → target[i], so first DRAFT_VOCAB
+        # positions get 42.0, the rest get -inf.
+        assert result[0, :DRAFT_VOCAB].eq(42.0).all()
+        assert result[0, DRAFT_VOCAB:].eq(float("-inf")).all()
+
+    def test_noop_when_no_mapping(self):
+        model = _make_scatter_model()
+        model.draft_id_to_target_id = None
+        logits = torch.randn(4, DRAFT_VOCAB)
+        result = model.scatter_logits_to_target(logits)
+        assert torch.equal(result, logits)
+
+    def test_noncontiguous_mapping(self):
+        """d2t offsets that scatter draft tokens to non-contiguous positions."""
+        model = _make_scatter_model(draft_vocab=3, target_vocab=8)
+        model.draft_id_to_target_id = nn.Parameter(
+            torch.tensor([2, 0, 1], dtype=torch.long), requires_grad=False
+        )
+
+        logits = torch.tensor([[10.0, 20.0, 30.0]])
+        result = model.scatter_logits_to_target(logits)
+        assert result.shape == (1, 8)
+        assert result[0, 2] == 10.0  # draft[0] → target[0+2]
+        assert result[0, 1] == 20.0  # draft[1] → target[1+0]
+        assert result[0, 3] == 30.0  # draft[2] → target[2+1]
+
+
+class TestDominoHeadDraftSpace:
+    @pytest.fixture
+    def domino_head(self):
+        """Standalone DominoHead using plain nn.Linear (no vLLM parallelism)."""
+        head = nn.Module()
+        head.gru_hidden_dim = GRU_HIDDEN
+        head.emb_dim = EMB_DIM
+        head.prefix_gru = nn.GRU(
+            input_size=HIDDEN_SIZE,
+            hidden_size=GRU_HIDDEN,
+            num_layers=1,
+            batch_first=True,
+            bias=False,
+        )
+        head.embed_proj = nn.Sequential(
+            nn.Linear(HIDDEN_SIZE + GRU_HIDDEN, EMB_DIM, bias=False),
+            nn.SiLU(),
+            nn.Linear(EMB_DIM, DRAFT_VOCAB, bias=False),
+        )
+
+        from vllm.model_executor.models.qwen3_dflash import DominoHead
+
+        head.compute_logits = DominoHead.compute_logits.__get__(head, type(head))
+        return head
+
+    def test_correction_same_space_as_base(self, domino_head):
+        batch = 2
+        hidden = torch.randn(batch, HIDDEN_SIZE)
+        gru_hidden = torch.randn(1, batch, GRU_HIDDEN)
+        base_logits = torch.randn(batch, DRAFT_VOCAB)
+
+        result = domino_head.compute_logits(hidden, gru_hidden, base_logits)
+        assert result.shape == (batch, DRAFT_VOCAB)
+
+    def test_correction_adds_to_base(self, domino_head):
+        batch = 1
+        hidden = torch.randn(batch, HIDDEN_SIZE)
+        gru_hidden = torch.zeros(1, batch, GRU_HIDDEN)
+        base_logits = torch.zeros(batch, DRAFT_VOCAB)
+
+        result = domino_head.compute_logits(hidden, gru_hidden, base_logits)
+        assert not result.eq(0.0).all(), "correction should modify base logits"
+
+
+class TestDominoPrunedVocabPipeline:
+    def test_full_pipeline(self):
+        """Draft logits → Domino correction → scatter: end-to-end shape check."""
+        model = _make_scatter_model()
+
+        draft_logits = torch.randn(4, DRAFT_VOCAB)
+        correction = torch.randn(4, DRAFT_VOCAB)
+        corrected = draft_logits + correction
+        final = model.scatter_logits_to_target(corrected)
+
+        assert final.shape == (4, TARGET_VOCAB)
+
+    def test_argmax_returns_target_space_ids(self):
+        """After scatter, argmax should return valid target-space token IDs."""
+        model = _make_scatter_model()
+
+        draft_logits = torch.randn(4, DRAFT_VOCAB)
+        final = model.scatter_logits_to_target(draft_logits)
+        token_ids = final.argmax(dim=-1)
+
+        assert token_ids.shape == (4,)
+        assert (token_ids < TARGET_VOCAB).all()
+        assert (token_ids >= 0).all()
+
+    @pytest.mark.parametrize("draft_eq_target", [True, False])
+    def test_same_tokens_with_and_without_scatter(self, draft_eq_target):
+        """When draft==target vocab, scatter is a noop; tokens should match."""
+        if draft_eq_target:
+            model = _make_scatter_model(draft_vocab=32, target_vocab=32)
+        else:
+            model = _make_scatter_model(draft_vocab=32, target_vocab=64)
+            model.draft_id_to_target_id = nn.Parameter(
+                torch.zeros(32, dtype=torch.long), requires_grad=False
+            )
+
+        logits = torch.randn(4, 32)
+        scattered = model.scatter_logits_to_target(logits)
+        tokens = scattered.argmax(dim=-1)
+        direct_tokens = logits.argmax(dim=-1)
+
+        assert torch.equal(tokens, direct_tokens)