@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0bae1d3848
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[MRV2][Spec Decode] DFlash (#44586)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai> Signed-off-by: Benjamin Chislett <bchislett@nvidia.com> Signed-off-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: Giancarlo Delfin <gdelfin@inferact.ai>

Benjamin Chislett committed Jun 10, 2026 at 11:47 UTC 0bae1d38480374365ad77bbea50be225237572ea
14 files changed +868 -47
tests/v1/e2e/spec_decode/test_spec_decode.py
+12 -2
@@ -1300,13 +1300,18 @@ def dflash_config():
1300 )
1301
1302
1303 -def test_dflash_acceptance_rates(dflash_config):
1303 +@pytest.mark.parametrize("use_mrv2", [False, True])
1304 +def test_dflash_acceptance_rates(
1305 + monkeypatch: pytest.MonkeyPatch, use_mrv2: bool, dflash_config
1306 +):
1307 """
1308 E2E test for DFlash (block diffusion) speculative decoding.
1309 Runs acceptance rate validation on GSM8k, MT-Bench, and HumanEval
1310 comparing against baseline results from the paper (Table 1).
1311 See https://github.com/z-lab/dflash/blob/main/benchmark_sglang.py for methodology.
1312 """
1313 + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0")
1314 +
1315 spec_llm = LLM(**dflash_config)
1316
1317 max_prompts_per_dataset = 200 # mt-bench has 80, humaneval has 164, truncates gsm8k
@@ -1414,11 +1419,16 @@ def test_synthetic_acceptance_rate():
1419 cleanup_dist_env_and_memory()
1420
1421
1417 -def test_dflash_correctness(dflash_config):
1422 +@pytest.mark.parametrize("use_mrv2", [False, True])
1423 +def test_dflash_correctness(
1424 + monkeypatch: pytest.MonkeyPatch, use_mrv2: bool, dflash_config
1425 +):
1426 """
1427 E2E test for DFlash (block diffusion) speculative decoding.
1428 Ensures output correctness on GSM8k, with cudagraphs and batching on.
1429 """
1430 + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0")
1431 +
1432 spec_llm = LLM(**dflash_config)
1433
1434 # Evaluate GSM8k accuracy (Qwen3-8B ref: ~87-92% on GSM8k)
vllm/config/vllm.py
+8 -4
@@ -2008,12 +2008,16 @@ class VllmConfig:
2008 # TODO: ngram / ngram_gpu are not supported by the v2 model runner yet
2009 if speculative_config.method in ("ngram", "ngram_gpu"):
2010 unsupported.append("ngram/ngram_gpu speculative decoding")
2011 - elif speculative_config.method not in ("eagle", "eagle3", "mtp"):
2011 + elif speculative_config.method not in ("eagle", "eagle3", "mtp", "dflash"):
2012 unsupported.append(f"speculative method '{speculative_config.method}'")
2013
2014 - # V2 EagleSpeculator does not support parallel_drafting (required by PEagle)
2015 - if speculative_config.parallel_drafting:
2016 - unsupported.append("parallel drafting for speculative decoding")
2014 + # V2 EagleSpeculator does not support parallel_drafting (for P-Eagle)
2015 + # DFlash uses parallel drafting natively in V2 via DFlashSpeculator.
2016 + if (
2017 + speculative_config.parallel_drafting
2018 + and speculative_config.method != "dflash"
2019 + ):
2020 + unsupported.append("parallel drafting for EAGLE speculative decoding")
2021
2022 if (
2023 speculative_config.method == "eagle3"
vllm/v1/worker/gpu/attn_utils.py
+2 -1
@@ -394,6 +394,7 @@ def build_attn_metadata(
394 positions: torch.Tensor | None = None,
395 model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
396 for_cudagraph_capture: bool = False,
397 + causal: bool = True,
398 ) -> dict[str, Any]:
399 seq_lens = seq_lens[:num_reqs]
400 if dcp_local_seq_lens is not None:
@@ -423,7 +424,7 @@ def build_attn_metadata(
424 max_query_len=max_query_len,
425 block_table_tensor=block_table,
426 slot_mapping=slot_mapping,
426 - causal=True,
427 + causal=causal,
428 dcp_local_seq_lens=dcp_local_seq_lens,
429 positions=positions,
430 **common_attn_metadata_extra_kwargs,
vllm/v1/worker/gpu/model_runner.py
+6 -3
@@ -193,11 +193,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
193 if self.is_last_pp_rank:
194 self.speculator = init_speculator(self.vllm_config, self.device)
195
196 - if self.speculative_config.method == "eagle3":
197 - # EAGLE3 may require auxiliary hidden states from target model outputs.
196 + if self.speculative_config.method in ("eagle3", "dflash"):
197 + # Drafting may require auxiliary hidden states from target model outputs
198 self.use_aux_hidden_state_outputs = True
199 if self.use_pp:
200 - raise ValueError("EAGLE3 with pipeline parallel is not supported.")
200 + raise ValueError(
201 + f"{self.speculative_config.method} with pipeline parallel "
202 + "is not supported."
203 + )
204
205 # Draft tokens propagation - for spec-dec + struct outputs.
206 self.draft_tokens_handler = DraftTokensHandler(self.device)
vllm/v1/worker/gpu/sample/gumbel.py
+12 -1
@@ -91,6 +91,7 @@ def gumbel_block_argmax(
91 vocab_size,
92 APPLY_TEMPERATURE: tl.constexpr,
93 USE_FP64: tl.constexpr,
94 + PER_TOKEN_COL: tl.constexpr = False,
95 ):
96 req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
97 temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
@@ -103,7 +104,10 @@ def gumbel_block_argmax(
104 if processed_logits_ptr is not None:
105 # Store the temperature-applied logits.
106 if processed_logits_col_ptr is not None:
106 - col = tl.load(processed_logits_col_ptr)
107 + if PER_TOKEN_COL:
108 + col = tl.load(processed_logits_col_ptr + token_idx)
109 + else:
110 + col = tl.load(processed_logits_col_ptr)
111 else:
112 col = 0
113 tl.store(
@@ -158,6 +162,7 @@ def _gumbel_sample_kernel(
162 BLOCK_SIZE: tl.constexpr,
163 APPLY_TEMPERATURE: tl.constexpr,
164 USE_FP64: tl.constexpr,
165 + PER_TOKEN_COL: tl.constexpr,
166 ):
167 token_idx = tl.program_id(0)
168 block_idx = tl.program_id(1)
@@ -185,6 +190,7 @@ def _gumbel_sample_kernel(
190 vocab_size,
191 APPLY_TEMPERATURE=APPLY_TEMPERATURE,
192 USE_FP64=USE_FP64,
193 + PER_TOKEN_COL=PER_TOKEN_COL,
194 )
195 token_id = block_idx * BLOCK_SIZE + idx
196 tl.store(local_argmax_ptr + token_idx * local_argmax_stride + block_idx, token_id)
@@ -208,6 +214,10 @@ def gumbel_sample(
214 local_argmax = logits.new_empty(num_tokens, num_blocks, dtype=torch.int64)
215 local_max_dtype = torch.float64 if use_fp64 else torch.float32
216 local_max = logits.new_empty(num_tokens, num_blocks, dtype=local_max_dtype)
217 + per_token_col = (
218 + output_processed_logits_col is not None
219 + and output_processed_logits_col.dim() > 0
220 + )
221 _gumbel_sample_kernel[(num_tokens, num_blocks)](
222 local_argmax,
223 local_argmax.stride(0),
@@ -226,6 +236,7 @@ def gumbel_sample(
236 BLOCK_SIZE=BLOCK_SIZE,
237 APPLY_TEMPERATURE=apply_temperature,
238 USE_FP64=use_fp64,
239 + PER_TOKEN_COL=per_token_col,
240 )
241 # NOTE(woosuk): Use int64 for later indexing.
242 max_block_idx = local_max.argmax(dim=-1, keepdim=True)
vllm/v1/worker/gpu/spec_decode/__init__.py
+7 -1
@@ -8,7 +8,13 @@ from vllm.config import VllmConfig
8 def init_speculator(vllm_config: VllmConfig, device: torch.device):
9 speculative_config = vllm_config.speculative_config
10 assert speculative_config is not None
11 - if speculative_config.use_gemma4_mtp():
11 + if speculative_config.method == "dflash":
12 + from vllm.v1.worker.gpu.spec_decode.dflash.speculator import (
13 + DFlashSpeculator,
14 + )
15 +
16 + return DFlashSpeculator(vllm_config, device)
17 + elif speculative_config.use_gemma4_mtp():
18 from vllm.v1.worker.gpu.spec_decode.gemma4.speculator import (
19 Gemma4Speculator,
20 )
vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
-29
@@ -18,7 +18,6 @@ from vllm.v1.worker.gpu.cudagraph_utils import (
18 )
19 from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
20 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
21 -from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
21 from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import (
22 DecodeSpeculatorCudaGraphManager,
23 PrefillSpeculatorCudaGraphManager,
@@ -279,34 +278,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
278
279 return self.draft_tokens[:num_reqs]
280
282 - def sample_draft(
283 - self,
284 - hidden_states: torch.Tensor,
285 - positions: torch.Tensor,
286 - idx_mapping: torch.Tensor,
287 - temperature: torch.Tensor,
288 - seeds: torch.Tensor,
289 - draft_step: torch.Tensor,
290 - draft_logits: torch.Tensor | None,
291 - ) -> torch.Tensor:
292 - logits = self.model.compute_logits(hidden_states)
293 - if draft_logits is not None:
294 - # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
295 - # used for draft and target sampling.
296 - return gumbel_sample(
297 - logits,
298 - idx_mapping,
299 - temperature,
300 - seeds,
301 - positions + 1,
302 - apply_temperature=True,
303 - output_processed_logits=draft_logits,
304 - output_processed_logits_col=draft_step,
305 - use_fp64=self.use_fp64_gumbel,
306 - )
307 - else:
308 - return logits.argmax(dim=-1)
309 -
281 @torch.inference_mode()
282 def _run_model(
283 self,
vllm/v1/worker/gpu/spec_decode/dflash/__init__.py
vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py new
+114
@@ -0,0 +1,114 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Callable
4 +
5 +import torch
6 +
7 +from vllm.config.compilation import CUDAGraphMode
8 +from vllm.v1.kv_cache_interface import KVCacheConfig
9 +from vllm.v1.worker.gpu.attn_utils import (
10 + build_attn_metadata,
11 + build_slot_mappings_by_layer,
12 +)
13 +from vllm.v1.worker.gpu.block_table import BlockTables
14 +from vllm.v1.worker.gpu.cudagraph_utils import (
15 + AttentionState,
16 + BatchExecutionDescriptor,
17 + CudaGraphManager,
18 +)
19 +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
20 +from vllm.v1.worker.utils import AttentionGroup
21 +
22 +
23 +def _prepare_dflash_inputs_to_capture(
24 + num_reqs: int,
25 + num_tokens: int,
26 + input_buffers: InputBuffers,
27 + block_tables: BlockTables,
28 + attn_groups: list[list[AttentionGroup]],
29 + kv_cache_config: KVCacheConfig,
30 + max_model_len: int,
31 + skip_attn: bool,
32 + causal: bool,
33 +) -> AttentionState:
34 + input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers)
35 + input_block_tables = block_tables.get_dummy_block_tables(num_reqs)
36 + slot_mappings = block_tables.get_dummy_slot_mappings(num_tokens)
37 + slot_mappings_by_layer = build_slot_mappings_by_layer(
38 + slot_mappings, kv_cache_config
39 + )
40 +
41 + attn_metadata = None
42 + if not skip_attn:
43 + query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
44 + attn_metadata = build_attn_metadata(
45 + attn_groups=attn_groups,
46 + num_reqs=num_reqs,
47 + num_tokens=num_tokens,
48 + query_start_loc_gpu=input_batch.query_start_loc,
49 + query_start_loc_cpu=query_start_loc_cpu,
50 + max_query_len=num_tokens // num_reqs,
51 + seq_lens=input_batch.seq_lens,
52 + max_seq_len=max_model_len,
53 + block_tables=input_block_tables,
54 + slot_mappings=slot_mappings,
55 + kv_cache_config=kv_cache_config,
56 + for_cudagraph_capture=True,
57 + causal=causal,
58 + )
59 + return AttentionState(attn_metadata, slot_mappings_by_layer)
60 +
61 +
62 +class DFlashCudaGraphManager(CudaGraphManager):
63 + """DFlash CudaGraphManager for the parallel-drafting query forward,
64 + building its own attention metadata from scratch."""
65 +
66 + def __init__(self, *args, causal: bool = False, **kwargs) -> None:
67 + super().__init__(*args, **kwargs)
68 + self.causal = causal
69 +
70 + def capture(
71 + self,
72 + forward_fn: Callable,
73 + input_buffers: InputBuffers,
74 + block_tables: BlockTables,
75 + attn_groups: list[list[AttentionGroup]],
76 + kv_cache_config: KVCacheConfig,
77 + max_model_len: int,
78 + progress_bar_desc: str = "Capturing CUDA graphs",
79 + ) -> None:
80 + def create_forward_fn(
81 + desc: BatchExecutionDescriptor,
82 + warmup: bool,
83 + ) -> tuple[Callable[[CUDAGraphMode], None], AttentionState]:
84 + num_tokens = desc.num_tokens
85 + num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs)
86 + num_tokens_across_dp = (
87 + torch.full((self.dp_size,), num_tokens, dtype=torch.int32, device="cpu")
88 + if self.dp_size > 1
89 + else None
90 + )
91 + attn_state = _prepare_dflash_inputs_to_capture(
92 + num_reqs,
93 + num_tokens,
94 + input_buffers,
95 + block_tables,
96 + attn_groups,
97 + kv_cache_config,
98 + max_model_len,
99 + skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE),
100 + causal=self.causal,
101 + )
102 + attn_metadata, slot_mappings = attn_state
103 +
104 + fwd = lambda cg_mode: forward_fn(
105 + num_reqs,
106 + num_tokens,
107 + attn_metadata,
108 + slot_mappings,
109 + num_tokens_across_dp,
110 + cg_mode,
111 + )
112 + return fwd, attn_state
113 +
114 + super().capture(create_forward_fn, progress_bar_desc)
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py new
+573
@@ -0,0 +1,573 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any
4 +
5 +import torch
6 +import torch.nn as nn
7 +
8 +from vllm.config import VllmConfig
9 +from vllm.config.compilation import CUDAGraphMode
10 +from vllm.forward_context import BatchDescriptor, set_forward_context
11 +from vllm.logger import init_logger
12 +from vllm.triton_utils import tl, triton
13 +from vllm.v1.attention.backends.utils import PAD_SLOT_ID
14 +from vllm.v1.kv_cache_interface import KVCacheConfig
15 +from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
16 +from vllm.v1.worker.gpu.block_table import BlockTables
17 +from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
18 +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
19 +from vllm.v1.worker.gpu.model_states.interface import ModelState
20 +from vllm.v1.worker.gpu.spec_decode.dflash.cudagraph import DFlashCudaGraphManager
21 +from vllm.v1.worker.gpu.spec_decode.dflash.utils import (
22 + get_dflash_causal,
23 + load_dflash_model,
24 +)
25 +from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
26 +from vllm.v1.worker.gpu.spec_decode.utils import get_parallel_drafting_token_id
27 +
28 +logger = init_logger(__name__)
29 +
30 +
31 +class DFlashSpeculator(DraftModelSpeculator):
32 + def __init__(self, vllm_config: VllmConfig, device: torch.device):
33 + super().__init__(vllm_config, device)
34 +
35 + self.hidden_states = torch.zeros(
36 + self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=device
37 + )
38 +
39 + # Multimodal inputs not currently supported.
40 + self.supports_mm_inputs = False
41 +
42 + # Each request emits exactly (bonus + N mask) query tokens per step.
43 + self.num_query_per_req = 1 + self.num_speculative_steps
44 +
45 + self.parallel_drafting_token_id = get_parallel_drafting_token_id(
46 + self.draft_model_config.hf_config
47 + )
48 +
49 + self.dflash_causal = get_dflash_causal(self.draft_model_config)
50 +
51 + # Buffers for context K/V precomputation. Populated by prepare_dflash_inputs,
52 + # and processed by the model's precompute_and_store_context_kv method.
53 + # NOT captured by CUDA graphs.
54 + self.context_positions = torch.zeros(
55 + self.max_num_tokens, dtype=torch.int64, device=device
56 + )
57 + self.context_slot_mapping = torch.zeros(
58 + self.max_num_tokens, dtype=torch.int64, device=device
59 + )
60 +
61 + # Per-mask-token sampling buffers. Flattened from (num_reqs, num_spec_tokens).
62 + max_num_sampled_tokens = self.max_num_reqs * self.num_speculative_steps
63 + self.sample_indices = torch.zeros(
64 + max_num_sampled_tokens, dtype=torch.int64, device=device
65 + )
66 + self.sample_pos = torch.zeros(
67 + max_num_sampled_tokens, dtype=torch.int64, device=device
68 + )
69 + self.sample_idx_mapping = torch.zeros(
70 + max_num_sampled_tokens, dtype=torch.int32, device=device
71 + )
72 + # [0, 1, ..., N-1, 0, 1, ..., N-1, ...] -> the per-token column index into
73 + # draft_logits[req, step, :].
74 + self.sample_col = torch.arange(
75 + self.num_speculative_steps, dtype=torch.int32, device=device
76 + ).repeat(self.max_num_reqs)
77 +
78 + self.query_cudagraph_manager: DFlashCudaGraphManager | None = None
79 + self.draft_kv_cache_group_id: int = -1
80 +
81 + def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
82 + # PIECEWISE cudagraphs are not supported for dflash
83 + if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL:
84 + cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
85 + else:
86 + cudagraph_mode = CUDAGraphMode.NONE
87 +
88 + self.query_cudagraph_manager = DFlashCudaGraphManager(
89 + self.vllm_config,
90 + self.device,
91 + cudagraph_mode,
92 + decode_query_len=self.num_query_per_req,
93 + causal=self.dflash_causal,
94 + )
95 +
96 + def capture(self, attn_states: dict | None = None) -> None:
97 + logger.info("Capturing model for DFlash speculator...")
98 + # Reset sampling indices to zero to prevent stale values from prior
99 + # dummy runs from being baked into the captured graph.
100 + self.sample_indices.zero_()
101 + self.sample_pos.zero_()
102 + self.sample_idx_mapping.zero_()
103 + assert self.query_cudagraph_manager is not None
104 + self.query_cudagraph_manager.capture(
105 + self._generate_draft,
106 + self.input_buffers,
107 + self.block_tables,
108 + self.attn_groups,
109 + self.kv_cache_config,
110 + self.max_model_len,
111 + progress_bar_desc="Capturing dflash CUDA graphs",
112 + )
113 +
114 + def load_draft_model(
115 + self,
116 + target_model: nn.Module,
117 + target_attn_layer_names: set[str],
118 + ) -> nn.Module:
119 + return load_dflash_model(target_model, self.vllm_config)
120 +
121 + def set_attn(
122 + self,
123 + model_state: ModelState,
124 + kv_cache_config: KVCacheConfig,
125 + block_tables: BlockTables,
126 + ) -> None:
127 + super().set_attn(model_state, kv_cache_config, block_tables)
128 +
129 + # DFlash precomputes context K/V with a single block_size; mixing
130 + # kv-cache groups would silently corrupt the cache for the non-matching group.
131 + draft_groups = [gid for gid, g in enumerate(self.attn_groups) if g]
132 + assert len(draft_groups) == 1, (
133 + "DFlash currently requires all draft attention layers to share "
134 + "a single kv-cache group."
135 + )
136 + self.draft_kv_cache_group_id = draft_groups[0]
137 + self.draft_block_size = self.block_tables.block_sizes[
138 + self.draft_kv_cache_group_id
139 + ]
140 +
141 + @torch.inference_mode()
142 + def _run_model(
143 + self,
144 + num_tokens: int,
145 + attn_metadata: dict[str, Any] | None,
146 + slot_mappings: dict[str, torch.Tensor] | None,
147 + num_tokens_across_dp: torch.Tensor | None,
148 + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
149 + ) -> torch.Tensor:
150 + batch_descriptor = BatchDescriptor(num_tokens=num_tokens)
151 + with set_forward_context(
152 + attn_metadata,
153 + self.vllm_config,
154 + num_tokens=num_tokens,
155 + cudagraph_runtime_mode=cudagraph_runtime_mode,
156 + num_tokens_across_dp=num_tokens_across_dp,
157 + slot_mapping=slot_mappings,
158 + batch_descriptor=batch_descriptor,
159 + ):
160 + last_hidden_states = self.model(
161 + input_ids=self.input_buffers.input_ids[:num_tokens],
162 + positions=self.input_buffers.positions[:num_tokens],
163 + inputs_embeds=None,
164 + )
165 + return last_hidden_states
166 +
167 + def _generate_draft(
168 + self,
169 + num_reqs: int,
170 + num_tokens_padded: int,
171 + attn_metadata: dict[str, Any] | None,
172 + slot_mappings: dict[str, torch.Tensor] | None,
173 + num_tokens_across_dp: torch.Tensor | None,
174 + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
175 + ) -> None:
176 + last_hidden_states = self._run_model(
177 + num_tokens_padded,
178 + attn_metadata,
179 + slot_mappings,
180 + num_tokens_across_dp,
181 + cudagraph_runtime_mode,
182 + )
183 +
184 + num_sample = num_reqs * self.num_speculative_steps
185 + sample_hidden_states = last_hidden_states[self.sample_indices[:num_sample]]
186 + draft_tokens = self.sample_draft(
187 + sample_hidden_states,
188 + self.sample_pos[:num_sample],
189 + self.sample_idx_mapping[:num_sample],
190 + self.temperature,
191 + self.seeds,
192 + self.sample_col[:num_sample],
193 + self.draft_logits,
194 + )
195 + self.draft_tokens[:num_reqs] = draft_tokens.view(
196 + num_reqs, self.num_speculative_steps
197 + )
198 +
199 + def _build_draft_attn_metadata(
200 + self,
201 + num_reqs: int,
202 + num_reqs_padded: int,
203 + num_tokens_padded: int,
204 + num_query_per_req: int | None = None,
205 + causal: bool = False,
206 + ) -> dict[str, Any] | None:
207 + if not self.draft_attn_layer_names:
208 + return None
209 + assert num_query_per_req is None # Omitted for DFlash, read from self instead
210 + return super()._build_draft_attn_metadata(
211 + num_reqs,
212 + num_reqs_padded,
213 + num_tokens_padded,
214 + num_query_per_req=self.num_query_per_req,
215 + causal=causal,
216 + )
217 +
218 + @torch.inference_mode()
219 + def propose(
220 + self,
221 + input_batch: InputBatch,
222 + attn_metadata: dict[str, Any],
223 + slot_mappings: dict[str, torch.Tensor],
224 + # [num_tokens, hidden_size]
225 + last_hidden_states: torch.Tensor,
226 + # num_layers x [num_tokens, hidden_size]
227 + aux_hidden_states: list[torch.Tensor] | None,
228 + # [num_reqs]
229 + num_sampled: torch.Tensor,
230 + # [num_reqs]
231 + num_rejected: torch.Tensor,
232 + # [max_num_reqs]
233 + last_sampled: torch.Tensor,
234 + # [max_num_reqs]
235 + next_prefill_tokens: torch.Tensor,
236 + # [max_num_reqs]
237 + temperature: torch.Tensor,
238 + # [max_num_reqs]
239 + seeds: torch.Tensor,
240 + num_tokens_across_dp: torch.Tensor | None = None,
241 + dummy_run: bool = False,
242 + skip_attn_for_dummy_run: bool = False,
243 + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
244 + is_profile: bool = False,
245 + ) -> torch.Tensor:
246 + num_reqs = input_batch.num_reqs
247 + num_target_tokens = input_batch.num_tokens
248 + num_query_tokens = num_reqs * self.num_query_per_req
249 + max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item()
250 + self.draft_max_seq_len = min(
251 + max_seq_len + self.num_query_per_req, self.max_model_len
252 + )
253 +
254 + # NOTE: To avoid CPU-GPU synchronization without CPU knowing the
255 + # number of rejected tokens, we maintain the size of input_ids and
256 + # hidden_states the same as the target model's. This means, we pad each
257 + # request's query length to include any rejected positions.
258 + if aux_hidden_states:
259 + hidden_states = self.model.combine_hidden_states(
260 + torch.cat(aux_hidden_states, dim=-1)
261 + )
262 + else:
263 + hidden_states = last_hidden_states
264 + self.hidden_states[:num_target_tokens].copy_(hidden_states[:num_target_tokens])
265 +
266 + self._copy_request_inputs(
267 + num_reqs,
268 + input_batch.idx_mapping,
269 + temperature,
270 + seeds,
271 + )
272 +
273 + if dummy_run and skip_attn_for_dummy_run:
274 + # Memory profiling path: block_tables / kv_cache_config are not initialized.
275 + # Since DFlash needs to build its own attention metadata, we must skip the
276 + # preparation in this path and run a minimal forward pass.
277 + self.model.precompute_and_store_context_kv(
278 + self.hidden_states[:num_target_tokens],
279 + self.context_positions[:num_target_tokens],
280 + )
281 + self._generate_draft(
282 + num_reqs,
283 + num_query_tokens,
284 + attn_metadata=None,
285 + slot_mappings=None,
286 + num_tokens_across_dp=num_tokens_across_dp,
287 + cudagraph_runtime_mode=CUDAGraphMode.NONE,
288 + )
289 + return self.draft_tokens[:num_reqs]
290 +
291 + # The query slot mapping is written into the shared BlockTables slot_mappings.
292 + # That buffer's address is what the captured CUDA graph reads from at replay.
293 + assert self.draft_kv_cache_group_id >= 0
294 + query_slot_mapping = self.block_tables.slot_mappings[
295 + self.draft_kv_cache_group_id
296 + ]
297 + prepare_dflash_inputs(
298 + self.input_buffers,
299 + query_slot_mapping,
300 + self.context_positions,
301 + self.context_slot_mapping,
302 + self.sample_indices,
303 + self.sample_pos,
304 + self.sample_idx_mapping,
305 + input_batch,
306 + num_sampled,
307 + num_rejected,
308 + last_sampled,
309 + next_prefill_tokens,
310 + self.block_tables.input_block_tables[self.draft_kv_cache_group_id],
311 + self.draft_block_size,
312 + self.parallel_drafting_token_id,
313 + self.num_query_per_req,
314 + self.num_speculative_steps,
315 + self.max_num_reqs,
316 + self.max_num_tokens,
317 + )
318 +
319 + # Pre-insert context K/V into the cache. Runs eagerly outside the captured graph
320 + # because the context shape varies per step. During dummy runs the block tables
321 + # are placeholders, so we skip the cache write to avoid clobbering real entries.
322 + self.model.precompute_and_store_context_kv(
323 + self.hidden_states[:num_target_tokens],
324 + self.context_positions[:num_target_tokens],
325 + context_slot_mapping=(
326 + None if dummy_run else self.context_slot_mapping[:num_target_tokens]
327 + ),
328 + )
329 +
330 + # Every DFlash step has exactly num_query_per_req tokens, so we can use FULL CGs
331 + batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
332 + self.query_cudagraph_manager,
333 + num_reqs,
334 + num_query_tokens,
335 + uniform_token_count=self.num_query_per_req,
336 + dp_size=self.dp_size,
337 + dp_rank=self.dp_rank,
338 + need_eager=is_profile,
339 + )
340 +
341 + num_reqs_padded = batch_desc.num_reqs or num_reqs
342 + num_tokens_padded = batch_desc.num_tokens
343 +
344 + # Rebuild the draft attention metadata even when replaying the FULL
345 + # graph so that any attention metadata builder state is updated.
346 + draft_attn_metadata = self._build_draft_attn_metadata(
347 + num_reqs=num_reqs,
348 + num_reqs_padded=num_reqs_padded,
349 + num_tokens_padded=num_tokens_padded,
350 + causal=self.dflash_causal,
351 + )
352 + draft_slot_mappings_by_layer = build_slot_mappings_by_layer(
353 + self.block_tables.slot_mappings[:, :num_tokens_padded],
354 + self.kv_cache_config,
355 + )
356 +
357 + if batch_desc.cg_mode == CUDAGraphMode.FULL:
358 + assert self.query_cudagraph_manager is not None
359 + self.query_cudagraph_manager.run_fullgraph(batch_desc)
360 + else:
361 + self._generate_draft(
362 + num_reqs_padded,
363 + num_tokens_padded,
364 + draft_attn_metadata,
365 + draft_slot_mappings_by_layer,
366 + num_tokens_across_dp=num_tokens_across_dp,
367 + cudagraph_runtime_mode=batch_desc.cg_mode,
368 + )
369 +
370 + return self.draft_tokens[:num_reqs]
371 +
372 +
373 +@triton.jit
374 +def _prepare_dflash_inputs_kernel(
375 + # Outputs
376 + out_input_ids_ptr,
377 + out_query_positions_ptr,
378 + out_query_start_loc_ptr,
379 + out_seq_lens_ptr,
380 + out_query_slot_mapping_ptr,
381 + out_context_positions_ptr,
382 + out_context_slot_mapping_ptr,
383 + out_sample_indices_ptr,
384 + out_sample_pos_ptr,
385 + out_sample_idx_mapping_ptr,
386 + # Inputs from target batch
387 + target_positions_ptr,
388 + target_query_start_loc_ptr,
389 + idx_mapping_ptr,
390 + last_sampled_ptr,
391 + next_prefill_tokens_ptr,
392 + num_sampled_ptr,
393 + num_rejected_ptr,
394 + # Block table for slot mapping lookup.
395 + block_table_ptr,
396 + block_table_stride,
397 + # Scalars
398 + parallel_drafting_token_id,
399 + block_size,
400 + num_query_per_req,
401 + num_speculative_steps,
402 + max_num_reqs,
403 + max_num_tokens,
404 + PAD_SLOT_ID: tl.constexpr,
405 + BLOCK_SIZE: tl.constexpr,
406 +):
407 + req_idx = tl.program_id(0)
408 + block_idx = tl.program_id(1)
409 + num_reqs = tl.num_programs(0)
410 + req_state_idx = tl.load(idx_mapping_ptr + req_idx)
411 +
412 + ctx_start = tl.load(target_query_start_loc_ptr + req_idx)
413 + ctx_end = tl.load(target_query_start_loc_ptr + req_idx + 1)
414 + num_ctx = ctx_end - ctx_start
415 +
416 + num_rejected = tl.load(num_rejected_ptr + req_idx)
417 + valid_ctx_end = ctx_end - num_rejected
418 +
419 + num_sampled = tl.load(num_sampled_ptr + req_idx)
420 + if num_sampled > 0:
421 + bonus_token = tl.load(last_sampled_ptr + req_state_idx).to(tl.int32)
422 + else:
423 + # Chunked prefilling: splice in the next prefill token.
424 + bonus_token = tl.load(next_prefill_tokens_ptr + req_state_idx).to(tl.int32)
425 +
426 + last_valid_pos = tl.load(target_positions_ptr + valid_ctx_end - 1)
427 + query_base = req_idx * num_query_per_req
428 +
429 + j = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
430 + is_ctx = j < num_ctx
431 + is_query = (j >= num_ctx) & (j < num_ctx + num_query_per_req)
432 + query_off = j - num_ctx
433 +
434 + # --- Context positions / slots ---
435 + ctx_pos_idx = ctx_start + tl.where(is_ctx, j, 0)
436 + ctx_pos = tl.load(target_positions_ptr + ctx_pos_idx, mask=is_ctx, other=0)
437 + ctx_block_num = ctx_pos // block_size
438 + ctx_block_num = tl.minimum(ctx_block_num, block_table_stride - 1)
439 + ctx_block_id = tl.load(
440 + block_table_ptr + req_idx * block_table_stride + ctx_block_num,
441 + mask=is_ctx,
442 + other=0,
443 + ).to(tl.int64)
444 + ctx_slot = ctx_block_id * block_size + (ctx_pos % block_size)
445 + tl.store(out_context_positions_ptr + ctx_start + j, ctx_pos, mask=is_ctx)
446 + tl.store(out_context_slot_mapping_ptr + ctx_start + j, ctx_slot, mask=is_ctx)
447 +
448 + # --- Query positions / input_ids / slots ---
449 + query_pos = last_valid_pos + 1 + query_off
450 + query_idx = query_base + query_off
451 + is_bonus = is_query & (query_off == 0)
452 + input_id = tl.where(is_bonus, bonus_token, parallel_drafting_token_id)
453 +
454 + q_block_num = query_pos // block_size
455 + q_block_num = tl.minimum(q_block_num, block_table_stride - 1)
456 + q_block_id = tl.load(
457 + block_table_ptr + req_idx * block_table_stride + q_block_num,
458 + mask=is_query,
459 + other=0,
460 + ).to(tl.int64)
461 + q_slot = q_block_id * block_size + (query_pos % block_size)
462 +
463 + tl.store(out_input_ids_ptr + query_idx, input_id, mask=is_query)
464 + tl.store(out_query_positions_ptr + query_idx, query_pos, mask=is_query)
465 + tl.store(out_query_slot_mapping_ptr + query_idx, q_slot, mask=is_query)
466 +
467 + # --- Sample indices / positions / idx_mapping (mask tokens only) ---
468 + is_sample = is_query & (query_off > 0)
469 + sample_idx = req_idx * num_speculative_steps + (query_off - 1)
470 + tl.store(out_sample_indices_ptr + sample_idx, query_idx, mask=is_sample)
471 + tl.store(out_sample_pos_ptr + sample_idx, query_pos, mask=is_sample)
472 + tl.store(out_sample_idx_mapping_ptr + sample_idx, req_state_idx, mask=is_sample)
473 +
474 + if block_idx == 0:
475 + tl.store(out_query_start_loc_ptr + req_idx, query_base)
476 + # seq_lens is the absolute sequence length the draft attention
477 + # reads up to (context + query), not just the count of accepted
478 + # tokens this step.
479 + tl.store(out_seq_lens_ptr + req_idx, last_valid_pos + 1 + num_query_per_req)
480 + if req_idx == num_reqs - 1:
481 + # Pad per-request buffers to max_num_reqs for CUDA graph safety.
482 + last_query_end = num_reqs * num_query_per_req
483 + for i in range(num_reqs, max_num_reqs + 1, BLOCK_SIZE):
484 + block = i + tl.arange(0, BLOCK_SIZE)
485 + mask = block < max_num_reqs + 1
486 + tl.store(out_query_start_loc_ptr + block, last_query_end, mask=mask)
487 + for i in range(num_reqs, max_num_reqs, BLOCK_SIZE):
488 + block = i + tl.arange(0, BLOCK_SIZE)
489 + mask = block < max_num_reqs
490 + tl.store(out_seq_lens_ptr + block, 0, mask=mask)
491 + # Padded sample slots point at query index 0 (a valid row in
492 + # last_hidden_states) so CG replay never reads OOB.
493 + pad_start = num_reqs * num_speculative_steps
494 + pad_end = max_num_reqs * num_speculative_steps
495 + for i in range(pad_start, pad_end, BLOCK_SIZE):
496 + block = i + tl.arange(0, BLOCK_SIZE)
497 + mask = block < pad_end
498 + tl.store(out_sample_indices_ptr + block, 0, mask=mask)
499 + tl.store(out_sample_pos_ptr + block, 0, mask=mask)
500 + tl.store(out_sample_idx_mapping_ptr + block, 0, mask=mask)
501 + # Pad query slot mappings past num_query_tokens with PAD so the
502 + # captured CG sees PAD slots (no K/V write) for replay sizes
503 + # larger than the current request count.
504 + q_pad_start = num_reqs * num_query_per_req
505 + for i in range(q_pad_start, max_num_tokens, BLOCK_SIZE):
506 + block = i + tl.arange(0, BLOCK_SIZE)
507 + mask = block < max_num_tokens
508 + tl.store(out_query_slot_mapping_ptr + block, PAD_SLOT_ID, mask=mask)
509 +
510 +
511 +def prepare_dflash_inputs(
512 + input_buffers: InputBuffers,
513 + query_slot_mapping: torch.Tensor,
514 + context_positions: torch.Tensor,
515 + context_slot_mapping: torch.Tensor,
516 + sample_indices: torch.Tensor,
517 + sample_pos: torch.Tensor,
518 + sample_idx_mapping: torch.Tensor,
519 + input_batch: InputBatch,
520 + # [num_reqs]
521 + num_sampled: torch.Tensor,
522 + # [num_reqs]
523 + num_rejected: torch.Tensor,
524 + # [max_num_reqs]
525 + last_sampled: torch.Tensor,
526 + # [max_num_reqs]
527 + next_prefill_tokens: torch.Tensor,
528 + # [max_num_reqs, max_num_blocks]
529 + block_table: torch.Tensor,
530 + block_size: int,
531 + parallel_drafting_token_id: int,
532 + num_query_per_req: int,
533 + num_speculative_steps: int,
534 + max_num_reqs: int,
535 + max_num_tokens: int,
536 +) -> None:
537 + num_reqs = input_batch.num_reqs
538 + assert num_reqs > 0
539 + # Cover the longest possible per-request span (ctx + query). Use the max
540 + # per-request query length, not the total token count across the batch.
541 + max_target_query_len = int(input_batch.num_scheduled_tokens.max())
542 + max_tokens_per_req = max_target_query_len + num_query_per_req
543 + BLOCK_SIZE = min(256, triton.next_power_of_2(max(1, max_tokens_per_req)))
544 + num_blocks = triton.cdiv(max_tokens_per_req, BLOCK_SIZE)
545 + _prepare_dflash_inputs_kernel[(num_reqs, num_blocks)](
546 + input_buffers.input_ids,
547 + input_buffers.positions,
548 + input_buffers.query_start_loc,
549 + input_buffers.seq_lens,
550 + query_slot_mapping,
551 + context_positions,
552 + context_slot_mapping,
553 + sample_indices,
554 + sample_pos,
555 + sample_idx_mapping,
556 + input_batch.positions,
557 + input_batch.query_start_loc,
558 + input_batch.idx_mapping,
559 + last_sampled,
560 + next_prefill_tokens,
561 + num_sampled,
562 + num_rejected,
563 + block_table,
564 + block_table.stride(0),
565 + parallel_drafting_token_id,
566 + block_size,
567 + num_query_per_req,
568 + num_speculative_steps,
569 + max_num_reqs,
570 + max_num_tokens,
571 + PAD_SLOT_ID=PAD_SLOT_ID,
572 + BLOCK_SIZE=BLOCK_SIZE,
573 + )
vllm/v1/worker/gpu/spec_decode/dflash/utils.py new
+70
@@ -0,0 +1,70 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch.nn as nn
4 +
5 +from vllm.config import ModelConfig, VllmConfig, replace
6 +from vllm.distributed.parallel_state import get_pp_group
7 +from vllm.model_executor.model_loader import get_model
8 +from vllm.v1.worker.gpu.spec_decode.eagle.utils import _should_share
9 +
10 +
11 +def get_dflash_causal(draft_model_config: ModelConfig) -> bool:
12 + """Whether the DFlash draft uses causal (vs non-causal) attention."""
13 + dflash_config = getattr(draft_model_config.hf_config, "dflash_config", None) or {}
14 + return dflash_config.get("causal", False)
15 +
16 +
17 +def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module:
18 + from vllm.compilation.backends import set_model_tag
19 +
20 + speculative_config = vllm_config.speculative_config
21 + assert speculative_config is not None
22 + draft_model_config = speculative_config.draft_model_config
23 + # Modify the attention config so that we select an attention backend that matches
24 + # the causal/non-causal mode of the dflash model.
25 + causal = get_dflash_causal(draft_model_config)
26 + draft_vllm_config = replace(
27 + vllm_config,
28 + attention_config=replace(
29 + vllm_config.attention_config, use_non_causal=not causal
30 + ),
31 + )
32 + with set_model_tag("dflash_head"):
33 + dflash_model = get_model(
34 + vllm_config=draft_vllm_config, model_config=draft_model_config
35 + )
36 +
37 + target_language_model = (
38 + target_model.get_language_model()
39 + if hasattr(target_model, "get_language_model")
40 + else target_model
41 + )
42 + target_inner = target_language_model.model
43 + draft_inner = dflash_model.model
44 +
45 + # Skip embedding sharing under PP — each rank owns its own embedding.
46 + if get_pp_group().world_size == 1:
47 + target_embed = getattr(target_inner, "embed_tokens", None) or getattr(
48 + target_inner, "embedding", None
49 + )
50 + draft_embed = getattr(draft_inner, "embed_tokens", None)
51 + if target_embed is not None and _should_share(
52 + dflash_model, "has_own_embed_tokens", draft_embed, target_embed
53 + ):
54 + if draft_embed is not None:
55 + del draft_inner.embed_tokens
56 + draft_inner.embed_tokens = target_embed
57 +
58 + # Share lm_head with the target unless the draft remaps vocab via
59 + # draft_id_to_target_id (in which case its own lm_head is required).
60 + target_lm_head = getattr(target_model, "lm_head", None)
61 + draft_lm_head = getattr(dflash_model, "lm_head", None)
62 + if (
63 + target_lm_head is not None
64 + and draft_lm_head is not None
65 + and getattr(dflash_model, "draft_id_to_target_id", None) is None
66 + ):
67 + del dflash_model.lm_head
68 + dflash_model.lm_head = target_lm_head
69 +
70 + return dflash_model
vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py
+6 -3
@@ -38,9 +38,12 @@ def get_eagle3_aux_layers_from_config(
38 if not (spec_config and spec_config.draft_model_config):
39 return None
40 hf_config = spec_config.draft_model_config.hf_config
41 - if not hasattr(hf_config, "eagle_aux_hidden_state_layer_ids"):
42 - return None
43 - layer_ids = hf_config.eagle_aux_hidden_state_layer_ids
41 + layer_ids = getattr(hf_config, "eagle_aux_hidden_state_layer_ids", None)
42 + if not layer_ids:
43 + dflash_config = getattr(hf_config, "dflash_config", None)
44 + if dflash_config and isinstance(dflash_config, dict):
45 + # Add 1 to convert DFlash's aux layer id semantics
46 + layer_ids = [i + 1 for i in (dflash_config.get("target_layer_ids") or [])]
47 if layer_ids and isinstance(layer_ids, (list, tuple)):
48 return tuple(layer_ids)
49 return None
vllm/v1/worker/gpu/spec_decode/speculator.py
+39 -3
@@ -21,6 +21,7 @@ from vllm.v1.worker.gpu.cudagraph_utils import (
21 )
22 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
23 from vllm.v1.worker.gpu.model_states.interface import ModelState
24 +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
25
26
27 class BaseSpeculator(ABC):
@@ -178,9 +179,15 @@ class DraftModelSpeculator(BaseSpeculator):
179 num_reqs: int,
180 num_reqs_padded: int,
181 num_tokens_padded: int,
182 + num_query_per_req: int = 1,
183 + causal: bool = True,
184 ) -> dict[str, Any] | None:
182 - query_start_loc_cpu = torch.clamp(
183 - self.arange[: num_reqs_padded + 1], max=num_reqs
185 + # Uniform query: query_start_loc[i] = min(i, num_reqs) * num_query_per_req.
186 + # Clamp keeps the series non-decreasing past num_reqs, which some
187 + # attention backends require.
188 + query_start_loc_cpu = (
189 + torch.clamp(self.arange[: num_reqs_padded + 1], max=num_reqs)
190 + * num_query_per_req
191 )
192 block_tables = [
193 x[:num_reqs_padded] for x in self.block_tables.input_block_tables
@@ -194,15 +201,44 @@ class DraftModelSpeculator(BaseSpeculator):
201 : num_reqs_padded + 1
202 ],
203 query_start_loc_cpu=query_start_loc_cpu,
197 - max_query_len=1,
204 + max_query_len=num_query_per_req,
205 seq_lens=self.input_buffers.seq_lens[:num_reqs_padded],
206 max_seq_len=self.draft_max_seq_len,
207 block_tables=block_tables,
208 slot_mappings=slot_mappings,
209 kv_cache_config=self.kv_cache_config,
210 + causal=causal,
211 )
212 return attn_metadata
213
214 + def sample_draft(
215 + self,
216 + hidden_states: torch.Tensor,
217 + positions: torch.Tensor,
218 + idx_mapping: torch.Tensor,
219 + temperature: torch.Tensor,
220 + seeds: torch.Tensor,
221 + draft_step: torch.Tensor,
222 + draft_logits: torch.Tensor | None,
223 + ) -> torch.Tensor:
224 + logits = self.model.compute_logits(hidden_states)
225 + if draft_logits is not None:
226 + # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
227 + # used for draft and target sampling.
228 + return gumbel_sample(
229 + logits,
230 + idx_mapping,
231 + temperature,
232 + seeds,
233 + positions + 1,
234 + apply_temperature=True,
235 + output_processed_logits=draft_logits,
236 + output_processed_logits_col=draft_step,
237 + use_fp64=self.use_fp64_gumbel,
238 + )
239 + else:
240 + return logits.argmax(dim=-1)
241 +
242 def _copy_request_inputs(
243 self,
244 num_reqs: int,
vllm/v1/worker/gpu/spec_decode/utils.py
+19
@@ -45,3 +45,22 @@ class DraftTokensHandler:
45 # This case only happens when async scheduling is disabled.
46 draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids]
47 return DraftTokenIds(self.req_ids, draft_token_ids)
48 +
49 +
50 +def get_parallel_drafting_token_id(hf_config) -> int:
51 + """Resolve the mask token id used for parallel drafting slots.
52 +
53 + Checks (in order): `dflash_config.mask_token_id`, `pard_token`,
54 + `ptd_token_id`. Raises ValueError if none are present.
55 + """
56 + dflash_config = getattr(hf_config, "dflash_config", None) or {}
57 + if "mask_token_id" in dflash_config:
58 + return int(dflash_config["mask_token_id"])
59 + if hasattr(hf_config, "pard_token"):
60 + return int(hf_config.pard_token)
61 + if hasattr(hf_config, "ptd_token_id"):
62 + return int(hf_config.ptd_token_id)
63 + raise ValueError(
64 + "Model config must specify `dflash_config.mask_token_id`,"
65 + " `pard_token`, or `ptd_token_id` for parallel drafting."
66 + )