| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Gemma 4 multimodal model (image + audio + video support). |
| 4 | |
| 5 | Adds vision tower, audio tower, and multimodal embedders on top of the |
| 6 | text-only Gemma4ForCausalLM. The vision/audio encoders are loaded via |
| 7 | AutoModel.from_config and run in eager mode while the language model uses |
| 8 | the vLLM-optimized path. |
| 9 | |
| 10 | Video support: Gemma4 does **not** have a native video tower. Videos are |
| 11 | decomposed into timestamped image frames (up to 32 frames at 70 soft tokens |
| 12 | each) and fed through the same vision tower as regular images. The |
| 13 | processor inserts ``mm:ss`` timestamps between frames so the model can |
| 14 | reason about temporal order. |
| 15 | """ |
| 16 | |
| 17 | import math |
| 18 | from collections.abc import Hashable, Iterable, Mapping, Sequence |
| 19 | from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal |
| 20 | |
| 21 | import numpy as np |
| 22 | import torch |
| 23 | from PIL import Image as PILImage |
| 24 | from torch import nn |
| 25 | from transformers import AutoModel, BatchFeature |
| 26 | from transformers.models.gemma4 import ( |
| 27 | Gemma4AudioFeatureExtractor, |
| 28 | Gemma4Config, |
| 29 | Gemma4Processor, |
| 30 | Gemma4VisionConfig, |
| 31 | ) |
| 32 | from transformers.models.gemma4.configuration_gemma4 import ( |
| 33 | Gemma4AudioConfig, |
| 34 | Gemma4TextConfig, |
| 35 | ) |
| 36 | |
| 37 | from vllm.config import VllmConfig |
| 38 | from vllm.config.model import get_served_model_name |
| 39 | from vllm.config.multimodal import MultiModalDummyOptions, VideoDummyOptions |
| 40 | from vllm.inputs import MultiModalDataDict |
| 41 | from vllm.logger import init_logger |
| 42 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 43 | from vllm.model_executor.layers.linear import ReplicatedLinear |
| 44 | from vllm.model_executor.models.gemma3n_mm import batch_audio_features |
| 45 | from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM |
| 46 | from vllm.model_executor.models.module_mapping import MultiModelKeys |
| 47 | from vllm.model_executor.models.transformers.utils import recursive_replace_linear |
| 48 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 49 | from vllm.multimodal.inputs import ( |
| 50 | MultiModalFieldConfig, |
| 51 | MultiModalKwargsItem, |
| 52 | MultiModalKwargsItems, |
| 53 | VideoItem, |
| 54 | ) |
| 55 | from vllm.multimodal.parse import ( |
| 56 | AudioProcessorItems, |
| 57 | ImageProcessorItems, |
| 58 | MultiModalDataItems, |
| 59 | MultiModalDataParser, |
| 60 | ) |
| 61 | from vllm.multimodal.processing import BaseDummyInputsBuilder |
| 62 | from vllm.multimodal.processing.processor import ( |
| 63 | BaseMultiModalProcessor, |
| 64 | BaseProcessingInfo, |
| 65 | PromptReplacement, |
| 66 | PromptUpdate, |
| 67 | PromptUpdateDetails, |
| 68 | ) |
| 69 | from vllm.sequence import IntermediateTensors |
| 70 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 71 | from vllm.utils.tensor_schema import TensorSchema, TensorShape |
| 72 | from vllm.utils.torch_utils import async_tensor_h2d |
| 73 | |
| 74 | from .interfaces import ( |
| 75 | MultiModalEmbeddings, |
| 76 | SupportsEagle3, |
| 77 | SupportsEncoderCudaGraph, |
| 78 | SupportsLoRA, |
| 79 | SupportsMultiModal, |
| 80 | SupportsPP, |
| 81 | SupportsQuant, |
| 82 | ) |
| 83 | from .utils import ( |
| 84 | AutoWeightsLoader, |
| 85 | WeightsMapper, |
| 86 | init_vllm_registered_model, |
| 87 | maybe_prefix, |
| 88 | ) |
| 89 | |
| 90 | if TYPE_CHECKING: |
| 91 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 92 | from vllm.v1.worker.encoder_cudagraph_defs import ( |
| 93 | EncoderCudaGraphCaptureInputs, |
| 94 | EncoderCudaGraphConfig, |
| 95 | EncoderCudaGraphReplayBuffers, |
| 96 | EncoderItemSpec, |
| 97 | ) |
| 98 | |
| 99 | logger = init_logger(__name__) |
| 100 | |
| 101 | # Video constants — match transformers Gemma4VideoProcessor defaults. |
| 102 | _SUPPORTED_SOFT_TOKENS = (70, 140, 280, 560, 1120) |
| 103 | _VIDEO_MAX_SOFT_TOKENS = 70 # soft tokens per video frame (vs 280 for images) |
| 104 | _VIDEO_MAX_FRAMES = 32 # max sampled frames per video |
| 105 | |
| 106 | |
| 107 | def _get_max_soft_tokens( |
| 108 | merged_kwargs: Mapping[str, object], |
| 109 | ) -> tuple[object | None, bool]: |
| 110 | """Return configured image max_soft_tokens and whether it is top-level.""" |
| 111 | val = merged_kwargs.get("max_soft_tokens") |
| 112 | if val is not None: |
| 113 | return val, True |
| 114 | |
| 115 | images_kwargs = merged_kwargs.get("images_kwargs") |
| 116 | if isinstance(images_kwargs, Mapping): |
| 117 | return images_kwargs.get("max_soft_tokens"), False |
| 118 | |
| 119 | return None, False |
| 120 | |
| 121 | |
| 122 | # --------------------------------------------------------------------------- |
| 123 | # Input schema |
| 124 | # --------------------------------------------------------------------------- |
| 125 | |
| 126 | |
| 127 | class Gemma4ImagePixelInputs(TensorSchema): |
| 128 | """Pre-patchified image inputs from the Gemma4 image processor. |
| 129 | |
| 130 | Dimensions: |
| 131 | - bn: Batch size * number of images |
| 132 | - np: Number of patches (max_patches = max_soft_tokens * pooling_kernel_size²) |
| 133 | - pp: Patch pixels (patch_size² * 3) |
| 134 | |
| 135 | The Gemma4 image processor outputs pixel_values as |
| 136 | (batch, max_patches, patch_pixels) — already patchified with |
| 137 | zero-padding for patches beyond the real image content. |
| 138 | pixel_position_ids provides (x, y) coordinates per patch, |
| 139 | with (-1, -1) for padding patches. |
| 140 | """ |
| 141 | |
| 142 | type: Literal["pixel_values"] = "pixel_values" |
| 143 | pixel_values: Annotated[ |
| 144 | torch.Tensor | list[torch.Tensor], |
| 145 | TensorShape("bn", "np", "pp", dynamic_dims={"np"}), |
| 146 | ] |
| 147 | pixel_position_ids: Annotated[ |
| 148 | torch.Tensor | list[torch.Tensor], |
| 149 | TensorShape("bn", "np", 2, dynamic_dims={"np"}), |
| 150 | ] |
| 151 | |
| 152 | |
| 153 | class Gemma4AudioInputs(TensorSchema): |
| 154 | """Dimensions: |
| 155 | - bn: Batch size * number of audios |
| 156 | - s: Sequence length (MEL spectrogram frames) |
| 157 | - f: Number of features (MEL bins) |
| 158 | """ |
| 159 | |
| 160 | type: Literal["audio"] = "audio" |
| 161 | input_features_padded: Annotated[ |
| 162 | torch.Tensor, TensorShape("bn", "s", "f", dynamic_dims={"s"}) |
| 163 | ] |
| 164 | input_features_mask: Annotated[ |
| 165 | torch.Tensor, TensorShape("bn", "s", dynamic_dims={"s"}) |
| 166 | ] |
| 167 | |
| 168 | |
| 169 | Gemma4ImageInputs = Gemma4ImagePixelInputs |
| 170 | |
| 171 | |
| 172 | class Gemma4VideoInputs(TensorSchema): |
| 173 | """Video frame inputs — same tensor format as image inputs. |
| 174 | |
| 175 | Gemma4 has no separate video tower; video frames are processed |
| 176 | through the vision tower at lower resolution (max_soft_tokens=70). |
| 177 | """ |
| 178 | |
| 179 | type: Literal["pixel_values_videos"] = "pixel_values_videos" |
| 180 | pixel_values_videos: Annotated[ |
| 181 | torch.Tensor, |
| 182 | TensorShape("bn", "np", "pp"), |
| 183 | ] |
| 184 | pixel_position_ids_videos: Annotated[ |
| 185 | torch.Tensor, |
| 186 | TensorShape("bn", "np", 2), |
| 187 | ] |
| 188 | video_frame_counts: Annotated[torch.Tensor, TensorShape("nv")] |
| 189 | |
| 190 | |
| 191 | # --------------------------------------------------------------------------- |
| 192 | # Processing info |
| 193 | # --------------------------------------------------------------------------- |
| 194 | |
| 195 | |
| 196 | class Gemma4ProcessingInfo(BaseProcessingInfo): |
| 197 | def get_hf_config(self): |
| 198 | return self.ctx.get_hf_config(Gemma4Config) |
| 199 | |
| 200 | def get_default_tok_params(self): |
| 201 | """Gemma4's chat template already embeds a literal ``<bos>`` token in |
| 202 | the rendered text. If ``add_special_tokens=True`` (the base-class |
| 203 | default), the tokenizer prepends *another* BOS, producing a |
| 204 | ``[2, 2, ...]`` double-BOS sequence that the model was not trained on. |
| 205 | |
| 206 | Setting ``add_special_tokens=False`` here prevents the duplicate and |
| 207 | ensures both ``llm.generate()`` and the chat/completions API behave |
| 208 | correctly for IT models. For PT models (without chat template), we |
| 209 | keep the default (True) to ensure BOS is added for raw prompts. |
| 210 | """ |
| 211 | tokenizer = self.ctx.get_tokenizer() |
| 212 | has_chat_template = getattr(tokenizer, "chat_template", None) is not None |
| 213 | |
| 214 | params = super().get_default_tok_params() |
| 215 | if has_chat_template: |
| 216 | params = params.with_kwargs(add_special_tokens=False) |
| 217 | return params |
| 218 | |
| 219 | def get_hf_processor(self, **kwargs: object) -> Gemma4Processor: |
| 220 | return self.ctx.get_hf_processor( |
| 221 | Gemma4Processor, |
| 222 | **kwargs, |
| 223 | ) |
| 224 | |
| 225 | def validate_num_items(self, modality: str, num_items: int) -> None: |
| 226 | if ( |
| 227 | modality == "audio" |
| 228 | and num_items > 0 |
| 229 | and self.get_hf_config().audio_config is None |
| 230 | ): |
| 231 | model_config = self.ctx.model_config |
| 232 | model = get_served_model_name( |
| 233 | model_config.model, model_config.served_model_name |
| 234 | ) |
| 235 | raise ValueError( |
| 236 | f"Audio input was provided but the model " |
| 237 | f"'{model}' does not have an audio tower. " |
| 238 | f"Audio inference is only supported for Gemma4 " |
| 239 | f"models that include an audio_config " |
| 240 | f"(i.e., models that include an audio_config)." |
| 241 | ) |
| 242 | super().validate_num_items(modality, num_items) |
| 243 | |
| 244 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 245 | limits: dict[str, int | None] = {"image": None} |
| 246 | if self.get_hf_config().audio_config is not None: |
| 247 | limits["audio"] = None |
| 248 | limits["video"] = None |
| 249 | return limits |
| 250 | |
| 251 | def get_mm_max_tokens_per_item( |
| 252 | self, seq_len: int, mm_counts: Mapping[str, int] |
| 253 | ) -> Mapping[str, int] | None: |
| 254 | config = self.get_hf_config() |
| 255 | # Upper bound: the pooler outputs max_soft_tokens slots per image. |
| 256 | # After padding is stripped the actual count is ≤ this value, but |
| 257 | # vLLM needs the max for memory planning. |
| 258 | tokens_per_image = config.vision_config.default_output_length |
| 259 | merged_kwargs = self.ctx.get_merged_mm_kwargs({}) |
| 260 | val, _ = _get_max_soft_tokens(merged_kwargs) |
| 261 | if isinstance(val, int) and val in _SUPPORTED_SOFT_TOKENS: |
| 262 | tokens_per_image = val |
| 263 | tokens: dict[str, int] = {"image": tokens_per_image} |
| 264 | if config.audio_config is not None: |
| 265 | # Audio max tokens from the processor's audio_seq_length. |
| 266 | processor = self.get_hf_processor() |
| 267 | tokens["audio"] = processor.audio_seq_length |
| 268 | # Video: each frame ≤ 70 soft tokens + boi + eoi + ~6 ts tokens. |
| 269 | num_frames = _VIDEO_MAX_FRAMES |
| 270 | mm_config = self.ctx.model_config.get_multimodal_config() |
| 271 | video_opts = mm_config.limit_per_prompt.get("video") |
| 272 | if video_opts is not None and video_opts.num_frames is not None: |
| 273 | num_frames = min(num_frames, video_opts.num_frames) |
| 274 | tokens["video"] = num_frames * (_VIDEO_MAX_SOFT_TOKENS + 2 + 6) |
| 275 | return tokens |
| 276 | |
| 277 | def get_data_parser(self) -> MultiModalDataParser: |
| 278 | config = self.get_hf_config() |
| 279 | kwargs: dict[str, Any] = {"video_needs_metadata": True} |
| 280 | if getattr(config, "audio_config", None) is not None: |
| 281 | processor = self.get_hf_processor() |
| 282 | kwargs["target_sr"] = processor.feature_extractor.sampling_rate |
| 283 | return MultiModalDataParser(**kwargs) |
| 284 | |
| 285 | def _compute_num_soft_tokens( |
| 286 | self, |
| 287 | image_width: int, |
| 288 | image_height: int, |
| 289 | max_soft_tokens: int | None = None, |
| 290 | ) -> int: |
| 291 | """Compute the number of soft tokens the vision tower produces |
| 292 | for an image of the given dimensions, after padding is stripped. |
| 293 | |
| 294 | Args: |
| 295 | image_width: Width of the image in pixels. |
| 296 | image_height: Height of the image in pixels. |
| 297 | max_soft_tokens: Override for the vision config's |
| 298 | ``default_output_length``. When *None*, the value from |
| 299 | the model config is used. |
| 300 | |
| 301 | """ |
| 302 | vision_cfg = self.get_hf_config().vision_config |
| 303 | patch_size = vision_cfg.patch_size |
| 304 | pooling_kernel_size = vision_cfg.pooling_kernel_size |
| 305 | |
| 306 | if max_soft_tokens is None: |
| 307 | max_soft_tokens = vision_cfg.default_output_length |
| 308 | |
| 309 | unit = patch_size * pooling_kernel_size |
| 310 | max_patches = max_soft_tokens * pooling_kernel_size**2 |
| 311 | num_patches_orig = (image_height / patch_size) * (image_width / patch_size) |
| 312 | scale = math.sqrt(max_patches / num_patches_orig) |
| 313 | target_h = max(unit, int(math.floor(image_height * scale / unit)) * unit) |
| 314 | target_w = max(unit, int(math.floor(image_width * scale / unit)) * unit) |
| 315 | num_patches = (target_h // patch_size) * (target_w // patch_size) |
| 316 | # Clamp to ``max_soft_tokens``: extreme aspect ratios (e.g. 3x900) |
| 317 | # cause the floor() above to round one dim up to ``unit`` while the |
| 318 | # other scales freely, which over-shoots ``max_patches``. The HF |
| 319 | # Gemma 4 image processor caps its vision-tower output at |
| 320 | # ``max_soft_tokens``, so without this clamp the prompt-side |
| 321 | # placeholder count exceeds the encoder output and |
| 322 | # ``_merge_multimodal_embeddings`` crashes. |
| 323 | return min(num_patches // (pooling_kernel_size**2), max_soft_tokens) |
| 324 | |
| 325 | def get_image_repl( |
| 326 | self, |
| 327 | *, |
| 328 | image_width: int, |
| 329 | image_height: int, |
| 330 | processor: Gemma4Processor | None, |
| 331 | max_soft_tokens: int | None = None, |
| 332 | ) -> PromptUpdateDetails: |
| 333 | """Return the dynamic image token sequence for this image. |
| 334 | |
| 335 | Computes the exact number of soft tokens the vision tower will |
| 336 | produce after stripping padding. |
| 337 | |
| 338 | Args: |
| 339 | image_width: Width of the image in pixels. |
| 340 | image_height: Height of the image in pixels. |
| 341 | processor: The Gemma4 processor, when already constructed. |
| 342 | max_soft_tokens: Override for the default token budget. |
| 343 | When *None*, falls back to the model config value. |
| 344 | |
| 345 | """ |
| 346 | if processor is None: |
| 347 | processor = self.get_hf_processor() |
| 348 | |
| 349 | num_soft = self._compute_num_soft_tokens( |
| 350 | image_width, |
| 351 | image_height, |
| 352 | max_soft_tokens=max_soft_tokens, |
| 353 | ) |
| 354 | config = self.get_hf_config() |
| 355 | token_ids = ( |
| 356 | [config.boi_token_id] |
| 357 | + [processor.image_token_id] * num_soft |
| 358 | + [config.eoi_token_id] |
| 359 | ) |
| 360 | return PromptUpdateDetails.select_token_id(token_ids, processor.image_token_id) |
| 361 | |
| 362 | @staticmethod |
| 363 | def _compute_audio_num_tokens( |
| 364 | num_samples: int, sampling_rate: int, audio_seq_length: int |
| 365 | ) -> int: |
| 366 | """Replicate the audio encoder's sequence-length arithmetic. |
| 367 | |
| 368 | Mirrors: mel framing (_unfold in Gemma4AudioFeatureExtractor) |
| 369 | followed by two Conv2d subsampling layers (kernel=3, stride=2, |
| 370 | semicausal padding top=1, bottom=1), capped at audio_seq_length. |
| 371 | """ |
| 372 | frame_length = int(round(sampling_rate * 20.0 / 1000.0)) |
| 373 | hop_length = int(round(sampling_rate * 10.0 / 1000.0)) |
| 374 | frame_size_for_unfold = frame_length + 1 |
| 375 | pad_left = frame_length // 2 |
| 376 | padded_samples = num_samples + pad_left |
| 377 | num_mel_frames = (padded_samples - frame_size_for_unfold) // hop_length + 1 |
| 378 | if num_mel_frames <= 0: |
| 379 | return 0 |
| 380 | t = num_mel_frames |
| 381 | for _ in range(2): |
| 382 | t = (t + 2 - 3) // 2 + 1 |
| 383 | return min(t, audio_seq_length) |
| 384 | |
| 385 | def get_audio_repl( |
| 386 | self, |
| 387 | *, |
| 388 | audio_len: int, |
| 389 | processor: Gemma4Processor | None, |
| 390 | ) -> PromptUpdateDetails: |
| 391 | """Return the dynamic audio token sequence for this audio. |
| 392 | |
| 393 | Computes the number of soft tokens from the audio waveform |
| 394 | length by replicating the audio encoder's sequence-length |
| 395 | arithmetic (mel framing + two Conv2d subsampling layers). |
| 396 | """ |
| 397 | if processor is None: |
| 398 | processor = self.get_hf_processor() |
| 399 | |
| 400 | sampling_rate = processor.feature_extractor.sampling_rate |
| 401 | num_tokens = self._compute_audio_num_tokens( |
| 402 | audio_len, sampling_rate, processor.audio_seq_length |
| 403 | ) |
| 404 | config = self.get_hf_config() |
| 405 | token_ids = ( |
| 406 | [config.boa_token_id] |
| 407 | + [processor.audio_token_id] * num_tokens |
| 408 | + [getattr(config, "eoa_token_id", config.eoa_token_index)] |
| 409 | ) |
| 410 | return PromptUpdateDetails.select_token_id(token_ids, processor.audio_token_id) |
| 411 | |
| 412 | def get_video_repl( |
| 413 | self, |
| 414 | *, |
| 415 | timestamps: list[float], |
| 416 | num_soft_tokens_per_frame: list[int], |
| 417 | processor: Gemma4Processor, |
| 418 | ) -> PromptUpdateDetails: |
| 419 | """Build the full token replacement for one video. |
| 420 | |
| 421 | Produces the same interleaved sequence as the HF Gemma4Processor: |
| 422 | mm:ss <boi><|video|>*N<eoi> mm:ss <boi><|video|>*N<eoi> ... |
| 423 | """ |
| 424 | tokenizer = self.ctx.get_tokenizer() |
| 425 | config = self.get_hf_config() |
| 426 | |
| 427 | boi_token_id = config.boi_token_id |
| 428 | eoi_token_id = config.eoi_token_id |
| 429 | video_token_id = processor.video_token_id |
| 430 | |
| 431 | all_token_ids: list[int] = [] |
| 432 | for i, (ts, n_tokens) in enumerate(zip(timestamps, num_soft_tokens_per_frame)): |
| 433 | # mm:ss timestamp — matches transformers: int-truncated, |
| 434 | # zero-padded. |
| 435 | minutes = int(ts // 60) |
| 436 | seconds = int(ts % 60) |
| 437 | ts_str = f"{minutes:02d}:{seconds:02d}" |
| 438 | |
| 439 | prefix = f" {ts_str} " if i > 0 else f"{ts_str} " |
| 440 | ts_token_ids = tokenizer.encode(prefix, add_special_tokens=False) |
| 441 | all_token_ids.extend(ts_token_ids) |
| 442 | |
| 443 | all_token_ids.append(boi_token_id) |
| 444 | all_token_ids.extend([video_token_id] * n_tokens) |
| 445 | all_token_ids.append(eoi_token_id) |
| 446 | |
| 447 | return PromptUpdateDetails.select_token_id(all_token_ids, video_token_id) |
| 448 | |
| 449 | |
| 450 | # --------------------------------------------------------------------------- |
| 451 | # Dummy inputs builder |
| 452 | # --------------------------------------------------------------------------- |
| 453 | |
| 454 | |
| 455 | def _dummy_audio_num_samples(processor: Any) -> int: |
| 456 | """Length in samples of a worst-case dummy audio item. |
| 457 | |
| 458 | ``fft_length`` is set in ``Gemma4AudioFeatureExtractor.__init__`` and is not |
| 459 | defined by ``Gemma4UnifiedAudioFeatureExtractor``, which the encoder-free |
| 460 | Gemma 4 variant uses. Reading it unconditionally therefore raises |
| 461 | ``AttributeError`` during memory profiling on that variant, which is reached |
| 462 | whenever audio is the modality with the largest per-item token count. |
| 463 | |
| 464 | For any other extractor, fall back to the audio budget the processor itself |
| 465 | declares -- the same pair of values used to warn about over-long audio |
| 466 | elsewhere in this file. |
| 467 | """ |
| 468 | feature_extractor = processor.feature_extractor |
| 469 | if isinstance(feature_extractor, Gemma4AudioFeatureExtractor): |
| 470 | return feature_extractor.fft_length |
| 471 | |
| 472 | sampling_rate = feature_extractor.sampling_rate |
| 473 | max_duration_s = processor.audio_seq_length * processor.audio_ms_per_token / 1000.0 |
| 474 | return int(max_duration_s * sampling_rate) |
| 475 | |
| 476 | |
| 477 | class Gemma4DummyInputsBuilder(BaseDummyInputsBuilder[Gemma4ProcessingInfo]): |
| 478 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 479 | num_images = mm_counts.get("image", 0) |
| 480 | num_audios = mm_counts.get("audio", 0) |
| 481 | num_videos = mm_counts.get("video", 0) |
| 482 | processor = self.info.get_hf_processor() |
| 483 | # Use image_token (<|image|>) with tab prefix — this is what the |
| 484 | # Gemma4 chat template inserts per image (\t<|image|>). |
| 485 | # _get_prompt_updates targets image_token and expands it to the |
| 486 | # full_image_sequence. |
| 487 | text = ("\t" + processor.image_token) * num_images |
| 488 | if num_audios > 0 and processor.audio_token: |
| 489 | text += processor.audio_token * num_audios |
| 490 | if num_videos > 0: |
| 491 | text += processor.video_token * num_videos |
| 492 | return text |
| 493 | |
| 494 | def get_dummy_mm_data( |
| 495 | self, |
| 496 | seq_len: int, |
| 497 | mm_counts: Mapping[str, int], |
| 498 | mm_options: MultiModalDummyOptions | None = None, |
| 499 | ) -> MultiModalDataDict: |
| 500 | num_audios = mm_counts.get("audio", 0) |
| 501 | num_videos = mm_counts.get("video", 0) |
| 502 | processor = self.info.get_hf_processor() |
| 503 | image_processor = processor.image_processor |
| 504 | # Use processor's configured image size for dummies. |
| 505 | # Gemma4ImageProcessor sets size=None (it uses patch_size / |
| 506 | # max_soft_tokens instead of the standard size dict), so we |
| 507 | # guard against None with `or {}`. |
| 508 | size = getattr(image_processor, "size", None) or {} |
| 509 | img_width = size.get("width", 224) |
| 510 | img_height = size.get("height", 224) |
| 511 | |
| 512 | data: dict[str, Any] = { |
| 513 | "image": self._get_dummy_images( |
| 514 | width=img_width, |
| 515 | height=img_height, |
| 516 | num_images=mm_counts.get("image", 0), |
| 517 | overrides=mm_options.get("image") if mm_options else None, |
| 518 | ), |
| 519 | } |
| 520 | |
| 521 | if num_audios > 0: |
| 522 | audio_len = _dummy_audio_num_samples(processor) |
| 523 | data["audio"] = self._get_dummy_audios( |
| 524 | length=audio_len, |
| 525 | num_audios=num_audios, |
| 526 | overrides=mm_options.get("audio") if mm_options else None, |
| 527 | ) |
| 528 | |
| 529 | if num_videos > 0: |
| 530 | data["video"] = self._get_dummy_videos( |
| 531 | width=img_width, |
| 532 | height=img_height, |
| 533 | num_frames=_VIDEO_MAX_FRAMES, |
| 534 | num_videos=num_videos, |
| 535 | overrides=mm_options.get("video") if mm_options else None, |
| 536 | ) |
| 537 | |
| 538 | return data |
| 539 | |
| 540 | def _get_dummy_videos( |
| 541 | self, |
| 542 | *, |
| 543 | width: int, |
| 544 | height: int, |
| 545 | num_frames: int, |
| 546 | num_videos: int, |
| 547 | overrides: VideoDummyOptions | None = None, |
| 548 | ) -> list[VideoItem]: |
| 549 | num_frames = max(num_frames, 2) |
| 550 | videos = super()._get_dummy_videos( |
| 551 | width=width, |
| 552 | height=height, |
| 553 | num_frames=num_frames, |
| 554 | num_videos=num_videos, |
| 555 | overrides=overrides, |
| 556 | ) |
| 557 | videos = [v.copy() for v in videos] |
| 558 | |
| 559 | video_items: list[VideoItem] = [] |
| 560 | for video in videos: |
| 561 | video_num_frames = video.shape[0] |
| 562 | video_metadata = { |
| 563 | "fps": 2.0, |
| 564 | "duration": video_num_frames / 2.0, |
| 565 | "total_num_frames": video_num_frames, |
| 566 | "frames_indices": list(range(video_num_frames)), |
| 567 | "video_backend": "opencv", |
| 568 | "do_sample_frames": False, |
| 569 | } |
| 570 | video_items.append((video, video_metadata)) |
| 571 | |
| 572 | return video_items |
| 573 | |
| 574 | |
| 575 | # --------------------------------------------------------------------------- |
| 576 | # Multimodal processor |
| 577 | # --------------------------------------------------------------------------- |
| 578 | |
| 579 | |
| 580 | class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]): |
| 581 | def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: |
| 582 | return self.dummy_inputs.get_dummy_text(mm_counts) |
| 583 | |
| 584 | def _apply_hf_processor_main( |
| 585 | self, |
| 586 | mm_items: MultiModalDataItems, |
| 587 | hf_kwargs: Mapping[str, object], |
| 588 | ) -> BatchFeature: |
| 589 | hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs( |
| 590 | mm_items, hf_kwargs |
| 591 | ) |
| 592 | |
| 593 | if not hf_data: |
| 594 | return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data) |
| 595 | |
| 596 | prompt_text = hf_data.pop("text") |
| 597 | assert isinstance(prompt_text, str) |
| 598 | |
| 599 | merged_kwargs = self.info.ctx.get_merged_mm_kwargs(hf_kwargs) |
| 600 | val, is_top_level_max_soft_tokens = _get_max_soft_tokens(merged_kwargs) |
| 601 | |
| 602 | if val is not None and val not in _SUPPORTED_SOFT_TOKENS: |
| 603 | raise ValueError( |
| 604 | f"Unsupported max_soft_tokens value: {val}. " |
| 605 | f"Valid values are {_SUPPORTED_SOFT_TOKENS}." |
| 606 | ) |
| 607 | |
| 608 | # ---- VIDEO HANDLING ---- |
| 609 | # Gemma4 decomposes video into timestamped image frames. |
| 610 | # Each frame is processed with max_soft_tokens=70 through the |
| 611 | # same vision tower, matching transformers processing_gemma4.py. |
| 612 | video_outputs: dict[str, Any] = {} |
| 613 | if videos := hf_data.pop("videos", []): |
| 614 | assert isinstance(videos, list) |
| 615 | processor = self.info.get_hf_processor() |
| 616 | |
| 617 | all_video_pixel_values: list[torch.Tensor] = [] |
| 618 | all_video_position_ids: list[torch.Tensor] = [] |
| 619 | video_num_soft_tokens_per_video: list[list[int]] = [] |
| 620 | video_timestamps_per_video: list[list[float]] = [] |
| 621 | video_frame_counts: list[int] = [] |
| 622 | |
| 623 | video_replacements: list[str] = [] |
| 624 | |
| 625 | for item in videos: |
| 626 | video_array, metadata = item |
| 627 | |
| 628 | # Convert frames to PIL images |
| 629 | if isinstance(video_array, np.ndarray): |
| 630 | frames = [ |
| 631 | PILImage.fromarray(video_array[i]) |
| 632 | for i in range(video_array.shape[0]) |
| 633 | ] |
| 634 | else: |
| 635 | frames = list(video_array) |
| 636 | |
| 637 | # Compute timestamps from metadata (same as transformers) |
| 638 | fps = metadata.get("fps") or 24 |
| 639 | frame_indices = metadata.get("frames_indices", list(range(len(frames)))) |
| 640 | timestamps = [idx / fps for idx in frame_indices] |
| 641 | |
| 642 | # Process frames as images with max_soft_tokens=70 |
| 643 | video_mm_kwargs = dict(hf_kwargs) |
| 644 | # Override configured image options too: these inputs are video |
| 645 | # frames, whose cache keys exclude images_kwargs. |
| 646 | video_mm_kwargs["images_kwargs"] = {} |
| 647 | video_mm_kwargs["max_soft_tokens"] = _VIDEO_MAX_SOFT_TOKENS |
| 648 | |
| 649 | dummy_prompt = ("\t" + processor.image_token) * len(frames) |
| 650 | |
| 651 | frame_outputs = self.info.ctx.call_hf_processor( |
| 652 | self.info.get_hf_processor(**video_mm_kwargs), |
| 653 | dict(text=dummy_prompt, **{"images": frames}), |
| 654 | video_mm_kwargs, |
| 655 | ) |
| 656 | |
| 657 | # Remap HF key name |
| 658 | if "image_position_ids" in frame_outputs: |
| 659 | frame_outputs["pixel_position_ids"] = frame_outputs.pop( |
| 660 | "image_position_ids" |
| 661 | ) |
| 662 | |
| 663 | all_video_pixel_values.append(frame_outputs["pixel_values"]) |
| 664 | all_video_position_ids.append(frame_outputs["pixel_position_ids"]) |
| 665 | |
| 666 | # Compute soft tokens per frame |
| 667 | num_soft_per_frame = [] |
| 668 | for img in frames: |
| 669 | w, h = img.size |
| 670 | n = self.info._compute_num_soft_tokens( |
| 671 | w, h, max_soft_tokens=_VIDEO_MAX_SOFT_TOKENS |
| 672 | ) |
| 673 | num_soft_per_frame.append(n) |
| 674 | |
| 675 | video_num_soft_tokens_per_video.append(num_soft_per_frame) |
| 676 | video_timestamps_per_video.append(timestamps) |
| 677 | video_frame_counts.append(len(frames)) |
| 678 | |
| 679 | # Build expanded replacement text for this video. |
| 680 | ts_strs = [f"{int(s // 60):02d}:{int(s % 60):02d}" for s in timestamps] |
| 681 | replacement = " ".join( |
| 682 | f"{t} {processor.boi_token}" |
| 683 | f"{processor.video_token * n}" |
| 684 | f"{processor.eoi_token}" |
| 685 | for t, n in zip(ts_strs, num_soft_per_frame) |
| 686 | ) |
| 687 | video_replacements.append(replacement) |
| 688 | |
| 689 | # Replace all <|video|> placeholders at once. We split on |
| 690 | # video_token to get N+1 parts, then interleave with the |
| 691 | # N replacement strings. This avoids the iterative |
| 692 | # split-replace bug where replacement text (which itself |
| 693 | # contains <|video|> tokens) collides with later splits. |
| 694 | vt = processor.video_token |
| 695 | parts = prompt_text.split(vt, len(video_replacements)) |
| 696 | |
| 697 | # NOTE: len(parts) <= len(video_replacements) + 1 |
| 698 | parts_with_repl: list[str] = [] |
| 699 | for part, repl in zip(parts, video_replacements): |
| 700 | parts_with_repl.extend([part, repl]) |
| 701 | parts_with_repl.extend(parts[len(video_replacements) :]) |
| 702 | |
| 703 | prompt_text = "".join(parts_with_repl) |
| 704 | |
| 705 | video_outputs = { |
| 706 | "pixel_values_videos": torch.cat(all_video_pixel_values, dim=0), |
| 707 | "pixel_position_ids_videos": torch.cat(all_video_position_ids, dim=0), |
| 708 | "video_frame_counts": torch.tensor(video_frame_counts), |
| 709 | "video_num_soft_tokens": video_num_soft_tokens_per_video, |
| 710 | "video_timestamps": video_timestamps_per_video, |
| 711 | } |
| 712 | |
| 713 | # Warn if any audio waveform exceeds the model's max duration. |
| 714 | if "audio" in hf_data: |
| 715 | processor = self.info.get_hf_processor() |
| 716 | sr = processor.feature_extractor.sampling_rate |
| 717 | max_tokens = processor.audio_seq_length |
| 718 | ms_per_tok = processor.audio_ms_per_token |
| 719 | max_duration_s = max_tokens * ms_per_tok / 1000.0 |
| 720 | audios = hf_data["audio"] |
| 721 | if not isinstance(audios, (list, tuple)): |
| 722 | audios = [audios] |
| 723 | for i, waveform in enumerate(audios): |
| 724 | duration_s = len(waveform) / sr |
| 725 | if duration_s > max_duration_s: |
| 726 | logger.warning( |
| 727 | "Audio duration exceeds max: %f > %f seconds", |
| 728 | duration_s, |
| 729 | max_duration_s, |
| 730 | ) |
| 731 | # vLLM's call_hf_processor (context.py) re-merges |
| 732 | # mm_processor_kwargs from the model config on every call via: |
| 733 | # config_kwargs | incoming_kwargs (right side wins) |
| 734 | # |
| 735 | # If we strip max_soft_tokens from incoming, the re-merge puts |
| 736 | # back the config's global default (e.g. 280), ignoring any |
| 737 | # per-prompt override. Instead, we keep it in the kwargs with |
| 738 | # the validated per-prompt value so it wins during the merge. |
| 739 | # |
| 740 | # NOTE: This requires a corresponding type annotation on the |
| 741 | # HF side (Gemma4ProcessorKwargs.images_kwargs) so that |
| 742 | # _merge_kwargs routes max_soft_tokens into images_kwargs. |
| 743 | patched_mm_kwargs = dict(hf_kwargs) |
| 744 | if val is not None and is_top_level_max_soft_tokens: |
| 745 | patched_mm_kwargs["max_soft_tokens"] = val |
| 746 | |
| 747 | processed_data = self.info.ctx.call_hf_processor( |
| 748 | self.info.get_hf_processor(**patched_mm_kwargs), |
| 749 | dict(text=prompt_text, **hf_data), |
| 750 | patched_mm_kwargs, |
| 751 | ) |
| 752 | |
| 753 | # HF uses 'image_position_ids'; vLLM uses 'pixel_position_ids'. |
| 754 | # Remap here to keep a single translation point. |
| 755 | if "image_position_ids" in processed_data: |
| 756 | processed_data["pixel_position_ids"] = processed_data.pop( |
| 757 | "image_position_ids" |
| 758 | ) |
| 759 | |
| 760 | if "input_features" in processed_data: |
| 761 | # Unpad per-item so each item's cache entry is |
| 762 | # self-contained. The batched() field config in |
| 763 | # _get_mm_fields_config stacks equal-length items and leaves |
| 764 | # the rest as a list; _process_audio_input re-pads that list |
| 765 | # to the batch's max length. |
| 766 | masks = processed_data["input_features_mask"] |
| 767 | unpadded_features = [ |
| 768 | f[mask] |
| 769 | for f, mask in zip( |
| 770 | processed_data["input_features"], |
| 771 | masks, |
| 772 | ) |
| 773 | ] |
| 774 | unpadded_masks = [mask[mask] for mask in masks] |
| 775 | processed_data["input_features"] = unpadded_features |
| 776 | processed_data["input_features_padded"] = unpadded_features |
| 777 | processed_data["input_features_mask"] = unpadded_masks |
| 778 | |
| 779 | processed_data.update(video_outputs) |
| 780 | return self._finalize_hf_mm_data( |
| 781 | hf_data, hf_kwargs, passthrough_data, processed_data |
| 782 | ) |
| 783 | |
| 784 | def _get_mm_fields_config( |
| 785 | self, |
| 786 | hf_inputs: BatchFeature, |
| 787 | hf_processor_mm_kwargs: Mapping[str, object], |
| 788 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 789 | fields = dict( |
| 790 | pixel_values=MultiModalFieldConfig.batched("image"), |
| 791 | pixel_position_ids=MultiModalFieldConfig.batched("image"), |
| 792 | input_features_padded=MultiModalFieldConfig.batched("audio"), |
| 793 | input_features_mask=MultiModalFieldConfig.batched("audio"), |
| 794 | ) |
| 795 | |
| 796 | # Video fields: frames stored flat, split per video by |
| 797 | # video_frame_counts. |
| 798 | video_frame_counts = hf_inputs.get("video_frame_counts") |
| 799 | if video_frame_counts is not None: |
| 800 | vfc = video_frame_counts |
| 801 | if not isinstance(vfc, torch.Tensor): |
| 802 | vfc = torch.tensor(vfc) |
| 803 | fields.update( |
| 804 | pixel_values_videos=( |
| 805 | MultiModalFieldConfig.flat_from_sizes("video", vfc) |
| 806 | ), |
| 807 | pixel_position_ids_videos=( |
| 808 | MultiModalFieldConfig.flat_from_sizes("video", vfc) |
| 809 | ), |
| 810 | video_frame_counts=MultiModalFieldConfig.batched( |
| 811 | "video", keep_on_cpu=True |
| 812 | ), |
| 813 | video_num_soft_tokens=MultiModalFieldConfig.batched( |
| 814 | "video", keep_on_cpu=True |
| 815 | ), |
| 816 | video_timestamps=MultiModalFieldConfig.batched( |
| 817 | "video", keep_on_cpu=True |
| 818 | ), |
| 819 | ) |
| 820 | |
| 821 | return fields |
| 822 | |
| 823 | def _get_prompt_updates( |
| 824 | self, |
| 825 | mm_items: MultiModalDataItems, |
| 826 | hf_processor_mm_kwargs: Mapping[str, Any], |
| 827 | out_mm_kwargs: MultiModalKwargsItems, |
| 828 | ) -> Sequence[PromptUpdate]: |
| 829 | hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) |
| 830 | |
| 831 | prompt_updates = [] |
| 832 | |
| 833 | if "image" in mm_items: |
| 834 | # Target image_token (<|image|>) — the single placeholder the |
| 835 | # Gemma4 chat template inserts once per image in the prompt. |
| 836 | # vLLM tokenizes the prompt without token expansion, so only |
| 837 | # one image_token exists per image in the token stream. |
| 838 | # The replacement expands it to the full image sequence |
| 839 | # (boi + N×image_token + eoi, where N = max_soft_tokens). |
| 840 | image_token_id = hf_processor.image_token_id |
| 841 | |
| 842 | def get_replacement_image(item_idx: int): |
| 843 | images = mm_items.get_items("image", ImageProcessorItems) |
| 844 | image_size = images.get_image_size(item_idx) |
| 845 | # Resolve the effective max_soft_tokens by merging |
| 846 | # per-prompt kwargs with the config-level defaults, |
| 847 | # consistent with how _apply_hf_processor_main resolves it. |
| 848 | # Without this merge, a missing per-prompt override |
| 849 | # would fall back to vision_cfg.default_output_length |
| 850 | # instead of the config's mm_processor_kwargs default. |
| 851 | merged_kwargs = self.info.ctx.get_merged_mm_kwargs( |
| 852 | hf_processor_mm_kwargs, |
| 853 | ) |
| 854 | val, _ = _get_max_soft_tokens(merged_kwargs) |
| 855 | max_soft_tokens = ( |
| 856 | val |
| 857 | if isinstance(val, int) and val in _SUPPORTED_SOFT_TOKENS |
| 858 | else None |
| 859 | ) |
| 860 | return self.info.get_image_repl( |
| 861 | image_width=image_size.width, |
| 862 | image_height=image_size.height, |
| 863 | processor=hf_processor, |
| 864 | max_soft_tokens=max_soft_tokens, |
| 865 | ) |
| 866 | |
| 867 | prompt_updates.append( |
| 868 | PromptReplacement( |
| 869 | modality="image", |
| 870 | target=[image_token_id], |
| 871 | replacement=get_replacement_image, |
| 872 | ) |
| 873 | ) |
| 874 | |
| 875 | if "video" in mm_items: |
| 876 | video_token_id = hf_processor.video_token_id |
| 877 | |
| 878 | def get_replacement_video(item_idx: int): |
| 879 | out_item = out_mm_kwargs["video"][item_idx] |
| 880 | timestamps_data = out_item["video_timestamps"].data |
| 881 | if isinstance(timestamps_data, torch.Tensor): |
| 882 | timestamps_data = timestamps_data.tolist() |
| 883 | if not isinstance(timestamps_data, list): |
| 884 | raise TypeError("video timestamps must be a list") |
| 885 | timestamps: list[float] = [] |
| 886 | for timestamp in timestamps_data: |
| 887 | if not isinstance(timestamp, (int, float)): |
| 888 | raise TypeError("each video timestamp must be numeric") |
| 889 | timestamps.append(float(timestamp)) |
| 890 | |
| 891 | num_soft_data = out_item["video_num_soft_tokens"].data |
| 892 | if isinstance(num_soft_data, torch.Tensor): |
| 893 | num_soft_data = num_soft_data.tolist() |
| 894 | if not isinstance(num_soft_data, list): |
| 895 | raise TypeError("video soft-token counts must be a list") |
| 896 | num_soft: list[int] = [] |
| 897 | for count in num_soft_data: |
| 898 | if not isinstance(count, int): |
| 899 | raise TypeError( |
| 900 | "each video soft-token count must be an integer" |
| 901 | ) |
| 902 | num_soft.append(count) |
| 903 | return self.info.get_video_repl( |
| 904 | timestamps=timestamps, |
| 905 | num_soft_tokens_per_frame=num_soft, |
| 906 | processor=hf_processor, |
| 907 | ) |
| 908 | |
| 909 | prompt_updates.append( |
| 910 | PromptReplacement( |
| 911 | modality="video", |
| 912 | target=[video_token_id], |
| 913 | replacement=get_replacement_video, |
| 914 | ) |
| 915 | ) |
| 916 | |
| 917 | if "audio" in mm_items: |
| 918 | audio_token_id = hf_processor.audio_token_id |
| 919 | |
| 920 | def get_replacement_audio(item_idx: int): |
| 921 | audios = mm_items.get_items("audio", AudioProcessorItems) |
| 922 | audio_len = audios.get_audio_length(item_idx) |
| 923 | return self.info.get_audio_repl( |
| 924 | audio_len=audio_len, |
| 925 | processor=hf_processor, |
| 926 | ) |
| 927 | |
| 928 | prompt_updates.append( |
| 929 | PromptReplacement( |
| 930 | modality="audio", |
| 931 | target=[audio_token_id], |
| 932 | replacement=get_replacement_audio, |
| 933 | ) |
| 934 | ) |
| 935 | |
| 936 | return prompt_updates |
| 937 | |
| 938 | # NOTE: Gemma3/Gemma3n override _apply_token_matches and |
| 939 | # _find_mm_placeholders to merge adjacent newline tokens that arise |
| 940 | # when full_image_sequence contains "\n\n" wrappers. Gemma4's |
| 941 | # full_image_sequence has NO newlines (just BOI + 280×image_token + |
| 942 | # EOI), so the base class implementations work correctly as-is. |
| 943 | |
| 944 | |
| 945 | # --------------------------------------------------------------------------- |
| 946 | # Multimodal embedder |
| 947 | # --------------------------------------------------------------------------- |
| 948 | |
| 949 | |
| 950 | class Gemma4MultimodalEmbedder(nn.Module): |
| 951 | """Projects vision/audio soft tokens into LM embedding space. |
| 952 | |
| 953 | Architecture: |
| 954 | inputs_embeds → embedding_projection → embedding_post_projection_norm |
| 955 | |
| 956 | Unlike Gemma3n which has separate hard/soft embedding paths with |
| 957 | per-path normalization and a learned embedding table, Gemma4 uses a |
| 958 | simplified 2-layer design: a linear projection followed by RMSNorm |
| 959 | (without learnable scale). The checkpoint confirms this — only |
| 960 | ``embedding_projection.weight`` exists; there is no embedding table |
| 961 | or pre-projection norm weights. |
| 962 | """ |
| 963 | |
| 964 | def __init__( |
| 965 | self, |
| 966 | multimodal_config: Gemma4VisionConfig | Gemma4AudioConfig, |
| 967 | text_config: Gemma4TextConfig, |
| 968 | *, |
| 969 | quant_config: "QuantizationConfig | None" = None, |
| 970 | prefix: str = "", |
| 971 | ): |
| 972 | super().__init__() |
| 973 | |
| 974 | self.eps = multimodal_config.rms_norm_eps |
| 975 | self.text_hidden_size = text_config.hidden_size |
| 976 | |
| 977 | # Audio tower uses output_proj_dims (1536) rather than hidden_size |
| 978 | # (1024); vision uses hidden_size (768) directly. |
| 979 | embedding_dim = ( |
| 980 | getattr(multimodal_config, "output_proj_dims", None) |
| 981 | or multimodal_config.hidden_size |
| 982 | ) |
| 983 | |
| 984 | self.embedding_pre_projection_norm = RMSNorm( |
| 985 | embedding_dim, |
| 986 | eps=self.eps, |
| 987 | has_weight=False, |
| 988 | ) |
| 989 | |
| 990 | self.embedding_projection = ReplicatedLinear( |
| 991 | embedding_dim, |
| 992 | self.text_hidden_size, |
| 993 | bias=False, |
| 994 | quant_config=quant_config, |
| 995 | prefix=maybe_prefix(prefix, "embedding_projection"), |
| 996 | ) |
| 997 | |
| 998 | def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor: |
| 999 | """Project soft tokens from a multimodal tower into LM space.""" |
| 1000 | embs_normed = self.embedding_pre_projection_norm(inputs_embeds) |
| 1001 | embs_proj, _ = self.embedding_projection(embs_normed) |
| 1002 | return embs_proj |
| 1003 | |
| 1004 | |
| 1005 | # --------------------------------------------------------------------------- |
| 1006 | # Main model |
| 1007 | # --------------------------------------------------------------------------- |
| 1008 | |
| 1009 | |
| 1010 | @MULTIMODAL_REGISTRY.register_processor( |
| 1011 | Gemma4MultiModalProcessor, |
| 1012 | info=Gemma4ProcessingInfo, |
| 1013 | dummy_inputs=Gemma4DummyInputsBuilder, |
| 1014 | ) |
| 1015 | class Gemma4ForConditionalGeneration( |
| 1016 | nn.Module, |
| 1017 | SupportsMultiModal, |
| 1018 | SupportsQuant, |
| 1019 | SupportsPP, |
| 1020 | SupportsLoRA, |
| 1021 | SupportsEagle3, |
| 1022 | SupportsEncoderCudaGraph, |
| 1023 | ): |
| 1024 | supports_encoder_cudagraph: ClassVar[Literal[True]] = True |
| 1025 | # Gemma4 clamps mm_prefix bidirectional ranges to the sliding window |
| 1026 | # in-kernel (HF's (causal OR blockwise) AND sliding_window). The model |
| 1027 | # runner reads this to keep image bidirectional ranges that exceed the |
| 1028 | # window instead of dropping them (which would make image attention |
| 1029 | # causal-only for images larger than the sliding window). |
| 1030 | mm_prefix_clamp_sliding_window: bool = True |
| 1031 | supports_tower_connector_lora = True |
| 1032 | |
| 1033 | packed_modules_mapping: dict[str, list[str]] = { |
| 1034 | "qkv_proj": [ |
| 1035 | "q_proj", |
| 1036 | "k_proj", |
| 1037 | "v_proj", |
| 1038 | ], |
| 1039 | "gate_up_proj": [ |
| 1040 | "gate_proj", |
| 1041 | "up_proj", |
| 1042 | ], |
| 1043 | } |
| 1044 | |
| 1045 | # Maps checkpoint prefixes to vLLM module paths. |
| 1046 | hf_to_vllm_mapper = WeightsMapper( |
| 1047 | orig_to_new_prefix={ |
| 1048 | # vision tower |
| 1049 | "model.vision_tower": "vision_tower", |
| 1050 | "model.embed_vision": "embed_vision", |
| 1051 | # audio tower |
| 1052 | "model.audio_tower.": "audio_tower.", |
| 1053 | "model.embed_audio.": "embed_audio.", |
| 1054 | # backbone |
| 1055 | "model.language_model.": "language_model.model.", |
| 1056 | "lm_head.": "language_model.lm_head.", |
| 1057 | "model": "language_model.model", |
| 1058 | }, |
| 1059 | ) |
| 1060 | |
| 1061 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 1062 | super().__init__() |
| 1063 | config = vllm_config.model_config.hf_config |
| 1064 | quant_config = vllm_config.quant_config |
| 1065 | multimodal_config = vllm_config.model_config.multimodal_config |
| 1066 | self.config = config |
| 1067 | self.quant_config = quant_config |
| 1068 | self.multimodal_config = multimodal_config |
| 1069 | self.model_dtype = vllm_config.model_config.dtype |
| 1070 | self.vllm_config = vllm_config |
| 1071 | lora_config = vllm_config.lora_config |
| 1072 | self._enable_mm_lora = bool( |
| 1073 | lora_config is not None and lora_config.enable_tower_connector_lora |
| 1074 | ) |
| 1075 | |
| 1076 | # Only quantize towers when the quant method supports their |
| 1077 | # dimensions. BNB/torchao handle arbitrary sizes; other methods |
| 1078 | # (Marlin, FP8, …) require dimensions divisible by 64, which |
| 1079 | # the vision tower (intermediate_size=4304) does not satisfy. |
| 1080 | # TODO(mgoin): remove this by fixing kernel padding. |
| 1081 | tower_quant: QuantizationConfig | None |
| 1082 | if quant_config and quant_config.get_name() in [ |
| 1083 | "bitsandbytes", |
| 1084 | "torchao", |
| 1085 | "compressed-tensors", |
| 1086 | ]: |
| 1087 | tower_quant = quant_config |
| 1088 | else: |
| 1089 | vision_cfg = config.vision_config |
| 1090 | quantizable = ( |
| 1091 | vision_cfg.hidden_size % 64 == 0 |
| 1092 | and vision_cfg.intermediate_size % 64 == 0 |
| 1093 | ) |
| 1094 | tower_quant = quant_config if quantizable else None |
| 1095 | |
| 1096 | # ---- Vision tower (shared by image and video) ---- |
| 1097 | with self._mark_tower_model(vllm_config, {"image", "video"}): |
| 1098 | self.vision_tower = AutoModel.from_config(config=config.vision_config) |
| 1099 | self.embed_vision = Gemma4MultimodalEmbedder( |
| 1100 | config.vision_config, |
| 1101 | config.text_config, |
| 1102 | quant_config=tower_quant, |
| 1103 | prefix=maybe_prefix(prefix, "embed_vision"), |
| 1104 | ) |
| 1105 | recursive_replace_linear( |
| 1106 | self.vision_tower, |
| 1107 | tower_quant, |
| 1108 | prefix=maybe_prefix(prefix, "vision_tower"), |
| 1109 | ) |
| 1110 | |
| 1111 | # ---- Audio tower (variants with audio_config) ---- |
| 1112 | self.embed_audio: Gemma4MultimodalEmbedder | None |
| 1113 | if config.audio_config is not None: |
| 1114 | with self._mark_tower_model(vllm_config, "audio"): |
| 1115 | self.audio_tower = AutoModel.from_config(config=config.audio_config) |
| 1116 | # AutoModel.from_config does NOT call post_init(), |
| 1117 | # which is needed to initialize buffers that are absent |
| 1118 | # from the checkpoint (e.g. inv_timescales for relative |
| 1119 | # position embeddings, softcap, gradient_clipping). |
| 1120 | self.audio_tower.post_init() |
| 1121 | self.embed_audio = Gemma4MultimodalEmbedder( |
| 1122 | config.audio_config, |
| 1123 | config.text_config, |
| 1124 | quant_config=tower_quant, |
| 1125 | prefix=maybe_prefix(prefix, "embed_audio"), |
| 1126 | ) |
| 1127 | recursive_replace_linear( |
| 1128 | self.audio_tower, |
| 1129 | tower_quant, |
| 1130 | prefix=maybe_prefix(prefix, "audio_tower"), |
| 1131 | ) |
| 1132 | else: |
| 1133 | self.audio_tower = None |
| 1134 | self.embed_audio = None |
| 1135 | |
| 1136 | # ---- Language model (vLLM optimised) ---- |
| 1137 | with self._mark_language_model(vllm_config): |
| 1138 | self.language_model: Gemma4ForCausalLM = init_vllm_registered_model( |
| 1139 | vllm_config=vllm_config, |
| 1140 | hf_config=config.text_config, |
| 1141 | prefix=maybe_prefix(prefix, "language_model"), |
| 1142 | architectures=["Gemma4ForCausalLM"], |
| 1143 | ) |
| 1144 | |
| 1145 | # Pre-allocate PLE buffer for CUDA graph compatibility. |
| 1146 | # Some variants have hidden_size_per_layer_input=None (no PLE). |
| 1147 | ple_dim = config.text_config.hidden_size_per_layer_input |
| 1148 | if ple_dim is not None and ple_dim > 0: |
| 1149 | embed = self.language_model.model.embed_tokens |
| 1150 | self.per_layer_embeddings = torch.zeros( |
| 1151 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 1152 | config.text_config.num_hidden_layers, |
| 1153 | ple_dim, |
| 1154 | device=next(embed.parameters()).device, |
| 1155 | dtype=vllm_config.model_config.dtype, |
| 1156 | ) |
| 1157 | else: |
| 1158 | self.per_layer_embeddings = None |
| 1159 | |
| 1160 | self.make_empty_intermediate_tensors = ( |
| 1161 | self.language_model.make_empty_intermediate_tensors |
| 1162 | ) |
| 1163 | |
| 1164 | # --- Precompute full-attention layer indices for bidi clearing --- |
| 1165 | self._full_attn_layer_idxs: frozenset[int] = frozenset() |
| 1166 | text_config = config.text_config |
| 1167 | if getattr(text_config, "use_bidirectional_attention", None) == "vision": |
| 1168 | layer_types = getattr(text_config, "layer_types", None) |
| 1169 | if layer_types: |
| 1170 | self._full_attn_layer_idxs = frozenset( |
| 1171 | i for i, lt in enumerate(layer_types) if lt != "sliding_attention" |
| 1172 | ) |
| 1173 | |
| 1174 | # --- MixtureOfExperts delegation to language_model --- |
| 1175 | self.moe_layers = self.language_model.moe_layers |
| 1176 | self.num_moe_layers = self.language_model.num_moe_layers |
| 1177 | self.num_logical_experts = self.language_model.num_logical_experts |
| 1178 | self.num_physical_experts = self.language_model.num_physical_experts |
| 1179 | self.num_local_physical_experts = self.language_model.num_local_physical_experts |
| 1180 | self.num_routed_experts = self.language_model.num_routed_experts |
| 1181 | self.num_expert_groups = self.language_model.num_expert_groups |
| 1182 | self.num_shared_experts = self.language_model.num_shared_experts |
| 1183 | self.num_redundant_experts = self.language_model.num_redundant_experts |
| 1184 | |
| 1185 | gen_cfg = vllm_config.model_config.try_get_generation_config() |
| 1186 | self._suppress_token_ids = gen_cfg.get("suppress_tokens") if gen_cfg else None |
| 1187 | |
| 1188 | # ------------------------------------------------------------------ # |
| 1189 | # Input parsing |
| 1190 | # ------------------------------------------------------------------ # |
| 1191 | |
| 1192 | def _parse_and_validate_image_input( |
| 1193 | self, **kwargs: object |
| 1194 | ) -> Gemma4ImageInputs | None: |
| 1195 | pixel_values = kwargs.pop("pixel_values", None) |
| 1196 | pixel_position_ids = kwargs.pop("pixel_position_ids", None) |
| 1197 | image_embeds = kwargs.pop("image_embeds", None) |
| 1198 | assert image_embeds is None, "Gemma4 does not support image_embeds." |
| 1199 | if pixel_values is None: |
| 1200 | return None |
| 1201 | return Gemma4ImagePixelInputs( |
| 1202 | pixel_values=pixel_values, |
| 1203 | pixel_position_ids=pixel_position_ids, |
| 1204 | ) |
| 1205 | |
| 1206 | def _parse_and_validate_audio_input( |
| 1207 | self, **kwargs: object |
| 1208 | ) -> Gemma4AudioInputs | None: |
| 1209 | input_features_padded = kwargs.pop("input_features_padded", None) |
| 1210 | if input_features_padded is None: |
| 1211 | return None |
| 1212 | input_features_mask = kwargs.pop("input_features_mask", None) |
| 1213 | if input_features_mask is None: |
| 1214 | return None |
| 1215 | return Gemma4AudioInputs( |
| 1216 | input_features_padded=input_features_padded, |
| 1217 | input_features_mask=input_features_mask, |
| 1218 | ) |
| 1219 | |
| 1220 | def _parse_and_validate_video_input( |
| 1221 | self, **kwargs: object |
| 1222 | ) -> Gemma4VideoInputs | None: |
| 1223 | pixel_values_videos = kwargs.pop("pixel_values_videos", None) |
| 1224 | pixel_position_ids_videos = kwargs.pop("pixel_position_ids_videos", None) |
| 1225 | video_frame_counts = kwargs.pop("video_frame_counts", None) |
| 1226 | if pixel_values_videos is None: |
| 1227 | return None |
| 1228 | return Gemma4VideoInputs( |
| 1229 | pixel_values_videos=pixel_values_videos, |
| 1230 | pixel_position_ids_videos=pixel_position_ids_videos, |
| 1231 | video_frame_counts=video_frame_counts, |
| 1232 | ) |
| 1233 | |
| 1234 | def _parse_and_validate_multimodal_inputs( |
| 1235 | self, **kwargs: object |
| 1236 | ) -> dict[str, Gemma4ImageInputs | Gemma4AudioInputs | Gemma4VideoInputs | None]: |
| 1237 | mm_input_by_modality: dict[ |
| 1238 | str, Gemma4ImageInputs | Gemma4AudioInputs | Gemma4VideoInputs | None |
| 1239 | ] = {} |
| 1240 | for input_key in list(kwargs): |
| 1241 | if ( |
| 1242 | input_key in ("pixel_values", "image_embeds") |
| 1243 | and "image" not in mm_input_by_modality |
| 1244 | ): |
| 1245 | mm_input_by_modality["image"] = self._parse_and_validate_image_input( |
| 1246 | **kwargs |
| 1247 | ) |
| 1248 | if ( |
| 1249 | input_key == "pixel_values_videos" |
| 1250 | and "video" not in mm_input_by_modality |
| 1251 | ): |
| 1252 | mm_input_by_modality["video"] = self._parse_and_validate_video_input( |
| 1253 | **kwargs |
| 1254 | ) |
| 1255 | if ( |
| 1256 | input_key == "input_features_padded" |
| 1257 | and "audio" not in mm_input_by_modality |
| 1258 | ): |
| 1259 | mm_input_by_modality["audio"] = self._parse_and_validate_audio_input( |
| 1260 | **kwargs |
| 1261 | ) |
| 1262 | return mm_input_by_modality |
| 1263 | |
| 1264 | @staticmethod |
| 1265 | def _encoder_chunk( |
| 1266 | patches_per_item: int, |
| 1267 | free_bytes: int, |
| 1268 | total_bytes: int, |
| 1269 | position_embedding_size: int, |
| 1270 | ) -> int: |
| 1271 | """Max chunk size whose F.one_hot transient fits in the budget. |
| 1272 | |
| 1273 | The dominant transient inside HF's ``Gemma4VisionPatchEmbedder. |
| 1274 | _position_embeddings`` is |
| 1275 | ``F.one_hot(clamped_positions, num_classes=position_embedding_size)`` |
| 1276 | with shape ``(chunk, patches, 2, position_embedding_size)``, |
| 1277 | int64, plus its simultaneous cast to the position embedding |
| 1278 | table dtype. That, not the encoder residual stream, sets peak |
| 1279 | memory. |
| 1280 | """ |
| 1281 | if patches_per_item <= 0: |
| 1282 | return 1 |
| 1283 | # Half of currently-free, capped at 10% of total so we leave room |
| 1284 | # for the rest of profile_run / the subsequent encoder + pooler. |
| 1285 | budget = min(free_bytes // 2, total_bytes // 10) |
| 1286 | if budget <= 0: |
| 1287 | return 1 |
| 1288 | # F.one_hot allocates (chunk, patches, 2, pos_emb_size) int64 |
| 1289 | # (the inner 2 is the (x, y) coordinate axis, 8 is sizeof(int64)). |
| 1290 | # Outer 2x covers the int64 buffer and its concurrent bf16 cast |
| 1291 | # plus the matmul output that live alongside it at peak. |
| 1292 | cost = patches_per_item * 4 * position_embedding_size * 8 |
| 1293 | return max(1, budget // cost) if cost > 0 else 1 |
| 1294 | |
| 1295 | # ------------------------------------------------------------------ # |
| 1296 | # Image processing |
| 1297 | # ------------------------------------------------------------------ # |
| 1298 | |
| 1299 | def _process_image_input( |
| 1300 | self, |
| 1301 | image_input: Gemma4ImageInputs, |
| 1302 | ) -> list[torch.Tensor]: |
| 1303 | """Batch-encode images through the vision tower. |
| 1304 | |
| 1305 | Groups images by patch count (resolution bucket) so each |
| 1306 | encoder call processes a uniform-shape batch with no |
| 1307 | cross-resolution padding. With MM LoRA enabled, all images are |
| 1308 | padded into one batch so the encoder call matches the tower mapping. |
| 1309 | """ |
| 1310 | pixel_values = image_input["pixel_values"] |
| 1311 | pixel_position_ids = image_input["pixel_position_ids"] |
| 1312 | |
| 1313 | vt = self.vision_tower |
| 1314 | vision_cfg = self.config.vision_config |
| 1315 | pooling_k2 = vision_cfg.pooling_kernel_size**2 |
| 1316 | |
| 1317 | # Concurrent requests with different image resolutions may |
| 1318 | # arrive as a list of per-image tensors, while same-resolution |
| 1319 | # batches may arrive as a stacked tensor. |
| 1320 | buckets: dict[int, list[tuple[int, torch.Tensor, torch.Tensor]]] = {} |
| 1321 | total_images = ( |
| 1322 | len(pixel_values) |
| 1323 | if isinstance(pixel_values, list) |
| 1324 | else pixel_values.shape[0] |
| 1325 | ) |
| 1326 | pool_position_ids = pixel_position_ids |
| 1327 | |
| 1328 | if self._enable_mm_lora: |
| 1329 | max_soft_tokens = vision_cfg.default_output_length |
| 1330 | mm_processor_kwargs = getattr( |
| 1331 | getattr(self, "multimodal_config", None), |
| 1332 | "mm_processor_kwargs", |
| 1333 | None, |
| 1334 | ) |
| 1335 | if isinstance(mm_processor_kwargs, Mapping): |
| 1336 | value, _ = _get_max_soft_tokens(mm_processor_kwargs) |
| 1337 | if isinstance(value, int) and value in _SUPPORTED_SOFT_TOKENS: |
| 1338 | max_soft_tokens = value |
| 1339 | |
| 1340 | max_patches = max_soft_tokens * pooling_k2 |
| 1341 | padded_position_ids: list[torch.Tensor] = [] |
| 1342 | for idx in range(total_images): |
| 1343 | pv = pixel_values[idx] |
| 1344 | pp = pixel_position_ids[idx] |
| 1345 | num_patches = pv.shape[0] |
| 1346 | if num_patches > max_patches: |
| 1347 | raise ValueError( |
| 1348 | f"Image {idx} has {num_patches} patches, which exceeds " |
| 1349 | f"the MM LoRA patch limit of {max_patches}." |
| 1350 | ) |
| 1351 | |
| 1352 | pad_len = max_patches - num_patches |
| 1353 | pv = torch.cat( |
| 1354 | (pv, pv.new_zeros((pad_len, *pv.shape[1:]))), |
| 1355 | dim=0, |
| 1356 | ) |
| 1357 | pp = torch.cat( |
| 1358 | (pp, pp.new_full((pad_len, *pp.shape[1:]), -1)), |
| 1359 | dim=0, |
| 1360 | ) |
| 1361 | buckets.setdefault(max_patches, []).append((idx, pv, pp)) |
| 1362 | padded_position_ids.append(pp) |
| 1363 | pool_position_ids = padded_position_ids |
| 1364 | else: |
| 1365 | for idx in range(total_images): |
| 1366 | pv = pixel_values[idx] |
| 1367 | pp = pixel_position_ids[idx] |
| 1368 | buckets.setdefault(pv.shape[0], []).append((idx, pv, pp)) |
| 1369 | |
| 1370 | # Encode each resolution bucket in memory-safe chunks. Re-read |
| 1371 | # free memory per bucket because the previous bucket's encoder |
| 1372 | # pass has already allocated activations we should account for. |
| 1373 | last_hidden_states_map: dict[int, torch.Tensor] = {} |
| 1374 | for patches, items in buckets.items(): |
| 1375 | if self._enable_mm_lora: |
| 1376 | max_batch_size = len(items) |
| 1377 | else: |
| 1378 | free, total = torch.accelerator.get_memory_info() |
| 1379 | max_batch_size = min( |
| 1380 | len(items), |
| 1381 | self._encoder_chunk( |
| 1382 | patches, free, total, vision_cfg.position_embedding_size |
| 1383 | ), |
| 1384 | ) |
| 1385 | |
| 1386 | for chunk_idx in range(0, len(items), max_batch_size): |
| 1387 | chunk_items = items[chunk_idx : chunk_idx + max_batch_size] |
| 1388 | |
| 1389 | pv_tensor = torch.cat( |
| 1390 | [item[1].unsqueeze(0) for item in chunk_items], dim=0 |
| 1391 | ) |
| 1392 | pp_tensor = torch.cat( |
| 1393 | [item[2].unsqueeze(0) for item in chunk_items], dim=0 |
| 1394 | ) |
| 1395 | pad_tensor = (pp_tensor == -1).all(dim=-1) |
| 1396 | |
| 1397 | inputs_embeds = vt.patch_embedder( |
| 1398 | pv_tensor, |
| 1399 | pp_tensor, |
| 1400 | pad_tensor, |
| 1401 | ).to(self.model_dtype) |
| 1402 | # HuggingFace's mask builder probes `padding_mask.all()` to |
| 1403 | # decide whether the mask can be skipped, which syncs. |
| 1404 | with gpu_sync_allowed(): |
| 1405 | encoder_outputs = vt.encoder( |
| 1406 | inputs_embeds=inputs_embeds, |
| 1407 | attention_mask=~pad_tensor, |
| 1408 | pixel_position_ids=pp_tensor, |
| 1409 | ) |
| 1410 | hidden_states = encoder_outputs.last_hidden_state |
| 1411 | |
| 1412 | for i, (orig_idx, _, _) in enumerate(chunk_items): |
| 1413 | last_hidden_states_map[orig_idx] = hidden_states[i] |
| 1414 | |
| 1415 | # Pool per image to strip padding and reduce spatial resolution. |
| 1416 | all_valid_states: list[torch.Tensor] = [None] * total_images # type: ignore[list-item] |
| 1417 | valid_lens = [0] * total_images |
| 1418 | |
| 1419 | for orig_idx in range(total_images): |
| 1420 | chunk_hidden = last_hidden_states_map[orig_idx] |
| 1421 | output_length = chunk_hidden.shape[0] // pooling_k2 |
| 1422 | |
| 1423 | single_hidden = chunk_hidden.unsqueeze(0) |
| 1424 | single_pos_ids = pool_position_ids[orig_idx].unsqueeze(0) |
| 1425 | padding_positions = (single_pos_ids == -1).all(dim=-1) |
| 1426 | |
| 1427 | # The pooler goes through HuggingFace's mask builder, which probes |
| 1428 | # `padding_mask.all()`, and the mask indexing below needs the |
| 1429 | # selected count on the host. |
| 1430 | with gpu_sync_allowed(): |
| 1431 | pooled_states, valid_mask = vt.pooler( |
| 1432 | hidden_states=single_hidden, |
| 1433 | pixel_position_ids=single_pos_ids, |
| 1434 | padding_positions=padding_positions, |
| 1435 | output_length=output_length, |
| 1436 | ) |
| 1437 | valid_states = pooled_states[valid_mask] |
| 1438 | |
| 1439 | if getattr(vt.config, "standardize", False): |
| 1440 | valid_states = (valid_states - vt.std_bias) * vt.std_scale |
| 1441 | |
| 1442 | all_valid_states[orig_idx] = valid_states |
| 1443 | valid_lens[orig_idx] = valid_states.shape[0] |
| 1444 | |
| 1445 | # Project all images in a single batched call. |
| 1446 | flat_valid_states = torch.cat(all_valid_states, dim=0).to(self.model_dtype) |
| 1447 | flat_proj_embs = self.embed_vision( |
| 1448 | inputs_embeds=flat_valid_states.unsqueeze(0) |
| 1449 | ).squeeze(0) |
| 1450 | |
| 1451 | # Split back into per-image tensors (slicing returns views). |
| 1452 | per_image_embeddings: list[torch.Tensor] = [] |
| 1453 | offset = 0 |
| 1454 | for length in valid_lens: |
| 1455 | per_image_embeddings.append(flat_proj_embs[offset : offset + length]) |
| 1456 | offset += length |
| 1457 | |
| 1458 | return per_image_embeddings |
| 1459 | |
| 1460 | # ------------------------------------------------------------------ # |
| 1461 | # Video processing (frames through vision tower) |
| 1462 | # ------------------------------------------------------------------ # |
| 1463 | |
| 1464 | def _process_video_input( |
| 1465 | self, |
| 1466 | video_input: Gemma4VideoInputs, |
| 1467 | ) -> list[torch.Tensor]: |
| 1468 | """Batch-encode video frames through the vision tower. |
| 1469 | |
| 1470 | Gemma4 has no separate video tower; video frames are images at |
| 1471 | lower resolution (max_soft_tokens=70). All frames across all |
| 1472 | videos in the batch are encoded together in chunks, then pooled |
| 1473 | and projected in a single batched call. |
| 1474 | |
| 1475 | Returns one concatenated embedding tensor per video (not per |
| 1476 | frame), matching the flat_from_sizes grouping that vLLM expects |
| 1477 | for embed_multimodal. |
| 1478 | """ |
| 1479 | pixel_values = video_input["pixel_values_videos"] |
| 1480 | pixel_position_ids = video_input["pixel_position_ids_videos"] |
| 1481 | frame_counts = video_input["video_frame_counts"] |
| 1482 | |
| 1483 | vt = self.vision_tower |
| 1484 | vision_cfg = self.config.vision_config |
| 1485 | pooling_k2 = vision_cfg.pooling_kernel_size**2 |
| 1486 | |
| 1487 | if isinstance(frame_counts, torch.Tensor): |
| 1488 | # Per-video frame counts drive the Python-level batching below. |
| 1489 | with gpu_sync_allowed(): |
| 1490 | fc_list = frame_counts.tolist() |
| 1491 | else: |
| 1492 | fc_list = list(frame_counts) |
| 1493 | |
| 1494 | total_frames = pixel_values.shape[0] |
| 1495 | free, total = torch.accelerator.get_memory_info() |
| 1496 | max_batch_size = min( |
| 1497 | total_frames, |
| 1498 | self._encoder_chunk( |
| 1499 | pixel_values.shape[1], |
| 1500 | free, |
| 1501 | total, |
| 1502 | vision_cfg.position_embedding_size, |
| 1503 | ), |
| 1504 | ) |
| 1505 | |
| 1506 | padding_positions = (pixel_position_ids == -1).all(dim=-1) |
| 1507 | |
| 1508 | # Encode frames in chunks bounded by _encoder_chunk. |
| 1509 | last_hidden_states_list: list[torch.Tensor] = [] |
| 1510 | for i in range(0, total_frames, max_batch_size): |
| 1511 | pv_chunk = pixel_values[i : i + max_batch_size] |
| 1512 | pp_chunk = pixel_position_ids[i : i + max_batch_size] |
| 1513 | pad_chunk = padding_positions[i : i + max_batch_size] |
| 1514 | |
| 1515 | inputs_embeds = vt.patch_embedder( |
| 1516 | pv_chunk, |
| 1517 | pp_chunk, |
| 1518 | pad_chunk, |
| 1519 | ).to(self.model_dtype) |
| 1520 | # HuggingFace's mask builder probes `padding_mask.all()`. |
| 1521 | with gpu_sync_allowed(): |
| 1522 | encoder_outputs = vt.encoder( |
| 1523 | inputs_embeds=inputs_embeds, |
| 1524 | attention_mask=~pad_chunk, |
| 1525 | pixel_position_ids=pp_chunk, |
| 1526 | ) |
| 1527 | last_hidden_states_list.append(encoder_outputs.last_hidden_state) |
| 1528 | |
| 1529 | last_hidden_states = torch.cat(last_hidden_states_list, dim=0) |
| 1530 | |
| 1531 | # Pool per frame to strip padding and reduce spatial resolution. |
| 1532 | output_length = pixel_values.shape[1] // pooling_k2 |
| 1533 | all_frame_valid_states: list[torch.Tensor] = [] |
| 1534 | frame_valid_lens: list[int] = [] |
| 1535 | |
| 1536 | for i in range(total_frames): |
| 1537 | single_hidden = last_hidden_states[i].unsqueeze(0) |
| 1538 | single_pos_ids = pixel_position_ids[i].unsqueeze(0) |
| 1539 | single_pad_pos = padding_positions[i].unsqueeze(0) |
| 1540 | |
| 1541 | # As above, plus mask indexing that needs the count on the host. |
| 1542 | with gpu_sync_allowed(): |
| 1543 | pooled_states, valid_mask = vt.pooler( |
| 1544 | hidden_states=single_hidden, |
| 1545 | pixel_position_ids=single_pos_ids, |
| 1546 | padding_positions=single_pad_pos, |
| 1547 | output_length=output_length, |
| 1548 | ) |
| 1549 | valid_states = pooled_states[valid_mask] |
| 1550 | |
| 1551 | if getattr(vt.config, "standardize", False): |
| 1552 | valid_states = (valid_states - vt.std_bias) * vt.std_scale |
| 1553 | |
| 1554 | all_frame_valid_states.append(valid_states) |
| 1555 | frame_valid_lens.append(valid_states.shape[0]) |
| 1556 | |
| 1557 | # Project all frames in a single batched call. |
| 1558 | flat_valid_states = torch.cat(all_frame_valid_states, dim=0).to( |
| 1559 | self.model_dtype |
| 1560 | ) |
| 1561 | flat_proj_embs = self.embed_vision( |
| 1562 | inputs_embeds=flat_valid_states.unsqueeze(0) |
| 1563 | ).squeeze(0) |
| 1564 | |
| 1565 | # Regroup into per-video tensors (slicing returns views). |
| 1566 | per_video_embeddings: list[torch.Tensor] = [] |
| 1567 | frame_idx = 0 |
| 1568 | offset = 0 |
| 1569 | for count in fc_list: |
| 1570 | video_tokens = sum(frame_valid_lens[frame_idx : frame_idx + count]) |
| 1571 | per_video_embeddings.append(flat_proj_embs[offset : offset + video_tokens]) |
| 1572 | offset += video_tokens |
| 1573 | frame_idx += count |
| 1574 | |
| 1575 | return per_video_embeddings |
| 1576 | |
| 1577 | # ------------------------------------------------------------------ # |
| 1578 | # Audio processing |
| 1579 | # ------------------------------------------------------------------ # |
| 1580 | |
| 1581 | def _process_audio_input( |
| 1582 | self, |
| 1583 | audio_input: Gemma4AudioInputs, |
| 1584 | ) -> list[torch.Tensor]: |
| 1585 | input_features, input_features_mask = batch_audio_features( |
| 1586 | audio_input["input_features_padded"], |
| 1587 | audio_input["input_features_mask"], |
| 1588 | ) |
| 1589 | |
| 1590 | assert self.audio_tower is not None |
| 1591 | assert self.embed_audio is not None |
| 1592 | # Run audio tower — mask convention: True=valid, False=padding. |
| 1593 | audio_outputs = self.audio_tower(input_features, input_features_mask) |
| 1594 | if isinstance(audio_outputs, tuple): |
| 1595 | audio_encodings, audio_mask = audio_outputs |
| 1596 | else: |
| 1597 | audio_encodings = audio_outputs.last_hidden_state |
| 1598 | audio_mask = audio_outputs.attention_mask |
| 1599 | |
| 1600 | # Project into LM embedding space. |
| 1601 | audio_features = self.embed_audio(inputs_embeds=audio_encodings) |
| 1602 | |
| 1603 | # Strip padding per-batch element: only keep valid (non-padding) |
| 1604 | # tokens. |
| 1605 | # Boolean-mask indexing needs the selected count on the host. |
| 1606 | per_audio = [] |
| 1607 | with gpu_sync_allowed(): |
| 1608 | for enc, mask in zip(audio_features, audio_mask, strict=True): |
| 1609 | per_audio.append(enc[mask]) # [num_real, hidden_size] |
| 1610 | |
| 1611 | return per_audio |
| 1612 | |
| 1613 | # ------------------------------------------------------------------ # |
| 1614 | # MultiModalEmbeddings interface |
| 1615 | # ------------------------------------------------------------------ # |
| 1616 | |
| 1617 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 1618 | mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs) |
| 1619 | multimodal_embeddings: list[torch.Tensor] = [] |
| 1620 | |
| 1621 | for modality, multimodal_input in mm_input_by_modality.items(): |
| 1622 | if multimodal_input is None: |
| 1623 | continue |
| 1624 | if modality == "image": |
| 1625 | assert isinstance(multimodal_input, Gemma4ImageInputs) |
| 1626 | multimodal_embeddings.extend( |
| 1627 | self._process_image_input(multimodal_input) |
| 1628 | ) |
| 1629 | elif modality == "video": |
| 1630 | assert isinstance(multimodal_input, Gemma4VideoInputs) |
| 1631 | multimodal_embeddings.extend( |
| 1632 | self._process_video_input(multimodal_input) |
| 1633 | ) |
| 1634 | elif modality == "audio": |
| 1635 | assert isinstance(multimodal_input, Gemma4AudioInputs) |
| 1636 | multimodal_embeddings.extend( |
| 1637 | self._process_audio_input(multimodal_input) |
| 1638 | ) |
| 1639 | |
| 1640 | return multimodal_embeddings |
| 1641 | |
| 1642 | # ------------------------------------------------------------------ # |
| 1643 | # EncoderCudaGraph protocol methods |
| 1644 | # ------------------------------------------------------------------ # |
| 1645 | |
| 1646 | def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig": |
| 1647 | from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig |
| 1648 | |
| 1649 | def pad_pixel_values(dst: torch.Tensor, src: torch.Tensor) -> None: |
| 1650 | dst.zero_() |
| 1651 | batch_size, num_patches = src.shape[0], src.shape[1] |
| 1652 | dst[:batch_size, :num_patches].copy_(src) |
| 1653 | |
| 1654 | def pad_pixel_position_ids(dst: torch.Tensor, src: torch.Tensor) -> None: |
| 1655 | dst.fill_(-1) |
| 1656 | batch_size, num_patches = src.shape[0], src.shape[1] |
| 1657 | dst[:batch_size, :num_patches].copy_(src) |
| 1658 | |
| 1659 | return EncoderCudaGraphConfig( |
| 1660 | modalities=["image", "video"], |
| 1661 | buffer_keys=[ |
| 1662 | "pixel_values", |
| 1663 | "pixel_position_ids", |
| 1664 | "gather_indices", |
| 1665 | ], |
| 1666 | out_hidden_size=self.config.text_config.hidden_size, |
| 1667 | max_frames_per_video=_VIDEO_MAX_FRAMES, |
| 1668 | padding_logics={ |
| 1669 | "pixel_values": pad_pixel_values, |
| 1670 | "pixel_position_ids": pad_pixel_position_ids, |
| 1671 | }, |
| 1672 | ) |
| 1673 | |
| 1674 | def get_encoder_cudagraph_budget_range( |
| 1675 | self, |
| 1676 | vllm_config: VllmConfig, |
| 1677 | ) -> tuple[int, int]: |
| 1678 | min_budget = _SUPPORTED_SOFT_TOKENS[0] |
| 1679 | max_budget = min( |
| 1680 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 1681 | vllm_config.model_config.max_model_len, |
| 1682 | ) |
| 1683 | return (min_budget, max_budget) |
| 1684 | |
| 1685 | def get_input_modality(self, mm_kwargs: dict[str, Any]) -> str: |
| 1686 | if "pixel_values" in mm_kwargs: |
| 1687 | return "image" |
| 1688 | elif "pixel_values_videos" in mm_kwargs: |
| 1689 | return "video" |
| 1690 | raise ValueError("Unsupported modality in mm_kwargs") |
| 1691 | |
| 1692 | def get_max_frames_per_video(self) -> int: |
| 1693 | return _VIDEO_MAX_FRAMES |
| 1694 | |
| 1695 | def get_encoder_cudagraph_item_specs( |
| 1696 | self, |
| 1697 | mm_kwargs: dict[str, Any], |
| 1698 | ) -> list["EncoderItemSpec"]: |
| 1699 | from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec |
| 1700 | |
| 1701 | vision_cfg = self.vision_tower.config |
| 1702 | pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2 |
| 1703 | |
| 1704 | modality = self.get_input_modality(mm_kwargs) |
| 1705 | if modality == "image": |
| 1706 | pixel_values = mm_kwargs["pixel_values"] |
| 1707 | if isinstance(pixel_values, list): |
| 1708 | return [ |
| 1709 | EncoderItemSpec( |
| 1710 | input_size=pv.shape[0], |
| 1711 | output_tokens=pv.shape[0] // pool_ratio, |
| 1712 | ) |
| 1713 | for pv in pixel_values |
| 1714 | ] |
| 1715 | else: |
| 1716 | return [ |
| 1717 | EncoderItemSpec( |
| 1718 | input_size=pixel_values.shape[1], |
| 1719 | output_tokens=pixel_values.shape[1] // pool_ratio, |
| 1720 | ) |
| 1721 | for _ in range(pixel_values.shape[0]) |
| 1722 | ] |
| 1723 | elif modality == "video": |
| 1724 | pixel_values_videos = mm_kwargs["pixel_values_videos"] |
| 1725 | video_frame_counts = mm_kwargs["video_frame_counts"] |
| 1726 | fc_list = ( |
| 1727 | video_frame_counts.tolist() |
| 1728 | if isinstance(video_frame_counts, torch.Tensor) |
| 1729 | else list(video_frame_counts) |
| 1730 | ) |
| 1731 | np_patches = pixel_values_videos.shape[1] |
| 1732 | return [ |
| 1733 | EncoderItemSpec( |
| 1734 | input_size=fc * np_patches, |
| 1735 | output_tokens=fc * (np_patches // pool_ratio), |
| 1736 | ) |
| 1737 | for fc in fc_list |
| 1738 | ] |
| 1739 | raise ValueError(f"Unknown modality: {modality}") |
| 1740 | |
| 1741 | def select_encoder_cudagraph_items( |
| 1742 | self, |
| 1743 | mm_kwargs: dict[str, Any], |
| 1744 | indices: list[int], |
| 1745 | ) -> dict[str, Any]: |
| 1746 | modality = self.get_input_modality(mm_kwargs) |
| 1747 | if modality == "image": |
| 1748 | pixel_values = mm_kwargs["pixel_values"] |
| 1749 | pixel_position_ids = mm_kwargs["pixel_position_ids"] |
| 1750 | if len(indices) == 0: |
| 1751 | is_pv_list = isinstance(pixel_values, list) |
| 1752 | is_pp_list = isinstance(pixel_position_ids, list) |
| 1753 | return { |
| 1754 | "pixel_values": ([] if is_pv_list else pixel_values[:0]), |
| 1755 | "pixel_position_ids": ( |
| 1756 | [] if is_pp_list else pixel_position_ids[:0] |
| 1757 | ), |
| 1758 | } |
| 1759 | if isinstance(pixel_values, list): |
| 1760 | return { |
| 1761 | "pixel_values": [pixel_values[i] for i in indices], |
| 1762 | "pixel_position_ids": [pixel_position_ids[i] for i in indices], |
| 1763 | } |
| 1764 | return { |
| 1765 | "pixel_values": pixel_values[indices], |
| 1766 | "pixel_position_ids": pixel_position_ids[indices], |
| 1767 | } |
| 1768 | elif modality == "video": |
| 1769 | pixel_values_videos = mm_kwargs["pixel_values_videos"] |
| 1770 | pixel_position_ids_videos = mm_kwargs["pixel_position_ids_videos"] |
| 1771 | video_frame_counts = mm_kwargs["video_frame_counts"] |
| 1772 | |
| 1773 | if len(indices) == 0: |
| 1774 | is_fc_tensor = isinstance(video_frame_counts, torch.Tensor) |
| 1775 | return { |
| 1776 | "pixel_values_videos": pixel_values_videos[:0], |
| 1777 | "pixel_position_ids_videos": pixel_position_ids_videos[:0], |
| 1778 | "video_frame_counts": ( |
| 1779 | video_frame_counts[:0] if is_fc_tensor else [] |
| 1780 | ), |
| 1781 | } |
| 1782 | |
| 1783 | fc_list = ( |
| 1784 | video_frame_counts.tolist() |
| 1785 | if isinstance(video_frame_counts, torch.Tensor) |
| 1786 | else list(video_frame_counts) |
| 1787 | ) |
| 1788 | cum_frames = [0] |
| 1789 | for fc in fc_list: |
| 1790 | cum_frames.append(cum_frames[-1] + fc) |
| 1791 | |
| 1792 | selected_pv = torch.cat( |
| 1793 | [ |
| 1794 | pixel_values_videos[cum_frames[i] : cum_frames[i + 1]] |
| 1795 | for i in indices |
| 1796 | ], |
| 1797 | dim=0, |
| 1798 | ) |
| 1799 | selected_pp = torch.cat( |
| 1800 | [ |
| 1801 | pixel_position_ids_videos[cum_frames[i] : cum_frames[i + 1]] |
| 1802 | for i in indices |
| 1803 | ], |
| 1804 | dim=0, |
| 1805 | ) |
| 1806 | selected_fc = ( |
| 1807 | video_frame_counts[indices] |
| 1808 | if isinstance(video_frame_counts, torch.Tensor) |
| 1809 | else [video_frame_counts[i] for i in indices] |
| 1810 | ) |
| 1811 | return { |
| 1812 | "pixel_values_videos": selected_pv, |
| 1813 | "pixel_position_ids_videos": selected_pp, |
| 1814 | "video_frame_counts": selected_fc, |
| 1815 | } |
| 1816 | raise ValueError(f"Unknown modality: {modality}") |
| 1817 | |
| 1818 | def prepare_encoder_cudagraph_capture_inputs( |
| 1819 | self, |
| 1820 | token_budget: int = 256, |
| 1821 | max_batch_size: int = 4, |
| 1822 | max_frames_per_batch: int = 1, |
| 1823 | device: torch.device | str = "cpu", |
| 1824 | dtype: torch.dtype | None = None, |
| 1825 | path: str = "default", |
| 1826 | axis_keys: tuple[Hashable, ...] | None = None, |
| 1827 | **kwargs: Any, |
| 1828 | ) -> "EncoderCudaGraphCaptureInputs": |
| 1829 | from vllm.v1.worker.encoder_cudagraph_defs import ( |
| 1830 | EncoderCudaGraphCaptureInputs, |
| 1831 | ) |
| 1832 | |
| 1833 | dtype = dtype or torch.float32 |
| 1834 | |
| 1835 | max_size = max(max_batch_size, max_frames_per_batch) |
| 1836 | |
| 1837 | vision_cfg = self.vision_tower.config |
| 1838 | pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2 |
| 1839 | |
| 1840 | # Retrieve the model's actual configured maximum tokens: |
| 1841 | configured_max_tokens = getattr( |
| 1842 | self.config.vision_config, |
| 1843 | "num_soft_tokens", |
| 1844 | _SUPPORTED_SOFT_TOKENS[2], |
| 1845 | ) |
| 1846 | # Dynamically compute the slot capacity per item bounded by both the |
| 1847 | # current graph budget and the user's maximum config: |
| 1848 | per_item_output = min(token_budget, configured_max_tokens) |
| 1849 | # Satisfy k^2 * per_item_output = per_item_patches |
| 1850 | per_item_patches = per_item_output * pool_ratio |
| 1851 | |
| 1852 | patch_size = self.vision_tower.config.patch_size |
| 1853 | num_channels = getattr(self.vision_tower.config, "num_channels", 3) |
| 1854 | patch_pixels = (patch_size**2) * num_channels |
| 1855 | |
| 1856 | dummy_pixel_values = torch.zeros( |
| 1857 | (max_size, per_item_patches, patch_pixels), |
| 1858 | device=device, |
| 1859 | dtype=dtype, |
| 1860 | ) |
| 1861 | dummy_pixel_position_ids = torch.full( |
| 1862 | (max_size, per_item_patches, 2), |
| 1863 | -1, |
| 1864 | device=device, |
| 1865 | dtype=torch.long, |
| 1866 | ) |
| 1867 | dummy_gather_indices = torch.zeros( |
| 1868 | (token_budget,), |
| 1869 | device=device, |
| 1870 | dtype=torch.long, |
| 1871 | ) |
| 1872 | |
| 1873 | return EncoderCudaGraphCaptureInputs( |
| 1874 | values={ |
| 1875 | "pixel_values": dummy_pixel_values, |
| 1876 | "pixel_position_ids": dummy_pixel_position_ids, |
| 1877 | "gather_indices": dummy_gather_indices, |
| 1878 | } |
| 1879 | ) |
| 1880 | |
| 1881 | def prepare_encoder_cudagraph_replay_buffers( |
| 1882 | self, |
| 1883 | mm_kwargs: dict[str, Any], |
| 1884 | max_batch_size: int = 4, |
| 1885 | max_frames_per_batch: int = 1, |
| 1886 | path: str = "default", |
| 1887 | **kwargs: Any, |
| 1888 | ) -> "EncoderCudaGraphReplayBuffers": |
| 1889 | from vllm.v1.worker.encoder_cudagraph_defs import ( |
| 1890 | EncoderCudaGraphReplayBuffers, |
| 1891 | ) |
| 1892 | |
| 1893 | modality = self.get_input_modality(mm_kwargs) |
| 1894 | if modality == "image": |
| 1895 | pixel_values = mm_kwargs["pixel_values"] |
| 1896 | pixel_position_ids = mm_kwargs["pixel_position_ids"] |
| 1897 | elif modality == "video": |
| 1898 | pixel_values = mm_kwargs["pixel_values_videos"] |
| 1899 | pixel_position_ids = mm_kwargs["pixel_position_ids_videos"] |
| 1900 | else: |
| 1901 | raise ValueError(f"Unsupported modality: {modality}") |
| 1902 | |
| 1903 | if isinstance(pixel_values, list): |
| 1904 | max_patches = max(pv.shape[0] for pv in pixel_values) |
| 1905 | batch_size = len(pixel_values) |
| 1906 | pv_tensor = torch.zeros( |
| 1907 | (batch_size, max_patches, pixel_values[0].shape[1]), |
| 1908 | dtype=pixel_values[0].dtype, |
| 1909 | device=pixel_values[0].device, |
| 1910 | ) |
| 1911 | pp_tensor = torch.full( |
| 1912 | (batch_size, max_patches, 2), |
| 1913 | -1, |
| 1914 | dtype=pixel_position_ids[0].dtype, |
| 1915 | device=pixel_position_ids[0].device, |
| 1916 | ) |
| 1917 | for i, (pv, pp) in enumerate(zip(pixel_values, pixel_position_ids)): |
| 1918 | pv_tensor[i, : pv.shape[0]].copy_(pv) |
| 1919 | pp_tensor[i, : pp.shape[0]].copy_(pp) |
| 1920 | pixel_values = pv_tensor |
| 1921 | pixel_position_ids = pp_tensor |
| 1922 | |
| 1923 | item_specs = self.get_encoder_cudagraph_item_specs(mm_kwargs) |
| 1924 | per_item_out_tokens = [spec.output_tokens for spec in item_specs] |
| 1925 | total_tokens = sum(per_item_out_tokens) |
| 1926 | |
| 1927 | device = pixel_values.device |
| 1928 | vision_cfg = self.vision_tower.config |
| 1929 | pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2 |
| 1930 | per_item_output = pixel_values.shape[1] // pool_ratio |
| 1931 | |
| 1932 | # ONLY allocate an array of exact size `total_tokens`. |
| 1933 | # DO NOT pad it. The upstream Graph Manager handles the padding securely. |
| 1934 | gather_indices = torch.zeros((total_tokens,), dtype=torch.long, device=device) |
| 1935 | |
| 1936 | if modality == "image": |
| 1937 | dst_offset = 0 |
| 1938 | for i, n_tok in enumerate(per_item_out_tokens): |
| 1939 | safe_n_tok = min(n_tok, per_item_output) |
| 1940 | src_start = i * per_item_output |
| 1941 | src_end = src_start + safe_n_tok |
| 1942 | gather_indices[dst_offset : dst_offset + safe_n_tok] = torch.arange( |
| 1943 | src_start, src_end, dtype=torch.long, device=device |
| 1944 | ) |
| 1945 | dst_offset += safe_n_tok |
| 1946 | elif modality == "video": |
| 1947 | video_frame_counts = mm_kwargs["video_frame_counts"] |
| 1948 | fc_list = ( |
| 1949 | video_frame_counts.tolist() |
| 1950 | if isinstance(video_frame_counts, torch.Tensor) |
| 1951 | else list(video_frame_counts) |
| 1952 | ) |
| 1953 | vision_cfg = self.vision_tower.config |
| 1954 | pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2 |
| 1955 | np_patches = pixel_values.shape[1] |
| 1956 | frame_output_tokens = np_patches // pool_ratio |
| 1957 | safe_frame_output_tokens = min(frame_output_tokens, per_item_output) |
| 1958 | |
| 1959 | dst_offset = 0 |
| 1960 | frame_idx = 0 |
| 1961 | for fc in fc_list: |
| 1962 | for f in range(fc): |
| 1963 | src_start = (frame_idx + f) * per_item_output |
| 1964 | src_end = src_start + safe_frame_output_tokens |
| 1965 | gather_indices[ |
| 1966 | dst_offset : dst_offset + safe_frame_output_tokens |
| 1967 | ] = torch.arange( |
| 1968 | src_start, src_end, dtype=torch.long, device=device |
| 1969 | ) |
| 1970 | dst_offset += safe_frame_output_tokens |
| 1971 | frame_idx += fc |
| 1972 | |
| 1973 | return EncoderCudaGraphReplayBuffers( |
| 1974 | values={ |
| 1975 | "pixel_values": pixel_values, |
| 1976 | "pixel_position_ids": pixel_position_ids, |
| 1977 | "gather_indices": gather_indices, |
| 1978 | } |
| 1979 | ) |
| 1980 | |
| 1981 | def encoder_cudagraph_forward( |
| 1982 | self, |
| 1983 | inputs: dict[str, torch.Tensor], |
| 1984 | path: str = "default", |
| 1985 | **kwargs: Any, |
| 1986 | ) -> torch.Tensor: |
| 1987 | pixel_values = inputs["pixel_values"] |
| 1988 | pixel_position_ids = inputs["pixel_position_ids"] |
| 1989 | gather_indices = inputs["gather_indices"] |
| 1990 | |
| 1991 | pad_tensor = (pixel_position_ids == -1).all(dim=-1) |
| 1992 | |
| 1993 | vt = self.vision_tower |
| 1994 | inputs_embeds = vt.patch_embedder( |
| 1995 | pixel_values, |
| 1996 | pixel_position_ids, |
| 1997 | pad_tensor, |
| 1998 | ).to(self.model_dtype) |
| 1999 | |
| 2000 | encoder_outputs = vt.encoder( |
| 2001 | inputs_embeds=inputs_embeds, |
| 2002 | attention_mask=~pad_tensor, |
| 2003 | pixel_position_ids=pixel_position_ids, |
| 2004 | ) |
| 2005 | hidden_states = encoder_outputs.last_hidden_state |
| 2006 | |
| 2007 | pool_ratio = getattr(vt.config, "pooling_kernel_size", 2) ** 2 |
| 2008 | per_item_output = pixel_values.shape[1] // pool_ratio |
| 2009 | |
| 2010 | pooled_states, _ = vt.pooler( |
| 2011 | hidden_states=hidden_states, |
| 2012 | pixel_position_ids=pixel_position_ids, |
| 2013 | padding_positions=pad_tensor, |
| 2014 | output_length=per_item_output, |
| 2015 | ) |
| 2016 | |
| 2017 | if getattr(vt.config, "standardize", False): |
| 2018 | pooled_states = (pooled_states - vt.std_bias) * vt.std_scale |
| 2019 | |
| 2020 | flat_pooled = pooled_states.reshape(-1, pooled_states.shape[-1]) |
| 2021 | gathered_states = flat_pooled[gather_indices] |
| 2022 | |
| 2023 | # Cast to the projection layer's dtype to resolve mixed-precision crash |
| 2024 | target_dtype = self.embed_vision.embedding_projection.weight.dtype |
| 2025 | gathered_states = gathered_states.to(target_dtype) |
| 2026 | |
| 2027 | flat_proj_embs = self.embed_vision( |
| 2028 | inputs_embeds=gathered_states.unsqueeze(0) |
| 2029 | ).squeeze(0) |
| 2030 | |
| 2031 | return flat_proj_embs |
| 2032 | |
| 2033 | def encoder_eager_forward( |
| 2034 | self, |
| 2035 | mm_kwargs: dict[str, Any], |
| 2036 | path: str = "default", |
| 2037 | **kwargs: Any, |
| 2038 | ) -> torch.Tensor: |
| 2039 | modality = self.get_input_modality(mm_kwargs) |
| 2040 | if modality == "image": |
| 2041 | image_input = self._parse_and_validate_image_input(**mm_kwargs) |
| 2042 | assert image_input is not None |
| 2043 | embeddings = self._process_image_input(image_input) |
| 2044 | elif modality == "video": |
| 2045 | video_input = self._parse_and_validate_video_input(**mm_kwargs) |
| 2046 | assert video_input is not None |
| 2047 | embeddings = self._process_video_input(video_input) |
| 2048 | else: |
| 2049 | raise ValueError(f"Unsupported modality: {modality}") |
| 2050 | |
| 2051 | return torch.cat(embeddings, dim=0) |
| 2052 | |
| 2053 | def embed_input_ids( |
| 2054 | self, |
| 2055 | input_ids: torch.Tensor, |
| 2056 | multimodal_embeddings: MultiModalEmbeddings | None = None, |
| 2057 | *, |
| 2058 | is_multimodal: torch.Tensor | None = None, |
| 2059 | ) -> torch.Tensor: |
| 2060 | # Cache per-layer embeddings (PLE) for the language model's |
| 2061 | # forward pass. During profiling embed_input_ids is not called, |
| 2062 | # so the pre-allocated zeros are used instead. |
| 2063 | if self.per_layer_embeddings is not None: |
| 2064 | # Mask multimodal tokens (image/audio) to 0 for PLE |
| 2065 | # computation (using token_type_ids == 0 as text_mask). |
| 2066 | # Replicate this: map image token positions to token 0. |
| 2067 | if is_multimodal is not None: |
| 2068 | ple_input_ids = torch.where( |
| 2069 | is_multimodal.to(input_ids.device, non_blocking=True), |
| 2070 | torch.zeros_like(input_ids), |
| 2071 | input_ids, |
| 2072 | ) |
| 2073 | else: |
| 2074 | ple_input_ids = input_ids |
| 2075 | |
| 2076 | per_layer_inputs = self.language_model.model.get_per_layer_inputs( |
| 2077 | ple_input_ids |
| 2078 | ) |
| 2079 | if per_layer_inputs is not None: |
| 2080 | per_layer_inputs = per_layer_inputs.reshape( |
| 2081 | -1, |
| 2082 | self.config.text_config.num_hidden_layers, |
| 2083 | self.config.text_config.hidden_size_per_layer_input, |
| 2084 | ) |
| 2085 | self.per_layer_embeddings[: per_layer_inputs.shape[0]].copy_( |
| 2086 | per_layer_inputs |
| 2087 | ) |
| 2088 | |
| 2089 | if multimodal_embeddings is None or is_multimodal is None: |
| 2090 | return super().embed_input_ids(input_ids) |
| 2091 | |
| 2092 | return super().embed_input_ids( |
| 2093 | input_ids, |
| 2094 | multimodal_embeddings=multimodal_embeddings, |
| 2095 | is_multimodal=is_multimodal, |
| 2096 | ) |
| 2097 | |
| 2098 | # ------------------------------------------------------------------ # |
| 2099 | # Forward |
| 2100 | # ------------------------------------------------------------------ # |
| 2101 | |
| 2102 | def forward( |
| 2103 | self, |
| 2104 | input_ids: torch.Tensor, |
| 2105 | positions: torch.Tensor, |
| 2106 | intermediate_tensors: IntermediateTensors | None = None, |
| 2107 | inputs_embeds: torch.Tensor | None = None, |
| 2108 | **kwargs: object, |
| 2109 | ) -> IntermediateTensors: |
| 2110 | if intermediate_tensors is not None: |
| 2111 | inputs_embeds = None |
| 2112 | |
| 2113 | # Select the pre-cached PLEs for this batch (None when PLE |
| 2114 | # is disabled for variants without PLE). |
| 2115 | per_layer_inputs = ( |
| 2116 | self.per_layer_embeddings[: inputs_embeds.shape[0]] |
| 2117 | if self.per_layer_embeddings is not None and inputs_embeds is not None |
| 2118 | else None |
| 2119 | ) |
| 2120 | |
| 2121 | # Gemma4 bidi: clear mm_prefix_range for full_attention layers. |
| 2122 | # Must run here (outside @support_torch_compile boundary) because |
| 2123 | # _run_decoder_layers is inside a compiled graph where Python |
| 2124 | # side effects are eliminated. |
| 2125 | self._clear_mm_prefix_for_full_attn_layers() |
| 2126 | |
| 2127 | hidden_states = self.language_model.model( |
| 2128 | input_ids, |
| 2129 | positions, |
| 2130 | per_layer_inputs=per_layer_inputs, |
| 2131 | intermediate_tensors=intermediate_tensors, |
| 2132 | inputs_embeds=inputs_embeds, |
| 2133 | **kwargs, |
| 2134 | ) |
| 2135 | |
| 2136 | return hidden_states |
| 2137 | |
| 2138 | def compute_logits( |
| 2139 | self, |
| 2140 | hidden_states: torch.Tensor, |
| 2141 | ) -> torch.Tensor | None: |
| 2142 | logits = self.language_model.compute_logits(hidden_states) |
| 2143 | if logits is not None and self._suppress_token_ids: |
| 2144 | # Cache a per-device index tensor for the (static) suppressed-token |
| 2145 | # set and use `index_fill_`, so neither the Python-list indices nor |
| 2146 | # the scalar fill value take a host roundtrip per call. |
| 2147 | cache = getattr(self, "_suppress_token_ids_cache", None) |
| 2148 | if cache is None: |
| 2149 | cache = {} |
| 2150 | self._suppress_token_ids_cache = cache |
| 2151 | suppress_idx = cache.get(logits.device) |
| 2152 | if suppress_idx is None: |
| 2153 | suppress_idx = async_tensor_h2d( |
| 2154 | self._suppress_token_ids, dtype=torch.long, device=logits.device |
| 2155 | ) |
| 2156 | cache[logits.device] = suppress_idx |
| 2157 | logits.index_fill_(1, suppress_idx, -float("inf")) |
| 2158 | return logits |
| 2159 | |
| 2160 | # ------------------------------------------------------------------ # |
| 2161 | # Bidirectional attention helpers |
| 2162 | # ------------------------------------------------------------------ # |
| 2163 | |
| 2164 | def _clear_mm_prefix_for_full_attn_layers(self) -> None: |
| 2165 | """Clear mm_prefix_range for non-sliding layers. |
| 2166 | |
| 2167 | Gemma4 with use_bidirectional_attention='vision' applies |
| 2168 | bidirectional attention only to sliding_attention layers. |
| 2169 | Full attention layers use plain causal masking. |
| 2170 | |
| 2171 | Uses _full_attn_layer_idxs (precomputed in __init__) for O(1) |
| 2172 | lookup instead of per-call regex parsing. |
| 2173 | """ |
| 2174 | if not self._full_attn_layer_idxs: |
| 2175 | return |
| 2176 | |
| 2177 | from vllm.forward_context import get_forward_context |
| 2178 | |
| 2179 | attn_metadata = get_forward_context().attn_metadata |
| 2180 | if attn_metadata is None: |
| 2181 | return |
| 2182 | |
| 2183 | def _process(metadata_dict: dict) -> None: |
| 2184 | for layer_name, metadata in metadata_dict.items(): |
| 2185 | if ".layers." not in layer_name: |
| 2186 | continue |
| 2187 | try: |
| 2188 | layer_idx = int(layer_name.split(".layers.")[1].split(".")[0]) |
| 2189 | except (ValueError, IndexError): |
| 2190 | continue |
| 2191 | if layer_idx in self._full_attn_layer_idxs: |
| 2192 | if hasattr(metadata, "mm_prefix_range"): |
| 2193 | metadata.mm_prefix_range = None |
| 2194 | if hasattr(metadata, "mm_prefix_range_tensor"): |
| 2195 | metadata.mm_prefix_range_tensor = None |
| 2196 | if hasattr(metadata, "mm_prefix_query_range_tensor"): |
| 2197 | metadata.mm_prefix_query_range_tensor = None |
| 2198 | |
| 2199 | if isinstance(attn_metadata, list): |
| 2200 | for ub_metadata in attn_metadata: |
| 2201 | _process(ub_metadata) |
| 2202 | elif isinstance(attn_metadata, dict): |
| 2203 | _process(attn_metadata) |
| 2204 | |
| 2205 | # ------------------------------------------------------------------ # |
| 2206 | # Weight loading |
| 2207 | # ------------------------------------------------------------------ # |
| 2208 | |
| 2209 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 2210 | # Some checkpoints have vestigial embed_vision.embedding and |
| 2211 | # embed_audio.embedding weights from the Gemma3n architecture |
| 2212 | # that are not used by Gemma4's MultimodalEmbedder (which only |
| 2213 | # has embedding_projection + embedding_post_projection_norm). |
| 2214 | ignore_prefixes = [ |
| 2215 | "embed_vision.embedding.", |
| 2216 | "embed_audio.embedding.", |
| 2217 | ] |
| 2218 | # Models without audio tower should skip audio weights entirely. |
| 2219 | if self.audio_tower is None: |
| 2220 | ignore_prefixes.extend( |
| 2221 | [ |
| 2222 | "audio_tower.", |
| 2223 | "embed_audio.", |
| 2224 | ] |
| 2225 | ) |
| 2226 | loader = AutoWeightsLoader( |
| 2227 | self, |
| 2228 | ignore_unexpected_prefixes=ignore_prefixes, |
| 2229 | ) |
| 2230 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 2231 | |
| 2232 | # ------------------------------------------------------------------ # |
| 2233 | # LoRA / multimodal mapping |
| 2234 | # ------------------------------------------------------------------ # |
| 2235 | |
| 2236 | def get_mm_mapping(self) -> MultiModelKeys: |
| 2237 | """Get the module prefix mapping for multimodal models.""" |
| 2238 | connectors = ["embed_vision"] |
| 2239 | tower_models = ["vision_tower"] |
| 2240 | if self.audio_tower is not None: |
| 2241 | connectors.append("embed_audio") |
| 2242 | tower_models.append("audio_tower") |
| 2243 | |
| 2244 | return MultiModelKeys.from_string_field( |
| 2245 | language_model="language_model", |
| 2246 | connector=connectors, |
| 2247 | tower_model=tower_models, |
| 2248 | ) |
| 2249 | |
| 2250 | def get_mm_lora_token_counts( |
| 2251 | self, |
| 2252 | *, |
| 2253 | modality: str, |
| 2254 | mm_kwargs: MultiModalKwargsItem | None, |
| 2255 | num_mm_embeds: int, |
| 2256 | ) -> tuple[int, int | None]: |
| 2257 | tower_tokens: int | None = None |
| 2258 | connector_tokens: int | None = None |
| 2259 | if modality in ("image", "video"): |
| 2260 | vision_config = self.config.vision_config |
| 2261 | pooling_k2 = vision_config.pooling_kernel_size**2 |
| 2262 | |
| 2263 | if modality == "image": |
| 2264 | pixel_values_key = "pixel_values" |
| 2265 | max_soft_tokens = int(vision_config.default_output_length) |
| 2266 | mm_processor_kwargs = getattr( |
| 2267 | getattr(self, "multimodal_config", None), |
| 2268 | "mm_processor_kwargs", |
| 2269 | None, |
| 2270 | ) |
| 2271 | if isinstance(mm_processor_kwargs, Mapping): |
| 2272 | val, _ = _get_max_soft_tokens(mm_processor_kwargs) |
| 2273 | if isinstance(val, int) and val in _SUPPORTED_SOFT_TOKENS: |
| 2274 | max_soft_tokens = val |
| 2275 | else: |
| 2276 | pixel_values_key = "pixel_values_videos" |
| 2277 | max_soft_tokens = _VIDEO_MAX_SOFT_TOKENS |
| 2278 | |
| 2279 | tower_tokens = max_soft_tokens * pooling_k2 if modality == "image" else None |
| 2280 | connector_tokens = num_mm_embeds |
| 2281 | if tower_tokens is None and mm_kwargs is not None: |
| 2282 | field = mm_kwargs.get(pixel_values_key) |
| 2283 | if field is not None: |
| 2284 | data = field.data |
| 2285 | if isinstance(data, torch.Tensor) and data.ndim >= 2: |
| 2286 | tower_tokens = int(math.prod(data.shape[:-1])) |
| 2287 | |
| 2288 | if tower_tokens is None: |
| 2289 | min_soft_tokens = min(_SUPPORTED_SOFT_TOKENS) |
| 2290 | tower_tokens = ( |
| 2291 | math.ceil(num_mm_embeds / min_soft_tokens) |
| 2292 | * max_soft_tokens |
| 2293 | * pooling_k2 |
| 2294 | ) |
| 2295 | |
| 2296 | elif modality == "audio": |
| 2297 | tower_tokens = num_mm_embeds |
| 2298 | connector_tokens = num_mm_embeds |
| 2299 | |
| 2300 | if mm_kwargs is not None: |
| 2301 | field = mm_kwargs.get("input_features_padded") |
| 2302 | if field is not None: |
| 2303 | data = field.data |
| 2304 | if isinstance(data, torch.Tensor) and data.ndim >= 2: |
| 2305 | batch_size = math.prod(data.shape[:-2]) |
| 2306 | audio_tokens = batch_size * math.ceil(data.shape[-2] / 4) |
| 2307 | tower_tokens = audio_tokens |
| 2308 | connector_tokens = audio_tokens |
| 2309 | |
| 2310 | else: |
| 2311 | raise ValueError(f"Unsupported modality: {modality}") |
| 2312 | |
| 2313 | assert tower_tokens is not None |
| 2314 | return tower_tokens, connector_tokens |
| 2315 | |
| 2316 | @classmethod |
| 2317 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 2318 | if modality == "image": |
| 2319 | return "<image_soft_token>" |
| 2320 | if modality == "audio": |
| 2321 | return "<audio_soft_token>" |
| 2322 | if modality == "video": |
| 2323 | return "<|video|>" |
| 2324 | raise ValueError(f"Unsupported modality: {modality}") |