[Model] Support EmbeddingGemma2 multimodal pooling architecture

[lucianommartins]

[Model] Support EmbeddingGemma 2 (Omni-Modal Embeddings)

Overview

This PR adds support for EmbeddingGemma 2 (google/embeddinggemma-2), Google's bidirectional omni-modal embedding model. EmbeddingGemma 2 maps text, image, audio, video, and interleaved multimodal inputs into a unified 768-dimensional L2-normalized vector space for cross-modal and multi-modal semantic search, retrieval, and classification.

Claims

• Omni-Modal Embedding Architecture: Supports single-modal and joint/interleaved inputs across Text, Image, Audio, and Video producing unit-normalized 768-dim embeddings in a single forward pass.
• Bidirectional Attention & RoPE: Implements bidirectional attention (ENCODER_ONLY) with an inclusive sliding window of 513 tokens (config.sliding_window + 1) and 512-dim rotary embeddings (RoPE) across all layers, matching Hugging Face / SentenceTransformers semantics.
• Upfront Dynamic Modality Loading: Supports bypassing unused modality encoders at engine startup via --limit-mm-per-prompt (e.g., image=0,audio=0,video=0). Omitting vision (167.8M params) and audio (305.6M params) loads a lightweight 271.0M param text-only model (~517 MiB VRAM), saving up to 64% VRAM with zero post-load deletion.
• Dedicated Video Sampling Backend: Adds EmbeddingGemma2VideoBackend to vllm/multimodal/video.py, providing 1 fps uniform sampling up to 32 frames with duplicate-frame expansion for short clips (< 1s), exactly matching EmbeddingGemma2VideoProcessor.
• Zero Generative Regressions: Implemented cleanly as a standalone model in vllm/model_executor/models/embedding_gemma2.py without mutating existing generative Gemma 4 model classes (gemma4.py, gemma4_mm.py).
• Validated Numerical Parity: Achieves >= 0.999 cosine similarity against Hugging Face / SentenceTransformers across all 16 multimodal evaluation cases (text, image, audio, video, and mixed inputs).

Validation

Unit and Parity Tests

All 20 unit and numerical parity tests pass on NVIDIA A100-SXM4-80GB:

# Run standalone unit and parity test suite
pytest tests/models/multimodal/pooling/test_embedding_gemma2.py -v

# Verify model registry loading
pytest tests/models/test_registry.py -k "EmbeddingGemma2"

CLI Serving Examples

# 1. Full Multimodal Serving (Text + Image + Audio + Video)
vllm serve google/embeddinggemma-2 \
    --runner pooling \
    --max-model-len 8192

# 2. Text-Only Deployment (Lightweight, ~517 MiB VRAM)
vllm serve google/embeddinggemma-2 \
    --runner pooling \
    --limit-mm-per-prompt image=0,audio=0,video=0 \
    --max-model-len 8192

# 3. Vision + Text Deployment (No Audio Tower)
vllm serve google/embeddinggemma-2 \
    --runner pooling \
    --limit-mm-per-prompt audio=0 \
    --max-model-len 8192

Python API Example

from vllm import LLM
from PIL import Image

# Initialize pooling model
llm = LLM(
    model="google/embeddinggemma-2",
    runner="pooling",
    max_model_len=8192,
)

# Text embedding
text_outputs = llm.embed(["Retrieve relevant documents for machine learning research"])
print("Text embedding shape:", len(text_outputs[0].outputs.embedding))  # 768

# Multimodal text + image embedding
image = Image.new("RGB", (224, 224), color=(73, 109, 137))
mm_outputs = llm.embed([{
    "prompt": "What is depicted in this image? <|image|>",
    "multi_modal_data": {"image": image},
}])
print("Multimodal embedding norm:", sum(x**2 for x in mm_outputs[0].outputs.embedding) ** 0.5)  # 1.0

Numerical Parity Test Results

1. Text Embeddings

TEST: Text Embedding (Shape & Norm)
  Prompt:    "Hello world"
  Embedding: [-0.0598, +0.0352, +0.0593, -0.0393, -0.0352, ... +0.0396] (dim=768)
  Norm:      1.0000 (unit L2 norm)
✅ PASSED: Shape (768,) and unit norm

TEST: Text Semantic Similarity
  Cosine Sim (A vs B): 0.8914  [Semantically similar: high]
  Cosine Sim (A vs C): 0.6065  [Different topic: low]
✅ PASSED: Similar > dissimilar (0.8914 > 0.6065)

TEST: Batch vs Individual Consistency
  Min Cosine Similarity: 0.999967
✅ PASSED: Batch matches individual encodings (cosine sim > 0.9999)

2. Image Embeddings

TEST: Image Embedding (Shape & Norm)
  Norm:      1.0000 (unit L2-normalized)
✅ PASSED: Valid 768-dim image embedding

TEST: Image Discrimination
  Cosine Sim: 0.6324  [Distinct images]
✅ PASSED: Different images have distinct embeddings (0.6324 < 0.99)

TEST: Text + Image Joint Embedding
  Norm:      1.0000
✅ PASSED: Valid joint text+image embedding

3. Audio Embeddings

TEST: Audio Embedding (Shape & Norm)
  Norm:      1.0000 (unit L2-normalized)
✅ PASSED: Valid 768-dim audio embedding

TEST: Audio Discrimination
  Cosine Sim: 0.7116  [Distinct audio signals]
✅ PASSED: Different audio inputs have distinct embeddings (0.7116 < 0.99)

TEST: Text + Audio Joint Embedding
  Norm:      1.0000
✅ PASSED: Valid joint text+audio embedding

4. Video Embeddings

