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 11, 22:30 UTC
[Bugfix][llama4_eagle] Fix missing 'lm_head' attribute (#29926)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
Dec 5, 2025 at 13:57 UTC
962d703818c038b6ea2ae18eab8061a6f2fb8f65
2 files changed
+16
-3
tests/v1/e2e/test_spec_decode.py
+5
-1
@@ -402,7 +402,11 @@ def test_eagle_correctness(
402
# Scout requires default backend selection
403
# because vision encoder has head_dim 88 being incompatible
404
# with FLASH_ATTN and needs to fall back to Flex Attn
405
- pass
405
+
406
+ # pass if not ROCm
407
+ if current_platform.is_rocm():
408
+ # TODO: Enable Flex Attn for spec_decode on ROCm
409
+ pytest.skip("Flex Attn for spec_decode not supported on ROCm currently")
410
else:
411
m.setenv("VLLM_MLA_DISABLE", "1")
412
m.setenv("VLLM_ATTENTION_BACKEND", attn_backend)
vllm/model_executor/models/llama4_eagle.py
+11
-2
@@ -28,7 +28,10 @@ from vllm.model_executor.layers.layernorm import RMSNorm
28
from vllm.model_executor.layers.logits_processor import LogitsProcessor
29
from vllm.model_executor.layers.quantization import QuantizationConfig
30
from vllm.model_executor.layers.quantization.torchao import TorchAOConfig
31
-from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
31
+from vllm.model_executor.layers.vocab_parallel_embedding import (
32
+ ParallelLMHead,
33
+ VocabParallelEmbedding,
34
+)
35
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
36
from vllm.model_executor.models.llama4 import Llama4DecoderLayer, Llama4ForCausalLM
37
from vllm.model_executor.models.utils import extract_layer_index
@@ -182,6 +185,12 @@ class EagleLlama4ForCausalLM(Llama4ForCausalLM):
185
self.config.vocab_size, scale=logit_scale
186
)
187
188
+ self.lm_head = ParallelLMHead(
189
+ self.config.draft_vocab_size,
190
+ self.config.hidden_size,
191
+ prefix=maybe_prefix(prefix, "lm_head"),
192
+ )
193
+
194
# Set MoE hyperparameters
195
self.set_moe_parameters()
196
@@ -211,6 +220,6 @@ class EagleLlama4ForCausalLM(Llama4ForCausalLM):
220
loader = AutoWeightsLoader(
221
self,
222
# lm_head is tied with target model (Llama4ForCausalLM)
214
- skip_prefixes=(["lm_head."]),
223
+ skip_prefixes=([]),
224
)
225
loader.load_weights(map(transform, weights))