@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c0c25e25fa
Fork 0

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, 19:45 UTC

[Model] Add support for Gemma 3 (#14660)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Signed-off-by: Roger Wang <ywang@roblox.com> Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: Roger Wang <ywang@roblox.com> Co-authored-by: DarkLight1337 <tlleungac@connect.ust.hk>

Woosuk Kwon committed Mar 12, 2025 at 08:36 UTC c0c25e25fa93ee7c3f279abbba5597c0fafa74ee
10 files changed +1071 -9
docs/source/models/supported_models.md
+39 -2
@@ -263,10 +263,15 @@ See [this page](#generative-models) for more information on how to use generativ
263 * ✅︎
264 * ✅︎
265 - * `Gemma2ForCausalLM`
266 - * Gemma2
266 + * Gemma 2
267 * `google/gemma-2-9b`, `google/gemma-2-27b`, etc.
268 * ✅︎
269 * ✅︎
270 +- * `Gemma3ForCausalLM`
271 + * Gemma 3
272 + * `google/gemma-3-1b-it`, etc.
273 + * ✅︎
274 + * ✅︎
275 - * `GlmForCausalLM`
276 * GLM-4
277 * `THUDM/glm-4-9b-chat-hf`, etc.
@@ -504,7 +509,7 @@ you should explicitly specify the task type to ensure that the model is used in
509 *
510 *
511 - * `Gemma2Model`
507 - * Gemma2-based
512 + * Gemma 2-based
513 * `BAAI/bge-multilingual-gemma2`, etc.
514 *
515 * ✅︎
@@ -752,6 +757,13 @@ See [this page](#generative-models) for more information on how to use generativ
757 *
758 * ✅︎
759 * ✅︎
760 +- * `Gemma3ForConditionalGeneration`
761 + * Gemma 3
762 + * T + I<sup>+</sup>
763 + * `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc.
764 + * ✅︎
765 + * ✅︎
766 + * ✅︎\*
767 - * `GLM4VForCausalLM`<sup>^</sup>
768 * GLM-4V
769 * T + I
@@ -937,6 +949,31 @@ For more details, please see: <gh-pr:4087#issuecomment-2250397630>
949 To use Qwen2.5-VL series models, you have to install Hugging Face Transformers library from source via `pip install git+https://github.com/huggingface/transformers`.
950 :::
951
952 +:::{note}
953 +To use Gemma3 series models, you have to install Hugging Face Transformers library from source via
954 +`pip install git+https://github.com/huggingface/transformers`.
955 +The earliest commit that supports this is [`50d3530aa04e7a7d003e6b255a98f79fd0447357`](https://github.com/huggingface/transformers/commit/50d3530aa04e7a7d003e6b255a98f79fd0447357).
956 +
957 +Both V0 and V1 support `Gemma3ForConditionalGeneration` for text-only inputs.
958 +However, there are differences in how they handle text + image inputs:
959 +
960 +V0 correctly implements the model's attention pattern:
961 +- Uses bidirectional attention between the image tokens corresponding to the same image
962 +- Uses causal attention for other tokens
963 +- Implemented via (naive) PyTorch SDPA with masking tensors
964 +- Note: May use significant memory for long prompts with image
965 +
966 +V1 currently uses a simplified attention pattern:
967 +- Uses causal attention for all tokens, including image tokens
968 +- Generates reasonable outputs but does not match the original model's attention for text + image inputs
969 +- Will be updated in the future to support the correct behavior
970 +
971 +This limitation exists because the model's mixed attention pattern (bidirectional for images, causal otherwise) is not yet supported by vLLM's attention backends.
972 +
973 +Additionally, vLLM's current Gemma 3 implementation does not support the pan-and-scan image pre-processing algorithm, which helps handle images with skewed aspect ratios by intelligently cropping them into multiple views.
974 +Without this feature, model performance may degrade when processing images that deviate significantly from square dimensions.
975 +:::
976 +
977 ### Pooling Models
978
979 See [this page](pooling-models) for more information on how to use pooling models.
examples/offline_inference/vision_language.py
+19 -1
@@ -118,6 +118,23 @@ def run_fuyu(questions: list[str], modality: str):
118 return llm, prompts, stop_token_ids
119
120
121 +# Gemma 3
122 +def run_gemma3(questions: list[str], modality: str):
123 + assert modality == "image"
124 + model_name = "google/gemma-3-4b-it"
125 +
126 + llm = LLM(model=model_name,
127 + max_model_len=2048,
128 + max_num_seqs=2,
129 + disable_mm_preprocessor_cache=args.disable_mm_preprocessor_cache)
130 +
131 + prompts = [("<bos><start_of_turn>user\n"
132 + f"<start_of_image>{question}<end_of_turn>\n"
133 + "<start_of_turn>model\n") for question in questions]
134 + stop_token_ids = None
135 + return llm, prompts, stop_token_ids
136 +
137 +
138 # GLM-4v
139 def run_glm4v(questions: list[str], modality: str):
140 assert modality == "image"
@@ -405,7 +422,7 @@ def run_mllama(questions: list[str], modality: str):
422 "type": "image"
423 }, {
424 "type": "text",
408 - "text": f"{question}"
425 + "text": question
426 }]
427 }] for question in questions]
428 prompts = tokenizer.apply_chat_template(messages,
@@ -664,6 +681,7 @@ model_example_map = {
681 "deepseek_vl_v2": run_deepseek_vl2,
682 "florence2": run_florence2,
683 "fuyu": run_fuyu,
684 + "gemma3": run_gemma3,
685 "glm4v": run_glm4v,
686 "h2ovl_chat": run_h2ovl,
687 "idefics3": run_idefics3,
examples/offline_inference/vision_language_multi_image.py
+37
@@ -80,6 +80,42 @@ def load_deepseek_vl2(question: str, image_urls: list[str]):
80 )
81
82
83 +def load_gemma3(question, image_urls: list[str]) -> ModelRequestData:
84 + model_name = "google/gemma-3-4b-it"
85 +
86 + llm = LLM(model=model_name,
87 + max_model_len=8192,
88 + max_num_seqs=2,
89 + limit_mm_per_prompt={"image": len(image_urls)})
90 +
91 + placeholders = [{"type": "image", "image": url} for url in image_urls]
92 + messages = [{
93 + "role":
94 + "user",
95 + "content": [
96 + *placeholders,
97 + {
98 + "type": "text",
99 + "text": question
100 + },
101 + ],
102 + }]
103 +
104 + processor = AutoProcessor.from_pretrained(model_name)
105 +
106 + prompt = processor.apply_chat_template(messages,
107 + tokenize=False,
108 + add_generation_prompt=True)
109 +
110 + return ModelRequestData(
111 + llm=llm,
112 + prompt=prompt,
113 + stop_token_ids=None,
114 + image_data=[fetch_image(url) for url in image_urls],
115 + chat_template=None,
116 + )
117 +
118 +
119 def load_h2ovl(question: str, image_urls: list[str]) -> ModelRequestData:
120 model_name = "h2oai/h2ovl-mississippi-800m"
121
@@ -496,6 +532,7 @@ def load_qwen2_5_vl(question, image_urls: list[str]) -> ModelRequestData:
532 model_example_map = {
533 "aria": load_aria,
534 "deepseek_vl_v2": load_deepseek_vl2,
535 + "gemma3": load_gemma3,
536 "h2ovl_chat": load_h2ovl,
537 "idefics3": load_idefics3,
538 "internvl_chat": load_internvl,
tests/models/multimodal/processing/test_common.py
+1
@@ -162,6 +162,7 @@ def _test_processing_correctness(
162 "deepseek-ai/deepseek-vl2-tiny",
163 "microsoft/Florence-2-base",
164 "adept/fuyu-8b",
165 + "google/gemma-3-4b-it",
166 "THUDM/glm-4v-9b",
167 "h2oai/h2ovl-mississippi-800m",
168 "OpenGVLab/InternVL2-1B",
tests/models/registry.py
+4
@@ -124,6 +124,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
124 "FalconForCausalLM": _HfExamplesInfo("tiiuae/falcon-7b"),
125 "GemmaForCausalLM": _HfExamplesInfo("google/gemma-2b"),
126 "Gemma2ForCausalLM": _HfExamplesInfo("google/gemma-2-9b"),
127 + "Gemma3ForCausalLM": _HfExamplesInfo("google/gemma-3-1b-it",
128 + min_transformers_version="4.50"),
129 "GlmForCausalLM": _HfExamplesInfo("THUDM/glm-4-9b-chat-hf"),
130 "GPT2LMHeadModel": _HfExamplesInfo("gpt2"),
131 "GPTBigCodeForCausalLM": _HfExamplesInfo("bigcode/starcoder"),
@@ -241,6 +243,8 @@ _MULTIMODAL_EXAMPLE_MODELS = {
243 "DeepseekVLV2ForCausalLM": _HfExamplesInfo("deepseek-ai/deepseek-vl2-tiny", # noqa: E501
244 hf_overrides={"architectures": ["DeepseekVLV2ForCausalLM"]}), # noqa: E501
245 "FuyuForCausalLM": _HfExamplesInfo("adept/fuyu-8b"),
246 + "Gemma3ForConditionalGeneration": _HfExamplesInfo("google/gemma-3-4b-it",
247 + min_transformers_version="4.50"),
248 "GLM4VForCausalLM": _HfExamplesInfo("THUDM/glm-4v-9b",
249 trust_remote_code=True,
250 hf_overrides={"architectures": ["GLM4VForCausalLM"]}), # noqa: E501
vllm/config.py
+9 -6
@@ -350,10 +350,11 @@ class ModelConfig:
350 if self.enforce_eager is None:
351 self.enforce_eager = False
352
353 + interleaved_attn_models = ["gemma2", "gemma3_text", "cohere2"]
354 sliding_window = getattr(self.hf_text_config, "sliding_window", None)
355 has_interleaved_attention = (sliding_window is not None) and (
356 isinstance(sliding_window, list) or
356 - (self.hf_text_config.model_type in ["gemma2", "cohere2"]))
357 + (self.hf_text_config.model_type in interleaved_attn_models))
358
359 if (not self.disable_sliding_window and has_interleaved_attention):
360 if (backend :=
@@ -2501,11 +2502,11 @@ def _get_and_verify_dtype(
2502 dtype = dtype.lower()
2503 if dtype == "auto":
2504 if config_dtype == torch.float32:
2504 - if config.model_type == "gemma2":
2505 + if config.model_type in ("gemma2", "gemma3", "gemma3_text"):
2506 logger.info(
2506 - "For Gemma 2, we downcast float32 to bfloat16 instead "
2507 - "of float16 by default. Please specify `dtype` if you "
2508 - "want to use float16.")
2507 + "For Gemma 2 and 3, we downcast float32 to bfloat16 "
2508 + "instead of float16 by default. Please specify `dtype` "
2509 + "if you want to use float16.")
2510 torch_dtype = torch.bfloat16
2511 else:
2512 # Following the common practice, we use float16 for float32
@@ -2637,7 +2638,9 @@ def _get_and_verify_max_len(
2638 derived_max_model_len = default_max_len
2639
2640 rope_scaling = getattr(hf_config, "rope_scaling", None)
2640 - if rope_scaling is not None:
2641 + # NOTE(woosuk): Gemma3's max_model_len (128K) is already scaled by RoPE
2642 + # scaling, so we skip applying the scaling factor again.
2643 + if rope_scaling is not None and "gemma3" not in hf_config.model_type:
2644 # No need to consider "type" key because of patch_rope_scaling when
2645 # loading HF config
2646 rope_type = rope_scaling["rope_type"]
vllm/entrypoints/chat_utils.py
+2
@@ -433,6 +433,8 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):
433 return "<image>"
434 if model_type == "aria":
435 return "<|fim_prefix|><|img|><|fim_suffix|>"
436 + if model_type == "gemma3":
437 + return "<start_of_image>"
438
439 raise TypeError(f"Unknown {modality} model type: {model_type}")
440 elif modality == "audio":
vllm/model_executor/models/gemma3.py new
+533
@@ -0,0 +1,533 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# Copyright 2025 The vLLM team.
3 +# Copyright 2025 Google Inc. HuggingFace Inc. team. All rights reserved.
4 +#
5 +#
6 +# Licensed under the Apache License, Version 2.0 (the "License");
7 +# you may not use this file except in compliance with the License.
8 +# You may obtain a copy of the License at
9 +#
10 +# http://www.apache.org/licenses/LICENSE-2.0
11 +#
12 +# Unless required by applicable law or agreed to in writing, software
13 +# distributed under the License is distributed on an "AS IS" BASIS,
14 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15 +# See the License for the specific language governing permissions and
16 +# limitations under the License.
17 +from typing import Iterable, Optional, Set, Tuple, Union
18 +
19 +import torch
20 +import torch.nn.functional as F
21 +from torch import nn
22 +from transformers import Gemma3TextConfig
23 +
24 +from vllm.attention import Attention
25 +from vllm.compilation.decorators import support_torch_compile
26 +from vllm.config import CacheConfig, VllmConfig
27 +from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
28 +from vllm.logger import init_logger
29 +from vllm.model_executor.layers.activation import GeluAndMul
30 +from vllm.model_executor.layers.layernorm import GemmaRMSNorm
31 +from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
32 + QKVParallelLinear,
33 + RowParallelLinear)
34 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
35 +from vllm.model_executor.layers.quantization import QuantizationConfig
36 +from vllm.model_executor.layers.rotary_embedding import get_rope
37 +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
38 +from vllm.model_executor.layers.vocab_parallel_embedding import (
39 + VocabParallelEmbedding)
40 +from vllm.model_executor.model_loader.weight_utils import (
41 + default_weight_loader, maybe_remap_kv_scale_name)
42 +from vllm.model_executor.sampling_metadata import SamplingMetadata
43 +from vllm.sequence import IntermediateTensors
44 +
45 +from .interfaces import SupportsLoRA, SupportsPP
46 +from .utils import (AutoWeightsLoader, extract_layer_index,
47 + is_pp_missing_parameter,
48 + make_empty_intermediate_tensors_factory, make_layers,
49 + maybe_prefix)
50 +
51 +logger = init_logger(__name__)
52 +
53 +
54 +class Gemma3MLP(nn.Module):
55 +
56 + def __init__(
57 + self,
58 + hidden_size: int,
59 + intermediate_size: int,
60 + hidden_activation: str,
61 + quant_config: Optional[QuantizationConfig] = None,
62 + ) -> None:
63 + super().__init__()
64 + self.gate_up_proj = MergedColumnParallelLinear(
65 + hidden_size, [intermediate_size] * 2,
66 + bias=False,
67 + quant_config=quant_config)
68 + self.down_proj = RowParallelLinear(intermediate_size,
69 + hidden_size,
70 + bias=False,
71 + quant_config=quant_config)
72 + if hidden_activation != "gelu_pytorch_tanh":
73 + raise ValueError(
74 + "Gemma3 uses `gelu_pytorch_tanh` as the hidden activation "
75 + "function. Please set `hidden_act` and `hidden_activation` to "
76 + "`gelu_pytorch_tanh`.")
77 + self.act_fn = GeluAndMul(approximate="tanh")
78 +
79 + def forward(self, x: torch.Tensor) -> torch.Tensor:
80 + gate_up, _ = self.gate_up_proj(x)
81 + x = self.act_fn(gate_up)
82 + x, _ = self.down_proj(x)
83 + return x
84 +
85 +
86 +class Gemma3Attention(nn.Module):
87 +
88 + def __init__(self,
89 + config: Gemma3TextConfig,
90 + hidden_size: int,
91 + num_heads: int,
92 + num_kv_heads: int,
93 + head_dim: int,
94 + max_position_embeddings: int,
95 + cache_config: Optional[CacheConfig] = None,
96 + quant_config: Optional[QuantizationConfig] = None,
97 + attn_logits_soft_cap: Optional[float] = None,
98 + prefix: str = "") -> None:
99 + super().__init__()
100 + self.config = config
101 + self.hidden_size = hidden_size
102 + tp_size = get_tensor_model_parallel_world_size()
103 + self.total_num_heads = num_heads
104 + assert self.total_num_heads % tp_size == 0
105 + self.num_heads = self.total_num_heads // tp_size
106 + self.total_num_kv_heads = num_kv_heads
107 + if self.total_num_kv_heads >= tp_size:
108 + # Number of KV heads is greater than TP size, so we partition
109 + # the KV heads across multiple tensor parallel GPUs.
110 + assert self.total_num_kv_heads % tp_size == 0
111 + else:
112 + # Number of KV heads is less than TP size, so we replicate
113 + # the KV heads across multiple tensor parallel GPUs.
114 + assert tp_size % self.total_num_kv_heads == 0
115 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
116 + self.head_dim = head_dim
117 + self.q_size = self.num_heads * self.head_dim
118 + self.kv_size = self.num_kv_heads * self.head_dim
119 + self.scaling = config.query_pre_attn_scalar**-0.5
120 +
121 + self.qkv_proj = QKVParallelLinear(
122 + hidden_size,
123 + self.head_dim,
124 + self.total_num_heads,
125 + self.total_num_kv_heads,
126 + bias=config.attention_bias,
127 + quant_config=quant_config,
128 + )
129 + self.o_proj = RowParallelLinear(
130 + self.total_num_heads * self.head_dim,
131 + hidden_size,
132 + bias=config.attention_bias,
133 + quant_config=quant_config,
134 + )
135 +
136 + self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
137 + self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
138 +
139 + # TODO(woosuk): Add reference to the original HF implementation.
140 + layer_idx = extract_layer_index(prefix)
141 + self.is_sliding = bool((layer_idx + 1) % config.sliding_window_pattern)
142 + # Initialize the rotary embedding.
143 + if self.is_sliding:
144 + # Local attention. Override the values in config.json.
145 + self.rope_theta = config.rope_local_base_freq
146 + self.rope_scaling = {"rope_type": "default"}
147 + self.sliding_window = config.interleaved_sliding_window
148 + else:
149 + # Global attention. Use the values in config.json.
150 + self.rope_theta = config.rope_theta
151 + self.rope_scaling = config.rope_scaling
152 + self.sliding_window = None
153 + self.rotary_emb = get_rope(
154 + self.head_dim,
155 + rotary_dim=self.head_dim,
156 + max_position=max_position_embeddings,
157 + base=self.rope_theta,
158 + is_neox_style=True,
159 + rope_scaling=self.rope_scaling,
160 + )
161 +
162 + # Initialize the attention.
163 + self.attn = Attention(self.num_heads,
164 + self.head_dim,
165 + self.scaling,
166 + num_kv_heads=self.num_kv_heads,
167 + cache_config=cache_config,
168 + quant_config=quant_config,
169 + logits_soft_cap=attn_logits_soft_cap,
170 + per_layer_sliding_window=self.sliding_window,
171 + prefix=f"{prefix}.attn")
172 +
173 + def forward(
174 + self,
175 + positions: torch.Tensor,
176 + hidden_states: torch.Tensor,
177 + **kwargs,
178 + ) -> torch.Tensor:
179 + qkv, _ = self.qkv_proj(hidden_states)
180 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
181 +
182 + q = q.unflatten(-1, (self.num_heads, self.head_dim))
183 + q = self.q_norm(q)
184 + q = q.flatten(-2, -1)
185 + k = k.unflatten(-1, (self.num_kv_heads, self.head_dim))
186 + k = self.k_norm(k)
187 + k = k.flatten(-2, -1)
188 +
189 + q, k = self.rotary_emb(positions, q, k)
190 + attn_output = self.attn(q, k, v)
191 +
192 + if not kwargs.get("has_images", False):
193 + # Fast path for text-only inputs. The performance for the text-only
194 + # inputs are not affected by the naive attention below.
195 + output, _ = self.o_proj(attn_output)
196 + return output
197 +
198 + # NOTE(woosuk): Gemma3 uses bidirectional attention between image tokens
199 + # that correspond to the same image while using causal attention
200 + # otherwise. Current attention backends cannot handle this pattern, so
201 + # we temporarily use a naive attention implementation with mask tensors.
202 +
203 + # We intentionally keep the attention backend as-is and only override
204 + # `attn_output` with the naive implementation's output. This minimizes
205 + # changes to existing model runners and attention backends. The call to
206 + # `self.attn(q, k, v)` is only used to populate the KV cache - its
207 + # output is discarded and overwritten below. While this duplicates
208 + # computation, it maintains compatibility.
209 + # TODO(woosuk): Optimize by implementing custom attention kernels.
210 + attn_output = self.naive_attn_with_masks(q,
211 + k,
212 + v,
213 + out=attn_output,
214 + **kwargs)
215 + output, _ = self.o_proj(attn_output)
216 + return output
217 +
218 + def naive_attn_with_masks(
219 + self,
220 + q: torch.Tensor,
221 + k: torch.Tensor,
222 + v: torch.Tensor,
223 + out: torch.Tensor,
224 + **kwargs,
225 + ) -> torch.Tensor:
226 + # NOTE(woosuk): As described in the comment above, this code is not
227 + # meant to be performant. It is only meant to be correct.
228 + q = q.view(-1, self.num_heads, self.head_dim)
229 + # Expand the key and value to handle GQA.
230 + num_queries_per_kv = self.num_heads // self.num_kv_heads
231 + k = k.view(-1, self.num_kv_heads, self.head_dim)
232 + k = k.repeat_interleave(num_queries_per_kv, dim=-2)
233 + v = v.view(-1, self.num_kv_heads, self.head_dim)
234 + v = v.repeat_interleave(num_queries_per_kv, dim=-2)
235 +
236 + if self.is_sliding:
237 + attn_masks = kwargs["local_attn_masks"]
238 + else:
239 + attn_masks = kwargs["global_attn_masks"]
240 +
241 + seq_lens = kwargs["seq_lens"]
242 + start_idx = 0
243 + for seq_len, attn_mask in zip(seq_lens, attn_masks):
244 + end_idx = start_idx + seq_len
245 + query = q[start_idx:end_idx].unsqueeze(0)
246 + key = k[start_idx:end_idx].unsqueeze(0)
247 + value = v[start_idx:end_idx].unsqueeze(0)
248 +
249 + # Transpose.
250 + query = query.transpose(1, 2)
251 + key = key.transpose(1, 2)
252 + value = value.transpose(1, 2)
253 +
254 + output = F.scaled_dot_product_attention(
255 + query,
256 + key,
257 + value,
258 + attn_mask,
259 + self.scaling,
260 + )
261 + output = output.transpose(1, 2).flatten(-2, -1)
262 + out[start_idx:end_idx] = output
263 + start_idx = end_idx
264 + return out
265 +
266 +
267 +class Gemma3DecoderLayer(nn.Module):
268 +
269 + def __init__(
270 + self,
271 + config: Gemma3TextConfig,
272 + cache_config: Optional[CacheConfig] = None,
273 + quant_config: Optional[QuantizationConfig] = None,
274 + prefix: str = "",
275 + ) -> None:
276 + super().__init__()
277 + self.hidden_size = config.hidden_size
278 + self.self_attn = Gemma3Attention(
279 + config=config,
280 + hidden_size=self.hidden_size,
281 + num_heads=config.num_attention_heads,
282 + num_kv_heads=config.num_key_value_heads,
283 + head_dim=config.head_dim,
284 + max_position_embeddings=config.max_position_embeddings,
285 + cache_config=cache_config,
286 + quant_config=quant_config,
287 + attn_logits_soft_cap=None,
288 + prefix=f"{prefix}.self_attn",
289 + )
290 + self.hidden_size = config.hidden_size
291 + self.mlp = Gemma3MLP(
292 + hidden_size=self.hidden_size,
293 + intermediate_size=config.intermediate_size,
294 + hidden_activation=config.hidden_activation,
295 + quant_config=quant_config,
296 + )
297 + self.input_layernorm = GemmaRMSNorm(config.hidden_size,
298 + eps=config.rms_norm_eps)
299 + self.post_attention_layernorm = GemmaRMSNorm(config.hidden_size,
300 + eps=config.rms_norm_eps)
301 + self.pre_feedforward_layernorm = GemmaRMSNorm(config.hidden_size,
302 + eps=config.rms_norm_eps)
303 + self.post_feedforward_layernorm = GemmaRMSNorm(config.hidden_size,
304 + eps=config.rms_norm_eps)
305 +
306 + def forward(
307 + self,
308 + positions: torch.Tensor,
309 + hidden_states: torch.Tensor,
310 + residual: Optional[torch.Tensor],
311 + **kwargs,
312 + ) -> Tuple[torch.Tensor, torch.Tensor]:
313 + if residual is None:
314 + residual = hidden_states
315 + hidden_states = self.input_layernorm(hidden_states)
316 + else:
317 + hidden_states, residual = self.input_layernorm(
318 + hidden_states, residual)
319 + hidden_states = self.self_attn(
320 + positions=positions,
321 + hidden_states=hidden_states,
322 + **kwargs,
323 + )
324 + hidden_states = self.post_attention_layernorm(hidden_states)
325 +
326 + hidden_states, residual = self.pre_feedforward_layernorm(
327 + hidden_states, residual)
328 + hidden_states = self.mlp(hidden_states)
329 + hidden_states = self.post_feedforward_layernorm(hidden_states)
330 + return hidden_states, residual
331 +
332 +
333 +@support_torch_compile
334 +class Gemma3Model(nn.Module):
335 +
336 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
337 + super().__init__()
338 + config = vllm_config.model_config.hf_config
339 + cache_config = vllm_config.cache_config
340 + quant_config = vllm_config.quant_config
341 + self.config = config
342 + self.quant_config = quant_config
343 +
344 + self.embed_tokens = VocabParallelEmbedding(
345 + config.vocab_size,
346 + config.hidden_size,
347 + )
348 + self.start_layer, self.end_layer, self.layers = make_layers(
349 + config.num_hidden_layers,
350 + lambda prefix: Gemma3DecoderLayer(
351 + config, cache_config, quant_config, prefix=prefix),
352 + prefix=f"{prefix}.layers")
353 + self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
354 +
355 + # Normalize the embedding by sqrt(hidden_size)
356 + # The normalizer's data type should be downcasted to the model's
357 + # data type such as bfloat16, not float32.
358 + # See https://github.com/huggingface/transformers/pull/29402
359 + normalizer = self.config.hidden_size**0.5
360 + self.register_buffer("normalizer", torch.tensor(normalizer))
361 + self.make_empty_intermediate_tensors = (
362 + make_empty_intermediate_tensors_factory(
363 + ["hidden_states", "residual"], config.hidden_size))
364 +
365 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
366 + # NOTE(woosuk): Only apply the normalizer to the output of
367 + # vocab embedding. Don't apply it to the vision embedding.
368 + return self.embed_tokens(input_ids) * self.normalizer
369 +
370 + def forward(
371 + self,
372 + input_ids: Optional[torch.Tensor],
373 + positions: torch.Tensor,
374 + intermediate_tensors: Optional[IntermediateTensors],
375 + inputs_embeds: Optional[torch.Tensor] = None,
376 + **kwargs,
377 + ) -> Union[torch.Tensor, IntermediateTensors]:
378 + if get_pp_group().is_first_rank:
379 + if inputs_embeds is not None:
380 + hidden_states = inputs_embeds
381 + else:
382 + hidden_states = self.get_input_embeddings(input_ids)
383 + residual = None
384 + else:
385 + assert intermediate_tensors is not None
386 + hidden_states = intermediate_tensors["hidden_states"]
387 + residual = intermediate_tensors["residual"]
388 + for layer in self.layers[self.start_layer:self.end_layer]:
389 + hidden_states, residual = layer(
390 + positions,
391 + hidden_states,
392 + residual,
393 + **kwargs,
394 + )
395 + if not get_pp_group().is_last_rank:
396 + return IntermediateTensors({
397 + "hidden_states": hidden_states,
398 + "residual": residual
399 + })
400 + hidden_states, _ = self.norm(hidden_states, residual)
401 + return hidden_states
402 +
403 + def load_weights(self, weights: Iterable[Tuple[str,
404 + torch.Tensor]]) -> Set[str]:
405 + stacked_params_mapping = [
406 + # (param_name, shard_name, shard_id)
407 + ("qkv_proj", "q_proj", "q"),
408 + ("qkv_proj", "k_proj", "k"),
409 + ("qkv_proj", "v_proj", "v"),
410 + ("gate_up_proj", "gate_proj", 0),
411 + ("gate_up_proj", "up_proj", 1),
412 + ]
413 + params_dict = dict(self.named_parameters())
414 + loaded_params: Set[str] = set()
415 + for name, loaded_weight in weights:
416 + if (self.quant_config is not None and
417 + (scale_name := self.quant_config.get_cache_scale(name))):
418 + # Loading kv cache scales for compressed-tensors quantization
419 + param = params_dict[scale_name]
420 + weight_loader = getattr(param, "weight_loader",
421 + default_weight_loader)
422 + loaded_weight = loaded_weight[0]
423 + weight_loader(param, loaded_weight)
424 + loaded_params.add(scale_name)
425 + continue
426 + for (param_name, shard_name, shard_id) in stacked_params_mapping:
427 + if shard_name not in name:
428 + continue
429 + name = name.replace(shard_name, param_name)
430 + # Skip loading extra bias for GPTQ models.
431 + if name.endswith(".bias") and name not in params_dict:
432 + continue
433 + if is_pp_missing_parameter(name, self):
434 + continue
435 + param = params_dict[name]
436 + weight_loader = param.weight_loader
437 + weight_loader(param, loaded_weight, shard_id)
438 + break
439 + else:
440 + # Skip loading extra bias for GPTQ models.
441 + if name.endswith(".bias") and name not in params_dict:
442 + continue
443 + # Remapping the name of FP8 kv-scale.
444 + name = maybe_remap_kv_scale_name(name, params_dict)
445 + if name is None:
446 + continue
447 + if is_pp_missing_parameter(name, self):
448 + continue
449 + param = params_dict[name]
450 + weight_loader = getattr(param, "weight_loader",
451 + default_weight_loader)
452 + weight_loader(param, loaded_weight)
453 + loaded_params.add(name)
454 +
455 + unloaded_params = params_dict.keys() - loaded_params
456 + if unloaded_params:
457 + logger.warning(
458 + "Some weights are not initialized from checkpoints: %s",
459 + unloaded_params)
460 + return loaded_params
461 +
462 +
463 +class Gemma3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
464 + packed_modules_mapping = {
465 + "qkv_proj": [
466 + "q_proj",
467 + "k_proj",
468 + "v_proj",
469 + ],
470 + "gate_up_proj": [
471 + "gate_proj",
472 + "up_proj",
473 + ],
474 + }
475 +
476 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
477 + config = vllm_config.model_config.hf_config
478 + quant_config = vllm_config.quant_config
479 + lora_config = vllm_config.lora_config
480 + del lora_config # Unused.
481 + super().__init__()
482 + self.config = config
483 + # currently all existing Gemma models have `tie_word_embeddings` enabled
484 + assert config.tie_word_embeddings
485 + self.quant_config = quant_config
486 + self.model = Gemma3Model(vllm_config=vllm_config,
487 + prefix=maybe_prefix(prefix, "model"))
488 + self.logits_processor = LogitsProcessor(
489 + config.vocab_size, soft_cap=config.final_logit_softcapping)
490 + self.sampler = get_sampler()
491 + self.make_empty_intermediate_tensors = (
492 + self.model.make_empty_intermediate_tensors)
493 +
494 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
495 + return self.model.get_input_embeddings(input_ids)
496 +
497 + def forward(
498 + self,
499 + input_ids: torch.Tensor,
500 + positions: torch.Tensor,
501 + intermediate_tensors: Optional[IntermediateTensors] = None,
502 + inputs_embeds: Optional[torch.Tensor] = None,
503 + **kwargs,
504 + ) -> Union[torch.Tensor, IntermediateTensors]:
505 + hidden_states = self.model(input_ids, positions, intermediate_tensors,
506 + inputs_embeds, **kwargs)
507 + return hidden_states
508 +
509 + def compute_logits(
510 + self,
511 + hidden_states: torch.Tensor,
512 + sampling_metadata: SamplingMetadata,
513 + ) -> Optional[torch.Tensor]:
514 + logits = self.logits_processor(self.model.embed_tokens, hidden_states,
515 + sampling_metadata)
516 + return logits
517 +
518 + def sample(
519 + self,
520 + logits: torch.Tensor,
521 + sampling_metadata: SamplingMetadata,
522 + ) -> Optional[SamplerOutput]:
523 + next_tokens = self.sampler(logits, sampling_metadata)
524 + return next_tokens
525 +
526 + def load_weights(self, weights: Iterable[Tuple[str,
527 + torch.Tensor]]) -> Set[str]:
528 + loader = AutoWeightsLoader(
529 + self,
530 + skip_prefixes=(["lm_head."]
531 + if self.config.tie_word_embeddings else None),
532 + )
533 + return loader.load_weights(weights)
vllm/model_executor/models/gemma3_mm.py new
+425
@@ -0,0 +1,425 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import (Any, Iterable, Literal, Mapping, Optional, Sequence, Set,
3 + Tuple, TypedDict, Union)
4 +
5 +import torch
6 +from torch import nn
7 +from transformers import BatchFeature, Gemma3Config, ProcessorMixin
8 +
9 +from vllm.config import VllmConfig
10 +from vllm.logger import init_logger
11 +from vllm.model_executor.layers.layernorm import GemmaRMSNorm
12 +from vllm.model_executor.layers.sampler import SamplerOutput
13 +from vllm.model_executor.sampling_metadata import SamplingMetadata
14 +from vllm.multimodal import MULTIMODAL_REGISTRY
15 +from vllm.multimodal.inputs import (MultiModalFieldConfig, MultiModalKwargs,
16 + NestedTensors)
17 +from vllm.multimodal.parse import ImageSize, MultiModalDataItems
18 +from vllm.multimodal.processing import (BaseMultiModalProcessor,
19 + BaseProcessingInfo, PromptReplacement,
20 + PromptUpdate, PromptUpdateDetails)
21 +from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs
22 +from vllm.sequence import IntermediateTensors
23 +
24 +from .interfaces import SupportsMultiModal, SupportsPP
25 +from .siglip import SiglipVisionModel
26 +from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,
27 + maybe_prefix, merge_multimodal_embeddings)
28 +
29 +logger = init_logger(__name__)
30 +
31 +
32 +class Gemma3ImagePixelInputs(TypedDict):
33 + type: Literal["pixel_values"]
34 + data: torch.Tensor
35 + """Shape: `(batch_size * num_images, num_channels, height, width)`"""
36 +
37 +
38 +Gemma3ImageInputs = Gemma3ImagePixelInputs
39 +
40 +
41 +class Gemma3ProcessingInfo(BaseProcessingInfo):
42 +
43 + def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
44 + return {"image": None}
45 +
46 + def get_mm_max_tokens_per_item(
47 + self,
48 + seq_len: int,
49 + mm_counts: Mapping[str, int],
50 + ) -> Mapping[str, int]:
51 + hf_config = self.ctx.get_hf_config()
52 + return {"image": hf_config.mm_tokens_per_image}
53 +
54 + def get_num_image_tokens(
55 + self,
56 + *,
57 + image_width: int,
58 + image_height: int,
59 + processor: Optional[ProcessorMixin],
60 + ) -> int:
61 + hf_config = self.ctx.get_hf_config()
62 + return hf_config.mm_tokens_per_image
63 +
64 + def get_image_size_with_most_features(self) -> ImageSize:
65 + # Result in the max possible feature size (h:w = 16:1)
66 + return ImageSize(height=8000, width=50)
67 +
68 +
69 +class Gemma3DummyInputsBuilder(BaseDummyInputsBuilder[Gemma3ProcessingInfo]):
70 +
71 + def get_dummy_processor_inputs(
72 + self,
73 + seq_len: int,
74 + mm_counts: Mapping[str, int],
75 + ) -> ProcessorInputs:
76 + tokenizer = self.info.get_tokenizer()
77 + boi_token = tokenizer.boi_token
78 +
79 + num_images = mm_counts.get("image", 0)
80 + target_width, target_height = \
81 + self.info.get_image_size_with_most_features()
82 +
83 + mm_data = {
84 + "image":
85 + self._get_dummy_images(width=target_width,
86 + height=target_height,
87 + num_images=num_images)
88 + }
89 + return ProcessorInputs(
90 + prompt_text=" ".join([boi_token] * num_images),
91 + mm_data=mm_data,
92 + )
93 +
94 +
95 +class Gemma3MultiModalProcessor(BaseMultiModalProcessor[Gemma3ProcessingInfo]):
96 +
97 + def _call_hf_processor(
98 + self,
99 + prompt: str,
100 + mm_data: Mapping[str, object],
101 + mm_kwargs: Mapping[str, object],
102 + ) -> BatchFeature:
103 + # TODO(woosuk): Support pan-and-scan.
104 + img_kwargs = mm_kwargs.get("images_kwargs", {})
105 + img_kwargs["do_pan_and_scan"] = False
106 + mm_kwargs["images_kwargs"] = img_kwargs
107 + return super()._call_hf_processor(
108 + prompt=prompt,
109 + mm_data=mm_data,
110 + mm_kwargs=mm_kwargs,
111 + )
112 +
113 + def _get_mm_fields_config(
114 + self,
115 + hf_inputs: BatchFeature,
116 + hf_processor_mm_kwargs: Mapping[str, object],
117 + ) -> Mapping[str, MultiModalFieldConfig]:
118 + return dict(pixel_values=MultiModalFieldConfig.batched("image"))
119 +
120 + def _get_prompt_updates(
121 + self,
122 + mm_items: MultiModalDataItems,
123 + hf_processor_mm_kwargs: Mapping[str, Any],
124 + out_mm_kwargs: MultiModalKwargs,
125 + ) -> Sequence[PromptUpdate]:
126 + tokenizer = self.info.get_tokenizer()
127 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
128 + hf_config = self.info.get_hf_config()
129 +
130 + boi_token = tokenizer.boi_token
131 + image_token = tokenizer.image_token
132 + mm_tokens_per_image = hf_config.mm_tokens_per_image
133 + image_tokens_expanded = "".join([image_token] * mm_tokens_per_image)
134 +
135 + def get_replacement_gemma3(item_idx: int):
136 + return PromptUpdateDetails(
137 + full=hf_processor.full_image_sequence,
138 + features=image_tokens_expanded,
139 + )
140 +
141 + return [
142 + PromptReplacement(
143 + modality="image",
144 + target=boi_token,
145 + replacement=get_replacement_gemma3,
146 + )
147 + ]
148 +
149 +
150 +class Gemma3MultiModalProjector(nn.Module):
151 +
152 + def __init__(self, config: Gemma3Config):
153 + super().__init__()
154 +
155 + self.mm_input_projection_weight = nn.Parameter(
156 + torch.zeros(config.vision_config.hidden_size,
157 + config.text_config.hidden_size))
158 +
159 + self.mm_soft_emb_norm = GemmaRMSNorm(
160 + config.vision_config.hidden_size,
161 + eps=config.vision_config.layer_norm_eps)
162 +
163 + self.patches_per_image = int(config.vision_config.image_size //
164 + config.vision_config.patch_size)
165 + self.tokens_per_side = int(config.mm_tokens_per_image**0.5)
166 + self.kernel_size = self.patches_per_image // self.tokens_per_side
167 + self.avg_pool = nn.AvgPool2d(kernel_size=self.kernel_size,
168 + stride=self.kernel_size)
169 +
170 + def forward(self, vision_outputs: torch.Tensor):
171 + batch_size, _, seq_length = vision_outputs.shape
172 +
173 + reshaped_vision_outputs = vision_outputs.transpose(1, 2)
174 + reshaped_vision_outputs = reshaped_vision_outputs.reshape(
175 + batch_size, seq_length, self.patches_per_image,
176 + self.patches_per_image)
177 + reshaped_vision_outputs = reshaped_vision_outputs.contiguous()
178 +
179 + pooled_vision_outputs = self.avg_pool(reshaped_vision_outputs)
180 + pooled_vision_outputs = pooled_vision_outputs.flatten(2)
181 + pooled_vision_outputs = pooled_vision_outputs.transpose(1, 2)
182 +
183 + normed_vision_outputs = self.mm_soft_emb_norm(pooled_vision_outputs)
184 +
185 + projected_vision_outputs = torch.matmul(
186 + normed_vision_outputs, self.mm_input_projection_weight)
187 + return projected_vision_outputs.type_as(vision_outputs)
188 +
189 +
190 +@MULTIMODAL_REGISTRY.register_processor(Gemma3MultiModalProcessor,
191 + info=Gemma3ProcessingInfo,
192 + dummy_inputs=Gemma3DummyInputsBuilder)
193 +class Gemma3ForConditionalGeneration(nn.Module, SupportsMultiModal,
194 + SupportsPP):
195 + packed_modules_mapping = {
196 + "qkv_proj": [
197 + "q_proj",
198 + "k_proj",
199 + "v_proj",
200 + ],
201 + "gate_up_proj": [
202 + "gate_proj",
203 + "up_proj",
204 + ],
205 + }
206 +
207 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
208 + super().__init__()
209 + config = vllm_config.model_config.hf_config
210 + quant_config = vllm_config.quant_config
211 + multimodal_config = vllm_config.model_config.multimodal_config
212 + self.config = config
213 + self.quant_config = quant_config
214 + self.multimodal_config = multimodal_config
215 + self.sliding_window = config.text_config.interleaved_sliding_window
216 +
217 + self.vision_tower = SiglipVisionModel(config.vision_config,
218 + quant_config,
219 + prefix=maybe_prefix(
220 + prefix, "vision_tower"))
221 + self.multi_modal_projector = Gemma3MultiModalProjector(config)
222 +
223 + self.language_model = init_vllm_registered_model(
224 + vllm_config=vllm_config,
225 + hf_config=config.text_config,
226 + prefix=maybe_prefix(prefix, "language_model"),
227 + architectures=["Gemma3ForCausalLM"],
228 + )
229 + logit_scale = getattr(config, "logit_scale", 1.0)
230 + self.language_model.logits_processor.scale *= logit_scale
231 +
232 + self.make_empty_intermediate_tensors = (
233 + self.language_model.make_empty_intermediate_tensors)
234 +
235 + @property
236 + def sampler(self):
237 + return self.language_model.sampler
238 +
239 + def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:
240 + h = w = self.config.vision_config.image_size
241 + expected_dims = (3, h, w)
242 +
243 + def _validate_shape(d: torch.Tensor):
244 + if d.shape != expected_dims:
245 + raise ValueError(
246 + "The expected shape of pixel values per image per batch "
247 + f"is {expected_dims}. You supplied {tuple(d.shape)}.")
248 +
249 + for d in data:
250 + _validate_shape(d)
251 +
252 + return data
253 +
254 + def _parse_and_validate_image_input(
255 + self, **kwargs: object) -> Optional[Gemma3ImageInputs]:
256 + pixel_values = kwargs.pop("pixel_values", None)
257 + image_embeds = kwargs.pop("image_embeds", None)
258 + assert image_embeds is None, "Gemma3 does not support image_embeds."
259 + if pixel_values is None:
260 + return None
261 +
262 + if not isinstance(pixel_values, (torch.Tensor, list[torch.Tensor])):
263 + raise ValueError("Incorrect type of pixel values. "
264 + f"Got type: {type(pixel_values)}")
265 +
266 + pixel_values = flatten_bn(pixel_values, concat=True)
267 + return Gemma3ImagePixelInputs(
268 + type="pixel_values",
269 + data=self._validate_pixel_values(pixel_values),
270 + )
271 +
272 + def _image_pixels_to_features(
273 + self,
274 + vision_tower: SiglipVisionModel,
275 + pixel_values: torch.Tensor,
276 + ) -> torch.Tensor:
277 + target_dtype = vision_tower.get_input_embeddings().weight.dtype
278 + image_features = vision_tower(pixel_values.to(dtype=target_dtype))
279 + return image_features
280 +
281 + def _process_image_input(
282 + self,
283 + image_input: Gemma3ImageInputs,
284 + ) -> torch.Tensor:
285 + assert self.vision_tower is not None
286 + pixel_values = image_input["data"]
287 + vision_outputs = self._image_pixels_to_features(
288 + self.vision_tower,
289 + pixel_values,
290 + )
291 + return self.multi_modal_projector(vision_outputs)
292 +
293 + def get_multimodal_embeddings(self, **kwargs) -> Optional[NestedTensors]:
294 + image_input = self._parse_and_validate_image_input(**kwargs)
295 + if image_input is None:
296 + return None
297 + vision_embeddings = self._process_image_input(image_input)
298 + return vision_embeddings
299 +
300 + def get_input_embeddings(
301 + self,
302 + input_ids: torch.Tensor,
303 + multimodal_embeddings: Optional[NestedTensors] = None,
304 + ) -> torch.Tensor:
305 + if multimodal_embeddings is None:
306 + inputs_embeds = self.language_model.get_input_embeddings(input_ids)
307 + else:
308 + inputs_embeds = self.language_model.get_input_embeddings(input_ids)
309 + inputs_embeds = merge_multimodal_embeddings(
310 + input_ids, inputs_embeds, multimodal_embeddings,
311 + self.config.image_token_index)
312 + return inputs_embeds
313 +
314 + def forward(self,
315 + input_ids: torch.Tensor,
316 + positions: torch.Tensor,
317 + intermediate_tensors: Optional[IntermediateTensors] = None,
318 + inputs_embeds: Optional[torch.Tensor] = None,
319 + **kwargs: object) -> Union[SamplerOutput, IntermediateTensors]:
320 + if intermediate_tensors is not None:
321 + inputs_embeds = None
322 +
323 + # NOTE: In v1, inputs_embeds is always generated at model runner, this
324 + # condition is for v0 compatibility.
325 + elif inputs_embeds is None:
326 + vision_embeddings = self.get_multimodal_embeddings(**kwargs)
327 +
328 + inputs_embeds = self.get_input_embeddings(input_ids,
329 + vision_embeddings)
330 + if vision_embeddings is not None:
331 + kwargs = self.prepare_attn_masks(
332 + input_ids,
333 + positions,
334 + mask_dtype=vision_embeddings.dtype,
335 + **kwargs)
336 + input_ids = None
337 +
338 + hidden_states = self.language_model.model(input_ids,
339 + positions,
340 + intermediate_tensors,
341 + inputs_embeds=inputs_embeds,
342 + **kwargs)
343 +
344 + return hidden_states
345 +
346 + def prepare_attn_masks(
347 + self,
348 + input_ids: torch.Tensor,
349 + positions: torch.Tensor,
350 + mask_dtype: torch.dtype,
351 + **kwargs,
352 + ):
353 + kwargs["has_images"] = True
354 + # NOTE(woosuk): Here, we distinguish the sequences by the position id 0.
355 + # This is a HACK. Fix this.
356 + start_idices = (positions == 0).cpu().nonzero()
357 + num_seqs = len(start_idices)
358 + seq_lens = []
359 + for i in range(num_seqs):
360 + start_idx = start_idices[i].item()
361 + if i < num_seqs - 1:
362 + end_idx = start_idices[i + 1].item()
363 + else:
364 + end_idx = len(input_ids)
365 + seq_lens.append(end_idx - start_idx)
366 + kwargs["seq_lens"] = seq_lens
367 +
368 + global_attn_masks = []
369 + local_attn_masks = []
370 + start_idx = 0
371 + for seq_len in seq_lens:
372 + end_idx = start_idx + seq_len
373 + input_token_ids = input_ids[start_idx:end_idx]
374 + start_idx = end_idx
375 + # Create a global causal mask.
376 + global_attn_mask = torch.empty(
377 + 1,
378 + 1,
379 + seq_len,
380 + seq_len,
381 + dtype=mask_dtype,
382 + device=input_ids.device,
383 + )
384 + global_attn_mask.fill_(float("-inf"))
385 + # Fill the lower triangle with 0.
386 + global_attn_mask = global_attn_mask.triu(diagonal=1)
387 +
388 + # Consider the bidirectional attention between image tokens.
389 + img_mask = torch.zeros_like(global_attn_mask)
390 + img_pos = (input_token_ids == self.config.image_token_index)
391 + img_mask[:, :, :, img_pos] += 1
392 + img_mask[:, :, img_pos, :] += 1
393 + global_attn_mask = torch.where(img_mask == 2, 0, global_attn_mask)
394 + global_attn_masks.append(global_attn_mask)
395 +
396 + # Create a local causal mask with sliding window (1024).
397 + local_attn_mask = torch.ones_like(global_attn_mask)
398 + local_attn_mask = torch.tril(local_attn_mask,
399 + diagonal=-self.sliding_window)
400 + local_attn_mask = torch.where(local_attn_mask == 0,
401 + global_attn_mask, float("-inf"))
402 + local_attn_masks.append(local_attn_mask)
403 + kwargs["global_attn_masks"] = global_attn_masks
404 + kwargs["local_attn_masks"] = local_attn_masks
405 + return kwargs
406 +
407 + def compute_logits(
408 + self,
409 + hidden_states: torch.Tensor,
410 + sampling_metadata: SamplingMetadata,
411 + ) -> Optional[torch.Tensor]:
412 + return self.language_model.compute_logits(hidden_states,
413 + sampling_metadata)
414 +
415 + def sample(
416 + self,
417 + logits: torch.Tensor,
418 + sampling_metadata: SamplingMetadata,
419 + ) -> Optional[SamplerOutput]:
420 + return self.language_model.sample(logits, sampling_metadata)
421 +
422 + def load_weights(self, weights: Iterable[Tuple[str,
423 + torch.Tensor]]) -> Set[str]:
424 + loader = AutoWeightsLoader(self)
425 + return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+2
@@ -53,6 +53,7 @@ _TEXT_GENERATION_MODELS = {
53 "Fairseq2LlamaForCausalLM": ("fairseq2_llama", "Fairseq2LlamaForCausalLM"),
54 "GemmaForCausalLM": ("gemma", "GemmaForCausalLM"),
55 "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),
56 + "Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"),
57 "GlmForCausalLM": ("glm", "GlmForCausalLM"),
58 "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
59 "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"),
@@ -161,6 +162,7 @@ _MULTIMODAL_MODELS = {
162 "ChameleonForConditionalGeneration": ("chameleon", "ChameleonForConditionalGeneration"), # noqa: E501
163 "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),
164 "FuyuForCausalLM": ("fuyu", "FuyuForCausalLM"),
165 + "Gemma3ForConditionalGeneration": ("gemma3_mm", "Gemma3ForConditionalGeneration"), # noqa: E501
166 "GLM4VForCausalLM": ("glm4v", "GLM4VForCausalLM"),
167 "H2OVLChatModel": ("h2ovl", "H2OVLChatModel"),
168 "InternVLChatModel": ("internvl", "InternVLChatModel"),