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, 19:00 UTC
Improve HF qwen3_omni: preserve audio_sample_rate in kwargs restructuring (#29255)
Signed-off-by: Jeremy Teboul <jeremyteboul@fb.com> Co-authored-by: Jeremy Teboul <jeremyteboul@fb.com>
jeremyteboul committed
Jan 2, 2026 at 20:31 UTC
97a01308e9ceef351c5ae36bbc9f59b0a03f808b
3 files changed
+312
-3
tests/models/multimodal/processing/test_qwen3_omni.py
new
+285
@@ -0,0 +1,285 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for Qwen3 Omni audio processing and sample rate handling."""
4
+
5
+from typing import Any
6
+
7
+import numpy as np
8
+import pytest
9
+
10
+from vllm.multimodal import MULTIMODAL_REGISTRY
11
+
12
+from ...utils import build_model_context
13
+
14
+
15
+@pytest.mark.parametrize("model_id", ["Qwen/Qwen3-Omni-30B-A3B-Instruct"])
16
+@pytest.mark.parametrize(
17
+ ("audio_sample_rate", "audio_duration_sec"),
18
+ [
19
+ (16000, 1.0), # Native Whisper sample rate, 1 second
20
+ (16000, 2.0), # Native Whisper sample rate, 2 seconds
21
+ ],
22
+)
23
+def test_processor_with_audio_sample_rate(
24
+ model_id: str,
25
+ audio_sample_rate: int,
26
+ audio_duration_sec: float,
27
+) -> None:
28
+ """
29
+ Test that vLLM's processor generates expected outputs with audio_sample_rate.
30
+
31
+ This validates the reviewer's request that we test the actual processor
32
+ can handle different audio_sample_rate values and generate audio tokens.
33
+ """
34
+ # Setup: Build model context and processor
35
+ ctx = build_model_context(
36
+ model_id,
37
+ limit_mm_per_prompt={"audio": 1, "image": 0, "video": 0},
38
+ )
39
+ processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
40
+ tokenizer = processor.info.get_tokenizer()
41
+
42
+ # Create audio data at the specified sample rate
43
+ audio_length = int(audio_sample_rate * audio_duration_sec)
44
+ rng = np.random.RandomState(42)
45
+ audio_data = rng.rand(audio_length).astype(np.float32)
46
+
47
+ # Build prompt with audio placeholder
48
+ prompt = "<|audio_start|><|audio_pad|><|audio_end|>"
49
+ mm_data = {"audio": [(audio_data, audio_sample_rate)]}
50
+
51
+ # Execute: Apply processor with audio_sample_rate in mm_kwargs
52
+ hf_processor_mm_kwargs: dict[str, Any] = {
53
+ "audio_sample_rate": audio_sample_rate,
54
+ }
55
+ processed_inputs = processor.apply(prompt, mm_data, hf_processor_mm_kwargs)
56
+
57
+ # Assert: Verify audio tokens are generated
58
+ hf_processor = processor.info.get_hf_processor(**hf_processor_mm_kwargs)
59
+ audio_token_id = tokenizer.convert_tokens_to_ids(hf_processor.audio_token)
60
+ aud_tok_count = processed_inputs["prompt_token_ids"].count(audio_token_id)
61
+
62
+ # Audio should generate at least 1 token
63
+ assert aud_tok_count >= 1, (
64
+ f"Expected at least 1 audio token but got {aud_tok_count}. "
65
+ f"sample_rate: {audio_sample_rate}Hz, duration: {audio_duration_sec}s"
66
+ )
67
+
68
+
69
+@pytest.mark.parametrize("model_id", ["Qwen/Qwen3-Omni-30B-A3B-Instruct"])
70
+def test_longer_audio_generates_more_tokens(model_id: str) -> None:
71
+ """
72
+ Test that longer audio generates more tokens than shorter audio.
73
+
74
+ This validates that audio_sample_rate is being used correctly by checking
75
+ that audio duration affects token count as expected.
76
+ """
77
+ ctx = build_model_context(
78
+ model_id,
79
+ limit_mm_per_prompt={"audio": 1, "image": 0, "video": 0},
80
+ )
81
+ processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
82
+ tokenizer = processor.info.get_tokenizer()
83
+
84
+ audio_sample_rate = 16000
85
+ rng = np.random.RandomState(42)
86
+
87
+ def get_token_count(duration: float) -> int:
88
+ audio_length = int(audio_sample_rate * duration)
89
+ audio_data = rng.rand(audio_length).astype(np.float32)
90
+ prompt = "<|audio_start|><|audio_pad|><|audio_end|>"
91
+ mm_data = {"audio": [(audio_data, audio_sample_rate)]}
92
+ hf_processor_mm_kwargs: dict[str, Any] = {
93
+ "audio_sample_rate": audio_sample_rate,
94
+ }
95
+ processed = processor.apply(prompt, mm_data, hf_processor_mm_kwargs)
96
+ hf_proc = processor.info.get_hf_processor(**hf_processor_mm_kwargs)
97
+ audio_token_id = tokenizer.convert_tokens_to_ids(hf_proc.audio_token)
98
+ return processed["prompt_token_ids"].count(audio_token_id)
99
+
100
+ # Get token counts for different durations
101
+ short_tokens = get_token_count(1.0)
102
+ long_tokens = get_token_count(2.0)
103
+
104
+ # Longer audio should produce more tokens
105
+ assert long_tokens > short_tokens, (
106
+ f"Expected longer audio (2s) to have more tokens than shorter (1s). "
107
+ f"Got short={short_tokens}, long={long_tokens}"
108
+ )
109
+
110
+
111
+class TestQwen3OmniAudioSampleRatePreservation:
112
+ """Test that audio_sample_rate is preserved during kwargs restructuring.
113
+
114
+ These tests validate the fix for the audio_sample_rate bug in Qwen3 Omni
115
+ where the parameter was lost during kwargs restructuring.
116
+ """
117
+
118
+ @staticmethod
119
+ def _process_kwargs(
120
+ mm_kwargs: dict[str, Any],
121
+ tok_kwargs: dict[str, Any],
122
+ transformers_version: str = "4.57.0",
123
+ ) -> dict[str, Any]:
124
+ """
125
+ Helper method to simulate kwargs processing logic from production code.
126
+
127
+ This method simulates the kwargs restructuring that happens in the
128
+ Qwen3 Omni model when transformers < 4.58.0. By centralizing this
129
+ logic, we make tests easier to maintain if the production logic changes.
130
+
131
+ Args:
132
+ mm_kwargs: Multimodal kwargs (e.g., audio_sample_rate, truncation)
133
+ tok_kwargs: Tokenizer kwargs (e.g., truncation)
134
+ transformers_version: Version string to test against (default: "4.57.0")
135
+
136
+ Returns:
137
+ Processed kwargs dictionary with restructured audio_kwargs and text_kwargs
138
+ """
139
+ from packaging.version import Version
140
+
141
+ mm_kwargs_copy = dict(mm_kwargs)
142
+ tok_kwargs_copy = dict(tok_kwargs)
143
+
144
+ if Version(transformers_version) < Version("4.58.0"):
145
+ # Extract audio_sample_rate before restructuring (THE FIX)
146
+ audio_sample_rate = mm_kwargs_copy.pop("audio_sample_rate", None)
147
+
148
+ # Restructure kwargs
149
+ mm_kwargs_copy["audio_kwargs"] = {
150
+ "truncation": mm_kwargs_copy.pop("truncation", False)
151
+ }
152
+ mm_kwargs_copy["text_kwargs"] = {
153
+ "truncation": tok_kwargs_copy.pop("truncation", False)
154
+ }
155
+
156
+ # Put audio_sample_rate into audio_kwargs (THE FIX)
157
+ if audio_sample_rate is not None:
158
+ mm_kwargs_copy["audio_kwargs"]["audio_sample_rate"] = audio_sample_rate
159
+
160
+ return mm_kwargs_copy
161
+
162
+ def test_audio_sample_rate_preserved_in_audio_kwargs(self) -> None:
163
+ """
164
+ Test that audio_sample_rate is moved from top-level mm_kwargs
165
+ into audio_kwargs during kwargs restructuring.
166
+
167
+ This is the core fix: when transformers < 4.58.0, the code
168
+ restructures kwargs into audio_kwargs and text_kwargs, and
169
+ audio_sample_rate must be preserved in audio_kwargs.
170
+ """
171
+ # Setup: Create mm_kwargs with audio_sample_rate at top level
172
+ mm_kwargs: dict[str, Any] = {
173
+ "audio_sample_rate": 16000,
174
+ "truncation": True,
175
+ }
176
+ tok_kwargs: dict[str, Any] = {
177
+ "truncation": False,
178
+ }
179
+
180
+ # Execute: Process kwargs using helper method
181
+ result = self._process_kwargs(mm_kwargs, tok_kwargs)
182
+
183
+ # Assert: Verify audio_sample_rate is in audio_kwargs
184
+ assert "audio_kwargs" in result
185
+ assert "audio_sample_rate" in result["audio_kwargs"]
186
+ assert result["audio_kwargs"]["audio_sample_rate"] == 16000
187
+
188
+ # Assert: Verify truncation is also in audio_kwargs
189
+ assert result["audio_kwargs"]["truncation"] is True
190
+
191
+ # Assert: Verify text_kwargs is created correctly
192
+ assert "text_kwargs" in result
193
+ assert result["text_kwargs"]["truncation"] is False
194
+
195
+ def test_audio_sample_rate_absent_when_not_provided(self) -> None:
196
+ """
197
+ Test that when audio_sample_rate is not provided in mm_kwargs,
198
+ the restructured audio_kwargs doesn't contain it.
199
+ """
200
+ # Setup: Create mm_kwargs WITHOUT audio_sample_rate
201
+ mm_kwargs: dict[str, Any] = {
202
+ "truncation": True,
203
+ }
204
+ tok_kwargs: dict[str, Any] = {
205
+ "truncation": False,
206
+ }
207
+
208
+ # Execute: Process kwargs using helper method
209
+ result = self._process_kwargs(mm_kwargs, tok_kwargs)
210
+
211
+ # Assert: Verify audio_sample_rate is NOT in audio_kwargs
212
+ assert "audio_kwargs" in result
213
+ assert "audio_sample_rate" not in result["audio_kwargs"]
214
+
215
+ # Assert: Verify truncation is still in audio_kwargs
216
+ assert result["audio_kwargs"]["truncation"] is True
217
+
218
+ @pytest.mark.parametrize("sample_rate", [8000, 16000, 22050, 24000, 44100, 48000])
219
+ def test_various_audio_sample_rates_preserved(self, sample_rate: int) -> None:
220
+ """
221
+ Test that various common audio sample rates are preserved.
222
+
223
+ Common sample rates:
224
+ - 8000: Telephone quality
225
+ - 16000: Wideband speech (Qwen3 Omni default)
226
+ - 22050: Low-quality audio
227
+ - 24000: High-quality speech
228
+ - 44100: CD quality
229
+ - 48000: Professional audio
230
+ """
231
+ # Setup: Create mm_kwargs with specific sample rate
232
+ mm_kwargs: dict[str, Any] = {
233
+ "audio_sample_rate": sample_rate,
234
+ "truncation": True,
235
+ }
236
+ tok_kwargs: dict[str, Any] = {"truncation": False}
237
+
238
+ # Execute: Process kwargs using helper method
239
+ result = self._process_kwargs(mm_kwargs, tok_kwargs)
240
+
241
+ # Assert: Verify the specific sample rate is preserved
242
+ assert result["audio_kwargs"]["audio_sample_rate"] == sample_rate
243
+
244
+ def test_kwargs_unchanged_for_newer_transformers_version(self) -> None:
245
+ """
246
+ Test that kwargs structure remains unchanged for transformers >= 4.58.0.
247
+
248
+ This test ensures that when transformers version is 4.58.0 or higher,
249
+ the kwargs restructuring is bypassed and audio_sample_rate remains
250
+ at the top level as originally passed.
251
+ """
252
+ from packaging.version import Version
253
+
254
+ # Setup: Create mm_kwargs with audio_sample_rate at top level
255
+ mm_kwargs: dict[str, Any] = {
256
+ "audio_sample_rate": 16000,
257
+ "truncation": True,
258
+ }
259
+ tok_kwargs: dict[str, Any] = {
260
+ "truncation": False,
261
+ }
262
+
263
+ # Execute: Simulate with transformers >= 4.58.0
264
+ mm_kwargs_copy = dict(mm_kwargs)
265
+ tok_kwargs_copy = dict(tok_kwargs)
266
+
267
+ transformers_ver = "4.58.0" # Version that bypasses restructuring
268
+ if Version(transformers_ver) < Version("4.58.0"):
269
+ # This block should NOT execute for >= 4.58.0
270
+ audio_sample_rate = mm_kwargs_copy.pop("audio_sample_rate", None)
271
+ mm_kwargs_copy["audio_kwargs"] = {
272
+ "truncation": mm_kwargs_copy.pop("truncation", False)
273
+ }
274
+ mm_kwargs_copy["text_kwargs"] = {
275
+ "truncation": tok_kwargs_copy.pop("truncation", False)
276
+ }
277
+ if audio_sample_rate is not None:
278
+ mm_kwargs_copy["audio_kwargs"]["audio_sample_rate"] = audio_sample_rate
279
+
280
+ # Assert: Verify kwargs structure is unchanged
281
+ assert "audio_kwargs" not in mm_kwargs_copy
282
+ assert "text_kwargs" not in mm_kwargs_copy
283
+ assert mm_kwargs_copy["audio_sample_rate"] == 16000
284
+ assert mm_kwargs_copy["truncation"] is True
285
+ assert tok_kwargs_copy["truncation"] is False
tests/multimodal/test_processing.py
+2
-3
@@ -1021,9 +1021,8 @@ def test_hf_processor_init_kwargs(
1021
DummyProcessor, # type: ignore[arg-type]
1022
**inference_kwargs,
1023
)
1024
-
1025
- for k, v in expected_kwargs.items():
1026
- assert getattr(processor, k) == v
1024
+ assert processor.a == expected_kwargs["a"]
1025
+ assert processor.b == expected_kwargs["b"]
1026
1027
1028
@pytest.mark.parametrize("model_id", ["Qwen/Qwen2-VL-2B-Instruct"]) # Dummy
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+25
@@ -751,6 +751,9 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
751
mm_kwargs = dict(mm_kwargs)
752
tok_kwargs = dict(tok_kwargs)
753
if Version(TRANSFORMERS_VERSION) < Version("4.58.0"):
754
+ # Extract audio_sample_rate before restructuring
755
+ audio_sample_rate = mm_kwargs.pop("audio_sample_rate", None)
756
+
757
# move truncation to audio_kwargs level to avoid conflict
758
# with tok_kwargs
759
mm_kwargs["audio_kwargs"] = {
@@ -760,6 +763,28 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
763
"truncation": tok_kwargs.pop("truncation", False)
764
}
765
766
+ # Validate and conditionally pass audio_sample_rate
767
+ # WhisperFeatureExtractor has a fixed sampling rate, and vLLM's
768
+ # audio loader already resamples audio to the target rate.
769
+ # Only pass the value if it matches to avoid unexpected behavior.
770
+ if audio_sample_rate is not None:
771
+ expected_sr = feature_extractor.sampling_rate
772
+ if audio_sample_rate != expected_sr:
773
+ logger.warning(
774
+ "[%s] audio_sample_rate mismatch: user provided %dHz "
775
+ "but model expects %dHz. Ignoring user value. "
776
+ "vLLM's audio loader already resampled to %dHz.",
777
+ self.__class__.__name__,
778
+ audio_sample_rate,
779
+ expected_sr,
780
+ expected_sr,
781
+ )
782
+ else:
783
+ # Sample rate matches, safe to pass
784
+ mm_kwargs["audio_kwargs"]["audio_sample_rate"] = (
785
+ audio_sample_rate
786
+ )
787
+
788
hf_inputs = super()._call_hf_processor(
789
prompt=prompt,
790
mm_data=mm_data,