[Bugfix][Model] Fix audio-in-video support for Qwen2.5-Omni and Qwen3-Omni

vllm/model_executor/models/qwen2_5_omni_thinker.py

@@ -1286,17 +1286,86 @@ def embed_input_ids(
         is_multimodal: torch.Tensor | None = None,
         handle_oov_mm_token: bool = False,
     ) -> torch.Tensor:
-        # This is to satisfy the type checker for each overload
+        from .utils import _merge_multimodal_embeddings
+
         if multimodal_embeddings is None or is_multimodal is None:
             return super().embed_input_ids(input_ids)
 
-        return super().embed_input_ids(
+        inputs_embeds = self._embed_text_input_ids(
             input_ids,
-            multimodal_embeddings=multimodal_embeddings,
+            self.get_language_model().embed_input_ids,
             is_multimodal=is_multimodal,
             handle_oov_mm_token=handle_oov_mm_token,
         )
 
+        if len(multimodal_embeddings) == 0:
+            return inputs_embeds
+
+        # Check for audio-in-video: interleaved video and audio tokens
+        # in the multimodal region. When use_audio_in_video=True, video
+        # and audio tokens are interleaved in the token sequence, but
+        # the embeddings are provided as separate contiguous tensors.
+        # A single masked_scatter_ would place them in the wrong order,
+        # so we scatter each modality separately using per-modality masks.
+        video_token_id = self.config.video_token_index
+        audio_token_id = self.config.audio_token_index
+
+        is_video = is_multimodal & (input_ids == video_token_id)
+        is_audio = is_multimodal & (input_ids == audio_token_id)
+
+        num_video = is_video.sum().item()
+        num_audio = is_audio.sum().item()
+
+        if num_video > 0 and num_audio > 0:
+            # Check if video and audio positions are actually interleaved
+            video_pos = is_video.nonzero(as_tuple=True)[0]
+            audio_pos = is_audio.nonzero(as_tuple=True)[0]
+
+            is_interleaved = (
+                video_pos[0].item() < audio_pos[-1].item()
+                and audio_pos[0].item() < video_pos[-1].item()
+            )
+
+            if is_interleaved:
+                # Match embeddings to modalities by exact token count
+                video_embeds: list[torch.Tensor] = []
+                audio_embeds: list[torch.Tensor] = []
+                other_embeds: list[torch.Tensor] = []
+                video_remaining = num_video
+                audio_remaining = num_audio
+
+                for emb in multimodal_embeddings:
+                    n = emb.shape[0]
+                    if video_remaining > 0 and n <= video_remaining:
+                        video_embeds.append(emb)
+                        video_remaining -= n
+                    elif audio_remaining > 0 and n <= audio_remaining:
+                        audio_embeds.append(emb)
+                        audio_remaining -= n
+                    else:
+                        other_embeds.append(emb)
+
+                if video_embeds:
+                    inputs_embeds = _merge_multimodal_embeddings(
+                        inputs_embeds, video_embeds, is_video
+                    )
+                if audio_embeds:
+                    inputs_embeds = _merge_multimodal_embeddings(
+                        inputs_embeds, audio_embeds, is_audio
+                    )
+                if other_embeds:
+                    other_mask = is_multimodal & ~is_video & ~is_audio
+                    inputs_embeds = _merge_multimodal_embeddings(
+                        inputs_embeds, other_embeds, other_mask
+                    )
+
+                return inputs_embeds
+
+        # Default: standard merge (no interleaving)
+        return _merge_multimodal_embeddings(
+            inputs_embeds, multimodal_embeddings, is_multimodal
+        )
+
     def forward(
         self,
         input_ids: torch.Tensor | None,

vllm/model_executor/models/qwen3_omni_moe_thinker.py

@@ -1780,6 +1780,24 @@ def embed_input_ids(
         if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
             return inputs_embeds
 
+        # Detect interleaved audio-in-video early, since it affects
+        # both the deepstack path and the final embedding merge.
+        video_token_id = self.config.video_token_id
+        audio_token_id = self.config.audio_token_id
+        is_video = is_multimodal & (input_ids == video_token_id)
+        is_audio = is_multimodal & (input_ids == audio_token_id)
+        num_video = is_video.sum().item()
+        num_audio = is_audio.sum().item()
+
+        is_interleaved = False
+        if num_video > 0 and num_audio > 0:
+            video_pos = is_video.nonzero(as_tuple=True)[0]
+            audio_pos = is_audio.nonzero(as_tuple=True)[0]
+            is_interleaved = (
+                video_pos[0].item() < audio_pos[-1].item()
+                and audio_pos[0].item() < video_pos[-1].item()
+            )
+
         deepstack_input_embeds = None
         # split the feat dim to obtain multi-scale visual feature
         has_vision_embeddings = [
@@ -1791,14 +1809,18 @@ def embed_input_ids(
         ):
             multiscale_len = len(self.visual.deepstack_visual_indexes)
             multimodal_embeddings_multiscale = []
-            is_vision = torch.zeros_like(is_multimodal)
-            mm_positions = torch.nonzero(is_multimodal, as_tuple=True)[0]
-            mm_position_idx = 0
+
+            if is_interleaved:
+                # Use input_ids-based mask for correct vision positions
+                # when audio and video tokens are interleaved.
+                is_vision = is_video.clone()
+            else:
+                is_vision = torch.zeros_like(is_multimodal)
+                mm_positions = torch.nonzero(is_multimodal, as_tuple=True)[0]
+                mm_position_idx = 0
+
             for index, embeddings in enumerate(multimodal_embeddings):
                 num_tokens = embeddings.shape[0]
-                current_positions = mm_positions[
-                    mm_position_idx : mm_position_idx + num_tokens
-                ]
 
                 # Vision embeddings
                 if embeddings.shape[-1] != self.config.text_config.hidden_size:
@@ -1809,13 +1831,22 @@ def embed_input_ids(
                     )
                     multimodal_embeddings[index] = embeddings_main
                     multimodal_embeddings_multiscale.append(embeddings_multiscale)
-                    is_vision[current_positions] = True
+                    if not is_interleaved:
+                        current_positions = mm_positions[
+                            mm_position_idx : mm_position_idx + num_tokens
+                        ]
+                        is_vision[current_positions] = True
 
                 # Audio embeddings
                 else:
-                    is_vision[current_positions] = False
+                    if not is_interleaved:
+                        current_positions = mm_positions[
+                            mm_position_idx : mm_position_idx + num_tokens
+                        ]
+                        is_vision[current_positions] = False
 
-                mm_position_idx += num_tokens
+                if not is_interleaved:
+                    mm_position_idx += num_tokens
 
             deepstack_input_embeds = inputs_embeds.new_zeros(
                 inputs_embeds.size(0), multiscale_len * inputs_embeds.size(1)
@@ -1834,6 +1865,42 @@ def embed_input_ids(
             )
             self._set_deepstack_input_embeds(deepstack_input_embeds)
 
+        if is_interleaved:
+            # Per-modality scatter for interleaved audio-in-video.
+            video_embeds: list[torch.Tensor] = []
+            audio_embeds: list[torch.Tensor] = []
+            other_embeds: list[torch.Tensor] = []
+            video_remaining = num_video
+            audio_remaining = num_audio
+
+            for emb in multimodal_embeddings:
+                n = emb.shape[0]
+                if video_remaining > 0 and n <= video_remaining:
+                    video_embeds.append(emb)
+                    video_remaining -= n
+                elif audio_remaining > 0 and n <= audio_remaining:
+                    audio_embeds.append(emb)
+                    audio_remaining -= n
+                else:
+                    other_embeds.append(emb)
+
+            if video_embeds:
+                inputs_embeds = _merge_multimodal_embeddings(
+                    inputs_embeds, video_embeds, is_video
+                )
+            if audio_embeds:
+                inputs_embeds = _merge_multimodal_embeddings(
+                    inputs_embeds, audio_embeds, is_audio
+                )
+            if other_embeds:
+                other_mask = is_multimodal & ~is_video & ~is_audio
+                inputs_embeds = _merge_multimodal_embeddings(
+                    inputs_embeds, other_embeds, other_mask
+                )
+
+            return inputs_embeds
+
+        # Default: standard merge (no interleaving)
         inputs_embeds = _merge_multimodal_embeddings(
             inputs_embeds=inputs_embeds,
             multimodal_embeddings=multimodal_embeddings,