vLLM/vllm/model_executor/models/embedding_gemma2.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""EmbeddingGemma2: bidirectional Gemma4-derived multimodal embedding model."""

import dataclasses
from collections.abc import Iterable, Mapping, Sequence
from typing import Any, cast

import torch
from torch import nn
from transformers import AutoModel, BatchFeature, PreTrainedConfig
from transformers.models.embedding_gemma2.configuration_embedding_gemma2 import (
    EmbeddingGemma2Config,
    EmbeddingGemma2TextConfig,
)
from transformers.models.embedding_gemma2.processing_embedding_gemma2 import (
    EmbeddingGemma2Processor,
)
from transformers.video_utils import VideoMetadata

from vllm.compilation.decorators import support_torch_compile
from vllm.config import VllmConfig
from vllm.distributed import get_tensor_model_parallel_world_size
from vllm.logger import init_logger
from vllm.model_executor.layers.activation import get_act_fn
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    QKVParallelLinear,
    ReplicatedLinear,
    RowParallelLinear,
)
from vllm.model_executor.layers.pooler import DispatchPooler
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.vocab_parallel_embedding import (
    VocabParallelEmbedding,
)
from vllm.model_executor.models.gemma4 import Gemma4MLP
from vllm.model_executor.models.gemma4_mm import (
    _SUPPORTED_SOFT_TOKENS,
    Gemma4DummyInputsBuilder,
    Gemma4ForConditionalGeneration,
    Gemma4MultimodalEmbedder,
    Gemma4MultiModalProcessor,
    Gemma4ProcessingInfo,
    _get_max_soft_tokens,
)
from vllm.model_executor.models.interfaces import SupportsMultiModal
from vllm.model_executor.models.interfaces_base import (
    VllmModelForPooling,
    default_pooling_type,
)
from vllm.model_executor.models.transformers.utils import (
    recursive_replace_linear,
)
from vllm.model_executor.models.utils import (
    AutoWeightsLoader,
    WeightsMapper,
    maybe_prefix,
)
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.inputs import (
    MultiModalKwargsItems,
    VideoItem,
)
from vllm.multimodal.parse import (
    AudioProcessorItems,
    ImageProcessorItems,
    MultiModalDataItems,
)
from vllm.multimodal.processing import (
    BaseProcessingInfo,
    PromptReplacement,
    PromptUpdate,
    PromptUpdateDetails,
)
from vllm.sequence import IntermediateTensors
from vllm.v1.attention.backend import AttentionType

logger = init_logger(__name__)

_WINDOW_OFFSET = 1  # HF masks |q-k| <= W; vLLM ENCODER_ONLY masks <= W-1.
_VIDEO_FRAME_WIDTH = 672
_VIDEO_FRAME_HEIGHT = 480
_VIDEO_MAX_FRAMES = 32
_VIDEO_MAX_SOFT_TOKENS = 140


class EmbeddingGemma2Attention(nn.Module):
    """Bidirectional attention with sliding window offset for EmbeddingGemma2."""

    def __init__(
        self,
        config: EmbeddingGemma2TextConfig,
        layer_cfg,
        layer_idx: int,
        cache_config,
        quant_config: QuantizationConfig | None,
        max_position: int,
        prefix: str,
    ):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        tp = get_tensor_model_parallel_world_size()
        self.num_heads = config.num_attention_heads // tp
        self.num_kv_heads = max(1, layer_cfg.num_key_value_heads // tp)
        self.head_dim = layer_cfg.head_dim
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim

        self.qkv_proj = QKVParallelLinear(
            config.hidden_size,
            self.head_dim,
            config.num_attention_heads,
            layer_cfg.num_key_value_heads,
            bias=config.attention_bias,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv_proj",
        )
        self.o_proj = RowParallelLinear(
            config.num_attention_heads * self.head_dim,
            config.hidden_size,
            bias=config.attention_bias,
            quant_config=quant_config,
            prefix=f"{prefix}.o_proj",
        )
        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.v_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps, has_weight=False)

        assert config.layer_types is not None
        layer_type = config.layer_types[layer_idx]
        self.is_sliding = layer_type == "sliding_attention"
        assert config.rope_parameters is not None
        rope_parameters = dict(config.rope_parameters[layer_type])

        self.rotary_emb = get_rope(
            self.head_dim,
            max_position=max_position,
            rope_parameters=rope_parameters,
            is_neox_style=True,
        )

        window = config.sliding_window if self.is_sliding else None
        self.attn = Attention(
            self.num_heads,
            self.head_dim,
            scale=1.0,
            num_kv_heads=self.num_kv_heads,
            cache_config=cache_config,
            quant_config=quant_config,
            per_layer_sliding_window=(window + _WINDOW_OFFSET)
            if window is not None
            else None,
            attn_type=AttentionType.ENCODER_ONLY,
            prefix=f"{prefix}.attn",
        )

    def forward(
        self, positions: torch.Tensor, hidden_states: torch.Tensor
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

        q = q.unflatten(-1, (self.num_heads, self.head_dim))
        q = self.q_norm(q)
        q = q.flatten(-2, -1)

        k = k.unflatten(-1, (self.num_kv_heads, self.head_dim))
        k = self.k_norm(k)
        k = k.flatten(-2, -1)
        q, k = self.rotary_emb(positions, q, k)

        v = v.unflatten(-1, (self.num_kv_heads, self.head_dim))
        v = self.v_norm(v)
        v = v.flatten(-2, -1)

        attn_output = self.attn(q, k, v)
        out, _ = self.o_proj(attn_output)
        return out


class EmbeddingGemma2PLEBlock(nn.Module):
    """Per-layer-embedding residual block (mirrors HF ple_block)."""

    def __init__(
        self,
        config: EmbeddingGemma2TextConfig,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.hidden_size_per_layer_input = config.hidden_size_per_layer_input

        self.per_layer_input_gate = ReplicatedLinear(
            self.hidden_size,
            self.hidden_size_per_layer_input,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.per_layer_input_gate",
            return_bias=False,
        )
        self.per_layer_projection = ReplicatedLinear(
            self.hidden_size_per_layer_input,
            self.hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.per_layer_projection",
            return_bias=False,
        )
        self.post_per_layer_input_norm = RMSNorm(
            self.hidden_size, eps=config.rms_norm_eps
        )
        self.act_fn = get_act_fn(config.hidden_activation)

    def forward(
        self, hidden_states: torch.Tensor, per_layer_input: torch.Tensor
    ) -> torch.Tensor:
        residual = hidden_states
        gate = self.per_layer_input_gate(hidden_states)
        gate = self.act_fn(gate)
        gated = gate * per_layer_input
        proj = self.per_layer_projection(gated)
        proj = self.post_per_layer_input_norm(proj)
        return residual + proj


class EmbeddingGemma2DecoderLayer(nn.Module):
    """Transformer decoder layer for EmbeddingGemma2 text model with PLE gating."""

    def __init__(
        self,
        config: EmbeddingGemma2TextConfig,
        idx: int,
        cache_config,
        quant_config: QuantizationConfig | None,
        max_position: int,
        prefix: str,
    ):
        super().__init__()
        lc = config.per_layer_config[idx]
        self.self_attn = EmbeddingGemma2Attention(
            config,
            lc,
            idx,
            cache_config,
            quant_config,
            max_position,
            f"{prefix}.self_attn",
        )
        self.mlp = Gemma4MLP(
            hidden_size=config.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_activation=config.hidden_activation,
            quant_config=quant_config,
            prefix=f"{prefix}.mlp",
        )
        self.ple_block = EmbeddingGemma2PLEBlock(
            config,
            quant_config=quant_config,
            prefix=f"{prefix}.ple_block",
        )
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.pre_feedforward_layernorm = RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.post_feedforward_layernorm = RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.layer_scalar: torch.Tensor
        self.register_buffer("layer_scalar", torch.ones(1), persistent=True)

    def forward(
        self,
        positions: torch.Tensor,
        h: torch.Tensor,
        per_layer_input: torch.Tensor,
    ) -> torch.Tensor:
        r = h
        h = (
            self.post_attention_layernorm(
                self.self_attn(positions, self.input_layernorm(h))
            )
            + r
        )
        r = h
        h = (
            self.post_feedforward_layernorm(self.mlp(self.pre_feedforward_layernorm(h)))
            + r
        )
        h = self.ple_block(h, per_layer_input)
        return h * cast(torch.Tensor, self.layer_scalar)


class EmbeddingGemma2TextPLE(nn.Module):
    """Per-layer embedding (PLE) projection and normalization module."""

    def __init__(
        self,
        config: EmbeddingGemma2TextConfig,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
        self.num_hidden_layers = config.num_hidden_layers
        self.hidden_size_per_layer_input = config.hidden_size_per_layer_input
        self.per_layer_model_projection = ColumnParallelLinear(
            config.hidden_size,
            config.num_hidden_layers * config.hidden_size_per_layer_input,
            bias=False,
            gather_output=True,
            quant_config=quant_config,
            prefix=f"{prefix}.per_layer_model_projection",
        )
        self.per_layer_model_projection_scale = config.hidden_size**-0.5
        self.per_layer_projection_norm = RMSNorm(
            config.hidden_size_per_layer_input, eps=config.rms_norm_eps
        )

    def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
        proj, _ = self.per_layer_model_projection(inputs_embeds)
        proj = proj * self.per_layer_model_projection_scale
        proj = proj.reshape(
            *inputs_embeds.shape[:-1],
            self.num_hidden_layers,
            self.hidden_size_per_layer_input,
        )
        return self.per_layer_projection_norm(proj)


@support_torch_compile
class EmbeddingGemma2TextModel(nn.Module):
    """Text embedding backbone for EmbeddingGemma2 with bidirectional attention."""

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        cfg = vllm_config.model_config.hf_config.text_config
        self.config = cfg
        max_pos = min(
            cfg.max_position_embeddings, vllm_config.model_config.max_model_len
        )
        self.embed_tokens = VocabParallelEmbedding(
            cfg.vocab_size,
            cfg.hidden_size,
            quant_config=vllm_config.quant_config,
            prefix=f"{prefix}.embed_tokens",
        )
        self.layers = nn.ModuleList(
            [
                EmbeddingGemma2DecoderLayer(
                    cfg,
                    i,
                    vllm_config.cache_config,
                    vllm_config.quant_config,
                    max_pos,
                    f"{prefix}.layers.{i}",
                )
                for i in range(cfg.num_hidden_layers)
            ]
        )
        self.norm = RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps)
        self.ple = EmbeddingGemma2TextPLE(
            cfg,
            quant_config=vllm_config.quant_config,
            prefix=f"{prefix}.ple",
        )
        self.embedding_projection = ReplicatedLinear(
            cfg.hidden_size,
            cfg.embedding_dim,
            bias=False,
            quant_config=vllm_config.quant_config,
            prefix=f"{prefix}.embedding_projection",
            return_bias=False,
        )
        self.embed_scale: torch.Tensor
        self.register_buffer(
            "embed_scale",
            torch.tensor(cfg.hidden_size**0.5),
            persistent=False,
        )

    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        embed_scale = cast(torch.Tensor, self.embed_scale)
        dtype = cast(torch.dtype, self.embed_tokens.weight.dtype)
        return self.embed_tokens(input_ids) * embed_scale.to(dtype)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if inputs_embeds is None:
            inputs_embeds = self.embed_input_ids(input_ids)
        per_layer = self.ple(inputs_embeds)
        h = inputs_embeds
        for i, layer in enumerate(self.layers):
            h = layer(positions, h, per_layer[..., i, :])
        h = self.norm(h)
        h = self.embedding_projection(h)
        return h


class EmbeddingGemma2ProcessingInfo(Gemma4ProcessingInfo):
    """Processing information and token budgeting for EmbeddingGemma2."""

    def get_hf_config(self) -> EmbeddingGemma2Config:
        return self.ctx.get_hf_config(EmbeddingGemma2Config)

    def get_hf_processor(  # type: ignore[override]
        self, **kwargs: object
    ) -> EmbeddingGemma2Processor:
        return self.ctx.get_hf_processor(EmbeddingGemma2Processor, **kwargs)

    def get_default_tok_params(self):
        # Base-class behavior: add_special_tokens=True (BOS + EOS)
        return BaseProcessingInfo.get_default_tok_params(self)

    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()
        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 + 2
        else:
            tokens_per_image = (
                getattr(config.vision_config, "default_output_length", 256) + 2
            )
        tokens: dict[str, int] = {"image": tokens_per_image}
        if config.audio_config is not None:
            processor = self.get_hf_processor()
            tokens["audio"] = processor.audio_seq_length + 2

        processor = self.get_hf_processor()
        video_proc = getattr(processor, "video_processor", None)
        max_frames = getattr(video_proc, "max_frames", 32)
        mm_config = self.ctx.model_config.get_multimodal_config()

        # The loader's max_frames cap is 128 when media_io_kwargs can raise it
        video_media_io = (
            getattr(mm_config, "media_io_kwargs", {}).get("video", {})
            if mm_config
            else {}
        )
        media_io_max_frames = video_media_io.get("max_frames")
        if media_io_max_frames is not None:
            max_frames = min(int(media_io_max_frames), 128)
        elif "max_frames" in merged_kwargs:
            max_frames = min(int(merged_kwargs["max_frames"]), 128)

        video_opts = mm_config.limit_per_prompt.get("video") if mm_config else None
        if video_opts is not None:
            num_frames_opt = getattr(video_opts, "num_frames", None)
            if isinstance(num_frames_opt, int):
                max_frames = min(max_frames, num_frames_opt)

        video_max_soft = merged_kwargs.get(
            "max_soft_tokens",
            getattr(video_proc, "max_soft_tokens", _VIDEO_MAX_SOFT_TOKENS),
        )
        if video_max_soft is not None and video_max_soft not in _SUPPORTED_SOFT_TOKENS:
            raise ValueError(
                f"Unsupported max_soft_tokens={video_max_soft} for video. "
                f"Valid values are {_SUPPORTED_SOFT_TOKENS}."
            )
        if not isinstance(video_max_soft, int):
            video_max_soft = _VIDEO_MAX_SOFT_TOKENS

        # Token allocation per video frame verified against HF EmbeddingGemma2Processor:
        # - Without timestamps: boi (1) + video_token * num_soft_tokens + eoi (1)
        #   Total = num_soft_tokens + 2
        # - With timestamps: timestamp prefix string (e.g. ' 01:23 ' = up to 7 tokens)
        #   Total = num_soft_tokens + 2 + 7
        add_ts = getattr(video_proc, "add_timestamps", False) or merged_kwargs.get(
            "add_timestamps", False
        )
        ts_slack = 7 if add_ts else 0
        tokens["video"] = max_frames * (video_max_soft + 2 + ts_slack)
        return tokens

    def get_image_repl(  # type: ignore[override]
        self,
        *,
        image_width: int,
        image_height: int,
        processor: EmbeddingGemma2Processor | None,
        max_soft_tokens: int | None = None,
    ) -> PromptUpdateDetails:
        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()
        assert config.boi_token_id is not None
        assert processor.image_token_id is not None
        assert config.eoi_token_id is not None
        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)

    def get_audio_repl(  # type: ignore[override]
        self,
        *,
        audio_len: int,
        processor: EmbeddingGemma2Processor | None,
    ) -> PromptUpdateDetails:
        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()
        assert config.boa_token_id is not None
        assert processor.audio_token_id is not None
        assert config.eoa_token_index is not None
        token_ids = (
            [config.boa_token_id]
            + [processor.audio_token_id] * num_tokens
            + [config.eoa_token_index]
        )
        return PromptUpdateDetails.select_token_id(token_ids, processor.audio_token_id)

    def get_video_repl(  # type: ignore[override]
        self,
        *,
        timestamps: list[float],
        num_soft_tokens_per_frame: list[int],
        processor: EmbeddingGemma2Processor,
        add_timestamps: bool = False,
    ) -> PromptUpdateDetails:
        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
        assert boi_token_id is not None
        assert eoi_token_id is not None
        assert video_token_id is not None

        all_token_ids: list[int] = []
        if not add_timestamps:
            for n_tokens in num_soft_tokens_per_frame:
                all_token_ids.append(boi_token_id)
                all_token_ids.extend([video_token_id] * n_tokens)
                all_token_ids.append(eoi_token_id)
        else:
            for i, (ts, n_tokens) in enumerate(
                zip(timestamps, num_soft_tokens_per_frame)
            ):
                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)


class EmbeddingGemma2DummyInputsBuilder(Gemma4DummyInputsBuilder):
    """Dummy input builder for EmbeddingGemma2 profiling and testing."""

    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()
        text = (
            (processor.image_token * num_images)
            + (processor.audio_token * num_audios if processor.audio_token else "")
            + (processor.video_token * num_videos)
        )
        return text

    def _get_dummy_videos(  # type: ignore[override]
        self,
        *,
        width: int,
        height: int,
        num_frames: int,
        num_videos: int,
        overrides=None,
    ) -> list[VideoItem]:
        num_frames = _VIDEO_MAX_FRAMES
        width = _VIDEO_FRAME_WIDTH
        height = _VIDEO_FRAME_HEIGHT
        videos = super(Gemma4DummyInputsBuilder, self)._get_dummy_videos(
            width=width,
            height=height,
            num_frames=num_frames,
            num_videos=num_videos,
            overrides=overrides,
        )
        video_items: list[VideoItem] = []
        for video in videos:
            video_num_frames = video.shape[0]
            video_metadata = {
                "fps": 1.0,
                "duration": float(video_num_frames),
                "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


class EmbeddingGemma2MultiModalProcessor(Gemma4MultiModalProcessor):
    """Multimodal input processor for EmbeddingGemma2."""

    info: EmbeddingGemma2ProcessingInfo

    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, BatchFeature()
            )

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

        processor = self.info.get_hf_processor()
        video_outputs: dict[str, Any] = {}

        if videos := hf_data.pop("videos", []):
            assert isinstance(videos, list)

            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] = []

            for item in videos:
                video_array, metadata = item

                if hf_kwargs.get("num_frames") not in (-1, None):
                    raise ValueError(
                        "EmbeddingGemma2 does not accept num_frames; "
                        "use max_frames/fps."
                    )

                if isinstance(metadata, dict):
                    if not metadata.get("do_sample_frames", True) and any(
                        k in hf_kwargs
                        for k in ("fps", "max_frames", "overflow_strategy")
                    ):
                        logger.warning_once(
                            "For file/URL-loaded videos, "
                            "fps/max_frames/overflow_strategy in mm_processor_kwargs "
                            "have no effect because frames are decoded and sampled by "
                            "media_io. Please specify these sampling options via "
                            "media_io_kwargs instead."
                        )
                    valid_meta_fields = {
                        f.name for f in dataclasses.fields(VideoMetadata)
                    }
                    meta_dict = {
                        k: v for k, v in metadata.items() if k in valid_meta_fields
                    }
                    metadata_obj = VideoMetadata(**meta_dict)
                else:
                    if getattr(metadata, "do_sample_frames", True) is False and any(
                        k in hf_kwargs
                        for k in ("fps", "max_frames", "overflow_strategy")
                    ):
                        logger.warning_once(
                            "For file/URL-loaded videos, "
                            "fps/max_frames/overflow_strategy in mm_processor_kwargs "
                            "have no effect because frames are decoded and sampled by "
                            "media_io. Please specify these sampling options via "
                            "media_io_kwargs instead."
                        )
                    metadata_obj = metadata

                video_mm_kwargs = dict(hf_kwargs)
                video_mm_kwargs.pop("num_frames", None)
                video_mm_kwargs["return_metadata"] = True
                if isinstance(metadata, dict) and "do_sample_frames" in metadata:
                    video_mm_kwargs["do_sample_frames"] = metadata["do_sample_frames"]
                elif hasattr(metadata, "do_sample_frames"):
                    video_mm_kwargs["do_sample_frames"] = metadata.do_sample_frames

                single_video_out = self.info.ctx.call_hf_processor(
                    self.info.get_hf_processor(**video_mm_kwargs),
                    dict(
                        text=processor.video_token,
                        videos=[video_array],
                        video_metadata=[metadata_obj],
                    ),
                    video_mm_kwargs,
                )

                all_video_pixel_values.append(single_video_out["pixel_values_videos"])
                all_video_position_ids.append(single_video_out["video_position_ids"])

                num_frames = int(single_video_out["num_frames_per_video"][0])
                video_frame_counts.append(num_frames)

                num_soft = single_video_out.get("num_soft_tokens_per_video")
                if num_soft is not None:
                    soft_val = num_soft[0]
                    if isinstance(soft_val, torch.Tensor):
                        soft_val = soft_val.item()
                    soft_tokens_list = [int(soft_val)] * num_frames
                elif "video_position_ids" in single_video_out and num_frames > 0:
                    pos = single_video_out["video_position_ids"]
                    pooling_k2 = (
                        getattr(processor.video_processor, "pooling_kernel_size", 3)
                        ** 2
                    )
                    valid_patches = (pos[0] != -1).all(dim=-1).sum().item()
                    soft_tokens_list = [valid_patches // pooling_k2] * num_frames
                elif "input_ids" in single_video_out and num_frames > 0:
                    vtok = (
                        (
                            single_video_out["input_ids"][0]
                            == processor.tokenizer.video_token_id
                        )
                        .sum()
                        .item()
                    )
                    soft_tokens_list = [vtok // num_frames] * num_frames
                else:
                    raise ValueError(
                        f"Unable to determine soft tokens for video frame: "
                        f"keys={list(single_video_out.keys())}"
                    )

                video_num_soft_tokens_per_video.append(soft_tokens_list)

                # Use returned video_metadata if available, else input metadata_obj
                ret_meta = single_video_out.get("video_metadata")
                meta_for_ts = (
                    ret_meta[0] if (ret_meta and len(ret_meta) > 0) else metadata_obj
                )

                add_ts = hf_kwargs.get(
                    "add_timestamps",
                    getattr(processor.video_processor, "add_timestamps", False),
                )
                if add_ts:
                    try:
                        ts = meta_for_ts.timestamps
                    except ValueError as e:
                        raise ValueError(
                            "Asked to build a prompt with frame timestamps, but no "
                            "valid fps/frames_indices was available in video metadata."
                        ) from e
                    video_timestamps_per_video.append([float(t) for t in ts])
                else:
                    try:
                        ts = meta_for_ts.timestamps
                        video_timestamps_per_video.append([float(t) for t in ts])
                    except ValueError:
                        video_timestamps_per_video.append([0.0] * num_frames)

            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, dtype=torch.long
                ),
                "video_num_soft_tokens": video_num_soft_tokens_per_video,
                "video_timestamps": video_timestamps_per_video,
            }

        if hf_data:
            remaining_counts: dict[str, int] = {}
            if "images" in hf_data:
                images_list = hf_data["images"]
                remaining_counts["image"] = (
                    len(images_list) if isinstance(images_list, list) else 1
                )
            if "audio" in hf_data:
                audio_list = hf_data["audio"]
                remaining_counts["audio"] = (
                    len(audio_list) if isinstance(audio_list, (list, tuple)) else 1
                )
            remaining_prompt_text = self.dummy_inputs.get_dummy_text(remaining_counts)

            call_kwargs = dict(hf_kwargs)
            call_kwargs.pop("num_frames", None)
            processed_data = self.info.ctx.call_hf_processor(
                self.info.get_hf_processor(**call_kwargs),
                dict(text=remaining_prompt_text, **hf_data),
                call_kwargs,
            )
        else:
            processed_data = BatchFeature()

        if "image_position_ids" in processed_data:
            processed_data["pixel_position_ids"] = processed_data.pop(
                "image_position_ids"
            )

        if "input_features" in processed_data:
            input_features = processed_data.pop("input_features")
            input_features_mask = processed_data.pop("input_features_mask")
            if not isinstance(input_features, list):
                unpadded_features = []
                unpadded_masks = []
                for i in range(input_features.shape[0]):
                    mask = input_features_mask[i]
                    valid_len = int(mask.sum().item())
                    unpadded_features.append(input_features[i, :valid_len])
                    unpadded_masks.append(mask[:valid_len])
                processed_data["input_features_padded"] = unpadded_features
                processed_data["input_features_mask"] = unpadded_masks
            else:
                processed_data["input_features_padded"] = input_features
                processed_data["input_features_mask"] = input_features_mask

        processed_data.update(video_outputs)
        processed_data.pop("num_frames_per_video", None)
        processed_data.pop("video_position_ids", None)

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

    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:
            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)
                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
            assert video_token_id is not None

            def get_replacement_video(item_idx: int):
                out_item = out_mm_kwargs["video"][item_idx]
                raw_ts = out_item["video_timestamps"].data
                if isinstance(raw_ts, torch.Tensor):
                    ts_list: list[Any] = raw_ts.tolist()
                elif isinstance(raw_ts, (list, tuple)):
                    ts_list = list(raw_ts)
                else:
                    ts_list = [raw_ts]
                timestamps = [float(t) for t in ts_list]

                raw_soft = out_item["video_num_soft_tokens"].data
                if isinstance(raw_soft, torch.Tensor):
                    soft_list: list[Any] = raw_soft.tolist()
                elif isinstance(raw_soft, (list, tuple)):
                    soft_list = list(raw_soft)
                else:
                    soft_list = [raw_soft]
                num_soft = [int(n) for n in soft_list]

                add_ts = hf_processor_mm_kwargs.get(
                    "add_timestamps",
                    getattr(hf_processor.video_processor, "add_timestamps", False),
                )
                return self.info.get_video_repl(
                    timestamps=timestamps,
                    num_soft_tokens_per_frame=num_soft,
                    processor=hf_processor,
                    add_timestamps=add_ts,
                )

            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
            assert audio_token_id is not None

            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


def _embedding_gemma2_weights_mapper(
    config_or_layers: PreTrainedConfig | int | None = None,
) -> WeightsMapper:
    return WeightsMapper(
        orig_to_new_stacked={
            ".self_attn.q_proj.weight": (".self_attn.qkv_proj.weight", "q"),
            ".self_attn.k_proj.weight": (".self_attn.qkv_proj.weight", "k"),
            ".self_attn.v_proj.weight": (".self_attn.qkv_proj.weight", "v"),
            ".self_attn.q_proj.bias": (".self_attn.qkv_proj.bias", "q"),
            ".self_attn.k_proj.bias": (".self_attn.qkv_proj.bias", "k"),
            ".self_attn.v_proj.bias": (".self_attn.qkv_proj.bias", "v"),
            ".mlp.gate_proj.weight": (".mlp.gate_up_proj.weight", 0),
            ".mlp.up_proj.weight": (".mlp.gate_up_proj.weight", 1),
            ".mlp.gate_proj.bias": (".mlp.gate_up_proj.bias", 0),
            ".mlp.up_proj.bias": (".mlp.gate_up_proj.bias", 1),
        }
    )


@MULTIMODAL_REGISTRY.register_processor(
    EmbeddingGemma2MultiModalProcessor,
    info=EmbeddingGemma2ProcessingInfo,
    dummy_inputs=EmbeddingGemma2DummyInputsBuilder,
)
@default_pooling_type(seq_pooling_type="MEAN", tok_pooling_type="ALL")
class EmbeddingGemma2Model(nn.Module, SupportsMultiModal, VllmModelForPooling):
    """Multimodal pooling embedding model supporting text, vision, and audio."""

    is_pooling_model = True

    hf_to_vllm_mapper = _embedding_gemma2_weights_mapper()
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

    # TODO(lucianomartins): Extract Gemma4EncoderMixin into a shared module
    # once generative model tests are migrated.
    # TODO(lucianomartins): Investigate @support_torch_compile integration for encoder.
    _parse_and_validate_image_input = (
        Gemma4ForConditionalGeneration._parse_and_validate_image_input
    )
    _parse_and_validate_video_input = (
        Gemma4ForConditionalGeneration._parse_and_validate_video_input
    )
    _parse_and_validate_audio_input = (
        Gemma4ForConditionalGeneration._parse_and_validate_audio_input
    )
    _parse_and_validate_multimodal_inputs = (
        Gemma4ForConditionalGeneration._parse_and_validate_multimodal_inputs
    )
    _process_image_input = Gemma4ForConditionalGeneration._process_image_input
    _process_video_input = Gemma4ForConditionalGeneration._process_video_input
    _process_audio_input = Gemma4ForConditionalGeneration._process_audio_input
    embed_multimodal = Gemma4ForConditionalGeneration.embed_multimodal
    _encoder_chunk = staticmethod(Gemma4ForConditionalGeneration._encoder_chunk)

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        nn.Module.__init__(self)
        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
        self.hf_to_vllm_mapper = _embedding_gemma2_weights_mapper(config)
        lora_config = vllm_config.lora_config
        self._enable_mm_lora = bool(
            lora_config is not None and lora_config.enable_tower_connector_lora
        )

        tower_quant: QuantizationConfig | None
        if quant_config and quant_config.get_name() in [
            "bitsandbytes",
            "torchao",
            "compressed-tensors",
        ]:
            tower_quant = quant_config
        else:
            vision_cfg = getattr(config, "vision_config", None)
            quantizable = (
                vision_cfg is not None
                and vision_cfg.hidden_size % 64 == 0
                and vision_cfg.intermediate_size % 64 == 0
            )
            tower_quant = quant_config if quantizable else None

        # ---- Vision tower (shared by image and video) ----
        self.embed_vision: Gemma4MultimodalEmbedder | None
        if getattr(config, "vision_config", None) is not None:
            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=tower_quant,
                    prefix=maybe_prefix(prefix, "embed_vision"),
                )
                recursive_replace_linear(
                    self.vision_tower,
                    tower_quant,
                    prefix=maybe_prefix(prefix, "vision_tower"),
                )
        else:
            self.vision_tower = None
            self.embed_vision = None

        # ---- Audio tower ----
        self.embed_audio: Gemma4MultimodalEmbedder | None
        if getattr(config, "audio_config", None) is not None:
            with self._mark_tower_model(vllm_config, "audio"):
                self.audio_tower = AutoModel.from_config(config=config.audio_config)
                self.audio_tower.post_init()
                self.embed_audio = Gemma4MultimodalEmbedder(
                    config.audio_config,
                    config.text_config,
                    quant_config=tower_quant,
                    prefix=maybe_prefix(prefix, "embed_audio"),
                )
                recursive_replace_linear(
                    self.audio_tower,
                    tower_quant,
                    prefix=maybe_prefix(prefix, "audio_tower"),
                )
        else:
            self.audio_tower = None
            self.embed_audio = None

        # ---- Language model ----
        with self._mark_language_model(vllm_config):
            self.language_model = EmbeddingGemma2TextModel(
                vllm_config=vllm_config,
                prefix=maybe_prefix(prefix, "language_model"),
            )

        assert vllm_config.model_config.pooler_config is not None
        self.pooler = DispatchPooler.for_embedding(
            vllm_config.model_config.pooler_config
        )

    def get_language_model(self) -> nn.Module:  # type: ignore[override]
        return self.language_model

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

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        ignore_prefixes = [
            "embed_vision.embedding.",
            "embed_audio.embedding.",
        ]
        if self.audio_tower is None:
            ignore_prefixes.extend(["audio_tower.", "embed_audio."])
        if self.vision_tower is None:
            ignore_prefixes.extend(["vision_tower.", "embed_vision."])
        return AutoWeightsLoader(
            self,
            ignore_unexpected_prefixes=ignore_prefixes,
        ).load_weights(weights, mapper=self.hf_to_vllm_mapper)

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