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, 05:30 UTC
Add: `SupportsEagle3` interface for explicit EAGLE3 support (#22642)
Signed-off-by: Rahul Tuli <rtuli@redhat.com>
Rahul Tuli committed
Aug 12, 2025 at 21:54 UTC
5a4b4b3729e1a1594bf56d38b7c8d3f556754634
5 files changed
+81
-8
tests/speculative_decoding/speculators/test_eagle3.py
+16
-2
@@ -3,12 +3,20 @@
3
import pytest
4
import torch
5
6
+from vllm.model_executor.models.interfaces import supports_eagle3
7
+
8
9
@pytest.mark.parametrize(
10
"model_path",
11
[("nm-testing/SpeculatorLlama3-1-8B-Eagle3-converted-0717-quantized")])
10
-def test_llama(vllm_runner, example_prompts, model_path):
12
+def test_llama(vllm_runner, example_prompts, model_path, monkeypatch):
13
+ # Set environment variable for V1 engine serialization
14
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
15
+
16
with vllm_runner(model_path, dtype=torch.bfloat16) as vllm_model:
17
+ eagle3_supported = vllm_model.apply_model(supports_eagle3)
18
+ assert eagle3_supported
19
+
20
vllm_outputs = vllm_model.generate_greedy(example_prompts,
21
max_tokens=20)
22
print(vllm_outputs)
@@ -18,8 +26,14 @@ def test_llama(vllm_runner, example_prompts, model_path):
26
@pytest.mark.parametrize(
27
"model_path",
28
[("nm-testing/Speculator-Qwen3-8B-Eagle3-converted-071-quantized")])
21
-def test_qwen(vllm_runner, example_prompts, model_path):
29
+def test_qwen(vllm_runner, example_prompts, model_path, monkeypatch):
30
+ # Set environment variable for V1 engine serialization
31
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
32
+
33
with vllm_runner(model_path, dtype=torch.bfloat16) as vllm_model:
34
+ eagle3_supported = vllm_model.apply_model(supports_eagle3)
35
+ assert eagle3_supported
36
+
37
vllm_outputs = vllm_model.generate_greedy(example_prompts,
38
max_tokens=20)
39
print(vllm_outputs)
vllm/model_executor/models/interfaces.py
+53
@@ -823,3 +823,56 @@ def supports_v0_only(
823
model: Union[type[object], object],
824
) -> Union[TypeIs[type[SupportsV0Only]], TypeIs[SupportsV0Only]]:
825
return getattr(model, "supports_v0_only", False)
826
+
827
+
828
+@runtime_checkable
829
+class SupportsEagle3(Protocol):
830
+ """The interface required for models that support
831
+ EAGLE3 speculative decoding."""
832
+
833
+ supports_eagle3: ClassVar[Literal[True]] = True
834
+ """
835
+ A flag that indicates this model supports EAGLE3
836
+ speculative decoding.
837
+
838
+ Note:
839
+ There is no need to redefine this flag if this class is in the
840
+ MRO of your model class.
841
+ """
842
+
843
+ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
844
+ """
845
+ Set which layers should output auxiliary
846
+ hidden states for EAGLE3.
847
+
848
+ Args:
849
+ layers: Tuple of layer indices that should output auxiliary
850
+ hidden states.
851
+ """
852
+ ...
853
+
854
+ def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
855
+ """
856
+ Get the layer indices that should output auxiliary hidden states
857
+ for EAGLE3.
858
+
859
+ Returns:
860
+ Tuple of layer indices for auxiliary hidden state outputs.
861
+ """
862
+ ...
863
+
864
+
865
+@overload
866
+def supports_eagle3(model: type[object]) -> TypeIs[type[SupportsEagle3]]:
867
+ ...
868
+
869
+
870
+@overload
871
+def supports_eagle3(model: object) -> TypeIs[SupportsEagle3]:
872
+ ...
873
+
874
+
875
+def supports_eagle3(
876
+ model: Union[type[object], object],
877
+) -> Union[TypeIs[type[SupportsEagle3]], TypeIs[SupportsEagle3]]:
878
+ return isinstance(model, SupportsEagle3)
vllm/model_executor/models/llama.py
+2
-2
@@ -49,7 +49,7 @@ from vllm.model_executor.model_loader.weight_utils import (
49
from vllm.model_executor.sampling_metadata import SamplingMetadata
50
from vllm.sequence import IntermediateTensors
51
52
-from .interfaces import SupportsLoRA, SupportsPP
52
+from .interfaces import SupportsEagle3, SupportsLoRA, SupportsPP
53
from .utils import (AutoWeightsLoader, PPMissingLayer, extract_layer_index,
54
is_pp_missing_parameter,
55
make_empty_intermediate_tensors_factory, make_layers,
@@ -463,7 +463,7 @@ class LlamaModel(nn.Module):
463
return loaded_params
464
465
466
-class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
466
+class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):
467
packed_modules_mapping = {
468
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
469
"gate_up_proj": ["gate_proj", "up_proj"]
vllm/model_executor/models/qwen3.py
+2
-2
@@ -44,7 +44,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
44
from vllm.model_executor.sampling_metadata import SamplingMetadata
45
from vllm.sequence import IntermediateTensors
46
47
-from .interfaces import SupportsLoRA, SupportsPP
47
+from .interfaces import SupportsEagle3, SupportsLoRA, SupportsPP
48
from .qwen2 import Qwen2MLP as Qwen3MLP
49
from .qwen2 import Qwen2Model
50
from .utils import (AutoWeightsLoader, PPMissingLayer, extract_layer_index,
@@ -261,7 +261,7 @@ class Qwen3Model(Qwen2Model):
261
decoder_layer_type=Qwen3DecoderLayer)
262
263
264
-class Qwen3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
264
+class Qwen3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):
265
packed_modules_mapping = {
266
"qkv_proj": [
267
"q_proj",
vllm/v1/worker/gpu_model_runner.py
+8
-2
@@ -35,6 +35,7 @@ from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaBase
35
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
36
from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader
37
from vllm.model_executor.models.interfaces import (is_mixture_of_experts,
38
+ supports_eagle3,
39
supports_transcription)
40
from vllm.model_executor.models.interfaces_base import (
41
VllmModelForPooling, is_pooling_model, is_text_generation_model)
@@ -1981,8 +1982,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1982
logger.info("Loading drafter model...")
1983
self.drafter.load_model(self.model)
1984
if self.use_aux_hidden_state_outputs:
1984
- self.model.set_aux_hidden_state_layers(
1985
- self.model.get_eagle3_aux_hidden_state_layers())
1985
+ if supports_eagle3(self.model):
1986
+ self.model.set_aux_hidden_state_layers(
1987
+ self.model.get_eagle3_aux_hidden_state_layers())
1988
+ else:
1989
+ raise RuntimeError(
1990
+ "Model does not support EAGLE3 interface but "
1991
+ "aux_hidden_state_outputs was requested")
1992
time_after_load = time.perf_counter()
1993
self.model_memory_usage = m.consumed_memory
1994
logger.info("Model loading took %.4f GiB and %.6f seconds",