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 8, 22:45 UTC
[Bugfix][Model] Fix audio-in-video support for Qwen2.5-Omni and Qwen3-Omni (#33605)
Signed-off-by: linyueqian <linyueqian@outlook.com> Signed-off-by: Roger Wang <hey@rogerw.io> Co-authored-by: Roger Wang <hey@rogerw.io>
Yueqian Lin committed
Feb 4, 2026 at 07:15 UTC
f8516a1ab95febcf131a37478914031f50fdd9db
2 files changed
+172
-12
vllm/model_executor/models/qwen2_5_omni_thinker.py
+123
-3
@@ -113,6 +113,95 @@ except (ImportError, ModuleNotFoundError):
113
logger = init_logger(__name__)
114
115
116
+def check_interleaved_audio_video(
117
+ is_video: torch.Tensor,
118
+ is_audio: torch.Tensor,
119
+ num_video: int,
120
+ num_audio: int,
121
+) -> bool:
122
+ """
123
+ Check if video and audio positions are interleaved in the multimodal region.
124
+
125
+ Returns:
126
+ True if video and audio tokens are interleaved, False otherwise.
127
+ """
128
+ if num_video == 0 or num_audio == 0:
129
+ return False
130
+
131
+ video_pos = is_video.nonzero(as_tuple=True)[0]
132
+ audio_pos = is_audio.nonzero(as_tuple=True)[0]
133
+
134
+ return (
135
+ video_pos[0].item() < audio_pos[-1].item()
136
+ and audio_pos[0].item() < video_pos[-1].item()
137
+ )
138
+
139
+
140
+def merge_interleaved_embeddings(
141
+ inputs_embeds: torch.Tensor,
142
+ multimodal_embeddings: "MultiModalEmbeddings",
143
+ is_video: torch.Tensor,
144
+ is_audio: torch.Tensor,
145
+ is_multimodal: torch.Tensor,
146
+ num_video: int,
147
+ num_audio: int,
148
+) -> torch.Tensor:
149
+ """
150
+ Merge embeddings for interleaved audio-in-video sequences.
151
+
152
+ When use_audio_in_video=True, video and audio tokens are interleaved in
153
+ the token sequence, but embeddings are provided as separate contiguous
154
+ tensors (video first, then audio). This function reorders video and audio
155
+ embeddings to match sequence position order and scatters them efficiently.
156
+
157
+ Args:
158
+ inputs_embeds: The input embeddings tensor to merge into.
159
+ multimodal_embeddings: List of embedding tensors (video, audio, other).
160
+ is_video: Boolean mask for video token positions.
161
+ is_audio: Boolean mask for audio token positions.
162
+ is_multimodal: Boolean mask for all multimodal token positions.
163
+ num_video: Total count of video tokens.
164
+ num_audio: Total count of audio tokens.
165
+
166
+ Returns:
167
+ The merged inputs_embeds tensor with multimodal embeddings scattered
168
+ to their correct positions.
169
+ """
170
+ # Categorize embeddings by modality based on token counts.
171
+ # Embeddings come grouped by modality but order varies (e.g., image, video, audio
172
+ # or video, audio depending on input kwargs order).
173
+ video_embeds: list[torch.Tensor] = []
174
+ audio_embeds: list[torch.Tensor] = []
175
+ other_embeds: list[torch.Tensor] = []
176
+ video_remaining = num_video
177
+ audio_remaining = num_audio
178
+
179
+ for emb in multimodal_embeddings:
180
+ n = emb.shape[0]
181
+ if video_remaining > 0 and n <= video_remaining:
182
+ video_embeds.append(emb)
183
+ video_remaining -= n
184
+ elif audio_remaining > 0 and n <= audio_remaining:
185
+ audio_embeds.append(emb)
186
+ audio_remaining -= n
187
+ else:
188
+ other_embeds.append(emb)
189
+
190
+ # Scatter each modality to its positions
191
+ if video_embeds:
192
+ video_positions = is_video.nonzero(as_tuple=True)[0]
193
+ inputs_embeds[video_positions] = torch.cat(video_embeds, dim=0)
194
+ if audio_embeds:
195
+ audio_positions = is_audio.nonzero(as_tuple=True)[0]
196
+ inputs_embeds[audio_positions] = torch.cat(audio_embeds, dim=0)
197
+ if other_embeds:
198
+ other_mask = is_multimodal & ~is_video & ~is_audio
199
+ other_positions = other_mask.nonzero(as_tuple=True)[0]
200
+ inputs_embeds[other_positions] = torch.cat(other_embeds, dim=0)
201
+
202
+ return inputs_embeds
203
+
204
+
205
class Qwen2_5OmniAudioFeatureInputs(TensorSchema):
206
"""
207
Dimensions:
@@ -1286,17 +1375,48 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
1375
is_multimodal: torch.Tensor | None = None,
1376
handle_oov_mm_token: bool = False,
1377
) -> torch.Tensor:
1289
- # This is to satisfy the type checker for each overload
1378
+ from .utils import _merge_multimodal_embeddings
1379
+
1380
if multimodal_embeddings is None or is_multimodal is None:
1381
return super().embed_input_ids(input_ids)
1382
1293
- return super().embed_input_ids(
1383
+ inputs_embeds = self._embed_text_input_ids(
1384
input_ids,
1295
- multimodal_embeddings=multimodal_embeddings,
1385
+ self.get_language_model().embed_input_ids,
1386
is_multimodal=is_multimodal,
1387
handle_oov_mm_token=handle_oov_mm_token,
1388
)
1389
1390
+ if len(multimodal_embeddings) == 0:
1391
+ return inputs_embeds
1392
+
1393
+ # Check for audio-in-video: interleaved video and audio tokens
1394
+ # in the multimodal region.
1395
+ video_token_id = self.config.video_token_index
1396
+ audio_token_id = self.config.audio_token_index
1397
+
1398
+ is_video = is_multimodal & (input_ids == video_token_id)
1399
+ is_audio = is_multimodal & (input_ids == audio_token_id)
1400
+
1401
+ num_video = is_video.sum().item()
1402
+ num_audio = is_audio.sum().item()
1403
+
1404
+ if check_interleaved_audio_video(is_video, is_audio, num_video, num_audio):
1405
+ return merge_interleaved_embeddings(
1406
+ inputs_embeds,
1407
+ multimodal_embeddings,
1408
+ is_video,
1409
+ is_audio,
1410
+ is_multimodal,
1411
+ num_video,
1412
+ num_audio,
1413
+ )
1414
+
1415
+ # Default: standard merge (no interleaving)
1416
+ return _merge_multimodal_embeddings(
1417
+ inputs_embeds, multimodal_embeddings, is_multimodal
1418
+ )
1419
+
1420
def forward(
1421
self,
1422
input_ids: torch.Tensor | None,
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+49
-9
@@ -92,6 +92,8 @@ from .qwen2_5_omni_thinker import (
92
Qwen2_5OmniConditionalGenerationMixin,
93
Qwen2_5OmniThinkerDummyInputsBuilder,
94
Qwen2_5OmniThinkerMultiModalProcessor,
95
+ check_interleaved_audio_video,
96
+ merge_interleaved_embeddings,
97
)
98
from .qwen2_5_vl import (
99
Qwen2_5_VisionAttention,
@@ -1780,6 +1782,19 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1782
if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
1783
return inputs_embeds
1784
1785
+ # Detect interleaved audio-in-video early, since it affects
1786
+ # both the deepstack path and the final embedding merge.
1787
+ video_token_id = self.config.video_token_id
1788
+ audio_token_id = self.config.audio_token_id
1789
+ is_video = is_multimodal & (input_ids == video_token_id)
1790
+ is_audio = is_multimodal & (input_ids == audio_token_id)
1791
+ num_video = is_video.sum().item()
1792
+ num_audio = is_audio.sum().item()
1793
+
1794
+ is_interleaved = check_interleaved_audio_video(
1795
+ is_video, is_audio, num_video, num_audio
1796
+ )
1797
+
1798
deepstack_input_embeds = None
1799
# split the feat dim to obtain multi-scale visual feature
1800
has_vision_embeddings = [
@@ -1791,14 +1806,18 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1806
):
1807
multiscale_len = len(self.visual.deepstack_visual_indexes)
1808
multimodal_embeddings_multiscale = []
1794
- is_vision = torch.zeros_like(is_multimodal)
1795
- mm_positions = torch.nonzero(is_multimodal, as_tuple=True)[0]
1796
- mm_position_idx = 0
1809
+
1810
+ if is_interleaved:
1811
+ # Use input_ids-based mask for correct vision positions
1812
+ # when audio and video tokens are interleaved.
1813
+ is_vision = is_video.clone()
1814
+ else:
1815
+ is_vision = torch.zeros_like(is_multimodal)
1816
+ mm_positions = torch.nonzero(is_multimodal, as_tuple=True)[0]
1817
+ mm_position_idx = 0
1818
+
1819
for index, embeddings in enumerate(multimodal_embeddings):
1820
num_tokens = embeddings.shape[0]
1799
- current_positions = mm_positions[
1800
- mm_position_idx : mm_position_idx + num_tokens
1801
- ]
1821
1822
# Vision embeddings
1823
if embeddings.shape[-1] != self.config.text_config.hidden_size:
@@ -1809,13 +1828,22 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1828
)
1829
multimodal_embeddings[index] = embeddings_main
1830
multimodal_embeddings_multiscale.append(embeddings_multiscale)
1812
- is_vision[current_positions] = True
1831
+ if not is_interleaved:
1832
+ current_positions = mm_positions[
1833
+ mm_position_idx : mm_position_idx + num_tokens
1834
+ ]
1835
+ is_vision[current_positions] = True
1836
1837
# Audio embeddings
1838
else:
1816
- is_vision[current_positions] = False
1839
+ if not is_interleaved:
1840
+ current_positions = mm_positions[
1841
+ mm_position_idx : mm_position_idx + num_tokens
1842
+ ]
1843
+ is_vision[current_positions] = False
1844
1818
- mm_position_idx += num_tokens
1845
+ if not is_interleaved:
1846
+ mm_position_idx += num_tokens
1847
1848
deepstack_input_embeds = inputs_embeds.new_zeros(
1849
inputs_embeds.size(0), multiscale_len * inputs_embeds.size(1)
@@ -1834,6 +1862,18 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1862
)
1863
self._set_deepstack_input_embeds(deepstack_input_embeds)
1864
1865
+ if is_interleaved:
1866
+ return merge_interleaved_embeddings(
1867
+ inputs_embeds,
1868
+ multimodal_embeddings,
1869
+ is_video,
1870
+ is_audio,
1871
+ is_multimodal,
1872
+ num_video,
1873
+ num_audio,
1874
+ )
1875
+
1876
+ # Default: standard merge (no interleaving)
1877
inputs_embeds = _merge_multimodal_embeddings(
1878
inputs_embeds=inputs_embeds,
1879
multimodal_embeddings=multimodal_embeddings,