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, 13:30 UTC
[Minor][Spec Decode] Add use_eagle to SpeculativeConfig (#17213)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Apr 25, 2025 at 21:08 UTC
1cf0719ebdcd67ecdd5801da9154e56683fa7360
3 files changed
+6
-5
vllm/config.py
+3
@@ -2566,6 +2566,9 @@ class SpeculativeConfig:
2566
"""
2567
return self.num_speculative_tokens
2568
2569
+ def use_eagle(self) -> bool:
2570
+ return self.method in ("eagle", "eagle3")
2571
+
2572
def __repr__(self) -> str:
2573
method = self.method
2574
model = None if method == "ngram" else self.draft_model_config.model
vllm/v1/core/sched/scheduler.py
+1
-1
@@ -126,7 +126,7 @@ class Scheduler(SchedulerInterface):
126
self.num_spec_tokens = self.num_lookahead_tokens = 0
127
if speculative_config:
128
self.num_spec_tokens = speculative_config.num_speculative_tokens
129
- if speculative_config.method in ("eagle", "eagle3"):
129
+ if speculative_config.use_eagle():
130
self.num_lookahead_tokens = self.num_spec_tokens
131
132
def schedule(self) -> SchedulerOutput:
vllm/v1/worker/gpu_model_runner.py
+2
-4
@@ -171,8 +171,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
171
if get_pp_group().is_last_rank:
172
if self.speculative_config.method == "ngram":
173
self.drafter = NgramProposer(self.vllm_config)
174
- elif self.speculative_config.method == "eagle" or \
175
- self.speculative_config.method == "eagle3":
174
+ elif self.speculative_config.use_eagle():
175
self.drafter = EagleProposer(self.vllm_config,
176
self.device) # type: ignore
177
if self.speculative_config.method == "eagle3":
@@ -1192,8 +1191,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1191
assert isinstance(self.drafter, NgramProposer)
1192
spec_token_ids = self.generate_draft_token_ids(
1193
valid_sampled_token_ids, sampling_metadata)
1195
- elif self.speculative_config.method == "eagle" or \
1196
- self.speculative_config.method == "eagle3":
1194
+ elif self.speculative_config.use_eagle():
1195
assert isinstance(self.drafter, EagleProposer)
1196
# TODO(woosuk): Refactor the loop.
1197
next_token_ids: list[int] = []