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, 09:45 UTC
[V1][Spec Decode][Bugfix] Load quantize weights for EAGLE (#18290)
Ekagra Ranjan committed
May 22, 2025 at 18:17 UTC
a04720bc36401d831cb048c3917b9e58173d9c1d
2 files changed
+9
-3
vllm/transformers_utils/configs/eagle.py
+4
-2
@@ -52,13 +52,15 @@ class EAGLEConfig(PretrainedConfig):
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
55
+ f"Eagle{arch}" if not arch.startswith("Eagle") \
56
+ else arch for arch in self.model.architectures
57
]
58
elif method == "eagle3":
59
assert self.model is not None, \
60
"model should not be None when method is eagle3"
61
kwargs["architectures"] = [
61
- f"Eagle3{arch}" for arch in self.model.architectures
62
+ f"Eagle3{arch}" if not arch.startswith("Eagle3") \
63
+ else arch for arch in self.model.architectures
64
]
65
else:
66
raise ValueError(f"Invalid method {method}. \
vllm/v1/spec_decode/eagle.py
+5
-1
@@ -9,7 +9,8 @@ from vllm.distributed.parallel_state import get_pp_group
9
from vllm.forward_context import set_forward_context
10
from vllm.logger import init_logger
11
from vllm.model_executor.model_loader import get_model_loader
12
-from vllm.model_executor.model_loader.utils import set_default_torch_dtype
12
+from vllm.model_executor.model_loader.utils import (
13
+ process_weights_after_loading, set_default_torch_dtype)
14
from vllm.model_executor.models import ModelRegistry
15
from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
16
from vllm.triton_utils import tl, triton
@@ -308,6 +309,9 @@ class EagleProposer:
309
loaded_weights = self.model.load_weights(
310
loader.get_all_weights(draft_model_config, self.model))
311
312
+ process_weights_after_loading(self.model, draft_model_config,
313
+ target_device)
314
+
315
# share embed_tokens with the target model if needed
316
if get_pp_group().world_size == 1:
317
assert "model.embed_tokens.weight" not in loaded_weights, \