[BugFix] Couple audio+video mm-cache for use_audio_in_video (#5308)

Signed-off-by: ZhengWG <zwg0606@gmail.com>

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -411,7 +411,6 @@ def test_audio_in_video_002(omni_server, openai_client) -> None:
     openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
-@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/5248")
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_audio_in_video_default_loader_sampling_regression(omni_server, openai_client) -> None:

tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py

@@ -215,6 +215,117 @@ def test_filter_video_use_audio_in_video_for_uncached_items_validates_request_ma
         )
 
 
+def test_get_audio_in_video_pairs_respects_per_video_mask():
+    from unittest.mock import MagicMock
+
+    mm_data_items = MagicMock()
+    mm_data_items.get_count.side_effect = lambda modality, strict=False: {
+        "video": 2,
+        "audio": 1,
+    }[modality]
+
+    pairs = Qwen2_5OmniThinkerMultiModalProcessor._get_audio_in_video_pairs(
+        SimpleNamespace(),
+        mm_data_items,
+        {"use_audio_in_video": [False, True]},
+    )
+
+    # Second video uses audio and pairs with audio[0], not audio[1].
+    assert pairs == [(1, 0)]
+
+
+def test_paired_cache_keys_change_when_paired_video_changes():
+    """Same audio + different video must not reuse the audio cache entry."""
+    pairs = [(0, 0)]
+    keys_a = Qwen2_5OmniThinkerMultiModalProcessor._paired_cache_keys(
+        {"video": ["video-a"], "audio": ["audio-same"]},
+        pairs,
+    )
+    keys_b = Qwen2_5OmniThinkerMultiModalProcessor._paired_cache_keys(
+        {"video": ["video-b"], "audio": ["audio-same"]},
+        pairs,
+    )
+
+    assert keys_a["audio"][0] != keys_b["audio"][0]
+    assert keys_a["video"][0] != keys_b["video"][0]
+
+
+def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss():
+    """Same audio + different video must reprocess both items together even when
+    the audio content hash alone would be a cache hit.
+    """
+    from unittest.mock import MagicMock, patch
+
+    mm_data_items = MagicMock()
+    mm_data_items.get_count.side_effect = lambda modality, strict=False: {
+        "video": 1,
+        "audio": 1,
+    }[modality]
+    mm_data_items.__getitem__ = MagicMock(side_effect=lambda modality: [f"{modality}-0"])
+
+    inputs = MagicMock()
+    inputs.mm_data_items = mm_data_items
+    inputs.hf_processor_mm_kwargs = {"use_audio_in_video": True}
+    inputs.tokenization_kwargs = {}
+    inputs.prompt = [1, 2, 3]
+    inputs.get_mm_hashes = MagicMock(return_value={"video": ["v0"], "audio": ["a0"]})
+
+    timing_ctx = MagicMock()
+    timing_ctx.record.return_value.__enter__ = MagicMock(return_value=None)
+    timing_ctx.record.return_value.__exit__ = MagicMock(return_value=False)
+
+    cache = MagicMock()
+    # Video miss, audio hit under pair-aware keys — expansion must force audio.
+    cache.is_cached.side_effect = lambda keys: [False] if keys[0].startswith("v") else [True]
+
+    captured = {}
+
+    def capture_merge(cache_arg, **kwargs):
+        captured["mm_is_cached"] = kwargs["mm_is_cached"]
+        captured["mm_cache_keys"] = kwargs["mm_hashes"]
+        return MagicMock(name="mm_kwargs"), {"video": [[]], "audio": [[]]}
+
+    fake_self = SimpleNamespace(cache=cache, info=SimpleNamespace(model_id="m"))
+    fake_self._get_hf_mm_data = MagicMock(return_value=({}, {}))
+    # Bound methods called via self inside _cached_apply_hf_processor.
+    fake_self._get_audio_in_video_pairs = lambda mm_data_items, hf_processor_mm_kwargs: (
+        Qwen2_5OmniThinkerMultiModalProcessor._get_audio_in_video_pairs(
+            fake_self,
+            mm_data_items,
+            hf_processor_mm_kwargs,
+        )
+    )
+    fake_self._paired_cache_keys = Qwen2_5OmniThinkerMultiModalProcessor._paired_cache_keys
+    fake_self.info.parse_mm_data = MagicMock(return_value=SimpleNamespace(name="missing"))
+    fake_self._apply_hf_processor_main = MagicMock(return_value=([9], {"video_grid_thw": object()}, False))
+    fake_self._get_mm_fields_config = MagicMock(return_value={})
+    fake_self._get_mm_prompt_updates = MagicMock(return_value={"video": [[object()]], "audio": [[object()]]})
+    fake_self._merge_mm_kwargs = capture_merge
+
+    with patch(
+        "vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.MultiModalKwargsItems.from_hf_inputs",
+        return_value={"video": [object()], "audio": [object()]},
+    ):
+        prompt_ids, mm_info, _ = Qwen2_5OmniThinkerMultiModalProcessor._cached_apply_hf_processor(
+            fake_self,
+            inputs,
+            timing_ctx,
+        )
+
+    assert prompt_ids == [9]
+    # Both modalities were reprocessed as one unit after pair-miss expansion.
+    assert fake_self.info.parse_mm_data.call_args.args[0] == {
+        "video": ["video-0"],
+        "audio": ["audio-0"],
+    }
+    assert captured["mm_is_cached"] == {"video": [False], "audio": [False]}
+    # Cache entries are stored under pair-aware keys...
+    assert captured["mm_cache_keys"]["video"][0].startswith("v0-aiv-")
+    assert captured["mm_cache_keys"]["audio"][0].startswith("a0-aiv-")
+    # ...while semantic hashes are exposed downstream.
+    assert mm_info.hashes == {"video": ["v0"], "audio": ["a0"]}
+
+
 def test_mm_fields_config_keeps_global_use_audio_in_video_shared():
     fake_self = _fake_processor_with_spatial_merge_size()
     hf_inputs = {

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py

@@ -1,5 +1,6 @@
 """Thin Omni wrapper: reuse upstream Qwen2.5-Omni thinker with minimal overrides."""
 
+import hashlib
 from collections.abc import Iterable, Iterator, Mapping, Sequence
 from functools import partial
 from typing import Any
@@ -16,6 +17,7 @@ from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
 )
 from vllm.config import VllmConfig
 from vllm.forward_context import set_forward_context