TEST: Video Embedding (Shape & Norm)
  Video Clip: synthetic_clip.mp4 (4.0s clip -> 4 sampled frames at fps=1)
  Norm:       1.0000 (unit L2-normalized)
✅ PASSED: Valid 768-dim video embedding

TEST: Text + Video Joint Embedding
  Text:       "A video clip showing conveyor belt operations"
  Norm:       1.0000
✅ PASSED: Valid joint text+video embedding

TEST: Video File Loader Input (MediaConnector / VideoMediaIO)
  Video File: file://synthetic_clip.mp4 (32 sampled frames, 3906 tokens)
  Norm:       1.0000
✅ PASSED: Valid video embedding via MediaConnector file loader

5. HuggingFace / SentenceTransformer vs vLLM Parity (16/16 Passed)

=========================================================================================================
Test Case                      | HF Norm  | vLLM Norm | HF Tok  | vLLM Tok | Cosine Sim  | Status
---------------------------------------------------------------------------------------------------------
Short text                     | 1.0028   | 1.0000    | 4       | 4        | 0.999946    | PASS
Sentence                       | 1.0000   | 1.0000    | 12      | 12       | 0.999836    | PASS
Technical paragraph            | 1.0037   | 1.0000    | 12      | 12       | 0.999955    | PASS
Task prompt (inlined)          | 1.0012   | 1.0000    | 15      | 15       | 0.999960    | PASS
Named prompt query             | 1.0012   | 1.0000    | 15      | 15       | 0.999960    | PASS
Long text > 512 tok            | 1.0015   | 1.0000    | 710     | 710      | 0.999883    | PASS
Long text > 1024 tok           | 1.0030   | 1.0000    | 1810    | 1810     | 0.999931    | PASS
Single image                   | 0.9991   | 1.0000    | 260     | 260      | 0.999985    | PASS
Text + Image                   | 0.9996   | 1.0000    | 265     | 265      | 0.999915    | PASS
Prompt + Image                 | 1.0016   | 1.0000    | 267     | 267      | 0.999944    | PASS
Prompt + Text + Image          | 1.0021   | 1.0000    | 273     | 273      | 0.999974    | PASS
Audio                          | 1.0023   | 1.0000    | 22      | 22       | 0.999933    | PASS
Text + Audio                   | 1.0027   | 1.0000    | 27      | 27       | 0.999907    | PASS
Video (frames + metadata)      | 1.0027   | 1.0000    | 368     | 368      | 0.999825    | PASS
Video (file loader)            | 1.0001   | 1.0000    | 3906    | 3906     | 0.999931    | PASS
Text + Image + Audio           | 1.0002   | 1.0000    | 284     | 284      | 0.999907    | PASS
=========================================================================================================
ALL 16 TEST CASES PASSED (token parity + cosine similarity >= 0.999)

Details

Architecture & Key Components

1. Standalone Model Architecture (vllm/model_executor/models/embedding_gemma2.py):

   • EmbeddingGemma2Model: Inherits from nn.Module, SupportsMultiModal, and VllmModelForPooling. Integrates vision and audio modality encoders via _mark_tower_model, language model backbone (EmbeddingGemma2TextModel), and a DispatchPooler.for_embedding pooling head with MEAN pooling and L2 normalization.
   • EmbeddingGemma2Attention: Implements bidirectional attention across all layers using ENCODER_ONLY attention. For sliding-window layers, sets sliding window size to config.sliding_window + 1 (513) to reconcile Hugging Face's |q - k| <= W condition with vLLM's |q - k| < W encoder attention kernel. Applies full 512-dim RoPE on global attention layers.
   • EmbeddingGemma2TextPLE & EmbeddingGemma2PLEBlock: Per-layer embedding projections matching the checkpoint weight schema.

2. Multimodal Processing & Tokenization:

   • EmbeddingGemma2ProcessingInfo: Retains standard <bos> and <eos> special tokens without applying chat templates, matching SentenceTransformers.
   • EmbeddingGemma2MultiModalProcessor: Dynamically budgets soft tokens (max_soft_tokens=140 per video frame) and suppresses unconditional timestamp prefixing (exclude_timestamps=True). Safely masks multimodal placeholder tokens exceeding vocab_size before embedding lookups.

3. Video Backend (vllm/multimodal/video.py):

   • Registers EmbeddingGemma2VideoBackend in VIDEO_LOADER_REGISTRY (embedding_gemma2).
   • Implements frame index computation mirroring EmbeddingGemma2VideoProcessor.sample_frames, with support for uniform overflow capping at 32 frames and duplicate frame expansion for short clips (< 1s).

4. Configuration & Attention Capping (vllm/model_executor/models/config.py, vllm/v1/attention/ops/triton_prefill_attention.py):

   • EmbeddingGemma2ModelConfig: Automatically caps default max_model_len to 8192 (down from the 262k theoretical context) and syncs scheduler token limits while respecting explicit user limits.
   • Triton Prefill Attention: Limits block size to 32 for head dimensions >= 512 to avoid shared memory exhaustion during bidirectional prefill attention.

5. Upfront Dynamic Modality Loading:

   • Integrated with vLLM's SupportsMultiModal lifecycle: when a modality is disabled via limit_mm_per_prompt (e.g., image=0,video=0 or audio=0), the corresponding encoder tower is created on torch.device("meta") and skipped during weight loading. Only active weights are loaded into VRAM.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

✅ @lucianommartins, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Your branch must contain every commit currently on its upstream target branch. Merge or rebase onto the latest target branch, then rerun the command. Append --allow-stale to a run command to test an outdated branch at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[DarkLight1337]

pre-commit is failing on main because of this PR, could you fix it ASAP?