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 15, 04:45 UTC
Gemma3n (Text-only) (#20134)
Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com> Signed-off-by: Roger Wang <hey@rogerw.me> Co-authored-by: Roger Wang <hey@rogerw.me>
Robert Shaw committed
Jun 27, 2025 at 03:16 UTC
d1c956dc0f8b64af7a43e23f9fc2850756dea1cb
5 files changed
+870
docs/models/supported_models.md
+4
@@ -336,6 +336,7 @@ Specified using `--task generate`.
336
| `GemmaForCausalLM` | Gemma | `google/gemma-2b`, `google/gemma-1.1-2b-it`, etc. | ✅︎ | ✅︎ | ✅︎ |
337
| `Gemma2ForCausalLM` | Gemma 2 | `google/gemma-2-9b`, `google/gemma-2-27b`, etc. | ✅︎ | ✅︎ | ✅︎ |
338
| `Gemma3ForCausalLM` | Gemma 3 | `google/gemma-3-1b-it`, etc. | ✅︎ | ✅︎ | ✅︎ |
339
+| `Gemma3nForConditionalGeneration` | Gemma 3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | | ✅︎ |
340
| `GlmForCausalLM` | GLM-4 | `THUDM/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ | ✅︎ |
341
| `Glm4ForCausalLM` | GLM-4-0414 | `THUDM/GLM-4-32B-0414`, etc. | ✅︎ | ✅︎ | ✅︎ |
342
| `GPT2LMHeadModel` | GPT-2 | `gpt2`, `gpt2-xl`, etc. | | ✅︎ | ✅︎ |
@@ -392,6 +393,9 @@ Specified using `--task generate`.
393
!!! note
394
Currently, the ROCm version of vLLM supports Mistral and Mixtral only for context lengths up to 4096.
395
396
+!!! note
397
+ Only text inputs are currently supported for `Gemma3nForConditionalGeneration`. To use this model, please upgrade Hugging Face Transformers to version 4.53.0.
398
+
399
### Pooling Models
400
401
See [this page](./pooling_models.md) for more information on how to use pooling models.
tests/models/registry.py
+2
@@ -164,6 +164,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
164
"GemmaForCausalLM": _HfExamplesInfo("google/gemma-1.1-2b-it"),
165
"Gemma2ForCausalLM": _HfExamplesInfo("google/gemma-2-9b"),
166
"Gemma3ForCausalLM": _HfExamplesInfo("google/gemma-3-1b-it"),
167
+ "Gemma3nForConditionalGeneration": _HfExamplesInfo("google/gemma-3n-E2B-it", # noqa: E501
168
+ min_transformers_version="4.53"),
169
"GlmForCausalLM": _HfExamplesInfo("THUDM/glm-4-9b-chat-hf"),
170
"Glm4ForCausalLM": _HfExamplesInfo("THUDM/GLM-4-9B-0414"),
171
"GPT2LMHeadModel": _HfExamplesInfo("openai-community/gpt2",
vllm/model_executor/layers/activation.py
+51
@@ -135,6 +135,57 @@ class MulAndSilu(CustomOp):
135
# def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
136
137
138
+@CustomOp.register("gelu_and_mul_sparse")
139
+class GeluAndMulSparse(CustomOp):
140
+ """An activation function for GeluAndMulSparse.
141
+ This activation function is used in Gemma3n. It computes:
142
+ up_proj = self.up_proj(x)
143
+ gate_proj = self.gate_proj(x)
144
+ gate_proj = self._gaussian_topk(gate_proj) # sparsity
145
+ activations = self.act_fn(gate_proj) # gelu
146
+ down_proj = self.down_proj(activations * up_proj)
147
+ Shapes:
148
+ x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
149
+ return: (num_tokens, d) or (batch_size, seq_len, d)
150
+ """
151
+
152
+ def __init__(self, activation_sparsity: float, approximate: str = "none"):
153
+ super().__init__()
154
+ # Gelu.
155
+ self.approximate = approximate
156
+ if approximate not in ("none", "tanh"):
157
+ raise ValueError(f"Unknown approximate mode: {approximate}")
158
+
159
+ # Sparsity.
160
+ if activation_sparsity == 0.0:
161
+ raise ValueError(
162
+ "activation_sparsity is 0.0. Please use GeluAndMul.")
163
+ target_sparsity_tensor = torch.tensor(activation_sparsity,
164
+ dtype=torch.float32)
165
+ normal_dist = torch.distributions.normal.Normal(0, 1)
166
+ self.std_multiplier = normal_dist.icdf(target_sparsity_tensor)
167
+
168
+ def _gaussian_topk(self, x: torch.Tensor) -> torch.Tensor:
169
+ """Get % sparse percentile of the Gaussian distribution."""
170
+ # NOTE(rob): for TP>1, we could all-gather to get the means/std.
171
+ # But we do not do this because in expectation they are the same
172
+ # and in practice the eval scores are good without gathering.
173
+ mean = torch.mean(x, dim=-1, keepdim=True)
174
+ std = torch.std(x, dim=-1, keepdim=True, unbiased=False)
175
+ cutoff_x = mean + std * self.std_multiplier
176
+ return nn.functional.relu(x - cutoff_x)
177
+
178
+ def forward_native(self, x: torch.Tensor) -> torch.Tensor:
179
+ """PyTorch-native implementation equivalent to forward()."""
180
+ d = x.shape[-1] // 2
181
+ out = self._gaussian_topk(x[..., :d])
182
+ out = F.gelu(out, approximate=self.approximate)
183
+ return out * x[..., d:]
184
+
185
+ def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
186
+ return self.forward_native(x)
187
+
188
+
189
@CustomOp.register("gelu_and_mul")
190
class GeluAndMul(CustomOp):
191
"""An activation function for GeGLU.
vllm/model_executor/models/gemma3n.py
new
+811
@@ -0,0 +1,811 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# Copyright 2025 The vLLM team.
4
+# Copyright 2025 Google Inc. HuggingFace Inc. team. All rights reserved.
5
+#
6
+#
7
+# Licensed under the Apache License, Version 2.0 (the "License");
8
+# you may not use this file except in compliance with the License.
9
+# You may obtain a copy of the License at
10
+#
11
+# http://www.apache.org/licenses/LICENSE-2.0
12
+#
13
+# Unless required by applicable law or agreed to in writing, software
14
+# distributed under the License is distributed on an "AS IS" BASIS,
15
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16
+# See the License for the specific language governing permissions and
17
+# limitations under the License.
18
+from collections.abc import Iterable
19
+from typing import Optional, Union
20
+
21
+import torch
22
+from torch import nn
23
+from transformers.models.gemma3n.configuration_gemma3n import Gemma3nTextConfig
24
+
25
+from vllm.attention import Attention
26
+from vllm.compilation.decorators import support_torch_compile
27
+from vllm.config import CacheConfig, VllmConfig
28
+from vllm.distributed import get_tensor_model_parallel_world_size
29
+from vllm.logger import init_logger
30
+from vllm.model_executor.layers.activation import (_ACTIVATION_REGISTRY,
31
+ GeluAndMul,
32
+ GeluAndMulSparse)
33
+from vllm.model_executor.layers.layernorm import RMSNorm
34
+from vllm.model_executor.layers.linear import (ColumnParallelLinear,
35
+ MergedColumnParallelLinear,
36
+ QKVParallelLinear,
37
+ ReplicatedLinear,
38
+ RowParallelLinear)
39
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
40
+from vllm.model_executor.layers.quantization import QuantizationConfig
41
+from vllm.model_executor.layers.rotary_embedding import get_rope
42
+from vllm.model_executor.layers.vocab_parallel_embedding import (
43
+ VocabParallelEmbedding)
44
+from vllm.model_executor.model_loader.weight_utils import (
45
+ default_weight_loader, maybe_remap_kv_scale_name)
46
+from vllm.model_executor.sampling_metadata import SamplingMetadata
47
+from vllm.sequence import IntermediateTensors
48
+
49
+from .utils import (AutoWeightsLoader, extract_layer_index,
50
+ is_pp_missing_parameter, make_layers, maybe_prefix)
51
+
52
+logger = init_logger(__name__)
53
+
54
+
55
+class Gemma3nAltUp(nn.Module):
56
+ """Alternating updates (Altup)
57
+ The AltUp module wraps transformer layers. The `predict` step modifies the
58
+ input to the transformer layer, and the `correct` step propagates the output
59
+ of the transformer layer to the sparsely updated dimensions.
60
+ See more in the research paper:
61
+ https://proceedings.neurips.cc/paper_files/paper/2023/file/f2059277ac6ce66e7e5543001afa8bb5-Paper-Conference.pdf
62
+ """
63
+
64
+ def __init__(
65
+ self,
66
+ hidden_size: int,
67
+ rms_norm_eps: float,
68
+ altup_num_inputs: int,
69
+ altup_coef_clip: float,
70
+ altup_active_idx: int,
71
+ prefix: str,
72
+ ):
73
+ super().__init__()
74
+
75
+ self.altup_num_inputs = altup_num_inputs
76
+ self.altup_active_idx = altup_active_idx
77
+ self.altup_coef_clip = altup_coef_clip
78
+
79
+ self.correction_coefs = ReplicatedLinear(
80
+ altup_num_inputs,
81
+ altup_num_inputs,
82
+ bias=False,
83
+ prefix=f"{prefix}.correction_coefs",
84
+ return_bias=False,
85
+ )
86
+ self.prediction_coefs = ReplicatedLinear(
87
+ altup_num_inputs,
88
+ altup_num_inputs**2,
89
+ bias=False,
90
+ prefix=f"{prefix}.prediction_coefs",
91
+ return_bias=False,
92
+ )
93
+ self.modality_router = ReplicatedLinear(
94
+ hidden_size,
95
+ altup_num_inputs,
96
+ bias=False,
97
+ prefix=f"{prefix}.modality_router",
98
+ return_bias=False,
99
+ )
100
+ self.router_norm = RMSNorm(
101
+ hidden_size=hidden_size,
102
+ eps=rms_norm_eps,
103
+ )
104
+ self.router_input_scale = torch.tensor(
105
+ hidden_size**-1.0, dtype=self.modality_router.weight.dtype)
106
+ self.correct_output_scale = nn.Parameter(
107
+ torch.zeros(hidden_size, dtype=torch.float32))
108
+
109
+ def _compute_router_modalities(self, x: torch.Tensor) -> torch.Tensor:
110
+ router_inputs = self.router_norm(x) * self.router_input_scale
111
+ routed = self.modality_router(router_inputs)
112
+ return torch.tanh(routed.float()).type_as(x)
113
+
114
+ def scale_corrected_output(self, corrected: torch.Tensor) -> torch.Tensor:
115
+ return (corrected.type_as(self.correct_output_scale) *
116
+ self.correct_output_scale).type_as(corrected)
117
+
118
+ def predict(self, hidden_states: torch.Tensor) -> torch.Tensor:
119
+ # hidden: [altup_num_inputs, num_tokens, hidden_size]
120
+ # modalities: [num_tokens, num_altup_inputs]
121
+ # all_coefs: [num_tokens, num_altup_inputs ** 2]
122
+ modalities = self._compute_router_modalities(
123
+ hidden_states[self.altup_active_idx])
124
+ all_coefs = self.prediction_coefs(modalities)
125
+
126
+ # Reshape and transpose the 2D matrix for the matmul.
127
+ # all_coefs_T: [num_tokens, num_altup_inputs, num_altup_inputs]
128
+ all_coefs_T = all_coefs.reshape(
129
+ -1,
130
+ self.altup_num_inputs,
131
+ self.altup_num_inputs,
132
+ ).permute(0, 2, 1)
133
+
134
+ # hidden_states to [num_tokens, hidden_size, altup_num_inputs]
135
+ predictions = torch.matmul(hidden_states.permute(1, 2, 0), all_coefs_T)
136
+ # [altup_num_inputs, num_tokens, hidden_size]
137
+ predictions = predictions.permute(2, 0, 1)
138
+ predictions += hidden_states
139
+ return predictions.contiguous()
140
+
141
+ def correct(self, predictions: torch.Tensor,
142
+ activated: torch.Tensor) -> torch.Tensor:
143
+ # predictions: [altup_num_inputs, num_tokens, hidden_size]
144
+ # activated: [num_tokens, hidden_size]
145
+ # modalities: [num_tokens, altup_num_inputs]
146
+ modalities = self._compute_router_modalities(activated)
147
+ # innovation: [num_tokens, altup_num_inputs]
148
+ innovation = activated - predictions[self.altup_active_idx]
149
+ # innovation: [altup_num_inputs, num_tokens, hidden_size]
150
+ innovation = innovation.repeat(self.altup_num_inputs, 1, 1)
151
+
152
+ # Permute to [altup_num_inputs, num_tokens] as the last dim
153
+ # is a scalar applied to each altup input and expand on
154
+ # num_tokens dim for broadcastability over hidden_size.
155
+ # all_coefs: [num_tokens, altup_num_inputs]
156
+ all_coefs = self.correction_coefs(modalities) + 1.0
157
+ # all_coefs: [altup_num_inputs, num_tokens, 1]
158
+ all_coefs = all_coefs.T.unsqueeze(-1)
159
+
160
+ # Elementwise (broadcast over hidden_size).
161
+ corrected = torch.mul(innovation, all_coefs)
162
+ corrected += predictions
163
+
164
+ return corrected.contiguous()
165
+
166
+
167
+class Gemma3nLaurelBlock(nn.Module):
168
+ """Learned Augmented Residual Layer"""
169
+
170
+ def __init__(self, hidden_size: int, laurel_rank: int, rms_norm_eps: float,
171
+ prefix: str):
172
+ super().__init__()
173
+
174
+ self.linear_left = ColumnParallelLinear(
175
+ hidden_size,
176
+ laurel_rank,
177
+ bias=False,
178
+ prefix=f"{prefix}.linear_left",
179
+ return_bias=False,
180
+ )
181
+ self.linear_right = RowParallelLinear(laurel_rank,
182
+ hidden_size,
183
+ bias=False,
184
+ prefix=f"{prefix}.linear_right",
185
+ return_bias=False)
186
+ self.post_laurel_norm = RMSNorm(
187
+ hidden_size=hidden_size,
188
+ eps=rms_norm_eps,
189
+ )
190
+
191
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
192
+ laurel_x = self.linear_left(x)
193
+ laurel_x = self.linear_right(laurel_x)
194
+ normed_laurel_x = self.post_laurel_norm(laurel_x)
195
+ return x + normed_laurel_x
196
+
197
+
198
+class Gemma3nMLP(nn.Module):
199
+
200
+ def __init__(
201
+ self,
202
+ hidden_size: int,
203
+ intermediate_size: int,
204
+ hidden_activation: str,
205
+ activation_sparsity: float = 0.0,
206
+ quant_config: Optional[QuantizationConfig] = None,
207
+ prefix: str = "",
208
+ ) -> None:
209
+ super().__init__()
210
+ self.gate_up_proj = MergedColumnParallelLinear(
211
+ hidden_size,
212
+ [intermediate_size] * 2,
213
+ bias=False,
214
+ quant_config=quant_config,
215
+ prefix=f"{prefix}.gate_up_proj",
216
+ )
217
+ self.down_proj = RowParallelLinear(
218
+ intermediate_size,
219
+ hidden_size,
220
+ bias=False,
221
+ quant_config=quant_config,
222
+ prefix=f"{prefix}.down_proj",
223
+ )
224
+ if hidden_activation != "gelu_pytorch_tanh":
225
+ raise ValueError(
226
+ "Gemma3 uses `gelu_pytorch_tanh` as the hidden activation "
227
+ "function. Please set `hidden_act` and `hidden_activation` to "
228
+ "`gelu_pytorch_tanh`.")
229
+
230
+ self.act_fn = GeluAndMulSparse(
231
+ activation_sparsity=activation_sparsity,
232
+ approximate="tanh") if activation_sparsity > 0.0 else GeluAndMul(
233
+ approximate="tanh")
234
+
235
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
236
+ gate_up, _ = self.gate_up_proj(x)
237
+ x = self.act_fn(gate_up)
238
+ x, _ = self.down_proj(x)
239
+ return x
240
+
241
+
242
+class Gemma3nAttention(nn.Module):
243
+
244
+ def __init__(self,
245
+ config: Gemma3nTextConfig,
246
+ hidden_size: int,
247
+ num_heads: int,
248
+ num_kv_heads: int,
249
+ head_dim: int,
250
+ max_position_embeddings: int,
251
+ cache_config: Optional[CacheConfig] = None,
252
+ quant_config: Optional[QuantizationConfig] = None,
253
+ prefix: str = "") -> None:
254
+ super().__init__()
255
+ self.config = config
256
+ self.hidden_size = hidden_size
257
+ tp_size = get_tensor_model_parallel_world_size()
258
+ self.total_num_heads = num_heads
259
+ assert self.total_num_heads % tp_size == 0
260
+ self.num_heads = self.total_num_heads // tp_size
261
+ self.total_num_kv_heads = num_kv_heads
262
+ if self.total_num_kv_heads >= tp_size:
263
+ # Number of KV heads is greater than TP size, so we partition
264
+ # the KV heads across multiple tensor parallel GPUs.
265
+ assert self.total_num_kv_heads % tp_size == 0
266
+ else:
267
+ # Number of KV heads is less than TP size, so we replicate
268
+ # the KV heads across multiple tensor parallel GPUs.
269
+ assert tp_size % self.total_num_kv_heads == 0
270
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
271
+ self.head_dim = head_dim
272
+ self.q_size = self.num_heads * self.head_dim
273
+ self.kv_size = self.num_kv_heads * self.head_dim
274
+
275
+ self.qkv_proj = QKVParallelLinear(
276
+ hidden_size,
277
+ self.head_dim,
278
+ self.total_num_heads,
279
+ self.total_num_kv_heads,
280
+ bias=config.attention_bias,
281
+ quant_config=quant_config,
282
+ prefix=f"{prefix}.qkv_proj",
283
+ )
284
+ self.o_proj = RowParallelLinear(
285
+ self.total_num_heads * self.head_dim,
286
+ hidden_size,
287
+ bias=config.attention_bias,
288
+ quant_config=quant_config,
289
+ prefix=f"{prefix}.o_proj",
290
+ )
291
+ self.q_norm = RMSNorm(hidden_size=self.head_dim,
292
+ eps=config.rms_norm_eps)
293
+ self.k_norm = RMSNorm(hidden_size=self.head_dim,
294
+ eps=config.rms_norm_eps)
295
+ self.v_norm = RMSNorm(hidden_size=self.head_dim,
296
+ eps=config.rms_norm_eps,
297
+ has_weight=False)
298
+
299
+ layer_idx = extract_layer_index(prefix)
300
+ if config.layer_types[layer_idx] == "sliding_attention":
301
+ self.sliding_window = config.sliding_window
302
+ rope_theta = config.rope_local_base_freq
303
+ rope_scaling = {"rope_type": "default"}
304
+ else:
305
+ self.sliding_window = None
306
+ rope_theta = config.rope_theta
307
+ rope_scaling = config.rope_scaling
308
+
309
+ first_kv_shared_layer_idx = (config.num_hidden_layers -
310
+ config.num_kv_shared_layers)
311
+ self.is_kv_shared = layer_idx >= first_kv_shared_layer_idx
312
+
313
+ if self.is_kv_shared:
314
+ # Last full attention layer is 1 before sharing
315
+ # Last sliding attention layer is 2 before sharing
316
+ offset = 2 if self.sliding_window is not None else 1
317
+ kv_shared_layer_index = first_kv_shared_layer_idx - offset
318
+ kv_sharing_target_layer_name = f"model.language_model.layers.{kv_shared_layer_index}.self_attn.attn" # noqa: E501
319
+ else:
320
+ kv_sharing_target_layer_name = None
321
+
322
+ self.rotary_emb = get_rope(
323
+ self.head_dim,
324
+ rotary_dim=self.head_dim,
325
+ max_position=max_position_embeddings,
326
+ base=rope_theta,
327
+ is_neox_style=True,
328
+ rope_scaling=rope_scaling,
329
+ )
330
+
331
+ self.attn = Attention(
332
+ num_heads=self.num_heads,
333
+ head_size=self.head_dim,
334
+ scale=1.0,
335
+ num_kv_heads=self.num_kv_heads,
336
+ cache_config=cache_config,
337
+ quant_config=quant_config,
338
+ per_layer_sliding_window=self.sliding_window,
339
+ kv_sharing_target_layer_name=kv_sharing_target_layer_name,
340
+ prefix=f"{prefix}.attn")
341
+
342
+ def forward(
343
+ self,
344
+ positions: torch.Tensor,
345
+ hidden_states: torch.Tensor,
346
+ **kwargs,
347
+ ) -> torch.Tensor:
348
+ qkv, _ = self.qkv_proj(hidden_states)
349
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
350
+
351
+ q = q.unflatten(-1, (self.num_heads, self.head_dim))
352
+ q = self.q_norm(q)
353
+ q = q.flatten(-2, -1)
354
+ k = k.unflatten(-1, (self.num_kv_heads, self.head_dim))
355
+ k = self.k_norm(k)
356
+ k = k.flatten(-2, -1)
357
+ v = v.unflatten(-1, (self.num_kv_heads, self.head_dim))
358
+ v = self.v_norm(v)
359
+ v = v.flatten(-2, -1)
360
+
361
+ q, k = self.rotary_emb(positions, q, k)
362
+ attn_output = self.attn(q, k, v)
363
+
364
+ output, _ = self.o_proj(attn_output)
365
+ return output
366
+
367
+
368
+class Gemma3nDecoderLayer(nn.Module):
369
+
370
+ def __init__(
371
+ self,
372
+ config: Gemma3nTextConfig,
373
+ cache_config: Optional[CacheConfig] = None,
374
+ quant_config: Optional[QuantizationConfig] = None,
375
+ prefix: str = "",
376
+ ) -> None:
377
+ super().__init__()
378
+ self.altup_active_idx = config.altup_active_idx
379
+ assert config.altup_correct_scale
380
+
381
+ self.altup = Gemma3nAltUp(
382
+ hidden_size=config.hidden_size,
383
+ rms_norm_eps=config.rms_norm_eps,
384
+ altup_num_inputs=config.altup_num_inputs,
385
+ altup_coef_clip=config.altup_coef_clip,
386
+ altup_active_idx=config.altup_active_idx,
387
+ prefix=f"{prefix}.altup",
388
+ )
389
+ self.self_attn = Gemma3nAttention(
390
+ config=config,
391
+ hidden_size=config.hidden_size,
392
+ num_heads=config.num_attention_heads,
393
+ num_kv_heads=config.num_key_value_heads,
394
+ head_dim=config.head_dim,
395
+ max_position_embeddings=config.max_position_embeddings,
396
+ cache_config=cache_config,
397
+ quant_config=quant_config,
398
+ prefix=f"{prefix}.self_attn",
399
+ )
400
+ self.mlp = Gemma3nMLP(
401
+ hidden_size=config.hidden_size,
402
+ # NOTE: Matformer https://github.com/huggingface/transformers/blob/a52478253bbe522a420e88ea3940d4d98a935300/src/transformers/models/gemma3n/modular_gemma3n.py#L258 # noqa: E501
403
+ intermediate_size=config.intermediate_size[extract_layer_index(
404
+ prefix)],
405
+ hidden_activation=config.hidden_activation,
406
+ quant_config=quant_config,
407
+ activation_sparsity=config.activation_sparsity_pattern[
408
+ extract_layer_index(prefix)],
409
+ prefix=f"{prefix}.mlp",
410
+ )
411
+ self.laurel = Gemma3nLaurelBlock(
412
+ hidden_size=config.hidden_size,
413
+ laurel_rank=config.laurel_rank,
414
+ rms_norm_eps=config.rms_norm_eps,
415
+ prefix=f"{prefix}.laurel",
416
+ )
417
+
418
+ # NOTE(rob): should be ColumnParallelLinear and RowParallelLinear
419
+ # But, we need to add per_layer_input_gate(x) to per_layer_input.
420
+ # per_layer_input cannot be sharded, so we replicate for now.
421
+ self.per_layer_input_gate = ReplicatedLinear(
422
+ config.hidden_size,
423
+ config.hidden_size_per_layer_input,
424
+ bias=False,
425
+ prefix=f"{prefix}.per_layer_input_gate",
426
+ return_bias=False,
427
+ )
428
+ self.per_layer_projection = ReplicatedLinear(
429
+ config.hidden_size_per_layer_input,
430
+ config.hidden_size,
431
+ bias=False,
432
+ prefix=f"{prefix}.per_layer_projection",
433
+ return_bias=False,
434
+ )
435
+
436
+ # LayerNorms.
437
+ self.input_layernorm = RMSNorm(
438
+ config.hidden_size,
439
+ eps=config.rms_norm_eps,
440
+ )
441
+ self.post_attention_layernorm = RMSNorm(
442
+ config.hidden_size,
443
+ eps=config.rms_norm_eps,
444
+ )
445
+ self.pre_feedforward_layernorm = RMSNorm(
446
+ config.hidden_size,
447
+ eps=config.rms_norm_eps,
448
+ )
449
+ self.post_feedforward_layernorm = RMSNorm(
450
+ config.hidden_size,
451
+ eps=config.rms_norm_eps,
452
+ )
453
+ self.post_per_layer_input_norm = RMSNorm(
454
+ config.hidden_size,
455
+ eps=config.rms_norm_eps,
456
+ )
457
+
458
+ self.act_fn = _ACTIVATION_REGISTRY[config.hidden_activation]
459
+
460
+ def forward(
461
+ self,
462
+ positions: torch.Tensor,
463
+ hidden_states: torch.Tensor,
464
+ per_layer_input: torch.Tensor,
465
+ **kwargs,
466
+ ) -> tuple[torch.Tensor, torch.Tensor]:
467
+
468
+ # ActUp (predict).
469
+ predictions = self.altup.predict(hidden_states)
470
+ active_prediction = predictions[self.altup_active_idx]
471
+ active_prediction_normed = self.input_layernorm(active_prediction)
472
+ laurel_output = self.laurel(active_prediction_normed)
473
+
474
+ # Attention.
475
+ attn = self.self_attn(
476
+ positions=positions,
477
+ hidden_states=active_prediction_normed,
478
+ **kwargs,
479
+ )
480
+ attn = self.post_attention_layernorm(attn)
481
+ attn_gated = attn + active_prediction
482
+ attn_laurel = (attn_gated + laurel_output) / torch.sqrt(
483
+ torch.tensor(2.0))
484
+
485
+ # MLP.
486
+ attn_norm = self.pre_feedforward_layernorm(attn_laurel)
487
+ attn_ffw = self.mlp(attn_norm)
488
+ attn_ffw_norm = self.post_feedforward_layernorm(attn_ffw)
489
+ attn_ffw_laurel_gated = attn_laurel + attn_ffw_norm
490
+
491
+ # ActUp (connect).
492
+ corrected_predictions = self.altup.correct(predictions,
493
+ attn_ffw_laurel_gated)
494
+ first_prediction = corrected_predictions[self.altup_active_idx]
495
+ first_prediction = self.altup.scale_corrected_output(first_prediction)
496
+
497
+ # per_layer_input_gate adapted from jax.numpy.einsum("btd,dp->btp", ...)
498
+ first_prediction = self.per_layer_input_gate(first_prediction)
499
+ first_prediction = self.act_fn(first_prediction)
500
+ first_prediction = torch.mul(first_prediction, per_layer_input)
501
+
502
+ # per_layer_projection adapted from jax.numpy.einsum("btp,pd->btd", ...)
503
+ first_prediction = self.per_layer_projection(first_prediction)
504
+ first_prediction = self.post_per_layer_input_norm(first_prediction)
505
+ corrected_predictions[1:] += first_prediction
506
+
507
+ return corrected_predictions
508
+
509
+
510
+@support_torch_compile
511
+class Gemma3nTextModel(nn.Module):
512
+
513
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
514
+ super().__init__()
515
+ config = vllm_config.model_config.hf_config.text_config
516
+ cache_config = vllm_config.cache_config
517
+ quant_config = vllm_config.quant_config
518
+ self.config = config
519
+ self.quant_config = quant_config
520
+
521
+ self.embed_tokens = VocabParallelEmbedding(
522
+ config.vocab_size,
523
+ config.hidden_size,
524
+ prefix=f"{prefix}.embed_tokens",
525
+ )
526
+ self.embed_scale = torch.tensor(
527
+ config.hidden_size**0.5,
528
+ dtype=self.embed_tokens.weight.dtype,
529
+ )
530
+ self.embed_tokens_per_layer = VocabParallelEmbedding(
531
+ config.vocab_size_per_layer_input,
532
+ config.num_hidden_layers * config.hidden_size_per_layer_input,
533
+ prefix=f"{prefix}.per_layer_embed_tokens",
534
+ )
535
+ self.embed_scale_per_layer = torch.tensor(
536
+ config.hidden_size_per_layer_input**0.5,
537
+ dtype=self.embed_tokens.weight.dtype,
538
+ )
539
+ self.per_layer_model_projection = ColumnParallelLinear(
540
+ config.hidden_size,
541
+ config.num_hidden_layers * config.hidden_size_per_layer_input,
542
+ bias=False,
543
+ gather_output=True,
544
+ return_bias=False,
545
+ prefix=f"{prefix}.per_layer_model_projection",
546
+ )
547
+ self.per_layer_projection_norm = RMSNorm(
548
+ hidden_size=config.hidden_size_per_layer_input,
549
+ eps=config.rms_norm_eps,
550
+ )
551
+ self.per_layer_input_scale = torch.rsqrt(torch.tensor(2.0)).to(
552
+ self.embed_tokens.weight.dtype)
553
+ self.per_layer_projection_scale = torch.tensor(
554
+ config.hidden_size**0.5,
555
+ dtype=self.embed_tokens.weight.dtype,
556
+ )
557
+ self.altup_projections = nn.ModuleList([
558
+ ColumnParallelLinear(
559
+ config.hidden_size,
560
+ config.hidden_size,
561
+ bias=False,
562
+ gather_output=True,
563
+ return_bias=False,
564
+ prefix=f"{prefix}.{idx-1}.altup_projections",
565
+ ) for idx in range(1, self.config.altup_num_inputs)
566
+ ])
567
+ self.altup_unembed_projections = nn.ModuleList([
568
+ ColumnParallelLinear(
569
+ config.hidden_size,
570
+ config.hidden_size,
571
+ bias=False,
572
+ gather_output=True,
573
+ return_bias=False,
574
+ prefix=f"{prefix}.{idx-1}.altup_unembed_projections",
575
+ ) for idx in range(1, self.config.altup_num_inputs)
576
+ ])
577
+
578
+ # Transformer blocks.
579
+ self.start_layer, self.end_layer, self.layers = make_layers(
580
+ config.num_hidden_layers,
581
+ lambda prefix: Gemma3nDecoderLayer(
582
+ config, cache_config, quant_config, prefix=prefix),
583
+ prefix=f"{prefix}.layers")
584
+ self.norm = RMSNorm(
585
+ config.hidden_size,
586
+ eps=config.rms_norm_eps,
587
+ )
588
+ self.eps = torch.tensor(torch.finfo().min)
589
+
590
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
591
+ return self.embed_tokens(input_ids) * self.embed_scale
592
+
593
+ def get_per_layer_input_embeddings(
594
+ self, input_ids: torch.Tensor) -> torch.Tensor:
595
+ # Deal with the fact that vocab_size_per_layer_input < vocab_size
596
+ # which causes us to have some out of vocab tokens by setting
597
+ # those token ids to 0. This matches the HF implementation.
598
+ per_layer_inputs_mask = torch.logical_and(
599
+ input_ids >= 0, input_ids < self.config.vocab_size_per_layer_input)
600
+ per_layer_inputs_tokens = torch.where(per_layer_inputs_mask, input_ids,
601
+ torch.zeros_like(input_ids))
602
+ return self.embed_tokens_per_layer(
603
+ per_layer_inputs_tokens) * self.embed_scale_per_layer
604
+
605
+ def forward(
606
+ self,
607
+ input_ids: Optional[torch.Tensor],
608
+ positions: torch.Tensor,
609
+ inputs_embeds: Optional[torch.Tensor] = None,
610
+ **kwargs,
611
+ ) -> Union[torch.Tensor, IntermediateTensors]:
612
+ if inputs_embeds is not None:
613
+ hidden_states_0 = inputs_embeds
614
+ else:
615
+ hidden_states_0 = self.get_input_embeddings(input_ids)
616
+
617
+ # Per layer inputs.
618
+ if input_ids is None:
619
+ raise ValueError("Passing None for input ids is not supported.")
620
+ per_layer_inputs = self.get_per_layer_input_embeddings(input_ids)
621
+ per_layer_inputs = per_layer_inputs.reshape(
622
+ -1, self.config.num_hidden_layers,
623
+ self.config.hidden_size_per_layer_input)
624
+ per_layer_projection = self.per_layer_model_projection(hidden_states_0)
625
+ per_layer_projection = per_layer_projection.reshape(
626
+ *hidden_states_0.shape[:-1],
627
+ self.config.num_hidden_layers,
628
+ self.config.hidden_size_per_layer_input,
629
+ )
630
+ per_layer_projection = self.per_layer_projection_norm(
631
+ per_layer_projection)
632
+ per_layer_inputs = per_layer_projection + per_layer_inputs
633
+ per_layer_inputs *= self.per_layer_input_scale
634
+
635
+ # Altup embed.
636
+ hidden_states = [hidden_states_0] * self.config.altup_num_inputs
637
+ target_magnitude = torch.mean(hidden_states_0**2, dim=-1,
638
+ keepdim=True)**0.5
639
+ for i in range(1, self.config.altup_num_inputs):
640
+ hidden_states[i] = self.altup_projections[i - 1](hidden_states[i])
641
+ new_magnitude = torch.mean(hidden_states[i]**2,
642
+ dim=-1,
643
+ keepdim=True)**0.5
644
+ hidden_states[i] *= target_magnitude / torch.maximum(
645
+ new_magnitude, self.eps)
646
+ hidden_states = torch.stack(hidden_states, dim=0)
647
+
648
+ # Transformer blocks.
649
+ for layer_idx, layer in enumerate(self.layers):
650
+ # [altup_num_inputs, num_tokens, hidden_size]
651
+ hidden_states = layer(
652
+ positions=positions,
653
+ hidden_states=hidden_states,
654
+ per_layer_input=per_layer_inputs[:, layer_idx, :],
655
+ **kwargs,
656
+ )
657
+
658
+ # Altup unembed.
659
+ target_magnitude = torch.mean(hidden_states[0]**2,
660
+ dim=-1,
661
+ keepdim=True)**0.5
662
+ for i in range(1, self.config.altup_num_inputs):
663
+ hidden_states[i] = self.altup_unembed_projections[i - 1](
664
+ hidden_states[i])
665
+ new_magnitude = torch.mean(hidden_states[i]**2,
666
+ dim=-1,
667
+ keepdim=True)**0.5
668
+ hidden_states[i] *= target_magnitude / torch.maximum(
669
+ new_magnitude, self.eps)
670
+ # [altup_num_inputs,num_tokens,hidden_size] -> [num_tokens,hidden_size]
671
+ hidden_states = torch.mean(hidden_states, dim=0)
672
+
673
+ return self.norm(hidden_states)
674
+
675
+ def load_weights(self, weights: Iterable[tuple[str,
676
+ torch.Tensor]]) -> set[str]:
677
+ stacked_params_mapping = [
678
+ # (param_name, shard_name, shard_id)
679
+ ("qkv_proj", "q_proj", "q"),
680
+ ("qkv_proj", "k_proj", "k"),
681
+ ("qkv_proj", "v_proj", "v"),
682
+ ("gate_up_proj", "gate_proj", 0),
683
+ ("gate_up_proj", "up_proj", 1),
684
+ ]
685
+ params_dict = dict(self.named_parameters())
686
+ loaded_params: set[str] = set()
687
+ for name, loaded_weight in weights:
688
+ if (self.quant_config is not None and
689
+ (scale_name := self.quant_config.get_cache_scale(name))):
690
+ # Loading kv cache scales for compressed-tensors quantization
691
+ param = params_dict[scale_name]
692
+ weight_loader = getattr(param, "weight_loader",
693
+ default_weight_loader)
694
+ loaded_weight = loaded_weight[0]
695
+ weight_loader(param, loaded_weight)
696
+ loaded_params.add(scale_name)
697
+ continue
698
+ for (param_name, shard_name, shard_id) in stacked_params_mapping:
699
+ if shard_name not in name:
700
+ continue
701
+ # Avoid spurious match with ".up_proj".
702
+ if "altup_projections" in name:
703
+ continue
704
+ name = name.replace(shard_name, param_name)
705
+ # Skip loading extra bias for GPTQ models.
706
+ if name.endswith(".bias") and name not in params_dict:
707
+ continue
708
+ if is_pp_missing_parameter(name, self):
709
+ continue
710
+ param = params_dict[name]
711
+ weight_loader = param.weight_loader
712
+ weight_loader(param, loaded_weight, shard_id)
713
+ break
714
+ else:
715
+ # Skip loading extra bias for GPTQ models.
716
+ if name.endswith(".bias") and name not in params_dict:
717
+ continue
718
+ # Remapping the name of FP8 kv-scale.
719
+ name = maybe_remap_kv_scale_name(name, params_dict)
720
+ if name is None:
721
+ continue
722
+ if is_pp_missing_parameter(name, self):
723
+ continue
724
+ param = params_dict[name]
725
+ weight_loader = getattr(param, "weight_loader",
726
+ default_weight_loader)
727
+ weight_loader(param, loaded_weight)
728
+ loaded_params.add(name)
729
+
730
+ return loaded_params
731
+
732
+
733
+class Gemma3nModel(nn.Module):
734
+
735
+ def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
736
+ super().__init__()
737
+ self.language_model = Gemma3nTextModel(vllm_config=vllm_config,
738
+ prefix=maybe_prefix(
739
+ prefix, "language_model"))
740
+
741
+ def forward(
742
+ self,
743
+ input_ids: Optional[torch.Tensor],
744
+ positions: torch.Tensor,
745
+ intermediate_tensors: Optional[IntermediateTensors] = None,
746
+ inputs_embeds: Optional[torch.Tensor] = None,
747
+ **kwargs,
748
+ ) -> torch.Tensor:
749
+ return self.language_model(input_ids=input_ids,
750
+ positions=positions,
751
+ inputs_embeds=inputs_embeds,
752
+ **kwargs)
753
+
754
+
755
+class Gemma3nForConditionalGeneration(nn.Module):
756
+ packed_modules_mapping = {
757
+ "qkv_proj": [
758
+ "q_proj",
759
+ "k_proj",
760
+ "v_proj",
761
+ ],
762
+ "gate_up_proj": [
763
+ "gate_proj",
764
+ "up_proj",
765
+ ],
766
+ }
767
+
768
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
769
+ config = vllm_config.model_config.hf_config
770
+ lora_config = vllm_config.lora_config
771
+ del lora_config # Unused.
772
+ super().__init__()
773
+ self.config = config
774
+ self.model = Gemma3nModel(vllm_config=vllm_config,
775
+ prefix=maybe_prefix(prefix, "model"))
776
+ self.logits_processor = LogitsProcessor(
777
+ config.text_config.vocab_size,
778
+ soft_cap=config.text_config.final_logit_softcapping)
779
+
780
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
781
+ return self.model.language_model.get_input_embeddings(input_ids)
782
+
783
+ def forward(
784
+ self,
785
+ input_ids: torch.Tensor,
786
+ positions: torch.Tensor,
787
+ intermediate_tensors: Optional[IntermediateTensors] = None,
788
+ inputs_embeds: Optional[torch.Tensor] = None,
789
+ **kwargs,
790
+ ) -> Union[torch.Tensor, IntermediateTensors]:
791
+ hidden_states = self.model(input_ids, positions, intermediate_tensors,
792
+ inputs_embeds, **kwargs)
793
+ return hidden_states
794
+
795
+ def compute_logits(
796
+ self,
797
+ hidden_states: torch.Tensor,
798
+ sampling_metadata: Optional[SamplingMetadata],
799
+ ) -> Optional[torch.Tensor]:
800
+ logits = self.logits_processor(self.model.language_model.embed_tokens,
801
+ hidden_states, sampling_metadata)
802
+ return logits
803
+
804
+ def load_weights(self, weights: Iterable[tuple[str,
805
+ torch.Tensor]]) -> set[str]:
806
+ loader = AutoWeightsLoader(self,
807
+ skip_substrs=([
808
+ "embed_audio.", "embed_vision.",
809
+ "audio_tower.", "vision_tower."
810
+ ]))
811
+ return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+2
@@ -58,6 +58,8 @@ _TEXT_GENERATION_MODELS = {
58
"GemmaForCausalLM": ("gemma", "GemmaForCausalLM"),
59
"Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),
60
"Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"),
61
+ #TODO(ywang96): Support multimodal gemma3n
62
+ "Gemma3nForConditionalGeneration": ("gemma3n", "Gemma3nForConditionalGeneration"), # noqa: E501
63
"GlmForCausalLM": ("glm", "GlmForCausalLM"),
64
"Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
65
"GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),