@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 494636b29d

Fork 0

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

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 19:45 UTC

[Feat][Spec Decode] DFlash (#36847)

Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>

Benjamin Chislett committed Mar 30, 2026 at 15:03 UTC 494636b29d3b3a7b35020e4becb6c6995e200f9d

17 files changed +1577 -107

tests/models/registry.py

@@ -1163,6 +1163,14 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {

1163 # "JackFram/llama-160m",

1164 # speculative_model="ibm-ai-platform/llama-160m-accelerator"

1165 # ),

1166 + # [DFlash]

1167 + "DFlashDraftModel": _HfExamplesInfo(

1168 + "Qwen/Qwen3.5-4B",

1169 + speculative_model="z-lab/Qwen3.5-4B-DFlash",

1170 + use_original_num_layers=True, # Need all layers since DFlash has >1 layer,

1171 + max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env

1172 + max_num_seqs=32,

1173 + ),

1174 # [Eagle]

1175 "EagleDeepSeekMTPModel": _HfExamplesInfo(

1176 "eagle618/deepseek-v3-random",

tests/v1/e2e/spec_decode/test_spec_decode.py

+166 -6

@@ -7,6 +7,7 @@ from typing import Any

8 import pytest

9 import torch

10 +from tqdm import tqdm

12 from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts, evaluate_gsm8k_offline

13 from tests.utils import (

@@ -1105,19 +1106,178 @@ def some_high_acceptance_metrics() -> dict:

1106 }

1107

1108

1108 -def compute_acceptance_rate(metrics: list[Metric]) -> float:

1109 +def compute_acceptance_rate(

1110 + metrics: list[Metric], prev_metrics: list[Metric] | None = None

1111 +) -> float:

1112 name2metric = {metric.name: metric for metric in metrics}

1110 - n_draft_toks = name2metric["vllm:spec_decode_num_draft_tokens"].value # type: ignore

1113 + n_draft_toks = name2metric["vllm:spec_decode_num_draft_tokens"].value

1114 if n_draft_toks == 0:

1115 return float("nan")

1113 - n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value # type: ignore

1116 + n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value

1117 + if prev_metrics is not None:

1118 + prev_name2metric = {metric.name: metric for metric in prev_metrics}

1119 + n_draft_toks -= prev_name2metric["vllm:spec_decode_num_draft_tokens"].value

1120 + n_accepted_toks -= prev_name2metric[

1121 + "vllm:spec_decode_num_accepted_tokens"

1122 + ].value

1123 + if n_draft_toks <= 0:

1124 + return float("nan")

1125 return n_accepted_toks / n_draft_toks

1126

1127

1117 -def compute_acceptance_len(metrics: list[Metric]) -> float:

1128 +def compute_acceptance_len(

1129 + metrics: list[Metric], prev_metrics: list[Metric] | None = None

1130 +) -> float:

1131 name2metric = {metric.name: metric for metric in metrics}

1119 - n_drafts = name2metric["vllm:spec_decode_num_drafts"].value # type: ignore

1120 - n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value # type: ignore

1132 + n_drafts = name2metric["vllm:spec_decode_num_drafts"].value

1133 + n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value

1134 if n_drafts == 0:

1135 return 1

1136 + if prev_metrics is not None:

1137 + prev_name2metric = {metric.name: metric for metric in prev_metrics}

1138 + n_drafts -= prev_name2metric["vllm:spec_decode_num_drafts"].value

1139 + n_accepted_toks -= prev_name2metric[

1140 + "vllm:spec_decode_num_accepted_tokens"

1141 + ].value

1142 + if n_drafts <= 0:

1143 + return 1

1144 return 1 + (n_accepted_toks / n_drafts)

1145 +

1146 +

1147 +# Datasets in the format used in DFlash validations

1148 +def load_and_process_dataset(data_name: str):

1149 + from datasets import load_dataset

1150 +

1151 + if data_name == "gsm8k":

1152 + dataset = load_dataset("openai/gsm8k", "main", split="test")

1153 + prompt_fmt = (

1154 + "{question}\nPlease reason step by step,"

1155 + " and put your final answer within \\boxed{{}}."

1156 + )

1157 + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})

1158 + elif data_name == "mt-bench":

1159 + dataset = load_dataset("HuggingFaceH4/mt_bench_prompts", split="train")

1160 + dataset = dataset.map(lambda x: {"turns": x["prompt"]})

1161 + elif data_name == "humaneval":

1162 + dataset = load_dataset("openai/openai_humaneval", split="test")

1163 + prompt_fmt = (

1164 + "Write a solution to the following problem and make sure"

1165 + " that it passes the tests:\n```python\n{prompt}\n```"

1166 + )

1167 + dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})

1168 +

1169 + return dataset

1170 +

1171 +

1172 +@pytest.fixture

1173 +def dflash_config():

1174 + target_model = "Qwen/Qwen3-8B"

1175 + draft_model = "z-lab/Qwen3-8B-DFlash-b16"

1176 +

1177 + return dict(

1178 + model=target_model,

1179 + trust_remote_code=True,

1180 + speculative_config={

1181 + "method": "dflash",

1182 + "model": draft_model,

1183 + "num_speculative_tokens": 16,

1184 + "max_model_len": 32768,

1185 + },

1186 + max_model_len=32768,

1187 + max_num_seqs=128,

1188 + gpu_memory_utilization=0.85,

1189 + enforce_eager=False,

1190 + disable_log_stats=False,

1191 + )

1192 +

1193 +

1194 +def test_dflash_acceptance_rates(dflash_config):

1195 + """

1196 + E2E test for DFlash (block diffusion) speculative decoding.

1197 + Runs acceptance rate validation on GSM8k, MT-Bench, and HumanEval

1198 + comparing against baseline results from the paper (Table 1).

1199 + See https://github.com/z-lab/dflash/blob/main/benchmark_sglang.py for methodology.

1200 + """

1201 + spec_llm = LLM(**dflash_config)

1202 +

1203 + max_prompts_per_dataset = 200 # mt-bench has 80, humaneval has 164, truncates gsm8k

1204 +

1205 + # All scores from Table 1 in https://arxiv.org/pdf/2602.06036

1206 + expected_acceptance_lengths = {

1207 + "mt-bench": 4.24,

1208 + "humaneval": 6.50,

1209 + "gsm8k": 6.54 * 0.95, # runs with a subset of prompts so extra wide tol here

1210 + }

1211 +

1212 + tokenizer = spec_llm.get_tokenizer()

1213 + for dataset_name, expected_len in expected_acceptance_lengths.items():

1214 + dataset = load_and_process_dataset(dataset_name)

1215 + prev_metrics = None

1216 + acceptance_lengths = []

1217 + for i in tqdm(

1218 + range(min(max_prompts_per_dataset, len(dataset))),

1219 + desc=f"Processing {dataset_name}",

1220 + ):

1221 + user_content = dataset[i]["turns"][0]

1222 + prompt_text = tokenizer.apply_chat_template(

1223 + [{"role": "user", "content": user_content}],

1224 + tokenize=False,

1225 + add_generation_prompt=True,

1226 + enable_thinking=False,

1227 + )

1228 +

1229 + # Temp=0, MaxTokens=2048 from the paper

1230 + spec_llm.generate(

1231 + [prompt_text],

1232 + SamplingParams(temperature=0, max_tokens=2048),

1233 + use_tqdm=False,

1234 + )

1235 + current_metrics = spec_llm.get_metrics()

1236 + acceptance_len = compute_acceptance_len(current_metrics, prev_metrics)

1237 + prev_metrics = current_metrics

1238 + acceptance_lengths.append(acceptance_len)

1239 +

1240 + mean_acceptance_length = sum(acceptance_lengths) / len(acceptance_lengths)

1241 + expected_len = expected_len * 0.9

1242 + print(

1243 + f"DFlash acceptance_len for {dataset_name}: {mean_acceptance_length:.2f}"

1244 + f" (expected at least {expected_len:.2f})"

1245 + )

1246 +

1247 + assert mean_acceptance_length >= expected_len, (

1248 + f"DFlash acceptance_len for {dataset_name} is below expected threshold:"

1249 + f"{mean_acceptance_length:.2f} < {expected_len:.2f}"

1250 + )

1251 +

1252 + del spec_llm

1253 + torch.accelerator.empty_cache()

1254 + cleanup_dist_env_and_memory()

1255 +

1256 +

1257 +def test_dflash_correctness(dflash_config):

1258 + """

1259 + E2E test for DFlash (block diffusion) speculative decoding.

1260 + Ensures output correctness on GSM8k, with cudagraphs and batching on.

1261 + """

1262 + spec_llm = LLM(**dflash_config)

1263 +

1264 + # Evaluate GSM8k accuracy (Qwen3-8B ref: ~87-92% on GSM8k)

1265 + evaluate_llm_for_gsm8k(spec_llm, expected_accuracy_threshold=0.8)

1266 +

1267 + current_metrics = spec_llm.get_metrics()

1268 + acceptance_len = compute_acceptance_len(current_metrics)

1269 +

1270 + # AR is thoroughly validated in test_dflash_acceptance_rates, in a manner consistent

1271 + # with the DFlash paper. However, that test measures AL per-request and thus runs

1272 + # with a batch size of 1. To ensure that AL does not collapse with large batch sizes

1273 + # we enforce a baseline on the AL over the full lm-eval-style GSM8k test.

1274 + expected_len = 3.5 # Measured is 3.9 to 4.0

1275 + print(f"DFlash GSM8k correctness test got AL {acceptance_len}")

1276 + assert acceptance_len >= expected_len, (

1277 + "DFlash correctness check failed with"

1278 + f" {acceptance_len=}, expected at least {expected_len}"

1279 + )

1280 +

1281 + del spec_llm

1282 + torch.accelerator.empty_cache()

1283 + cleanup_dist_env_and_memory()

tests/v1/spec_decode/test_eagle.py

+152 -3

@@ -27,6 +27,7 @@ from vllm.config.load import LoadConfig

27 from vllm.model_executor.models.llama import LlamaForCausalLM

28 from vllm.platforms import current_platform

29 from vllm.v1.attention.backends.registry import AttentionBackendEnum

30 +from vllm.v1.spec_decode.dflash import DFlashProposer

31 from vllm.v1.spec_decode.draft_model import DraftModelProposer

32 from vllm.v1.spec_decode.eagle import EagleProposer

33 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata

@@ -36,6 +37,8 @@ model_dir = "meta-llama/Llama-3.1-8B-Instruct"

37 eagle_dir = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"

38 eagle3_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"

39 ar_draft_model_dir = "amd/PARD-Llama-3.2-1B" # Compatible with parallel and AR drafting

40 +dflash_target_dir = "Qwen/Qwen3-8B"

41 +dflash_dir = "z-lab/Qwen3-8B-DFlash-b16"

43 BLOCK_SIZE = 16

@@ -47,18 +50,29 @@ def _create_proposer(

50 speculative_token_tree: list[tuple[int, ...]] | None = None,

51 parallel_drafting: bool = False,

52 ) -> EagleProposer:

50 - model_config = ModelConfig(model=model_dir, runner="generate", max_model_len=100)

51 -

53 # Method-dependent setup

54 if method == "eagle":

55 + target_model_dir = model_dir

56 draft_model_dir = eagle_dir

57 elif method == "eagle3":

58 + target_model_dir = model_dir

59 draft_model_dir = eagle3_dir

60 elif method == "draft_model":

61 + target_model_dir = model_dir

62 draft_model_dir = ar_draft_model_dir

63 + elif method == "dflash":

64 + target_model_dir = dflash_target_dir

65 + draft_model_dir = dflash_dir

66 else:

67 raise ValueError(f"Unknown method: {method}")

69 + model_config = ModelConfig(

70 + model=target_model_dir,

71 + runner="generate",

72 + max_model_len=100,

73 + trust_remote_code=(method == "dflash"),

74 + )

