Fix gemma4 unified multi audio

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

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)

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