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
[Speculators][Speculative Decoding] Fix gpt-oss eagle3 accuracy issue (#25406)
Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com>
jiahanc committed
Sep 23, 2025 at 12:44 UTC
d5944d5146d9a6847badc7018b328de079cc7e80
6 files changed
+79
-17
tests/v1/spec_decode/test_eagle.py
+4
@@ -534,6 +534,8 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
534
proposer.runner.attn_groups.append([mock.MagicMock()])
535
proposer.runner.attn_groups[0][0].get_metadata_builder.return_value = \
536
attn_metadata_builder
537
+ proposer._get_attention_metadata_builder = mock.MagicMock(
538
+ return_value=attn_metadata_builder)
539
540
result = proposer.propose(target_token_ids=target_token_ids,
541
target_positions=target_positions,
@@ -660,6 +662,8 @@ def test_propose_tree(spec_token_tree):
662
proposer.runner.attn_groups.append([mock.MagicMock()])
663
proposer.runner.attn_groups[0][0].get_metadata_builder.return_value = \
664
attn_metadata_builder
665
+ proposer._get_attention_metadata_builder = mock.MagicMock(
666
+ return_value=attn_metadata_builder)
667
668
# Setup inputs for the proposer.
669
target_token_ids = torch.randint(0,
vllm/config/model.py
+1
@@ -1003,6 +1003,7 @@ class ModelConfig:
1003
self.quantization = quantization_override
1004
break
1005
1006
+ quant_method = quant_method if quant_method != "" else None
1007
# Verify quantization configurations.
1008
if self.quantization is None:
1009
self.quantization = quant_method
vllm/model_executor/models/llama_eagle.py
+5
@@ -134,6 +134,11 @@ class EagleLlamaForCausalLM(LlamaForCausalLM):
134
nn.Module.__init__(self)
135
self.config = vllm_config. \
136
speculative_config.draft_model_config.hf_config
137
+ # Ensure draft_vocab_size is set
138
+ # default to the base vocab size when absent
139
+ if getattr(self.config, "draft_vocab_size", None) is None:
140
+ base_vocab_size = getattr(self.config, "vocab_size", None)
141
+ self.config.draft_vocab_size = base_vocab_size
142
target_layer_num = vllm_config.model_config.get_num_layers(
143
vllm_config.parallel_config)
144
self.model = LlamaModel(vllm_config=vllm_config,
vllm/model_executor/models/llama_eagle3.py
+5
@@ -203,6 +203,11 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):
203
nn.Module.__init__(self)
204
self.config = vllm_config. \
205
speculative_config.draft_model_config.hf_config
206
+ # Ensure draft_vocab_size is set
207
+ # default to the base vocab size when absent
208
+ if getattr(self.config, "draft_vocab_size", None) is None:
209
+ base_vocab_size = getattr(self.config, "vocab_size", None)
210
+ self.config.draft_vocab_size = base_vocab_size
211
target_layer_num = vllm_config.model_config.get_num_layers(
212
vllm_config.parallel_config)
213
vllm/v1/spec_decode/eagle.py
+56
-14
@@ -9,6 +9,7 @@ import numpy as np
9
import torch
10
import torch.nn as nn
11
12
+from vllm.attention.backends.abstract import AttentionMetadataBuilder
13
from vllm.attention.layer import Attention
14
from vllm.config import (CompilationLevel, VllmConfig,
15
get_layers_from_vllm_config)
@@ -77,6 +78,8 @@ class EagleProposer:
78
self.is_multimodal_model = vllm_config.model_config \
79
.is_multimodal_model
80
81
+ self.attn_metadata_builder: Optional[AttentionMetadataBuilder] = None
82
+
83
self.use_cuda_graph = (self.vllm_config.compilation_config.level
84
== CompilationLevel.PIECEWISE and
85
not self.vllm_config.model_config.enforce_eager)
@@ -117,7 +120,7 @@ class EagleProposer:
120
with_numpy=True)
121
122
# Determine allowed attention backends once during initialization.
120
- self.allowed_attn_types: tuple[type[EagleAttentionMetadata], ...]
123
+ self.allowed_attn_types: tuple[type, ...]
124
if current_platform.is_rocm():
125
rocm_types = [TritonAttentionMetadata, FlashAttentionMetadata]
126
# vllm.v1.attention.backends.rocm_aiter_fa is an optional backend
@@ -190,10 +193,12 @@ class EagleProposer:
193
194
assert self.runner is not None
195
193
- # FIXME: need to consider multiple kv_cache_groups
194
- attn_metadata_builder = \
195
- self.runner.attn_groups[0][0].get_metadata_builder()
196
- attn_metadata = attn_metadata_builder.build_for_drafting(
196
+ # Select the correct attention metadata builders for EAGLE layers.
197
+ # Get the attention metadata builders once and reuse for later.
198
+ builder = (self._get_attention_metadata_builder()
199
+ if self.attn_metadata_builder is None else
200
+ self.attn_metadata_builder)
201
+ attn_metadata = builder.build_for_drafting(
202
common_attn_metadata=common_attn_metadata, draft_index=0)
203
204
# At this moment, we assume all eagle layers belong to the same KV
@@ -327,11 +332,9 @@ class EagleProposer:
332
exceeds_max_model_len, PADDING_SLOT_ID)
333
334
# Rebuild attention metadata
330
- attn_metadata_builder = \
331
- self.runner.attn_groups[0][0].get_metadata_builder()
332
- attn_metadata = attn_metadata_builder\
333
- .build_for_drafting(common_attn_metadata=common_attn_metadata,
334
- draft_index=token_index + 1)
335
+ attn_metadata = builder.build_for_drafting(
336
+ common_attn_metadata=common_attn_metadata,
337
+ draft_index=token_index + 1)
338
for layer_name in self.attn_layer_names:
339
per_layer_attn_metadata[layer_name] = attn_metadata
340
@@ -851,10 +854,24 @@ class EagleProposer:
854
# share lm_head with the target model if needed
855
# some model definition do not define lm_head explicitly
856
# and reuse embed_tokens for lm_head, e.g., CohereForCausalLM
854
- if self.vllm_config.speculative_config.method != "eagle3" and \
855
- hasattr(target_language_model, "lm_head"):
856
- logger.info("Loading EAGLE LM head weights from the target model.")
857
- self.model.lm_head = target_language_model.lm_head
857
+ if self.vllm_config.speculative_config.method != "eagle3":
858
+ if hasattr(target_language_model, "lm_head"):
859
+ logger.info(
860
+ "Loading EAGLE LM head weights from the target model.")
861
+ self.model.lm_head = target_language_model.lm_head
862
+ else:
863
+ if (hasattr(self.model, "lm_head")
864
+ and hasattr(target_language_model, "lm_head")
865
+ and self.model.lm_head.weight.shape
866
+ == target_language_model.lm_head.weight.shape):
867
+ logger.info("Assuming the EAGLE head shares the same lm_head"
868
+ " with the target model.")
869
+ del self.model.lm_head
870
+ self.model.lm_head = target_language_model.lm_head
871
+ else:
872
+ logger.info(
873
+ "The EAGLE head's lm_head will be loaded separately"
874
+ " from the target model.")
875
876
@torch.inference_mode()
877
def dummy_run(
@@ -877,6 +894,31 @@ class EagleProposer:
894
inputs_embeds=inputs_embeds,
895
)
896
897
+ def _get_attention_metadata_builder(
898
+ self) -> list[AttentionMetadataBuilder]:
899
+ """Find and return the attention metadata builders for EAGLE layers.
900
+
901
+ Returns:
902
+ The metadata builders for EAGLE layers.
903
+
904
+ Raises:
905
+ AssertionError: If no metadata builders are found for EAGLE layers.
906
+ """
907
+ builder = None
908
+ chosen_layer = self.attn_layer_names[0]
909
+
910
+ for kv_cache_group in self.runner.attn_groups:
911
+ for attn_group in kv_cache_group:
912
+ if chosen_layer in attn_group.layer_names:
913
+ builder = attn_group.get_metadata_builder()
914
+ break
915
+ if builder is not None:
916
+ break
917
+
918
+ assert builder is not None, (
919
+ "Failed to find attention metadata builder for EAGLE layers.")
920
+ return builder
921
+
922
def validate_same_kv_cache_group(self,
923
kv_cache_config: KVCacheConfig) -> None:
924
"""
vllm/v1/worker/gpu_model_runner.py
+8
-3
@@ -1177,9 +1177,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1177
encoder_seq_lens=encoder_seq_lens,
1178
)
1179
1180
- if self.speculative_config and \
1181
- spec_decode_common_attn_metadata is None:
1182
- spec_decode_common_attn_metadata = common_attn_metadata
1180
+ if (self.speculative_config
1181
+ and spec_decode_common_attn_metadata is None):
1182
+ if isinstance(self.drafter, EagleProposer):
1183
+ if (self.drafter.attn_layer_names[0]
1184
+ in kv_cache_group_spec.layer_names):
1185
+ spec_decode_common_attn_metadata = common_attn_metadata
1186
+ else:
1187
+ spec_decode_common_attn_metadata = common_attn_metadata
1188
1189
for attn_group in self.attn_groups[kv_cache_group_id]:
1190
# Prepare for cascade attention if enabled & beneficial.