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, 08:30 UTC
[Spec Decode] Make EAGLE3 draft token ID mapping optional (#18488)
Signed-off-by: Benjamin Chislett <benjamin.chislett@centml.ai> Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Benjamin Chislett committed
May 22, 2025 at 23:17 UTC
583507d13075783a12ccbd774575974d10ca4959
1 file changed
+9
-1
vllm/model_executor/models/llama_eagle3.py
+9
-1
@@ -214,6 +214,9 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):
214
) -> Optional[torch.Tensor]:
215
logits = self.logits_processor(self.lm_head, hidden_states,
216
sampling_metadata)
217
+ if self.draft_id_to_target_id is None:
218
+ return logits
219
+
220
base = torch.arange(self.config.draft_vocab_size, device=logits.device)
221
targets = base + self.draft_id_to_target_id
222
logits_new = logits.new_full((
@@ -246,4 +249,9 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):
249
name = "model." + name
250
model_weights[name] = loaded_weight
251
249
- return loader.load_weights(model_weights.items())
252
+ loaded_weights = loader.load_weights(model_weights.items())
253
+
254
+ if 'd2t' not in loaded_weights:
255
+ self.draft_id_to_target_id = None
256
+
257
+ return loaded_weights