From b9b8282cf5a41726081a6516bf974af7ebdba70d Mon Sep 17 00:00:00 2001 From: Zheng Wengang Date: Mon, 27 Jul 2026 12:29:53 +0800 Subject: [PATCH] [BugFix] Couple audio+video mm-cache for use_audio_in_video (#5308) Signed-off-by: ZhengWG --- .../test_qwen3_omni_expansion.py | 1 - .../qwen2_5_omni/test_per_video_audio_mask.py | 111 ++++++++++++++++++ .../qwen2_5_omni/qwen2_5_omni_thinker.py | 86 +++++++++++++- 3 files changed, 191 insertions(+), 7 deletions(-) diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index d9fa2e6e2..3ded285ce 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/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: diff --git a/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py b/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py index e17cdd6f2..4e325bbdc 100644 --- a/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py +++ b/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 = { diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py index 03edcfa10..827f355a6 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py +++ b/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,