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 7, 16:30 UTC
[Bugfix] Fix Qwen2.5-Omni and Qwen3-Omni mixed-modality embed regression (#35368)
Signed-off-by: linyueqian <linyueqian@outlook.com>
Yueqian Lin committed
Feb 26, 2026 at 06:58 UTC
c0615a296d44ce1963d795ea65dcff6172b4ae8d
3 files changed
+379
-21
tests/models/multimodal/processing/test_qwen2_5_omni_embed.py
new
+358
@@ -0,0 +1,358 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Unit tests for Qwen2.5-Omni embed_input_ids to verify embeddings are
5
+correctly assigned to audio/image/video token positions.
6
+
7
+Regression test for: https://github.com/vllm-project/vllm/issues/34506
8
+ - Non-interleaved mixed modalities (audio + image + video) should correctly
9
+ assign audio embeddings to audio positions, image to image, video to video.
10
+ - Interleaved (use_audio_in_video) should also work correctly.
11
+"""
12
+
13
+from unittest.mock import Mock
14
+
15
+import pytest
16
+import torch
17
+
18
+from vllm.model_executor.models.qwen2_5_omni_thinker import (
19
+ check_interleaved_audio_video,
20
+ merge_interleaved_embeddings,
21
+)
22
+
23
+# Fake token IDs
24
+AUDIO_TOKEN_ID = 1001
25
+IMAGE_TOKEN_ID = 1002
26
+VIDEO_TOKEN_ID = 1003
27
+TEXT_TOKEN_ID = 0
28
+
29
+
30
+# ---------------------------------------------------------------------------
31
+# Helpers
32
+# ---------------------------------------------------------------------------
33
+
34
+
35
+def make_token_seq(
36
+ audio_n: int, image_n: int, video_n: int, text_prefix: int = 3, text_sep: int = 2
37
+):
38
+ """
39
+ Build a flat token sequence:
40
+ [text_prefix] [AUDIO * audio_n] [text_sep] [IMAGE * image_n]
41
+ [text_sep] [VIDEO * video_n] [text_sep]
42
+ Returns (input_ids tensor, is_multimodal mask, positions dict).
43
+ """
44
+ tokens = (
45
+ [TEXT_TOKEN_ID] * text_prefix
46
+ + [AUDIO_TOKEN_ID] * audio_n
47
+ + [TEXT_TOKEN_ID] * text_sep
48
+ + [IMAGE_TOKEN_ID] * image_n
49
+ + [TEXT_TOKEN_ID] * text_sep
50
+ + [VIDEO_TOKEN_ID] * video_n
51
+ + [TEXT_TOKEN_ID] * text_sep
52
+ )
53
+ input_ids = torch.tensor(tokens)
54
+ is_multimodal = (
55
+ (input_ids == AUDIO_TOKEN_ID)
56
+ | (input_ids == IMAGE_TOKEN_ID)
57
+ | (input_ids == VIDEO_TOKEN_ID)
58
+ )
59
+ return input_ids, is_multimodal
60
+
61
+
62
+def make_interleaved_seq(
63
+ video_chunks: list[int], audio_chunks: list[int], text_prefix: int = 2
64
+):
65
+ """
66
+ Build an interleaved sequence like use_audio_in_video:
67
+ [text] [V*v0] [A*a0] [V*v1] [A*a1] ...
68
+ """
69
+ tokens = [TEXT_TOKEN_ID] * text_prefix
70
+ for v, a in zip(video_chunks, audio_chunks):
71
+ tokens += [VIDEO_TOKEN_ID] * v + [AUDIO_TOKEN_ID] * a
72
+ input_ids = torch.tensor(tokens)
73
+ is_multimodal = (input_ids == VIDEO_TOKEN_ID) | (input_ids == AUDIO_TOKEN_ID)
74
+ return input_ids, is_multimodal
75
+
76
+
77
+# ---------------------------------------------------------------------------
78
+# Tests for check_interleaved_audio_video
79
+# ---------------------------------------------------------------------------
80
+
81
+
82
+class TestCheckInterleavedAudioVideo:
83
+ def test_non_interleaved_audio_then_video(self):
84
+ """Audio entirely before video → not interleaved."""
85
+ input_ids, is_multimodal = make_token_seq(5, 0, 4)
86
+ is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
87
+ is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
88
+ assert not check_interleaved_audio_video(
89
+ is_video, is_audio, is_video.sum().item(), is_audio.sum().item()
90
+ )
91
+
92
+ def test_non_interleaved_with_image(self):
93
+ """Audio + image + video (the mixed_modalities case) → not interleaved."""
94
+ input_ids, is_multimodal = make_token_seq(5, 4, 6)
95
+ is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
96
+ is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
97
+ assert not check_interleaved_audio_video(
98
+ is_video, is_audio, is_video.sum().item(), is_audio.sum().item()
99
+ )
100
+
101
+ def test_no_audio(self):
102
+ """Video only → not interleaved."""
103
+ input_ids, is_multimodal = make_token_seq(0, 0, 6)
104
+ is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
105
+ is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
106
+ assert not check_interleaved_audio_video(
107
+ is_video, is_audio, is_video.sum().item(), is_audio.sum().item()
108
+ )
109
+
110
+ def test_interleaved(self):
111
+ """V A V A interleaved → True."""
112
+ input_ids, is_multimodal = make_interleaved_seq([4, 4], [3, 3])
113
+ is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
114
+ is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
115
+ assert check_interleaved_audio_video(
116
+ is_video, is_audio, is_video.sum().item(), is_audio.sum().item()
117
+ )
118
+
119
+
120
+# ---------------------------------------------------------------------------
121
+# Tests for embed_input_ids via a minimal mock
122
+# ---------------------------------------------------------------------------
123
+
124
+
125
+def make_mock_model(hidden: int = 8):
126
+ """
127
+ Return a minimal mock of Qwen2_5OmniThinkerForConditionalGeneration
128
+ that has enough structure to run embed_input_ids.
129
+ """
130
+ from vllm.model_executor.models.qwen2_5_omni_thinker import (
131
+ Qwen2_5OmniThinkerForConditionalGeneration,
132
+ )
133
+
134
+ model = Mock(spec=Qwen2_5OmniThinkerForConditionalGeneration)
135
+
136
+ # Config with token IDs
137
+ cfg = Mock()
138
+ cfg.video_token_index = VIDEO_TOKEN_ID
139
+ cfg.audio_token_index = AUDIO_TOKEN_ID
140
+ model.config = cfg
141
+
142
+ # embed_input_ids: simply embed each token as a one-hot-like vector
143
+ # token_id * ones so we can verify which embedding ends up where.
144
+ def fake_lm_embed(ids: torch.Tensor) -> torch.Tensor:
145
+ # Use .clone() so the tensor is contiguous (expand() creates a strided
146
+ # view with shared memory, which masked_scatter_ cannot handle).
147
+ return ids.float().unsqueeze(-1).expand(-1, hidden).clone()
148
+
149
+ lang_model = Mock()
150
+ lang_model.embed_input_ids = fake_lm_embed
151
+ model.get_language_model = Mock(return_value=lang_model)
152
+
153
+ # _embed_text_input_ids: delegate to SupportsMultiModal's implementation
154
+ from vllm.model_executor.models.interfaces import SupportsMultiModal
155
+
156
+ model._embed_text_input_ids = (
157
+ lambda *a, **kw: SupportsMultiModal._embed_text_input_ids(model, *a, **kw)
158
+ )
159
+
160
+ # super().embed_input_ids → use SupportsMultiModal.embed_input_ids
161
+ def fake_super_embed(
162
+ ids, mm_embs=None, *, is_multimodal=None, handle_oov_mm_token=False
163
+ ):
164
+ return SupportsMultiModal.embed_input_ids(
165
+ model,
166
+ ids,
167
+ mm_embs,
168
+ is_multimodal=is_multimodal,
169
+ handle_oov_mm_token=handle_oov_mm_token,
170
+ )
171
+
172
+ # Bind embed_input_ids as the real method
173
+ model.embed_input_ids = (
174
+ lambda *a, **kw: Qwen2_5OmniThinkerForConditionalGeneration.embed_input_ids(
175
+ model, *a, **kw
176
+ )
177
+ )
178
+
179
+ # Store super-embed for use inside the method
180
+ model._super_embed_input_ids = fake_super_embed
181
+
182
+ return model, hidden
183
+
184
+
185
+def build_mm_embeds(
186
+ audio_n, image_n, video_n, hidden, audio_val=10.0, image_val=20.0, video_val=30.0
187
+):
188
+ """
189
+ Build multimodal_embeddings list in position order (audio, image, video).
190
+ Each embedding is filled with a distinct constant so we can verify placement.
191
+ """
192
+ embs = []
193
+ if audio_n:
194
+ embs.append(torch.full((audio_n, hidden), audio_val))
195
+ if image_n:
196
+ embs.append(torch.full((image_n, hidden), image_val))
197
+ if video_n:
198
+ embs.append(torch.full((video_n, hidden), video_val))
199
+ return embs
200
+
201
+
202
+class TestEmbedInputIds:
203
+ def _run(self, audio_n, image_n, video_n, hidden=8):
204
+ """
205
+ Run embed_input_ids for a non-interleaved mixed-modality sequence.
206
+ Returns (result_embeds, input_ids, is_multimodal).
207
+ """
208
+ input_ids, is_multimodal = make_token_seq(audio_n, image_n, video_n)
209
+ mm_embeds = build_mm_embeds(audio_n, image_n, video_n, hidden)
210
+
211
+ model, _ = make_mock_model(hidden)
212
+ result = model.embed_input_ids(
213
+ input_ids, mm_embeds, is_multimodal=is_multimodal
214
+ )
215
+ return result, input_ids, is_multimodal
216
+
217
+ def test_audio_only(self):
218
+ """Audio-only: audio positions get audio embeddings."""
219
+ audio_n, hidden = 5, 8
220
+ audio_val = 10.0
221
+ result, input_ids, is_multimodal = self._run(audio_n, 0, 0, hidden)
222
+
223
+ audio_pos = (input_ids == AUDIO_TOKEN_ID).nonzero(as_tuple=True)[0]
224
+ assert result[audio_pos].allclose(torch.full((audio_n, hidden), audio_val)), (
225
+ "Audio positions should get audio embeddings"
226
+ )
227
+
228
+ def test_video_only(self):
229
+ """Video-only: video positions get video embeddings."""
230
+ video_n, hidden = 6, 8
231
+ video_val = 30.0
232
+ result, input_ids, is_multimodal = self._run(0, 0, video_n, hidden)
233
+
234
+ video_pos = (input_ids == VIDEO_TOKEN_ID).nonzero(as_tuple=True)[0]
235
+ assert result[video_pos].allclose(torch.full((video_n, hidden), video_val)), (
236
+ "Video positions should get video embeddings"
237
+ )
238
+
239
+ def test_mixed_modalities_audio_goes_to_audio_pos(self):
240
+ """
241
+ Regression test for GitHub issue #34506:
242
+ With audio + image + video (non-interleaved), audio positions must
243
+ receive audio embeddings (not image or video embeddings).
244
+ """
245
+ audio_n, image_n, video_n, hidden = 5, 4, 6, 8
246
+ audio_val, image_val, video_val = 10.0, 20.0, 30.0
247
+
248
+ result, input_ids, is_multimodal = self._run(audio_n, image_n, video_n, hidden)
249
+
250
+ audio_pos = (input_ids == AUDIO_TOKEN_ID).nonzero(as_tuple=True)[0]
251
+ image_pos = (input_ids == IMAGE_TOKEN_ID).nonzero(as_tuple=True)[0]
252
+ video_pos = (input_ids == VIDEO_TOKEN_ID).nonzero(as_tuple=True)[0]
253
+
254
+ mean_a = result[audio_pos].mean().item()
255
+ assert result[audio_pos].allclose(torch.full((audio_n, hidden), audio_val)), (
256
+ f"Audio emb wrong: expected {audio_val}, got mean={mean_a:.1f}"
257
+ )
258
+
259
+ mean_i = result[image_pos].mean().item()
260
+ assert result[image_pos].allclose(torch.full((image_n, hidden), image_val)), (
261
+ f"Image emb wrong: expected {image_val}, got mean={mean_i:.1f}"
262
+ )
263
+
264
+ mean_v = result[video_pos].mean().item()
265
+ assert result[video_pos].allclose(torch.full((video_n, hidden), video_val)), (
266
+ f"Video emb wrong: expected {video_val}, got mean={mean_v:.1f}"
267
+ )
268
+
269
+ def test_text_positions_unchanged(self):
270
+ """Text positions should keep their text embeddings."""
271
+ audio_n, image_n, video_n, hidden = 3, 2, 4, 8
272
+ result, input_ids, is_multimodal = self._run(audio_n, image_n, video_n, hidden)
273
+
274
+ text_pos = (~is_multimodal).nonzero(as_tuple=True)[0]
275
+ # Text tokens have value TEXT_TOKEN_ID=0, so embed → 0.0
276
+ assert result[text_pos].allclose(torch.zeros(len(text_pos), hidden)), (
277
+ "Text positions should keep text embeddings"
278
+ )
279
+
280
+ def test_interleaved_use_audio_in_video(self):
281
+ """
282
+ Interleaved (use_audio_in_video): video chunks interleaved with audio.
283
+ Video embeddings must go to video positions, audio to audio positions.
284
+ """
285
+ hidden = 8
286
+ audio_val, video_val = 10.0, 30.0
287
+ # Two video chunks of 4, two audio chunks of 3
288
+ video_chunks = [4, 4]
289
+ audio_chunks = [3, 3]
290
+ input_ids, is_multimodal = make_interleaved_seq(video_chunks, audio_chunks)
291
+
292
+ video_n = sum(video_chunks) # 8
293
+ audio_n = sum(audio_chunks) # 6
294
+
295
+ # mm_embeds come in [video, audio] order (video feature first in
296
+ # mm_features when positions are the same for use_audio_in_video)
297
+ mm_embeds = [
298
+ torch.full((video_n, hidden), video_val),
299
+ torch.full((audio_n, hidden), audio_val),
300
+ ]
301
+
302
+ model, _ = make_mock_model(hidden)
303
+ result = model.embed_input_ids(
304
+ input_ids, mm_embeds, is_multimodal=is_multimodal
305
+ )
306
+
307
+ video_pos = (input_ids == VIDEO_TOKEN_ID).nonzero(as_tuple=True)[0]
308
+ audio_pos = (input_ids == AUDIO_TOKEN_ID).nonzero(as_tuple=True)[0]
309
+
310
+ assert result[video_pos].allclose(torch.full((video_n, hidden), video_val)), (
311
+ "Interleaved: video positions should get video embeddings"
312
+ )
313
+
314
+ assert result[audio_pos].allclose(torch.full((audio_n, hidden), audio_val)), (
315
+ "Interleaved: audio positions should get audio embeddings"
316
+ )
317
+
318
+
319
+# ---------------------------------------------------------------------------
320
+# Tests for merge_interleaved_embeddings helper
321
+# ---------------------------------------------------------------------------
322
+
323
+
324
+class TestMergeInterleavedEmbeddings:
325
+ def test_basic_interleaved(self):
326
+ """Video chunks + audio chunks scattered to correct positions."""
327
+ hidden = 4
328
+ input_ids, is_multimodal = make_interleaved_seq([3, 3], [2, 2])
329
+
330
+ is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
331
+ is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
332
+ num_video = is_video.sum().item() # 6
333
+ num_audio = is_audio.sum().item() # 4
334
+
335
+ inputs_embeds = torch.zeros(len(input_ids), hidden)
336
+ mm_embeds = [
337
+ torch.full((num_video, hidden), 30.0),
338
+ torch.full((num_audio, hidden), 10.0),
339
+ ]
340
+
341
+ result = merge_interleaved_embeddings(
342
+ inputs_embeds,
343
+ mm_embeds,
344
+ is_video,
345
+ is_audio,
346
+ is_multimodal,
347
+ num_video,
348
+ num_audio,
349
+ )
350
+
351
+ video_pos = is_video.nonzero(as_tuple=True)[0]
352
+ audio_pos = is_audio.nonzero(as_tuple=True)[0]
353
+ assert result[video_pos].allclose(torch.full((num_video, hidden), 30.0))
354
+ assert result[audio_pos].allclose(torch.full((num_audio, hidden), 10.0))
355
+
356
+
357
+if __name__ == "__main__":
358
+ pytest.main([__file__, "-v"])
vllm/model_executor/models/qwen2_5_omni_thinker.py
+14
-16
@@ -1376,23 +1376,12 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
1376
is_multimodal: torch.Tensor | None = None,
1377
handle_oov_mm_token: bool = False,
1378
) -> torch.Tensor:
1379
- from .utils import _merge_multimodal_embeddings
1380
-
1379
if multimodal_embeddings is None or is_multimodal is None:
1380
return super().embed_input_ids(input_ids)
1381
1384
- inputs_embeds = self._embed_text_input_ids(
1385
- input_ids,
1386
- self.get_language_model().embed_input_ids,
1387
- is_multimodal=is_multimodal,
1388
- handle_oov_mm_token=handle_oov_mm_token,
1389
- )
1390
-
1391
- if len(multimodal_embeddings) == 0:
1392
- return inputs_embeds
1393
-
1382
# Check for audio-in-video: interleaved video and audio tokens
1395
- # in the multimodal region.
1383
+ # in the multimodal region. Only use the interleaved path when
1384
+ # needed; otherwise fall back to the default parent implementation.
1385
video_token_id = self.config.video_token_index
1386
audio_token_id = self.config.audio_token_index
1387
@@ -1403,6 +1392,12 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
1392
num_audio = is_audio.sum().item()
1393
1394
if check_interleaved_audio_video(is_video, is_audio, num_video, num_audio):
1395
+ inputs_embeds = self._embed_text_input_ids(
1396
+ input_ids,
1397
+ self.get_language_model().embed_input_ids,
1398
+ is_multimodal=is_multimodal,
1399
+ handle_oov_mm_token=handle_oov_mm_token,
1400
+ )
1401
return merge_interleaved_embeddings(
1402
inputs_embeds,
1403
multimodal_embeddings,
@@ -1413,9 +1408,12 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
1408
num_audio,
1409
)
1410
1416
- # Default: standard merge (no interleaving)
1417
- return _merge_multimodal_embeddings(
1418
- inputs_embeds, multimodal_embeddings, is_multimodal
1411
+ # Default: standard merge (no interleaving), same as parent class
1412
+ return super().embed_input_ids(
1413
+ input_ids,
1414
+ multimodal_embeddings=multimodal_embeddings,
1415
+ is_multimodal=is_multimodal,
1416
+ handle_oov_mm_token=handle_oov_mm_token,
1417
)
1418
1419
def forward(
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+7
-5
@@ -1904,15 +1904,17 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1904
num_audio,
1905
)
1906
1907
- # Default: standard merge (no interleaving)
1908
- inputs_embeds = _merge_multimodal_embeddings(
1909
- inputs_embeds=inputs_embeds,
1907
+ # Default: standard merge (no interleaving), same as parent class.
1908
+ # multimodal_embeddings may have been updated above (deepstack
1909
+ # main-scale). Use super() to stay consistent with the parent
1910
+ # implementation and avoid issues seen in Qwen2.5-Omni (#34506).
1911
+ return super().embed_input_ids(
1912
+ input_ids,
1913
multimodal_embeddings=multimodal_embeddings,
1914
is_multimodal=is_multimodal,
1915
+ handle_oov_mm_token=handle_oov_mm_token,
1916
)
1917
1914
- return inputs_embeds
1915
-
1918
def forward(
1919
self,
1920
input_ids: torch.Tensor | None,