75 +

76 spec_token_tree_str = None

77 if speculative_token_tree is not None:

78 assert num_speculative_tokens == len(speculative_token_tree)

@@ -92,7 +106,9 @@ def _create_proposer(

106 attention_config=AttentionConfig(backend=attention_backend),

107 )

108

95 - if "eagle" in method:

109 + if method == "dflash":

110 + proposer = DFlashProposer(vllm_config=vllm_config, device=device)

111 + elif "eagle" in method:

112 proposer = EagleProposer(vllm_config=vllm_config, device=device)

113 else:

114 proposer = DraftModelProposer(vllm_config=vllm_config, device=device)

@@ -1152,3 +1168,136 @@ def test_propose_tree(spec_token_tree):

1168

1169 # Verify that the draft tokens match our expectations.

1170 assert torch.equal(result, expected_tokens)

1171 +

1172 +

1173 +def test_set_inputs_first_pass_dflash():

1174 + """

1175 + Test for DFlash set_inputs_first_pass.

1176 +

1177 + DFlash uses cross-attention: context tokens become K/V and only

1178 + query tokens (bonus + mask) are Q. This tests the DFlash-specific

1179 + input preparation where:

1180 + - Context hidden states are stored by reference (no copy)

1181 + - Query input_ids are [next_token, mask, mask, ...] per request

1182 + - Context and query positions are written to separate buffers

1183 + - token_indices_to_sample points to mask token positions only

1184 + - A new CommonAttentionMetadata is returned with causal=False

1185 +

1186 + Setup:

1187 + - 3 requests with query_lens [3, 2, 4]

1188 + - num_speculative_tokens = 3

1189 + - num_query_per_req = 4 (1 bonus + 3 mask tokens)

1190 + - next_token_ids: [100, 200, 300]

1191 +

1192 + Expected output layout (query tokens only, 12 total):

1193 + Request 0 (indices 0-3): [100, mask, mask, mask]

1194 + Request 1 (indices 4-7): [200, mask, mask, mask]

1195 + Request 2 (indices 8-11): [300, mask, mask, mask]

1196 +

1197 + Expected positions layout (separate buffers):

1198 + Context (_context_positions_buffer, 9 tokens): copied from target_positions

1199 + Query (positions, 12 tokens):

1200 + Request 0: last_pos=9, query=[10, 11, 12, 13]

1201 + Request 1: last_pos=7, query=[8, 9, 10, 11]

1202 + Request 2: last_pos=11, query=[12, 13, 14, 15]

1203 + """

1204 + device = torch.device(current_platform.device_type)

1205 +

1206 + num_speculative_tokens = 3

1207 + proposer = _create_proposer("dflash", num_speculative_tokens)

1208 + mask_token_id = proposer.parallel_drafting_token_id

1209 +

1210 + # Setup batch with 3 requests

1211 + batch_spec = BatchSpec(

1212 + seq_lens=[10, 8, 12],

1213 + query_lens=[3, 2, 4],

1214 + )

1215 +

1216 + common_attn_metadata = create_common_attn_metadata(

1217 + batch_spec,

1218 + block_size=BLOCK_SIZE,

1219 + device=device,

1220 + arange_block_indices=True,

1221 + )

1222 +

1223 + # Input tensors

1224 + # Request 0: tokens [10, 11, 12] at positions [7, 8, 9]

1225 + # Request 1: tokens [20, 21] at positions [6, 7]

1226 + # Request 2: tokens [30, 31, 32, 33] at positions [8, 9, 10, 11]

1227 + target_token_ids = torch.tensor(

1228 + [10, 11, 12, 20, 21, 30, 31, 32, 33], dtype=torch.int32, device=device

1229 + )

1230 + target_positions = torch.tensor(

1231 + [7, 8, 9, 6, 7, 8, 9, 10, 11], dtype=torch.int64, device=device

1232 + )

1233 + target_hidden_states = torch.randn(

1234 + 9, proposer.hidden_size, dtype=proposer.dtype, device=device

1235 + )

1236 + next_token_ids = torch.tensor([100, 200, 300], dtype=torch.int32, device=device)

1237 +

1238 + num_tokens, token_indices_to_sample, output_cad = proposer.set_inputs_first_pass(

1239 + target_token_ids=target_token_ids,

1240 + next_token_ids=next_token_ids,

1241 + target_positions=target_positions,

1242 + target_hidden_states=target_hidden_states,

1243 + token_indices_to_sample=None,

1244 + cad=common_attn_metadata,

1245 + num_rejected_tokens_gpu=None,

1246 + )

1247 +

1248 + num_query_per_req = 1 + num_speculative_tokens # 4

1249 + num_context = 9

1250 +

1251 + # num_tokens is the query-only count

1252 + assert num_tokens == 3 * num_query_per_req # 12

1253 +

1254 + # Verify input_ids (query tokens only)

1255 + # Each request: [next_token, mask, mask, mask]

1256 + M = mask_token_id

1257 + expected_input_ids = torch.tensor(

1258 + [100, M, M, M, 200, M, M, M, 300, M, M, M],

1259 + dtype=torch.int32,

1260 + device=device,

1261 + )

1262 + assert torch.equal(proposer.input_ids[:num_tokens], expected_input_ids)

1263 +

1264 + # Verify context positions (separate buffer): copied from target_positions

1265 + assert torch.equal(

1266 + proposer._context_positions_buffer[:num_context], target_positions

1267 + )

1268 +

1269 + # Verify query positions (separate buffer, starts at index 0):

1270 + # req0: last_pos=9, query=[10, 11, 12, 13]

1271 + # req1: last_pos=7, query=[8, 9, 10, 11]

1272 + # req2: last_pos=11, query=[12, 13, 14, 15]

1273 + expected_query_positions = torch.tensor(

1274 + [10, 11, 12, 13, 8, 9, 10, 11, 12, 13, 14, 15],

1275 + dtype=torch.int64,

1276 + device=device,

1277 + )

1278 + assert torch.equal(

1279 + proposer.positions[:num_tokens],

1280 + expected_query_positions,

1281 + )

1282 +

1283 + # Verify token_indices_to_sample (mask tokens only, skip bonus at offset 0)

1284 + # req0: query indices 0-3, mask at 1,2,3

1285 + # req1: query indices 4-7, mask at 5,6,7

1286 + # req2: query indices 8-11, mask at 9,10,11

1287 + expected_token_indices_to_sample = torch.tensor(

1288 + [1, 2, 3, 5, 6, 7, 9, 10, 11], dtype=torch.int32, device=device

1289 + )

1290 + assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample)

1291 +

1292 + # Verify the new CAD has DFlash-specific properties

1293 + assert output_cad.causal is False # DFlash requires non-causal attention

1294 + assert output_cad.num_actual_tokens == num_tokens # query-only count

1295 + assert output_cad.max_query_len == num_query_per_req

1296 +

1297 + expected_query_start_loc = torch.tensor(

1298 + [0, 4, 8, 12], dtype=torch.int32, device=device

1299 + )

1300 + assert torch.equal(output_cad.query_start_loc, expected_query_start_loc)

1301 +

1302 + # Verify hidden states (stored by reference, not copied)

1303 + assert proposer._dflash_hidden_states is target_hidden_states

vllm/config/speculative.py

+21 -6

@@ -47,8 +47,11 @@ MTPModelTypes = Literal[

47 "pangu_ultra_moe_mtp",

48 "step3p5_mtp",

49 ]

50 -EagleModelTypes = Literal["eagle", "eagle3", "extract_hidden_states", MTPModelTypes]

50 NgramGPUTypes = Literal["ngram_gpu"]

51 +DFlashModelTypes = Literal["dflash"]

52 +EagleModelTypes = Literal[

53 + "eagle", "eagle3", "extract_hidden_states", MTPModelTypes, DFlashModelTypes

54 +]

55 SpeculativeMethod = Literal[

56 "ngram",

57 "medusa",

@@ -206,7 +209,11 @@ class SpeculativeConfig:

209 factors: list[Any] = []

210 # Eagle3 and extract_hidden_states affect the computation graph because

211 # they return intermediate hidden states in addition to the final hidden state.

209 - uses_aux_hidden_states = self.method in ("eagle3", "extract_hidden_states")

212 + uses_aux_hidden_states = self.method in (

213 + "eagle3",

214 + "extract_hidden_states",

215 + "dflash",

216 + )

217 factors.append(uses_aux_hidden_states)

218

219 # The specific layers used also affect the computation graph

@@ -490,7 +497,7 @@ class SpeculativeConfig:

497 )

498

499 # Automatically detect the method

493 - if self.method in ("eagle", "eagle3"):

500 + if self.method in ("eagle", "eagle3", "dflash"):

501 pass

502 # examples:

503 # yuhuili/EAGLE-LLaMA3-Instruct-8B

@@ -500,6 +507,8 @@ class SpeculativeConfig:

507 self.method = "eagle"

508 elif "eagle3" in self.draft_model_config.model.lower():

509 self.method = "eagle3"

510 + elif "dflash" in self.draft_model_config.model.lower():

511 + self.method = "dflash"

512 elif self.draft_model_config.hf_config.model_type == "medusa":

513 self.method = "medusa"

514 elif self.draft_model_config.hf_config.model_type == "mlp_speculator":

@@ -532,7 +541,7 @@ class SpeculativeConfig:

541 )

542

543 # Replace hf_config for EAGLE draft_model

535 - if self.method in ("eagle", "eagle3"):

544 + if self.method in ("eagle", "eagle3", "dflash"):

545 from vllm.transformers_utils.configs.eagle import EAGLEConfig

546 from vllm.transformers_utils.configs.speculators import (

547 SpeculatorsConfig,

@@ -552,6 +561,9 @@ class SpeculativeConfig:

561 self.draft_model_config.hf_config = eagle_config

562 self.update_arch_()

563

564 + if self.method == "dflash":

565 + self.parallel_drafting = True

566 +

567 if self.num_speculative_tokens is not None and hasattr(

568 self.draft_model_config.hf_config, "num_lookahead_tokens"

569 ):

@@ -807,7 +819,7 @@ class SpeculativeConfig:

819 "kimi_k25",

820 ]

