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 15, 12:15 UTC
[V1][Spec Decode] Make Eagle model arch config driven (#17323)
Ekagra Ranjan committed
Apr 28, 2025 at 22:22 UTC
e1360005956d4887d101bbd675b7b0574a0afc45
3 files changed
+26
-13
vllm/config.py
+2
-1
@@ -2401,7 +2401,8 @@ class SpeculativeConfig:
2401
pass
2402
else:
2403
eagle_config = EAGLEConfig(
2404
- self.draft_model_config.hf_config)
2404
+ self.draft_model_config.hf_config,
2405
+ method=self.method)
2406
self.draft_model_config.hf_config = eagle_config
2407
2408
if (self.num_speculative_tokens is not None
vllm/transformers_utils/configs/eagle.py
+18
-1
@@ -15,6 +15,7 @@ class EAGLEConfig(PretrainedConfig):
15
def __init__(self,
16
model: Union[PretrainedConfig, dict, None] = None,
17
truncated_vocab_size: Optional[int] = None,
18
+ method: Optional[str] = 'eagle',
19
**kwargs):
20
21
model_config: Union[PretrainedConfig, DeepseekV2Config, None]
@@ -45,7 +46,23 @@ class EAGLEConfig(PretrainedConfig):
46
if not envs.VLLM_USE_V1:
47
kwargs["architectures"] = ["EAGLEModel"]
48
else:
48
- kwargs["architectures"] = ["EagleLlamaForCausalLM"]
49
+ # Eagle model name should follow naming convention of
50
+ # LlamaForCausalLM -> EagleLlamaForCausalLM
51
+ if method == "eagle":
52
+ assert self.model is not None, \
53
+ "model should not be None when method is eagle"
54
+ kwargs["architectures"] = [
55
+ f"Eagle{arch}" for arch in self.model.architectures
56
+ ]
57
+ elif method == "eagle3":
58
+ assert self.model is not None, \
59
+ "model should not be None when method is eagle3"
60
+ kwargs["architectures"] = [
61
+ f"Eagle3{arch}" for arch in self.model.architectures
62
+ ]
63
+ else:
64
+ raise ValueError(f"Invalid method {method}. \
65
+ Supported methods are eagle and eagle3.")
66
67
super().__init__(**kwargs)
68
vllm/v1/spec_decode/eagle.py
+6
-11
@@ -9,8 +9,7 @@ from vllm.forward_context import set_forward_context
9
from vllm.logger import init_logger
10
from vllm.model_executor.model_loader.loader import get_model_loader
11
from vllm.model_executor.model_loader.utils import set_default_torch_dtype
12
-from vllm.model_executor.models.llama_eagle import EagleLlamaForCausalLM
13
-from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
12
+from vllm.model_executor.models import ModelRegistry
13
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
14
from vllm.v1.sample.metadata import SamplingMetadata
15
@@ -225,15 +224,11 @@ class EagleProposer:
224
with set_default_torch_dtype(
225
draft_model_config.dtype), set_current_vllm_config(
226
self.vllm_config):
228
- if self.vllm_config.speculative_config.method == "eagle":
229
- self.model = EagleLlamaForCausalLM(
230
- model_config=draft_model_config,
231
- start_layer_id=target_layer_num).to(target_device)
232
- else:
233
- assert self.vllm_config.speculative_config.method == "eagle3"
234
- self.model = Eagle3LlamaForCausalLM(
235
- model_config=draft_model_config,
236
- start_layer_id=target_layer_num).to(target_device)
227
+ draft_model_cls, arch = ModelRegistry.resolve_model_cls(
228
+ draft_model_config.architectures)
229
+ self.model = draft_model_cls(
230
+ model_config=draft_model_config,
231
+ start_layer_id=target_layer_num).to(target_device)
232
233
loaded_weights = self.model.load_weights(
234
loader.get_all_weights(