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 10, 06:30 UTC
[BugFix]Fix eagle draft_model_config and add tests (#31753)
Signed-off-by: Xingyu Liu <charlotteliu12x@gmail.com>
Xingyu Liu committed
Jan 12, 2026 at 23:09 UTC
80221e1884bd36751e8ae0308acf71f42946a05e
4 files changed
+40
-4
tests/config/draft_model_arch_groundtruth.json
+3
-3
@@ -38,7 +38,7 @@
38
"EagleDeepSeekMTPModel"
39
],
40
"model_type": "eagle",
41
- "text_model_type": "deepseek_mtp",
41
+ "text_model_type": "eagle",
42
"hidden_size": 2560,
43
"total_num_hidden_layers": 1,
44
"total_num_attention_heads": 32,
@@ -55,7 +55,7 @@
55
"EagleLlamaForCausalLM"
56
],
57
"model_type": "eagle",
58
- "text_model_type": "llama",
58
+ "text_model_type": "eagle",
59
"hidden_size": 4096,
60
"total_num_hidden_layers": 1,
61
"total_num_attention_heads": 32,
@@ -72,7 +72,7 @@
72
"Eagle3LlamaForCausalLM"
73
],
74
"model_type": "eagle",
75
- "text_model_type": "llama",
75
+ "text_model_type": "eagle",
76
"hidden_size": 4096,
77
"total_num_hidden_layers": 1,
78
"total_num_attention_heads": 32,
tests/test_config.py
+22
@@ -13,8 +13,10 @@ from vllm.compilation.backends import VllmBackend
13
from vllm.config import (
14
CompilationConfig,
15
ModelConfig,
16
+ ParallelConfig,
17
PoolerConfig,
18
SchedulerConfig,
19
+ SpeculativeConfig,
20
VllmConfig,
21
update_config,
22
)
@@ -1105,3 +1107,23 @@ def test_needs_dp_coordination(
1107
vllm_config = VllmConfig(model_config=model_config, parallel_config=parallel_config)
1108
1109
assert vllm_config.needs_dp_coordinator == expected_needs_coordinator
1110
+
1111
+
1112
+def test_eagle_draft_model_config():
1113
+ """Test that EagleDraft model config is correctly set."""
1114
+ target_model_config = ModelConfig(
1115
+ "meta-llama/Meta-Llama-3-8B-Instruct", trust_remote_code=True
1116
+ )
1117
+ speculative_config = SpeculativeConfig(
1118
+ model="yuhuili/EAGLE-LLaMA3-Instruct-8B",
1119
+ num_speculative_tokens=1,
1120
+ target_model_config=target_model_config,
1121
+ target_parallel_config=ParallelConfig(),
1122
+ )
1123
+ draft_model_config = speculative_config.draft_model_config
1124
+ assert draft_model_config.hf_config.architectures == ["EagleLlamaForCausalLM"]
1125
+ assert draft_model_config.hf_text_config.architectures == ["EagleLlamaForCausalLM"]
1126
+ assert draft_model_config.hf_config.model_type == "eagle"
1127
+ assert draft_model_config.hf_text_config.model_type == "eagle"
1128
+ assert draft_model_config.architectures == ["EagleLlamaForCausalLM"]
1129
+ assert draft_model_config.architecture == "EagleLlamaForCausalLM"
vllm/config/speculative.py
+14
@@ -12,6 +12,7 @@ from vllm.config.model import ModelConfig
12
from vllm.config.parallel import ParallelConfig
13
from vllm.config.utils import config
14
from vllm.logger import init_logger
15
+from vllm.transformers_utils.config import get_hf_text_config
16
from vllm.utils.hashing import safe_hash
17
from vllm.utils.import_utils import LazyLoader, has_arctic_inference
18
@@ -409,10 +410,23 @@ class SpeculativeConfig:
410
method=self.method,
411
model_type="eagle",
412
)
413
+ # EAGLEConfig primarily updates architectures, so update
414
+ # all architectures-related fields in draft_model_config
415
self.draft_model_config.hf_config = eagle_config
416
+ self.draft_model_config.hf_text_config = get_hf_text_config(
417
+ self.draft_model_config.hf_config
418
+ )
419
self.draft_model_config.model_arch_config = (
420
self.draft_model_config.get_model_arch_config()
421
)
422
+ model_info, arch = (
423
+ self.draft_model_config.registry.inspect_model_cls(
424
+ self.draft_model_config.architectures,
425
+ self.draft_model_config,
426
+ )
427
+ )
428
+ self.draft_model_config._model_info = model_info
429
+ self.draft_model_config._architecture = arch
430
431
if self.num_speculative_tokens is not None and hasattr(
432
self.draft_model_config.hf_config, "num_lookahead_tokens"
vllm/transformers_utils/model_arch_config_convertor.py
+1
-1
@@ -201,7 +201,7 @@ class ModelArchConfigConvertorBase:
201
# underlying architecture
202
return (
203
self.hf_text_config.model.model_type
204
- in ("deepseek_v2", "deepseek_v3", "deepseek_v32")
204
+ in ("deepseek_v2", "deepseek_v3", "deepseek_v32", "deepseek_mtp")
205
and self.hf_text_config.kv_lora_rank is not None
206
)
207
return False