821 if (

810 - self.method in ("eagle3", "extract_hidden_states")

822 + self.method in ("eagle3", "extract_hidden_states", "dflash")

823 and self.target_model_config

824 and not any(

825 supported_model in self.target_model_config.hf_text_config.model_type

@@ -855,7 +867,10 @@ class SpeculativeConfig:

867 return slots_per_req

868

869 def use_eagle(self) -> bool:

858 - return self.method in ("eagle", "eagle3", "mtp")

870 + return self.method in ("eagle", "eagle3", "mtp", "dflash")

871 +

872 + def use_dflash(self) -> bool:

873 + return self.method == "dflash"

874

875 def uses_draft_model(self) -> bool:

876 return self.method == "draft_model"

vllm/config/vllm.py

+20

@@ -1327,6 +1327,26 @@ class VllmConfig:

1327 max_num_batched_tokens - scheduled_token_delta

1328 )

1329

1330 + if self.scheduler_config.max_num_scheduled_tokens <= 0:

1331 + raise ValueError(

1332 + "max_num_scheduled_tokens is set to"

1333 + f" {self.scheduler_config.max_num_scheduled_tokens} based on"

1334 + " the speculative decoding settings, which does not allow"

1335 + " any tokens to be scheduled. Increase max_num_batched_tokens"

1336 + " to accommodate the additional draft token slots, or decrease"

1337 + " num_speculative_tokens or max_num_seqs."

1338 + )

1339 + if self.scheduler_config.max_num_scheduled_tokens < 8192:

1340 + logger.warning_once(

1341 + "max_num_scheduled_tokens is set to"

1342 + f" {self.scheduler_config.max_num_scheduled_tokens} based on"

1343 + " the speculative decoding settings. This may lead to suboptimal"

1344 + " performance. Consider increasing max_num_batched_tokens to"

1345 + " accommodate the additional draft token slots, or decrease"

1346 + " num_speculative_tokens or max_num_seqs.",

1347 + scope="local",

1348 + )

1349 +

1350 max_num_scheduled_tokens = self.scheduler_config.max_num_scheduled_tokens

1351 if max_num_batched_tokens < max_num_scheduled_tokens + (

1352 self.speculative_config.max_num_new_slots_for_drafting

vllm/model_executor/models/qwen3.py

@@ -285,6 +285,7 @@ class Qwen3ForCausalLM(

285

286 self.config = config

287

288 + self.vllm_config = vllm_config

289 self.quant_config = quant_config

290 self.model = Qwen3Model(

291 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

vllm/model_executor/models/qwen3_dflash.py new

+619

@@ -0,0 +1,619 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from collections.abc import Iterable

5 +

6 +import torch

7 +import torch.nn.functional as F

8 +from torch import nn

9 +from transformers import Qwen3Config

10 +

11 +from vllm import _custom_ops as ops

12 +from vllm.compilation.decorators import support_torch_compile

13 +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config

14 +from vllm.distributed import get_tensor_model_parallel_world_size

15 +from vllm.logger import init_logger

16 +from vllm.model_executor.layers.attention import Attention

17 +from vllm.model_executor.layers.layernorm import RMSNorm

18 +from vllm.model_executor.layers.linear import (

19 + QKVParallelLinear,

20 + ReplicatedLinear,

21 + RowParallelLinear,

22 +)

23 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

24 +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig

25 +from vllm.model_executor.layers.rotary_embedding import get_rope

26 +from vllm.model_executor.layers.vocab_parallel_embedding import (

27 + ParallelLMHead,

28 + VocabParallelEmbedding,

29 +)

30 +from vllm.model_executor.model_loader.weight_utils import (

31 + default_weight_loader,

32 + maybe_remap_kv_scale_name,

33 +)

34 +from vllm.multimodal.inputs import NestedTensors

35 +from vllm.transformers_utils.config import set_default_rope_theta

36 +from vllm.v1.attention.backend import AttentionType

37 +

38 +from .qwen2 import Qwen2MLP as Qwen3MLP

39 +from .qwen3 import Qwen3ForCausalLM

40 +from .utils import (

41 + AutoWeightsLoader,

42 + get_draft_quant_config,

43 + maybe_prefix,

44 + process_eagle_weight,

45 +)

46 +

47 +logger = init_logger(__name__)

48 +

49 +

50 +class DFlashQwen3Attention(nn.Module):

51 + """Attention for DFlash speculative decoding.

52 +

53 + Context KVs are pre-inserted into the KV cache before the forward pass.

54 + This layer handles only query tokens via standard attention.

55 + Adapted from Qwen3Attention."""

56 +

57 + def __init__(

58 + self,

59 + hidden_size: int,

60 + num_heads: int,

61 + num_kv_heads: int,

62 + rope_parameters: dict,

63 + max_position: int = 4096 * 32,

64 + head_dim: int | None = None,

65 + rms_norm_eps: float = 1e-06,

66 + attention_bias: bool = False,

67 + cache_config: CacheConfig | None = None,

68 + quant_config: QuantizationConfig | None = None,

69 + prefix: str = "",

70 + attn_type: str = AttentionType.DECODER,

71 + ) -> None:

72 + super().__init__()

73 + self.layer_name = prefix

74 + self.hidden_size = hidden_size

75 + tp_size = get_tensor_model_parallel_world_size()

76 + self.total_num_heads = num_heads

77 + assert self.total_num_heads % tp_size == 0

78 + self.num_heads = self.total_num_heads // tp_size

79 + self.total_num_kv_heads = num_kv_heads

80 + if self.total_num_kv_heads >= tp_size:

81 + assert self.total_num_kv_heads % tp_size == 0

82 + else:

83 + assert tp_size % self.total_num_kv_heads == 0

84 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

85 + self.head_dim = head_dim or hidden_size // self.total_num_heads

86 + self.q_size = self.num_heads * self.head_dim

87 + self.kv_size = self.num_kv_heads * self.head_dim

88 + self.scaling = self.head_dim**-0.5

89 +

90 + self.qkv_proj = QKVParallelLinear(

91 + hidden_size,

92 + self.head_dim,

93 + self.total_num_heads,

94 + self.total_num_kv_heads,

95 + bias=attention_bias,

96 + quant_config=quant_config,

97 + prefix=f"{prefix}.qkv_proj",

98 + )

99 + self.o_proj = RowParallelLinear(

100 + self.total_num_heads * self.head_dim,

101 + hidden_size,

102 + bias=attention_bias, # DFlash has o_proj bias when using attention bias

103 + quant_config=quant_config,

104 + prefix=f"{prefix}.o_proj",

105 + )

106 +

107 + self.rotary_emb = get_rope(

108 + self.head_dim,

109 + max_position=max_position,

110 + rope_parameters=rope_parameters,

111 + )

112 + self.attn = Attention(

113 + self.num_heads,

114 + self.head_dim,

115 + self.scaling,

116 + num_kv_heads=self.num_kv_heads,

117 + cache_config=cache_config,

118 + quant_config=quant_config,

119 + prefix=f"{prefix}.attn",

120 + attn_type=attn_type,

121 + )

122 + self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)

123 + self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)

124 +

125 + def forward(

126 + self,

127 + positions: torch.Tensor,

128 + hidden_states: torch.Tensor,

129 + ) -> torch.Tensor:

130 + """DFlash attention assumes that the KV cache is already populated

131 + with the context K/V from the target model's hidden states. This forward op

132 + computes attention for the query tokens only.

133 + See also: precompute_and_store_context_kv"""

134 + qkv = F.linear(hidden_states, self.qkv_proj.weight, self.qkv_proj.bias)

135 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

136 +

137 + # Per-head RMSNorm

138 + q_shape, k_shape = q.shape, k.shape

139 + q = self.q_norm(

140 + q.view(*q_shape[:-1], q_shape[-1] // self.head_dim, self.head_dim)

141 + ).view(q_shape)

142 + k = self.k_norm(

143 + k.view(*k_shape[:-1], k_shape[-1] // self.head_dim, self.head_dim)

144 + ).view(k_shape)

145 +

146 + q, k = self.rotary_emb(positions, q, k)

147 +

148 + attn_output = self.attn(q, k, v)

149 + output, _ = self.o_proj(attn_output)

150 + return output

151 +

152 +

153 +class DFlashQwen3DecoderLayer(nn.Module):

154 + def __init__(

155 + self,

156 + vllm_config: VllmConfig,

157 + *,

158 + config: Qwen3Config,

159 + cache_config: CacheConfig | None = None,

160 + quant_config: QuantizationConfig | None = None,

161 + prefix: str = "",

162 + ) -> None:

163 + super().__init__()

164 + self.hidden_size = config.hidden_size

165 + set_default_rope_theta(config, default_theta=1000000)

166 + attn_type = AttentionType.DECODER

167 +

168 + self.self_attn = DFlashQwen3Attention(

169 + hidden_size=self.hidden_size,

170 + num_heads=config.num_attention_heads,

171 + max_position=config.max_position_embeddings,

172 + num_kv_heads=config.num_key_value_heads,

173 + rms_norm_eps=config.rms_norm_eps,

174 + attention_bias=getattr(config, "attention_bias", False),

175 + head_dim=getattr(config, "head_dim", None),

176 + cache_config=cache_config,

177 + quant_config=quant_config,

178 + rope_parameters=config.rope_parameters,

179 + prefix=f"{prefix}.self_attn",

180 + attn_type=attn_type,

181 + )

182 + self.mlp = Qwen3MLP(

183 + hidden_size=self.hidden_size,

184 + intermediate_size=config.intermediate_size,

185 + hidden_act=config.hidden_act,

186 + quant_config=quant_config,

187 + prefix=f"{prefix}.mlp",

188 + )

189 + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

190 + self.post_attention_layernorm = RMSNorm(

191 + config.hidden_size, eps=config.rms_norm_eps

192 + )

193 +

194 + def forward(

195 + self,

196 + positions: torch.Tensor,

197 + hidden_states: torch.Tensor,

198 + residual: torch.Tensor | None,

199 + ) -> tuple[torch.Tensor, torch.Tensor]:

200 + if residual is not None:

201 + hidden_states, residual = self.input_layernorm(hidden_states, residual)

202 + else:

203 + residual = hidden_states

204 + hidden_states = self.input_layernorm(hidden_states)

205 +

206 + hidden_states = self.self_attn(

207 + positions=positions,

208 + hidden_states=hidden_states,

209 + )

210 +

211 + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)

212 + hidden_states = self.mlp(hidden_states)

213 + return hidden_states, residual

214 +

215 +

216 +@support_torch_compile

217 +class DFlashQwen3Model(nn.Module):

218 + def __init__(

219 + self,

220 + *,

221 + vllm_config: VllmConfig,

222 + start_layer_id: int = 0,

223 + prefix: str = "",

224 + ) -> None:

225 + super().__init__()

226 + self.config = vllm_config.speculative_config.draft_model_config.hf_config

227 + self.vocab_size = self.config.vocab_size

228 + self.quant_config = get_draft_quant_config(vllm_config)

229 +

230 + drafter_config = getattr(self.config, "eagle_config", {})

231 + drafter_config.update(getattr(self.config, "dflash_config", {}))

232 +

233 + if drafter_config is not None and "use_aux_hidden_state" in drafter_config:

234 + self.use_aux_hidden_state = drafter_config["use_aux_hidden_state"]

235 + else:

236 + self.use_aux_hidden_state = True

237 +

238 + current_vllm_config = get_current_vllm_config()

239 +

240 + self.embed_tokens = VocabParallelEmbedding(

241 + self.config.vocab_size,

242 + self.config.hidden_size,

243 + prefix=maybe_prefix(prefix, "embed_tokens"),

244 + )

245 +

246 + self.layers = nn.ModuleList(

247 + [

248 + DFlashQwen3DecoderLayer(

249 + current_vllm_config,

250 + prefix=maybe_prefix(prefix, f"layers.{layer_idx + start_layer_id}"),

251 + config=self.config,

252 + )

253 + for layer_idx in range(self.config.num_hidden_layers)

254 + ]

255 + )

256 + if self.use_aux_hidden_state:

257 + num_features_to_use = self.config.num_hidden_layers

258 + if "target_layer_ids" in drafter_config:

259 + num_features_to_use = len(drafter_config["target_layer_ids"])

260 + elif "layer_ids" in drafter_config:

261 + num_features_to_use = len(drafter_config["layer_ids"])

262 + if hasattr(self.config, "target_hidden_size"):

263 + fc_input_size = self.config.target_hidden_size * num_features_to_use

264 + else:

265 + fc_input_size = self.config.hidden_size * num_features_to_use

266 + self.fc = ReplicatedLinear(

267 + input_size=fc_input_size,

268 + output_size=self.config.hidden_size,

269 + bias=False,

270 + params_dtype=vllm_config.model_config.dtype,

271 + quant_config=self.quant_config,

272 + prefix=maybe_prefix(prefix, "fc"),

273 + return_bias=False,

274 + )

275 + self.hidden_norm = RMSNorm(

276 + self.config.hidden_size,

277 + eps=self.config.rms_norm_eps,

278 + )

279 + self.norm = RMSNorm(

280 + self.config.hidden_size,

281 + eps=self.config.rms_norm_eps,

282 + )

283 +

284 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

285 + return self.embed_tokens(input_ids)

286 +

287 + def _build_fused_kv_buffers(self) -> None:

288 + """Build fused weight buffers for precompute_and_store_context_kv.

289 +

290 + Must be called after weights are loaded. Stacks the KV-projection

291 + weights, K-norm weights, and RoPE parameters from every attention

292 + layer so that precompute_and_store_context_kv can run one fused

293 + GEMM for all layers at once. Also aliases the weight of the hidden_norm.

294 + """

295 + layers_attn = [layer.self_attn for layer in self.layers]

296 + attn0 = layers_attn[0]

297 + has_bias = attn0.qkv_proj.bias is not None

298 +

299 + self._hidden_norm_weight = self.hidden_norm.weight.data

300 +

301 + # KV projection weights: [num_layers * 2 * kv_size, hidden_size]

302 + kv_weights = [a.qkv_proj.weight[a.q_size :] for a in layers_attn]

303 + self._fused_kv_weight = torch.cat(kv_weights, dim=0)

304 + if has_bias:

305 + kv_biases = [a.qkv_proj.bias[a.q_size :] for a in layers_attn]

306 + self._fused_kv_bias: torch.Tensor | None = torch.cat(kv_biases, dim=0)

307 + else:

308 + self._fused_kv_bias = None

309 +

310 + # K-norm weights: list of [head_dim] tensors, one per layer.

311 + self._k_norm_weights = [a.k_norm.weight.data for a in layers_attn]

312 +

313 + # RoPE parameters

314 + self._rope_head_size = attn0.rotary_emb.head_size

315 + self._rope_cos_sin_cache = attn0.rotary_emb.cos_sin_cache

316 + self._rope_is_neox = attn0.rotary_emb.is_neox_style

317 + # Validation that RoPE params are the same across all layers

318 + for attn in layers_attn[1:]:

319 + assert (

320 + attn.rotary_emb.head_size == self._rope_head_size

321 + and attn.rotary_emb.is_neox_style == self._rope_is_neox

322 + ), "All layers must have the same RoPE parameters for DFlash precomputation"

323 +

324 + # Layer metadata

325 + self._num_attn_layers = len(layers_attn)

326 + self._kv_size = attn0.kv_size

327 + self._head_dim = attn0.head_dim

328 + self._num_kv_heads = attn0.num_kv_heads

329 + self._rms_norm_eps = attn0.q_norm.variance_epsilon

330 + # Validation that all layers have the same attention config

331 + for attn in layers_attn[1:]:

332 + assert (

333 + attn.kv_size == self._kv_size

334 + and attn.head_dim == self._head_dim

335 + and attn.num_kv_heads == self._num_kv_heads

336 + and attn.q_norm.variance_epsilon == self._rms_norm_eps

337 + ), "All layers must have the same attn config for DFlash precomputation"

338 +

339 + # References to inner Attention layers for direct cache writes

340 + self._attn_layers = [layer.self_attn.attn for layer in self.layers]

341 +

342 + def precompute_and_store_context_kv(

343 + self,

344 + context_states: torch.Tensor,

345 + context_positions: torch.Tensor,

346 + context_slot_mapping: torch.Tensor | None = None,

347 + ) -> None:

348 + """Precompute K/V for context states write them into each layer's KV cache.

349 +

350 + Input context states are projected to K/V, normed, and have RoPE applied.

351 + Since the context shape is different than the query shape, we can't rely on the

352 + regular forward pass to apply torch.compile and CUDA graphs to this section.

353 + As such, this function is optimized to minimize the number of torch ops present:

354 + we use fused vLLM kernels for RMSNorm and RoPE, fuse the GEMM into one

355 + large projection, and avoid cloning buffers (with .contiguous()) where possible.

356 +

357 + When context_slot_mapping is None (e.g. during dummy_run) only

358 + the computation runs, and no K/V is written to cache.

359 + """

360 + if not hasattr(self, "_num_attn_layers"):

361 + logger.warning_once(

362 + "DFlash buffer initialization was skipped. If dummy weights are not "

363 + "in use, this may indicate an error in weight loading."

364 + )

365 + self._build_fused_kv_buffers()

366 +

367 + num_ctx = context_states.shape[0]

368 + L = self._num_attn_layers

369 + kv = self._kv_size

370 + hd = self._head_dim

371 + nkv = self._num_kv_heads

372 +

373 + # --- Fused KV projection (one GEMM for all layers) ---

374 + normed_context_states = torch.empty_like(context_states)

375 + ops.rms_norm(

376 + normed_context_states,

377 + context_states,

378 + self._hidden_norm_weight,

379 + self._rms_norm_eps,

380 + )

381 + all_kv_flat = F.linear(

382 + normed_context_states, self._fused_kv_weight, self._fused_kv_bias

383 + )

384 + # Single contiguous copy that separates K/V and transposes to

385 + # layer-major layout. Result: [2, L, num_ctx, nkv, hd] contiguous.

386 + # Indexing dim-0 gives contiguous [L, num_ctx, nkv, hd] for K and V.

387 + all_kv = (

388 + all_kv_flat.view(num_ctx, L, 2, nkv, hd).permute(2, 1, 0, 3, 4).contiguous()

389 + )

390 + all_k = all_kv[0] # [L, num_ctx, nkv, hd], contiguous

391 + all_v = all_kv[1] # [L, num_ctx, nkv, hd], contiguous

392 +

393 + # --- Per-layer RMSNorm K (3D: [num_ctx, nkv, hd] per layer) ---

394 + all_k_normed = torch.empty_like(all_k)

395 + for i in range(L):

396 + ops.rms_norm(

397 + all_k_normed[i],

398 + all_k[i],

399 + self._k_norm_weights[i],

400 + self._rms_norm_eps,

401 + )

402 +

403 + # --- Fused RoPE across all layers ---

404 + # View as [L * num_ctx, kv] so RoPE sees one big batch (no copy).

405 + # In-place RoPE: pass K as the "query" arg with key=None.

406 + all_k_flat = all_k_normed.view(L * num_ctx, kv)

407 + positions_repeated = context_positions.repeat(L)

408 + cos_sin_cache = self._rope_cos_sin_cache

409 + if cos_sin_cache.dtype != all_k_flat.dtype:

410 + cos_sin_cache = cos_sin_cache.to(dtype=all_k_flat.dtype)

411 + ops.rotary_embedding(

412 + positions_repeated,

413 + all_k_flat,

414 + None,

415 + self._rope_head_size,

416 + cos_sin_cache,

417 + self._rope_is_neox,

418 + )

419 +

420 + if context_slot_mapping is None:

421 + return

422 +

423 + # --- Per-layer cache insert ---

424 + all_k_final = all_k_flat.view(L, num_ctx, nkv, hd)

425 + for i in range(L):

426 + attn = self._attn_layers[i]

427 + kv_cache = attn.kv_cache

428 + attn.impl.do_kv_cache_update(

429 + attn,

430 + all_k_final[i],

431 + all_v[i],

432 + kv_cache,

433 + context_slot_mapping,

434 + )

435 +

436 + def forward(

437 + self,

438 + input_ids: torch.Tensor,

439 + positions: torch.Tensor,

440 + input_embeds: torch.Tensor | None = None,

441 + ) -> torch.Tensor:

442 + if input_embeds is None:

443 + input_embeds = self.embed_input_ids(input_ids)

444 +

445 + hidden_states = input_embeds

446 +

447 + residual = None

448 + for layer in self.layers:

449 + hidden_states, residual = layer(

450 + positions=positions,

451 + hidden_states=hidden_states,

452 + residual=residual,

453 + )

454 + hidden_states, _ = self.norm(hidden_states, residual)

455 + return hidden_states

456 +

457 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

458 + stacked_params_mapping = [

459 + (".qkv_proj", ".q_proj", "q"),

460 + (".qkv_proj", ".k_proj", "k"),

461 + (".qkv_proj", ".v_proj", "v"),

462 + (".gate_up_proj", ".gate_proj", 0),

463 + (".gate_up_proj", ".up_proj", 1),

464 + ]

465 + params_dict = dict(self.named_parameters())

466 + loaded_params: set[str] = set()

467 + for name, loaded_weight in weights:

468 + if "midlayer." in name:

469 + name = name.replace("midlayer.", "layers.0.")

470 + if self.quant_config is not None and (

471 + scale_name := self.quant_config.get_cache_scale(name)

472 + ):

473 + param = params_dict[scale_name]

474 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

475 + loaded_weight = (

476 + loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]

477 + )

478 + weight_loader(param, loaded_weight)

479 + loaded_params.add(scale_name)

480 + continue

481 + if "scale" in name:

482 + name = maybe_remap_kv_scale_name(name, params_dict)

483 + if name is None:

484 + continue

485 + for param_name, weight_name, shard_id in stacked_params_mapping:

486 + if weight_name not in name:

487 + continue

488 + name = name.replace(weight_name, param_name)

489 + param = params_dict[name]

490 + weight_loader = param.weight_loader

491 + weight_loader(param, loaded_weight, shard_id)

492 + break

493 + else:

494 + param = params_dict[name]

495 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

496 + weight_loader(param, loaded_weight)

497 + loaded_params.add(name)

498 + return loaded_params

499 +

500 +

501 +class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):

502 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

503 + nn.Module.__init__(self)

504 + self.config = vllm_config.speculative_config.draft_model_config.hf_config

505 + if getattr(self.config, "draft_vocab_size", None) is None:

506 + self.config.draft_vocab_size = getattr(self.config, "vocab_size", None)

507 + target_layer_num = vllm_config.model_config.get_num_layers(

508 + vllm_config.parallel_config

509 + )

510 + self.config.target_layer_count = target_layer_num

511 + self.model = DFlashQwen3Model(

512 + vllm_config=vllm_config,

513 + prefix="model",

514 + start_layer_id=target_layer_num,

515 + )

516 +

517 + logit_scale = getattr(self.config, "logit_scale", 1.0)

518 + self.lm_head = ParallelLMHead(

519 + self.config.draft_vocab_size,

520 + self.config.hidden_size,

521 + prefix=maybe_prefix(prefix, "lm_head"),

522 + )

523 + self.logits_processor = LogitsProcessor(

524 + self.config.draft_vocab_size, scale=logit_scale

525 + )

526 + self.draft_id_to_target_id = None

527 +

528 + def embed_input_ids(

529 + self,

530 + input_ids: torch.Tensor,

531 + multimodal_embeddings: NestedTensors | None = None,

532 + is_multimodal: torch.Tensor | None = None,

533 + ) -> torch.Tensor:

534 + return self.model.embed_input_ids(input_ids)

535 +

536 + def forward(

537 + self,

538 + input_ids: torch.Tensor,

539 + positions: torch.Tensor,

540 + inputs_embeds: torch.Tensor | None = None,

541 + ) -> torch.Tensor:

542 + return self.model(input_ids, positions, inputs_embeds)

543 +

544 + def compute_logits(

545 + self,

546 + hidden_states: torch.Tensor,

547 + ) -> torch.Tensor | None:

548 + logits = self.logits_processor(self.lm_head, hidden_states)

549 + if self.draft_id_to_target_id is None:

550 + return logits

551 +

552 + base = torch.arange(self.config.draft_vocab_size, device=logits.device)

553 + targets = base + self.draft_id_to_target_id

554 + logits_new = logits.new_full(

555 + (logits.shape[0], self.config.vocab_size),

556 + float("-inf"),

557 + )

558 + logits_new[:, targets] = logits

559 + return logits_new

560 +

561 + def precompute_and_store_context_kv(

562 + self,

563 + context_states: torch.Tensor,

564 + context_positions: torch.Tensor,

565 + context_slot_mapping: torch.Tensor | None = None,

566 + ) -> None:

567 + """Precompute projected + RoPE'd K/V and write to cache."""

568 + self.model.precompute_and_store_context_kv(

569 + context_states, context_positions, context_slot_mapping

570 + )

571 +

572 + def combine_hidden_states(

573 + self,

574 + hidden_states: torch.Tensor,

575 + ) -> torch.Tensor:

576 + if not self.model.use_aux_hidden_state:

577 + return hidden_states

578 + needs_squeeze = hidden_states.dim() == 1

579 + if needs_squeeze:

580 + hidden_states = hidden_states.unsqueeze(0)

581 + result = self.model.fc(hidden_states)

582 + if needs_squeeze:

583 + result = result.squeeze(0)

584 + return result

585 +

586 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

587 + model_weights = {}

588 + includes_draft_id_mapping = False

589 + includes_embed_tokens = False

590 + for name, loaded_weight in weights:

591 + assert "mask_hidden" not in name, (

592 + "DFlash should use mask_token_id to embed the padding hidden state"

593 + )

594 + if "t2d" in name:

595 + continue

596 + if "d2t" in name:

597 + name = name.replace("d2t", "draft_id_to_target_id")

598 + includes_draft_id_mapping = True

599 + elif "lm_head" not in name:

600 + name = "model." + name

601 + if "embed_tokens" in name:

602 + includes_embed_tokens = True

603 + model_weights[name] = loaded_weight

604 + process_eagle_weight(self, name)

605 +

606 + skip_substrs = []

607 + if not includes_draft_id_mapping:

608 + skip_substrs.append("draft_id_to_target_id")

609 + if not includes_embed_tokens:

610 + skip_substrs.append("embed_tokens")

611 + if not self.model.use_aux_hidden_state:

612 + skip_substrs.append("fc.")

613 + loader = AutoWeightsLoader(

614 + self,

615 + skip_prefixes=None,

616 + skip_substrs=skip_substrs,

617 + )

618 + loader.load_weights(model_weights.items())

619 + self.model._build_fused_kv_buffers()

vllm/model_executor/models/qwen3_next.py

+6 -8

@@ -56,6 +56,7 @@ from vllm.sequence import IntermediateTensors

56 from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig

58 from .interfaces import (

59 + EagleModelMixin,

60 HasInnerState,

61 IsHybrid,

62 MixtureOfExperts,

@@ -454,7 +455,7 @@ class Qwen3NextDecoderLayer(nn.Module):

455

456

457 @support_torch_compile

457 -class Qwen3NextModel(nn.Module):

458 +class Qwen3NextModel(nn.Module, EagleModelMixin):

459 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

460 super().__init__()

461

@@ -492,8 +493,6 @@ class Qwen3NextModel(nn.Module):

493 else:

494 self.norm = PPMissingLayer()

495

495 - self.aux_hidden_state_layers: tuple[int, ...] = ()

496 -

496 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

497 return self.embed_tokens(input_ids)

498

@@ -515,20 +514,19 @@ class Qwen3NextModel(nn.Module):

514 hidden_states = intermediate_tensors["hidden_states"]

515 residual = intermediate_tensors["residual"]

516

518 - aux_hidden_states = []

517 + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)

518 for layer_idx, layer in enumerate(

519 islice(self.layers, self.start_layer, self.end_layer),

520 start=self.start_layer,

521 ):

523 - if layer_idx in self.aux_hidden_state_layers:

524 - aux_hidden_states.append(

525 - hidden_states + residual if residual is not None else hidden_states

526 - )

522 hidden_states, residual = layer(

523 positions=positions,

524 hidden_states=hidden_states,

525 residual=residual,

526 )

527 + self._maybe_add_hidden_state(

528 + aux_hidden_states, layer_idx + 1, hidden_states, residual

529 + )

530

531 if not get_pp_group().is_last_rank:

532 return IntermediateTensors(

vllm/model_executor/models/registry.py

@@ -546,6 +546,7 @@ _SPECULATIVE_DECODING_MODELS = {

546 "EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"),

547 "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),

548 "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),

549 + "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),

550 "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),

551 "LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"),

552 "Eagle3Qwen2_5vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),

vllm/transformers_utils/configs/eagle.py

+12 -1

@@ -62,9 +62,20 @@ class EAGLEConfig(PretrainedConfig):

62 else f"Eagle3{arch}"

63 for arch in self.model.architectures

64 ]

