[Spec Decode] EAGLE 3: Micro-optimization of Draft to Target logit conversion

vllm/model_executor/models/llama_eagle.py

@@ -18,6 +18,7 @@
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.llama import (LlamaDecoderLayer,
                                               LlamaForCausalLM)
+from vllm.v1.spec_decode.eagle_mixin import EagleMixin
 
 from .utils import AutoWeightsLoader, maybe_prefix
 
@@ -133,7 +134,7 @@ def load_weights(self, weights: Iterable[tuple[str,
         return loaded_params
 
 
-class EagleLlamaForCausalLM(LlamaForCausalLM):
+class EagleLlamaForCausalLM(EagleMixin, LlamaForCausalLM):
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         nn.Module.__init__(self)

vllm/model_executor/models/llama_eagle3.py

@@ -20,6 +20,7 @@
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.llama import (LlamaDecoderLayer,
                                               LlamaForCausalLM)
+from vllm.v1.spec_decode.eagle_mixin import Eagle3Mixin
 
 from .utils import AutoWeightsLoader, maybe_prefix
 
@@ -207,7 +208,7 @@ def load_weights(self, weights: Iterable[tuple[str,
         return loaded_params
 
 
-class Eagle3LlamaForCausalLM(LlamaForCausalLM):
+class Eagle3LlamaForCausalLM(Eagle3Mixin, LlamaForCausalLM):
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         nn.Module.__init__(self)
@@ -254,26 +255,6 @@ def forward(
     ) -> tuple[torch.Tensor, torch.Tensor]:
         return self.model(input_ids, positions, hidden_states, inputs_embeds)
 
-    def compute_logits(
-        self,
-        hidden_states: torch.Tensor,
-    ) -> Optional[torch.Tensor]:
-        logits = self.logits_processor(self.lm_head, hidden_states)
-        if self.draft_id_to_target_id is None:
-            assert logits.shape[1] == self.config.vocab_size, \
-                "Expected logits to have shape " \
-                f"(*, {self.config.vocab_size}), but got {logits.shape}"
-            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((
-            logits.shape[0],
-            self.config.vocab_size,
-        ), float('-inf'))
-        logits_new[:, targets] = logits
-        return logits_new
-
     def combine_hidden_states(
         self,
         hidden_states: torch.Tensor,

vllm/v1/spec_decode/eagle.py

@@ -255,11 +255,10 @@ def propose(
             else:
                 last_hidden_states, hidden_states = ret_hidden_states
         sample_hidden_states = last_hidden_states[last_token_indices]
-        logits = self.model.compute_logits(sample_hidden_states)
 
         # Early exit if there is only one draft token to be generated.
         if self.num_speculative_tokens == 1:
-            draft_token_ids = logits.argmax(dim=-1)
+            draft_token_ids = self.model.propose_chain(sample_hidden_states)
             return draft_token_ids.view(-1, 1)
 
         if self.uses_mrope:
@@ -273,6 +272,7 @@ def propose(
 
         if isinstance(attn_metadata, TreeAttentionMetadata):
             # Draft using tree attention.
+            logits = self.model.compute_logits(sample_hidden_states, None)
             draft_token_ids_list = self.propose_tree(
                 batch_size=batch_size,
                 logits=logits,
@@ -283,7 +283,7 @@ def propose(
             # [batch_size, num_tree_tokens]
             return torch.cat(draft_token_ids_list, dim=1)
 
-        draft_token_ids = logits.argmax(dim=-1)
+        draft_token_ids = self.model.propose_chain(sample_hidden_states)
 
         if self.allowed_attn_types is not None and \
             not isinstance(attn_metadata, self.allowed_attn_types):
@@ -406,8 +406,8 @@ def propose(
                 else:
                     last_hidden_states, hidden_states = ret_hidden_states
             hidden_states = hidden_states[:batch_size]
-            logits = self.model.compute_logits(last_hidden_states[:batch_size])
-            draft_token_ids = logits.argmax(dim=-1)
+            draft_token_ids = self.model.propose_chain(
+                last_hidden_states[:batch_size])
             draft_token_ids_list.append(draft_token_ids)
 
         # [batch_size, num_speculative_tokens]

vllm/v1/spec_decode/eagle_mixin.py

@@ -0,0 +1,45 @@
+import torch
+from typing import Optional
+
+class EagleMixin:
+    def propose_chain(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        logits = self.compute_logits(hidden_states)
+        draft_token_ids = logits.argmax(dim=-1)
+        return draft_token_ids
+
+class Eagle3Mixin:
+    def compute_draft_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> Optional[torch.Tensor]:
+        # Computes draft logits from hidden states.
+        logits = self.logits_processor(self.lm_head, hidden_states)
+        return logits
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> Optional[torch.Tensor]:
+        # Maps draft logits to target logits.
+        logits = self.compute_draft_logits(hidden_states)
+        if self.draft_id_to_target_id is None:
+            assert logits.shape[1] == self.config.vocab_size, \
+                "Expected logits to have shape " \
+                f"(*, {self.config.vocab_size}), but got {logits.shape}"
+            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((
+            logits.shape[0],
+            self.config.vocab_size,
+        ), float('-inf'))
+        logits_new[:, targets] = logits
+        return logits_new
+
+    def propose_chain(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        logits = self.compute_draft_logits(hidden_states)
+        draft_token_ids = logits.argmax(dim=-1)
+        offset = self.draft_id_to_target_id[draft_token_ids]
+        draft_token_ids += offset
+        return draft_token_ids