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 12, 14:45 UTC
[Spec Decode] Add support for EAGLE3 heads that do not use_aux_hidden_states (#27688)
Signed-off-by: hjjq <hanjieq@nvidia.com> Signed-off-by: Benjamin Chislett <bchislett@nvidia.com> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
Hanjie Qiu committed
Nov 24, 2025 at 20:13 UTC
5f9679a43bf92fc0fc8610f0ba5cc9c857148ccf
3 files changed
+46
-15
vllm/model_executor/models/llama_eagle3.py
+24
-14
@@ -142,6 +142,12 @@ class LlamaModel(nn.Module):
142
# Get drafter's quantization config
143
self.quant_config = get_draft_quant_config(vllm_config)
144
145
+ eagle_config = getattr(self.config, "eagle_config", None)
146
+ if eagle_config is not None and "use_aux_hidden_state" in eagle_config:
147
+ self.use_aux_hidden_state = eagle_config["use_aux_hidden_state"]
148
+ else:
149
+ self.use_aux_hidden_state = True
150
+
151
current_vllm_config = get_current_vllm_config()
152
153
self.embed_tokens = VocabParallelEmbedding(
@@ -161,20 +167,20 @@ class LlamaModel(nn.Module):
167
for layer_idx in range(self.config.num_hidden_layers)
168
]
169
)
164
- if hasattr(self.config, "target_hidden_size"):
165
- fc_input_size = self.config.target_hidden_size * 3
166
- else:
167
- fc_input_size = self.config.hidden_size * 3
168
- self.fc = ReplicatedLinear(
169
- input_size=fc_input_size,
170
- output_size=self.config.hidden_size,
171
- bias=False,
172
- params_dtype=vllm_config.model_config.dtype,
173
- quant_config=self.quant_config,
174
- prefix=maybe_prefix(prefix, "fc"),
175
- return_bias=False,
176
- )
177
-
170
+ if self.use_aux_hidden_state:
171
+ if hasattr(self.config, "target_hidden_size"):
172
+ fc_input_size = self.config.target_hidden_size * 3
173
+ else:
174
+ fc_input_size = self.config.hidden_size * 3
175
+ self.fc = ReplicatedLinear(
176
+ input_size=fc_input_size,
177
+ output_size=self.config.hidden_size,
178
+ bias=False,
179
+ params_dtype=vllm_config.model_config.dtype,
180
+ quant_config=self.quant_config,
181
+ prefix=maybe_prefix(prefix, "fc"),
182
+ return_bias=False,
183
+ )
184
self.norm = RMSNorm(
185
self.config.hidden_size,
186
eps=self.config.rms_norm_eps,
@@ -332,6 +338,8 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):
338
self,
339
hidden_states: torch.Tensor,
340
) -> torch.Tensor:
341
+ if not self.model.use_aux_hidden_state:
342
+ return hidden_states
343
# combine multiple auxiliary hidden states returned by eagle3
344
return self.model.fc(hidden_states)
345
@@ -357,6 +365,8 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):
365
skip_substrs.append("draft_id_to_target_id")
366
if not includes_embed_tokens:
367
skip_substrs.append("embed_tokens")
368
+ if not self.model.use_aux_hidden_state:
369
+ skip_substrs.append("fc.")
370
loader = AutoWeightsLoader(
371
self,
372
skip_prefixes=None,
vllm/v1/spec_decode/eagle.py
+19
@@ -83,6 +83,9 @@ class EagleProposer:
83
self.draft_indexer_metadata_builder: AttentionMetadataBuilder | None = None
84
self.attn_layer_names: list[str] = []
85
self.indexer_layer_names: list[str] = []
86
+ self.eagle3_use_aux_hidden_state: bool = (
87
+ self._get_eagle3_use_aux_hidden_state_from_config()
88
+ )
89
90
self.use_cuda_graph = False
91
@@ -1169,6 +1172,22 @@ class EagleProposer:
1172
)
1173
return builder
1174
1175
+ def _get_eagle3_use_aux_hidden_state_from_config(self) -> bool:
1176
+ """
1177
+ Some eagle3 heads (e.g., nvidia/gpt-oss-120b-Eagle3-v2) do not use auxiliary
1178
+ hidden states and directly uses the last layer output just like eagle1.
1179
+ They might indicate this by setting "use_aux_hidden_state" to False
1180
+ inside the "eagle_config" dict of their hf_config.
1181
+ """
1182
+ if self.method != "eagle3":
1183
+ return False
1184
+ # Assume that eagle3 heads use aux hidden states by default
1185
+ use_aux_hidden_state = True
1186
+ eagle_config = getattr(self.draft_model_config.hf_config, "eagle_config", None)
1187
+ if eagle_config is not None:
1188
+ use_aux_hidden_state = eagle_config.get("use_aux_hidden_state", True)
1189
+ return use_aux_hidden_state
1190
+
1191
def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:
1192
"""
1193
Validate that all eagle layers belong to the same KVCacheGroup.
vllm/v1/worker/gpu_model_runner.py
+3
-1
@@ -375,7 +375,9 @@ class GPUModelRunner(
375
elif self.speculative_config.use_eagle():
376
self.drafter = EagleProposer(self.vllm_config, self.device, self)
377
if self.speculative_config.method == "eagle3":
378
- self.use_aux_hidden_state_outputs = True
378
+ self.use_aux_hidden_state_outputs = (
379
+ self.drafter.eagle3_use_aux_hidden_state
380
+ )
381
elif self.speculative_config.method == "medusa":
382
self.drafter = MedusaProposer(
383
vllm_config=self.vllm_config, device=self.device