+from vllm.inputs import MultiModalHashes
 from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import (
     MultiModalEmbeddings,
@@ -555,26 +557,98 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
 
         return mm_processed_data
 
+    def _get_audio_in_video_pairs(
+        self,
+        mm_data_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> list[tuple[int, int]]:
+        """``(video_idx, audio_idx)`` pairs the HF processor handles as one
+        indivisible unit when the audio track is interleaved into the video.
+
+        Pairing follows omni's per-video mask convention: the j-th video with
+        ``use_audio_in_video=True`` pairs with ``audio[j]``.
+        """
+        num_videos = mm_data_items.get_count("video", strict=False)
+        num_audios = mm_data_items.get_count("audio", strict=False)
+        if num_videos == 0 or num_audios == 0:
+            return []
+
+        video_use_audio_in_video = _get_request_video_use_audio_in_video(
+            hf_processor_mm_kwargs,
+            num_videos,
+        )
+        videos_using_audio = [video_idx for video_idx, uses_audio in enumerate(video_use_audio_in_video) if uses_audio]
+        return list(zip(videos_using_audio, range(num_audios)))
+
+    @staticmethod
+    def _paired_cache_keys(
+        mm_hashes: MultiModalHashes,
+        aiv_pairs: list[tuple[int, int]],
+    ) -> MultiModalHashes:
+        """Processor-cache keys in which each paired video/audio key
+        incorporates the hash of both members, so a pair is reused only when
+        both match and any change misses the pair together. Semantic
+        ``mm_hashes`` are left untouched.
+        """
+        if not aiv_pairs:
+            return mm_hashes
+
+        mm_cache_keys = {modality: list(hashes) for modality, hashes in mm_hashes.items()}
+        for video_idx, audio_idx in aiv_pairs:
+            video_hash = mm_hashes["video"][video_idx]
+            audio_hash = mm_hashes["audio"][audio_idx]
+            pair_token = hashlib.sha256(f"{video_hash}\0{audio_hash}".encode()).hexdigest()[:16]
+            mm_cache_keys["video"][video_idx] = f"{video_hash}-aiv-{pair_token}"
+            mm_cache_keys["audio"][audio_idx] = f"{audio_hash}-aiv-{pair_token}"
+        return mm_cache_keys
+
     def _cached_apply_hf_processor(
         self,
         inputs: ProcessorInputs,
         timing_ctx: TimingContext,
     ) -> tuple[list[int], MultiModalProcessingInfo, bool]:
+        """Cache-aware: processes video/audio pairs as units,
+        using pair-specific cache keys. Falls back to standard cache if no pairs."""
+
         cache = self.cache
 
         _, passthrough_data = self._get_hf_mm_data(inputs.mm_data_items)
         if cache is None or passthrough_data:
             return self._apply_hf_processor(inputs, timing_ctx)
 
+        aiv_pairs = self._get_audio_in_video_pairs(
+            inputs.mm_data_items,
+            inputs.hf_processor_mm_kwargs,
+        )
+
         with timing_ctx.record("get_mm_hashes"):
             mm_hashes = inputs.get_mm_hashes(self.info.model_id)
 
+        mm_cache_keys = self._paired_cache_keys(mm_hashes, aiv_pairs)
+
         with timing_ctx.record("get_cache_missing_items"):
-            mm_is_cached, mm_missing_data_items = self._get_cache_missing_items(
-                cache=cache,
-                mm_data_items=inputs.mm_data_items,
-                mm_hashes=mm_hashes,
-            )
+            mm_is_cached = {modality: list(cache.is_cached(keys)) for modality, keys in mm_cache_keys.items()}
+            # A pair is one HF processing unit: if either member misses the
+            # cache (including LRU evicting just one), reprocess both.
+            for video_idx, audio_idx in aiv_pairs:
+                if not (mm_is_cached["video"][video_idx] and mm_is_cached["audio"][audio_idx]):
+                    mm_is_cached["video"][video_idx] = False
+                    mm_is_cached["audio"][audio_idx] = False
+
+            # TODO: Inlined from BaseMultiModalProcessor._get_cache_missing_items
+            # (pair expansion must run after is_cached); keep in sync with vLLM.
+            mm_missing_data = {}
+            for modality, items_is_cached in mm_is_cached.items():
+                missing_modality_data = []
+                for idx, item_is_cached in enumerate(items_is_cached):
+                    if item_is_cached:
+                        continue
+                    data = inputs.mm_data_items[modality][idx]
+                    if data is None:
+                        raise ValueError(f"Cache miss for {modality} at index {idx} but data is not provided.")
+                    missing_modality_data.append(data)
+                mm_missing_data[modality] = missing_modality_data
+            mm_missing_data_items = self.info.parse_mm_data(mm_missing_data, validate=False)
 
         hf_processor_mm_kwargs = _filter_video_use_audio_in_video_for_uncached_items(
             inputs.hf_processor_mm_kwargs,
@@ -611,7 +685,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
         with timing_ctx.record("merge_mm_kwargs"):
             mm_kwargs, mm_prompt_updates = self._merge_mm_kwargs(
                 cache,
-                mm_hashes=mm_hashes,
+                mm_hashes=mm_cache_keys,
                 mm_is_cached=mm_is_cached,
                 mm_missing_kwargs=mm_missing_kwargs,
                 mm_missing_prompt_updates=mm_missing_prompt_updates,