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 14, 05:30 UTC
Enable Eagle3 speculative decoding for GPT-OSS model (#25246)
Signed-off-by: Eldar Kurtic <8884008+eldarkurtic@users.noreply.github.com>
Eldar Kurtić committed
Sep 22, 2025 at 10:50 UTC
21467f9a1c6219cbd66640e539b1f23221cff375
3 files changed
+41
-12
vllm/config/speculative.py
+1
-1
@@ -527,7 +527,7 @@ class SpeculativeConfig:
527
"speculative decoding is > 1, but got "
528
f"{self.disable_by_batch_size=}")
529
530
- eagle3_target_supported = ["llama", "qwen"]
530
+ eagle3_target_supported = ["llama", "qwen", "gpt_oss"]
531
if self.method == "eagle3" and self.target_model_config and not any(
532
supported_model in
533
self.target_model_config.hf_text_config.model_type
vllm/model_executor/models/gpt_oss.py
+17
-2
@@ -27,7 +27,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
27
from vllm.sequence import IntermediateTensors
28
from vllm.utils import cdiv
29
30
-from .interfaces import SupportsPP
30
+from .interfaces import SupportsEagle3, SupportsPP
31
from .utils import (AutoWeightsLoader, WeightsMapper, extract_layer_index,
32
is_pp_missing_parameter,
33
make_empty_intermediate_tensors_factory, make_layers,
@@ -238,6 +238,7 @@ class GptOssModel(nn.Module):
238
self.make_empty_intermediate_tensors = (
239
make_empty_intermediate_tensors_factory(
240
["hidden_states", "residual"], self.config.hidden_size))
241
+ self.aux_hidden_state_layers = tuple[int, ...]()
242
243
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
244
return self.embedding(input_ids)
@@ -261,8 +262,12 @@ class GptOssModel(nn.Module):
262
x = intermediate_tensors["hidden_states"]
263
residual = intermediate_tensors["residual"]
264
265
+ aux_hidden_states = []
266
for i in range(self.start_layer, self.end_layer):
267
layer = self.layers[i]
268
+ if i in self.aux_hidden_state_layers:
269
+ aux_hidden_states.append(x if residual is None else x +
270
+ residual)
271
x, residual = layer(x, positions, residual)
272
if not get_pp_group().is_last_rank:
273
return IntermediateTensors({
@@ -270,6 +275,9 @@ class GptOssModel(nn.Module):
275
"residual": residual
276
})
277
x, _ = self.norm(x, residual)
278
+
279
+ if len(aux_hidden_states) > 0:
280
+ return x, aux_hidden_states
281
return x
282
283
def _load_weights_mxfp4(
@@ -610,7 +618,7 @@ class GptOssModel(nn.Module):
618
weights, stacked_params_mapping)
619
620
613
-class GptOssForCausalLM(nn.Module, SupportsPP):
621
+class GptOssForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
622
packed_modules_mapping = {"qkv": ["q_proj", "k_proj", "v_proj"]}
623
624
hf_to_vllm_mapper = WeightsMapper(
@@ -658,6 +666,13 @@ class GptOssForCausalLM(nn.Module, SupportsPP):
666
self.make_empty_intermediate_tensors = (
667
self.model.make_empty_intermediate_tensors)
668
669
+ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
670
+ self.model.aux_hidden_state_layers = layers
671
+
672
+ def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
673
+ num_layers = len(self.model.layers)
674
+ return (2, num_layers // 2, num_layers - 3)
675
+
676
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
677
return self.model.get_input_embeddings(input_ids)
678
vllm/v1/spec_decode/eagle.py
+23
-9
@@ -823,15 +823,29 @@ class EagleProposer:
823
else:
824
target_language_model = target_model
825
# share embed_tokens with the target model if needed
826
- if get_pp_group().world_size == 1 \
827
- and self.model.model.embed_tokens.weight.shape \
828
- == target_language_model.model.embed_tokens.weight.shape:
829
- logger.info(
830
- "Assuming the EAGLE head shares the same vocab embedding"
831
- " with the target model.")
832
- del self.model.model.embed_tokens
833
- self.model.model.embed_tokens = (
834
- target_language_model.model.embed_tokens)
826
+ if get_pp_group().world_size == 1:
827
+ if hasattr(target_language_model.model, 'embed_tokens'):
828
+ target_embed_tokens = target_language_model.model.embed_tokens
829
+ elif hasattr(target_language_model.model, 'embedding'):
830
+ target_embed_tokens = target_language_model.model.embedding
831
+ else:
832
+ raise AttributeError(
833
+ "Target model does not have 'embed_tokens' or 'embedding' "
834
+ "attribute")
835
+
836
+ # Check if shapes match and we found the embedding
837
+ eagle_shape = self.model.model.embed_tokens.weight.shape
838
+ target_shape = target_embed_tokens.weight.shape
839
+ if eagle_shape == target_shape:
840
+ logger.info(
841
+ "Assuming the EAGLE head shares the same vocab embedding"
842
+ " with the target model.")
843
+ del self.model.model.embed_tokens
844
+ self.model.model.embed_tokens = target_embed_tokens
845
+ else:
846
+ logger.info(
847
+ "The EAGLE head's vocab embedding will be loaded separately"
848
+ " from the target model.")
849
else:
850
logger.info(
851
"The EAGLE head's vocab embedding will be loaded separately"