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, 22:35 UTC
[Model] [Quantization] Support quantization for Gemma3n (#21974)
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Kyle Sayers committed
Aug 1, 2025 at 01:45 UTC
0f46a780d4f53b8564a37370f9f068cdf4e69604
1 file changed
+12
-4
vllm/model_executor/models/gemma3n.py
+12
-4
@@ -46,6 +46,7 @@ from vllm.model_executor.model_loader.weight_utils import (
46
from vllm.model_executor.sampling_metadata import SamplingMetadata
47
from vllm.sequence import IntermediateTensors
48
49
+from .interfaces import SupportsQuant
50
from .utils import (AutoWeightsLoader, extract_layer_index,
51
is_pp_missing_parameter, make_layers, maybe_prefix)
52
@@ -68,6 +69,7 @@ class Gemma3nAltUp(nn.Module):
69
altup_num_inputs: int,
70
altup_coef_clip: float,
71
altup_active_idx: int,
72
+ quant_config: QuantizationConfig,
73
prefix: str,
74
):
75
super().__init__()
@@ -80,6 +82,7 @@ class Gemma3nAltUp(nn.Module):
82
altup_num_inputs,
83
altup_num_inputs,
84
bias=False,
85
+ quant_config=quant_config,
86
prefix=f"{prefix}.correction_coefs",
87
return_bias=False,
88
)
@@ -87,6 +90,7 @@ class Gemma3nAltUp(nn.Module):
90
altup_num_inputs,
91
altup_num_inputs**2,
92
bias=False,
93
+ quant_config=quant_config,
94
prefix=f"{prefix}.prediction_coefs",
95
return_bias=False,
96
)
@@ -94,6 +98,7 @@ class Gemma3nAltUp(nn.Module):
98
hidden_size,
99
altup_num_inputs,
100
bias=False,
101
+ quant_config=quant_config,
102
prefix=f"{prefix}.modality_router",
103
return_bias=False,
104
)
@@ -400,6 +405,7 @@ class Gemma3nDecoderLayer(nn.Module):
405
altup_num_inputs=config.altup_num_inputs,
406
altup_coef_clip=config.altup_coef_clip,
407
altup_active_idx=config.altup_active_idx,
408
+ quant_config=quant_config,
409
prefix=f"{prefix}.altup",
410
)
411
self.self_attn = Gemma3nAttention(
@@ -527,7 +533,7 @@ class Gemma3nDecoderLayer(nn.Module):
533
534
535
@support_torch_compile
530
-class Gemma3nTextModel(nn.Module):
536
+class Gemma3nTextModel(nn.Module, SupportsQuant):
537
538
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
539
super().__init__()
@@ -540,6 +546,7 @@ class Gemma3nTextModel(nn.Module):
546
self.embed_tokens = VocabParallelEmbedding(
547
config.vocab_size,
548
config.hidden_size,
549
+ quant_config=quant_config,
550
prefix=f"{prefix}.embed_tokens",
551
)
552
self.embed_scale = torch.tensor(
@@ -549,6 +556,7 @@ class Gemma3nTextModel(nn.Module):
556
self.embed_tokens_per_layer = VocabParallelEmbedding(
557
config.vocab_size_per_layer_input,
558
config.num_hidden_layers * config.hidden_size_per_layer_input,
559
+ quant_config=quant_config,
560
prefix=f"{prefix}.per_layer_embed_tokens",
561
)
562
self.embed_scale_per_layer = torch.tensor(
@@ -582,7 +590,7 @@ class Gemma3nTextModel(nn.Module):
590
gather_output=True,
591
return_bias=False,
592
quant_config=quant_config,
585
- prefix=f"{prefix}.{idx-1}.altup_projections",
593
+ prefix=f"{prefix}.altup_projections.{idx-1}",
594
) for idx in range(1, self.config.altup_num_inputs)
595
])
596
self.altup_unembed_projections = nn.ModuleList([
@@ -593,7 +601,7 @@ class Gemma3nTextModel(nn.Module):
601
gather_output=True,
602
return_bias=False,
603
quant_config=quant_config,
596
- prefix=f"{prefix}.{idx-1}.altup_unembed_projections",
604
+ prefix=f"{prefix}.altup_unembed_projections.{idx-1}",
605
) for idx in range(1, self.config.altup_num_inputs)
606
])
607
@@ -774,7 +782,7 @@ class Gemma3nModel(nn.Module):
782
**kwargs)
783
784
777
-class Gemma3nForConditionalGeneration(nn.Module):
785
+class Gemma3nForConditionalGeneration(nn.Module, SupportsQuant):
786
packed_modules_mapping = {
787
"qkv_proj": [
788
"q_proj",