[Bugfix][Model] Gemma3n/Gemma4: pad variable-length audio batches (#50958)
Signed-off-by: TrainToGPB <kaiba0514@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
TrainToGPB committed
Aug 5, 2026 at 13:46 UTC
eb3dce975758f19c37f2cc892a6effd35a30a228
4 files changed
+165
-11
tests/models/multimodal/generation/test_gemma_ragged_audio.py
new
+96
@@ -0,0 +1,96 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from tests.models.registry import HF_EXAMPLE_MODELS
7
+from vllm.assets.audio import AudioAsset
8
+
9
+MODELS = {
10
+ "Gemma3nForConditionalGeneration": "google/gemma-3n-E2B-it",
11
+ "Gemma4ForConditionalGeneration": "google/gemma-4-E2B-it",
12
+}
13
+
14
+# Seconds. Spread wide enough that no two clips can produce the same number of
15
+# mel frames: equal-length clips stack normally and never reach the ragged path.
16
+CLIP_DURATIONS = (1.0, 3.5, 7.0, 12.0)
17
+
18
+
19
+@pytest.fixture(scope="module")
20
+def ragged_clips():
21
+ """One asset sliced to several durations, so only length varies."""
22
+ audio, sample_rate = AudioAsset("mary_had_lamb").audio_and_sample_rate
23
+
24
+ clips = []
25
+ for duration in CLIP_DURATIONS:
26
+ num_samples = int(duration * sample_rate)
27
+ if num_samples > len(audio):
28
+ pytest.skip(f"asset is shorter than {duration}s")
29
+ clips.append((audio[:num_samples], sample_rate))
30
+
31
+ assert len({len(clip) for clip, _ in clips}) == len(clips)
32
+ return clips
33
+
34
+
35
+@pytest.mark.core_model
36
+@pytest.mark.parametrize("arch", sorted(MODELS))
37
+@pytest.mark.parametrize("max_tokens", [8])
38
+def test_variable_length_audio_batching(
39
+ vllm_runner,
40
+ ragged_clips,
41
+ arch: str,
42
+ max_tokens: int,
43
+) -> None:
44
+ """Test batching of requests with different audio durations.
45
+
46
+ Audio features are unpadded per item so a multimodal cache entry does not
47
+ depend on the batch it was first processed in, and
48
+ `MultiModalFieldConfig.batched` stacks items only when their shapes agree.
49
+ A batch of differently sized clips therefore reaches the model as a list,
50
+ which used to take EngineCore down with
51
+
52
+ AttributeError: 'list' object has no attribute 'squeeze'
53
+
54
+ failing every in-flight request with EngineDeadError.
55
+ """
56
+ from transformers import AutoProcessor
57
+
58
+ model_name = MODELS[arch]
59
+ model_info = HF_EXAMPLE_MODELS.get_hf_info(arch)
60
+ model_info.check_available_online(on_fail="skip")
61
+ model_info.check_transformers_version(on_fail="skip")
62
+
63
+ # The two architectures spell the audio placeholder differently; ask the
64
+ # processor rather than hardcoding either.
65
+ processor = AutoProcessor.from_pretrained(model_name)
66
+ prompt = processor.apply_chat_template(
67
+ [
68
+ {
69
+ "role": "user",
70
+ "content": [
71
+ {"type": "audio", "audio": None},
72
+ {"type": "text", "text": "Transcribe this audio."},
73
+ ],
74
+ }
75
+ ],
76
+ tokenize=False,
77
+ add_generation_prompt=True,
78
+ )
79
+
80
+ with vllm_runner(
81
+ model_name,
82
+ dtype="bfloat16",
83
+ enforce_eager=True,
84
+ limit_mm_per_prompt={"audio": 1},
85
+ ) as vllm_model:
86
+ # One call so the clips are scheduled into the same forward pass;
87
+ # issued separately they never reach the ragged path.
88
+ outputs = vllm_model.generate_greedy(
89
+ [prompt] * len(ragged_clips),
90
+ max_tokens,
91
+ audios=[[clip] for clip in ragged_clips],
92
+ )
93
+
94
+ assert len(outputs) == len(ragged_clips)
95
+ for output in outputs:
96
+ assert len(output[1]) > 0, "Expected non-empty output"
vllm/model_executor/models/gemma3n_mm.py
+56
-4
@@ -83,6 +83,52 @@ class Gemma3nImagePixelInputs(TensorSchema):
83
pixel_values: Annotated[torch.Tensor, TensorShape("bn", 3, "h", "w")]
84
85
86
+def batch_audio_features(
87
+ input_features: torch.Tensor | list[torch.Tensor],
88
+ input_features_mask: torch.Tensor | list[torch.Tensor],
89
+) -> tuple[torch.Tensor, torch.Tensor]:
90
+ """Return mel features and their validity mask as batched tensors.
91
+
92
+ Audio features are unpadded per item so that a multimodal cache entry does
93
+ not depend on the batch it was first processed in.
94
+ [`MultiModalFieldConfig.batched`][vllm.multimodal.inputs.MultiModalFieldConfig.batched]
95
+ stacks items only when their shapes agree, so a batch of clips with
96
+ differing durations reaches the model as a list and must be re-padded here.
97
+
98
+ Padded frames are zero-filled and marked invalid. Audio towers consume the
99
+ mask, and callers keep only masked-in positions, so padding never reaches
100
+ the language model.
101
+
102
+ Gemma4 shares this helper: its audio path has the same unpad/re-pad
103
+ contract and the same fields.
104
+
105
+ Args:
106
+ input_features: `(bn, s, f)` tensor, or a list of `(s_i, f)` tensors
107
+ when clip durations differ.
108
+ input_features_mask: Matching `(bn, s)` tensor, or list of `(s_i,)`
109
+ tensors. `True` marks a valid frame.
110
+
111
+ Returns:
112
+ The `(bn, s_max, f)` features and their `(bn, s_max)` mask.
113
+ """
114
+ if isinstance(input_features, torch.Tensor):
115
+ return input_features.squeeze(1), input_features_mask.squeeze(1)
116
+
117
+ max_len = max(features.shape[0] for features in input_features)
118
+ batched_features = input_features[0].new_zeros(
119
+ (len(input_features), max_len, input_features[0].shape[-1])
120
+ )
121
+ batched_mask = input_features_mask[0].new_zeros(
122
+ (len(input_features_mask), max_len), dtype=torch.bool
123
+ )
124
+ for i, (features, mask) in enumerate(
125
+ zip(input_features, input_features_mask, strict=True)
126
+ ):
127
+ batched_features[i, : features.shape[0]] = features
128
+ batched_mask[i, : mask.shape[0]] = mask
129
+ return batched_features, batched_mask
130
+
131
+
132
class Gemma3nAudioInputs(TensorSchema):
133
"""
134
Dimensions:
@@ -92,8 +138,12 @@ class Gemma3nAudioInputs(TensorSchema):
138
"""
139
140
type: Literal["audio"] = "audio"
95
- input_features_padded: Annotated[torch.Tensor, TensorShape("bn", "s", "f")]
96
- input_features_mask: Annotated[torch.Tensor, TensorShape("bn", "s")]
141
+ input_features_padded: Annotated[
142
+ torch.Tensor, TensorShape("bn", "s", "f", dynamic_dims={"s"})
143
+ ]
144
+ input_features_mask: Annotated[
145
+ torch.Tensor, TensorShape("bn", "s", dynamic_dims={"s"})
146
+ ]
147
148
149
Gemma3nImageInputs = Gemma3nImagePixelInputs
@@ -615,8 +665,10 @@ class Gemma3nForConditionalGeneration(
665
audio_input: Gemma3nAudioInputs,
666
) -> list[torch.Tensor]:
667
# Run on padded features to enable batching
618
- input_features = audio_input["input_features_padded"].squeeze(1)
619
- input_features_mask = audio_input["input_features_mask"].squeeze(1)
668
+ input_features, input_features_mask = batch_audio_features(
669
+ audio_input["input_features_padded"],
670
+ audio_input["input_features_mask"],
671
+ )
672
audio_outputs = self.audio_tower(input_features, ~input_features_mask)
673
audio_encodings = audio_outputs.last_hidden_state
674
audio_mask = audio_outputs.audio_mel_mask
vllm/model_executor/models/gemma4_mm.py
+8
-5
@@ -40,6 +40,7 @@ from vllm.inputs import MultiModalDataDict
40
from vllm.logger import init_logger
41
from vllm.model_executor.layers.layernorm import RMSNorm
42
from vllm.model_executor.layers.linear import ReplicatedLinear
43
+from vllm.model_executor.models.gemma3n_mm import batch_audio_features
44
from vllm.model_executor.models.gemma4 import (
45
_GEMMA4_EXPERT_PARENT_MAPPER,
46
Gemma4ForCausalLM,
@@ -750,9 +751,9 @@ class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]):
751
if "input_features" in processed_outputs:
752
# Unpad per-item so each item's cache entry is
753
# self-contained. The batched() field config in
753
- # _get_mm_fields_config will re-pad all fields to the
754
- # batch's max length at batch time, ensuring consistent
755
- # padding regardless of cache history.
754
+ # _get_mm_fields_config stacks equal-length items and leaves
755
+ # the rest as a list; _process_audio_input re-pads that list
756
+ # to the batch's max length.
757
masks = processed_outputs["input_features_mask"]
758
unpadded_features = [
759
f[mask]
@@ -1487,8 +1488,10 @@ class Gemma4ForConditionalGeneration(
1488
self,
1489
audio_input: Gemma4AudioInputs,
1490
) -> list[torch.Tensor]:
1490
- input_features = audio_input["input_features_padded"].squeeze(1)
1491
- input_features_mask = audio_input["input_features_mask"].squeeze(1)
1491
+ input_features, input_features_mask = batch_audio_features(
1492
+ audio_input["input_features_padded"],
1493
+ audio_input["input_features_mask"],
1494
+ )
1495
1496
# Run audio tower — mask convention: True=valid, False=padding.
1497
audio_outputs = self.audio_tower(input_features, input_features_mask)
vllm/model_executor/models/gemma4_unified.py
+5
-2
@@ -30,6 +30,7 @@ from transformers.models.gemma4_unified.processing_gemma4_unified import (
30
from vllm.config import VllmConfig
31
from vllm.config.multimodal import VideoDummyOptions
32
from vllm.model_executor.layers.linear import ColumnParallelLinear
33
+from vllm.model_executor.models.gemma3n_mm import batch_audio_features
34
from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
35
from vllm.model_executor.models.gemma4_mm import (
36
_SUPPORTED_SOFT_TOKENS,
@@ -423,8 +424,10 @@ class Gemma4UnifiedForConditionalGeneration(Gemma4ForConditionalGeneration):
424
No audio tower: the per-frame raw features are passed straight
425
through the multimodal embedder, then padding is stripped.
426
"""
426
- input_features = audio_input["input_features_padded"].squeeze(1)
427
- input_features_mask = audio_input["input_features_mask"].squeeze(1)
427
+ input_features, input_features_mask = batch_audio_features(
428
+ audio_input["input_features_padded"],
429
+ audio_input["input_features_mask"],
430
+ )
431
432
target_dtype = self.embed_audio.embedding_projection.weight.dtype
433
audio_features = self.embed_audio(input_features.to(target_dtype))