65 + elif method == "dflash":

66 + assert self.model is not None, (

67 + "model should not be None when method is dflash"

68 + )

69 + kwargs["architectures"] = [

70 + arch

71 + if arch.startswith("DFlash") or arch.endswith("DFlash")

72 + else f"DFlash{arch}"

73 + for arch in self.model.architectures

74 + ]

75 else:

76 raise ValueError(

67 - f"Invalid method {method}. Supported methods are eagle and eagle3."

77 + f"Invalid method {method}. Supported methods are "

78 + "eagle, eagle3, and dflash."

79 )

81 super().__init__(**kwargs)

vllm/v1/attention/backend.py

+14

@@ -220,6 +220,17 @@ class AttentionBackend(ABC):

220 def supports_per_head_quant_scales(cls) -> bool:

221 return False

222

223 + @classmethod

224 + def supports_non_causal(cls) -> bool:

225 + """Check if backend supports non-causal (bidirectional) attention

226 + for decoder models.

227 +

228 + Unlike ENCODER_ONLY attention type which implies a different

229 + execution model, this refers to non-causal attention within the

230 + standard paged-KV-cache decoder path.

231 + """

232 + return False

233 +

234 @classmethod

235 def supports_attn_type(cls, attn_type: str) -> bool:

236 """Check if backend supports a given attention type.

@@ -261,6 +272,7 @@ class AttentionBackend(ABC):

272 use_per_head_quant_scales: bool,

273 device_capability: "DeviceCapability",

274 attn_type: str,

275 + use_non_causal: bool = False,

276 ) -> list[str]:

277 invalid_reasons = []

278 if not cls.supports_head_size(head_size):

@@ -293,6 +305,8 @@ class AttentionBackend(ABC):

305 invalid_reasons.append("compute capability not supported")

306 if not cls.supports_attn_type(attn_type):

307 invalid_reasons.append(f"attention type {attn_type} not supported")

308 + if use_non_causal and not cls.supports_non_causal():

309 + invalid_reasons.append("non-causal attention not supported")

310 combination_reason = cls.supports_combination(

311 head_size,

312 dtype,

vllm/v1/attention/backends/flash_attn.py

@@ -101,6 +101,10 @@ class FlashAttentionBackend(AttentionBackend):

101 def get_name() -> str:

102 return "FLASH_ATTN"

103

104 + @classmethod

105 + def supports_non_causal(cls) -> bool:

106 + return True

107 +

108 @classmethod

109 def supports_attn_type(cls, attn_type: str) -> bool:

110 """FlashAttention supports all attention types."""

vllm/v1/attention/selector.py

+9 -1

@@ -29,6 +29,7 @@ class AttentionSelectorConfig(NamedTuple):

29 use_mm_prefix: bool = False

30 use_per_head_quant_scales: bool = False

31 attn_type: str = AttentionType.DECODER

32 + use_non_causal: bool = False

34 def __repr__(self):

35 return (

@@ -41,7 +42,8 @@ class AttentionSelectorConfig(NamedTuple):

42 f"use_sparse={self.use_sparse}, "

43 f"use_mm_prefix={self.use_mm_prefix}, "

44 f"use_per_head_quant_scales={self.use_per_head_quant_scales}, "

44 - f"attn_type={self.attn_type})"

45 + f"attn_type={self.attn_type}, "

46 + f"use_non_causal={self.use_non_causal})"

47 )

@@ -76,6 +78,11 @@ def get_attn_backend(

78 else:

79 block_size = None

81 + speculative_config = vllm_config.speculative_config

82 + use_non_causal = (

83 + speculative_config is not None and speculative_config.method == "dflash"

84 + )

85 +

86 attn_selector_config = AttentionSelectorConfig(

87 head_size=head_size,

88 dtype=dtype,

@@ -87,6 +94,7 @@ def get_attn_backend(

94 use_mm_prefix=use_mm_prefix,

95 use_per_head_quant_scales=use_per_head_quant_scales,

96 attn_type=attn_type or AttentionType.DECODER,

97 + use_non_causal=use_non_causal,

98 )

100 return _cached_get_attn_backend(

vllm/v1/spec_decode/dflash.py new

+282

@@ -0,0 +1,282 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from typing import Any

5 +

6 +import torch

7 +from typing_extensions import override

8 +

9 +from vllm.config import VllmConfig

10 +from vllm.forward_context import set_forward_context

11 +from vllm.logger import init_logger

12 +from vllm.triton_utils import triton

13 +from vllm.v1.attention.backend import CommonAttentionMetadata

14 +from vllm.v1.spec_decode.eagle import SpecDecodeBaseProposer

15 +from vllm.v1.spec_decode.utils import copy_and_expand_dflash_inputs_kernel

16 +

17 +logger = init_logger(__name__)

18 +

19 +

20 +class DFlashProposer(SpecDecodeBaseProposer):

21 + def __init__(

22 + self,

23 + vllm_config: VllmConfig,

24 + device: torch.device,

25 + runner=None,

26 + ):

27 + assert vllm_config.speculative_config is not None

28 + assert vllm_config.speculative_config.method == "dflash"

29 + super().__init__(

30 + vllm_config=vllm_config,

31 + device=device,

32 + pass_hidden_states_to_model=True,

33 + runner=runner,

34 + )

35 +

36 + # Only next_token_ids and mask tokens are query tokens, all other context is K/V

37 + self.max_query_tokens = self.max_batch_size * (1 + self.num_speculative_tokens)

38 + # Positions covers both context states + query states

39 + self.max_positions = self.max_num_tokens + self.max_query_tokens

40 +

41 + # Separate context buffers to keep query buffer addresses stable for CUDA graphs

42 + self._context_slot_mapping_buffer = torch.zeros(

43 + self.max_num_tokens,

44 + dtype=torch.int64,

45 + device=device,

46 + )

47 + self._slot_mapping_buffer = torch.zeros(

48 + self.max_query_tokens,

49 + dtype=torch.int64,

50 + device=device,

51 + )

52 + self._context_positions_buffer = torch.zeros(

53 + self.max_num_tokens,

54 + dtype=torch.int64,

55 + device=device,

56 + )

57 + self.positions = torch.zeros(

58 + self.max_query_tokens,

59 + dtype=torch.int64,

60 + device=device,

61 + )

62 +

63 + self.arange = torch.arange(

64 + self.max_positions + 1, device=device, dtype=torch.int32

65 + )

66 +

67 + # For DFlash we use the input embeddings to embed the mask token

68 + self.parallel_drafting_hidden_state_tensor = None

69 +

70 + @override

71 + def _raise_if_multimodal(self):

72 + # Override to allow multimodal inputs since DFlash supports Qwen3.5 models

73 + # Support for multimodal inputs has not been tested.

74 + pass

75 +

76 + @override

77 + def set_inputs_first_pass(

78 + self,

79 + target_token_ids: torch.Tensor,

80 + next_token_ids: torch.Tensor,

81 + target_positions: torch.Tensor,

82 + target_hidden_states: torch.Tensor,

83 + token_indices_to_sample: torch.Tensor | None,

84 + cad: CommonAttentionMetadata,

85 + num_rejected_tokens_gpu: torch.Tensor | None,

86 + ) -> tuple[int, torch.Tensor, CommonAttentionMetadata]:

87 + # DFlash cross-attention: context K/V from target hidden states,

88 + # Q from query embeddings (bonus + mask tokens).

89 + batch_size = cad.batch_size()

90 + num_context = target_token_ids.shape[0]

91 + num_query_per_req = 1 + self.num_speculative_tokens

92 + num_query_total = batch_size * num_query_per_req

93 +

94 + # Store for build_model_inputs_first_pass to use

95 + self._dflash_num_context = num_context

96 +

97 + # We don't need to copy into a buffer here since the context preprocessing

98 + # does not run in a CUDA graph

99 + self._dflash_hidden_states = target_hidden_states

100 +

101 + token_indices_to_sample = torch.empty(

102 + batch_size * self.num_speculative_tokens,

103 + dtype=torch.int32,

104 + device=self.device,

105 + )

106 +

107 + # Launch fused triton kernel for input_ids, positions, slot_mapping,

108 + # and token_indices_to_sample

109 + max_ctx_per_req = cad.max_query_len

110 + max_tokens_per_req = max_ctx_per_req + num_query_per_req

111 + BLOCK_SIZE = min(256, triton.next_power_of_2(max_tokens_per_req))

112 + num_blocks = triton.cdiv(max_tokens_per_req, BLOCK_SIZE)

113 + grid = (batch_size, num_blocks)

114 +

115 + has_num_rejected = num_rejected_tokens_gpu is not None

116 + copy_and_expand_dflash_inputs_kernel[grid](

117 + # Inputs

118 + next_token_ids_ptr=next_token_ids,

119 + target_positions_ptr=target_positions,

120 + # Outputs

121 + out_input_ids_ptr=self.input_ids,

122 + out_context_positions_ptr=self._context_positions_buffer,

123 + out_query_positions_ptr=self.positions,

124 + out_context_slot_mapping_ptr=self._context_slot_mapping_buffer,

125 + out_query_slot_mapping_ptr=self._slot_mapping_buffer,

126 + out_token_indices_ptr=token_indices_to_sample,

127 + # Block table

128 + block_table_ptr=cad.block_table_tensor,

129 + block_table_stride=cad.block_table_tensor.stride(0),

130 + # Metadata

131 + query_start_loc_ptr=cad.query_start_loc,

132 + num_rejected_tokens_ptr=(

133 + num_rejected_tokens_gpu if has_num_rejected else 0

134 + ),

135 + # Scalars

136 + parallel_drafting_token_id=self.parallel_drafting_token_id,

137 + block_size=self.block_size,

138 + num_query_per_req=num_query_per_req,

139 + num_speculative_tokens=self.num_speculative_tokens,

140 + total_input_tokens=num_context,

141 + BLOCK_SIZE=BLOCK_SIZE,

142 + HAS_NUM_REJECTED=has_num_rejected,

143 + )

144 +

145 + query_slot_mapping = self._slot_mapping_buffer[:num_query_total]

146 + new_query_start_loc = self.arange[: batch_size + 1] * num_query_per_req

147 +

148 + # In padded mode, cad.seq_lens includes rejected tokens. Subtract

149 + # them so attention only sees the valid prefix of context states.

150 + effective_seq_lens = cad.seq_lens

151 + if has_num_rejected:

152 + effective_seq_lens = effective_seq_lens - num_rejected_tokens_gpu

153 +

154 + new_cad = CommonAttentionMetadata(

155 + query_start_loc=new_query_start_loc,

156 + seq_lens=effective_seq_lens + num_query_per_req,

157 + query_start_loc_cpu=(

158 + torch.from_numpy(self.token_arange_np[: batch_size + 1]).clone()

159 + * num_query_per_req

160 + ),

161 + _seq_lens_cpu=None,

162 + _num_computed_tokens_cpu=None,

163 + num_reqs=cad.num_reqs,

164 + num_actual_tokens=num_query_total,

165 + max_query_len=num_query_per_req,

166 + max_seq_len=cad.max_seq_len + num_query_per_req,

167 + block_table_tensor=cad.block_table_tensor,

168 + slot_mapping=query_slot_mapping,

169 + causal=False, # Non-causal attention is required for DFlash

170 + )

171 +

172 + return num_query_total, token_indices_to_sample, new_cad

173 +

174 + @override

175 + @torch.inference_mode()

176 + def dummy_run(

177 + self,

178 + num_tokens: int,

179 + use_cudagraphs: bool = True,

180 + is_graph_capturing: bool = False,

181 + slot_mappings: dict[str, torch.Tensor] | None = None,

182 + ) -> None:

183 + """

