A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 09:00 UTC
[Bugfix] Fix integer overflow in Gemma3n audio processing (#31657)
Signed-off-by: Jeremy Teboul <jeremyte@meta.com>
Jeremy Teboul committed
Jan 10, 2026 at 01:52 UTC
07286ec5a6edf3a7a3623ddfc8db6a24b1d3c70a
3 files changed
+239
-32
tests/models/multimodal/processing/test_gemma3.py
+141
-1
@@ -2,14 +2,154 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import pytest
5
+import torch
6
7
+from vllm.model_executor.models.gemma3n_audio_utils import (
8
+ adjust_audio_features_to_expected_length,
9
+)
10
from vllm.multimodal import MULTIMODAL_REGISTRY
11
12
from ....conftest import ImageTestAssets
13
from ...utils import build_model_context
14
15
+# Gemma3 (image) model
16
+GEMMA3_MODEL_ID = "google/gemma-3-4b-it"
17
12
-@pytest.mark.parametrize("model_id", ["google/gemma-3-4b-it"])
18
+# Gemma3n (multimodal with audio) model
19
+GEMMA3N_MODEL_ID = "google/gemma-3n-E2B-it"
20
+
21
+# Expected audio tokens for Gemma3n (audio_soft_tokens_per_image)
22
+GEMMA3N_EXPECTED_AUDIO_TOKENS = 188
23
+
24
+
25
+class TestGemma3nAudioTensorLogic:
26
+ """CPU-based tests for Gemma3n audio feature tensor manipulation.
27
+
28
+ These tests validate the padding/truncation logic in
29
+ adjust_audio_features_to_expected_length() which fixes the
30
+ integer overflow in _process_audio_input when audio_seq_len > 188.
31
+ """
32
+
33
+ def test_padding_when_audio_short(self):
34
+ """Test that short audio is padded to expected length."""
35
+ batch_size, seq_len, embed_dim = 1, 100, 256
36
+ expected_tokens = GEMMA3N_EXPECTED_AUDIO_TOKENS
37
+
38
+ audio_features = torch.randn(batch_size, seq_len, embed_dim)
39
+ padding_embs = torch.zeros(1, 1, embed_dim)
40
+
41
+ result, tokens_truncated = adjust_audio_features_to_expected_length(
42
+ audio_features, expected_tokens, padding_embs
43
+ )
44
+
45
+ assert result.shape == (batch_size, expected_tokens, embed_dim)
46
+ assert tokens_truncated == 0
47
+ # First 100 tokens should be original, rest should be padding (zeros)
48
+ assert torch.allclose(result[:, :seq_len, :], audio_features)
49
+ assert torch.allclose(
50
+ result[:, seq_len:, :],
51
+ torch.zeros(batch_size, expected_tokens - seq_len, embed_dim),
52
+ )
53
+
54
+ def test_truncation_when_audio_long(self):
55
+ """Test that long audio is truncated to expected length.
56
+
57
+ This is the key test for the overflow fix. Previously, when
58
+ audio_seq_len > expected_tokens, the code would compute a negative
59
+ padding value causing: RuntimeError: numel: integer multiplication overflow
60
+ """
61
+ batch_size, seq_len, embed_dim = 1, 192, 256 # 192 > 188
62
+ expected_tokens = GEMMA3N_EXPECTED_AUDIO_TOKENS
63
+
64
+ audio_features = torch.randn(batch_size, seq_len, embed_dim)
65
+ padding_embs = torch.zeros(1, 1, embed_dim)
66
+
67
+ result, tokens_truncated = adjust_audio_features_to_expected_length(
68
+ audio_features, expected_tokens, padding_embs
69
+ )
70
+
71
+ assert result.shape == (batch_size, expected_tokens, embed_dim)
72
+ assert tokens_truncated == seq_len - expected_tokens # 192 - 188 = 4
73
+ # Result should be first 188 tokens of original
74
+ assert torch.allclose(result, audio_features[:, :expected_tokens, :])
75
+
76
+ def test_no_change_when_exact_length(self):
77
+ """Test that exact-length audio passes through unchanged."""
78
+ batch_size, embed_dim = 1, 256
79
+ expected_tokens = GEMMA3N_EXPECTED_AUDIO_TOKENS
80
+
81
+ audio_features = torch.randn(batch_size, expected_tokens, embed_dim)
82
+ padding_embs = torch.zeros(1, 1, embed_dim)
83
+
84
+ result, tokens_truncated = adjust_audio_features_to_expected_length(
85
+ audio_features, expected_tokens, padding_embs
86
+ )
87
+
88
+ assert result.shape == audio_features.shape
89
+ assert tokens_truncated == 0
90
+ assert torch.allclose(result, audio_features)
91
+
92
+ def test_original_bug_would_fail(self):
93
+ """Verify the original buggy implementation would cause overflow.
94
+
95
+ The original code always tried to pad, which fails when
96
+ audio_seq_len > expected_tokens because expand() gets negative size.
97
+ """
98
+ batch_size, seq_len, embed_dim = 1, 192, 256
99
+ expected_tokens = GEMMA3N_EXPECTED_AUDIO_TOKENS
100
+
101
+ padding_embs = torch.zeros(1, 1, embed_dim)
102
+
103
+ # Original buggy logic (always pads, never truncates)
104
+ extra_padding_tokens = expected_tokens - seq_len # = -4 (negative!)
105
+
106
+ with pytest.raises(RuntimeError):
107
+ # This should fail with negative size error
108
+ padding_embs.expand(batch_size, extra_padding_tokens, embed_dim)
109
+
110
+ @pytest.mark.parametrize(
111
+ "seq_len",
112
+ [50, 100, 150, 187, 188, 189, 192, 200, 300],
113
+ )
114
+ def test_various_audio_lengths(self, seq_len: int):
115
+ """Test padding/truncation with various audio lengths."""
116
+ batch_size, embed_dim = 1, 256
117
+ expected_tokens = GEMMA3N_EXPECTED_AUDIO_TOKENS
118
+
119
+ audio_features = torch.randn(batch_size, seq_len, embed_dim)
120
+ padding_embs = torch.zeros(1, 1, embed_dim)
121
+
122
+ # Should not raise any errors
123
+ result, tokens_truncated = adjust_audio_features_to_expected_length(
124
+ audio_features, expected_tokens, padding_embs
125
+ )
126
+
127
+ # Output should always be expected_tokens length
128
+ assert result.shape == (batch_size, expected_tokens, embed_dim)
129
+
130
+ # Verify truncation count is correct
131
+ if seq_len > expected_tokens:
132
+ assert tokens_truncated == seq_len - expected_tokens
133
+ else:
134
+ assert tokens_truncated == 0
135
+
136
+ def test_batch_processing(self):
137
+ """Test that batch processing works correctly."""
138
+ batch_size, seq_len, embed_dim = 4, 192, 256
139
+ expected_tokens = GEMMA3N_EXPECTED_AUDIO_TOKENS
140
+
141
+ audio_features = torch.randn(batch_size, seq_len, embed_dim)
142
+ padding_embs = torch.zeros(1, 1, embed_dim)
143
+
144
+ result, tokens_truncated = adjust_audio_features_to_expected_length(
145
+ audio_features, expected_tokens, padding_embs
146
+ )
147
+
148
+ assert result.shape == (batch_size, expected_tokens, embed_dim)
149
+ assert tokens_truncated == seq_len - expected_tokens
150
+
151
+
152
+@pytest.mark.parametrize("model_id", [GEMMA3_MODEL_ID])
153
def test_get_image_size_with_most_features(
154
image_assets: ImageTestAssets, model_id: str
155
):
vllm/model_executor/models/gemma3n_audio_utils.py
new
+57
@@ -0,0 +1,57 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Lightweight utility functions for Gemma3n audio processing.
5
+
6
+This module is separate from gemma3n_mm.py to avoid heavy CUDA dependencies,
7
+making it testable without a full vLLM build.
8
+"""
9
+
10
+import torch
11
+
12
+
13
+def adjust_audio_features_to_expected_length(
14
+ audio_features: torch.Tensor,
15
+ expected_tokens: int,
16
+ audio_padding_embs: torch.Tensor,
17
+) -> tuple[torch.Tensor, int]:
18
+ """Adjust audio features to expected token length via padding or truncation.
19
+
20
+ The Gemma3nProcessor expects all audio will be ~30s in length and inserts
21
+ a fixed number of audio soft tokens into the text. However, the audio
22
+ preprocessing and encoder do not guarantee they will produce exactly that
23
+ many soft tokens; they may produce fewer tokens (for shorter audio) or more
24
+ tokens (for longer audio or due to BOA/EOA special tokens).
25
+
26
+ This function handles both cases:
27
+ - If fewer tokens: pad with the provided padding embeddings
28
+ - If more tokens: truncate to the expected count
29
+
30
+ Args:
31
+ audio_features: Audio embeddings tensor of shape
32
+ (batch_size, seq_len, embed_dim)
33
+ expected_tokens: The expected number of audio tokens (e.g., 188)
34
+ audio_padding_embs: Padding embeddings tensor of shape (1, 1, embed_dim)
35
+
36
+ Returns:
37
+ Tuple of:
38
+ - adjusted_features: Audio features adjusted to expected_tokens length
39
+ - tokens_truncated: Number of tokens truncated (0 if padding was applied)
40
+ """
41
+ audio_batch_size, audio_seq_len, audio_embed_dim = audio_features.shape
42
+ tokens_truncated = 0
43
+
44
+ if audio_seq_len < expected_tokens:
45
+ # Pad to expected length with padding embeddings
46
+ extra_padding_tokens = expected_tokens - audio_seq_len
47
+ extra_padding_features = audio_padding_embs.expand(
48
+ audio_batch_size, extra_padding_tokens, audio_embed_dim
49
+ )
50
+ audio_features = torch.cat((audio_features, extra_padding_features), dim=1)
51
+ elif audio_seq_len > expected_tokens:
52
+ # Truncate to expected length (audio encoder produced more tokens
53
+ # than expected, e.g., due to longer audio or placeholder mismatch)
54
+ tokens_truncated = audio_seq_len - expected_tokens
55
+ audio_features = audio_features[:, :expected_tokens, :]
56
+
57
+ return audio_features, tokens_truncated
vllm/model_executor/models/gemma3n_mm.py
+41
-31
@@ -1,11 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Iterable, Mapping, Sequence
4
-from typing import Annotated, Any, Literal, Optional, Union, cast
4
+from typing import Annotated, Any, Literal, cast
5
6
import numpy as np
7
import torch
8
-
8
from torch import nn
9
from transformers import AutoModel, BatchFeature
10
from transformers.models.gemma3n import (
@@ -26,6 +25,9 @@ from vllm.model_executor.layers.layernorm import RMSNorm
25
from vllm.model_executor.layers.linear import RowParallelLinear
26
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
27
from vllm.model_executor.models.gemma3n import Gemma3nForCausalLM
28
+from vllm.model_executor.models.gemma3n_audio_utils import (
29
+ adjust_audio_features_to_expected_length,
30
+)
31
from vllm.model_executor.models.module_mapping import MultiModelKeys
32
from vllm.model_executor.models.whisper import ISO639_1_SUPPORTED_LANGS
33
from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -105,12 +107,12 @@ class Gemma3nProcessingInfo(BaseProcessingInfo):
107
def get_hf_processor(self, **kwargs: object):
108
return self.ctx.get_hf_processor(Gemma3nProcessor, **kwargs)
109
108
- def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
110
+ def get_supported_mm_limits(self) -> Mapping[str, int | None]:
111
return {"image": None, "audio": None}
112
113
def get_max_tokens_per_item(
114
self, seq_len: int, mm_counts: Mapping[str, int]
113
- ) -> Optional[Mapping[str, int]]:
115
+ ) -> Mapping[str, int] | None:
116
return {"image": TOKENS_PER_IMAGE, "audio": TOKENS_PER_AUDIO}
117
118
def get_image_repl(
@@ -118,7 +120,7 @@ class Gemma3nProcessingInfo(BaseProcessingInfo):
120
*,
121
image_width: int,
122
image_height: int,
121
- processor: Optional[Gemma3nProcessor],
123
+ processor: Gemma3nProcessor | None,
124
) -> str:
125
"""
126
Get the replacement text for image tokens.
@@ -136,7 +138,7 @@ class Gemma3nProcessingInfo(BaseProcessingInfo):
138
def get_audio_repl(
139
self,
140
*,
139
- processor: Optional[Gemma3nProcessor],
141
+ processor: Gemma3nProcessor | None,
142
) -> str:
143
"""
144
Get the replacement text for audio tokens.
@@ -168,7 +170,7 @@ class Gemma3nDummyInputsBuilder(BaseDummyInputsBuilder[Gemma3nProcessingInfo]):
170
self,
171
seq_len: int,
172
mm_counts: Mapping[str, int],
171
- mm_options: Optional[Mapping[str, BaseDummyOptions]] = None,
173
+ mm_options: Mapping[str, BaseDummyOptions] | None = None,
174
) -> MultiModalDataDict:
175
num_images = mm_counts.get("image", 0)
176
num_audios = mm_counts.get("audio", 0)
@@ -387,7 +389,7 @@ class Gemma3nMultimodalEmbedder(nn.Module):
389
390
def __init__(
391
self,
390
- multimodal_config: Union[Gemma3nAudioConfig, Gemma3nVisionConfig],
392
+ multimodal_config: Gemma3nAudioConfig | Gemma3nVisionConfig,
393
text_config: Gemma3nTextConfig,
394
):
395
super().__init__()
@@ -427,8 +429,8 @@ class Gemma3nMultimodalEmbedder(nn.Module):
429
430
def forward(
431
self,
430
- input_ids: Optional[torch.LongTensor] = None,
431
- inputs_embeds: Optional[torch.Tensor] = None,
432
+ input_ids: torch.LongTensor | None = None,
433
+ inputs_embeds: torch.Tensor | None = None,
434
) -> torch.Tensor:
435
"""Embeds token ids or soft tokens for multimodal content into language model space.
436
@@ -529,7 +531,7 @@ class Gemma3nForConditionalGeneration(
531
532
def _parse_and_validate_image_input(
533
self, **kwargs: object
532
- ) -> Optional[Gemma3nImageInputs]:
534
+ ) -> Gemma3nImageInputs | None:
535
pixel_values = kwargs.pop("pixel_values", None)
536
image_embeds = kwargs.pop("image_embeds", None)
537
# TODO is this the case?
@@ -541,7 +543,7 @@ class Gemma3nForConditionalGeneration(
543
544
def _parse_and_validate_audio_input(
545
self, **kwargs: object
544
- ) -> Optional[Gemma3nAudioInputs]:
546
+ ) -> Gemma3nAudioInputs | None:
547
input_features_padded = kwargs.pop("input_features_padded", None)
548
if input_features_padded is None:
549
return None
@@ -616,12 +618,15 @@ class Gemma3nForConditionalGeneration(
618
)
619
audio_features = self.embed_audio(inputs_embeds=audio_outputs)
620
619
- # ruff: noqa
620
- # The Gemma3nProcessor expects all audio will be 30s in length and inserts 188 audio soft tokens into the
621
- # text to account for this. However, the audio preprocessing and encoder do not guarantee they will
622
- # produce 188 soft tokens; they will produce at most that many tokens, but they may produce fewer tokens
623
- # depending on the length of the longest audio input in the batch. When we encounter this situation, we pad
624
- # the audio feature out to 188 soft tokens with the embedding of the last token in the embed_audio vocab.
621
+ # The Gemma3nProcessor expects all audio will be 30s in length and
622
+ # inserts 188 audio soft tokens into the text to account for this.
623
+ # However, the audio preprocessing and encoder do not guarantee they
624
+ # will produce exactly 188 soft tokens; they may produce fewer tokens
625
+ # (for shorter audio) or more tokens (for longer audio or due to
626
+ # BOA/EOA special tokens in the placeholder sequence).
627
+ # We handle both cases:
628
+ # - If fewer tokens: pad with the embedding of the last vocab token
629
+ # - If more tokens: truncate to the expected count
630
# TODO precompute and cache padding
631
audio_padding_toks = torch.tensor(
632
[[self.vocab_size - 1]], dtype=torch.long, device=audio_features.device
@@ -631,13 +636,18 @@ class Gemma3nForConditionalGeneration(
636
audio_mask.unsqueeze(-1), audio_padding_embs, audio_features
637
)
638
634
- audio_batch_size, audio_seq_len, audio_embed_dim = audio_features.shape
635
- extra_padding_tokens = self.config.audio_soft_tokens_per_image - audio_seq_len # noqa: E501
636
- extra_padding_features = audio_padding_embs.expand(
637
- audio_batch_size, extra_padding_tokens, audio_embed_dim
639
+ expected_tokens = self.config.audio_soft_tokens_per_image
640
+ audio_features, tokens_truncated = adjust_audio_features_to_expected_length(
641
+ audio_features, expected_tokens, audio_padding_embs
642
)
643
+ if tokens_truncated > 0:
644
+ logger.warning(
645
+ "Gemma3n audio encoder produced %d extra tokens. "
646
+ "Truncating to match placeholder count of %d.",
647
+ tokens_truncated,
648
+ expected_tokens,
649
+ )
650
640
- audio_features = torch.cat((audio_features, extra_padding_features), dim=1)
651
# Return a list of embeddings instead of a batched tensor
652
return audio_features.unbind(0)
653
@@ -666,9 +676,9 @@ class Gemma3nForConditionalGeneration(
676
def embed_input_ids(
677
self,
678
input_ids: torch.Tensor,
669
- multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
679
+ multimodal_embeddings: MultiModalEmbeddings | None = None,
680
*,
671
- is_multimodal: Optional[torch.Tensor] = None,
681
+ is_multimodal: torch.Tensor | None = None,
682
handle_oov_mm_token: bool = False,
683
) -> torch.Tensor:
684
# NOTE (NickLucche) Each pass needs tokens to compute PLE so we cache
@@ -701,8 +711,8 @@ class Gemma3nForConditionalGeneration(
711
self,
712
input_ids: torch.Tensor,
713
positions: torch.Tensor,
704
- intermediate_tensors: Optional[IntermediateTensors] = None,
705
- inputs_embeds: Optional[torch.Tensor] = None,
714
+ intermediate_tensors: IntermediateTensors | None = None,
715
+ inputs_embeds: torch.Tensor | None = None,
716
**kwargs: object,
717
) -> IntermediateTensors:
718
if intermediate_tensors is not None:
@@ -729,7 +739,7 @@ class Gemma3nForConditionalGeneration(
739
def compute_logits(
740
self,
741
hidden_states: torch.Tensor,
732
- ) -> Optional[torch.Tensor]:
742
+ ) -> torch.Tensor | None:
743
return self.language_model.compute_logits(hidden_states)
744
745
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -747,7 +757,7 @@ class Gemma3nForConditionalGeneration(
757
)
758
759
@classmethod
750
- def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:
760
+ def get_placeholder_str(cls, modality: str, i: int) -> str | None:
761
if modality == "image":
762
return "<image_soft_token>"
763
elif modality == "audio":
@@ -761,10 +771,10 @@ class Gemma3nForConditionalGeneration(
771
audio: np.ndarray,
772
stt_config: SpeechToTextConfig,
773
model_config: ModelConfig,
764
- language: Optional[str],
774
+ language: str | None,
775
task_type: Literal["transcribe", "translate"],
776
request_prompt: str,
767
- to_language: Optional[str],
777
+ to_language: str | None,
778
) -> PromptType:
779
"""
780
Gemma3n supports "free-form" transcription.