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 6, 00:45 UTC
[Bugfix] Pass drafter quant_config to ParallelLMHead in Eagle3 (#37280)
Signed-off-by: Matthias Gehre <matthias.gehre@amd.com>
Matthias Gehre committed
Mar 25, 2026 at 12:42 UTC
a889b7f58426e612b55aa67d7d243dbbceaf1184
3 files changed
+48
tests/model_executor/test_eagle_quantization.py
+45
@@ -139,6 +139,51 @@ def test_maybe_remap_kv_scale_name():
139
assert remapped in params_dict or remapped == name or remapped is None
140
141
142
+def test_eagle3_lm_head_receives_quant_config():
143
+ """Eagle3LlamaForCausalLM must pass quant_config to ParallelLMHead.
144
+
145
+ Without quant_config, quantized lm_head weights (e.g. INT8 per-channel)
146
+ in Eagle3 drafter checkpoints fail to load because ParallelLMHead doesn't
147
+ expect weight_packed tensors.
148
+ """
149
+ from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
150
+
151
+ mock_quant_config = Mock()
152
+
153
+ mock_hf_config = Mock()
154
+ mock_hf_config.draft_vocab_size = 1000
155
+ mock_hf_config.hidden_size = 256
156
+ mock_hf_config.vocab_size = 32000
157
+ mock_hf_config.logit_scale = 1.0
158
+
159
+ mock_vllm_config = Mock()
160
+ mock_vllm_config.speculative_config.draft_model_config.hf_config = mock_hf_config
161
+ mock_vllm_config.model_config.get_num_layers.return_value = 32
162
+ mock_vllm_config.speculative_config.parallel_drafting = False
163
+
164
+ with (
165
+ patch("vllm.model_executor.models.llama_eagle3.LlamaModel") as MockModel,
166
+ patch("vllm.model_executor.models.llama_eagle3.ParallelLMHead") as MockLMHead,
167
+ patch("vllm.model_executor.models.llama_eagle3.LogitsProcessor"),
168
+ patch(
169
+ "vllm.model_executor.models.llama_eagle3.get_draft_quant_config",
170
+ return_value=mock_quant_config,
171
+ ),
172
+ ):
173
+ MockModel.return_value.use_aux_hidden_state = True
174
+
175
+ Eagle3LlamaForCausalLM(vllm_config=mock_vllm_config)
176
+
177
+ MockLMHead.assert_called_once()
178
+ call_kwargs = MockLMHead.call_args.kwargs
179
+ assert "quant_config" in call_kwargs, (
180
+ "ParallelLMHead must receive quant_config for quantized lm_head weights"
181
+ )
182
+ assert call_kwargs["quant_config"] is mock_quant_config, (
183
+ "ParallelLMHead must receive the draft model's quant_config"
184
+ )
185
+
186
+
187
def test_load_weights_kv_scale_handling():
188
kv_scale_param = Mock()
189
kv_scale_param.weight_loader = Mock()
vllm/model_executor/models/llama_eagle3.py
+1
@@ -294,6 +294,7 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):
294
self.lm_head = ParallelLMHead(
295
self.config.draft_vocab_size,
296
self.config.hidden_size,
297
+ quant_config=get_draft_quant_config(vllm_config),
298
prefix=maybe_prefix(prefix, "lm_head"),
299
)
300
self.logits_processor = LogitsProcessor(
vllm/v1/spec_decode/eagle.py
+2
@@ -1390,6 +1390,8 @@ class SpecDecodeBaseProposer:
1390
)
1391
elif (
1392
hasattr(target_language_model, "lm_head")
1393
+ and hasattr(target_language_model.lm_head, "weight")
1394
+ and hasattr(self.model.lm_head, "weight")
1395
and isinstance(target_language_model.lm_head.weight, torch.Tensor)
1396
and isinstance(self.model.lm_head.weight, torch.Tensor)
1397
# TODO: Offload to CPU for comparison to avoid extra GPU memory