From eb3dce975758f19c37f2cc892a6effd35a30a228 Mon Sep 17 00:00:00 2001 From: TrainToGPB <114007062+TrainToGPB@users.noreply.github.com> Date: Wed, 5 Aug 2026 13:46:00 +0900 Subject: [PATCH] [Bugfix][Model] Gemma3n/Gemma4: pad variable-length audio batches (#50958) Signed-off-by: TrainToGPB Co-authored-by: Claude --- .../generation/test_gemma_ragged_audio.py | 96 +++++++++++++++++++ vllm/model_executor/models/gemma3n_mm.py | 60 +++++++++++- vllm/model_executor/models/gemma4_mm.py | 13 ++- vllm/model_executor/models/gemma4_unified.py | 7 +- 4 files changed, 165 insertions(+), 11 deletions(-) create mode 100644 tests/models/multimodal/generation/test_gemma_ragged_audio.py diff --git a/tests/models/multimodal/generation/test_gemma_ragged_audio.py b/tests/models/multimodal/generation/test_gemma_ragged_audio.py new file mode 100644 index 0000000000..e51eb002ef --- /dev/null +++ b/tests/models/multimodal/generation/test_gemma_ragged_audio.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from tests.models.registry import HF_EXAMPLE_MODELS +from vllm.assets.audio import AudioAsset + +MODELS = { + "Gemma3nForConditionalGeneration": "google/gemma-3n-E2B-it", + "Gemma4ForConditionalGeneration": "google/gemma-4-E2B-it", +} + +# Seconds. Spread wide enough that no two clips can produce the same number of +# mel frames: equal-length clips stack normally and never reach the ragged path. +CLIP_DURATIONS = (1.0, 3.5, 7.0, 12.0) + + +@pytest.fixture(scope="module") +def ragged_clips(): + """One asset sliced to several durations, so only length varies.""" + audio, sample_rate = AudioAsset("mary_had_lamb").audio_and_sample_rate + + clips = [] + for duration in CLIP_DURATIONS: + num_samples = int(duration * sample_rate) + if num_samples > len(audio): + pytest.skip(f"asset is shorter than {duration}s") + clips.append((audio[:num_samples], sample_rate)) + + assert len({len(clip) for clip, _ in clips}) == len(clips) + return clips + + +@pytest.mark.core_model +@pytest.mark.parametrize("arch", sorted(MODELS)) +@pytest.mark.parametrize("max_tokens", [8]) +def test_variable_length_audio_batching( + vllm_runner, + ragged_clips, + arch: str, + max_tokens: int, +) -> None: + """Test batching of requests with different audio durations. + + Audio features are unpadded per item so a multimodal cache entry does not + depend on the batch it was first processed in, and + `MultiModalFieldConfig.batched` stacks items only when their shapes agree. + A batch of differently sized clips therefore reaches the model as a list, + which used to take EngineCore down with + + AttributeError: 'list' object has no attribute 'squeeze' + + failing every in-flight request with EngineDeadError. + """ + from transformers import AutoProcessor + + model_name = MODELS[arch] + model_info = HF_EXAMPLE_MODELS.get_hf_info(arch) + model_info.check_available_online(on_fail="skip") + model_info.check_transformers_version(on_fail="skip") + + # The two architectures spell the audio placeholder differently; ask the + # processor rather than hardcoding either. + processor = AutoProcessor.from_pretrained(model_name) + prompt = processor.apply_chat_template( + [ + { + "role": "user", + "content": [ + {"type": "audio", "audio": None}, + {"type": "text", "text": "Transcribe this audio."}, + ], + } + ], + tokenize=False, + add_generation_prompt=True, + ) + + with vllm_runner( + model_name, + dtype="bfloat16", + enforce_eager=True, + limit_mm_per_prompt={"audio": 1}, + ) as vllm_model: + # One call so the clips are scheduled into the same forward pass; + # issued separately they never reach the ragged path. + outputs = vllm_model.generate_greedy( + [prompt] * len(ragged_clips), + max_tokens, + audios=[[clip] for clip in ragged_clips], + ) + + assert len(outputs) == len(ragged_clips) + for output in outputs: + assert len(output[1]) > 0, "Expected non-empty output" diff --git a/vllm/model_executor/models/gemma3n_mm.py b/vllm/model_executor/models/gemma3n_mm.py index 1dd44313c1..7762ddd3a3 100644 --- a/vllm/model_executor/models/gemma3n_mm.py +++ b/vllm/model_executor/models/gemma3n_mm.py @@ -83,6 +83,52 @@ class Gemma3nImagePixelInputs(TensorSchema): pixel_values: Annotated[torch.Tensor, TensorShape("bn", 3, "h", "w")] +def batch_audio_features( + input_features: torch.Tensor | list[torch.Tensor], + input_features_mask: torch.Tensor | list[torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor]: + """Return mel features and their validity mask as batched tensors. + + Audio features are unpadded per item so that a multimodal cache entry does + not depend on the batch it was first processed in. + [`MultiModalFieldConfig.batched`][vllm.multimodal.inputs.MultiModalFieldConfig.batched] + stacks items only when their shapes agree, so a batch of clips with + differing durations reaches the model as a list and must be re-padded here. + + Padded frames are zero-filled and marked invalid. Audio towers consume the + mask, and callers keep only masked-in positions, so padding never reaches + the language model. + + Gemma4 shares this helper: its audio path has the same unpad/re-pad + contract and the same fields. + + Args: + input_features: `(bn, s, f)` tensor, or a list of `(s_i, f)` tensors + when clip durations differ. + input_features_mask: Matching `(bn, s)` tensor, or list of `(s_i,)` + tensors. `True` marks a valid frame. + + Returns: + The `(bn, s_max, f)` features and their `(bn, s_max)` mask. + """ + if isinstance(input_features, torch.Tensor): + return input_features.squeeze(1), input_features_mask.squeeze(1) + + max_len = max(features.shape[0] for features in input_features) + batched_features = input_features[0].new_zeros( + (len(input_features), max_len, input_features[0].shape[-1]) + ) + batched_mask = input_features_mask[0].new_zeros( + (len(input_features_mask), max_len), dtype=torch.bool + ) + for i, (features, mask) in enumerate( + zip(input_features, input_features_mask, strict=True) + ): + batched_features[i, : features.shape[0]] = features + batched_mask[i, : mask.shape[0]] = mask + return batched_features, batched_mask + + class Gemma3nAudioInputs(TensorSchema): """ Dimensions: @@ -92,8 +138,12 @@ class Gemma3nAudioInputs(TensorSchema): """ type: Literal["audio"] = "audio" - input_features_padded: Annotated[torch.Tensor, TensorShape("bn", "s", "f")] - input_features_mask: Annotated[torch.Tensor, TensorShape("bn", "s")] + input_features_padded: Annotated[ + torch.Tensor, TensorShape("bn", "s", "f", dynamic_dims={"s"}) + ] + input_features_mask: Annotated[ + torch.Tensor, TensorShape("bn", "s", dynamic_dims={"s"}) + ] Gemma3nImageInputs = Gemma3nImagePixelInputs @@ -615,8 +665,10 @@ class Gemma3nForConditionalGeneration( audio_input: Gemma3nAudioInputs, ) -> list[torch.Tensor]: # Run on padded features to enable batching - input_features = audio_input["input_features_padded"].squeeze(1) - input_features_mask = audio_input["input_features_mask"].squeeze(1) + input_features, input_features_mask = batch_audio_features( + audio_input["input_features_padded"], + audio_input["input_features_mask"], + ) audio_outputs = self.audio_tower(input_features, ~input_features_mask) audio_encodings = audio_outputs.last_hidden_state audio_mask = audio_outputs.audio_mel_mask diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index e158a78501..48136e7c2f 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -40,6 +40,7 @@ from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.models.gemma3n_mm import batch_audio_features from vllm.model_executor.models.gemma4 import ( _GEMMA4_EXPERT_PARENT_MAPPER, Gemma4ForCausalLM, @@ -750,9 +751,9 @@ class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]): 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. + # _get_mm_fields_config stacks equal-length items and leaves + # the rest as a list; _process_audio_input re-pads that list + # to the batch's max length. masks = processed_outputs["input_features_mask"] unpadded_features = [ f[mask] @@ -1487,8 +1488,10 @@ class Gemma4ForConditionalGeneration( 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, input_features_mask = batch_audio_features( + audio_input["input_features_padded"], + 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 64c084b616..f41a36adfd 100644 --- a/vllm/model_executor/models/gemma4_unified.py +++ b/vllm/model_executor/models/gemma4_unified.py @@ -30,6 +30,7 @@ from transformers.models.gemma4_unified.processing_gemma4_unified import ( from vllm.config import VllmConfig from vllm.config.multimodal import VideoDummyOptions from vllm.model_executor.layers.linear import ColumnParallelLinear +from vllm.model_executor.models.gemma3n_mm import batch_audio_features from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM from vllm.model_executor.models.gemma4_mm import ( _SUPPORTED_SOFT_TOKENS, @@ -423,8 +424,10 @@ class Gemma4UnifiedForConditionalGeneration(Gemma4ForConditionalGeneration): 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, input_features_mask = batch_audio_features( + audio_input["input_features_padded"], + audio_input["input_features_mask"], + ) target_dtype = self.embed_audio.embedding_projection.weight.dtype audio_features = self.embed_audio(input_features.to(target_dtype))