[Bugfix] [Gemma4] Fix Gemma4 MTP draft model layers ignoring quant_config (#47091)
Signed-off-by: Ayushman Singh <40520701+ayush1399@users.noreply.github.com> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
Ayushman Singh committed
Jul 6, 2026 at 09:04 UTC
7a90eb98ab70dcca8360214eb3cc08c999137861
1 file changed
+13
-4
vllm/model_executor/models/gemma4_mtp.py
+13
-4
@@ -51,6 +51,7 @@ from .utils import (
51
AutoWeightsLoader,
52
WeightsMapper,
53
extract_layer_index,
54
+ get_draft_quant_config,
55
maybe_prefix,
56
)
57
@@ -182,14 +183,14 @@ class Gemma4MTPAttention(nn.Module):
183
hidden_size,
184
self.total_num_heads * self.head_dim,
185
bias=config.attention_bias,
185
- quant_config=None,
186
+ quant_config=quant_config,
187
prefix=f"{prefix}.q_proj",
188
)
189
self.o_proj = RowParallelLinear(
190
self.total_num_heads * self.head_dim,
191
hidden_size,
192
bias=config.attention_bias,
192
- quant_config=None,
193
+ quant_config=quant_config,
194
prefix=f"{prefix}.o_proj",
195
)
196
self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
@@ -304,7 +305,7 @@ class Gemma4MTPDecoderLayer(nn.Module):
305
hidden_size=self.hidden_size,
306
intermediate_size=text_config.intermediate_size,
307
hidden_activation=text_config.hidden_activation,
307
- quant_config=None,
308
+ quant_config=quant_config,
309
prefix=f"{prefix}.mlp",
310
)
311
@@ -357,7 +358,9 @@ class Gemma4MultiTokenPredictor(nn.Module):
358
359
config = vllm_config.speculative_config.draft_model_config.hf_config
360
text_config = _get_text_config(config)
361
+ quant_config = get_draft_quant_config(vllm_config)
362
self.config = text_config
363
+ self.quant_config = quant_config
364
365
self.hidden_size = text_config.hidden_size
366
self.backbone_hidden_size = getattr(
@@ -369,6 +372,8 @@ class Gemma4MultiTokenPredictor(nn.Module):
372
self.embed_tokens = VocabParallelEmbedding(
373
self.vocab_size,
374
self.hidden_size,
375
+ quant_config=quant_config,
376
+ prefix=f"{prefix}.embed_tokens",
377
)
378
379
self.pre_projection = ColumnParallelLinear(
@@ -376,6 +381,7 @@ class Gemma4MultiTokenPredictor(nn.Module):
381
self.hidden_size,
382
bias=False,
383
gather_output=True,
384
+ quant_config=quant_config,
385
prefix=f"{prefix}.pre_projection",
386
)
387
@@ -384,6 +390,7 @@ class Gemma4MultiTokenPredictor(nn.Module):
390
self.backbone_hidden_size,
391
bias=False,
392
input_is_parallel=False,
393
+ quant_config=quant_config,
394
prefix=f"{prefix}.post_projection",
395
)
396
@@ -391,7 +398,7 @@ class Gemma4MultiTokenPredictor(nn.Module):
398
Gemma4MTPDecoderLayer(
399
text_config,
400
cache_config=vllm_config.cache_config,
394
- quant_config=vllm_config.quant_config,
401
+ quant_config=quant_config,
402
prefix=f"{prefix}.layers.{idx}",
403
)
404
for idx in range(self.num_mtp_layers)
@@ -473,6 +480,7 @@ class Gemma4MTP(nn.Module):
480
super().__init__()
481
config = vllm_config.speculative_config.draft_model_config.hf_config
482
text_config = _get_text_config(config)
483
+ self.quant_config = get_draft_quant_config(vllm_config)
484
self.config = config
485
self._stable_full_lm_head_weight: torch.Tensor | None = None
486
@@ -488,6 +496,7 @@ class Gemma4MTP(nn.Module):
496
self.lm_head = ParallelLMHead(
497
text_config.vocab_size,
498
text_config.hidden_size,
499
+ quant_config=self.quant_config,
500
prefix=maybe_prefix(prefix, "lm_head"),
501
)
502
if getattr(config, "tie_word_embeddings", True):