184 + Key differences to default dummy_run:

185 + - Only one forward pass due to parallel drafting

186 + - DFlash uses context states as unpadded metadata, so hidden_states will

187 + use the unpadded num_tokens instead of num_input_tokens

188 + - max_query_tokens is quite small, DFlash only sees spec tokens as queries

189 + - Multimodal inputs are not currently supported

190 + """

191 + num_query_tokens = min(num_tokens, self.max_query_tokens)

192 + cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (

193 + self._determine_batch_execution_and_padding(

194 + num_query_tokens, use_cudagraphs=use_cudagraphs

195 + )

196 + )

197 +

198 + # Slot mapping sized to num_input_tokens (query only), matching

199 + # the K/V tensor size from the model forward. Context KVs are

200 + # pre-inserted separately and don't flow through the model.

201 + if (

202 + self._draft_attn_layer_names

203 + and slot_mappings is not None

204 + and next(iter(self._draft_attn_layer_names)) in slot_mappings

205 + ):

206 + slot_mapping_dict = self._get_slot_mapping(num_input_tokens)

207 + else:

208 + slot_mapping_dict = slot_mappings or {}

209 +

210 + # Context and query positions use separate buffers; no copy needed.

211 + context_positions = self._context_positions_buffer[:num_tokens]

212 + # Context states will be passed directly to the precomputation without

213 + # going through the buffer, since no CUDA graph is used for the precomputation.

214 + # For the dummy run, we use the dummy buffer.

215 + context_states = self.hidden_states[:num_tokens]

216 +

217 + # Run the KV projection (GEMM + norms + RoPE) for memory profiling,

218 + self.model.precompute_and_store_context_kv(context_states, context_positions)

219 + with set_forward_context(

220 + None,

221 + self.vllm_config,

222 + num_tokens=num_input_tokens,

223 + num_tokens_across_dp=num_tokens_across_dp,

224 + cudagraph_runtime_mode=cudagraph_runtime_mode,

225 + slot_mapping=slot_mapping_dict,

226 + ):

227 + self.model(

228 + input_ids=self.input_ids[:num_input_tokens],

229 + positions=self._get_positions(num_input_tokens),

230 + inputs_embeds=None,

231 + )

232 +

233 + @override

234 + def build_model_inputs_first_pass(

235 + self,

236 + num_tokens: int,

237 + num_input_tokens: int,

238 + mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None,

239 + ) -> tuple[dict[str, Any], int]:

240 + # Context and query positions/slots were written to separate

241 + # buffers by the kernel — no copy needed.

242 + num_context = self._dflash_num_context

243 +

244 + # Pre-insert context KVs directly into cache

245 + self.model.precompute_and_store_context_kv(

246 + self._dflash_hidden_states, # Shape is already [num_context, hidden_size]

247 + self._context_positions_buffer[:num_context],

248 + self._context_slot_mapping_buffer[:num_context],

249 + )

250 + return (

251 + dict(

252 + input_ids=self.input_ids[:num_input_tokens],

253 + positions=self._get_positions(num_input_tokens),

254 + inputs_embeds=None,

255 + ),

256 + num_input_tokens,

257 + )

258 +

259 + @override

260 + def build_per_layer_attn_metadata(

261 + self, cad: CommonAttentionMetadata, draft_index: int = 0

262 + ) -> dict[str, object]:

263 + per_layer_attention_metadata = super().build_per_layer_attn_metadata(

264 + cad, draft_index

265 + )

266 + for layer_name, attn_metadata in per_layer_attention_metadata.items():

267 + assert getattr(attn_metadata, "causal", None) is False, (

268 + f"Attention metadata for layer {layer_name} does not have"

269 + " non-causal support, which is required for DFlash."

270 + " Consider using a different attention backend, such as FlashAttention."

271 + )

272 + return per_layer_attention_metadata

273 +

274 + @override

275 + def _get_eagle3_use_aux_hidden_state_from_config(self):

276 + use_aux_hidden_state = True

277 + dflash_config = getattr(

278 + self.draft_model_config.hf_config, "dflash_config", None

279 + )

280 + if dflash_config is not None:

281 + use_aux_hidden_state = dflash_config.get("use_aux_hidden_state", True)

282 + return use_aux_hidden_state

vllm/v1/spec_decode/eagle.py

+119 -72

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import ast

4 from importlib.util import find_spec

5 -from typing import cast

5 +from typing import Any, cast

7 import numpy as np

8 import torch

@@ -23,6 +23,7 @@ from vllm.model_executor.models import supports_multimodal

23 from vllm.model_executor.models.deepseek_eagle3 import Eagle3DeepseekV2ForCausalLM

24 from vllm.model_executor.models.interfaces import SupportsMultiModal

25 from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM

26 +from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM

27 from vllm.multimodal import MULTIMODAL_REGISTRY

28 from vllm.platforms import current_platform

29 from vllm.triton_utils import triton

@@ -83,13 +84,15 @@ class SpecDecodeBaseProposer:

84 self.hidden_size = self.draft_model_config.get_hidden_size()

85 self.inputs_embeds_size = self.draft_model_config.get_inputs_embeds_size()

86 - # Unifying eagle, draft model, and parallel drafting support

87 + # Unifying eagle, draft model, and parallel drafting support.

88 + # DFlash always uses parallel drafting (all tokens in one pass),

89 + # but has an additional slot for the next_token_id (does not shift like EAGLE)

90 self.parallel_drafting: bool = self.speculative_config.parallel_drafting

91 self.extra_slots_per_request = (

92 1 if not self.parallel_drafting else self.num_speculative_tokens

93 )

94 self.net_num_new_slots_per_request = self.extra_slots_per_request - (

92 - 1 if self.pass_hidden_states_to_model else 0

95 + 1 if (self.pass_hidden_states_to_model and self.method != "dflash") else 0

96 )

97 self.needs_extra_input_slots = self.net_num_new_slots_per_request > 0

@@ -101,10 +104,14 @@ class SpecDecodeBaseProposer:

104 self.speculative_config.use_local_argmax_reduction

105 )

106

104 - max_batch_size = vllm_config.scheduler_config.max_num_seqs

107 + self.max_batch_size = vllm_config.scheduler_config.max_num_seqs

108 self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens

109 self.token_arange_np = np.arange(self.max_num_tokens)

110

111 + # Can be specialized by methods like DFlash to reduce the limit

112 + self.max_query_tokens = self.max_num_tokens

113 + self.max_positions = self.max_num_tokens

114 +

115 # Multi-modal data support

116 self.mm_registry = MULTIMODAL_REGISTRY

117 self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(

@@ -146,18 +153,20 @@ class SpecDecodeBaseProposer:

153 # 1D-RoPE.

154 # See page 5 of https://arxiv.org/abs/2409.12191

155 self.mrope_positions = torch.zeros(

149 - (3, self.max_num_tokens + 1), dtype=torch.int64, device=device

156 + (3, self.max_positions + 1), dtype=torch.int64, device=device

157 )

158 elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:

159 self.xdrope_positions = torch.zeros(

153 - (self.uses_xdrope_dim, self.max_num_tokens + 1),

160 + (self.uses_xdrope_dim, self.max_positions + 1),

161 dtype=torch.int64,

162 device=device,

163 )

164 else:

165 # RoPE need (max_num_tokens,)

166 self.positions = torch.zeros(

160 - self.max_num_tokens, dtype=torch.int64, device=device

167 + self.max_positions,

168 + dtype=torch.int64,

169 + device=device,

170 )

171 self.hidden_states = torch.zeros(

172 (self.max_num_tokens, self.hidden_size), dtype=self.dtype, device=device

@@ -168,7 +177,7 @@ class SpecDecodeBaseProposer:

177

178 # We need +1 here because the arange is used to set query_start_loc,

179 # which has one more element than batch_size.

171 - max_num_slots_for_arange = max(max_batch_size + 1, self.max_num_tokens)

180 + max_num_slots_for_arange = max(self.max_batch_size + 1, self.max_num_tokens)

181 self.arange = torch.arange(

182 max_num_slots_for_arange, device=device, dtype=torch.int32

183 )

@@ -200,7 +209,7 @@ class SpecDecodeBaseProposer:

209 )

210

211 self.backup_next_token_ids = CpuGpuBuffer(

203 - max_batch_size,

212 + self.max_batch_size,

213 dtype=torch.int32,

214 pin_memory=is_pin_memory_available(),

215 device=device,

@@ -208,7 +217,9 @@ class SpecDecodeBaseProposer:

217 )

218

219 self._slot_mapping_buffer = torch.zeros(

211 - self.max_num_tokens, dtype=torch.int64, device=device

220 + self.max_positions,

221 + dtype=torch.int64,

222 + device=device,

223 )

224

225 # Determine allowed attention backends once during initialization.

@@ -275,7 +286,7 @@ class SpecDecodeBaseProposer:

286 # Precompute draft position offsets in flattened tree.

287 self.tree_draft_pos_offsets = torch.arange(

288 1, len(self.tree_choices) + 1, device=device, dtype=torch.int32

278 - ).repeat(max_batch_size, 1)

289 + ).repeat(self.max_batch_size, 1)

290

291 def _raise_if_padded_drafter_batch_disabled(self):

292 if self.speculative_config.disable_padded_drafter_batch:

@@ -305,14 +316,19 @@ class SpecDecodeBaseProposer:

316 # for those masked slots.

317

318 model_hf_config = self.draft_model_config.hf_config

308 - if hasattr(model_hf_config, "pard_token"):

319 + # DFlash stores mask_token_id in dflash_config

320 + dflash_config = getattr(model_hf_config, "dflash_config", None)

321 + if dflash_config and "mask_token_id" in dflash_config:

322 + self.parallel_drafting_token_id = dflash_config["mask_token_id"]

323 + elif hasattr(model_hf_config, "pard_token"):

324 self.parallel_drafting_token_id = model_hf_config.pard_token

325 elif hasattr(model_hf_config, "ptd_token_id"):

326 self.parallel_drafting_token_id = model_hf_config.ptd_token_id

327 else:

328 raise ValueError(

329 "For parallel drafting, the draft model config must have "

315 - "`pard_token` or `ptd_token_id` specified in its config.json."

330 + "`pard_token`, `ptd_token_id`, or "

331 + "`dflash_config.mask_token_id` specified in its config.json."

332 )

333

334 if self.pass_hidden_states_to_model:

@@ -402,9 +418,14 @@ class SpecDecodeBaseProposer:

418 ) -> torch.Tensor:

419 batch_size = common_attn_metadata.batch_size()

420

405 - if self.method == "eagle3":

421 + if self.method in ("eagle3", "dflash"):

422 assert isinstance(

407 - self.model, (Eagle3LlamaForCausalLM, Eagle3DeepseekV2ForCausalLM)

423 + self.model,

424 + (

425 + Eagle3LlamaForCausalLM,

426 + Eagle3DeepseekV2ForCausalLM,

427 + DFlashQwen3ForCausalLM,

428 + ),

429 )

430 target_hidden_states = self.model.combine_hidden_states(

431 target_hidden_states

@@ -423,40 +444,17 @@ class SpecDecodeBaseProposer:

444 )

445 )

446

426 - per_layer_attn_metadata: dict[str, object] = {}

427 - for attn_group in self.draft_attn_groups:

428 - attn_metadata = attn_group.get_metadata_builder().build_for_drafting(

429 - common_attn_metadata=common_attn_metadata, draft_index=0

430 - )

431 - for layer_name in attn_group.layer_names:

432 - per_layer_attn_metadata[layer_name] = attn_metadata

447 + per_layer_attn_metadata = self.build_per_layer_attn_metadata(

448 + common_attn_metadata

449 + )

450

451 cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (

452 self._determine_batch_execution_and_padding(num_tokens)

453 )

454

438 - if self.supports_mm_inputs:

439 - mm_embeds, is_mm_embed = mm_embed_inputs or (None, None)

440 -

441 - self.inputs_embeds[:num_tokens] = self.model.embed_input_ids(

442 - self.input_ids[:num_tokens],

443 - multimodal_embeddings=mm_embeds,

444 - is_multimodal=is_mm_embed,

445 - )

446 -

447 - input_ids = None

448 - inputs_embeds = self.inputs_embeds[:num_input_tokens]

449 - else:

450 - input_ids = self.input_ids[:num_input_tokens]

451 - inputs_embeds = None

452 -

453 - model_kwargs = {

454 - "input_ids": input_ids,

455 - "positions": self._get_positions(num_input_tokens),

456 - "inputs_embeds": inputs_embeds,

457 - }

458 - if self.pass_hidden_states_to_model:

459 - model_kwargs["hidden_states"] = self.hidden_states[:num_input_tokens]

455 + model_kwargs, slot_mapping_size = self.build_model_inputs_first_pass(

456 + num_tokens, num_input_tokens, mm_embed_inputs

457 + )

458

459 with set_forward_context(

460 per_layer_attn_metadata,

@@ -465,7 +463,7 @@ class SpecDecodeBaseProposer:

463 num_tokens_across_dp=num_tokens_across_dp,

464 cudagraph_runtime_mode=cudagraph_runtime_mode,

465 slot_mapping=self._get_slot_mapping(

468 - num_input_tokens, common_attn_metadata.slot_mapping

466 + slot_mapping_size, common_attn_metadata.slot_mapping

467 ),

468 ):

469 ret_hidden_states = self.model(**model_kwargs)

@@ -488,7 +486,10 @@ class SpecDecodeBaseProposer:

486 positions = self.positions[token_indices_to_sample]

487 hidden_states = hidden_states[token_indices_to_sample]

488

491 - if isinstance(attn_metadata, TreeAttentionMetadata):

489 + if any(

490 + isinstance(attn_metadata, TreeAttentionMetadata)

491 + for attn_metadata in per_layer_attn_metadata.values()

492 + ):

493 # Draft using tree attention - requires full logits for top-k

494 logits = self.model.compute_logits(sample_hidden_states)

495 draft_token_ids_list = self.propose_tree(

@@ -504,15 +505,16 @@ class SpecDecodeBaseProposer:

505

506 draft_token_ids = self._greedy_sample(sample_hidden_states)

507

507 - if self.allowed_attn_types is not None and not isinstance(

508 - attn_metadata, self.allowed_attn_types

509 - ):

510 - raise ValueError(

511 - f"Unsupported attention metadata type for speculative "

512 - "decoding with num_speculative_tokens > 1: "

513 - f"{type(attn_metadata)}. Supported types are: "

514 - f"{self.allowed_attn_types}"

515 - )

508 + for attn_metadata in per_layer_attn_metadata.values():

509 + if self.allowed_attn_types is not None and not isinstance(

510 + attn_metadata, self.allowed_attn_types

511 + ):

512 + raise ValueError(

513 + f"Unsupported attention metadata type for speculative "

514 + "decoding with num_speculative_tokens > 1: "

515 + f"{type(attn_metadata)}. Supported types are: "

516 + f"{self.allowed_attn_types}"

517 + )

518

519 # Generate the remaining draft tokens.

520 draft_token_ids_list = [draft_token_ids]

@@ -593,13 +595,9 @@ class SpecDecodeBaseProposer:

595 common_attn_metadata._num_computed_tokens_cpu += 1

596

597 # Rebuild attention metadata

596 - for attn_group in self.draft_attn_groups:

597 - attn_metadata = attn_group.get_metadata_builder().build_for_drafting(

598 - common_attn_metadata=common_attn_metadata,

599 - draft_index=token_index + 1,

600 - )

601 - for layer_name in attn_group.layer_names:

602 - per_layer_attn_metadata[layer_name] = attn_metadata

598 + per_layer_attn_metadata = self.build_per_layer_attn_metadata(

599 + common_attn_metadata, draft_index=token_index + 1

600 + )

601

602 # copy inputs to buffer for cudagraph

603 self.input_ids[:batch_size] = input_ids

@@ -780,8 +778,51 @@ class SpecDecodeBaseProposer:

778

779 return total_num_output_tokens, token_indices_to_sample, new_cad

780

781 + def build_model_inputs_first_pass(

782 + self,

783 + num_tokens: int,

784 + num_input_tokens: int,

785 + mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None,

786 + ) -> tuple[dict[str, Any], int]:

787 + if self.supports_mm_inputs:

788 + mm_embeds, is_mm_embed = mm_embed_inputs or (None, None)

789 +

790 + self.inputs_embeds[:num_tokens] = self.model.embed_input_ids(

791 + self.input_ids[:num_tokens],

792 + multimodal_embeddings=mm_embeds,

793 + is_multimodal=is_mm_embed,

794 + )

795 +

796 + input_ids = None

797 + inputs_embeds = self.inputs_embeds[:num_input_tokens]

798 + else:

799 + input_ids = self.input_ids[:num_input_tokens]

800 + inputs_embeds = None

801 +

802 + model_kwargs = {

803 + "input_ids": input_ids,

804 + "positions": self._get_positions(num_input_tokens),

805 + "inputs_embeds": inputs_embeds,

806 + }

807 + if self.pass_hidden_states_to_model:

808 + model_kwargs["hidden_states"] = self.hidden_states[:num_input_tokens]

809 +

810 + return model_kwargs, num_input_tokens

811 +

812 + def build_per_layer_attn_metadata(

813 + self, common_attn_metadata: CommonAttentionMetadata, draft_index: int = 0

814 + ) -> dict[str, object]:

815 + per_layer_attn_metadata: dict[str, object] = {}

816 + for attn_group in self.draft_attn_groups:

817 + attn_metadata = attn_group.get_metadata_builder().build_for_drafting(

818 + common_attn_metadata=common_attn_metadata, draft_index=draft_index

819 + )

820 + for layer_name in attn_group.layer_names:

821 + per_layer_attn_metadata[layer_name] = attn_metadata

822 + return per_layer_attn_metadata

823 +

824 def model_returns_tuple(self) -> bool:

784 - return self.method not in ("mtp", "draft_model")

825 + return self.method not in ("mtp", "draft_model", "dflash")

826

827 def prepare_next_token_ids_cpu(

828 self,

@@ -1310,15 +1351,20 @@ class SpecDecodeBaseProposer:

1351 self._maybe_share_embeddings(target_language_model)

1352 self._maybe_share_lm_head(target_language_model)

1353

1313 - if self.parallel_drafting and self.pass_hidden_states_to_model:

1314 - assert self.parallel_drafting_hidden_state_tensor is not None

1315 - self.parallel_drafting_hidden_state_tensor.copy_(

1316 - self.model.combine_hidden_states(

1317 - self.model.mask_hidden.view(3 * self.hidden_size)

1354 + if (

1355 + self.parallel_drafting

1356 + and self.pass_hidden_states_to_model

1357 + and self.parallel_drafting_hidden_state_tensor is not None

1358 + ):

1359 + flat_mask = self.model.mask_hidden.view(-1)

1360 + if self.eagle3_use_aux_hidden_state:

1361 + # EAGLE3: mask_hidden stores all aux hidden states,

1362 + # project through combine_hidden_states

1363 + self.parallel_drafting_hidden_state_tensor.copy_(

1364 + self.model.combine_hidden_states(flat_mask)

1365 )

1319 - if self.eagle3_use_aux_hidden_state

1320 - else self.model.mask_hidden.view(self.hidden_size)

1321 - )

1366 + else:

1367 + self.parallel_drafting_hidden_state_tensor.copy_(flat_mask)

1368

1369 def _maybe_share_embeddings(self, target_language_model: nn.Module) -> None:

1370 """

