diff --git a/vllm/model_executor/models/llama_eagle3.py b/vllm/model_executor/models/llama_eagle3.py index 2302d1352de6..98e78e8cc0da 100644 --- a/vllm/model_executor/models/llama_eagle3.py +++ b/vllm/model_executor/models/llama_eagle3.py @@ -213,6 +213,9 @@ def compute_logits( ) -> Optional[torch.Tensor]: logits = self.logits_processor(self.lm_head, hidden_states, sampling_metadata) + if self.draft_id_to_target_id is None: + return logits + base = torch.arange(self.config.draft_vocab_size, device=logits.device) targets = base + self.draft_id_to_target_id logits_new = logits.new_full(( @@ -245,4 +248,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): name = "model." + name model_weights[name] = loaded_weight - return loader.load_weights(model_weights.items()) + loaded_weights = loader.load_weights(model_weights.items()) + + if 'd2t' not in loaded_weights: + self.draft_id_to_target_id = None + + return loaded_weights