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 11, 15:00 UTC
Add Eagle and Eagle3 support to Transformers modeling backend (#30340)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Dec 11, 2025 at 17:02 UTC
8781cd6b88ad264a01886a05e698b5e036fb4eb9
2 files changed
+94
-8
tests/v1/e2e/test_spec_decode.py
+34
-2
@@ -280,9 +280,20 @@ def test_speculators_model_integration(
280
281
282
@pytest.mark.parametrize(
283
- ["model_setup", "mm_enabled", "enable_chunked_prefill"],
283
+ ["model_setup", "mm_enabled", "enable_chunked_prefill", "model_impl"],
284
[
285
- (("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1), False, False),
285
+ (
286
+ ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
287
+ False,
288
+ False,
289
+ "auto",
290
+ ),
291
+ (
292
+ ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
293
+ False,
294
+ False,
295
+ "transformers",
296
+ ),
297
pytest.param(
298
(
299
"eagle3",
@@ -292,6 +303,7 @@ def test_speculators_model_integration(
303
),
304
False,
305
False,
306
+ "auto",
307
marks=pytest.mark.skip(
308
reason="architecture of its eagle3 is LlamaForCausalLMEagle3"
309
),
@@ -305,6 +317,7 @@ def test_speculators_model_integration(
317
),
318
False,
319
False,
320
+ "auto",
321
marks=pytest.mark.skip(
322
reason="Skipping due to its head_dim not being a a multiple of 32"
323
),
@@ -318,6 +331,7 @@ def test_speculators_model_integration(
331
),
332
False,
333
True,
334
+ "auto",
335
marks=large_gpu_mark(min_gb=40),
336
), # works on 4x H100
337
(
@@ -329,6 +343,7 @@ def test_speculators_model_integration(
343
),
344
False,
345
False,
346
+ "auto",
347
),
348
pytest.param(
349
(
@@ -339,6 +354,7 @@ def test_speculators_model_integration(
354
),
355
False,
356
False,
357
+ "auto",
358
marks=large_gpu_mark(min_gb=80),
359
), # works on 4x H100
360
pytest.param(
@@ -350,6 +366,7 @@ def test_speculators_model_integration(
366
),
367
True,
368
True,
369
+ "auto",
370
marks=large_gpu_mark(min_gb=80),
371
), # works on 4x H100
372
(
@@ -361,10 +378,12 @@ def test_speculators_model_integration(
378
),
379
False,
380
False,
381
+ "auto",
382
),
383
],
384
ids=[
385
"qwen3_eagle3",
386
+ "qwen3_eagle3-transformers",
387
"qwen3_vl_eagle3",
388
"qwen2_5_vl_eagle3",
389
"llama3_eagle",
@@ -381,6 +400,7 @@ def test_eagle_correctness(
400
model_setup: tuple[str, str, str, int],
401
mm_enabled: bool,
402
enable_chunked_prefill: bool,
403
+ model_impl: str,
404
attn_backend: str,
405
):
406
if attn_backend == "TREE_ATTN":
@@ -389,6 +409,17 @@ def test_eagle_correctness(
409
"TREE_ATTN is flaky in the test disable for now until it can be "
410
"resolved (see https://github.com/vllm-project/vllm/issues/22922)"
411
)
412
+ if model_impl == "transformers":
413
+ import transformers
414
+ from packaging.version import Version
415
+
416
+ installed = Version(transformers.__version__)
417
+ required = Version("5.0.0.dev")
418
+ if installed < required:
419
+ pytest.skip(
420
+ "Eagle3 with the Transformers modeling backend requires "
421
+ f"transformers>={required}, but got {installed}"
422
+ )
423
424
# Generate test prompts inside the function instead of using fixture
425
test_prompts = get_test_prompts(mm_enabled)
@@ -448,6 +479,7 @@ def test_eagle_correctness(
479
max_model_len=max_model_len,
480
max_num_batched_tokens=max_num_batched_tokens,
481
enable_chunked_prefill=enable_chunked_prefill,
482
+ model_impl=model_impl,
483
)
484
spec_outputs = spec_llm.chat(test_prompts, sampling_config)
485
matches = 0
vllm/model_executor/models/transformers/base.py
+60
-6
@@ -36,6 +36,8 @@ from vllm.distributed.utils import get_pp_indices
36
from vllm.logger import init_logger
37
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
38
from vllm.model_executor.models.interfaces import (
39
+ SupportsEagle,
40
+ SupportsEagle3,
41
SupportsLoRA,
42
SupportsPP,
43
SupportsQuant,
@@ -92,7 +94,15 @@ def vllm_flash_attention_forward(
94
ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_flash_attention_forward
95
96
95
-class Base(nn.Module, VllmModel, SupportsQuant, SupportsLoRA, SupportsPP):
97
+class Base(
98
+ nn.Module,
99
+ VllmModel,
100
+ SupportsQuant,
101
+ SupportsLoRA,
102
+ SupportsPP,
103
+ SupportsEagle,
104
+ SupportsEagle3,
105
+):
106
embedding_modules = ["embed_tokens"] # TODO transformers will have a util to get it
107
hf_to_vllm_mapper = WeightsMapper(
108
orig_to_new_prefix={
@@ -131,17 +141,24 @@ class Base(nn.Module, VllmModel, SupportsQuant, SupportsLoRA, SupportsPP):
141
self.pp_group = get_pp_group()
142
self.tp_group = get_tp_group()
143
134
- # Weights to skip in `self.load_weights`
144
+ # Attrs for weight loading (see self.load_weights)
145
self.skip_prefixes: list[str] = []
146
"""Skip loading weights whose qualname starts with these prefixes."""
147
self.skip_substrs: list[str] = []
148
"""Skip loading weights whose qualname contains these substrings."""
149
self.ignore_unexpected_prefixes: list[str] = []
140
- """Ignore unexpected weights whose qualname starts with these prefixes.
141
- """
150
+ """Ignore unexpected weights whose qualname starts with these prefixes."""
151
self.ignore_unexpected_suffixes: list[str] = []
152
"""Ignore unexpected weights whose qualname ends with these suffixes."""
153
154
+ # Attrs for Eagle3 (see self.set_aux_hidden_state_layers)
155
+ self._target_class: type[nn.Module] = nn.Module
156
+ """Target class for Eagle3 aux hidden state recording."""
157
+ self._layer_names: dict[int, str] = {}
158
+ """Mapping from layer index to layer name for Eagle3."""
159
+ self._output_aux_hidden_states_kwargs: dict[str, bool] = {}
160
+ """Kwargs to pass to model forward for Eagle3 aux hidden states."""
161
+
162
if self.quant_config:
163
quant_method_name = self.quant_config.get_name()
164
# Check for unsupported quantization methods.
@@ -278,6 +295,15 @@ class Base(nn.Module, VllmModel, SupportsQuant, SupportsLoRA, SupportsPP):
295
for child_name, child_module in module.named_children():
296
new_module = child_module
297
qual_name = maybe_prefix(prefix, child_name)
298
+ # Populate Eagle3 attrs
299
+ if (
300
+ isinstance(module, nn.ModuleList)
301
+ and len(module) == self.text_config.num_hidden_layers
302
+ ):
303
+ self._target_class = type(child_module)
304
+ layer_name = qual_name.removeprefix("model.")
305
+ self._layer_names[int(child_name)] = layer_name
306
+ # Replace modules as needed
307
if isinstance(child_module, nn.Linear):
308
generator = (p for p in tp_plan if re.match(p, qual_name))
309
pattern = next(generator, None)
@@ -425,19 +451,26 @@ class Base(nn.Module, VllmModel, SupportsQuant, SupportsLoRA, SupportsPP):
451
else:
452
position_ids = positions[None, ...]
453
428
- hidden_states = self.model(
454
+ outputs = self.model(
455
input_ids=input_ids,
456
inputs_embeds=inputs_embeds,
457
use_cache=False,
458
position_ids=position_ids,
459
attention_instances=self.attention_instances,
460
return_dict=False,
461
+ **self._output_aux_hidden_states_kwargs,
462
**kwargs,
436
- )[0][0, ...] # we remove batch dimension for now
463
+ )
464
+ # We must remove the batch dimension from these outputs
465
+ hidden_states = outputs[0][0, ...]
466
+ if self._output_aux_hidden_states_kwargs:
467
+ aux_hidden_states = [x[0][0, ...] for x in outputs[1:]]
468
469
if not self.pp_group.is_last_rank:
470
return IntermediateTensors({"hidden_states": hidden_states})
471
472
+ if self._output_aux_hidden_states_kwargs and len(aux_hidden_states) > 0:
473
+ return hidden_states, aux_hidden_states
474
return hidden_states
475
476
def load_weights(
@@ -462,3 +495,24 @@ class Base(nn.Module, VllmModel, SupportsQuant, SupportsLoRA, SupportsPP):
495
f"Transformers modeling backend requires transformers>={required} "
496
f"for {feature}, but got {installed}"
497
)
498
+
499
+ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
500
+ self.check_version("5.0.0.dev0", "Eagle3 support")
501
+ from transformers.utils.generic import OutputRecorder
502
+
503
+ # The default value in PreTrainedModel is None
504
+ if self.model._can_record_outputs is None:
505
+ self.model._can_record_outputs = {}
506
+
507
+ target_class = self._target_class
508
+ for layer in layers:
509
+ # layer - 1 because we want the input to the layer
510
+ layer_name = self._layer_names[layer - 1]
511
+ layer_key = f"aux_hidden_state_{layer}"
512
+ aux_hidden_state_i = OutputRecorder(target_class, layer_name=layer_name)
513
+ self.model._can_record_outputs[layer_key] = aux_hidden_state_i
514
+ self._output_aux_hidden_states_kwargs[f"output_{layer_key}"] = True
515
+
516
+ def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
517
+ num_layers = self.text_config.num_hidden_layers
518
+ return (2, num_layers // 2, num_layers - 3)