@@ -1493,8 +1539,9 @@ class SpecDecodeBaseProposer:

1539 ) -> None:

1540 # FIXME: when using tree-based specdec, adjust number of forward-passes

1541 # according to the depth of the tree.

1542 + only_one_forward_pass = is_graph_capturing or self.parallel_drafting

1543 for fwd_idx in range(

1497 - self.num_speculative_tokens if not is_graph_capturing else 1

1544 + 1 if only_one_forward_pass else self.num_speculative_tokens

1545 ):

1546 if fwd_idx <= 1:

1547 cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (

vllm/v1/spec_decode/utils.py

+108

@@ -441,6 +441,114 @@ def copy_and_expand_eagle_inputs_kernel(

441 )

442

443

444 +@triton.jit

445 +def copy_and_expand_dflash_inputs_kernel(

446 + # Inputs

447 + next_token_ids_ptr, # [num_reqs]

448 + target_positions_ptr, # [num_context]

449 + # Outputs

450 + out_input_ids_ptr, # [num_query_total] (output)

451 + out_context_positions_ptr, # [num_context] (output)

452 + out_query_positions_ptr, # [num_query_total] (output)

453 + out_context_slot_mapping_ptr, # [num_context] (output)

454 + out_query_slot_mapping_ptr, # [num_query_total] (output)

455 + out_token_indices_ptr, # [num_reqs * num_speculative_tokens] (output)

456 + # Block table

457 + block_table_ptr, # [max_reqs, max_blocks]

458 + block_table_stride, # stride of block_table dim 0 (in elements)

459 + # Metadata

460 + query_start_loc_ptr, # [num_reqs + 1]

461 + num_rejected_tokens_ptr, # [num_reqs] or null (0) when not padded

462 + # Scalars

463 + parallel_drafting_token_id, # tl.int32

464 + block_size, # tl.int32

465 + num_query_per_req, # tl.int32

466 + num_speculative_tokens, # tl.int32

467 + total_input_tokens, # tl.int32

468 + BLOCK_SIZE: tl.constexpr,

469 + HAS_NUM_REJECTED: tl.constexpr = False,

470 +):

471 + """

472 + Fused kernel for DFlash first-pass input setup.

473 +

474 + Per request, this kernel:

475 + 1. Copies context positions from target_positions to

476 + out_context_positions.

477 + 2. Computes query positions (last_target_pos + 1 + offset) and writes

478 + them to out_query_positions.

479 + 3. Writes input_ids for query tokens: [next_token, mask, mask, ...].

480 + 4. Computes slot_mapping for context and query positions into separate

481 + buffers via block_table lookup.

482 + 5. Writes token_indices_to_sample for the mask (speculative) tokens.

483 + """

484 + req_idx = tl.program_id(axis=0)

485 + block_idx = tl.program_id(axis=1)

486 +

487 + # Load context token range for this request

488 + ctx_start = tl.load(query_start_loc_ptr + req_idx)

489 + ctx_end = tl.load(query_start_loc_ptr + req_idx + 1)

490 + num_ctx = ctx_end - ctx_start

491 + total_tokens = num_ctx + num_query_per_req

492 +

493 + j = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

494 + in_bounds = j < total_tokens

495 + is_ctx = j < num_ctx

496 + is_query = (~is_ctx) & in_bounds

497 + query_off = j - num_ctx # offset within query portion (0-indexed)

498 +

499 + # --- Positions ---

500 + # Context: load from target_positions

501 + ctx_pos_idx = tl.minimum(ctx_start + j, total_input_tokens - 1)

502 + ctx_pos = tl.load(target_positions_ptr + ctx_pos_idx, mask=is_ctx, other=0)

503 +

504 + # Query: last_valid_pos + 1 + query_off

505 + # In padded mode, ctx_end includes rejected tokens; use valid_ctx_end

506 + # to find the last accepted context position.

507 + if HAS_NUM_REJECTED:

508 + num_rejected = tl.load(num_rejected_tokens_ptr + req_idx)

509 + valid_ctx_end = ctx_end - num_rejected

510 + else:

511 + valid_ctx_end = ctx_end

512 + last_pos = tl.load(target_positions_ptr + valid_ctx_end - 1)

513 + query_pos = last_pos + 1 + query_off

514 +

515 + positions = tl.where(is_ctx, ctx_pos, query_pos)

516 +

517 + # Context and query positions go to separate buffers.

518 + ctx_pos_out = ctx_start + j

519 + tl.store(out_context_positions_ptr + ctx_pos_out, ctx_pos, mask=is_ctx)

520 + query_out = req_idx * num_query_per_req + query_off

521 + tl.store(out_query_positions_ptr + query_out, query_pos, mask=is_query)

522 +

523 + # --- Slot mapping (block_table lookup for all positions) ---

524 + block_num = positions // block_size

525 + # # Clamp block_number to avoid OOB when position is at max

526 + block_num = tl.minimum(block_num, block_table_stride - 1)

527 + block_id = tl.load(

528 + block_table_ptr + req_idx * block_table_stride + block_num,

529 + mask=in_bounds,

530 + other=0,

531 + ).to(tl.int64)

532 + slot = block_id * block_size + (positions % block_size)

533 + tl.store(out_context_slot_mapping_ptr + ctx_pos_out, slot, mask=is_ctx)

534 + tl.store(out_query_slot_mapping_ptr + query_out, slot, mask=is_query)

535 +

536 + # --- Input IDs (query tokens only) ---

537 + bonus_token = tl.load(next_token_ids_ptr + req_idx)

538 + is_bonus = is_query & (query_off == 0)

539 + input_id = tl.where(is_bonus, bonus_token, parallel_drafting_token_id)

540 + tl.store(out_input_ids_ptr + query_out, input_id, mask=is_query)

541 +

542 + # --- Token indices to sample (mask tokens, skip the bonus token) ---

543 + is_sample = is_query & (query_off > 0)

544 + sample_out_idx = req_idx * num_speculative_tokens + (query_off - 1)

545 + tl.store(

546 + out_token_indices_ptr + sample_out_idx,

547 + query_out,

548 + mask=is_sample,

549 + )

550 +

551 +

552 @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)

