vLLM/vllm/model_executor/models/gemma4_mm.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Gemma 4 multimodal model (image + audio + video support).

Adds vision tower, audio tower, and multimodal embedders on top of the
text-only Gemma4ForCausalLM.  The vision/audio encoders are loaded via
AutoModel.from_config and run in eager mode while the language model uses
the vLLM-optimized path.

Video support:  Gemma4 does **not** have a native video tower.  Videos are
decomposed into timestamped image frames (up to 32 frames at 70 soft tokens
each) and fed through the same vision tower as regular images.  The
processor inserts ``mm:ss`` timestamps between frames so the model can
reason about temporal order.
"""

import math
from collections.abc import Hashable, Iterable, Mapping, Sequence
from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal

import numpy as np
import torch
from PIL import Image as PILImage
from torch import nn
from transformers import AutoModel, BatchFeature
from transformers.models.gemma4 import (
    Gemma4AudioFeatureExtractor,
    Gemma4Config,
    Gemma4Processor,
    Gemma4VisionConfig,
)
from transformers.models.gemma4.configuration_gemma4 import (
    Gemma4AudioConfig,
    Gemma4TextConfig,
)

from vllm.config import VllmConfig
from vllm.config.model import get_served_model_name
from vllm.config.multimodal import MultiModalDummyOptions, VideoDummyOptions
from vllm.inputs import MultiModalDataDict
from vllm.logger import init_logger
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import ReplicatedLinear
from vllm.model_executor.models.gemma3n_mm import batch_audio_features
from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
from vllm.model_executor.models.module_mapping import MultiModelKeys
from vllm.model_executor.models.transformers.utils import recursive_replace_linear
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.inputs import (
    MultiModalFieldConfig,
    MultiModalKwargsItem,
    MultiModalKwargsItems,
    VideoItem,
)
from vllm.multimodal.parse import (
    AudioProcessorItems,
    ImageProcessorItems,
    MultiModalDataItems,
    MultiModalDataParser,
)
from vllm.multimodal.processing import BaseDummyInputsBuilder
from vllm.multimodal.processing.processor import (
    BaseMultiModalProcessor,
    BaseProcessingInfo,
    PromptReplacement,
    PromptUpdate,
    PromptUpdateDetails,
)
from vllm.sequence import IntermediateTensors
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
from vllm.utils.tensor_schema import TensorSchema, TensorShape
from vllm.utils.torch_utils import async_tensor_h2d

from .interfaces import (
    MultiModalEmbeddings,
    SupportsEagle3,
    SupportsEncoderCudaGraph,
    SupportsLoRA,
    SupportsMultiModal,
    SupportsPP,
    SupportsQuant,
)
from .utils import (
    AutoWeightsLoader,
    WeightsMapper,
    init_vllm_registered_model,
    maybe_prefix,
)

if TYPE_CHECKING:
    from vllm.model_executor.layers.quantization import QuantizationConfig
    from vllm.v1.worker.encoder_cudagraph_defs import (
        EncoderCudaGraphCaptureInputs,
        EncoderCudaGraphConfig,
        EncoderCudaGraphReplayBuffers,
        EncoderItemSpec,
    )

logger = init_logger(__name__)


def _get_tower_quant_config(
    tower_config: Any,
    quant_config: "QuantizationConfig | None",
) -> "QuantizationConfig | None":
    # BNB/torchao/compressed-tensors handle arbitrary sizes; other methods
    # (Marlin, FP8, …) require dimensions divisible by 64.
    # TODO(mgoin): remove this by fixing kernel padding.
    if quant_config and quant_config.get_name() in [
        "bitsandbytes",
        "torchao",
        "compressed-tensors",
    ]:
        return quant_config
    # The audio config has no intermediate_size; treat it as always divisible.
    intermediate_size = getattr(tower_config, "intermediate_size", None)
    quantizable = tower_config.hidden_size % 64 == 0 and (
        intermediate_size is None or intermediate_size % 64 == 0
    )
    return quant_config if quantizable else None


# Video constants — match transformers Gemma4VideoProcessor defaults.
_SUPPORTED_SOFT_TOKENS = (70, 140, 280, 560, 1120)
_VIDEO_MAX_SOFT_TOKENS = 70  # soft tokens per video frame (vs 280 for images)
_VIDEO_MAX_FRAMES = 32  # max sampled frames per video


def _get_max_soft_tokens(
    merged_kwargs: Mapping[str, object],
) -> tuple[object | None, bool]:
    """Return configured image max_soft_tokens and whether it is top-level."""
    val = merged_kwargs.get("max_soft_tokens")
    if val is not None:
        return val, True

    images_kwargs = merged_kwargs.get("images_kwargs")
    if isinstance(images_kwargs, Mapping):
        return images_kwargs.get("max_soft_tokens"), False

    return None, False


# ---------------------------------------------------------------------------
# Input schema
# ---------------------------------------------------------------------------


class Gemma4ImagePixelInputs(TensorSchema):
    """Pre-patchified image inputs from the Gemma4 image processor.

    Dimensions:
        - bn: Batch size * number of images
        - np: Number of patches (max_patches = max_soft_tokens * pooling_kernel_size²)
        - pp: Patch pixels (patch_size² * 3)

    The Gemma4 image processor outputs pixel_values as
    (batch, max_patches, patch_pixels) — already patchified with
    zero-padding for patches beyond the real image content.
    pixel_position_ids provides (x, y) coordinates per patch,
    with (-1, -1) for padding patches.
    """

    type: Literal["pixel_values"] = "pixel_values"
    pixel_values: Annotated[
        torch.Tensor | list[torch.Tensor],
        TensorShape("bn", "np", "pp", dynamic_dims={"np"}),
    ]
    pixel_position_ids: Annotated[
        torch.Tensor | list[torch.Tensor],
        TensorShape("bn", "np", 2, dynamic_dims={"np"}),
    ]


class Gemma4AudioInputs(TensorSchema):
    """Dimensions:
    - bn: Batch size * number of audios
    - s: Sequence length (MEL spectrogram frames)
    - f: Number of features (MEL bins)
    """

    type: Literal["audio"] = "audio"
    input_features_padded: Annotated[
        torch.Tensor, TensorShape("bn", "s", "f", dynamic_dims={"s"})
    ]
    input_features_mask: Annotated[
        torch.Tensor, TensorShape("bn", "s", dynamic_dims={"s"})
    ]


Gemma4ImageInputs = Gemma4ImagePixelInputs


class Gemma4VideoInputs(TensorSchema):
    """Video frame inputs — same tensor format as image inputs.

    Gemma4 has no separate video tower; video frames are processed
    through the vision tower at lower resolution (max_soft_tokens=70).
    """

    type: Literal["pixel_values_videos"] = "pixel_values_videos"
    pixel_values_videos: Annotated[
        torch.Tensor,
        TensorShape("bn", "np", "pp"),
    ]
    pixel_position_ids_videos: Annotated[
        torch.Tensor,
        TensorShape("bn", "np", 2),
    ]
    video_frame_counts: Annotated[torch.Tensor, TensorShape("nv")]


# ---------------------------------------------------------------------------
# Processing info
# ---------------------------------------------------------------------------


class Gemma4ProcessingInfo(BaseProcessingInfo):
    def get_hf_config(self):
        return self.ctx.get_hf_config(Gemma4Config)

    def get_default_tok_params(self):
        """Gemma4's chat template already embeds a literal ``<bos>`` token in
        the rendered text.  If ``add_special_tokens=True`` (the base-class
        default), the tokenizer prepends *another* BOS, producing a
        ``[2, 2, ...]`` double-BOS sequence that the model was not trained on.

        Setting ``add_special_tokens=False`` here prevents the duplicate and
        ensures both ``llm.generate()`` and the chat/completions API behave
        correctly for IT models. For PT models (without chat template), we
        keep the default (True) to ensure BOS is added for raw prompts.
        """
        tokenizer = self.ctx.get_tokenizer()
        has_chat_template = getattr(tokenizer, "chat_template", None) is not None

        params = super().get_default_tok_params()
        if has_chat_template:
            params = params.with_kwargs(add_special_tokens=False)
        return params

    def get_hf_processor(self, **kwargs: object) -> Gemma4Processor:
        return self.ctx.get_hf_processor(
            Gemma4Processor,
            **kwargs,
        )

    def has_audio(self) -> bool:
        return self.get_hf_config().audio_config is not None

    def validate_num_items(self, modality: str, num_items: int) -> None:
        if modality == "audio" and num_items > 0 and not self.has_audio():
            model_config = self.ctx.model_config
            model = get_served_model_name(
                model_config.model, model_config.served_model_name
            )
            raise ValueError(
                f"Audio input was provided but the model "
                f"'{model}' does not have an audio tower. "
                f"Audio inference is only supported for Gemma4 "
                f"models that include an audio_config "
                f"(i.e., models that include an audio_config)."
            )
        super().validate_num_items(modality, num_items)

    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
        limits: dict[str, int | None] = {"image": None}
        if self.has_audio():
            limits["audio"] = None
        limits["video"] = None
        return limits

    def get_mm_max_tokens_per_item(
        self, seq_len: int, mm_counts: Mapping[str, int]
    ) -> Mapping[str, int] | None:
        config = self.get_hf_config()
        # Upper bound: the pooler outputs max_soft_tokens slots per image.
        # After padding is stripped the actual count is ≤ this value, but
        # vLLM needs the max for memory planning.
        tokens_per_image = config.vision_config.default_output_length
        merged_kwargs = self.ctx.get_merged_mm_kwargs({})
        val, _ = _get_max_soft_tokens(merged_kwargs)
        if isinstance(val, int) and val in _SUPPORTED_SOFT_TOKENS:
            tokens_per_image = val
        tokens: dict[str, int] = {"image": tokens_per_image}
        if self.has_audio():
            # Audio max tokens from the processor's audio_seq_length.
            processor = self.get_hf_processor()
            tokens["audio"] = processor.audio_seq_length
        # Video: each frame ≤ 70 soft tokens + boi + eoi + ~6 ts tokens.
        num_frames = _VIDEO_MAX_FRAMES
        mm_config = self.ctx.model_config.get_multimodal_config()
        video_opts = mm_config.limit_per_prompt.get("video")
        if video_opts is not None and video_opts.num_frames is not None:
            num_frames = min(num_frames, video_opts.num_frames)
        tokens["video"] = num_frames * (_VIDEO_MAX_SOFT_TOKENS + 2 + 6)
        return tokens

    def get_data_parser(self) -> MultiModalDataParser:
        kwargs: dict[str, Any] = {"video_needs_metadata": True}
        if self.has_audio():
            processor = self.get_hf_processor()
            kwargs["target_sr"] = processor.feature_extractor.sampling_rate
        return MultiModalDataParser(**kwargs)

    def _compute_num_soft_tokens(
        self,
        image_width: int,
        image_height: int,
        max_soft_tokens: int | None = None,
    ) -> int:
        """Compute the number of soft tokens the vision tower produces
        for an image of the given dimensions, after padding is stripped.

        Args:
            image_width: Width of the image in pixels.
            image_height: Height of the image in pixels.
            max_soft_tokens: Override for the vision config's
                ``default_output_length``.  When *None*, the value from
                the model config is used.

        """
        vision_cfg = self.get_hf_config().vision_config
        patch_size = vision_cfg.patch_size
        pooling_kernel_size = vision_cfg.pooling_kernel_size

        if max_soft_tokens is None:
            max_soft_tokens = vision_cfg.default_output_length

        unit = patch_size * pooling_kernel_size
        max_patches = max_soft_tokens * pooling_kernel_size**2
        num_patches_orig = (image_height / patch_size) * (image_width / patch_size)
        scale = math.sqrt(max_patches / num_patches_orig)
        target_h = max(unit, int(math.floor(image_height * scale / unit)) * unit)
        target_w = max(unit, int(math.floor(image_width * scale / unit)) * unit)
        num_patches = (target_h // patch_size) * (target_w // patch_size)
        # Clamp to ``max_soft_tokens``: extreme aspect ratios (e.g. 3x900)
        # cause the floor() above to round one dim up to ``unit`` while the
        # other scales freely, which over-shoots ``max_patches``. The HF
        # Gemma 4 image processor caps its vision-tower output at
        # ``max_soft_tokens``, so without this clamp the prompt-side
        # placeholder count exceeds the encoder output and
        # ``_merge_multimodal_embeddings`` crashes.
        return min(num_patches // (pooling_kernel_size**2), max_soft_tokens)

    def get_image_repl(
        self,
        *,
        image_width: int,
        image_height: int,
        processor: Gemma4Processor | None,
        max_soft_tokens: int | None = None,
    ) -> PromptUpdateDetails:
        """Return the dynamic image token sequence for this image.

        Computes the exact number of soft tokens the vision tower will
        produce after stripping padding.

        Args:
            image_width: Width of the image in pixels.
            image_height: Height of the image in pixels.
            processor: The Gemma4 processor, when already constructed.
            max_soft_tokens: Override for the default token budget.
                When *None*, falls back to the model config value.

        """
        if processor is None:
            processor = self.get_hf_processor()

        num_soft = self._compute_num_soft_tokens(
            image_width,
            image_height,
            max_soft_tokens=max_soft_tokens,
        )
        config = self.get_hf_config()
        token_ids = (
            [config.boi_token_id]
            + [processor.image_token_id] * num_soft
            + [config.eoi_token_id]
        )
        return PromptUpdateDetails.select_token_id(token_ids, processor.image_token_id)

    @staticmethod
    def _compute_audio_num_tokens(
        num_samples: int, sampling_rate: int, audio_seq_length: int
    ) -> int:
        """Replicate the audio encoder's sequence-length arithmetic.

        Mirrors: mel framing (_unfold in Gemma4AudioFeatureExtractor)
        followed by two Conv2d subsampling layers (kernel=3, stride=2,
        semicausal padding top=1, bottom=1), capped at audio_seq_length.
        """
        frame_length = int(round(sampling_rate * 20.0 / 1000.0))
        hop_length = int(round(sampling_rate * 10.0 / 1000.0))
        frame_size_for_unfold = frame_length + 1
        pad_left = frame_length // 2
        padded_samples = num_samples + pad_left
        num_mel_frames = (padded_samples - frame_size_for_unfold) // hop_length + 1
        if num_mel_frames <= 0:
            return 0
        t = num_mel_frames
        for _ in range(2):
            t = (t + 2 - 3) // 2 + 1
        return min(t, audio_seq_length)

    def get_audio_repl(
        self,
        *,
        audio_len: int,
        processor: Gemma4Processor | None,
    ) -> PromptUpdateDetails:
        """Return the dynamic audio token sequence for this audio.

        Computes the number of soft tokens from the audio waveform
        length by replicating the audio encoder's sequence-length
        arithmetic (mel framing + two Conv2d subsampling layers).
        """
        if processor is None:
            processor = self.get_hf_processor()

        sampling_rate = processor.feature_extractor.sampling_rate
        num_tokens = self._compute_audio_num_tokens(
            audio_len, sampling_rate, processor.audio_seq_length
        )
        config = self.get_hf_config()
        token_ids = (
            [config.boa_token_id]
            + [processor.audio_token_id] * num_tokens
            + [getattr(config, "eoa_token_id", config.eoa_token_index)]
        )
        return PromptUpdateDetails.select_token_id(token_ids, processor.audio_token_id)

    def get_video_repl(
        self,
        *,
        timestamps: list[float],
        num_soft_tokens_per_frame: list[int],
        processor: Gemma4Processor,
    ) -> PromptUpdateDetails:
        """Build the full token replacement for one video.

        Produces the same interleaved sequence as the HF Gemma4Processor:
            mm:ss <boi><|video|>*N<eoi> mm:ss <boi><|video|>*N<eoi> ...
        """
        tokenizer = self.ctx.get_tokenizer()
        config = self.get_hf_config()

        boi_token_id = config.boi_token_id
        eoi_token_id = config.eoi_token_id
        video_token_id = processor.video_token_id

        all_token_ids: list[int] = []
        for i, (ts, n_tokens) in enumerate(zip(timestamps, num_soft_tokens_per_frame)):
            # mm:ss timestamp — matches transformers: int-truncated,
            # zero-padded.
            minutes = int(ts // 60)
            seconds = int(ts % 60)
            ts_str = f"{minutes:02d}:{seconds:02d}"

            prefix = f" {ts_str} " if i > 0 else f"{ts_str} "
            ts_token_ids = tokenizer.encode(prefix, add_special_tokens=False)
            all_token_ids.extend(ts_token_ids)

            all_token_ids.append(boi_token_id)
            all_token_ids.extend([video_token_id] * n_tokens)
            all_token_ids.append(eoi_token_id)

        return PromptUpdateDetails.select_token_id(all_token_ids, video_token_id)


# ---------------------------------------------------------------------------
# Dummy inputs builder
# ---------------------------------------------------------------------------


def _dummy_audio_num_samples(processor: Any) -> int:
    """Length in samples of a worst-case dummy audio item.

    ``fft_length`` is set in ``Gemma4AudioFeatureExtractor.__init__`` and is not
    defined by ``Gemma4UnifiedAudioFeatureExtractor``, which the encoder-free
    Gemma 4 variant uses. Reading it unconditionally therefore raises
    ``AttributeError`` during memory profiling on that variant, which is reached
    whenever audio is the modality with the largest per-item token count.

    For any other extractor, fall back to the audio budget the processor itself
    declares -- the same pair of values used to warn about over-long audio
    elsewhere in this file.
    """
    feature_extractor = processor.feature_extractor
    if isinstance(feature_extractor, Gemma4AudioFeatureExtractor):
        return feature_extractor.fft_length

    sampling_rate = feature_extractor.sampling_rate
    max_duration_s = processor.audio_seq_length * processor.audio_ms_per_token / 1000.0
    return int(max_duration_s * sampling_rate)


class Gemma4DummyInputsBuilder(BaseDummyInputsBuilder[Gemma4ProcessingInfo]):
    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
        num_images = mm_counts.get("image", 0)
        num_audios = mm_counts.get("audio", 0)
        num_videos = mm_counts.get("video", 0)
        processor = self.info.get_hf_processor()
        # Use image_token (<|image|>) with tab prefix — this is what the
        # Gemma4 chat template inserts per image (\t<|image|>).
        # _get_prompt_updates targets image_token and expands it to the
        # full_image_sequence.
        text = ("\t" + processor.image_token) * num_images
        if num_audios > 0 and processor.audio_token:
            text += processor.audio_token * num_audios
        if num_videos > 0:
            text += processor.video_token * num_videos
        return text

    def get_dummy_mm_data(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
        mm_options: MultiModalDummyOptions | None = None,
    ) -> MultiModalDataDict:
        num_audios = mm_counts.get("audio", 0)
        num_videos = mm_counts.get("video", 0)
        processor = self.info.get_hf_processor()
        image_processor = processor.image_processor
        # Use processor's configured image size for dummies.
        # Gemma4ImageProcessor sets size=None (it uses patch_size /
        # max_soft_tokens instead of the standard size dict), so we
        # guard against None with `or {}`.
        size = getattr(image_processor, "size", None) or {}
        img_width = size.get("width", 224)
        img_height = size.get("height", 224)

        data: dict[str, Any] = {
            "image": self._get_dummy_images(
                width=img_width,
                height=img_height,
                num_images=mm_counts.get("image", 0),
                overrides=mm_options.get("image") if mm_options else None,
            ),
        }

        if num_audios > 0:
            audio_len = _dummy_audio_num_samples(processor)
            data["audio"] = self._get_dummy_audios(
                length=audio_len,
                num_audios=num_audios,
                overrides=mm_options.get("audio") if mm_options else None,
            )

        if num_videos > 0:
            data["video"] = self._get_dummy_videos(
                width=img_width,
                height=img_height,
                num_frames=_VIDEO_MAX_FRAMES,
                num_videos=num_videos,
                overrides=mm_options.get("video") if mm_options else None,
            )

        return data

    def _get_dummy_videos(
        self,
        *,
        width: int,
        height: int,
        num_frames: int,
        num_videos: int,
        overrides: VideoDummyOptions | None = None,
    ) -> list[VideoItem]:
        num_frames = max(num_frames, 2)
        videos = super()._get_dummy_videos(
            width=width,
            height=height,
            num_frames=num_frames,
            num_videos=num_videos,
            overrides=overrides,
        )
        videos = [v.copy() for v in videos]

        video_items: list[VideoItem] = []
        for video in videos:
            video_num_frames = video.shape[0]
            video_metadata = {
                "fps": 2.0,
                "duration": video_num_frames / 2.0,
                "total_num_frames": video_num_frames,
                "frames_indices": list(range(video_num_frames)),
                "video_backend": "opencv",
                "do_sample_frames": False,
            }
            video_items.append((video, video_metadata))

        return video_items


# ---------------------------------------------------------------------------
# Multimodal processor
# ---------------------------------------------------------------------------


class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]):
    def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
        return self.dummy_inputs.get_dummy_text(mm_counts)

    def _apply_hf_processor_main(
        self,
        mm_items: MultiModalDataItems,
        hf_kwargs: Mapping[str, object],
    ) -> BatchFeature:
        hf_data, hf_kwargs, passthrough_data = self._get_hf_mm_inputs(
            mm_items, hf_kwargs
        )

        if not hf_data:
            return self._finalize_hf_mm_data(hf_data, hf_kwargs, passthrough_data)

        prompt_text = hf_data.pop("text")
        assert isinstance(prompt_text, str)

        merged_kwargs = self.info.ctx.get_merged_mm_kwargs(hf_kwargs)
        val, is_top_level_max_soft_tokens = _get_max_soft_tokens(merged_kwargs)

        if val is not None and val not in _SUPPORTED_SOFT_TOKENS:
            raise ValueError(
                f"Unsupported max_soft_tokens value: {val}. "
                f"Valid values are {_SUPPORTED_SOFT_TOKENS}."
            )

        # ---- VIDEO HANDLING ----
        # Gemma4 decomposes video into timestamped image frames.
        # Each frame is processed with max_soft_tokens=70 through the
        # same vision tower, matching transformers processing_gemma4.py.
        video_outputs: dict[str, Any] = {}
        if videos := hf_data.pop("videos", []):
            assert isinstance(videos, list)
            processor = self.info.get_hf_processor()

            all_video_pixel_values: list[torch.Tensor] = []
            all_video_position_ids: list[torch.Tensor] = []
            video_num_soft_tokens_per_video: list[list[int]] = []
            video_timestamps_per_video: list[list[float]] = []
            video_frame_counts: list[int] = []

            video_replacements: list[str] = []

            for item in videos:
                video_array, metadata = item

                # Convert frames to PIL images
                if isinstance(video_array, np.ndarray):
                    frames = [
                        PILImage.fromarray(video_array[i])
                        for i in range(video_array.shape[0])
                    ]
                else:
                    frames = list(video_array)

                # Compute timestamps from metadata (same as transformers)
                fps = metadata.get("fps") or 24
                frame_indices = metadata.get("frames_indices", list(range(len(frames))))
                timestamps = [idx / fps for idx in frame_indices]

                # Process frames as images with max_soft_tokens=70
                video_mm_kwargs = dict(hf_kwargs)
                # Override configured image options too: these inputs are video
                # frames, whose cache keys exclude images_kwargs.
                video_mm_kwargs["images_kwargs"] = {}
                video_mm_kwargs["max_soft_tokens"] = _VIDEO_MAX_SOFT_TOKENS

                dummy_prompt = ("\t" + processor.image_token) * len(frames)

                frame_outputs = self.info.ctx.call_hf_processor(
                    self.info.get_hf_processor(**video_mm_kwargs),
                    dict(text=dummy_prompt, **{"images": frames}),
                    video_mm_kwargs,
                )

                # Remap HF key name
                if "image_position_ids" in frame_outputs:
                    frame_outputs["pixel_position_ids"] = frame_outputs.pop(
                        "image_position_ids"
                    )

                all_video_pixel_values.append(frame_outputs["pixel_values"])
                all_video_position_ids.append(frame_outputs["pixel_position_ids"])

                # Compute soft tokens per frame
                num_soft_per_frame = []
                for img in frames:
                    w, h = img.size
                    n = self.info._compute_num_soft_tokens(
                        w, h, max_soft_tokens=_VIDEO_MAX_SOFT_TOKENS
                    )
                    num_soft_per_frame.append(n)

                video_num_soft_tokens_per_video.append(num_soft_per_frame)
                video_timestamps_per_video.append(timestamps)
                video_frame_counts.append(len(frames))

                # Build expanded replacement text for this video.
                ts_strs = [f"{int(s // 60):02d}:{int(s % 60):02d}" for s in timestamps]
                replacement = " ".join(
                    f"{t} {processor.boi_token}"
                    f"{processor.video_token * n}"
                    f"{processor.eoi_token}"
                    for t, n in zip(ts_strs, num_soft_per_frame)
                )
                video_replacements.append(replacement)

            # Replace all <|video|> placeholders at once. We split on
            # video_token to get N+1 parts, then interleave with the
            # N replacement strings. This avoids the iterative
            # split-replace bug where replacement text (which itself
            # contains <|video|> tokens) collides with later splits.
            vt = processor.video_token
            parts = prompt_text.split(vt, len(video_replacements))

            # NOTE: len(parts) <= len(video_replacements) + 1
            parts_with_repl: list[str] = []
            for part, repl in zip(parts, video_replacements):
                parts_with_repl.extend([part, repl])
            parts_with_repl.extend(parts[len(video_replacements) :])

            prompt_text = "".join(parts_with_repl)

            video_outputs = {
                "pixel_values_videos": torch.cat(all_video_pixel_values, dim=0),
                "pixel_position_ids_videos": torch.cat(all_video_position_ids, dim=0),
                "video_frame_counts": torch.tensor(video_frame_counts),
                "video_num_soft_tokens": video_num_soft_tokens_per_video,
                "video_timestamps": video_timestamps_per_video,
            }

        # Warn if any audio waveform exceeds the model's max duration.
        if "audio" in hf_data:
            processor = self.info.get_hf_processor()
            sr = processor.feature_extractor.sampling_rate
            max_tokens = processor.audio_seq_length
            ms_per_tok = processor.audio_ms_per_token
            max_duration_s = max_tokens * ms_per_tok / 1000.0
            audios = hf_data["audio"]
            if not isinstance(audios, (list, tuple)):
                audios = [audios]
            for i, waveform in enumerate(audios):
                duration_s = len(waveform) / sr
                if duration_s > max_duration_s:
                    logger.warning(
                        "Audio duration exceeds max: %f > %f seconds",
                        duration_s,
                        max_duration_s,
                    )
        # vLLM's call_hf_processor (context.py) re-merges
        # mm_processor_kwargs from the model config on every call via:
        #   config_kwargs | incoming_kwargs  (right side wins)
        #
        # If we strip max_soft_tokens from incoming, the re-merge puts
        # back the config's global default (e.g. 280), ignoring any
        # per-prompt override.  Instead, we keep it in the kwargs with
        # the validated per-prompt value so it wins during the merge.
        #
        # NOTE: This requires a corresponding type annotation on the
        # HF side (Gemma4ProcessorKwargs.images_kwargs) so that
        # _merge_kwargs routes max_soft_tokens into images_kwargs.
        patched_mm_kwargs = dict(hf_kwargs)
        if val is not None and is_top_level_max_soft_tokens:
            patched_mm_kwargs["max_soft_tokens"] = val

        processed_data = self.info.ctx.call_hf_processor(
            self.info.get_hf_processor(**patched_mm_kwargs),
            dict(text=prompt_text, **hf_data),
            patched_mm_kwargs,
        )

        # HF uses 'image_position_ids'; vLLM uses 'pixel_position_ids'.
        # Remap here to keep a single translation point.
        if "image_position_ids" in processed_data:
            processed_data["pixel_position_ids"] = processed_data.pop(
                "image_position_ids"
            )

        if "input_features" in processed_data:
            # Unpad per-item so each item's cache entry is
            # self-contained. The batched() field config in
            # _get_mm_fields_config stacks equal-length items and leaves
            # the rest as a list; _process_audio_input re-pads that list
            # to the batch's max length.
            masks = processed_data["input_features_mask"]
            unpadded_features = [
                f[mask]
                for f, mask in zip(
                    processed_data["input_features"],
                    masks,
                )
            ]
            unpadded_masks = [mask[mask] for mask in masks]
            processed_data["input_features"] = unpadded_features
            processed_data["input_features_padded"] = unpadded_features
            processed_data["input_features_mask"] = unpadded_masks

        processed_data.update(video_outputs)
        return self._finalize_hf_mm_data(
            hf_data, hf_kwargs, passthrough_data, processed_data
        )

    def _get_mm_fields_config(
        self,
        hf_inputs: BatchFeature,
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
        fields = dict(
            pixel_values=MultiModalFieldConfig.batched("image"),
            pixel_position_ids=MultiModalFieldConfig.batched("image"),
            input_features_padded=MultiModalFieldConfig.batched("audio"),
            input_features_mask=MultiModalFieldConfig.batched("audio"),
        )

        # Video fields: frames stored flat, split per video by
        # video_frame_counts.
        video_frame_counts = hf_inputs.get("video_frame_counts")
        if video_frame_counts is not None:
            vfc = video_frame_counts
            if not isinstance(vfc, torch.Tensor):
                vfc = torch.tensor(vfc)
            fields.update(
                pixel_values_videos=(
                    MultiModalFieldConfig.flat_from_sizes("video", vfc)
                ),
                pixel_position_ids_videos=(
                    MultiModalFieldConfig.flat_from_sizes("video", vfc)
                ),
                video_frame_counts=MultiModalFieldConfig.batched(
                    "video", keep_on_cpu=True
                ),
                video_num_soft_tokens=MultiModalFieldConfig.batched(
                    "video", keep_on_cpu=True
                ),
                video_timestamps=MultiModalFieldConfig.batched(
                    "video", keep_on_cpu=True
                ),
            )

        return fields

    def _get_prompt_updates(
        self,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, Any],
        out_mm_kwargs: MultiModalKwargsItems,
    ) -> Sequence[PromptUpdate]:
        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

        prompt_updates = []

        if "image" in mm_items:
            # Target image_token (<|image|>) — the single placeholder the
            # Gemma4 chat template inserts once per image in the prompt.
            # vLLM tokenizes the prompt without token expansion, so only
            # one image_token exists per image in the token stream.
            # The replacement expands it to the full image sequence
            # (boi + N×image_token + eoi, where N = max_soft_tokens).
            image_token_id = hf_processor.image_token_id

            def get_replacement_image(item_idx: int):
                images = mm_items.get_items("image", ImageProcessorItems)
                image_size = images.get_image_size(item_idx)
                # Resolve the effective max_soft_tokens by merging
                # per-prompt kwargs with the config-level defaults,
                # consistent with how _apply_hf_processor_main resolves it.
                # Without this merge, a missing per-prompt override
                # would fall back to vision_cfg.default_output_length
                # instead of the config's mm_processor_kwargs default.
                merged_kwargs = self.info.ctx.get_merged_mm_kwargs(
                    hf_processor_mm_kwargs,
                )
                val, _ = _get_max_soft_tokens(merged_kwargs)
                max_soft_tokens = (
                    val
                    if isinstance(val, int) and val in _SUPPORTED_SOFT_TOKENS
                    else None
                )
                return self.info.get_image_repl(
                    image_width=image_size.width,
                    image_height=image_size.height,
                    processor=hf_processor,
                    max_soft_tokens=max_soft_tokens,
                )

            prompt_updates.append(
                PromptReplacement(
                    modality="image",
                    target=[image_token_id],
                    replacement=get_replacement_image,
                )
            )

        if "video" in mm_items:
            video_token_id = hf_processor.video_token_id

            def get_replacement_video(item_idx: int):
                out_item = out_mm_kwargs["video"][item_idx]
                timestamps_data = out_item["video_timestamps"].data
                if isinstance(timestamps_data, torch.Tensor):
                    timestamps_data = timestamps_data.tolist()
                if not isinstance(timestamps_data, list):
                    raise TypeError("video timestamps must be a list")
                timestamps: list[float] = []
                for timestamp in timestamps_data:
                    if not isinstance(timestamp, (int, float)):
                        raise TypeError("each video timestamp must be numeric")
                    timestamps.append(float(timestamp))

                num_soft_data = out_item["video_num_soft_tokens"].data
                if isinstance(num_soft_data, torch.Tensor):
                    num_soft_data = num_soft_data.tolist()
                if not isinstance(num_soft_data, list):
                    raise TypeError("video soft-token counts must be a list")
                num_soft: list[int] = []
                for count in num_soft_data:
                    if not isinstance(count, int):
                        raise TypeError(
                            "each video soft-token count must be an integer"
                        )
                    num_soft.append(count)
                return self.info.get_video_repl(
                    timestamps=timestamps,
                    num_soft_tokens_per_frame=num_soft,
                    processor=hf_processor,
                )

            prompt_updates.append(
                PromptReplacement(
                    modality="video",
                    target=[video_token_id],
                    replacement=get_replacement_video,
                )
            )

        if "audio" in mm_items:
            audio_token_id = hf_processor.audio_token_id

            def get_replacement_audio(item_idx: int):
                audios = mm_items.get_items("audio", AudioProcessorItems)
                audio_len = audios.get_audio_length(item_idx)
                return self.info.get_audio_repl(
                    audio_len=audio_len,
                    processor=hf_processor,
                )

            prompt_updates.append(
                PromptReplacement(
                    modality="audio",
                    target=[audio_token_id],
                    replacement=get_replacement_audio,
                )
            )

        return prompt_updates

    # NOTE: Gemma3/Gemma3n override _apply_token_matches and
    # _find_mm_placeholders to merge adjacent newline tokens that arise
    # when full_image_sequence contains "\n\n" wrappers.  Gemma4's
    # full_image_sequence has NO newlines (just BOI + 280×image_token +
    # EOI), so the base class implementations work correctly as-is.


# ---------------------------------------------------------------------------
# Multimodal embedder
# ---------------------------------------------------------------------------


class Gemma4MultimodalEmbedder(nn.Module):
    """Projects vision/audio soft tokens into LM embedding space.

    Architecture:
        inputs_embeds → embedding_projection → embedding_post_projection_norm

    Unlike Gemma3n which has separate hard/soft embedding paths with
    per-path normalization and a learned embedding table, Gemma4 uses a
    simplified 2-layer design: a linear projection followed by RMSNorm
    (without learnable scale).  The checkpoint confirms this — only
    ``embedding_projection.weight`` exists; there is no embedding table
    or pre-projection norm weights.
    """

    def __init__(
        self,
        multimodal_config: Gemma4VisionConfig | Gemma4AudioConfig,
        text_config: Gemma4TextConfig,
        *,
        quant_config: "QuantizationConfig | None" = None,
        prefix: str = "",
    ):
        super().__init__()

        self.eps = multimodal_config.rms_norm_eps
        self.text_hidden_size = text_config.hidden_size

        # Audio tower uses output_proj_dims (1536) rather than hidden_size
        # (1024); vision uses hidden_size (768) directly.
        embedding_dim = (
            getattr(multimodal_config, "output_proj_dims", None)
            or multimodal_config.hidden_size
        )

        self.embedding_pre_projection_norm = RMSNorm(
            embedding_dim,
            eps=self.eps,
            has_weight=False,
        )

        self.embedding_projection = ReplicatedLinear(
            embedding_dim,
            self.text_hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=maybe_prefix(prefix, "embedding_projection"),
        )

    def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
        """Project soft tokens from a multimodal tower into LM space."""
        embs_normed = self.embedding_pre_projection_norm(inputs_embeds)
        embs_proj, _ = self.embedding_projection(embs_normed)
        return embs_proj


# ---------------------------------------------------------------------------
# Main model
# ---------------------------------------------------------------------------


@MULTIMODAL_REGISTRY.register_processor(
    Gemma4MultiModalProcessor,
    info=Gemma4ProcessingInfo,
    dummy_inputs=Gemma4DummyInputsBuilder,
)
class Gemma4ForConditionalGeneration(
    nn.Module,
    SupportsMultiModal,
    SupportsQuant,
    SupportsPP,
    SupportsLoRA,
    SupportsEagle3,
    SupportsEncoderCudaGraph,
):
    supports_encoder_cudagraph: ClassVar[Literal[True]] = True
    # Gemma4 clamps mm_prefix bidirectional ranges to the sliding window
    # in-kernel (HF's (causal OR blockwise) AND sliding_window). The model
    # runner reads this to keep image bidirectional ranges that exceed the
    # window instead of dropping them (which would make image attention
    # causal-only for images larger than the sliding window).
    mm_prefix_clamp_sliding_window: bool = True
    supports_tower_connector_lora = True

    packed_modules_mapping: dict[str, list[str]] = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

    # Maps checkpoint prefixes to vLLM module paths.
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={
            # vision tower
            "model.vision_tower": "vision_tower",
            "model.embed_vision": "embed_vision",
            # audio tower
            "model.audio_tower.": "audio_tower.",
            "model.embed_audio.": "embed_audio.",
            # backbone
            "model.language_model.": "language_model.model.",
            "lm_head.": "language_model.lm_head.",
            "model": "language_model.model",
        },
    )

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        multimodal_config = vllm_config.model_config.multimodal_config
        self.config = config
        self.quant_config = quant_config
        self.multimodal_config = multimodal_config
        self.model_dtype = vllm_config.model_config.dtype
        self.vllm_config = vllm_config
        lora_config = vllm_config.lora_config
        self._enable_mm_lora = bool(
            lora_config is not None and lora_config.enable_tower_connector_lora
        )

        # Apply the dimension check independently because the vision and audio
        # towers do not necessarily have the same hidden and intermediate sizes.
        vision_tower_quant = _get_tower_quant_config(config.vision_config, quant_config)

        # ---- Vision tower (shared by image and video) ----
        with self._mark_tower_model(vllm_config, {"image", "video"}):
            self.vision_tower = AutoModel.from_config(config=config.vision_config)
            self.embed_vision = Gemma4MultimodalEmbedder(
                config.vision_config,
                config.text_config,
                quant_config=vision_tower_quant,
                prefix=maybe_prefix(prefix, "embed_vision"),
            )
            recursive_replace_linear(
                self.vision_tower,
                vision_tower_quant,
                prefix=maybe_prefix(prefix, "vision_tower"),
            )

        # ---- Audio tower (variants with audio_config) ----
        self.embed_audio: Gemma4MultimodalEmbedder | None
        if config.audio_config is not None:
            audio_tower_quant = _get_tower_quant_config(
                config.audio_config, quant_config
            )
            with self._mark_tower_model(vllm_config, "audio"):
                self.audio_tower = AutoModel.from_config(config=config.audio_config)
                # AutoModel.from_config does NOT call post_init(),
                # which is needed to initialize buffers that are absent
                # from the checkpoint (e.g. inv_timescales for relative
                # position embeddings, softcap, gradient_clipping).
                self.audio_tower.post_init()
                self.embed_audio = Gemma4MultimodalEmbedder(
                    config.audio_config,
                    config.text_config,
                    quant_config=audio_tower_quant,
                    prefix=maybe_prefix(prefix, "embed_audio"),
                )
                recursive_replace_linear(
                    self.audio_tower,
                    audio_tower_quant,
                    prefix=maybe_prefix(prefix, "audio_tower"),
                )
        else:
            self.audio_tower = None
            self.embed_audio = None

        # ---- Language model (vLLM optimised) ----
        with self._mark_language_model(vllm_config):
            self.language_model: Gemma4ForCausalLM = init_vllm_registered_model(
                vllm_config=vllm_config,
                hf_config=config.text_config,
                prefix=maybe_prefix(prefix, "language_model"),
                architectures=["Gemma4ForCausalLM"],
            )

            # Pre-allocate PLE buffer for CUDA graph compatibility.
            # Some variants have hidden_size_per_layer_input=None (no PLE).
            ple_dim = config.text_config.hidden_size_per_layer_input
            if ple_dim is not None and ple_dim > 0:
                embed = self.language_model.model.embed_tokens
                self.per_layer_embeddings = torch.zeros(
                    vllm_config.scheduler_config.max_num_batched_tokens,
                    config.text_config.num_hidden_layers,
                    ple_dim,
                    device=next(embed.parameters()).device,
                    dtype=vllm_config.model_config.dtype,
                )
            else:
                self.per_layer_embeddings = None

        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors
        )

        # --- Precompute full-attention layer indices for bidi clearing ---
        self._full_attn_layer_idxs: frozenset[int] = frozenset()
        text_config = config.text_config
        if getattr(text_config, "use_bidirectional_attention", None) == "vision":
            layer_types = getattr(text_config, "layer_types", None)
            if layer_types:
                self._full_attn_layer_idxs = frozenset(
                    i for i, lt in enumerate(layer_types) if lt != "sliding_attention"
                )

        # --- MixtureOfExperts delegation to language_model ---
        self.moe_layers = self.language_model.moe_layers
        self.num_moe_layers = self.language_model.num_moe_layers
        self.num_logical_experts = self.language_model.num_logical_experts
        self.num_physical_experts = self.language_model.num_physical_experts
        self.num_local_physical_experts = self.language_model.num_local_physical_experts
        self.num_routed_experts = self.language_model.num_routed_experts
        self.num_expert_groups = self.language_model.num_expert_groups
        self.num_shared_experts = self.language_model.num_shared_experts
        self.num_redundant_experts = self.language_model.num_redundant_experts

        gen_cfg = vllm_config.model_config.try_get_generation_config()
        self._suppress_token_ids = gen_cfg.get("suppress_tokens") if gen_cfg else None

    # ------------------------------------------------------------------ #
    # Input parsing
    # ------------------------------------------------------------------ #

    def _parse_and_validate_image_input(
        self, **kwargs: object
    ) -> Gemma4ImageInputs | None:
        pixel_values = kwargs.pop("pixel_values", None)
        pixel_position_ids = kwargs.pop("pixel_position_ids", None)
        image_embeds = kwargs.pop("image_embeds", None)
        assert image_embeds is None, "Gemma4 does not support image_embeds."
        if pixel_values is None:
            return None
        return Gemma4ImagePixelInputs(
            pixel_values=pixel_values,
            pixel_position_ids=pixel_position_ids,
        )

    def _parse_and_validate_audio_input(
        self, **kwargs: object
    ) -> Gemma4AudioInputs | None:
        input_features_padded = kwargs.pop("input_features_padded", None)
        if input_features_padded is None:
            return None
        input_features_mask = kwargs.pop("input_features_mask", None)
        if input_features_mask is None:
            return None
        return Gemma4AudioInputs(
            input_features_padded=input_features_padded,
            input_features_mask=input_features_mask,
        )

    def _parse_and_validate_video_input(
        self, **kwargs: object
    ) -> Gemma4VideoInputs | None:
        pixel_values_videos = kwargs.pop("pixel_values_videos", None)
        pixel_position_ids_videos = kwargs.pop("pixel_position_ids_videos", None)
        video_frame_counts = kwargs.pop("video_frame_counts", None)
        if pixel_values_videos is None:
            return None
        return Gemma4VideoInputs(
            pixel_values_videos=pixel_values_videos,
            pixel_position_ids_videos=pixel_position_ids_videos,
            video_frame_counts=video_frame_counts,
        )

    def _parse_and_validate_multimodal_inputs(
        self, **kwargs: object
    ) -> dict[str, Gemma4ImageInputs | Gemma4AudioInputs | Gemma4VideoInputs | None]:
        mm_input_by_modality: dict[
            str, Gemma4ImageInputs | Gemma4AudioInputs | Gemma4VideoInputs | None
        ] = {}
        for input_key in list(kwargs):
            if (
                input_key in ("pixel_values", "image_embeds")
                and "image" not in mm_input_by_modality
            ):
                mm_input_by_modality["image"] = self._parse_and_validate_image_input(
                    **kwargs
                )
            if (
                input_key == "pixel_values_videos"
                and "video" not in mm_input_by_modality
            ):
                mm_input_by_modality["video"] = self._parse_and_validate_video_input(
                    **kwargs
                )
            if (
                input_key == "input_features_padded"
                and "audio" not in mm_input_by_modality
            ):
                mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(
                    **kwargs
                )
        return mm_input_by_modality

    @staticmethod
    def _encoder_chunk(
        patches_per_item: int,
        free_bytes: int,
        total_bytes: int,
        position_embedding_size: int,
    ) -> int:
        """Max chunk size whose F.one_hot transient fits in the budget.

        The dominant transient inside HF's ``Gemma4VisionPatchEmbedder.
        _position_embeddings`` is
        ``F.one_hot(clamped_positions, num_classes=position_embedding_size)``
        with shape ``(chunk, patches, 2, position_embedding_size)``,
        int64, plus its simultaneous cast to the position embedding
        table dtype. That, not the encoder residual stream, sets peak
        memory.
        """
        if patches_per_item <= 0:
            return 1
        # Half of currently-free, capped at 10% of total so we leave room
        # for the rest of profile_run / the subsequent encoder + pooler.
        budget = min(free_bytes // 2, total_bytes // 10)
        if budget <= 0:
            return 1
        # F.one_hot allocates (chunk, patches, 2, pos_emb_size) int64
        # (the inner 2 is the (x, y) coordinate axis, 8 is sizeof(int64)).
        # Outer 2x covers the int64 buffer and its concurrent bf16 cast
        # plus the matmul output that live alongside it at peak.
        cost = patches_per_item * 4 * position_embedding_size * 8
        return max(1, budget // cost) if cost > 0 else 1

    # ------------------------------------------------------------------ #
    # Image processing
    # ------------------------------------------------------------------ #

    def _process_image_input(
        self,
        image_input: Gemma4ImageInputs,
    ) -> list[torch.Tensor]:
        """Batch-encode images through the vision tower.

        Groups images by patch count (resolution bucket) so each
        encoder call processes a uniform-shape batch with no
        cross-resolution padding. With MM LoRA enabled, all images are
        padded into one batch so the encoder call matches the tower mapping.
        """
        pixel_values = image_input["pixel_values"]
        pixel_position_ids = image_input["pixel_position_ids"]

        vt = self.vision_tower
        vision_cfg = self.config.vision_config
        pooling_k2 = vision_cfg.pooling_kernel_size**2

        # Concurrent requests with different image resolutions may
        # arrive as a list of per-image tensors, while same-resolution
        # batches may arrive as a stacked tensor.
        buckets: dict[int, list[tuple[int, torch.Tensor, torch.Tensor]]] = {}
        total_images = (
            len(pixel_values)
            if isinstance(pixel_values, list)
            else pixel_values.shape[0]
        )
        pool_position_ids = pixel_position_ids

        if self._enable_mm_lora:
            max_soft_tokens = vision_cfg.default_output_length
            mm_processor_kwargs = getattr(
                getattr(self, "multimodal_config", None),
                "mm_processor_kwargs",
                None,
            )
            if isinstance(mm_processor_kwargs, Mapping):
                value, _ = _get_max_soft_tokens(mm_processor_kwargs)
                if isinstance(value, int) and value in _SUPPORTED_SOFT_TOKENS:
                    max_soft_tokens = value

            max_patches = max_soft_tokens * pooling_k2
            padded_position_ids: list[torch.Tensor] = []
            for idx in range(total_images):
                pv = pixel_values[idx]
                pp = pixel_position_ids[idx]
                num_patches = pv.shape[0]
                if num_patches > max_patches:
                    raise ValueError(
                        f"Image {idx} has {num_patches} patches, which exceeds "
                        f"the MM LoRA patch limit of {max_patches}."
                    )

                pad_len = max_patches - num_patches
                pv = torch.cat(
                    (pv, pv.new_zeros((pad_len, *pv.shape[1:]))),
                    dim=0,
                )
                pp = torch.cat(
                    (pp, pp.new_full((pad_len, *pp.shape[1:]), -1)),
                    dim=0,
                )
                buckets.setdefault(max_patches, []).append((idx, pv, pp))
                padded_position_ids.append(pp)
            pool_position_ids = padded_position_ids
        else:
            for idx in range(total_images):
                pv = pixel_values[idx]
                pp = pixel_position_ids[idx]
                buckets.setdefault(pv.shape[0], []).append((idx, pv, pp))

        # Encode each resolution bucket in memory-safe chunks. Re-read
        # free memory per bucket because the previous bucket's encoder
        # pass has already allocated activations we should account for.
        last_hidden_states_map: dict[int, torch.Tensor] = {}
        for patches, items in buckets.items():
            if self._enable_mm_lora:
                max_batch_size = len(items)
            else:
                free, total = torch.accelerator.get_memory_info()
                max_batch_size = min(
                    len(items),
                    self._encoder_chunk(
                        patches, free, total, vision_cfg.position_embedding_size
                    ),
                )

            for chunk_idx in range(0, len(items), max_batch_size):
                chunk_items = items[chunk_idx : chunk_idx + max_batch_size]

                pv_tensor = torch.cat(
                    [item[1].unsqueeze(0) for item in chunk_items], dim=0
                )
                pp_tensor = torch.cat(
                    [item[2].unsqueeze(0) for item in chunk_items], dim=0
                )
                pad_tensor = (pp_tensor == -1).all(dim=-1)

                inputs_embeds = vt.patch_embedder(
                    pv_tensor,
                    pp_tensor,
                    pad_tensor,
                ).to(self.model_dtype)
                # HuggingFace's mask builder probes `padding_mask.all()` to
                # decide whether the mask can be skipped, which syncs.
                with gpu_sync_allowed():
                    encoder_outputs = vt.encoder(
                        inputs_embeds=inputs_embeds,
                        attention_mask=~pad_tensor,
                        pixel_position_ids=pp_tensor,
                    )
                hidden_states = encoder_outputs.last_hidden_state

                for i, (orig_idx, _, _) in enumerate(chunk_items):
                    last_hidden_states_map[orig_idx] = hidden_states[i]

        # Pool per image to strip padding and reduce spatial resolution.
        all_valid_states: list[torch.Tensor] = [None] * total_images  # type: ignore[list-item]
        valid_lens = [0] * total_images

        for orig_idx in range(total_images):
            chunk_hidden = last_hidden_states_map[orig_idx]
            output_length = chunk_hidden.shape[0] // pooling_k2

            single_hidden = chunk_hidden.unsqueeze(0)
            single_pos_ids = pool_position_ids[orig_idx].unsqueeze(0)
            padding_positions = (single_pos_ids == -1).all(dim=-1)

            # The pooler goes through HuggingFace's mask builder, which probes
            # `padding_mask.all()`, and the mask indexing below needs the
            # selected count on the host.
            with gpu_sync_allowed():
                pooled_states, valid_mask = vt.pooler(
                    hidden_states=single_hidden,
                    pixel_position_ids=single_pos_ids,
                    padding_positions=padding_positions,
                    output_length=output_length,
                )
                valid_states = pooled_states[valid_mask]

            if getattr(vt.config, "standardize", False):
                valid_states = (valid_states - vt.std_bias) * vt.std_scale

            all_valid_states[orig_idx] = valid_states
            valid_lens[orig_idx] = valid_states.shape[0]

        # Project all images in a single batched call.
        flat_valid_states = torch.cat(all_valid_states, dim=0).to(self.model_dtype)
        flat_proj_embs = self.embed_vision(
            inputs_embeds=flat_valid_states.unsqueeze(0)
        ).squeeze(0)

        # Split back into per-image tensors (slicing returns views).
        per_image_embeddings: list[torch.Tensor] = []
        offset = 0
        for length in valid_lens:
            per_image_embeddings.append(flat_proj_embs[offset : offset + length])
            offset += length

        return per_image_embeddings

    # ------------------------------------------------------------------ #
    # Video processing (frames through vision tower)
    # ------------------------------------------------------------------ #

    def _process_video_input(
        self,
        video_input: Gemma4VideoInputs,
    ) -> list[torch.Tensor]:
        """Batch-encode video frames through the vision tower.

        Gemma4 has no separate video tower; video frames are images at
        lower resolution (max_soft_tokens=70).  All frames across all
        videos in the batch are encoded together in chunks, then pooled
        and projected in a single batched call.

        Returns one concatenated embedding tensor per video (not per
        frame), matching the flat_from_sizes grouping that vLLM expects
        for embed_multimodal.
        """
        pixel_values = video_input["pixel_values_videos"]
        pixel_position_ids = video_input["pixel_position_ids_videos"]
        frame_counts = video_input["video_frame_counts"]

        vt = self.vision_tower
        vision_cfg = self.config.vision_config
        pooling_k2 = vision_cfg.pooling_kernel_size**2

        if isinstance(frame_counts, torch.Tensor):
            # Per-video frame counts drive the Python-level batching below.
            with gpu_sync_allowed():
                fc_list = frame_counts.tolist()
        else:
            fc_list = list(frame_counts)

        total_frames = pixel_values.shape[0]
        free, total = torch.accelerator.get_memory_info()
        max_batch_size = min(
            total_frames,
            self._encoder_chunk(
                pixel_values.shape[1],
                free,
                total,
                vision_cfg.position_embedding_size,
            ),
        )

        padding_positions = (pixel_position_ids == -1).all(dim=-1)

        # Encode frames in chunks bounded by _encoder_chunk.
        last_hidden_states_list: list[torch.Tensor] = []
        for i in range(0, total_frames, max_batch_size):
            pv_chunk = pixel_values[i : i + max_batch_size]
            pp_chunk = pixel_position_ids[i : i + max_batch_size]
            pad_chunk = padding_positions[i : i + max_batch_size]

            inputs_embeds = vt.patch_embedder(
                pv_chunk,
                pp_chunk,
                pad_chunk,
            ).to(self.model_dtype)
            # HuggingFace's mask builder probes `padding_mask.all()`.
            with gpu_sync_allowed():
                encoder_outputs = vt.encoder(
                    inputs_embeds=inputs_embeds,
                    attention_mask=~pad_chunk,
                    pixel_position_ids=pp_chunk,
                )
            last_hidden_states_list.append(encoder_outputs.last_hidden_state)

        last_hidden_states = torch.cat(last_hidden_states_list, dim=0)

        # Pool per frame to strip padding and reduce spatial resolution.
        output_length = pixel_values.shape[1] // pooling_k2
        all_frame_valid_states: list[torch.Tensor] = []
        frame_valid_lens: list[int] = []

        for i in range(total_frames):
            single_hidden = last_hidden_states[i].unsqueeze(0)
            single_pos_ids = pixel_position_ids[i].unsqueeze(0)
            single_pad_pos = padding_positions[i].unsqueeze(0)

            # As above, plus mask indexing that needs the count on the host.
            with gpu_sync_allowed():
                pooled_states, valid_mask = vt.pooler(
                    hidden_states=single_hidden,
                    pixel_position_ids=single_pos_ids,
                    padding_positions=single_pad_pos,
                    output_length=output_length,
                )
                valid_states = pooled_states[valid_mask]

            if getattr(vt.config, "standardize", False):
                valid_states = (valid_states - vt.std_bias) * vt.std_scale

            all_frame_valid_states.append(valid_states)
            frame_valid_lens.append(valid_states.shape[0])

        # Project all frames in a single batched call.
        flat_valid_states = torch.cat(all_frame_valid_states, dim=0).to(
            self.model_dtype
        )
        flat_proj_embs = self.embed_vision(
            inputs_embeds=flat_valid_states.unsqueeze(0)
        ).squeeze(0)

        # Regroup into per-video tensors (slicing returns views).
        per_video_embeddings: list[torch.Tensor] = []
        frame_idx = 0
        offset = 0
        for count in fc_list:
            video_tokens = sum(frame_valid_lens[frame_idx : frame_idx + count])
            per_video_embeddings.append(flat_proj_embs[offset : offset + video_tokens])
            offset += video_tokens
            frame_idx += count

        return per_video_embeddings

    # ------------------------------------------------------------------ #
    # Audio processing
    # ------------------------------------------------------------------ #

    def _process_audio_input(
        self,
        audio_input: Gemma4AudioInputs,
    ) -> list[torch.Tensor]:
        input_features, input_features_mask = batch_audio_features(
            audio_input["input_features_padded"],
            audio_input["input_features_mask"],
        )

        assert self.audio_tower is not None
        assert self.embed_audio is not None
        # Run audio tower — mask convention: True=valid, False=padding.
        audio_outputs = self.audio_tower(input_features, input_features_mask)
        if isinstance(audio_outputs, tuple):
            audio_encodings, audio_mask = audio_outputs
        else:
            audio_encodings = audio_outputs.last_hidden_state
            audio_mask = audio_outputs.attention_mask

        # Project into LM embedding space.
        audio_features = self.embed_audio(inputs_embeds=audio_encodings)

        # Strip padding per-batch element: only keep valid (non-padding)
        # tokens.
        # Boolean-mask indexing needs the selected count on the host.
        per_audio = []
        with gpu_sync_allowed():
            for enc, mask in zip(audio_features, audio_mask, strict=True):
                per_audio.append(enc[mask])  # [num_real, hidden_size]

        return per_audio

    # ------------------------------------------------------------------ #
    # MultiModalEmbeddings interface
    # ------------------------------------------------------------------ #

    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
        mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
        multimodal_embeddings: list[torch.Tensor] = []

        for modality, multimodal_input in mm_input_by_modality.items():
            if multimodal_input is None:
                continue
            if modality == "image":
                assert isinstance(multimodal_input, Gemma4ImageInputs)
                multimodal_embeddings.extend(
                    self._process_image_input(multimodal_input)
                )
            elif modality == "video":
                assert isinstance(multimodal_input, Gemma4VideoInputs)
                multimodal_embeddings.extend(
                    self._process_video_input(multimodal_input)
                )
            elif modality == "audio":
                assert isinstance(multimodal_input, Gemma4AudioInputs)
                multimodal_embeddings.extend(
                    self._process_audio_input(multimodal_input)
                )

        return multimodal_embeddings

    # ------------------------------------------------------------------ #
    # EncoderCudaGraph protocol methods
    # ------------------------------------------------------------------ #

    def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig":
        from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig

        def pad_pixel_values(dst: torch.Tensor, src: torch.Tensor) -> None:
            dst.zero_()
            batch_size, num_patches = src.shape[0], src.shape[1]
            dst[:batch_size, :num_patches].copy_(src)

        def pad_pixel_position_ids(dst: torch.Tensor, src: torch.Tensor) -> None:
            dst.fill_(-1)
            batch_size, num_patches = src.shape[0], src.shape[1]
            dst[:batch_size, :num_patches].copy_(src)

        configured_max_tokens = self.config.vision_config.default_output_length
        # One graph per row width (soft tokens of one image or video frame),
        # so a budget's rows are not all padded to the widest image.
        row_widths = tuple(
            t for t in _SUPPORTED_SOFT_TOKENS if t <= configured_max_tokens
        )

        return EncoderCudaGraphConfig(
            modalities=["image", "video"],
            buffer_keys=[
                "pixel_values",
                "pixel_position_ids",
                "gather_indices",
            ],
            out_hidden_size=self.config.text_config.hidden_size,
            max_frames_per_video=_VIDEO_MAX_FRAMES,
            padding_logics={
                "pixel_values": pad_pixel_values,
                "pixel_position_ids": pad_pixel_position_ids,
            },
            capture_axes=(row_widths,),
        )

    def get_encoder_cudagraph_budget_range(
        self,
        vllm_config: VllmConfig,
    ) -> tuple[int, int]:
        min_budget = _SUPPORTED_SOFT_TOKENS[0]
        max_budget = min(
            vllm_config.scheduler_config.max_num_batched_tokens,
            vllm_config.model_config.max_model_len,
        )
        return (min_budget, max_budget)

    def get_input_modality(self, mm_kwargs: dict[str, Any]) -> str:
        if "pixel_values" in mm_kwargs:
            return "image"
        elif "pixel_values_videos" in mm_kwargs:
            return "video"
        raise ValueError("Unsupported modality in mm_kwargs")

    def get_max_frames_per_video(self) -> int:
        return _VIDEO_MAX_FRAMES

    def get_encoder_cudagraph_item_specs(
        self,
        mm_kwargs: dict[str, Any],
    ) -> list["EncoderItemSpec"]:
        from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec

        vision_cfg = self.vision_tower.config
        pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2

        modality = self.get_input_modality(mm_kwargs)
        if modality == "image":
            pixel_values = mm_kwargs["pixel_values"]
            if isinstance(pixel_values, list):
                # A batch pads its images to the widest one, so pack by that.
                row_tokens = (
                    max((pv.shape[0] for pv in pixel_values), default=0) // pool_ratio
                )
                return [
                    EncoderItemSpec(
                        input_size=pv.shape[0],
                        output_tokens=pv.shape[0] // pool_ratio,
                        path_output_tokens={"default": row_tokens},
                    )
                    for pv in pixel_values
                ]
            else:
                return [
                    EncoderItemSpec(
                        input_size=pixel_values.shape[1],
                        output_tokens=pixel_values.shape[1] // pool_ratio,
                    )
                    for _ in range(pixel_values.shape[0])
                ]
        elif modality == "video":
            pixel_values_videos = mm_kwargs["pixel_values_videos"]
            video_frame_counts = mm_kwargs["video_frame_counts"]
            fc_list = (
                video_frame_counts.tolist()
                if isinstance(video_frame_counts, torch.Tensor)
                else list(video_frame_counts)
            )
            np_patches = pixel_values_videos.shape[1]
            return [
                EncoderItemSpec(
                    input_size=fc * np_patches,
                    output_tokens=fc * (np_patches // pool_ratio),
                )
                for fc in fc_list
            ]
        raise ValueError(f"Unknown modality: {modality}")

    def select_encoder_cudagraph_items(
        self,
        mm_kwargs: dict[str, Any],
        indices: list[int],
    ) -> dict[str, Any]:
        from vllm.v1.worker.encoder_cudagraph_defs import (
            ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG,
        )

        vision_cfg = self.vision_tower.config
        pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2

        modality = self.get_input_modality(mm_kwargs)
        if modality == "image":
            pixel_values = mm_kwargs["pixel_values"]
            pixel_position_ids = mm_kwargs["pixel_position_ids"]
            if isinstance(pixel_values, list):
                row_patches = max(
                    (pixel_values[i].shape[0] for i in indices), default=0
                )
            else:
                row_patches = pixel_values.shape[1]
            axis_keys = (row_patches // pool_ratio,)

            if len(indices) == 0:
                is_pv_list = isinstance(pixel_values, list)
                is_pp_list = isinstance(pixel_position_ids, list)
                return {
                    "pixel_values": ([] if is_pv_list else pixel_values[:0]),
                    "pixel_position_ids": (
                        [] if is_pp_list else pixel_position_ids[:0]
                    ),
                    ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG: axis_keys,
                }
            if isinstance(pixel_values, list):
                return {
                    "pixel_values": [pixel_values[i] for i in indices],
                    "pixel_position_ids": [pixel_position_ids[i] for i in indices],
                    ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG: axis_keys,
                }
            return {
                "pixel_values": pixel_values[indices],
                "pixel_position_ids": pixel_position_ids[indices],
                ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG: axis_keys,
            }
        elif modality == "video":
            pixel_values_videos = mm_kwargs["pixel_values_videos"]
            pixel_position_ids_videos = mm_kwargs["pixel_position_ids_videos"]
            video_frame_counts = mm_kwargs["video_frame_counts"]
            axis_keys = (pixel_values_videos.shape[1] // pool_ratio,)

            if len(indices) == 0:
                is_fc_tensor = isinstance(video_frame_counts, torch.Tensor)
                return {
                    "pixel_values_videos": pixel_values_videos[:0],
                    "pixel_position_ids_videos": pixel_position_ids_videos[:0],
                    "video_frame_counts": (
                        video_frame_counts[:0] if is_fc_tensor else []
                    ),
                    ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG: axis_keys,
                }

            fc_list = (
                video_frame_counts.tolist()
                if isinstance(video_frame_counts, torch.Tensor)
                else list(video_frame_counts)
            )
            cum_frames = [0]
            for fc in fc_list:
                cum_frames.append(cum_frames[-1] + fc)

            selected_pv = torch.cat(
                [
                    pixel_values_videos[cum_frames[i] : cum_frames[i + 1]]
                    for i in indices
                ],
                dim=0,
            )
            selected_pp = torch.cat(
                [
                    pixel_position_ids_videos[cum_frames[i] : cum_frames[i + 1]]
                    for i in indices
                ],
                dim=0,
            )
            selected_fc = (
                video_frame_counts[indices]
                if isinstance(video_frame_counts, torch.Tensor)
                else [video_frame_counts[i] for i in indices]
            )
            return {
                "pixel_values_videos": selected_pv,
                "pixel_position_ids_videos": selected_pp,
                "video_frame_counts": selected_fc,
                ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG: axis_keys,
            }
        raise ValueError(f"Unknown modality: {modality}")

    def prepare_encoder_cudagraph_capture_inputs(
        self,
        token_budget: int = 256,
        max_batch_size: int = 4,
        max_frames_per_batch: int = 1,
        device: torch.device | str = "cpu",
        dtype: torch.dtype | None = None,
        path: str = "default",
        axis_keys: tuple[Hashable, ...] | None = None,
        **kwargs: Any,
    ) -> "EncoderCudaGraphCaptureInputs":
        from vllm.v1.worker.encoder_cudagraph_defs import (
            EncoderCudaGraphCaptureInputs,
        )

        dtype = dtype or torch.float32

        assert axis_keys is not None
        (per_item_output,) = axis_keys
        assert isinstance(per_item_output, int)

        # Every row of a batch replayed on this graph is an image or video
        # frame of per_item_output tokens, so a batch that fits token_budget
        # never needs more rows than this. Capture needs at least one row.
        max_size = min(
            max(max_batch_size, max_frames_per_batch),
            max(1, token_budget // per_item_output),
        )

        vision_cfg = self.vision_tower.config
        pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2

        # Satisfy k^2 * per_item_output = per_item_patches
        per_item_patches = per_item_output * pool_ratio

        patch_size = self.vision_tower.config.patch_size
        num_channels = getattr(self.vision_tower.config, "num_channels", 3)
        patch_pixels = (patch_size**2) * num_channels

        dummy_pixel_values = torch.zeros(
            (max_size, per_item_patches, patch_pixels),
            device=device,
            dtype=dtype,
        )
        dummy_pixel_position_ids = torch.full(
            (max_size, per_item_patches, 2),
            -1,
            device=device,
            dtype=torch.long,
        )
        dummy_gather_indices = torch.zeros(
            (token_budget,),
            device=device,
            dtype=torch.long,
        )

        return EncoderCudaGraphCaptureInputs(
            values={
                "pixel_values": dummy_pixel_values,
                "pixel_position_ids": dummy_pixel_position_ids,
                "gather_indices": dummy_gather_indices,
            }
        )

    def prepare_encoder_cudagraph_replay_buffers(
        self,
        mm_kwargs: dict[str, Any],
        max_batch_size: int = 4,
        max_frames_per_batch: int = 1,
        path: str = "default",
        **kwargs: Any,
    ) -> "EncoderCudaGraphReplayBuffers":
        from vllm.v1.worker.encoder_cudagraph_defs import (
            EncoderCudaGraphReplayBuffers,
        )

        modality = self.get_input_modality(mm_kwargs)
        if modality == "image":
            pixel_values = mm_kwargs["pixel_values"]
            pixel_position_ids = mm_kwargs["pixel_position_ids"]
        elif modality == "video":
            pixel_values = mm_kwargs["pixel_values_videos"]
            pixel_position_ids = mm_kwargs["pixel_position_ids_videos"]
        else:
            raise ValueError(f"Unsupported modality: {modality}")

        if isinstance(pixel_values, list):
            max_patches = max(pv.shape[0] for pv in pixel_values)
            batch_size = len(pixel_values)
            pv_tensor = torch.zeros(
                (batch_size, max_patches, pixel_values[0].shape[1]),
                dtype=pixel_values[0].dtype,
                device=pixel_values[0].device,
            )
            pp_tensor = torch.full(
                (batch_size, max_patches, 2),
                -1,
                dtype=pixel_position_ids[0].dtype,
                device=pixel_position_ids[0].device,
            )
            for i, (pv, pp) in enumerate(zip(pixel_values, pixel_position_ids)):
                pv_tensor[i, : pv.shape[0]].copy_(pv)
                pp_tensor[i, : pp.shape[0]].copy_(pp)
            pixel_values = pv_tensor
            pixel_position_ids = pp_tensor

        item_specs = self.get_encoder_cudagraph_item_specs(mm_kwargs)
        per_item_out_tokens = [spec.output_tokens for spec in item_specs]
        total_tokens = sum(per_item_out_tokens)

        device = pixel_values.device
        vision_cfg = self.vision_tower.config
        pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2
        per_item_output = pixel_values.shape[1] // pool_ratio

        # ONLY allocate an array of exact size `total_tokens`.
        # DO NOT pad it. The upstream Graph Manager handles the padding securely.
        gather_indices = torch.zeros((total_tokens,), dtype=torch.long, device=device)

        if modality == "image":
            dst_offset = 0
            for i, n_tok in enumerate(per_item_out_tokens):
                safe_n_tok = min(n_tok, per_item_output)
                src_start = i * per_item_output
                src_end = src_start + safe_n_tok
                gather_indices[dst_offset : dst_offset + safe_n_tok] = torch.arange(
                    src_start, src_end, dtype=torch.long, device=device
                )
                dst_offset += safe_n_tok
        elif modality == "video":
            video_frame_counts = mm_kwargs["video_frame_counts"]
            fc_list = (
                video_frame_counts.tolist()
                if isinstance(video_frame_counts, torch.Tensor)
                else list(video_frame_counts)
            )
            vision_cfg = self.vision_tower.config
            pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2
            np_patches = pixel_values.shape[1]
            frame_output_tokens = np_patches // pool_ratio
            safe_frame_output_tokens = min(frame_output_tokens, per_item_output)

            dst_offset = 0
            frame_idx = 0
            for fc in fc_list:
                for f in range(fc):
                    src_start = (frame_idx + f) * per_item_output
                    src_end = src_start + safe_frame_output_tokens
                    gather_indices[
                        dst_offset : dst_offset + safe_frame_output_tokens
                    ] = torch.arange(
                        src_start, src_end, dtype=torch.long, device=device
                    )
                    dst_offset += safe_frame_output_tokens
                frame_idx += fc

        return EncoderCudaGraphReplayBuffers(
            values={
                "pixel_values": pixel_values,
                "pixel_position_ids": pixel_position_ids,
                "gather_indices": gather_indices,
            }
        )

    def encoder_cudagraph_forward(
        self,
        inputs: dict[str, torch.Tensor],
        path: str = "default",
        **kwargs: Any,
    ) -> torch.Tensor:
        pixel_values = inputs["pixel_values"]
        pixel_position_ids = inputs["pixel_position_ids"]
        gather_indices = inputs["gather_indices"]

        pad_tensor = (pixel_position_ids == -1).all(dim=-1)

        vt = self.vision_tower
        inputs_embeds = vt.patch_embedder(
            pixel_values,
            pixel_position_ids,
            pad_tensor,
        ).to(self.model_dtype)

        encoder_outputs = vt.encoder(
            inputs_embeds=inputs_embeds,
            attention_mask=~pad_tensor,
            pixel_position_ids=pixel_position_ids,
        )
        hidden_states = encoder_outputs.last_hidden_state

        pool_ratio = getattr(vt.config, "pooling_kernel_size", 2) ** 2
        per_item_output = pixel_values.shape[1] // pool_ratio

        pooled_states, _ = vt.pooler(
            hidden_states=hidden_states,
            pixel_position_ids=pixel_position_ids,
            padding_positions=pad_tensor,
            output_length=per_item_output,
        )

        if getattr(vt.config, "standardize", False):
            pooled_states = (pooled_states - vt.std_bias) * vt.std_scale

        flat_pooled = pooled_states.reshape(-1, pooled_states.shape[-1])
        gathered_states = flat_pooled[gather_indices]

        # Cast to the projection layer's dtype to resolve mixed-precision crash
        target_dtype = self.embed_vision.embedding_projection.weight.dtype
        gathered_states = gathered_states.to(target_dtype)

        flat_proj_embs = self.embed_vision(
            inputs_embeds=gathered_states.unsqueeze(0)
        ).squeeze(0)

        return flat_proj_embs

    def encoder_eager_forward(
        self,
        mm_kwargs: dict[str, Any],
        path: str = "default",
        **kwargs: Any,
    ) -> torch.Tensor:
        modality = self.get_input_modality(mm_kwargs)
        if modality == "image":
            image_input = self._parse_and_validate_image_input(**mm_kwargs)
            assert image_input is not None
            embeddings = self._process_image_input(image_input)
        elif modality == "video":
            video_input = self._parse_and_validate_video_input(**mm_kwargs)
            assert video_input is not None
            embeddings = self._process_video_input(video_input)
        else:
            raise ValueError(f"Unsupported modality: {modality}")

        return torch.cat(embeddings, dim=0)

    def embed_input_ids(
        self,
        input_ids: torch.Tensor,
        multimodal_embeddings: MultiModalEmbeddings | None = None,
        *,
        is_multimodal: torch.Tensor | None = None,
    ) -> torch.Tensor:
        # Cache per-layer embeddings (PLE) for the language model's
        # forward pass.  During profiling embed_input_ids is not called,
        # so the pre-allocated zeros are used instead.
        if self.per_layer_embeddings is not None:
            # Mask multimodal tokens (image/audio) to 0 for PLE
            # computation (using token_type_ids == 0 as text_mask).
            # Replicate this: map image token positions to token 0.
            if is_multimodal is not None:
                ple_input_ids = torch.where(
                    is_multimodal.to(input_ids.device, non_blocking=True),
                    torch.zeros_like(input_ids),
                    input_ids,
                )
            else:
                ple_input_ids = input_ids

            per_layer_inputs = self.language_model.model.get_per_layer_inputs(
                ple_input_ids
            )
            if per_layer_inputs is not None:
                per_layer_inputs = per_layer_inputs.reshape(
                    -1,
                    self.config.text_config.num_hidden_layers,
                    self.config.text_config.hidden_size_per_layer_input,
                )
                self.per_layer_embeddings[: per_layer_inputs.shape[0]].copy_(
                    per_layer_inputs
                )

        if multimodal_embeddings is None or is_multimodal is None:
            return super().embed_input_ids(input_ids)

        return super().embed_input_ids(
            input_ids,
            multimodal_embeddings=multimodal_embeddings,
            is_multimodal=is_multimodal,
        )

    # ------------------------------------------------------------------ #
    # Forward
    # ------------------------------------------------------------------ #

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
        **kwargs: object,
    ) -> IntermediateTensors:
        if intermediate_tensors is not None:
            inputs_embeds = None

        # Select the pre-cached PLEs for this batch (None when PLE
        # is disabled for variants without PLE).
        per_layer_inputs = (
            self.per_layer_embeddings[: inputs_embeds.shape[0]]
            if self.per_layer_embeddings is not None and inputs_embeds is not None
            else None
        )

        # Gemma4 bidi: clear mm_prefix_range for full_attention layers.
        # Must run here (outside @support_torch_compile boundary) because
        # _run_decoder_layers is inside a compiled graph where Python
        # side effects are eliminated.
        self._clear_mm_prefix_for_full_attn_layers()

        hidden_states = self.language_model.model(
            input_ids,
            positions,
            per_layer_inputs=per_layer_inputs,
            intermediate_tensors=intermediate_tensors,
            inputs_embeds=inputs_embeds,
            **kwargs,
        )

        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor | None:
        logits = self.language_model.compute_logits(hidden_states)
        if logits is not None and self._suppress_token_ids:
            # Cache a per-device index tensor for the (static) suppressed-token
            # set and use `index_fill_`, so neither the Python-list indices nor
            # the scalar fill value take a host roundtrip per call.
            cache = getattr(self, "_suppress_token_ids_cache", None)
            if cache is None:
                cache = {}
                self._suppress_token_ids_cache = cache
            suppress_idx = cache.get(logits.device)
            if suppress_idx is None:
                suppress_idx = async_tensor_h2d(
                    self._suppress_token_ids, dtype=torch.long, device=logits.device
                )
                cache[logits.device] = suppress_idx
            logits.index_fill_(1, suppress_idx, -float("inf"))
        return logits

    # ------------------------------------------------------------------ #
    # Bidirectional attention helpers
    # ------------------------------------------------------------------ #

    def _clear_mm_prefix_for_full_attn_layers(self) -> None:
        """Clear mm_prefix_range for non-sliding layers.

        Gemma4 with use_bidirectional_attention='vision' applies
        bidirectional attention only to sliding_attention layers.
        Full attention layers use plain causal masking.

        Uses _full_attn_layer_idxs (precomputed in __init__) for O(1)
        lookup instead of per-call regex parsing.
        """
        if not self._full_attn_layer_idxs:
            return

        from vllm.forward_context import get_forward_context

        attn_metadata = get_forward_context().attn_metadata
        if attn_metadata is None:
            return

        def _process(metadata_dict: dict) -> None:
            for layer_name, metadata in metadata_dict.items():
                if ".layers." not in layer_name:
                    continue
                try:
                    layer_idx = int(layer_name.split(".layers.")[1].split(".")[0])
                except (ValueError, IndexError):
                    continue
                if layer_idx in self._full_attn_layer_idxs:
                    if hasattr(metadata, "mm_prefix_range"):
                        metadata.mm_prefix_range = None
                    if hasattr(metadata, "mm_prefix_range_tensor"):
                        metadata.mm_prefix_range_tensor = None
                    if hasattr(metadata, "mm_prefix_query_range_tensor"):
                        metadata.mm_prefix_query_range_tensor = None

        if isinstance(attn_metadata, list):
            for ub_metadata in attn_metadata:
                _process(ub_metadata)
        elif isinstance(attn_metadata, dict):
            _process(attn_metadata)

    # ------------------------------------------------------------------ #
    # Weight loading
    # ------------------------------------------------------------------ #

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        # Some checkpoints have vestigial embed_vision.embedding and
        # embed_audio.embedding weights from the Gemma3n architecture
        # that are not used by Gemma4's MultimodalEmbedder (which only
        # has embedding_projection + embedding_post_projection_norm).
        ignore_prefixes = [
            "embed_vision.embedding.",
            "embed_audio.embedding.",
        ]
        # Models without audio tower should skip audio weights entirely.
        if self.audio_tower is None:
            ignore_prefixes.extend(
                [
                    "audio_tower.",
                    "embed_audio.",
                ]
            )
        loader = AutoWeightsLoader(
            self,
            ignore_unexpected_prefixes=ignore_prefixes,
        )
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

    # ------------------------------------------------------------------ #
    # LoRA / multimodal mapping
    # ------------------------------------------------------------------ #

    def get_mm_mapping(self) -> MultiModelKeys:
        """Get the module prefix mapping for multimodal models."""
        connectors = ["embed_vision"]
        tower_models = ["vision_tower"]
        if self.audio_tower is not None:
            connectors.append("embed_audio")
            tower_models.append("audio_tower")

        return MultiModelKeys.from_string_field(
            language_model="language_model",
            connector=connectors,
            tower_model=tower_models,
        )

    def get_mm_lora_token_counts(
        self,
        *,
        modality: str,
        mm_kwargs: MultiModalKwargsItem | None,
        num_mm_embeds: int,
    ) -> tuple[int, int | None]:
        tower_tokens: int | None = None
        connector_tokens: int | None = None
        if modality in ("image", "video"):
            vision_config = self.config.vision_config
            pooling_k2 = vision_config.pooling_kernel_size**2

            if modality == "image":
                pixel_values_key = "pixel_values"
                max_soft_tokens = int(vision_config.default_output_length)
                mm_processor_kwargs = getattr(
                    getattr(self, "multimodal_config", None),
                    "mm_processor_kwargs",
                    None,
                )
                if isinstance(mm_processor_kwargs, Mapping):
                    val, _ = _get_max_soft_tokens(mm_processor_kwargs)
                    if isinstance(val, int) and val in _SUPPORTED_SOFT_TOKENS:
                        max_soft_tokens = val
            else:
                pixel_values_key = "pixel_values_videos"
                max_soft_tokens = _VIDEO_MAX_SOFT_TOKENS

            tower_tokens = max_soft_tokens * pooling_k2 if modality == "image" else None
            connector_tokens = num_mm_embeds
            if tower_tokens is None and mm_kwargs is not None:
                field = mm_kwargs.get(pixel_values_key)
                if field is not None:
                    data = field.data
                    if isinstance(data, torch.Tensor) and data.ndim >= 2:
                        tower_tokens = int(math.prod(data.shape[:-1]))

            if tower_tokens is None:
                min_soft_tokens = min(_SUPPORTED_SOFT_TOKENS)
                tower_tokens = (
                    math.ceil(num_mm_embeds / min_soft_tokens)
                    * max_soft_tokens
                    * pooling_k2
                )

        elif modality == "audio":
            tower_tokens = num_mm_embeds
            connector_tokens = num_mm_embeds

            if mm_kwargs is not None:
                field = mm_kwargs.get("input_features_padded")
                if field is not None:
                    data = field.data
                    if isinstance(data, torch.Tensor) and data.ndim >= 2:
                        batch_size = math.prod(data.shape[:-2])
                        audio_tokens = batch_size * math.ceil(data.shape[-2] / 4)
                        tower_tokens = audio_tokens
                        connector_tokens = audio_tokens

        else:
            raise ValueError(f"Unsupported modality: {modality}")

        assert tower_tokens is not None
        return tower_tokens, connector_tokens

    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
        if modality == "image":
            return "<image_soft_token>"
        if modality == "audio":
            return "<audio_soft_token>"
        if modality == "video":
            return "<|video|>"
        raise ValueError(f"Unsupported modality: {modality}")