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 14, 05:30 UTC
[Gemma3n] Fix audio batching (#24052)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Sep 2, 2025 at 16:23 UTC
0a74e9d0f2367cc121547aa8e21e13b04d4cad30
2 files changed
+63
-7
examples/online_serving/openai_chat_completion_client_for_multimodal.py
+42
@@ -266,10 +266,52 @@ def run_audio(model: str) -> None:
266
print("Chat completion output from base64 encoded audio:", result)
267
268
269
+def run_multi_audio(model: str) -> None:
270
+ from vllm.assets.audio import AudioAsset
271
+
272
+ # Two different audios to showcase batched inference.
273
+ audio_url = AudioAsset("winning_call").url
274
+ audio_base64 = encode_base64_content_from_url(audio_url)
275
+ audio_url2 = AudioAsset("azacinto_foscolo").url
276
+ audio_base64_2 = encode_base64_content_from_url(audio_url2)
277
+
278
+ # OpenAI-compatible schema (`input_audio`)
279
+ chat_completion_from_base64 = client.chat.completions.create(
280
+ messages=[
281
+ {
282
+ "role": "user",
283
+ "content": [
284
+ {"type": "text", "text": "Are these two audios the same?"},
285
+ {
286
+ "type": "input_audio",
287
+ "input_audio": {
288
+ "data": audio_base64,
289
+ "format": "wav",
290
+ },
291
+ },
292
+ {
293
+ "type": "input_audio",
294
+ "input_audio": {
295
+ "data": audio_base64_2,
296
+ "format": "wav",
297
+ },
298
+ },
299
+ ],
300
+ }
301
+ ],
302
+ model=model,
303
+ max_completion_tokens=64,
304
+ )
305
+
306
+ result = chat_completion_from_base64.choices[0].message.content
307
+ print("Chat completion output from input audio:", result)
308
+
309
+
310
example_function_map = {
311
"text-only": run_text_only,
312
"single-image": run_single_image,
313
"multi-image": run_multi_image,
314
+ "multi-audio": run_multi_audio,
315
"video": run_video,
316
"audio": run_audio,
317
}
vllm/model_executor/models/gemma3n_mm.py
+21
-7
@@ -5,6 +5,7 @@ from typing import Any, Literal, Optional, TypedDict, Union, cast
5
6
import numpy as np
7
import torch
8
+# yapf: disable
9
from torch import nn
10
from transformers import AutoModel, BatchFeature
11
from transformers.models.gemma3n import (Gemma3nAudioConfig,
@@ -30,7 +31,6 @@ from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
31
MultiModalKwargsItems)
32
from vllm.multimodal.parse import (ImageProcessorItems, MultiModalDataItems,
33
MultiModalDataParser)
33
-# yapf: disable
34
from vllm.multimodal.processing import (BaseMultiModalProcessor,
35
BaseProcessingInfo,
36
MultiModalPromptUpdates,
@@ -62,7 +62,8 @@ class Gemma3nImagePixelInputs(TypedDict):
62
63
64
class Gemma3nAudioInputs(TypedDict):
65
- input_features: torch.Tensor
65
+ input_features: Union[torch.Tensor, list[torch.Tensor]]
66
+ input_features_padded: torch.Tensor
67
"""Shape: `(batch_size * num_audio, seq_length, num_features)`"""
68
input_features_mask: torch.Tensor
69
"""Shape: `(batch_size * num_audio, seq_length)`"""
@@ -188,8 +189,13 @@ class Gemma3nMultiModalProcessor(BaseMultiModalProcessor[Gemma3nProcessingInfo]
189
mm_kwargs,
190
tok_kwargs,
191
)
192
+
193
if 'input_features' in processed_outputs:
192
- # Avoid padding since we need the output of each item to be
194
+ # Padding enables audio_tower to run in batched mode
195
+ processed_outputs["input_features_padded"] = \
196
+ processed_outputs["input_features"]
197
+
198
+ # Unpad features here since we need the output of each item to be
199
# independent of other items for the cache to work correctly
200
unpadded_features = [
201
f[mask] for f, mask in zip(
@@ -206,9 +212,11 @@ class Gemma3nMultiModalProcessor(BaseMultiModalProcessor[Gemma3nProcessingInfo]
212
hf_processor_mm_kwargs: Mapping[str, object],
213
) -> Mapping[str, MultiModalFieldConfig]:
214
209
- return dict(pixel_values=MultiModalFieldConfig.batched("image"),
210
- input_features=MultiModalFieldConfig.batched("audio"),
211
- input_features_mask=MultiModalFieldConfig.batched("audio"))
215
+ return dict(
216
+ pixel_values=MultiModalFieldConfig.batched("image"),
217
+ input_features=MultiModalFieldConfig.batched("audio"),
218
+ input_features_padded=MultiModalFieldConfig.batched("audio"),
219
+ input_features_mask=MultiModalFieldConfig.batched("audio"))
220
221
def _get_prompt_updates(
222
self,
@@ -516,9 +524,14 @@ class Gemma3nForConditionalGeneration(nn.Module, SupportsMultiModal,
524
if input_features_mask is None:
525
return None
526
527
+ input_features_padded = kwargs.pop("input_features_padded", None)
528
+ if input_features_padded is None:
529
+ return None
530
+
531
return Gemma3nAudioInputs(
532
input_features=input_features,
533
input_features_mask=input_features_mask,
534
+ input_features_padded=input_features_padded,
535
)
536
537
def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
@@ -564,7 +577,8 @@ class Gemma3nForConditionalGeneration(nn.Module, SupportsMultiModal,
577
audio_input: Gemma3nAudioInputs,
578
) -> list[torch.Tensor]:
579
assert self.audio_tower is not None
567
- input_features = audio_input["input_features"].squeeze(1)
580
+ # Run on padded features to enable batching
581
+ input_features = audio_input["input_features_padded"].squeeze(1)
582
input_features_mask = audio_input["input_features_mask"].squeeze(1)
583
audio_outputs, audio_mask = self.audio_tower(input_features,
584
~input_features_mask)