553 def update_num_computed_tokens_for_batch_change(

554 num_computed_tokens: torch.Tensor,

vllm/v1/worker/gpu_model_runner.py

+35 -10

@@ -160,6 +160,7 @@ from vllm.v1.sample.logits_processor.interface import LogitsProcessor

160 from vllm.v1.sample.metadata import SamplingMetadata

161 from vllm.v1.sample.rejection_sampler import RejectionSampler

162 from vllm.v1.sample.sampler import Sampler

163 +from vllm.v1.spec_decode.dflash import DFlashProposer

164 from vllm.v1.spec_decode.draft_model import DraftModelProposer

165 from vllm.v1.spec_decode.eagle import EagleProposer

166 from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer

@@ -515,6 +516,7 @@ class GPUModelRunner(

516 | NgramProposerGPU

517 | SuffixDecodingProposer

518 | EagleProposer

519 + | DFlashProposer

520 | DraftModelProposer

521 | MedusaProposer

522 | ExtractHiddenStatesProposer

@@ -546,6 +548,9 @@ class GPUModelRunner(

548 self._ngram_pinned_val_buf = torch.zeros(

549 self.max_num_reqs, dtype=torch.int32, pin_memory=True

550 )

551 + elif self.speculative_config.use_dflash():

552 + self.drafter = DFlashProposer(self.vllm_config, self.device, self)

553 + self.use_aux_hidden_state_outputs = True

554 elif self.speculative_config.method == "suffix":

555 self.drafter = SuffixDecodingProposer(self.vllm_config)

556 elif self.speculative_config.use_eagle():

@@ -2289,7 +2294,7 @@ class GPUModelRunner(

2294 cm.slot_mapping = slot_mappings[kv_cache_gid]

2295

2296 if self.speculative_config and spec_decode_common_attn_metadata is None:

2292 - if isinstance(self.drafter, EagleProposer):

2297 + if isinstance(self.drafter, (EagleProposer, DFlashProposer)):

2298 if self.drafter.kv_cache_gid == kv_cache_gid:

2299 spec_decode_common_attn_metadata = cm

2300 else:

@@ -4202,7 +4207,10 @@ class GPUModelRunner(

4207 # as inputs, and does not need to wait for bookkeeping to finish.

4208 assert isinstance(

4209 self.drafter,

4205 - EagleProposer | DraftModelProposer | ExtractHiddenStatesProposer,

4210 + EagleProposer

4211 + | DFlashProposer

4212 + | DraftModelProposer

4213 + | ExtractHiddenStatesProposer,

4214 )

4215 sampled_token_ids = sampler_output.sampled_token_ids

4216 if input_fits_in_drafter:

@@ -4589,8 +4597,14 @@ class GPUModelRunner(

4597 next_token_ids, valid_sampled_tokens_count

4598 )

4599

4592 - elif spec_config.use_eagle() or spec_config.uses_draft_model():

4593 - assert isinstance(self.drafter, EagleProposer | DraftModelProposer)

4600 + elif (

4601 + spec_config.use_eagle()

4602 + or spec_config.use_dflash()

4603 + or spec_config.uses_draft_model()

4604 + ):

4605 + assert isinstance(

4606 + self.drafter, EagleProposer | DFlashProposer | DraftModelProposer

4607 + )

4608

4609 if spec_config.disable_padded_drafter_batch:

4610 # When padded-batch is disabled, the sampled_token_ids should be

@@ -4889,10 +4903,13 @@ class GPUModelRunner(

4903 return None

4904

4905 hf_config = self.speculative_config.draft_model_config.hf_config

4892 - if not hasattr(hf_config, "eagle_aux_hidden_state_layer_ids"):

4893 - return None

4906

4895 - layer_ids = hf_config.eagle_aux_hidden_state_layer_ids

4907 + layer_ids = getattr(hf_config, "eagle_aux_hidden_state_layer_ids", None)

4908 + if not layer_ids:

4909 + dflash_config = getattr(hf_config, "dflash_config", None)

4910 + if dflash_config and isinstance(dflash_config, dict):

4911 + layer_ids = dflash_config.get("target_layer_ids")

4912 +

4913 if layer_ids and isinstance(layer_ids, (list, tuple)):

4914 return tuple(layer_ids)

4915

@@ -5479,7 +5496,10 @@ class GPUModelRunner(

5496 ):

5497 assert isinstance(

5498 self.drafter,

5482 - EagleProposer | DraftModelProposer | ExtractHiddenStatesProposer,

5499 + EagleProposer

5500 + | DFlashProposer

5501 + | DraftModelProposer

5502 + | ExtractHiddenStatesProposer,

5503 )

5504 assert self.speculative_config is not None

5505 # Eagle currently only supports PIECEWISE cudagraphs.

@@ -6236,7 +6256,9 @@ class GPUModelRunner(

6256 self.speculative_config.use_eagle()

6257 or self.speculative_config.uses_draft_model()

6258 ):

6239 - assert isinstance(self.drafter, EagleProposer | DraftModelProposer)

6259 + assert isinstance(

6260 + self.drafter, EagleProposer | DFlashProposer | DraftModelProposer

6261 + )

6262 self.drafter.initialize_attn_backend(kv_cache_config, kernel_block_sizes)

6263

6264 def _check_and_update_cudagraph_mode(

@@ -6420,7 +6442,10 @@ class GPUModelRunner(

6442 self.speculative_config.use_eagle()

6443 or self.speculative_config.uses_extract_hidden_states()

6444 ):

6423 - assert isinstance(self.drafter, EagleProposer | ExtractHiddenStatesProposer)

6445 + assert isinstance(

6446 + self.drafter,

6447 + EagleProposer | DFlashProposer | ExtractHiddenStatesProposer,

6448 + )

6449 self.drafter.initialize_cudagraph_keys(cudagraph_mode)

6450

6451 def calculate_reorder_batch_threshold(self) -> None: