[Spec Decode] Add Qwen3 architecture support for EAGLE3 (#43132)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Benjamin Chislett committed
Jun 21, 2026 at 16:55 UTC
89bd2c14d39075a6109ff188b896b600732dc348
6 files changed
+523
-2
tests/models/registry.py
+19
@@ -1530,6 +1530,16 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1530
"Qwen/Qwen3-VL-8B-Instruct",
1531
speculative_model="taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3",
1532
),
1533
+ "Eagle3Qwen3ForCausalLM": _HfExamplesInfo(
1534
+ "Qwen/Qwen3-8B",
1535
+ trust_remote_code=True,
1536
+ speculative_model=(
1537
+ "inference-optimization/"
1538
+ "Qwen3-8B-from-Qwen3-8B_regen-speculators.eagle3-qwen3arch-ckpt1"
1539
+ ),
1540
+ tokenizer="Qwen/Qwen3-8B",
1541
+ use_original_num_layers=True,
1542
+ ),
1543
# [PEagle]
1544
"PEagleDraftModel": _HfExamplesInfo(
1545
"Qwen/Qwen3-8B",
@@ -1545,6 +1555,15 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1555
tokenizer="Qwen/Qwen3-8B",
1556
use_original_num_layers=True,
1557
),
1558
+ "PeagleQwen3ForCausalLM": _HfExamplesInfo(
1559
+ "Qwen/Qwen3-8B",
1560
+ trust_remote_code=True,
1561
+ speculative_model=(
1562
+ "inference-optimization/Qwen3-8B-speculators.peagle-qwen3arch-ckpt4"
1563
+ ),
1564
+ tokenizer="Qwen/Qwen3-8B",
1565
+ use_original_num_layers=True,
1566
+ ),
1567
# [MTP]
1568
"DeepSeekMTPModel": _HfExamplesInfo(
1569
"luccafong/deepseek_mtp_main_random",
tests/v1/spec_decode/test_speculators_correctness.py
+31
@@ -53,9 +53,39 @@ PEAGLE_CONFIG = SpeculatorTestConfig(
53
parallel_drafting=True,
54
)
55
56
+QWEN3_EAGLE3_CONFIG = SpeculatorTestConfig(
57
+ model_path=(
58
+ "inference-optimization/"
59
+ "Qwen3-8B-from-Qwen3-8B_regen-speculators.eagle3-qwen3arch-ckpt1"
60
+ ),
61
+ method="eagle3",
62
+ display_name="Qwen3 Eagle3",
63
+ expected_gsm8k_accuracy=0.88,
64
+ accuracy_rtol=0.05,
65
+ expected_acceptance_len=2.67,
66
+ acceptance_len_rtol=0.10,
67
+ expected_per_pos_acceptance_rates=(0.76, 0.55, 0.36),
68
+ per_pos_rtol=0.10,
69
+)
70
+
71
+QWEN3_PEAGLE_CONFIG = SpeculatorTestConfig(
72
+ model_path="inference-optimization/Qwen3-8B-speculators.peagle-qwen3arch-ckpt4",
73
+ method="eagle3",
74
+ display_name="Qwen3 PEagle",
75
+ expected_gsm8k_accuracy=0.88,
76
+ accuracy_rtol=0.05,
77
+ expected_acceptance_len=3.42,
78
+ acceptance_len_rtol=0.15,
79
+ expected_per_pos_acceptance_rates=(0.78, 0.59, 0.43, 0.29, 0.18, 0.10, 0.05),
80
+ per_pos_rtol=0.10,
81
+ parallel_drafting=True,
82
+)
83
+
84
SPECULATOR_CONFIGS = [
85
pytest.param(DFLASH_CONFIG, id="dflash"),
86
pytest.param(PEAGLE_CONFIG, id="peagle"),
87
+ pytest.param(QWEN3_EAGLE3_CONFIG, id="qwen3arch_eagle3"),
88
+ pytest.param(QWEN3_PEAGLE_CONFIG, id="qwen3arch_peagle"),
89
]
90
91
@@ -176,6 +206,7 @@ def test_speculators_correctness(monkeypatch, config):
206
207
results = evaluate_gsm8k_offline(spec_llm)
208
accuracy = results["accuracy"]
209
+ print(f"GSM8K Accuracy: {accuracy:.4f}")
210
accuracy_threshold = config.expected_gsm8k_accuracy * (1 - config.accuracy_rtol)
211
assert accuracy >= accuracy_threshold, (
212
f"Expected GSM8K accuracy >= {accuracy_threshold:.3f}, got {accuracy:.3f}"
vllm/model_executor/models/qwen3_eagle3.py
new
+453
@@ -0,0 +1,453 @@
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 as nn
8
+from transformers import Qwen3Config
9
+
10
+from vllm.compilation.decorators import support_torch_compile
11
+from vllm.config import VllmConfig, get_current_vllm_config
12
+from vllm.logger import init_logger
13
+from vllm.model_executor.layers.layernorm import RMSNorm
14
+from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear
15
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
16
+from vllm.model_executor.layers.vocab_parallel_embedding import (
17
+ ParallelLMHead,
18
+ VocabParallelEmbedding,
19
+)
20
+from vllm.model_executor.model_loader.weight_utils import (
21
+ default_weight_loader,
22
+ maybe_remap_kv_scale_name,
23
+)
24
+from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer, Qwen3ForCausalLM
25
+from vllm.multimodal.inputs import NestedTensors
26
+
27
+from .utils import (
28
+ AutoWeightsLoader,
29
+ get_draft_quant_config,
30
+ maybe_prefix,
31
+ process_eagle_weight,
32
+)
33
+
34
+logger = init_logger(__name__)
35
+
36
+
37
+class Qwen3Eagle3DecoderLayer(Qwen3DecoderLayer):
38
+ def __init__(
39
+ self,
40
+ vllm_config: VllmConfig,
41
+ prefix: str = "",
42
+ config: Qwen3Config | None = None,
43
+ layer_idx: int = 0,
44
+ ) -> None:
45
+ config = config or vllm_config.model_config.hf_config
46
+ cache_config = vllm_config.cache_config
47
+ quant_config = get_draft_quant_config(vllm_config)
48
+
49
+ super().__init__(
50
+ config=config,
51
+ cache_config=cache_config,
52
+ quant_config=quant_config,
53
+ prefix=prefix,
54
+ )
55
+
56
+ # First layer uses 2*hidden_size (embeds + hidden_states concatenated)
57
+ # Subsequent layers use hidden_size (only hidden_states, no embeds)
58
+ qkv_input_size = 2 * self.hidden_size if layer_idx == 0 else self.hidden_size
59
+
60
+ # Parallel drafting checkpoints may have attention bias enabled
61
+ qkv_bias = getattr(config, "attention_bias", False)
62
+
63
+ # Override qkv_proj with correct input size and bias setting
64
+ self.self_attn.qkv_proj = QKVParallelLinear(
65
+ qkv_input_size,
66
+ self.self_attn.head_dim,
67
+ self.self_attn.total_num_heads,
68
+ self.self_attn.total_num_kv_heads,
69
+ bias=qkv_bias,
70
+ quant_config=quant_config,
71
+ prefix=maybe_prefix(prefix, "qkv_proj"),
72
+ )
73
+
74
+ self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
75
+ self.layer_idx = layer_idx
76
+
77
+ if getattr(config, "norm_before_residual", False):
78
+ self._residual_norm = self._norm_before_residual
79
+ else:
80
+ self._residual_norm = self._norm_after_residual
81
+
82
+ def _norm_before_residual(
83
+ self, hidden_states: torch.Tensor
84
+ ) -> tuple[torch.Tensor, torch.Tensor]:
85
+ hidden_states = self.hidden_norm(hidden_states)
86
+ residual = hidden_states
87
+ return hidden_states, residual
88
+
89
+ def _norm_after_residual(
90
+ self, hidden_states: torch.Tensor
91
+ ) -> tuple[torch.Tensor, torch.Tensor]:
92
+ residual = hidden_states
93
+ hidden_states = self.hidden_norm(hidden_states)
94
+ return hidden_states, residual
95
+
96
+ def forward(
97
+ self,
98
+ positions: torch.Tensor,
99
+ embeds: torch.Tensor,
100
+ hidden_states: torch.Tensor,
101
+ residual: torch.Tensor | None,
102
+ ) -> tuple[torch.Tensor, torch.Tensor]:
103
+ if self.layer_idx == 0:
104
+ # First layer: concatenate embeds with hidden_states
105
+ embeds = self.input_layernorm(embeds)
106
+ hidden_states, residual = self._residual_norm(hidden_states=hidden_states)
107
+ hidden_states = torch.cat([embeds, hidden_states], dim=-1)
108
+ else:
109
+ # Subsequent layers: process hidden_states and residuals only
110
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
111
+
112
+ # Self Attention
113
+ hidden_states = self.self_attn(
114
+ positions=positions,
115
+ hidden_states=hidden_states,
116
+ )
117
+
118
+ hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
119
+
120
+ # Fully Connected
121
+ hidden_states = self.mlp(hidden_states)
122
+
123
+ return hidden_states, residual
124
+
125
+
126
+@support_torch_compile(
127
+ dynamic_arg_dims={
128
+ "input_ids": 0,
129
+ "positions": -1,
130
+ "hidden_states": 0,
131
+ "input_embeds": 0,
132
+ }
133
+)
134
+class Qwen3Eagle3Model(nn.Module):
135
+ def __init__(
136
+ self,
137
+ *,
138
+ vllm_config: VllmConfig,
139
+ start_layer_id: int = 0,
140
+ prefix: str = "",
141
+ ) -> None:
142
+ super().__init__()
143
+ self.config = vllm_config.speculative_config.draft_model_config.hf_config
144
+ self.vocab_size = self.config.vocab_size
145
+
146
+ # Get drafter's quantization config
147
+ self.quant_config = get_draft_quant_config(vllm_config)
148
+
149
+ eagle_config = getattr(self.config, "eagle_config", None) or {}
150
+ if "use_aux_hidden_state" in eagle_config:
151
+ self.use_aux_hidden_state = eagle_config["use_aux_hidden_state"]
152
+ else:
153
+ self.use_aux_hidden_state = True
154
+ self.norm_before_fc = bool(
155
+ eagle_config.get(
156
+ "norm_before_fc", getattr(self.config, "norm_before_fc", False)
157
+ )
158
+ )
159
+ self.fc_input_size = self.config.hidden_size
160
+
161
+ current_vllm_config = get_current_vllm_config()
162
+
163
+ self.embed_tokens = VocabParallelEmbedding(
164
+ self.config.vocab_size,
165
+ self.config.hidden_size,
166
+ prefix=maybe_prefix(prefix, "embed_tokens"),
167
+ )
168
+
169
+ self.layers = nn.ModuleList(
170
+ [
171
+ Qwen3Eagle3DecoderLayer(
172
+ current_vllm_config,
173
+ prefix=maybe_prefix(prefix, f"layers.{layer_idx + start_layer_id}"),
174
+ config=self.config,
175
+ layer_idx=layer_idx,
176
+ )
177
+ for layer_idx in range(self.config.num_hidden_layers)
178
+ ]
179
+ )
180
+ if self.use_aux_hidden_state:
181
+ num_aux_features = getattr(self.config, "num_aux_layers", None)
182
+ if num_aux_features is None:
183
+ num_aux_features = getattr(self.config, "num_aux_hidden_states", None)
184
+ if num_aux_features is None:
185
+ aux_ids = getattr(
186
+ self.config, "eagle_aux_hidden_state_layer_ids", None
187
+ ) or eagle_config.get("eagle_aux_hidden_state_layer_ids")
188
+ num_aux_features = len(aux_ids) if aux_ids else 3
189
+ self.num_aux_layers = num_aux_features
190
+ target_hidden_size = getattr(
191
+ self.config, "target_hidden_size", self.config.hidden_size
192
+ )
193
+ self.fc_input_size = target_hidden_size * num_aux_features
194
+ if self.norm_before_fc:
195
+ self.input_norm = RMSNorm(
196
+ self.fc_input_size,
197
+ eps=self.config.rms_norm_eps,
198
+ )
199
+ else:
200
+ self.input_norm = None
201
+
202
+ use_fc_norm = getattr(self.config, "fc_norm", False)
203
+ if use_fc_norm:
204
+ self.fc_norm = nn.ModuleList(
205
+ [
206
+ RMSNorm(target_hidden_size, eps=self.config.rms_norm_eps)
207
+ for _ in range(num_aux_features)
208
+ ]
209
+ )
210
+ else:
211
+ self.fc_norm = None
212
+
213
+ self.fc = ReplicatedLinear(
214
+ input_size=self.fc_input_size,
215
+ output_size=self.config.hidden_size,
216
+ bias=False,
217
+ params_dtype=vllm_config.model_config.dtype,
218
+ quant_config=self.quant_config,
219
+ prefix=maybe_prefix(prefix, "fc"),
220
+ return_bias=False,
221
+ )
222
+
223
+ self.norm_output = getattr(self.config, "norm_output", False)
224
+ self.norm = RMSNorm(
225
+ self.config.hidden_size,
226
+ eps=self.config.rms_norm_eps,
227
+ )
228
+
229
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
230
+ return self.embed_tokens(input_ids)
231
+
232
+ def forward(
233
+ self,
234
+ input_ids: torch.Tensor,
235
+ positions: torch.Tensor,
236
+ hidden_states: torch.Tensor,
237
+ input_embeds: torch.Tensor | None = None,
238
+ ) -> tuple[torch.Tensor, torch.Tensor]:
239
+ if input_embeds is None:
240
+ input_embeds = self.embed_input_ids(input_ids)
241
+ assert hidden_states.shape[-1] == input_embeds.shape[-1]
242
+
243
+ residual = None
244
+ for layer in self.layers:
245
+ hidden_states, residual = layer(
246
+ positions=positions,
247
+ embeds=input_embeds,
248
+ hidden_states=hidden_states,
249
+ residual=residual,
250
+ )
251
+ hidden_states, hidden_prenorm = self.norm(hidden_states, residual)
252
+
253
+ # norm_output variant uses the post-norm hidden states.
254
+ aux_output = hidden_states if self.norm_output else hidden_prenorm
255
+
256
+ return hidden_states, aux_output
257
+
258
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
259
+ stacked_params_mapping = [
260
+ # (param_name, shard_name, shard_id)
261
+ (".qkv_proj", ".q_proj", "q"),
262
+ (".qkv_proj", ".k_proj", "k"),
263
+ (".qkv_proj", ".v_proj", "v"),
264
+ (".gate_up_proj", ".gate_proj", 0),
265
+ (".gate_up_proj", ".up_proj", 1),
266
+ ]
267
+ params_dict = dict(self.named_parameters())
268
+ loaded_params: set[str] = set()
269
+ for name, loaded_weight in weights:
270
+ if "midlayer." in name:
271
+ name = name.replace("midlayer.", "layers.0.")
272
+ # Remapping the name FP8 kv-scale or zero point.
273
+ if "scale" in name or "zero_point" in name:
274
+ name = maybe_remap_kv_scale_name(name, params_dict)
275
+ if name is None:
276
+ continue
277
+ for param_name, weight_name, shard_id in stacked_params_mapping:
278
+ if weight_name not in name:
279
+ continue
280
+ name = name.replace(weight_name, param_name)
281
+ param = params_dict[name]
282
+ weight_loader = param.weight_loader
283
+ weight_loader(param, loaded_weight, shard_id)
284
+ break
285
+ else:
286
+ param = params_dict[name]
287
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
288
+ weight_loader(param, loaded_weight)
289
+ loaded_params.add(name)
290
+ return loaded_params
291
+
292
+
293
+class Eagle3Qwen3ForCausalLM(Qwen3ForCausalLM):
294
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
295
+ nn.Module.__init__(self)
296
+ self.config = vllm_config.speculative_config.draft_model_config.hf_config
297
+ # Ensure draft_vocab_size is set
298
+ # default to the base vocab size when absent
299
+ if getattr(self.config, "draft_vocab_size", None) is None:
300
+ base_vocab_size = getattr(self.config, "vocab_size", None)
301
+ self.config.draft_vocab_size = base_vocab_size
302
+ target_layer_num = vllm_config.model_config.get_num_layers(
303
+ vllm_config.parallel_config
304
+ )
305
+
306
+ # Store target layer count in draft config for
307
+ # proper layer_types indexing in draft models
308
+ self.config.target_layer_count = target_layer_num
309
+ self.model = Qwen3Eagle3Model(
310
+ vllm_config=vllm_config,
311
+ prefix=maybe_prefix(prefix, "model"),
312
+ start_layer_id=target_layer_num,
313
+ )
314
+
315
+ logit_scale = getattr(self.config, "logit_scale", 1.0)
316
+ self.lm_head = ParallelLMHead(
317
+ self.config.draft_vocab_size,
318
+ self.config.hidden_size,
319
+ quant_config=get_draft_quant_config(vllm_config),
320
+ prefix=maybe_prefix(prefix, "lm_head"),
321
+ )
322
+ self.logits_processor = LogitsProcessor(
323
+ self.config.draft_vocab_size, scale=logit_scale
324
+ )
325
+ self.draft_id_to_target_id = nn.Parameter(
326
+ torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
327
+ requires_grad=False,
328
+ )
329
+
330
+ self.use_parallel_drafting = vllm_config.speculative_config.parallel_drafting
331
+
332
+ if self.use_parallel_drafting:
333
+ self.register_buffer(
334
+ "mask_hidden",
335
+ torch.zeros(1, self.model.fc_input_size),
336
+ persistent=False,
337
+ )
338
+
339
+ def embed_input_ids(
340
+ self,
341
+ input_ids: torch.Tensor,
342
+ multimodal_embeddings: NestedTensors | None = None,
343
+ is_multimodal: torch.Tensor | None = None,
344
+ ) -> torch.Tensor:
345
+ return self.model.embed_input_ids(input_ids)
346
+
347
+ def forward(
348
+ self,
349
+ input_ids: torch.Tensor,
350
+ positions: torch.Tensor,
351
+ hidden_states: torch.Tensor,
352
+ inputs_embeds: torch.Tensor | None = None,
353
+ ) -> tuple[torch.Tensor, torch.Tensor]:
354
+ return self.model(input_ids, positions, hidden_states, inputs_embeds)
355
+
356
+ def compute_logits(
357
+ self,
358
+ hidden_states: torch.Tensor,
359
+ ) -> torch.Tensor | None:
360
+ logits = self.logits_processor(self.lm_head, hidden_states)
361
+ if self.draft_id_to_target_id is None:
362
+ assert logits.shape[1] == self.config.vocab_size, (
363
+ "Expected logits to have shape "
364
+ f"(*, {self.config.vocab_size}), but got {logits.shape}"
365
+ )
366
+ return logits
367
+
368
+ base = torch.arange(self.config.draft_vocab_size, device=logits.device)
369
+ targets = base + self.draft_id_to_target_id
370
+ logits_new = logits.new_full(
371
+ (
372
+ logits.shape[0],
373
+ self.config.vocab_size,
374
+ ),
375
+ float("-inf"),
376
+ )
377
+ logits_new[:, targets] = logits
378
+ return logits_new
379
+
380
+ def combine_hidden_states(
381
+ self,
382
+ hidden_states: torch.Tensor,
383
+ ) -> torch.Tensor:
384
+ if not self.model.use_aux_hidden_state:
385
+ return hidden_states
386
+ # combine multiple auxiliary hidden states returned by eagle3
387
+
388
+ if self.model.norm_before_fc:
389
+ hidden_states = self.model.input_norm(hidden_states)
390
+
391
+ # `norm_before_fc` adds a single RMSNorm before the FC layer, whereas `fc_norm`
392
+ # applies separate RMSNorms to each chunk of the hidden states.
393
+ if self.model.fc_norm is not None:
394
+ chunks = hidden_states.chunk(self.model.num_aux_layers, dim=-1)
395
+ hidden_states = torch.cat(
396
+ [norm(chunk) for norm, chunk in zip(self.model.fc_norm, chunks)],
397
+ dim=-1,
398
+ )
399
+
400
+ return self.model.fc(hidden_states)
401
+
402
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
403
+ model_weights = {}
404
+ includes_draft_id_mapping = False
405
+ includes_embed_tokens = False
406
+ includes_mask_hidden = False
407
+ for name, loaded_weight in weights:
408
+ if "t2d" in name:
409
+ continue
410
+ if "d2t" in name:
411
+ name = name.replace("d2t", "draft_id_to_target_id")
412
+ includes_draft_id_mapping = True
413
+ elif "mask_hidden" in name:
414
+ # Load mask_hidden directly into buffer
415
+ if not self.use_parallel_drafting:
416
+ logger.warning(
417
+ "mask_hidden found in weights but "
418
+ "model is not configured for parallel drafting. "
419
+ "Skipping loading mask_hidden."
420
+ )
421
+ continue
422
+ self.mask_hidden.copy_(loaded_weight.view(1, -1))
423
+ includes_mask_hidden = True
424
+ continue
425
+ elif "lm_head" not in name:
426
+ name = "model." + name
427
+ if "embed_tokens" in name:
428
+ includes_embed_tokens = True
429
+ model_weights[name] = loaded_weight
430
+ process_eagle_weight(self, name)
431
+
432
+ if not includes_mask_hidden and self.use_parallel_drafting:
433
+ raise ValueError(
434
+ "mask_hidden not found in weights but "
435
+ "model is configured for parallel drafting. "
436
+ "Please provide mask_hidden in the weights."
437
+ )
438
+
439
+ skip_substrs = ["mask_hidden"]
440
+ if not includes_draft_id_mapping:
441
+ skip_substrs.append("draft_id_to_target_id")
442
+ if not includes_embed_tokens:
443
+ skip_substrs.append("embed_tokens")
444
+ if not self.model.use_aux_hidden_state:
445
+ skip_substrs.append("fc.")
446
+ if not self.model.norm_before_fc:
447
+ skip_substrs.append("input_norm.")
448
+ loader = AutoWeightsLoader(
449
+ self,
450
+ skip_prefixes=None,
451
+ skip_substrs=skip_substrs,
452
+ )
453
+ loader.load_weights(model_weights.items())
vllm/model_executor/models/registry.py
+2
@@ -617,6 +617,8 @@ _SPECULATIVE_DECODING_MODELS = {
617
"LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
618
"Eagle3Qwen2_5vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
619
"Eagle3Qwen3vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
620
+ "Eagle3Qwen3ForCausalLM": ("qwen3_eagle3", "Eagle3Qwen3ForCausalLM"),
621
+ "PeagleQwen3ForCausalLM": ("qwen3_eagle3", "Eagle3Qwen3ForCausalLM"),
622
"EagleMistralForCausalLM": ("mistral_eagle", "EagleMistralForCausalLM"),
623
"EagleMistralLarge3ForCausalLM": (
624
"mistral_large_3_eagle",
vllm/transformers_utils/configs/speculators/algos.py
+16
-2
@@ -36,7 +36,14 @@ def update_eagle3(config_dict: dict, pre_trained_config: dict) -> None:
36
"norm_before_residual", True
37
)
38
pre_trained_config["norm_before_fc"] = config_dict.get("norm_before_fc", False)
39
- pre_trained_config["architectures"] = ["Eagle3LlamaForCausalLM"]
39
+ eagle3_arch_map = {
40
+ "qwen3": "Eagle3Qwen3ForCausalLM",
41
+ "llama": "Eagle3LlamaForCausalLM",
42
+ }
43
+ model_type = pre_trained_config.get("model_type", "llama")
44
+ if model_type not in eagle3_arch_map:
45
+ raise ValueError(f"Unsupported model_type {model_type} for Eagle3 speculator")
46
+ pre_trained_config["architectures"] = [eagle3_arch_map[model_type]]
47
if config_dict.get("eagle_aux_hidden_state_layer_ids"):
48
pre_trained_config["eagle_aux_hidden_state_layer_ids"] = config_dict[
49
"eagle_aux_hidden_state_layer_ids"
@@ -59,7 +66,6 @@ def update_peagle(config_dict: dict, pre_trained_config: dict) -> None:
66
- eagle_aux_hidden_state_layer_ids: Layer indices from the target model
67
whose intermediate hidden states are used as auxiliary inputs
68
"""
62
- pre_trained_config["architectures"] = ["PeagleLlamaForCausalLM"]
69
pre_trained_config["draft_vocab_size"] = config_dict.get("draft_vocab_size")
70
if config_dict.get("target_hidden_size") is not None:
71
pre_trained_config["target_hidden_size"] = config_dict["target_hidden_size"]
@@ -67,6 +73,14 @@ def update_peagle(config_dict: dict, pre_trained_config: dict) -> None:
73
"norm_before_residual", False
74
)
75
pre_trained_config["norm_before_fc"] = config_dict.get("norm_before_fc", False)
76
+ peagle_arch_map = {
77
+ "qwen3": "PeagleQwen3ForCausalLM",
78
+ "llama": "PeagleLlamaForCausalLM",
79
+ }
80
+ model_type = pre_trained_config.get("model_type", "llama")
81
+ if model_type not in peagle_arch_map:
82
+ raise ValueError(f"Unsupported model_type {model_type} for PEagle speculator")
83
+ pre_trained_config["architectures"] = [peagle_arch_map[model_type]]
84
pre_trained_config["pard_token"] = config_dict["mask_token_id"]
85
if config_dict.get("eagle_aux_hidden_state_layer_ids"):
86
pre_trained_config["eagle_aux_hidden_state_layer_ids"] = config_dict[
vllm/v1/spec_decode/llm_base_proposer.py
+2
@@ -24,6 +24,7 @@ from vllm.model_executor.models.deepseek_eagle3 import Eagle3DeepseekV2ForCausal
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.model_executor.models.qwen3_eagle3 import Eagle3Qwen3ForCausalLM
28
from vllm.multimodal import MULTIMODAL_REGISTRY
29
from vllm.platforms import current_platform
30
from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d
@@ -473,6 +474,7 @@ class SpecDecodeBaseProposer:
474
Eagle3LlamaForCausalLM,
475
Eagle3DeepseekV2ForCausalLM,
476
DFlashQwen3ForCausalLM,
477
+ Eagle3Qwen3ForCausalLM,
478
),
479
)
480
target_hidden_states = self.model.combine_hidden_states(