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 14, 22:35 UTC
[Speculators][Speculative Decoding] Add Qwen Eagle3 Support (#21835)
Signed-off-by: Dipika Sikka <dipikasikka1@gmail.com>
Dipika Sikka committed
Aug 1, 2025 at 22:43 UTC
9f9c38c392476fd35b9154221c00a2255dcfd010
4 files changed
+46
-11
tests/speculative_decoding/speculators/test_eagle3.py
+12
-2
@@ -6,11 +6,21 @@ import torch
6
7
@pytest.mark.parametrize(
8
"model_path",
9
- [("nm-testing/SpeculatorLlama3-1-8B-Eagle3-converted-0717"),
10
- ("nm-testing/SpeculatorLlama3-1-8B-Eagle3-converted-0717-quantized")])
9
+ [("nm-testing/SpeculatorLlama3-1-8B-Eagle3-converted-0717-quantized")])
10
def test_llama(vllm_runner, example_prompts, model_path):
11
with vllm_runner(model_path, dtype=torch.bfloat16) as vllm_model:
12
vllm_outputs = vllm_model.generate_greedy(example_prompts,
13
max_tokens=20)
14
print(vllm_outputs)
15
assert vllm_outputs
16
+
17
+
18
+@pytest.mark.parametrize(
19
+ "model_path",
20
+ [("nm-testing/Speculator-Qwen3-8B-Eagle3-converted-071-quantized")])
21
+def test_qwen(vllm_runner, example_prompts, model_path):
22
+ with vllm_runner(model_path, dtype=torch.bfloat16) as vllm_model:
23
+ vllm_outputs = vllm_model.generate_greedy(example_prompts,
24
+ max_tokens=20)
25
+ print(vllm_outputs)
26
+ assert vllm_outputs
vllm/config.py
+12
-3
@@ -3175,10 +3175,19 @@ class SpeculativeConfig:
3175
"speculative decoding is > 1, but got "
3176
f"{self.disable_by_batch_size=}")
3177
3178
- if self.method == "eagle3" and self.target_model_config and \
3179
- "llama" not in self.target_model_config.hf_text_config.model_type:
3178
+ from vllm.transformers_utils.configs import SpeculatorsConfig
3179
+
3180
+ eagle3_target_supported = ["llama"]
3181
+ if self.draft_model_config and isinstance(
3182
+ self.draft_model_config.hf_config, SpeculatorsConfig):
3183
+ eagle3_target_supported.append("qwen")
3184
+
3185
+ if self.method == "eagle3" and self.target_model_config and not any(
3186
+ supported_model in
3187
+ self.target_model_config.hf_text_config.model_type
3188
+ for supported_model in eagle3_target_supported):
3189
raise ValueError(
3181
- "Eagle3 is only supported for Llama models. "
3190
+ f"Eagle3 is only supported for {eagle3_target_supported} models. " # noqa: E501
3191
f"Got {self.target_model_config.hf_text_config.model_type=}")
3192
3193
return self
vllm/model_executor/models/qwen2.py
+15
-6
@@ -330,6 +330,8 @@ class Qwen2Model(nn.Module):
330
else:
331
self.norm = PPMissingLayer()
332
333
+ self.aux_hidden_state_layers: tuple[int] = tuple()
334
+
335
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
336
return self.embed_tokens(input_ids)
337
@@ -350,18 +352,25 @@ class Qwen2Model(nn.Module):
352
assert intermediate_tensors is not None
353
hidden_states = intermediate_tensors["hidden_states"]
354
residual = intermediate_tensors["residual"]
353
- for layer in self.layers[self.start_layer:self.end_layer]:
354
- hidden_states, residual = layer(
355
- positions,
356
- hidden_states,
357
- residual,
358
- )
355
+
356
+ aux_hidden_states = []
357
+ for idx, layer in enumerate(
358
+ self.layers[self.start_layer:self.end_layer]):
359
+ if idx in self.aux_hidden_state_layers:
360
+ aux_hidden_states.append(hidden_states + residual)
361
+ hidden_states, residual = layer(positions, hidden_states, residual)
362
+
363
if not get_pp_group().is_last_rank:
364
return IntermediateTensors({
365
"hidden_states": hidden_states,
366
"residual": residual
367
})
368
+
369
hidden_states, _ = self.norm(hidden_states, residual)
370
+
371
+ if len(aux_hidden_states) > 0:
372
+ return hidden_states, aux_hidden_states
373
+
374
return hidden_states
375
376
def load_weights(self, weights: Iterable[tuple[str,
vllm/model_executor/models/qwen3.py
+7
@@ -288,6 +288,13 @@ class Qwen3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
288
self.make_empty_intermediate_tensors = (
289
self.model.make_empty_intermediate_tensors)
290
291
+ def set_aux_hidden_state_layers(self, layers: tuple[int]) -> None:
292
+ self.model.aux_hidden_state_layers = layers
293
+
294
+ def get_eagle3_aux_hidden_state_layers(self) -> tuple[int]:
295
+ num_layers = len(self.model.layers)
296
+ return (2, num_layers // 2, num_layers - 3)
297
+
298
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
299
return self.model.get_input_embeddings(input_ids)
300