diff --git a/tests/models/multimodal/processing/test_gemma4_unified.py b/tests/models/multimodal/processing/test_gemma4_unified.py index 67a81ddb7b41..bdecae62e28e 100644 --- a/tests/models/multimodal/processing/test_gemma4_unified.py +++ b/tests/models/multimodal/processing/test_gemma4_unified.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math from collections.abc import Mapping +import numpy as np import pytest import torch from PIL import Image as PILImage @@ -204,3 +206,103 @@ def test_limit_mm_per_prompt( mm_items=processor.info.parse_mm_data(mm_data), hf_processor_mm_kwargs={}, ) + + +def test_audio_field_batching_repads_ragged_lengths(): + """Audios with different frame counts arrive at the model as a plain + list (batched() does not re-pad); _pad_ragged_audio_features must + normalize them back to padded tensors.""" + features = [torch.randn(3, 640), torch.randn(6, 640)] + masks = [torch.ones(3, dtype=torch.bool), torch.ones(6, dtype=torch.bool)] + + field = MultiModalFieldConfig.batched("audio").field + reduced_features = field.reduce_data( + list(field.build_elems("audio", "input_features_padded", features)) + ) + reduced_masks = field.reduce_data( + list(field.build_elems("audio", "input_features_mask", masks)) + ) + assert isinstance(reduced_features, list) + + padded_features, padded_masks = _pad_ragged_audio_features( + reduced_features, reduced_masks + ) + assert padded_features.shape == torch.Size([2, 6, 640]) + assert padded_masks.shape == torch.Size([2, 6]) + assert padded_masks.sum(-1).tolist() == [3, 6] + assert torch.equal(padded_features[0, :3], features[0]) + assert torch.equal(padded_features[1], features[1]) + + Gemma4AudioInputs( + input_features_padded=padded_features, + input_features_mask=padded_masks, + ) + + +@pytest.mark.parametrize( + "residue", + # residues 1..160 are exactly where the tower-based mel/conv formula + # undercounts ceil(L / 640) by one + [0, 1, 160, 161, 639], +) +@pytest.mark.parametrize("model_id", [GEMMA4_UNIFIED_MODEL_ID]) +def test_audio_repl_token_count_matches_hf_processor(model_id: str, residue: int): + """Unified audio placeholder count must equal the HF processor's + per-audio valid frame count, ceil(L / 640).""" + ctx = build_model_context( + model_id, + limit_mm_per_prompt={"audio": 1}, + ) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + hf_processor = processor.info.get_hf_processor() + + audio_len = 640 * 5 + residue + repl = processor.info.get_audio_repl( + audio_len=audio_len, + processor=hf_processor, + ) + num_repl_tokens = repl.full.count(hf_processor.audio_token_id) + + hf_features = hf_processor.feature_extractor( + [np.zeros(audio_len, dtype=np.float32)] + ) + num_hf_tokens = int(hf_features["input_features_mask"][0].sum()) + + assert num_hf_tokens == math.ceil(audio_len / 640) + assert num_repl_tokens == num_hf_tokens + + +@pytest.mark.parametrize("model_id", [GEMMA4_UNIFIED_MODEL_ID]) +def test_multi_audio_apply_different_lengths(model_id: str): + """Two audios of different lengths must produce self-contained + unpadded per-item features whose frame counts match the placeholders.""" + ctx = build_model_context( + model_id, + limit_mm_per_prompt={"audio": 2}, + ) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + hf_processor = processor.info.get_hf_processor() + + sr = hf_processor.feature_extractor.sampling_rate + audios = [ + (np.zeros(640 * 3, dtype=np.float32), sr), + (np.zeros(640 * 5 + 100, dtype=np.float32), sr), + ] + expected_frames = [3, 6] + + prompt = hf_processor.audio_token * 2 + processed_inputs = processor( + prompt, + mm_items=processor.info.parse_mm_data({"audio": audios}), + hf_processor_mm_kwargs={}, + ) + + audio_placeholders = processed_inputs["mm_placeholders"]["audio"] + assert [p.get_num_embeds() for p in audio_placeholders] == expected_frames + + mm_items = processed_inputs["mm_kwargs"]["audio"] + for item, num_frames in zip(mm_items, expected_frames, strict=True): + assert item["input_features_padded"].data.shape == torch.Size( + [num_frames, 640] + ) + assert int(item["input_features_mask"].data.sum()) == num_frames diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index e158a7850164..fa2c3e4fbb92 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -22,6 +22,7 @@ import torch from PIL import Image as PILImage from torch import nn +from torch.nn.utils.rnn import pad_sequence from transformers import AutoModel, BatchFeature from transformers.models.gemma4 import ( Gemma4Config, @@ -166,6 +167,26 @@ class Gemma4AudioInputs(TensorSchema): ] +def _pad_ragged_audio_features( + input_features_padded: torch.Tensor | list[torch.Tensor], + input_features_mask: torch.Tensor | list[torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor]: + """Re-pad per-audio features to the batch max sequence length. + + `_call_hf_processor` stores each audio's features unpadded so that + cache entries are independent of the batch they were processed with. + `MultiModalFieldConfig.batched()` stacks items only when their shapes + are all equal and otherwise yields a plain list, so audios with + different frame counts arrive here as ragged lists. + """ + if isinstance(input_features_padded, torch.Tensor): + return input_features_padded, input_features_mask + return ( + pad_sequence(list(input_features_padded), batch_first=True), + pad_sequence(list(input_features_mask), batch_first=True, padding_value=False), + ) + + Gemma4ImageInputs = Gemma4ImagePixelInputs @@ -749,10 +770,11 @@ def _call_hf_processor( if "input_features" in processed_outputs: # Unpad per-item so each item's cache entry is - # self-contained. The batched() field config in - # _get_mm_fields_config will re-pad all fields to the - # batch's max length at batch time, ensuring consistent - # padding regardless of cache history. + # self-contained regardless of which audios it was + # co-processed with. The batched() field config does NOT + # re-pad ragged items at batch time (it yields a plain + # list), so _parse_and_validate_audio_input re-pads to the + # batch max via _pad_ragged_audio_features. masks = processed_outputs["input_features_mask"] unpadded_features = [ f[mask] @@ -1173,6 +1195,9 @@ def _parse_and_validate_audio_input( input_features_mask = kwargs.pop("input_features_mask", None) if input_features_mask is None: return None + input_features_padded, input_features_mask = _pad_ragged_audio_features( + input_features_padded, input_features_mask + ) return Gemma4AudioInputs( input_features_padded=input_features_padded, input_features_mask=input_features_mask, @@ -1487,8 +1512,8 @@ def _process_audio_input( self, audio_input: Gemma4AudioInputs, ) -> list[torch.Tensor]: - input_features = audio_input["input_features_padded"].squeeze(1) - input_features_mask = audio_input["input_features_mask"].squeeze(1) + input_features = audio_input["input_features_padded"] + input_features_mask = audio_input["input_features_mask"] # Run audio tower — mask convention: True=valid, False=padding. audio_outputs = self.audio_tower(input_features, input_features_mask) diff --git a/vllm/model_executor/models/gemma4_unified.py b/vllm/model_executor/models/gemma4_unified.py index 64c084b6161b..4535fd57730e 100644 --- a/vllm/model_executor/models/gemma4_unified.py +++ b/vllm/model_executor/models/gemma4_unified.py @@ -47,6 +47,7 @@ ) from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.processing.processor import PromptUpdateDetails from .utils import ( AutoWeightsLoader, @@ -183,6 +184,36 @@ def get_mm_max_tokens_per_item( tokens["video"] = num_frames * (_VIDEO_MAX_SOFT_TOKENS + 2 + 6) return tokens + def get_audio_repl( + self, + *, + audio_len: int, + processor: Gemma4UnifiedProcessor | None, + ) -> PromptUpdateDetails[list[int]]: + """Return the dynamic audio token sequence for this audio. + + The unified variant is encoder-free: the feature extractor pads + the waveform up to a multiple of ``audio_samples_per_token`` + (640) and emits one soft token per frame with no downsampling, + so the token count is ``ceil(audio_len / 640)`` — matching + ``Gemma4UnifiedProcessor.replace_audio_token`` (mask.sum()). + The parent's mel + conv subsampling arithmetic applies only to + the tower-based variant and undercounts by one whenever + ``audio_len % 640`` falls in [1, 160]. + """ + if processor is None: + processor = self.get_hf_processor() + + samples_per_token = processor.feature_extractor.audio_samples_per_token + num_tokens = math.ceil(audio_len / samples_per_token) + config = self.get_hf_config() + token_ids = ( + [config.boa_token_id] + + [processor.audio_token_id] * num_tokens + + [getattr(config, "eoa_token_id", config.eoa_token_index)] + ) + return PromptUpdateDetails.select_token_id(token_ids, processor.audio_token_id) + def _compute_num_soft_tokens( self, image_width: int, @@ -423,8 +454,8 @@ def _process_audio_input( No audio tower: the per-frame raw features are passed straight through the multimodal embedder, then padding is stripped. """ - input_features = audio_input["input_features_padded"].squeeze(1) - input_features_mask = audio_input["input_features_mask"].squeeze(1) + input_features = audio_input["input_features_padded"] + input_features_mask = audio_input["input_features_mask"] target_dtype = self.embed_audio.embedding_projection.weight.dtype audio_features = self.embed_audio(input_features.to(target_dtype))