[Bugfix][Model] Gemma3n/Gemma4: pad variable-length audio batches (#50958)

Signed-off-by: TrainToGPB <kaiba0514@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>

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"

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

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)

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))