diff --git a/tests/models/multimodal/pooling/test_embedding_gemma2.py b/tests/models/multimodal/pooling/test_embedding_gemma2.py new file mode 100644 index 000000000000..701d0914ddbc --- /dev/null +++ b/tests/models/multimodal/pooling/test_embedding_gemma2.py @@ -0,0 +1,687 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import ast +import inspect +import os +import textwrap + +import pytest +import torch + +MODEL = os.environ.get("EG2_MODEL_PATH") +needs_model = pytest.mark.skipif(not MODEL, reason="EG2_MODEL_PATH not set") +TEXT_COS, MM_COS = 0.999, 0.999 + + +def _cos(a, b): + return torch.nn.functional.cosine_similarity( + torch.tensor(a, dtype=torch.float32), + torch.tensor(b, dtype=torch.float32), + dim=-1, + ) + + +def _ensure_synthetic_video( + path: str, duration: float, fps: float, width: int, height: int +) -> str: + if os.path.exists(path): + os.remove(path) + import cv2 + import numpy as np + + fourcc_fn = getattr(cv2, "VideoWriter_fourcc", cv2.VideoWriter.fourcc) + fourcc = fourcc_fn(*"mp4v") + num_frames = int(duration * fps) + writer = cv2.VideoWriter(path, fourcc, fps, (width, height)) + y_coords, x_coords = np.mgrid[0:height, 0:width] + for i in range(num_frames): + # Spatially varying color pattern that evolves dynamically over frames + frame: np.ndarray = np.zeros((height, width, 3), dtype=np.uint8) + frame[..., 0] = ((x_coords * 255 // max(1, width)) + i * 17) % 256 + frame[..., 1] = ((y_coords * 255 // max(1, height)) + i * 29) % 256 + frame[..., 2] = (x_coords + y_coords + i * 13) % 256 + writer.write(frame) + writer.release() + return path + + +@pytest.fixture(scope="module") +def st_model(): + from sentence_transformers import SentenceTransformer + + return SentenceTransformer( + MODEL, device="cuda", model_kwargs={"torch_dtype": torch.bfloat16} + ) + + +@needs_model +def test_structure(vllm_runner, monkeypatch): + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + + def probe(model): + lm = model.language_model + win = [layer.self_attn.attn.impl.sliding_window for layer in lm.layers] + return win, model.config.text_config.sliding_window + + win, cfg_w = vm.apply_model(probe)[0] + assert (cfg_w, cfg_w) in win # (512, 512) == W-1 with W = 513 + + +@needs_model +@pytest.mark.parametrize( + "prompt", + [ + "short query", + "task: search result | query: what is vllm", + "word " * 700, # > 512: crosses sliding window + "word " * 1500, # > 1024 + ], +) +def test_text_parity(vllm_runner, st_model, prompt): + ref = st_model.encode([prompt], normalize_embeddings=True) + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed([prompt]) + sim = _cos(out, ref).min().item() + assert sim >= TEXT_COS, ( + f"Text cosine {sim} < {TEXT_COS} for prompt len {len(prompt)}" + ) + + +@needs_model +def test_mm_recipe(vllm_runner, st_model): + from PIL import Image + + img = Image.new("RGB", (224, 224), color=(255, 0, 0)) + prompt = "Describe this image: <|image|>" + + # SentenceTransformer multimodal encode + ref = st_model.encode([{"image": img, "text": prompt}], normalize_embeddings=True) + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed([{"prompt": prompt, "multi_modal_data": {"image": img}}]) + sim = _cos(out, ref).min().item() + assert sim >= MM_COS, f"Image cosine {sim} < {MM_COS}" + + +@needs_model +def test_natural_image_parity(vllm_runner, st_model): + from PIL import Image + + img_path = "/projects/gemma4-vllm/images/cat.jpg" + if not os.path.exists(img_path): + pytest.skip(f"Test asset {img_path} not found") + img = Image.open(img_path) + prompt = "What is this image? <|image|>" + ref = st_model.encode([{"image": img, "text": prompt}], normalize_embeddings=True) + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed([{"prompt": prompt, "multi_modal_data": {"image": img}}]) + sim = _cos(out, ref).min().item() + assert sim >= MM_COS, f"Natural image cosine {sim} < {MM_COS}" + + +@needs_model +def test_audio_parity(vllm_runner, st_model): + import torchaudio + + audio_path = "/projects/gemma4-vllm/audio/tellMeAboutTheSun.wav" + if not os.path.exists(audio_path): + pytest.skip(f"Test asset {audio_path} not found") + wav, sr = torchaudio.load(audio_path) + mono_wav = wav.mean(dim=0).numpy() + prompt = "Describe this audio: <|audio|>" + ref = st_model.encode( + [{"audio": audio_path, "text": prompt}], normalize_embeddings=True + ) + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed( + [{"prompt": prompt, "multi_modal_data": {"audio": (mono_wav, sr)}}] + ) + sim = _cos(out, ref).min().item() + assert sim >= MM_COS, f"Audio cosine {sim} < {MM_COS}" + + +@needs_model +def test_mixed_parity(vllm_runner, st_model): + import torchaudio + from PIL import Image + + img_path = "/projects/gemma4-vllm/images/cat.jpg" + audio_path = "/projects/gemma4-vllm/audio/tellMeAboutTheSun.wav" + if not os.path.exists(img_path) or not os.path.exists(audio_path): + pytest.skip("Test assets for mixed parity not found") + img = Image.open(img_path) + wav, sr = torchaudio.load(audio_path) + mono_wav = wav.mean(dim=0).numpy() + prompt = "Describe image <|image|> and audio <|audio|>" + ref = st_model.encode( + [{"image": img, "audio": audio_path, "text": prompt}], + normalize_embeddings=True, + ) + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed( + [ + { + "prompt": prompt, + "multi_modal_data": { + "image": img, + "audio": (mono_wav, sr), + }, + } + ] + ) + sim = _cos(out, ref).min().item() + assert sim >= MM_COS, f"Mixed cosine {sim} < {MM_COS}" + + +@needs_model +def test_video_parity_uniform_overflow(vllm_runner, st_model): + from transformers.video_utils import load_video + + video_path = _ensure_synthetic_video( + "/tmp/test_40s.mp4", duration=40.0, fps=10.0, width=160, height=120 + ) + frames, meta = load_video(video_path, backend="opencv") + prompt = "Describe this video: <|video|>" + ref = st_model.encode( + [{"video": video_path, "text": prompt}], normalize_embeddings=True + ) + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed( + [{"prompt": prompt, "multi_modal_data": {"video": (frames, meta)}}] + ) + sim = _cos(out, ref).min().item() + assert sim >= MM_COS, f"40s video cosine {sim} < {MM_COS}" + + +@needs_model +def test_video_parity_portrait(vllm_runner, st_model): + from transformers.video_utils import load_video + + video_path = _ensure_synthetic_video( + "/tmp/test_portrait.mp4", duration=3.0, fps=10.0, width=120, height=160 + ) + frames, meta = load_video(video_path, backend="opencv") + prompt = "Describe this video: <|video|>" + ref = st_model.encode( + [{"video": video_path, "text": prompt}], normalize_embeddings=True + ) + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed( + [{"prompt": prompt, "multi_modal_data": {"video": (frames, meta)}}] + ) + sim = _cos(out, ref).min().item() + assert sim >= MM_COS, f"Portrait video cosine {sim} < {MM_COS}" + + +@needs_model +def test_tower_skip(vllm_runner, monkeypatch): + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + + def nparams(m): + return sum(p.numel() for p in m.parameters()) / 1e6 + + with vllm_runner( + MODEL, + runner="pooling", + max_model_len=2048, + limit_mm_per_prompt={"image": 0, "audio": 0, "video": 0}, + ) as vm: + p_text = vm.apply_model(nparams)[0] + assert abs(p_text - 271.00) <= 0.05, ( + f"Expected 271.00M text params, got {p_text}M" + ) + + with vllm_runner(MODEL, runner="pooling", max_model_len=2048) as vm: + p_full = vm.apply_model(nparams)[0] + assert abs(p_full - 744.37) <= 0.05, ( + f"Expected 744.37M full params, got {p_full}M" + ) + + +def test_encoder_contract(): + from vllm.model_executor.models.embedding_gemma2 import EmbeddingGemma2Model + from vllm.model_executor.models.gemma4_mm import ( + Gemma4ForConditionalGeneration as G, + ) + + borrowed_methods = [ + "_parse_and_validate_multimodal_inputs", + "_parse_and_validate_image_input", + "_parse_and_validate_video_input", + "_parse_and_validate_audio_input", + "_process_image_input", + "_process_video_input", + "_process_audio_input", + "embed_multimodal", + "_encoder_chunk", + ] + + for name in borrowed_methods: + # Check method identity (L-F6 + T-F4) + assert getattr(EmbeddingGemma2Model, name) is getattr(G, name) + + # AST checks: verify no zero-arg super() calls in borrowed methods + func = getattr(EmbeddingGemma2Model, name) + src = textwrap.dedent(inspect.getsource(func)) + tree = ast.parse(src) + for node in ast.walk(tree): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "super" + ): + assert len(node.args) > 0, ( + f"Zero-arg super() call found in borrowed method {name}" + ) + + # AST check: all self. non-callable accesses are assigned in __init__ + init_src = inspect.getsource(EmbeddingGemma2Model.__init__) + for name in borrowed_methods: + func = getattr(EmbeddingGemma2Model, name) + src = textwrap.dedent(inspect.getsource(func)) + tree = ast.parse(src) + for node in ast.walk(tree): + if ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "self" + ): + attr = node.attr + assert hasattr(EmbeddingGemma2Model, attr) or attr in init_src, ( + f"Attribute {attr} accessed on self in {name} is neither " + f"defined on class nor in __init__" + ) + + +@needs_model +def test_instance_attribute_access(vllm_runner): + """Check self. accesses on a real model instance.""" + from vllm.model_executor.models.embedding_gemma2 import EmbeddingGemma2Model + + borrowed_methods = [ + "_parse_and_validate_multimodal_inputs", + "_parse_and_validate_image_input", + "_parse_and_validate_video_input", + "_parse_and_validate_audio_input", + "_process_image_input", + "_process_video_input", + "_process_audio_input", + "embed_multimodal", + "_encoder_chunk", + ] + attrs_accessed = set() + for name in borrowed_methods: + func = getattr(EmbeddingGemma2Model, name) + src = textwrap.dedent(inspect.getsource(func)) + tree = ast.parse(src) + for node in ast.walk(tree): + if ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "self" + ): + attrs_accessed.add(node.attr) + + def check_attrs(model): + for attr in attrs_accessed: + assert hasattr(model, attr), ( + f"Real instance of EmbeddingGemma2Model is missing attribute '{attr}' " + f"accessed by borrowed method" + ) + return True + + with vllm_runner(MODEL, runner="pooling", max_model_len=2048) as vm: + res = vm.apply_model(check_attrs) + assert res[0] is True + + +@needs_model +def test_video_loader_real_video(vllm_runner, st_model): + """Decode a real MP4 file through EmbeddingGemma2VideoBackend.load_bytes + + and VideoMediaIO / MediaConnector with default num_frames=32. + """ + from vllm.multimodal.media import MediaConnector + from vllm.multimodal.media.video import VideoMediaIO + from vllm.multimodal.video import VIDEO_LOADER_REGISTRY + + video_path = "/projects/gemma4-vllm/videos/test_intel.mp4" + backend_cls = VIDEO_LOADER_REGISTRY.name2class["embedding_gemma2"] + with open(video_path, "rb") as f: + data = f.read() + + # Direct backend load + frames, meta = backend_cls.load_bytes(data) + + # VideoMediaIO with default num_frames=32 (VideoMediaIO passes num_frames=32) + video_io = VideoMediaIO(image_io=None, video_backend="embedding_gemma2") + loaded = video_io.load_bytes(data) + frames_io, meta_io = loaded.media + assert len(frames_io) == len(frames) + assert meta_io.get("do_sample_frames") is False + + # MediaConnector with video_processor="EmbeddingGemma2VideoProcessor" + connector = MediaConnector(allowed_local_media_path="/projects/gemma4-vllm/videos") + fetched = connector.fetch_video( + f"file://{video_path}", + video_processor="EmbeddingGemma2VideoProcessor", + ) + frames_conn, meta_conn = fetched.media + assert len(frames_conn) == 32 + assert meta_conn.get("do_sample_frames") is False + + prompt = "Describe this video: <|video|>" + ref = st_model.encode( + [{"video": video_path, "text": prompt}], normalize_embeddings=True + ) + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed( + [ + { + "prompt": prompt, + "multi_modal_data": {"video": (frames_conn, meta_conn)}, + } + ] + ) + sim = _cos(out, ref).min().item() + assert sim >= MM_COS, f"Real video loader cosine {sim} < {MM_COS}" + + +@needs_model +def test_video_loader_duplicate_expansion_e2e(vllm_runner, st_model): + """End-to-end duplicate expansion test with low-fps clip.""" + import numpy as np + + from vllm.multimodal.video import VIDEO_LOADER_REGISTRY + + backend_cls = VIDEO_LOADER_REGISTRY.name2class["embedding_gemma2"] + # 4s clip at 0.5 fps = 2 unique frames + lowfps_path = "/tmp/test_lowfps_05.mp4" + _ensure_synthetic_video(lowfps_path, duration=4.0, fps=0.5, width=64, height=48) + with open(lowfps_path, "rb") as f: + data = f.read() + # Default fps=1.0 -> 4 frames expanding 2 unique frames to [0, 0, 1, 1] + frames, meta = backend_cls.load_bytes(data) + assert frames.shape[0] == 4 + # Compare returned frames directly, not only indices + assert np.array_equal(frames[0], frames[1]) + assert not np.array_equal(frames[0], frames[2]) + assert np.array_equal(frames[2], frames[3]) + + prompt = "Describe this video: <|video|>" + with vllm_runner(MODEL, runner="pooling", max_model_len=8192) as vm: + out = vm.embed( + [{"prompt": prompt, "multi_modal_data": {"video": (frames, meta)}}] + ) + ref = st_model.encode( + [{"video": lowfps_path, "text": prompt}], normalize_embeddings=True + ) + sim = _cos(out, ref).min().item() + assert sim >= MM_COS, f"Duplicate expansion cosine {sim} < {MM_COS}" + assert len(out) == 1 + assert len(out[0]) == 768 + + +def test_video_loader_indices(): + from transformers.models.embedding_gemma2.video_processing_embedding_gemma2 import ( + EmbeddingGemma2VideoProcessor, + ) + from transformers.video_utils import VideoMetadata + + from vllm.multimodal.video import ( + VIDEO_LOADER_REGISTRY, + VideoSourceMetadata, + VideoTargetMetadata, + ) + + backend_cls = VIDEO_LOADER_REGISTRY.name2class["embedding_gemma2"] + hf_proc = EmbeddingGemma2VideoProcessor() + + # Case 1: clip shorter than 32s (e.g. 5s @ 24fps = 120 frames) + src_5s = VideoSourceMetadata(total_frames_num=120, original_fps=24.0, duration=5.0) + tgt = VideoTargetMetadata(num_frames=-1, fps=1.0, max_duration=300) + vllm_indices = backend_cls.compute_frames_index_to_sample(src_5s, tgt) + hf_meta_5s = VideoMetadata(fps=24.0, duration=5.0, total_num_frames=120) + hf_indices_5s = hf_proc.sample_frames( + hf_meta_5s, fps=1.0, max_frames=32, overflow_strategy="uniform" + ).tolist() + assert vllm_indices == hf_indices_5s + + # Case 2: clip shorter than 1s (e.g. 0.5s @ 24fps = 12 frames) + src_sub1s = VideoSourceMetadata( + total_frames_num=12, original_fps=24.0, duration=0.5 + ) + vllm_sub1s = backend_cls.compute_frames_index_to_sample(src_sub1s, tgt) + hf_meta_sub1s = VideoMetadata(fps=24.0, duration=0.5, total_num_frames=12) + hf_indices_sub1s = hf_proc.sample_frames( + hf_meta_sub1s, fps=1.0, max_frames=32, overflow_strategy="uniform" + ).tolist() + assert vllm_sub1s == hf_indices_sub1s + + # Case 3: clip longer than 32s (e.g. 40s @ 10fps = 400 frames) + # -> uniform overflow cap 32 frames + src_40s = VideoSourceMetadata( + total_frames_num=400, original_fps=10.0, duration=40.0 + ) + vllm_40s = backend_cls.compute_frames_index_to_sample(src_40s, tgt) + hf_meta_40s = VideoMetadata(fps=10.0, duration=40.0, total_num_frames=400) + hf_indices_40s = hf_proc.sample_frames( + hf_meta_40s, fps=1.0, max_frames=32, overflow_strategy="uniform" + ).tolist() + assert vllm_40s == hf_indices_40s + assert len(vllm_40s) == 32 + + # Case 4: default target fps (target.fps <= 0 defaults to 1.0) + tgt_default_fps = VideoTargetMetadata(num_frames=-1, fps=-1.0, max_duration=300) + vllm_def = backend_cls.compute_frames_index_to_sample(src_5s, tgt_default_fps) + assert vllm_def == hf_indices_5s + + # Case 5 (NEW-2): 60 fps / 300 s (18000 frames) vs HF sample_frames + src_300s = VideoSourceMetadata( + total_frames_num=18000, original_fps=60.0, duration=300.0 + ) + v_300s = backend_cls.compute_frames_index_to_sample( + src_300s, VideoTargetMetadata(num_frames=-1, fps=60.0, max_duration=300) + ) + h_300s = hf_proc.sample_frames( + VideoMetadata(fps=60.0, duration=300.0, total_num_frames=18000), + fps=60.0, + max_frames=32, + overflow_strategy="uniform", + ).tolist() + assert v_300s == h_300s, f"60fps/300s mismatch: {v_300s[-3:]} vs {h_300s[-3:]}" + + # Case 6 (NEW-2): 1 fps / 4 h (432000 frames @ 30 native fps) vs HF sample_frames + src_4h = VideoSourceMetadata( + total_frames_num=432000, original_fps=30.0, duration=14400.0 + ) + v_4h = backend_cls.compute_frames_index_to_sample( + src_4h, VideoTargetMetadata(num_frames=-1, fps=1.0, max_duration=300) + ) + h_4h = hf_proc.sample_frames( + VideoMetadata(fps=30.0, duration=14400.0, total_num_frames=432000), + fps=1.0, + max_frames=32, + overflow_strategy="uniform", + ).tolist() + assert v_4h == h_4h, f"1fps/4h mismatch: {v_4h[-2:]} vs {h_4h[-2:]}" + + +def test_video_loader_duplicate_expansion(): + from transformers.models.embedding_gemma2.video_processing_embedding_gemma2 import ( + EmbeddingGemma2VideoProcessor, + ) + from transformers.video_utils import VideoMetadata + + from vllm.multimodal.video import ( + VIDEO_LOADER_REGISTRY, + VideoSourceMetadata, + VideoTargetMetadata, + ) + + backend_cls = VIDEO_LOADER_REGISTRY.name2class["embedding_gemma2"] + hf_proc = EmbeddingGemma2VideoProcessor() + + # 4s clip @ 0.5fps = 2 frames total. Requested target fps = 1.0 -> 4 frames. + # Source has fewer frames than requested sampled frames, requiring duplicates. + src_lowfps = VideoSourceMetadata(total_frames_num=2, original_fps=0.5, duration=4.0) + tgt = VideoTargetMetadata(num_frames=-1, fps=1.0, max_duration=300) + vllm_indices = backend_cls.compute_frames_index_to_sample(src_lowfps, tgt) + hf_meta = VideoMetadata(fps=0.5, duration=4.0, total_num_frames=2) + hf_indices = hf_proc.sample_frames( + hf_meta, fps=1.0, max_frames=32, overflow_strategy="uniform" + ).tolist() + assert vllm_indices == hf_indices + assert len(vllm_indices) == 4 + assert vllm_indices == [0, 0, 1, 1] + + +def test_video_media_io_embedding_gemma2(): + """Verify VideoMediaIO with default num_frames=32 loads video bytes for + + EmbeddingGemma2 without error. + """ + import os + from unittest.mock import MagicMock + + from vllm.multimodal.media import MediaConnector + from vllm.multimodal.media.video import VideoMediaIO + + video_path = "/projects/gemma4-vllm/videos/test_intel.mp4" + if not os.path.exists(video_path): + return + + with open(video_path, "rb") as f: + data = f.read() + + # VideoMediaIO defaults to num_frames=32 + video_io = VideoMediaIO(image_io=MagicMock(), video_backend="embedding_gemma2") + loaded = video_io.load_bytes(data) + frames, meta = loaded.media + assert len(frames) == 32 + assert meta.get("do_sample_frames") is False + + # MediaConnector path + connector = MediaConnector(allowed_local_media_path="/projects/gemma4-vllm/videos") + fetched = connector.fetch_video( + f"file://{video_path}", + video_processor="EmbeddingGemma2VideoProcessor", + ) + frames_conn, meta_conn = fetched.media + assert len(frames_conn) == 32 + assert meta_conn.get("do_sample_frames") is False + + +def test_config_capping_and_explicit_limits(caplog): + """Verify EmbeddingGemma2ModelConfig capping and explicit user overrides.""" + import logging + from types import SimpleNamespace + from unittest.mock import patch + + from vllm.model_executor.models.config import ( + EmbeddingGemma2ModelConfig, + Gemma4Config, + ) + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + # (a) Defaults: uncapped max_model_len=262144 -> capped to 8192, + # and derived scheduler limits (>= 262144) lowered to 8192. + cfg_default = SimpleNamespace( + attention_config=SimpleNamespace(backend=None), + model_config=SimpleNamespace( + model="dummy", max_model_len=262144, original_max_model_len=None + ), + scheduler_config=SimpleNamespace( + max_num_batched_tokens=262144, + max_num_encoder_input_tokens=262144, + encoder_cache_size=262144, + max_num_seqs=256, + verify_max_model_len=lambda x: None, + ), + ) + with patch.object(Gemma4Config, "verify_and_update_config"): + EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_default) + + assert cfg_default.attention_config.backend == AttentionBackendEnum.TRITON_ATTN + assert cfg_default.model_config.max_model_len == 8192 + assert cfg_default.scheduler_config.max_num_batched_tokens == 8192 + assert cfg_default.scheduler_config.max_num_encoder_input_tokens == 8192 + assert cfg_default.scheduler_config.encoder_cache_size == 8192 + + # (b) Reachable explicit limits: explicit max_model_len=8192 and explicit + # max_num_batched_tokens=32768 stays 32768. + # Note: omitting max_model_len while passing max_num_batched_tokens < 262144 + # is rejected earlier by SchedulerConfig validation (32768 < 262144) before + # this hook runs. + cfg_explicit_batched = SimpleNamespace( + attention_config=SimpleNamespace(backend=None), + model_config=SimpleNamespace( + model="dummy", max_model_len=8192, original_max_model_len=8192 + ), + scheduler_config=SimpleNamespace( + max_num_batched_tokens=32768, + max_num_encoder_input_tokens=32768, + encoder_cache_size=32768, + max_num_seqs=256, + verify_max_model_len=lambda x: None, + ), + ) + with patch.object(Gemma4Config, "verify_and_update_config"): + EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_batched) + + assert cfg_explicit_batched.model_config.max_model_len == 8192 + assert cfg_explicit_batched.scheduler_config.max_num_batched_tokens == 32768 + assert cfg_explicit_batched.scheduler_config.max_num_encoder_input_tokens == 32768 + assert cfg_explicit_batched.scheduler_config.encoder_cache_size == 32768 + + # (c) Explicit value > uncapped_len (e.g. 300000) with uncapped model len is lowered + # to 8192 and logs a warning because it exceeds uncapped_len. + cfg_explicit_large = SimpleNamespace( + attention_config=SimpleNamespace(backend=None), + model_config=SimpleNamespace( + model="dummy", max_model_len=262144, original_max_model_len=None + ), + scheduler_config=SimpleNamespace( + max_num_batched_tokens=300000, + max_num_encoder_input_tokens=300000, + encoder_cache_size=300000, + max_num_seqs=256, + verify_max_model_len=lambda x: None, + ), + ) + with ( + caplog.at_level(logging.WARNING), + patch.object(Gemma4Config, "verify_and_update_config"), + ): + EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_large) + + assert cfg_explicit_large.model_config.max_model_len == 8192 + assert cfg_explicit_large.scheduler_config.max_num_batched_tokens == 8192 + assert cfg_explicit_large.scheduler_config.max_num_encoder_input_tokens == 8192 + assert cfg_explicit_large.scheduler_config.encoder_cache_size == 8192 + assert any( + "lowering explicitly set scheduler_config.max_num_batched_tokens" + " from 300000 to 8192" in rec.message + for rec in caplog.records + ) + + # (d) Explicit max_model_len=16384 is respected and not capped to 8192 + cfg_explicit_model_len = SimpleNamespace( + attention_config=SimpleNamespace(backend=None), + model_config=SimpleNamespace( + model="dummy", max_model_len=16384, original_max_model_len=16384 + ), + scheduler_config=SimpleNamespace( + max_num_batched_tokens=16384, + max_num_encoder_input_tokens=16384, + encoder_cache_size=16384, + max_num_seqs=256, + verify_max_model_len=lambda x: None, + ), + ) + with patch.object(Gemma4Config, "verify_and_update_config"): + EmbeddingGemma2ModelConfig.verify_and_update_config(cfg_explicit_model_len) + + assert cfg_explicit_model_len.model_config.max_model_len == 16384 + assert cfg_explicit_model_len.scheduler_config.max_num_batched_tokens == 16384 diff --git a/tests/models/registry.py b/tests/models/registry.py index b764cad595a8..95b424ede580 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -581,6 +581,11 @@ def check_available_online( trust_remote_code=True, hf_overrides={"model_type": "deepseek_v3", "auto_map": None}, ), + "EmbeddingGemma2Model": _HfExamplesInfo( + "google/embeddinggemma-2", + min_transformers_version="5.18.0", + is_available_online=False, + ), "Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"), "Gemma3TextModel": _HfExamplesInfo("google/embeddinggemma-300m"), "GteModel": _HfExamplesInfo( diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index d80663731d1e..97b9bfa0d38b 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -282,6 +282,111 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: ) +class EmbeddingGemma2ModelConfig(Gemma4Config): + @staticmethod + def verify_and_update_config(vllm_config: "VllmConfig") -> None: + Gemma4Config.verify_and_update_config(vllm_config) + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + attn_config = vllm_config.attention_config + if attn_config.backend is None: + attn_config.backend = AttentionBackendEnum.TRITON_ATTN + logger.info( + "EmbeddingGemma2: attention backend not specified; defaulting to " + "TRITON_ATTN (validated backend for heterogeneous head dimensions)." + ) + else: + logger.info( + "EmbeddingGemma2: using attention backend %s", attn_config.backend + ) + + model_config = vllm_config.model_config + orig_len = getattr(model_config, "original_max_model_len", None) + if orig_len in (None, -1): + default_len = 8192 + # Check sentence_bert_config.json if available + try: + import json + from pathlib import Path + + cfg_path = Path(model_config.model) / "sentence_bert_config.json" + if cfg_path.exists(): + data = json.loads(cfg_path.read_text()) + if "max_seq_length" in data and isinstance( + data["max_seq_length"], int + ): + default_len = data["max_seq_length"] + except Exception: + pass + + if model_config.max_model_len > default_len: + uncapped_len = model_config.max_model_len + logger.info( + "EmbeddingGemma2: max_model_len not explicitly set; capping from " + "%d to %d.", + model_config.max_model_len, + default_len, + ) + model_config.max_model_len = default_len + scheduler_config = vllm_config.scheduler_config + if scheduler_config is not None: + new_batched_tokens = max( + model_config.max_model_len, scheduler_config.max_num_seqs + ) + if scheduler_config.max_num_batched_tokens >= uncapped_len: + if scheduler_config.max_num_batched_tokens > uncapped_len: + logger.warning( + "EmbeddingGemma2: lowering explicitly set " + "scheduler_config.max_num_batched_tokens from %d to %d " + "to match capped max_model_len.", + scheduler_config.max_num_batched_tokens, + new_batched_tokens, + ) + else: + logger.info( + "EmbeddingGemma2: lowering " + "scheduler_config.max_num_batched_tokens from %d to %d " + "to match capped max_model_len.", + scheduler_config.max_num_batched_tokens, + new_batched_tokens, + ) + scheduler_config.max_num_batched_tokens = new_batched_tokens + if hasattr(vllm_config, "_set_compile_ranges"): + vllm_config._set_compile_ranges() + if scheduler_config.max_num_encoder_input_tokens >= uncapped_len: + if scheduler_config.max_num_encoder_input_tokens > uncapped_len: + logger.warning( + "EmbeddingGemma2: lowering explicitly set " + "scheduler_config.max_num_encoder_input_tokens " + "from %d to %d to match capped max_model_len.", + scheduler_config.max_num_encoder_input_tokens, + new_batched_tokens, + ) + scheduler_config.max_num_encoder_input_tokens = ( + new_batched_tokens + ) + if scheduler_config.encoder_cache_size >= uncapped_len: + if scheduler_config.encoder_cache_size > uncapped_len: + logger.warning( + "EmbeddingGemma2: lowering explicitly set " + "scheduler_config.encoder_cache_size from %d to %d " + "to match capped max_model_len.", + scheduler_config.encoder_cache_size, + new_batched_tokens, + ) + scheduler_config.encoder_cache_size = new_batched_tokens + if hasattr(scheduler_config, "verify_max_model_len"): + scheduler_config.verify_max_model_len( + model_config.max_model_len + ) + elif model_config.max_model_len > 32768: + logger.warning( + "EmbeddingGemma2: max_model_len=%d is large; " + "consider --max-model-len 8192 to bound encoder memory.", + model_config.max_model_len, + ) + + class DiffusionGemmaModelForBlockDiffusionConfig(VerifyAndUpdateConfig): @classmethod def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: @@ -1025,6 +1130,7 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: "DeepseekV41ForCausalLM": DeepseekV4ForCausalLMConfig, "DeepseekV32ForCausalLM": DeepseekV32ForCausalLM, "DiffusionGemmaForBlockDiffusion": DiffusionGemmaModelForBlockDiffusionConfig, # noqa: E501 + "EmbeddingGemma2Model": EmbeddingGemma2ModelConfig, "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501 "FalconMambaForCausalLM": MambaModelConfig, "Gemma3TextModel": Gemma3TextModelConfig, diff --git a/vllm/model_executor/models/embedding_gemma2.py b/vllm/model_executor/models/embedding_gemma2.py new file mode 100644 index 000000000000..2d0012714eab --- /dev/null +++ b/vllm/model_executor/models/embedding_gemma2.py @@ -0,0 +1,1142 @@ +# 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.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, + gemma4_layer_config, +) +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 = gemma4_layer_config(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) + + +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 = ""): + 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 + 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}") diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 48b831bbdf0a..e019d674f502 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -246,6 +246,7 @@ # [Multimodal] "CLIPModel": ("clip", "CLIPEmbeddingModel"), "ColPaliForRetrieval": ("colpali", "ColPaliModel"), + "EmbeddingGemma2Model": ("embedding_gemma2", "EmbeddingGemma2Model"), "LlamaNemotronVLModel": ("nemotron_vl", "LlamaNemotronVLForEmbedding"), "LlavaNextForConditionalGeneration": ( "llava_next", diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index 44a8e9452dc3..8f22bffe958c 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -405,6 +405,184 @@ def load_bytes( ) +@VIDEO_LOADER_REGISTRY.register( + "embedding_gemma2", + video_processor="EmbeddingGemma2VideoProcessor", +) +class EmbeddingGemma2VideoBackend(VideoBackend): + """Samples frames with the same indices as HF EmbeddingGemma2VideoProcessor.""" + + PROCESSOR_NAMES = ("EmbeddingGemma2Processor", "EmbeddingGemma2VideoProcessor") + + _MAX_ALLOWED_FRAMES: ClassVar[int] = 128 + _MAX_ALLOWED_FPS: ClassVar[float] = 60.0 + + @classmethod + def _get_processor_defaults(cls, **kwargs) -> tuple[float, int]: + """Read default fps and max_frames from video_processor if available.""" + proc = kwargs.get("video_processor") or kwargs.get("processor") + if proc is not None: + vp = getattr(proc, "video_processor", proc) + fps = getattr(vp, "fps", None) + max_frames = getattr(vp, "max_frames", None) + if fps is not None and max_frames is not None: + return float(fps), int(max_frames) + try: + from transformers.models.embedding_gemma2 import ( + EmbeddingGemma2VideoProcessor, + ) + + vp = EmbeddingGemma2VideoProcessor() + return float(getattr(vp, "fps", 1.0)), int(getattr(vp, "max_frames", 32)) + except Exception: + return 1.0, 32 + + @classmethod + def compute_frames_index_to_sample( + cls, + source: VideoSourceMetadata, + target: VideoTargetMetadata, + **kwargs, + ) -> list[int]: + # Ignores target.num_frames (EmbeddingGemma2 does not accept num_frames) + default_fps, default_max_frames = cls._get_processor_defaults(**kwargs) + + fps = kwargs.get("fps") + if fps is None: + fps = target.fps if target.fps > 0 else default_fps + if fps is not None: + if fps <= 0: + raise ValueError(f"fps must be positive, got {fps}") + if fps > cls._MAX_ALLOWED_FPS: + raise ValueError( + f"fps {fps} exceeds maximum allowed {cls._MAX_ALLOWED_FPS}" + ) + + max_frames = kwargs.get("max_frames", default_max_frames) + if max_frames is not None: + if max_frames <= 0: + raise ValueError(f"max_frames must be positive, got {max_frames}") + if max_frames > cls._MAX_ALLOWED_FRAMES: + raise ValueError( + f"max_frames {max_frames} exceeds maximum allowed " + f"{cls._MAX_ALLOWED_FRAMES}" + ) + + overflow_strategy = kwargs.get("overflow_strategy", "uniform") + + total_frames = source.total_frames_num + original_fps = source.original_fps + duration = source.duration + + if fps is not None and (original_fps <= 0 or duration <= 0): + fps = None + + if overflow_strategy is not None: + if max_frames is None: + raise ValueError( + "You must pass max_frames when requesting an " + f"overflow_strategy={overflow_strategy}!" + ) + if overflow_strategy not in ("truncate", "uniform"): + raise ValueError( + f"You passed overflow_strategy={overflow_strategy} but " + "expected one of ['truncate', 'uniform']" + ) + + if fps is None: + if overflow_strategy is None or total_frames <= max_frames: + indices = np.arange(total_frames, dtype=int) + elif overflow_strategy == "truncate": + indices = np.arange(max_frames, dtype=int) + elif overflow_strategy == "uniform": + indices = np.linspace(0, total_frames - 1, max_frames, dtype=int) + else: + step = original_fps / fps + num_sampled = max(1, int(duration * fps)) + + if overflow_strategy is None or num_sampled <= max_frames: + i_arr = np.arange(num_sampled, dtype=float) + indices = np.minimum(total_frames - 1, (i_arr * step).astype(int)) + elif overflow_strategy == "truncate": + i_arr = np.arange(max_frames, dtype=float) + indices = np.minimum(total_frames - 1, (i_arr * step).astype(int)) + elif overflow_strategy == "uniform": + linspace_idx = np.linspace(0, num_sampled - 1, max_frames, dtype=int) + indices = np.minimum( + total_frames - 1, (linspace_idx * step).astype(int) + ) + + return indices.tolist() + + @classmethod + def load_bytes( + cls, + data: bytes, + num_frames: int = -1, + fps: int | float = -1, + max_duration: int = 300, + frame_recovery: bool = False, + *, + backend: VideoDecoderBackend = "opencv", + **kwargs: Any, + ) -> tuple[DecodedFrames, dict[str, Any]]: + # Accept num_frames as required by VideoLoader interface (VideoMediaIO + # passes default num_frames=32), but ignore it since EmbeddingGemma2 uses + # fps/max_frames. Force num_frames=-1 so compute_frames_index_to_sample + # drives sampling via fps/duration. + kwargs.pop("num_frames", None) + frames, metadata = super().load_bytes( + data, + num_frames=-1, + fps=fps, # type: ignore[arg-type] + max_duration=max_duration, + frame_recovery=frame_recovery, + backend=backend, + **kwargs, + ) + + source = VideoSourceMetadata( + total_frames_num=metadata["total_num_frames"], + original_fps=metadata["fps"], + duration=metadata["duration"], + ) + target = VideoTargetMetadata(num_frames=-1, fps=fps, max_duration=max_duration) + sampling_kwargs, _ = resolve_video_backend_kwargs(backend, kwargs) + requested_indices = cls.compute_frames_index_to_sample( + source, target, **sampling_kwargs + ) + + # Expand deduplicated frames back to requested_indices if needed + # (e.g., OpenCV deduplicates frame_idx using set(frame_idx)) + valid_frame_indices = metadata.get("frames_indices") + if valid_frame_indices is not None and requested_indices != valid_frame_indices: + idx_to_pos = {idx: pos for pos, idx in enumerate(valid_frame_indices)} + if all(idx in idx_to_pos for idx in requested_indices): + positions = [idx_to_pos[idx] for idx in requested_indices] + if isinstance(frames, np.ndarray): + frames = frames[positions] + else: + frames = frames[torch.tensor(positions, device=frames.device)] + metadata["frames_indices"] = list(requested_indices) + + return frames, metadata + + @classmethod + def create_hf_metadata( + cls, + source: VideoSourceMetadata, + valid_frame_indices: list[int], + video_backend: str, + ) -> dict[str, Any]: + meta = super().create_hf_metadata( + source=source, + valid_frame_indices=valid_frame_indices, + video_backend=video_backend, + ) + meta["do_sample_frames"] = False + return meta + + @VIDEO_LOADER_REGISTRY.register( "qwen2_vl", video_processor="Qwen2VLVideoProcessor", diff --git a/vllm/v1/attention/ops/triton_prefill_attention.py b/vllm/v1/attention/ops/triton_prefill_attention.py index 1ed55c3e28b5..9e8482b5ecdd 100644 --- a/vllm/v1/attention/ops/triton_prefill_attention.py +++ b/vllm/v1/attention/ops/triton_prefill_attention.py @@ -243,9 +243,10 @@ def context_attention_fwd( b_seq_len: [b] out: [b * s, head, head_dim] """ - BLOCK = get_block_size(q.dtype) - Lq, Lk, _ = q.shape[-1], k.shape[-1], v.shape[-1] + BLOCK = get_block_size(q.dtype) + if Lk >= 512: + BLOCK = min(BLOCK, 32) sm_scale = 1.0 / (Lq**0.5) if softmax_scale is None else softmax_scale # rescale with 1/ln(2) for triton exp2