diff --git a/docs/features/speculative_decoding/README.md b/docs/features/speculative_decoding/README.md index ceb25dbfd02f..b000cdd4199d 100644 --- a/docs/features/speculative_decoding/README.md +++ b/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 \ }' ``` +#### 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 diff --git a/docs/features/speculative_decoding/dflash.md b/docs/features/speculative_decoding/dflash.md new file mode 100644 index 000000000000..3e1927aea335 --- /dev/null +++ b/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`. diff --git a/tests/v1/spec_decode/test_domino_vocab.py b/tests/v1/spec_decode/test_domino_vocab.py new file mode 100644 index 000000000000..3658a5795f68 --- /dev/null +++ b/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) diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 90b2f48faa82..3cbb020c9837 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -661,6 +661,86 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ) +class DominoHead(nn.Module): + def __init__( + self, *, config, vllm_config: VllmConfig, quant_config, prefix: str = "" + ): + super().__init__() + dflash_config = getattr(config, "dflash_config", {}) or {} + + self.gru_hidden_dim = int(dflash_config["gru_hidden_dim"]) + self.emb_dim = int(dflash_config["emb_dim"]) + + self.prefix_gru = nn.GRU( + input_size=config.hidden_size, + hidden_size=self.gru_hidden_dim, + num_layers=1, + batch_first=True, + bias=False, + ) + + self.embed_proj = nn.Sequential( + ReplicatedLinear( + input_size=config.hidden_size + self.gru_hidden_dim, + output_size=self.emb_dim, + bias=False, + params_dtype=vllm_config.model_config.dtype, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "embed_proj.0"), + return_bias=False, + ), + nn.SiLU(), + ReplicatedLinear( + input_size=self.emb_dim, + output_size=getattr(config, "draft_vocab_size", config.vocab_size), + bias=False, + params_dtype=vllm_config.model_config.dtype, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "embed_proj.2"), + return_bias=False, + ), + ) + + def init_state(self, embed_input_ids, prefix_token_ids: torch.Tensor): + prefix_embeds = embed_input_ids(prefix_token_ids) + _, gru_hidden = self.prefix_gru(prefix_embeds) + return gru_hidden + + def advance_state( + self, + embed_input_ids, + token_ids: torch.Tensor, + gru_hidden: torch.Tensor, + ): + if token_ids.dim() == 1: + token_ids = token_ids.unsqueeze(-1) + token_embeds = embed_input_ids(token_ids) + _, gru_hidden = self.prefix_gru(token_embeds, gru_hidden) + return gru_hidden + + def compute_logits( + self, + parallel_hidden: torch.Tensor, + gru_hidden: torch.Tensor, + base_logits: torch.Tensor, + ): + squeeze_time_dim = parallel_hidden.dim() == 2 + if squeeze_time_dim: + parallel_hidden = parallel_hidden.unsqueeze(1) + + if base_logits.dim() == 2: + base_logits = base_logits.unsqueeze(1) + + state = gru_hidden.transpose(0, 1) + correction_input = torch.cat([parallel_hidden, state], dim=-1) + correction_bias = self.embed_proj(correction_input) + + logits = base_logits + correction_bias + if squeeze_time_dim: + logits = logits.squeeze(1) + return logits + + class DFlashQwen3ForCausalLM(Qwen3ForCausalLM): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): nn.Module.__init__(self) @@ -695,6 +775,20 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): else: self.draft_id_to_target_id = None + dflash_config = getattr(self.config, "dflash_config", {}) or {} + self.projector_type = dflash_config.get("projector_type") + self.is_domino = self.projector_type == "domino" + + if self.is_domino: + self.domino_head = DominoHead( + config=self.config, + vllm_config=vllm_config, + quant_config=self.model.quant_config, + prefix=maybe_prefix(prefix, "domino_head"), + ) + else: + self.domino_head = None + def embed_input_ids( self, input_ids: torch.Tensor, @@ -719,14 +813,18 @@ def get_draft_attn_causal(self) -> list[bool]: get_draft_kv_cache_layer_names.""" return [layer.self_attn.causal for layer in self.model.layers] - def compute_logits( + def compute_draft_logits( self, hidden_states: torch.Tensor, ) -> torch.Tensor | None: - logits = self.logits_processor(self.lm_head, hidden_states) + return self.logits_processor(self.lm_head, hidden_states) + + def scatter_logits_to_target( + self, + logits: torch.Tensor, + ) -> torch.Tensor: if self.draft_id_to_target_id is None: return logits - base = torch.arange(self.config.draft_vocab_size, device=logits.device) targets = base + self.draft_id_to_target_id logits_new = logits.new_full( @@ -736,6 +834,15 @@ def compute_logits( logits_new[:, targets] = logits return logits_new + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + logits = self.compute_draft_logits(hidden_states) + if logits is None: + return None + return self.scatter_logits_to_target(logits) + def precompute_and_store_context_kv( self, context_states: torch.Tensor, @@ -773,12 +880,27 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): model_weights = {} includes_draft_id_mapping = False includes_embed_tokens = False + + # Domino-only weights should be loaded into self.domino_head, not into the + # DFlash backbone. + domino_weights = [] + for name, loaded_weight in weights: assert "mask_hidden" not in name, ( "DFlash embeds masked slots via mask_token_id (optionally " "overridden by a mask_embedding.pt file); it should not ship a " "mask_hidden weight." ) + + if getattr(self, "is_domino", False) and ( + name.startswith("prefix_gru.") or name.startswith("embed_proj.") + ): + domino_weights.append((f"domino_head.{name}", loaded_weight)) + continue + if getattr(self, "is_domino", False) and name.startswith("domino_head."): + domino_weights.append((name, loaded_weight)) + continue + if "t2d" in name: continue if "d2t" in name: @@ -813,6 +935,25 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): skip_substrs=skip_substrs, ) loader.load_weights(model_weights.items()) + + if getattr(self, "is_domino", False): + loader = AutoWeightsLoader(self) + loader.load_weights(domino_weights) + + domino_param_names = { + name + for name, _ in self.named_parameters() + if name.startswith("domino_head.") + } + loaded_domino_names = {name for name, _ in domino_weights} + missing_domino_names = domino_param_names - loaded_domino_names + + if missing_domino_names: + raise RuntimeError( + "Domino weight loading is incomplete. Missing: " + f"{sorted(missing_domino_names)}" + ) + self.model._build_fused_kv_buffers() def _read_mask_embedding(self) -> torch.Tensor | None: @@ -853,3 +994,35 @@ def _read_mask_embedding(self) -> torch.Tensor | None: MASK_EMBEDDING_FILENAME, ) return state.reshape(-1) + + def init_domino_state(self, prefix_token_ids: torch.Tensor) -> torch.Tensor: + if self.domino_head is None: + raise RuntimeError("Domino head is not enabled.") + return self.domino_head.init_state(self.embed_input_ids, prefix_token_ids) + + def advance_domino_state( + self, + token_ids: torch.Tensor, + gru_hidden: torch.Tensor, + ) -> torch.Tensor: + if self.domino_head is None: + raise RuntimeError("Domino head is not enabled.") + return self.domino_head.advance_state( + self.embed_input_ids, + token_ids, + gru_hidden, + ) + + def compute_domino_logits( + self, + parallel_hidden: torch.Tensor, + gru_hidden: torch.Tensor, + base_logits: torch.Tensor, + ) -> torch.Tensor: + if self.domino_head is None: + raise RuntimeError("Domino head is not enabled.") + return self.domino_head.compute_logits( + parallel_hidden, + gru_hidden, + base_logits, + ) diff --git a/vllm/transformers_utils/configs/speculators/algos.py b/vllm/transformers_utils/configs/speculators/algos.py index fbd2c4518196..fb25876bbf7b 100644 --- a/vllm/transformers_utils/configs/speculators/algos.py +++ b/vllm/transformers_utils/configs/speculators/algos.py @@ -119,16 +119,25 @@ def update_dflash(config_dict: dict, pre_trained_config: dict) -> None: pre_trained_config["eagle_aux_hidden_state_layer_ids"] = aux_layer_ids # DFlash configs use different indexing for the target layers, see #40727 - pre_trained_config["dflash_config"] = { + dflash_config = { "mask_token_id": config_dict["mask_token_id"], "target_layer_ids": [i - 1 for i in aux_layer_ids], "sample_from_anchor": config_dict.get("sample_from_anchor", False), } # Enable causal masking in SWA for vllm-project/speculators models - pre_trained_config["dflash_config"]["causal"] = not config_dict.get( + dflash_config["causal"] = not config_dict.get( "sliding_window_non_causal", True ) + # Domino projector fields + for key in ("projector_type", "shift_label", "pure_draft_prefix_len", + "gru_hidden_dim", "emb_dim"): + if key in config_dict: + pre_trained_config[key] = config_dict[key] + dflash_config[key] = config_dict[key] + + pre_trained_config["dflash_config"] = dflash_config + @register_speculator("dspark") def update_dspark(config_dict: dict, pre_trained_config: dict) -> None: diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 645469e611f4..9e946115b26e 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Mapping from typing import Any import torch @@ -74,6 +73,9 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.context_positions = torch.zeros( self.max_num_tokens, dtype=torch.int64, device=device ) + self.context_slot_mapping = torch.zeros( + self.max_num_tokens, dtype=torch.int64, device=device + ) # Per-mask-token sampling buffers. Flattened from (num_reqs, num_spec_tokens). max_num_sampled_tokens = self.max_num_reqs * self.num_speculative_steps @@ -147,7 +149,7 @@ def capture(self) -> None: self.attn_groups, self.kv_cache_config, self.max_model_len, - causal=self._group_causal, + causal=self.dflash_causal, progress_bar_desc=f"Capturing {self._speculator_name.lower()} CUDA graphs", ) @@ -174,18 +176,12 @@ def set_attn( target_attn_groups, ) - self.draft_kv_cache_group_ids = [ - gid for gid, g in enumerate(self.attn_groups) if g - ] - assert self.draft_kv_cache_group_ids, "No draft attention groups found." - self.draft_kv_cache_group_id = self.draft_kv_cache_group_ids[0] - - # Per-group context slot buffers for the precompute (one row per group). - self._context_slot_mappings = torch.zeros( - len(self.draft_kv_cache_group_ids), - self.max_num_tokens, - dtype=torch.int64, - device=self.device, + # DFlash precomputes context K/V with a single block_size; mixing + # kv-cache groups would silently corrupt the cache for the non-matching group. + draft_groups = [gid for gid, g in enumerate(self.attn_groups) if g] + assert len(draft_groups) == 1, ( + "DFlash currently requires all draft attention layers to share " + "a single kv-cache group." ) # Map each draft decoder layer to the index (within draft_kv_cache_group_ids) @@ -239,6 +235,109 @@ def _run_model( ) return last_hidden_states + def _sample_domino_logits( + self, + logits: torch.Tensor, + positions: torch.Tensor, + idx_mapping: torch.Tensor, + draft_step: torch.Tensor, + ) -> torch.Tensor: + return logits.argmax(dim=-1) + + def _generate_domino_draft( + self, + last_hidden_states: torch.Tensor, + num_reqs: int, + ) -> None: + dflash_config = getattr(self.model.config, "dflash_config", {}) or {} + shift_label = bool(dflash_config.get("shift_label", False)) + prefix_len = min( + int(dflash_config.get("pure_draft_prefix_len", 0)), + self.num_speculative_steps, + ) + + num_sample = num_reqs * self.num_speculative_steps + sample_hidden_states = last_hidden_states[ + self.sample_indices[:num_sample] + ].view(num_reqs, self.num_speculative_steps, -1) + sample_pos = self.sample_pos[:num_sample].view( + num_reqs, self.num_speculative_steps + ) + sample_idx_mapping = self.sample_idx_mapping[:num_sample].view( + num_reqs, self.num_speculative_steps + ) + sample_col = self.sample_col[:num_sample].view( + num_reqs, self.num_speculative_steps + ) + + anchor_indices = ( + torch.arange(num_reqs, device=self.device, dtype=torch.long) + * self.num_query_per_req + ) + anchor_token_ids = self.input_buffers.input_ids[anchor_indices].unsqueeze(-1) + + if shift_label: + sample_hidden_states = torch.cat( + [last_hidden_states[anchor_indices].unsqueeze(1), sample_hidden_states], + dim=1, + )[:, : self.num_speculative_steps, :] + sample_pos = torch.cat( + [self.input_buffers.positions[anchor_indices].unsqueeze(1), sample_pos], + dim=1, + )[:, : self.num_speculative_steps] + sample_idx_mapping = torch.cat( + [sample_idx_mapping[:, :1], sample_idx_mapping], + dim=1, + )[:, : self.num_speculative_steps] + + draft_logits = self.model.compute_draft_logits( + sample_hidden_states.reshape(num_reqs * self.num_speculative_steps, -1) + ) + if draft_logits is None: + raise RuntimeError( + "Domino speculative decoding requires full draft logits on this rank." + ) + draft_logits = draft_logits.view(num_reqs, self.num_speculative_steps, -1) + + if prefix_len > 0: + prefix_logits = draft_logits[:, :prefix_len, :] + prefix_logits_flat = self.model.scatter_logits_to_target( + prefix_logits.reshape(-1, prefix_logits.shape[-1]) + ) + prefix_token_ids = self._sample_domino_logits( + prefix_logits_flat, + sample_pos[:, :prefix_len].reshape(-1), + sample_idx_mapping[:, :prefix_len].reshape(-1), + sample_col[:, :prefix_len].reshape(-1), + ).view(num_reqs, prefix_len) + self.draft_tokens[:num_reqs, :prefix_len] = prefix_token_ids + realized_prefix_ids = torch.cat([anchor_token_ids, prefix_token_ids], dim=1) + else: + realized_prefix_ids = anchor_token_ids + + gru_hidden = self.model.init_domino_state(realized_prefix_ids) + correction_start = prefix_len + + for step in range(prefix_len, self.num_speculative_steps): + base_step_logits = draft_logits[:, step, :] + if step < correction_start: + logits = base_step_logits + else: + hidden_step = sample_hidden_states[:, step, :] + logits = self.model.compute_domino_logits( + hidden_step, gru_hidden, base_step_logits + ) + logits = self.model.scatter_logits_to_target(logits) + token = self._sample_domino_logits( + logits, + sample_pos[:, step], + sample_idx_mapping[:, step], + sample_col[:, step], + ) + self.draft_tokens[:num_reqs, step] = token + if step + 1 < self.num_speculative_steps: + gru_hidden = self.model.advance_domino_state(token, gru_hidden) + def _generate_draft( self, num_reqs: int, @@ -256,13 +355,16 @@ def _generate_draft( cudagraph_runtime_mode, ) + projector_type = getattr(self.model, "projector_type", None) + if projector_type == "domino": + self._generate_domino_draft(last_hidden_states, num_reqs) + return + num_sample = num_reqs * self.num_speculative_steps sample_hidden_states = last_hidden_states[self.sample_indices[:num_sample]] - # sample_pos is the predicted token's position Q; verification keys - # Gumbel by the predecessor (Q-1). sample_draft adds +1, so pass Q-2. draft_tokens = self.sample_draft( sample_hidden_states, - self.sample_pos[:num_sample] - 2, + self.sample_pos[:num_sample], self.sample_idx_mapping[:num_sample], self.temperature, self.seeds, @@ -281,7 +383,7 @@ def _build_draft_attn_metadata( seq_lens_cpu_upper_bound: torch.Tensor, step: int, num_query_per_req: int | None = None, - causal: bool | Mapping[int, bool] = False, + causal: bool = False, ) -> dict[str, Any] | None: if not self.draft_attn_layer_names: return None @@ -359,9 +461,6 @@ def propose( self.hidden_states[:num_target_tokens], self.context_positions[:num_target_tokens], ) - # DFlash processes all speculative tokens in one forward pass, - # so the real token count is num_query_tokens. - self._prepare_eplb_forward(num_query_tokens) self._generate_draft( num_reqs, num_query_tokens, @@ -375,49 +474,40 @@ def propose( # The query slot mapping is written into the shared BlockTables slot_mappings. # That buffer's address is what the captured CUDA graph reads from at replay. assert self.draft_kv_cache_group_id >= 0 - # Support multiple draft KV cache groups by preparing inputs once for each - for i, gid in enumerate(self.draft_kv_cache_group_ids): - prepare_dflash_inputs( - self.input_buffers, - self.block_tables.slot_mappings[gid], - self.context_positions, - self._context_slot_mappings[i], - self.sample_indices, - self.sample_pos, - self.sample_idx_mapping, - input_batch, - num_sampled, - num_rejected, - last_sampled, - next_prefill_tokens, - self.block_tables.input_block_tables[gid], - self.block_tables.kernel_block_sizes[gid], - self.parallel_drafting_token_id, - self.num_query_per_req, - self.num_speculative_steps, - self.max_num_reqs, - self.max_num_tokens, - self.max_model_len, - self.sample_from_anchor, - ) + query_slot_mapping = self.block_tables.slot_mappings[ + self.draft_kv_cache_group_id + ] + prepare_dflash_inputs( + self.input_buffers, + query_slot_mapping, + self.context_positions, + self.context_slot_mapping, + self.sample_indices, + self.sample_pos, + self.sample_idx_mapping, + input_batch, + num_sampled, + num_rejected, + last_sampled, + next_prefill_tokens, + self.block_tables.input_block_tables[self.draft_kv_cache_group_id], + self.draft_block_size, + self.parallel_drafting_token_id, + self.num_query_per_req, + self.num_speculative_steps, + self.max_num_reqs, + self.max_num_tokens, + ) # Pre-insert context K/V into the cache. Runs eagerly outside the captured graph # because the context shape varies per step. During dummy runs the block tables # are placeholders, so we skip the cache write to avoid clobbering real entries. - # Each layer uses the context slots of its own kv-cache group. - if dummy_run: - context_slots: torch.Tensor | list[torch.Tensor | None] | None = None - elif self._layer_group_idx is not None: - context_slots = [ - self._context_slot_mappings[gidx][:num_target_tokens] - for gidx in self._layer_group_idx - ] - else: - context_slots = self._context_slot_mappings[0][:num_target_tokens] self.model.precompute_and_store_context_kv( self.hidden_states[:num_target_tokens], self.context_positions[:num_target_tokens], - context_slots, + context_slot_mapping=( + None if dummy_run else self.context_slot_mapping[:num_target_tokens] + ), ) # Every DFlash step has exactly num_query_per_req tokens, so we can use FULL CGs @@ -449,16 +539,12 @@ def propose( self.kv_cache_config, ) - # DFlash processes all speculative tokens in one forward pass, - # so the real token count is num_query_tokens. - self._prepare_eplb_forward(num_query_tokens) - if batch_desc.cg_mode == CUDAGraphMode.FULL: assert self.query_cudagraph_manager is not None self.query_cudagraph_manager.run_fullgraph(batch_desc) else: self._generate_draft( - num_reqs, + num_reqs_padded, num_tokens_padded, draft_attn_metadata, draft_slot_mappings_by_layer, @@ -500,8 +586,6 @@ def _prepare_dflash_inputs_kernel( num_speculative_steps, max_num_reqs, max_num_tokens, - max_model_len, - SAMPLE_FROM_ANCHOR: tl.constexpr, PAD_SLOT_ID: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): @@ -562,21 +646,14 @@ def _prepare_dflash_inputs_kernel( q_slot = q_block_id * block_size + (query_pos % block_size) tl.store(out_input_ids_ptr + query_idx, input_id, mask=is_query) - clamped_query_pos = tl.minimum(query_pos, max_model_len - 1) - tl.store(out_query_positions_ptr + query_idx, clamped_query_pos, mask=is_query) + tl.store(out_query_positions_ptr + query_idx, query_pos, mask=is_query) tl.store(out_query_slot_mapping_ptr + query_idx, q_slot, mask=is_query) - # --- Sample indices / positions / idx_mapping --- - # When SAMPLE_FROM_ANCHOR (DSpark), so we sample at EVERY query position - # and each position k predicts the NEXT token (sampled position = query_pos + 1). - # Otherwise (DFlash default) the anchor is the bonus token and only the mask tokens - # at offsets > 0 are sampled from, each AT its own position. - sample_off = 0 if SAMPLE_FROM_ANCHOR else 1 - is_sample = is_query & (query_off >= sample_off) - sample_idx = req_idx * num_speculative_steps + (query_off - sample_off) - sample_pos = query_pos + 1 if SAMPLE_FROM_ANCHOR else query_pos + # --- Sample indices / positions / idx_mapping (mask tokens only) --- + is_sample = is_query & (query_off > 0) + sample_idx = req_idx * num_speculative_steps + (query_off - 1) tl.store(out_sample_indices_ptr + sample_idx, query_idx, mask=is_sample) - tl.store(out_sample_pos_ptr + sample_idx, sample_pos, mask=is_sample) + tl.store(out_sample_pos_ptr + sample_idx, query_pos, mask=is_sample) tl.store(out_sample_idx_mapping_ptr + sample_idx, req_state_idx, mask=is_sample) if block_idx == 0: @@ -597,9 +674,7 @@ def _prepare_dflash_inputs_kernel( mask = block < max_num_reqs tl.store(out_seq_lens_ptr + block, 0, mask=mask) # Padded sample slots point at query index 0 (a valid row in - # last_hidden_states) so CG replay never reads OOB. Padded - # sample idx mappings point to -1, which is ignored during - # sampling to prevent writing stale values to draft logits. + # last_hidden_states) so CG replay never reads OOB. pad_start = num_reqs * num_speculative_steps pad_end = max_num_reqs * num_speculative_steps for i in range(pad_start, pad_end, BLOCK_SIZE): @@ -607,7 +682,7 @@ def _prepare_dflash_inputs_kernel( mask = block < pad_end tl.store(out_sample_indices_ptr + block, 0, mask=mask) tl.store(out_sample_pos_ptr + block, 0, mask=mask) - tl.store(out_sample_idx_mapping_ptr + block, -1, mask=mask) + tl.store(out_sample_idx_mapping_ptr + block, 0, mask=mask) # Pad query slot mappings past num_query_tokens with PAD so the # captured CG sees PAD slots (no K/V write) for replay sizes # larger than the current request count. @@ -643,8 +718,6 @@ def prepare_dflash_inputs( num_speculative_steps: int, max_num_reqs: int, max_num_tokens: int, - max_model_len: int, - sample_from_anchor: bool = False, ) -> None: num_reqs = input_batch.num_reqs assert num_reqs > 0 @@ -680,8 +753,6 @@ def prepare_dflash_inputs( num_speculative_steps, max_num_reqs, max_num_tokens, - max_model_len, - SAMPLE_FROM_ANCHOR=sample_from_anchor, PAD_SLOT_ID=PAD_SLOT_ID, BLOCK_SIZE=BLOCK_SIZE, )