[Spec Decode] Add support for EAGLE3 heads that do not use_aux_hidden_states

vllm/config/speculative.py

@@ -166,9 +166,15 @@ def compute_hash(self) -> str:
         the final hidden states.
         """
         factors: list[Any] = []
-        # Eagle3 affects the computation graph because it returns intermediate
-        # hidden states in addition to the final hidden state.
-        factors.append(self.method == "eagle3")
+        if self.method == "eagle3":
+            # Eagle3 affects the computation graph because it returns intermediate
+            # hidden states in addition to the final hidden state, and
+            # some eagle3 draft models may have different intermediate_size
+            if self.draft_model_config is not None:
+                factors.append(self.draft_model_config.compute_hash())
+            else:
+                factors.append("None")
+
         hash_str = hashlib.md5(str(factors).encode(), usedforsecurity=False).hexdigest()
         return hash_str
 

vllm/v1/spec_decode/eagle.py

@@ -83,6 +83,9 @@ def __init__(
         self.draft_indexer_metadata_builder: AttentionMetadataBuilder | None = None
         self.attn_layer_names: list[str] = []
         self.indexer_layer_names: list[str] = []
+        self.eagle3_use_aux_hidden_state: bool = (
+            self._get_eagle3_use_aux_hidden_state_from_config()
+        )
 
         self.use_cuda_graph = False
 
@@ -225,9 +228,11 @@ def propose(
 
         if self.method == "eagle3":
             assert isinstance(self.model, Eagle3LlamaForCausalLM)
-            target_hidden_states = self.model.combine_hidden_states(
-                target_hidden_states
-            )
+            # Do not combine hidden states if eagle3 head does not use aux hidden states
+            if self.eagle3_use_aux_hidden_state:
+                target_hidden_states = self.model.combine_hidden_states(
+                    target_hidden_states
+                )
             assert target_hidden_states.shape[-1] == self.hidden_size
         # Shift the input ids by one token.
         # E.g., [a1, b1, b2, c1, c2, c3] -> [b1, b2, c1, c2, c3, c3]
@@ -1148,6 +1153,22 @@ def _get_attention_metadata_builder(self) -> AttentionMetadataBuilder:
         )
         return builder
 
+    def _get_eagle3_use_aux_hidden_state_from_config(self) -> bool:
+        """
+        Some eagle3 heads (e.g., nvidia/gpt-oss-120b-Eagle3-v2) do not use auxiliary
+        hidden states and directly uses the last layer output just like eagle1.
+        They might indicate this by setting "use_aux_hidden_state" to False
+        inside the "eagle_config" dict of their hf_config.
+        """
+        if self.method != "eagle3":
+            return False
+        # Assume that eagle3 heads use aux hidden states by default
+        use_aux_hidden_state = True
+        eagle_config = getattr(self.draft_model_config.hf_config, "eagle_config", None)
+        if eagle_config is not None:
+            use_aux_hidden_state = eagle_config.get("use_aux_hidden_state", True)
+        return use_aux_hidden_state
+
     def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:
         """
         Validate that all eagle layers belong to the same KVCacheGroup.

vllm/v1/worker/gpu_model_runner.py

@@ -365,7 +365,9 @@ def __init__(
             elif self.speculative_config.use_eagle():
                 self.drafter = EagleProposer(self.vllm_config, self.device, self)
                 if self.speculative_config.method == "eagle3":
-                    self.use_aux_hidden_state_outputs = True
+                    self.use_aux_hidden_state_outputs = (
+                        self.drafter.eagle3_use_aux_hidden_state
+                    )
             elif self.speculative_config.method == "medusa":
                 self.drafter = MedusaProposer(
                     vllm_config=self.vllm_config, device=self.device