From 10d3a25916561aab0ab4f1f2bb191ce1c8cdd3ec Mon Sep 17 00:00:00 2001 From: Wei <48665416+hbhflw2000@users.noreply.github.com> Date: Mon, 20 Jul 2026 11:39:59 +0800 Subject: [PATCH] [Feature] Support per-video audio masks for Qwen Omni (#4656) Signed-off-by: hbhflw2000 <417911774@qq.com> --- .../qwen2_5_omni/test_per_video_audio_mask.py | 708 ++++++++++++++++++ .../qwen2_5_omni/qwen2_5_omni_thinker.py | 562 ++++++++++++-- .../qwen3_omni/qwen3_omni_moe_thinker.py | 193 ++--- vllm_omni/outputs/__init__.py | 12 + 4 files changed, 1338 insertions(+), 137 deletions(-) create mode 100644 tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py 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 new file mode 100644 index 000000000..e17cdd6f2 --- /dev/null +++ b/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py @@ -0,0 +1,708 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch +from vllm.multimodal.inputs import MultiModalBatchedField, MultiModalSharedField +from vllm.multimodal.processing.processor import PlaceholderFeaturesInfo + +from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( + _PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY, + Qwen2_5OmniThinkerMultiModalProcessor, + _coerce_use_audio_in_video_for_hf_processor, + _filter_video_use_audio_in_video_for_uncached_items, + _normalize_use_audio_in_video, +) +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( + Qwen3OmniMoeThinkerForConditionalGeneration, + Qwen3OmniMoeThinkerMultiModalProcessor, +) +from vllm_omni.outputs import OmniModelRunnerOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +AUDIO_TOKEN_ID = 10 +VIDEO_TOKEN_ID = 20 +AUDIO_BOS_TOKEN_ID = 30 +AUDIO_EOS_TOKEN_ID = 31 +IMAGE_TOKEN_ID = 40 + + +class _Feature: + def __init__(self, data: torch.Tensor) -> None: + self.data = data + + +def _fake_processor() -> SimpleNamespace: + tokenizer = SimpleNamespace(get_vocab=lambda: {"<|audio_pad|>": AUDIO_TOKEN_ID}) + processor = SimpleNamespace(audio_token="<|audio_pad|>") + config = SimpleNamespace(audio_token_id=AUDIO_TOKEN_ID) + info = SimpleNamespace( + get_hf_config=lambda: config, + get_hf_processor=lambda: processor, + get_tokenizer=lambda: tokenizer, + ) + return SimpleNamespace(info=info) + + +def _fake_qwen2_prompt_processor() -> SimpleNamespace: + tokenizer = SimpleNamespace( + get_vocab=lambda: { + "<|audio_pad|>": AUDIO_TOKEN_ID, + "<|video_pad|>": VIDEO_TOKEN_ID, + "<|image_pad|>": IMAGE_TOKEN_ID, + } + ) + processor = SimpleNamespace( + audio_token="<|audio_pad|>", + video_token="<|video_pad|>", + image_token="<|image_pad|>", + ) + image_processor = SimpleNamespace(merge_size=1) + config = SimpleNamespace(audio_token_id=AUDIO_TOKEN_ID) + info = SimpleNamespace( + get_hf_config=lambda: config, + get_hf_processor=lambda **_kwargs: processor, + get_image_processor=lambda **_kwargs: image_processor, + get_tokenizer=lambda: tokenizer, + ) + + def omni_get_updates_use_audio_in_video(**_kwargs): + return [ + AUDIO_BOS_TOKEN_ID, + VIDEO_TOKEN_ID, + AUDIO_TOKEN_ID, + AUDIO_EOS_TOKEN_ID, + ] + + return SimpleNamespace( + info=info, + omni_get_updates_use_audio_in_video=omni_get_updates_use_audio_in_video, + ) + + +def _fake_processor_with_spatial_merge_size() -> SimpleNamespace: + hf_config = SimpleNamespace(vision_config=SimpleNamespace(spatial_merge_size=2)) + return SimpleNamespace(info=SimpleNamespace(get_hf_config=lambda: hf_config)) + + +def _fake_qwen3_processor() -> SimpleNamespace: + tokenizer = SimpleNamespace( + get_vocab=lambda: { + "<|audio_pad|>": AUDIO_TOKEN_ID, + "<|video_pad|>": VIDEO_TOKEN_ID, + } + ) + processor = SimpleNamespace( + audio_token="<|audio_pad|>", + video_token="<|video_pad|>", + ) + info = SimpleNamespace( + get_hf_processor=lambda: processor, + get_tokenizer=lambda: tokenizer, + ) + return SimpleNamespace(info=info) + + +def _fake_qwen3_prompt_processor() -> SimpleNamespace: + tokenizer = SimpleNamespace( + get_vocab=lambda: { + "<|audio_pad|>": AUDIO_TOKEN_ID, + "<|video_pad|>": VIDEO_TOKEN_ID, + "<|image_pad|>": IMAGE_TOKEN_ID, + } + ) + processor = SimpleNamespace( + audio_token="<|audio_pad|>", + video_token="<|video_pad|>", + image_token="<|image_pad|>", + ) + image_processor = SimpleNamespace(merge_size=1) + config = SimpleNamespace( + audio_token_id=AUDIO_TOKEN_ID, + video_token_id=VIDEO_TOKEN_ID, + audio_start_token_id=AUDIO_BOS_TOKEN_ID, + audio_end_token_id=AUDIO_EOS_TOKEN_ID, + vision_config=SimpleNamespace(spatial_merge_size=1), + position_id_per_seconds=1, + ) + info = SimpleNamespace( + get_hf_config=lambda: config, + get_hf_processor=lambda **_kwargs: processor, + get_image_processor=lambda **_kwargs: image_processor, + get_tokenizer=lambda: tokenizer, + ) + return SimpleNamespace(info=info) + + +def test_normalize_use_audio_in_video_accepts_bool_and_per_video_mask(): + assert _normalize_use_audio_in_video(True, 2) == [True, True] + assert _normalize_use_audio_in_video(False, 2) == [False, False] + assert _normalize_use_audio_in_video([True, False], 2) == [True, False] + assert _normalize_use_audio_in_video(torch.tensor([1, 0]), 2) == [ + True, + False, + ] + + +def test_normalize_use_audio_in_video_rejects_wrong_length(): + with pytest.raises(ValueError, match="one boolean per video"): + _normalize_use_audio_in_video([True], 2) + + +def test_coerce_use_audio_in_video_for_hf_processor_keeps_global_bool(): + mm_kwargs = {"use_audio_in_video": True} + + result = _coerce_use_audio_in_video_for_hf_processor({"videos": [object(), object()]}, mm_kwargs) + + assert result is mm_kwargs + assert result["use_audio_in_video"] is True + + +def test_coerce_use_audio_in_video_for_hf_processor_hides_per_video_mask_from_hf(): + mm_kwargs = {"use_audio_in_video": [True, False], "other": "kept"} + + result = _coerce_use_audio_in_video_for_hf_processor({"videos": [object(), object()]}, mm_kwargs) + + assert result is not mm_kwargs + assert result["use_audio_in_video"] is False + assert result[_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY] == [True, False] + assert result["other"] == "kept" + assert mm_kwargs["use_audio_in_video"] == [True, False] + + +def test_coerce_use_audio_in_video_for_hf_processor_does_not_validate_against_empty_mm_data(): + result = _coerce_use_audio_in_video_for_hf_processor( + {}, + {"use_audio_in_video": [True, False]}, + ) + + assert result["use_audio_in_video"] is False + assert result[_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY] == [True, False] + + +def test_coerce_use_audio_in_video_for_hf_processor_does_not_validate_against_partial_mm_data(): + result = _coerce_use_audio_in_video_for_hf_processor( + {"videos": [object()]}, + {"use_audio_in_video": [True, False]}, + ) + + assert result["use_audio_in_video"] is False + assert result[_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY] == [True, False] + + +def test_filter_video_use_audio_in_video_for_uncached_items_aligns_partial_cache_mask(): + result = _filter_video_use_audio_in_video_for_uncached_items( + { + "use_audio_in_video": [True, False, True], + "second_per_grid_ts": [0.5, 1.0, 1.5], + }, + [True, False, False], + ) + + assert result["use_audio_in_video"] == [False, True] + assert result[_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY] == [False, True] + assert result["second_per_grid_ts"] == [1.0, 1.5] + + +def test_filter_video_use_audio_in_video_for_uncached_items_validates_request_mask_length(): + with pytest.raises(ValueError, match="one boolean per video"): + _filter_video_use_audio_in_video_for_uncached_items( + {"use_audio_in_video": [True]}, + [False, False], + ) + + +def test_mm_fields_config_keeps_global_use_audio_in_video_shared(): + fake_self = _fake_processor_with_spatial_merge_size() + hf_inputs = { + "video_grid_thw": torch.tensor([[1, 4, 4], [1, 4, 4]]), + "use_audio_in_video": torch.tensor(False), + } + + result = Qwen2_5OmniThinkerMultiModalProcessor._get_mm_fields_config( + fake_self, + hf_inputs, + {}, + ) + + assert isinstance(result["use_audio_in_video"].field, MultiModalSharedField) + + +def test_mm_fields_config_uses_batched_field_for_per_video_audio_mask(): + fake_self = _fake_processor_with_spatial_merge_size() + hf_inputs = { + "video_grid_thw": torch.tensor([[1, 4, 4], [1, 4, 4]]), + "use_audio_in_video": torch.tensor([True, False]), + } + + result = Qwen2_5OmniThinkerMultiModalProcessor._get_mm_fields_config( + fake_self, + hf_inputs, + {_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY: [True, False]}, + ) + + assert isinstance(result["use_audio_in_video"].field, MultiModalBatchedField) + + +def test_mm_fields_config_uses_batched_field_for_vector_hf_audio_mask(): + fake_self = _fake_processor_with_spatial_merge_size() + hf_inputs = { + "video_grid_thw": torch.tensor([[1, 4, 4], [1, 4, 4]]), + "use_audio_in_video": torch.tensor([True, False]), + } + + result = Qwen2_5OmniThinkerMultiModalProcessor._get_mm_fields_config( + fake_self, + hf_inputs, + {}, + ) + + assert isinstance(result["use_audio_in_video"].field, MultiModalBatchedField) + + +def test_get_video_use_audio_in_video_preserves_explicit_false(): + fake_self = _fake_processor() + mm_kwargs = { + "video": [ + {"use_audio_in_video": _Feature(torch.tensor(True))}, + {"use_audio_in_video": _Feature(torch.tensor(False))}, + ] + } + + result = Qwen2_5OmniThinkerMultiModalProcessor._get_video_use_audio_in_video( + fake_self, + mm_kwargs, + {}, + ) + + assert result == [True, False] + + +def test_get_video_use_audio_in_video_treats_missing_item_as_false_with_explicit_mask(): + fake_self = _fake_processor() + mm_kwargs = { + "video": [ + {"use_audio_in_video": _Feature(torch.tensor(True))}, + None, + ] + } + + result = Qwen2_5OmniThinkerMultiModalProcessor._get_video_use_audio_in_video( + fake_self, + mm_kwargs, + {}, + ) + + assert result == [True, False] + + +def test_get_video_use_audio_in_video_falls_back_for_partial_cache_hit_with_explicit_mask(): + fake_self = _fake_processor() + update_with_audio = SimpleNamespace(content=SimpleNamespace(full=[VIDEO_TOKEN_ID, AUDIO_TOKEN_ID])) + mm_kwargs = { + "video": [ + {"use_audio_in_video": _Feature(torch.tensor(True))}, + None, + ] + } + + result = Qwen2_5OmniThinkerMultiModalProcessor._get_video_use_audio_in_video( + fake_self, + mm_kwargs, + {"video": [[object()], [update_with_audio]]}, + ) + + assert result == [True, True] + + +def test_qwen2_5_prompt_updates_apply_audio_mask_per_video(): + fake_self = _fake_qwen2_prompt_processor() + out_mm_kwargs = SimpleNamespace( + get_data=lambda **_kwargs: { + "audio_feature_lengths": torch.tensor([100]), + "video_grid_thw": torch.tensor([[1, 1, 2], [1, 1, 2]]), + } + ) + mm_items = SimpleNamespace( + get_all_counts=lambda: {"video": 2, "audio": 1}, + get_items=lambda *_args, **_kwargs: SimpleNamespace(get=lambda _idx: [0] * 100), + ) + + updates = Qwen2_5OmniThinkerMultiModalProcessor._get_prompt_updates( + fake_self, + mm_items, + {"use_audio_in_video": [True, False]}, + out_mm_kwargs, + ) + video_update = next(update for update in updates if update.modality == "video") + + assert video_update.resolve(0).content.full == [ + AUDIO_BOS_TOKEN_ID, + VIDEO_TOKEN_ID, + AUDIO_TOKEN_ID, + AUDIO_EOS_TOKEN_ID, + ] + assert video_update.resolve(1).content.full == [VIDEO_TOKEN_ID, VIDEO_TOKEN_ID] + + +def test_qwen2_5_prompt_updates_prefer_hf_second_per_grid_ts_from_outputs(): + captured = {} + fake_self = _fake_qwen2_prompt_processor() + + def omni_get_updates_use_audio_in_video(**kwargs): + captured["video_second_per_grid_t"] = kwargs["video_second_per_grid_t"] + return [VIDEO_TOKEN_ID] + + fake_self.omni_get_updates_use_audio_in_video = omni_get_updates_use_audio_in_video + out_mm_kwargs = SimpleNamespace( + get_data=lambda **_kwargs: { + "audio_feature_lengths": torch.tensor([100]), + "video_grid_thw": torch.tensor([[1, 1, 1]]), + "second_per_grid_ts": torch.tensor([0.25]), + } + ) + mm_items = SimpleNamespace( + get_all_counts=lambda: {"video": 1, "audio": 1}, + get_items=lambda *_args, **_kwargs: SimpleNamespace(get=lambda _idx: [0] * 100), + ) + + updates = Qwen2_5OmniThinkerMultiModalProcessor._get_prompt_updates( + fake_self, + mm_items, + {"use_audio_in_video": [True], "second_per_grid_ts": torch.tensor([0.5])}, + out_mm_kwargs, + ) + video_update = next(update for update in updates if update.modality == "video") + video_update.resolve(0) + + assert captured["video_second_per_grid_t"] == 0.25 + + +def test_qwen2_5_prompt_updates_accept_tensor_second_per_grid_ts_from_kwargs(): + captured = {} + fake_self = _fake_qwen2_prompt_processor() + + def omni_get_updates_use_audio_in_video(**kwargs): + captured["video_second_per_grid_t"] = kwargs["video_second_per_grid_t"] + return [VIDEO_TOKEN_ID] + + fake_self.omni_get_updates_use_audio_in_video = omni_get_updates_use_audio_in_video + out_mm_kwargs = SimpleNamespace( + get_data=lambda **_kwargs: { + "audio_feature_lengths": torch.tensor([100]), + "video_grid_thw": torch.tensor([[1, 1, 1]]), + } + ) + mm_items = SimpleNamespace( + get_all_counts=lambda: {"video": 1, "audio": 1}, + get_items=lambda *_args, **_kwargs: SimpleNamespace(get=lambda _idx: [0] * 100), + ) + + updates = Qwen2_5OmniThinkerMultiModalProcessor._get_prompt_updates( + fake_self, + mm_items, + {"use_audio_in_video": [True], "second_per_grid_ts": torch.tensor([0.75, 1.0])}, + out_mm_kwargs, + ) + video_update = next(update for update in updates if update.modality == "video") + video_update.resolve(0) + + assert captured["video_second_per_grid_t"] == 0.75 + + +def test_qwen2_5_prompt_updates_validate_mask_against_request_video_count(): + fake_self = _fake_qwen2_prompt_processor() + out_mm_kwargs = SimpleNamespace( + get_data=lambda **_kwargs: { + "video_grid_thw": torch.tensor([[1, 1, 2], [1, 1, 2]]), + } + ) + mm_items = SimpleNamespace(get_all_counts=lambda: {"video": 2}) + + with pytest.raises(ValueError, match="one boolean per video"): + Qwen2_5OmniThinkerMultiModalProcessor._get_prompt_updates( + fake_self, + mm_items, + {"use_audio_in_video": [True]}, + out_mm_kwargs, + ) + + +def test_qwen2_5_mm_only_dummy_counts_subtract_only_videos_using_audio(): + captured_counts = {} + fake_self = SimpleNamespace( + dummy_inputs=SimpleNamespace(get_dummy_text=lambda counts: captured_counts.update(counts) or "dummy"), + _apply_hf_processor_text_mm=lambda **_kwargs: ([], {}, False), + ) + mm_items = SimpleNamespace(get_all_counts=lambda: {"video": 2, "audio": 1}) + + Qwen2_5OmniThinkerMultiModalProcessor._apply_hf_processor_mm_only( + fake_self, + mm_items, + {"use_audio_in_video": [True, False]}, + {}, + ) + + assert captured_counts == {"video": 2, "audio": 0} + + +def test_get_video_use_audio_in_video_falls_back_to_prompt_updates_for_cache_hits(): + fake_self = _fake_processor() + update_with_audio = SimpleNamespace(content=SimpleNamespace(full=[VIDEO_TOKEN_ID, AUDIO_TOKEN_ID])) + update_without_audio = SimpleNamespace(content=SimpleNamespace(full=[VIDEO_TOKEN_ID])) + + result = Qwen2_5OmniThinkerMultiModalProcessor._get_video_use_audio_in_video( + fake_self, + {"video": [None, None]}, + {"video": [[update_with_audio], [update_without_audio]]}, + ) + + assert result == [True, False] + + +def test_derive_audio_from_video_placeholders_only_pairs_true_videos(): + fake_self = _fake_processor() + video_placeholders = [ + PlaceholderFeaturesInfo( + modality="video", + item_idx=0, + start_idx=5, + tokens=[VIDEO_TOKEN_ID, AUDIO_TOKEN_ID, VIDEO_TOKEN_ID], + is_embed=torch.tensor([True, True, True]), + ), + PlaceholderFeaturesInfo( + modality="video", + item_idx=1, + start_idx=12, + tokens=[VIDEO_TOKEN_ID, VIDEO_TOKEN_ID], + is_embed=torch.tensor([True, True]), + ), + ] + + result = Qwen2_5OmniThinkerMultiModalProcessor._derive_audio_from_video_placeholders( + fake_self, + {"video": video_placeholders}, + {"audio": [[object()]]}, + [True, False], + ) + + assert len(result["audio"]) == 1 + assert result["audio"][0].item_idx == 0 + assert result["audio"][0].start_idx == 5 + assert result["audio"][0].is_embed.tolist() == [False, True, False] + + assert len(result["video"]) == 2 + assert result["video"][0].item_idx == 0 + assert result["video"][0].is_embed.tolist() == [True, False, True] + assert result["video"][1].item_idx == 1 + assert result["video"][1].is_embed.tolist() == [True, True] + + +def test_qwen3_processor_inherits_vllm_omni_per_video_helpers(): + assert issubclass( + Qwen3OmniMoeThinkerMultiModalProcessor, + Qwen2_5OmniThinkerMultiModalProcessor, + ) + + +def test_qwen3_prompt_updates_prefer_hf_second_per_grid_ts_from_outputs(): + captured = {} + fake_self = _fake_qwen3_prompt_processor() + + def get_updates_use_audio_in_video(**kwargs): + captured["video_second_per_grid_t"] = kwargs["video_second_per_grid_t"] + return [VIDEO_TOKEN_ID] + + fake_self.get_updates_use_audio_in_video = get_updates_use_audio_in_video + out_mm_kwargs = SimpleNamespace( + get_data=lambda **_kwargs: { + "audio_feature_lengths": torch.tensor([100]), + "video_grid_thw": torch.tensor([[1, 1, 1]]), + "second_per_grid_ts": torch.tensor([0.25]), + } + ) + mm_items = SimpleNamespace( + get_items=lambda *_args, **_kwargs: SimpleNamespace(get=lambda _idx: [0] * 100), + ) + + updates = Qwen3OmniMoeThinkerMultiModalProcessor._get_prompt_updates( + fake_self, + mm_items, + {"use_audio_in_video": [True], "second_per_grid_ts": torch.tensor([0.5])}, + out_mm_kwargs, + ) + video_update = next(update for update in updates if update.modality == "video") + video_update.resolve(0) + + assert captured["video_second_per_grid_t"] == 0.25 + + +def test_qwen3_prompt_updates_accept_tensor_second_per_grid_ts_from_kwargs(): + captured = {} + fake_self = _fake_qwen3_prompt_processor() + + def get_updates_use_audio_in_video(**kwargs): + captured["video_second_per_grid_t"] = kwargs["video_second_per_grid_t"] + return [VIDEO_TOKEN_ID] + + fake_self.get_updates_use_audio_in_video = get_updates_use_audio_in_video + out_mm_kwargs = SimpleNamespace( + get_data=lambda **_kwargs: { + "audio_feature_lengths": torch.tensor([100]), + "video_grid_thw": torch.tensor([[1, 1, 1]]), + } + ) + mm_items = SimpleNamespace( + get_items=lambda *_args, **_kwargs: SimpleNamespace(get=lambda _idx: [0] * 100), + ) + + updates = Qwen3OmniMoeThinkerMultiModalProcessor._get_prompt_updates( + fake_self, + mm_items, + {"use_audio_in_video": [True], "second_per_grid_ts": torch.tensor([0.75, 1.0])}, + out_mm_kwargs, + ) + video_update = next(update for update in updates if update.modality == "video") + video_update.resolve(0) + + assert captured["video_second_per_grid_t"] == 0.75 + + +def test_qwen3_derive_audio_from_video_placeholders_excludes_audio_wrappers_from_embeds(): + fake_self = _fake_qwen3_processor() + video_placeholders = [ + PlaceholderFeaturesInfo( + modality="video", + item_idx=0, + start_idx=5, + tokens=[ + AUDIO_BOS_TOKEN_ID, + VIDEO_TOKEN_ID, + AUDIO_TOKEN_ID, + AUDIO_EOS_TOKEN_ID, + VIDEO_TOKEN_ID, + ], + is_embed=torch.ones(5, dtype=torch.bool), + ) + ] + + result = Qwen3OmniMoeThinkerMultiModalProcessor._derive_audio_from_video_placeholders( + fake_self, + {"video": video_placeholders}, + {"audio": [[object()]]}, + [True], + ) + + assert result["audio"][0].is_embed.tolist() == [ + False, + False, + True, + False, + False, + ] + assert result["video"][0].is_embed.tolist() == [ + False, + True, + False, + False, + True, + ] + + +def test_qwen3_mrope_positions_match_plain_video_placeholder_length(): + fake_self = SimpleNamespace( + iter_mm_features=lambda _: iter( + [ + ( + 2, + "video", + { + "grid_t": 1, + "grid_h": 2, + "grid_w": 2, + "t_factor": 1.0, + "use_audio_in_video": False, + "placeholder_len": 4, + }, + ) + ] + ) + ) + + positions, _ = Qwen3OmniMoeThinkerForConditionalGeneration.get_mrope_input_positions( + fake_self, + list(range(8)), + [], + ) + + assert positions.shape == (3, 8) + + +def test_qwen3_mrope_positions_count_single_interleaved_audio_wrappers(): + fake_self = SimpleNamespace( + iter_mm_features=lambda _: iter( + [ + ( + 2, + "video", + { + "grid_t": 1, + "grid_h": 2, + "grid_w": 2, + "t_factor": 1.0, + "use_audio_in_video": True, + "placeholder_len": 8, + }, + ) + ] + ), + _compute_interleaved_positions=lambda _start_idx, _data: ( + torch.zeros((3, 6), dtype=torch.long).numpy(), + 6, + ), + ) + + positions, _ = Qwen3OmniMoeThinkerForConditionalGeneration.get_mrope_input_positions( + fake_self, + list(range(10)), + [], + ) + + assert positions.shape == (3, 10) + + +def test_qwen3_interleaved_mrope_prefers_placeholder_length_for_audio_count(): + fake_self = SimpleNamespace( + _compute_audio_token_count=lambda _audio_feature_length: 100, + ) + data = { + "grid_t": 1, + "grid_h": 2, + "grid_w": 2, + "t_factor": 1.0, + "audio_feature_length": 999, + "placeholder_len": 8, + } + + positions, token_count = Qwen3OmniMoeThinkerForConditionalGeneration._compute_interleaved_positions( + fake_self, + 0, + data, + ) + + assert token_count == 6 + assert positions.shape == (3, 6) + + +def test_omni_model_runner_output_can_carry_kv_connector_output_only(): + marker = object() + + output = OmniModelRunnerOutput.with_kv_conn_output_only(marker) + + assert output.req_ids == [] + assert output.req_id_to_index == {} + assert output.sampled_token_ids == [] + assert output.kv_connector_output is marker 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 d0df1c2ea..03edcfa10 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,6 +1,7 @@ """Thin Omni wrapper: reuse upstream Qwen2.5-Omni thinker with minimal overrides.""" -from collections.abc import Iterable, Iterator, Mapping +from collections.abc import Iterable, Iterator, Mapping, Sequence +from functools import partial from typing import Any import numpy as np @@ -28,6 +29,7 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import ( Qwen2_5OmniAudioFeatureInputs, Qwen2_5OmniThinkerDummyInputsBuilder, check_interleaved_audio_video, + create_qwen2_5_omni_thinker_field_factory, merge_interleaved_embeddings, ) from vllm.model_executor.models.qwen2_5_omni_thinker import ( @@ -51,6 +53,9 @@ from vllm.model_executor.models.qwen2_5_vl import ( Qwen2_5_VLVideoInputs, Qwen2_5_VLVideoPixelInputs, ) +from vllm.model_executor.models.qwen2_audio import ( + _get_feat_extract_output_lengths, +) from vllm.model_executor.models.utils import ( AutoWeightsLoader, WeightsMapper, @@ -60,12 +65,19 @@ from vllm.model_executor.models.utils import ( from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFeatureSpec, + MultiModalFieldConfig, MultiModalKwargsItems, ) -from vllm.multimodal.parse import MultiModalDataItems, VideoProcessorItems +from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems, VideoProcessorItems +from vllm.multimodal.processing.context import TimingContext +from vllm.multimodal.processing.inputs import ProcessorInputs from vllm.multimodal.processing.processor import ( + MultiModalProcessingInfo, MultiModalPromptUpdates, PlaceholderFeaturesInfo, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, ) from vllm.sequence import IntermediateTensors from vllm.utils.collection_utils import is_list_of @@ -80,6 +92,158 @@ except (ImportError, ModuleNotFoundError): flash_attn = None logger = init_logger(__name__) +_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY = "_vllm_omni_per_video_use_audio_in_video" + + +def _normalize_use_audio_in_video( + value: object, + num_videos: int, +) -> list[bool]: + if isinstance(value, torch.Tensor): + if value.numel() == 1: + return [bool(value.item())] * num_videos + values = value.flatten().tolist() + elif isinstance(value, np.ndarray): + if value.size == 1: + return [bool(value.item())] * num_videos + values = value.flatten().tolist() + elif isinstance(value, bool): + return [value] * num_videos + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + values = list(value) + else: + return [bool(value)] * num_videos + + if len(values) != num_videos: + raise ValueError( + "use_audio_in_video must be a boolean or contain one boolean per " + f"video, but found {len(values)} values for {num_videos} videos." + ) + return [bool(v) for v in values] + + +def _normalize_per_video_use_audio_in_video(value: object) -> list[bool]: + if isinstance(value, torch.Tensor): + values = value.flatten().tolist() + elif isinstance(value, np.ndarray): + values = value.flatten().tolist() + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + values = list(value) + else: + values = [value] + return [bool(v) for v in values] + + +def _is_per_video_use_audio_in_video(value: object) -> bool: + if isinstance(value, bool): + return False + if isinstance(value, torch.Tensor): + return value.numel() != 1 + if isinstance(value, np.ndarray): + return value.size != 1 + return isinstance(value, Sequence) and not isinstance(value, (str, bytes)) + + +def _get_request_video_use_audio_in_video( + hf_processor_mm_kwargs: Mapping[str, object], + num_videos: int, +) -> list[bool]: + value = hf_processor_mm_kwargs.get( + _PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY, + hf_processor_mm_kwargs.get("use_audio_in_video", False), + ) + return _normalize_use_audio_in_video(value, num_videos) + + +def _prompt_update_has_audio_token( + mm_prompt_updates: MultiModalPromptUpdates, + item_idx: int, + audio_token_id: int, +) -> bool: + updates = mm_prompt_updates.get("video", []) + if item_idx >= len(updates): + return False + return any( + audio_token_id in update.content.full for update in updates[item_idx] if isinstance(update.content.full, list) + ) + + +def _get_video_second_per_grid_t( + out_mm_data: Mapping[str, object], + hf_processor_mm_kwargs: Mapping[str, object], + item_idx: int, + default: float, +) -> float: + second_per_grid_ts = out_mm_data.get("second_per_grid_ts") + if second_per_grid_ts is None: + second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None) + if second_per_grid_ts is None: + return default + + if isinstance(second_per_grid_ts, torch.Tensor): + if second_per_grid_ts.numel() == 1: + return float(second_per_grid_ts.item()) + return float(second_per_grid_ts.flatten()[item_idx].item()) + + if isinstance(second_per_grid_ts, np.ndarray): + if second_per_grid_ts.size == 1: + return float(second_per_grid_ts.item()) + return float(second_per_grid_ts.flatten()[item_idx].item()) + + if isinstance(second_per_grid_ts, Sequence) and not isinstance(second_per_grid_ts, (str, bytes)): + return float(second_per_grid_ts[item_idx]) + + return float(second_per_grid_ts) + + +def _filter_video_use_audio_in_video_for_uncached_items( + hf_processor_mm_kwargs: Mapping[str, object], + video_is_cached: Sequence[bool] | None, +) -> Mapping[str, object]: + use_audio_in_video = hf_processor_mm_kwargs.get( + _PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY, + hf_processor_mm_kwargs.get("use_audio_in_video"), + ) + if video_is_cached is None or not _is_per_video_use_audio_in_video(use_audio_in_video): + return hf_processor_mm_kwargs + + mask = _normalize_use_audio_in_video(use_audio_in_video, len(video_is_cached)) + missing_mask = [use_audio for use_audio, is_cached in zip(mask, video_is_cached) if not is_cached] + + filtered_kwargs = dict(hf_processor_mm_kwargs) + filtered_kwargs["use_audio_in_video"] = missing_mask + filtered_kwargs[_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY] = missing_mask + second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts") + if ( + isinstance(second_per_grid_ts, Sequence) + and not isinstance(second_per_grid_ts, (str, bytes)) + and len(second_per_grid_ts) == len(video_is_cached) + ): + filtered_kwargs["second_per_grid_ts"] = [ + second_per_grid_t + for second_per_grid_t, is_cached in zip(second_per_grid_ts, video_is_cached) + if not is_cached + ] + return filtered_kwargs + + +def _coerce_use_audio_in_video_for_hf_processor( + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], +) -> Mapping[str, object]: + use_audio_in_video = mm_kwargs.get("use_audio_in_video") + if not _is_per_video_use_audio_in_video(use_audio_in_video): + return mm_kwargs + + video_use_audio_in_video = _normalize_per_video_use_audio_in_video(use_audio_in_video) + + hf_mm_kwargs = dict(mm_kwargs) + # HF processors only support a global bool. For per-video masks, vLLM + # consumes the list while building prompt updates and placeholders. + hf_mm_kwargs["use_audio_in_video"] = False + hf_mm_kwargs[_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY] = video_use_audio_in_video + return hf_mm_kwargs + def _presampled_videos_hf_kwargs( mm_data: Mapping[str, object], @@ -87,12 +251,9 @@ def _presampled_videos_hf_kwargs( ) -> Mapping[str, object]: """Adjust HF video kwargs for videos pre-sampled by vLLM's video loader. - When ``video_metadata`` is present (emitted by - ``Qwen2_5OmniVideoProcessorItems``), the frames were already sampled - according to ``media_io_kwargs``, so the HF processor must not re-sample - them, and it needs the sampled fps (instead of its default) to compute - ``video_second_per_grid`` correctly. Otherwise the audio/video temporal - alignment breaks under ``use_audio_in_video=True``. + When ``video_metadata`` is present, the frames were already sampled + according to ``media_io_kwargs``. The HF processor should not resample + them and should use the sampled fps when computing temporal metadata. """ video_metadata = mm_data.get("video_metadata") if not video_metadata: @@ -109,18 +270,17 @@ def _presampled_videos_hf_kwargs( return None return len(indices) / float(duration) - # An explicit user-provided fps takes precedence. if "fps" not in videos_kwargs: fps_values = [_compute_sampled_video_fps(m) for m in video_metadata] - # HF accepts a single fps per call and uses it for every video's - # video_second_per_grid. known_fps = [fps for fps in fps_values if fps is not None] unique_fps = set(known_fps) if len(unique_fps) == 1: videos_kwargs["fps"] = known_fps[0] elif len(unique_fps) > 1: logger.warning( - f"Mixed sampled FPS {sorted(unique_fps)} in one request; HF accepts a single fps, using {known_fps[0]}." + "Mixed sampled FPS %s in one request; HF accepts a single fps, using %s.", + sorted(unique_fps), + known_fps[0], ) videos_kwargs["fps"] = known_fps[0] @@ -129,12 +289,7 @@ def _presampled_videos_hf_kwargs( class Qwen2_5OmniVideoProcessorItems(VideoProcessorItems): - """Video items that carry the loader's ``(frames, metadata)`` tuples. - - The tuples are kept in ``data`` so that mm hashing and cache-miss - re-parsing retain the metadata; ``get_processor_data`` unpacks them into - the ``videos`` + ``video_metadata`` arguments of the HF processor. - """ + """Video items that carry the loader's ``(frames, metadata)`` tuples.""" def get_processor_data(self) -> Mapping[str, object]: from transformers.video_utils import VideoMetadata @@ -151,7 +306,6 @@ class Qwen2_5OmniThinkerMultiModalDataParser(_Qwen2_5OmniThinkerMultiModalDataPa if data is None or isinstance(data, dict) or self.is_embeddings(data): return super()._parse_video_data(data) - # Normalize to a list of per-video items (same as the base parser). if (is_list_of(data, PILImage.Image) and len(data) > 0) or ( isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 4 ): @@ -164,11 +318,6 @@ class Qwen2_5OmniThinkerMultiModalDataParser(_Qwen2_5OmniThinkerMultiModalDataPa data_items = data videos_with_metadata = [self._get_video_with_metadata(item) for item in data_items] - - # Metadata is optional: videos fetched by vLLM's media connector - # arrive as (frames, metadata) tuples, while plain arrays (e.g. - # offline `multi_modal_data` or dummy profiling inputs) carry no - # metadata and parse as before. if any(metadata is None for _, metadata in videos_with_metadata): return super()._parse_video_data(data) @@ -195,14 +344,6 @@ class Qwen2_5OmniThinkerMultiModalProcessor( ): """Override to fix use_audio_in_video detection when mm cache returns None.""" - def _cached_apply_hf_processor(self, inputs, timing_ctx): - # If use_audio_in_video, process video and audio together; otherwise, - # skip cache to avoid errors. - # Prevents partial cache hits from causing processing failures. - if inputs.hf_processor_mm_kwargs.get("use_audio_in_video"): - return self._apply_hf_processor(inputs, timing_ctx) - return super()._cached_apply_hf_processor(inputs, timing_ctx) - def _call_hf_processor( self, prompt: str, @@ -211,12 +352,341 @@ class Qwen2_5OmniThinkerMultiModalProcessor( tok_kwargs: Mapping[str, object], ): mm_kwargs = _presampled_videos_hf_kwargs(mm_data, mm_kwargs) - return super()._call_hf_processor( + mm_kwargs = _coerce_use_audio_in_video_for_hf_processor(mm_data, mm_kwargs) + hf_inputs = super()._call_hf_processor( prompt=prompt, mm_data=mm_data, mm_kwargs=mm_kwargs, tok_kwargs=tok_kwargs, ) + per_video_mask = mm_kwargs.get(_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY) + if per_video_mask is not None: + hf_inputs["use_audio_in_video"] = torch.tensor(per_video_mask) + return hf_inputs + + def _get_mm_fields_config( + self, + hf_inputs, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + fields_config = dict( + create_qwen2_5_omni_thinker_field_factory(self.info.get_hf_config().vision_config.spatial_merge_size)( + hf_inputs + ) + ) + use_audio_in_video = hf_inputs.get("use_audio_in_video") + is_batched_mask = isinstance(use_audio_in_video, torch.Tensor) and use_audio_in_video.numel() > 1 + if _PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY in hf_processor_mm_kwargs or is_batched_mask: + fields_config["use_audio_in_video"] = MultiModalFieldConfig.batched("video") + return fields_config + + def _get_video_use_audio_in_video( + self, + mm_kwargs: MultiModalKwargsItems, + mm_prompt_updates: MultiModalPromptUpdates, + ) -> list[bool]: + video_kwargs = mm_kwargs.get("video", []) + if not video_kwargs: + return [] + + audio_token_id = self.info.get_hf_config().audio_token_id + video_use_audio_in_video = [] + has_use_audio_in_video = any(item is not None and "use_audio_in_video" in item for item in video_kwargs) + for item_idx, item in enumerate(video_kwargs): + if has_use_audio_in_video: + if item is None or "use_audio_in_video" not in item: + video_use_audio_in_video.append( + _prompt_update_has_audio_token( + mm_prompt_updates, + item_idx, + audio_token_id, + ) + ) + continue + use_audio_in_video_tensor = item["use_audio_in_video"].data + if use_audio_in_video_tensor.numel() > 0: + video_use_audio_in_video.append(bool(use_audio_in_video_tensor.item())) + continue + video_use_audio_in_video.append(False) + continue + + video_use_audio_in_video.append( + _prompt_update_has_audio_token( + mm_prompt_updates, + item_idx, + audio_token_id, + ) + ) + + return video_use_audio_in_video + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, Any], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + tokenizer = self.info.get_tokenizer() + image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs) + vocab = tokenizer.get_vocab() + + audio_token = processor.audio_token + image_token = processor.image_token + video_token = processor.video_token + audio_token_id = vocab[audio_token] + image_token_id = vocab[image_token] + video_token_id = vocab[video_token] + + out_mm_data = out_mm_kwargs.get_data() + audio_feature_lengths = out_mm_data.get("audio_feature_lengths") + feature_attention_mask = out_mm_data.get("feature_attention_mask") + if audio_feature_lengths is None and feature_attention_mask is None: + audio_output_lengths = [] + elif audio_feature_lengths is not None: + _, audio_output_lens = _get_feat_extract_output_lengths(audio_feature_lengths) + audio_output_lengths = audio_output_lens.tolist() + elif feature_attention_mask is not None: + assert isinstance(feature_attention_mask, torch.Tensor) + _, audio_output_lens = _get_feat_extract_output_lengths(feature_attention_mask.sum(-1)) + audio_output_lengths = audio_output_lens.tolist() + + audio_in_video_item_idx = 0 + + def get_replacement_qwen2_audio(item_idx: int): + item_idx += audio_in_video_item_idx + + num_features = audio_output_lengths[item_idx] + if num_features == 0: + audios = mm_items.get_items("audio", AudioProcessorItems) + audio = audios.get(item_idx) + raise ValueError( + f"The audio {audio} (len={len(audio)}) is too short to be represented inside the model" + ) + + return [audio_token_id] * num_features + + def get_replacement_qwen2_vision(item_idx: int, modality: str): + grid_thw = out_mm_data[f"{modality}_grid_thw"][item_idx] + assert isinstance(grid_thw, torch.Tensor) + merge_length = image_processor.merge_size**2 + + token_id = image_token_id if modality == "image" else video_token_id + return [token_id] * (int(grid_thw.prod()) // merge_length) + + num_videos = mm_items.get_all_counts().get("video", 0) + video_use_audio_in_video = _get_request_video_use_audio_in_video( + hf_processor_mm_kwargs, + num_videos, + ) + thinker_config = self.info.get_hf_config() + + def get_replacement_qwen2_video(item_idx: int): + nonlocal audio_in_video_item_idx + + if not video_use_audio_in_video[item_idx]: + return get_replacement_qwen2_vision(item_idx, modality="video") + + audio_num_features = audio_output_lengths[audio_in_video_item_idx] + video_grid_thw = out_mm_data["video_grid_thw"][item_idx] + + audio_in_video_item_idx += 1 + + video_second_per_grid_t = _get_video_second_per_grid_t( + out_mm_data, + hf_processor_mm_kwargs, + item_idx, + default=1.0, + ) + + updates = self.omni_get_updates_use_audio_in_video( + thinker_config=thinker_config, + audio_len=audio_num_features, + video_grid_thw=video_grid_thw, + video_second_per_grid_t=video_second_per_grid_t, + ) + + return PromptUpdateDetails.select_token_id( + seq=updates, + embed_token_id=video_token_id, + ) + + return [ + PromptReplacement( + modality="audio", + target=audio_token, + replacement=get_replacement_qwen2_audio, + ), + PromptReplacement( + modality="image", + target=image_token, + replacement=partial(get_replacement_qwen2_vision, modality="image"), + ), + PromptReplacement( + modality="video", + target=video_token, + replacement=get_replacement_qwen2_video, + ), + ] + + def _apply_hf_processor_mm_only( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ): + mm_counts = mm_items.get_all_counts() + + if "video" in mm_counts: + video_use_audio_in_video = _get_request_video_use_audio_in_video( + hf_processor_mm_kwargs, + mm_counts["video"], + ) + if any(video_use_audio_in_video): + assert "audio" in mm_counts + mm_counts["audio"] -= sum(video_use_audio_in_video) + + _, mm_processed_data, _ = self._apply_hf_processor_text_mm( + prompt_text=self.dummy_inputs.get_dummy_text(mm_counts), + mm_items=mm_items, + hf_processor_mm_kwargs=hf_processor_mm_kwargs, + tokenization_kwargs=tokenization_kwargs, + ) + + return mm_processed_data + + def _cached_apply_hf_processor( + self, + inputs: ProcessorInputs, + timing_ctx: TimingContext, + ) -> tuple[list[int], MultiModalProcessingInfo, bool]: + 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) + + with timing_ctx.record("get_mm_hashes"): + mm_hashes = inputs.get_mm_hashes(self.info.model_id) + + 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, + ) + + hf_processor_mm_kwargs = _filter_video_use_audio_in_video_for_uncached_items( + inputs.hf_processor_mm_kwargs, + mm_is_cached.get("video"), + ) + + with timing_ctx.record("apply_hf_processor"): + ( + prompt_ids, + mm_missing_processed_data, + is_update_applied, + ) = self._apply_hf_processor_main( + prompt=inputs.prompt, + mm_items=mm_missing_data_items, + hf_processor_mm_kwargs=hf_processor_mm_kwargs, + tokenization_kwargs=inputs.tokenization_kwargs, + enable_hf_prompt_update=False, + ) + + mm_missing_kwargs = MultiModalKwargsItems.from_hf_inputs( + mm_missing_processed_data, + self._get_mm_fields_config( + mm_missing_processed_data, + hf_processor_mm_kwargs, + ), + ) + + mm_missing_prompt_updates = self._get_mm_prompt_updates( + mm_missing_data_items, + hf_processor_mm_kwargs, + mm_missing_kwargs, + ) + + with timing_ctx.record("merge_mm_kwargs"): + mm_kwargs, mm_prompt_updates = self._merge_mm_kwargs( + cache, + mm_hashes=mm_hashes, + mm_is_cached=mm_is_cached, + mm_missing_kwargs=mm_missing_kwargs, + mm_missing_prompt_updates=mm_missing_prompt_updates, + ) + + mm_info = MultiModalProcessingInfo( + kwargs=mm_kwargs, + hashes=mm_hashes, + prompt_updates=mm_prompt_updates, + ) + + return prompt_ids, mm_info, is_update_applied + + def _derive_audio_from_video_placeholders( + self, + placeholders: Mapping[str, list[PlaceholderFeaturesInfo]], + mm_prompt_updates: MultiModalPromptUpdates, + video_use_audio_in_video: Sequence[bool] | None = None, + ) -> Mapping[str, list[PlaceholderFeaturesInfo]]: + if "video" not in placeholders: + return placeholders + + num_videos = len(placeholders["video"]) + if video_use_audio_in_video is None: + video_use_audio_in_video = [True] * num_videos + elif len(video_use_audio_in_video) != num_videos: + raise ValueError( + "use_audio_in_video must contain one boolean per video, " + f"but found {len(video_use_audio_in_video)} values for " + f"{num_videos} videos." + ) + + num_audio_in_video = sum(video_use_audio_in_video) + num_audios = len(mm_prompt_updates.get("audio", [])) + if num_audios != num_audio_in_video: + raise ValueError( + "use_audio_in_video requires equal number of audio and video " + f"items using audio, got {num_audios=}, {num_audio_in_video=}" + ) + + tokenizer = self.info.get_tokenizer() + processor = self.info.get_hf_processor() + audio_token_id = tokenizer.get_vocab()[processor.audio_token] + + result_placeholders = dict(placeholders) + audio_placeholders = [] + video_placeholders = [] + + audio_idx = 0 + for video_idx, video_placeholder in enumerate(placeholders["video"]): + audio_is_embed = torch.tensor(video_placeholder.tokens) == audio_token_id + + if video_use_audio_in_video[video_idx]: + audio_placeholder = PlaceholderFeaturesInfo( + modality="audio", + item_idx=audio_idx, + start_idx=video_placeholder.start_idx, + tokens=video_placeholder.tokens, + is_embed=audio_is_embed, + ) + audio_placeholders.append(audio_placeholder) + audio_idx += 1 + + video_placeholder_with_mask = PlaceholderFeaturesInfo( + modality="video", + item_idx=video_idx, + start_idx=video_placeholder.start_idx, + tokens=video_placeholder.tokens, + is_embed=~audio_is_embed, + ) + video_placeholders.append(video_placeholder_with_mask) + + result_placeholders["audio"] = audio_placeholders + result_placeholders["video"] = video_placeholders + return result_placeholders def _maybe_apply_prompt_updates( self, @@ -230,24 +700,8 @@ class Qwen2_5OmniThinkerMultiModalProcessor( self._validate_mm_kwargs(mm_kwargs, mm_item_counts) self._validate_mm_updates(mm_prompt_updates, mm_item_counts) - # Detect use_audio_in_video from mm_kwargs - use_audio_in_video = False - if "video" in mm_kwargs: - for item in mm_kwargs["video"]: - if item and item.get("use_audio_in_video"): - use_audio_in_video_tensor = item["use_audio_in_video"].data - if use_audio_in_video_tensor.numel() > 0: - use_audio_in_video = bool(use_audio_in_video_tensor.item()) - break - # for mutilmodality cache - if any(item is None for item in mm_kwargs["video"]): - video_token_id = self.info.get_hf_config().video_token_id - audio_token_id = self.info.get_hf_config().audio_token_id - video_audio_item_num = sum(id in (video_token_id, audio_token_id) for id in prompt_ids) - audio_updates_num = len(mm_prompt_updates.get("audio", [])) - video_updates_num = len(mm_prompt_updates.get("video", [])) - if video_audio_item_num != video_updates_num + audio_updates_num: - use_audio_in_video = True + video_use_audio_in_video = self._get_video_use_audio_in_video(mm_kwargs, mm_prompt_updates) + use_audio_in_video = any(video_use_audio_in_video) if is_update_applied: mm_placeholders = self._find_mm_placeholders( @@ -265,7 +719,11 @@ class Qwen2_5OmniThinkerMultiModalProcessor( prompt_ids, filtered_updates, ) - mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates) + mm_placeholders = self._derive_audio_from_video_placeholders( + mm_placeholders, + mm_prompt_updates, + video_use_audio_in_video, + ) else: prompt_ids, mm_placeholders = self._apply_prompt_updates( prompt_ids, diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index 011496f86..fde493deb 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -121,6 +121,8 @@ from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( Qwen2_5OmniConditionalGenerationMixin, Qwen2_5OmniThinkerMultiModalDataParser, Qwen2_5OmniThinkerMultiModalProcessor, + _get_request_video_use_audio_in_video, + _get_video_second_per_grid_t, _presampled_videos_hf_kwargs, ) from vllm_omni.quantization.component_config import ( @@ -626,9 +628,6 @@ class Qwen3OmniMoeThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProc def get_data_parser(self): feature_extractor = self.get_feature_extractor() - # Install the Omni parser, which keeps video metadata from vLLM's - # video loader (see Qwen2_5OmniVideoProcessorItems); the parser - # inherited from Qwen2AudioProcessingInfo would drop it. return Qwen2_5OmniThinkerMultiModalDataParser( spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size, target_sr=feature_extractor.sampling_rate, @@ -686,10 +685,8 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( tok_kwargs: Mapping[str, object], ) -> BatchFeature: mm_data = dict(mm_data) - mm_kwargs = dict(mm_kwargs) - audios = mm_data.pop("audios", []) - mm_kwargs = dict(_presampled_videos_hf_kwargs(mm_data, mm_kwargs)) + audios = mm_data.pop("audios", []) def pad_to_hop_length(x: np.ndarray, hop_length: int) -> np.ndarray: length = x.shape[-1] @@ -787,22 +784,8 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( mm_item_counts = mm_items.get_all_counts() self._validate_mm_kwargs(mm_kwargs, mm_item_counts) - use_audio_in_video = False - if "video" in mm_kwargs: - for item in mm_kwargs["video"]: - if item and item["use_audio_in_video"].data: - use_audio_in_video = True - else: - use_audio_in_video = False - # for mutilmodality cache - if any(item is None for item in mm_kwargs["video"]): - video_token_id = self.info.get_hf_config().video_token_id - audio_token_id = self.info.get_hf_config().audio_token_id - video_audio_item_num = sum(id in (video_token_id, audio_token_id) for id in prompt_ids) - audio_updates_num = len(mm_prompt_updates.get("audio", [])) - video_updates_num = len(mm_prompt_updates.get("video", [])) - if video_audio_item_num != video_updates_num + audio_updates_num: - use_audio_in_video = True + video_use_audio_in_video = self._get_video_use_audio_in_video(mm_kwargs, mm_prompt_updates) + use_audio_in_video = any(video_use_audio_in_video) # normal case with `use_audio_in_video=False` if is_update_applied: @@ -821,7 +804,11 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( prompt_ids, filtered_updates, ) - mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates) + mm_placeholders = self._derive_audio_from_video_placeholders( + mm_placeholders, + mm_prompt_updates, + video_use_audio_in_video, + ) else: prompt_ids, mm_placeholders = self._apply_prompt_updates( prompt_ids, @@ -935,25 +922,30 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( token_id = image_token_id if modality == "image" else video_token_id return [token_id] * (int(grid_thw.prod()) // merge_length) - use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False) + num_videos = len(out_mm_data.get("video_grid_thw", [])) + video_use_audio_in_video = _get_request_video_use_audio_in_video( + hf_processor_mm_kwargs, + num_videos, + ) thinker_config = self.info.get_hf_config() - def get_replacement_qwen2_use_audio_in_video(item_idx: int): + def get_replacement_qwen2_video(item_idx: int): nonlocal audio_in_video_item_idx + + if not video_use_audio_in_video[item_idx]: + return get_replacement_qwen2_vision(item_idx, modality="video") + audio_num_features = audio_output_lengths[audio_in_video_item_idx] video_grid_thw = out_mm_data["video_grid_thw"][item_idx] audio_in_video_item_idx += 1 - # Prefer the HF-computed value (temporal_patch_size / sampled fps) - # over request kwargs; - second_per_grid_ts = out_mm_data.get("second_per_grid_ts") - if second_per_grid_ts is None: - second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None) - if second_per_grid_ts is not None: - video_second_per_grid_t = float(second_per_grid_ts[item_idx]) - else: - video_second_per_grid_t = 2.0 + video_second_per_grid_t = _get_video_second_per_grid_t( + out_mm_data, + hf_processor_mm_kwargs, + item_idx, + default=2.0, + ) placeholder = self.get_updates_use_audio_in_video( thinker_config=thinker_config, @@ -963,12 +955,6 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( ) return PromptUpdateDetails.select_token_id(placeholder, embed_token_id=video_token_id) - video_replacement_fn = ( - get_replacement_qwen2_use_audio_in_video - if use_audio_in_video - else partial(get_replacement_qwen2_vision, modality="video") - ) - return [ PromptReplacement( modality="audio", @@ -983,7 +969,7 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( PromptReplacement( modality="video", target=video_token, - replacement=video_replacement_fn, + replacement=get_replacement_qwen2_video, ), ] @@ -991,6 +977,7 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( self, placeholders: Mapping[str, list[PlaceholderFeaturesInfo]], mm_prompt_updates: MultiModalPromptUpdates, + video_use_audio_in_video: Sequence[bool] | None = None, ) -> Mapping[str, list[PlaceholderFeaturesInfo]]: """ Helper to derive audio placeholders from video placeholders when @@ -1000,34 +987,62 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( return placeholders num_videos = len(placeholders["video"]) - num_audios = len(mm_prompt_updates.get("audio", [])) - if num_audios != num_videos: + if video_use_audio_in_video is None: + video_use_audio_in_video = [True] * num_videos + elif len(video_use_audio_in_video) != num_videos: raise ValueError( - f"use_audio_in_video requires equal number of audio and video items, got {num_audios=}, {num_videos=}" + "use_audio_in_video must contain one boolean per video, " + f"but found {len(video_use_audio_in_video)} values for " + f"{num_videos} videos." + ) + + num_audio_in_video = sum(video_use_audio_in_video) + num_audios = len(mm_prompt_updates.get("audio", [])) + if num_audios != num_audio_in_video: + raise ValueError( + "use_audio_in_video requires equal number of audio and video " + f"items using audio, got {num_audios=}, {num_audio_in_video=}" ) tokenizer = self.info.get_tokenizer() processor = self.info.get_hf_processor() audio_token_id = tokenizer.get_vocab()[processor.audio_token] + video_token_id = tokenizer.get_vocab()[processor.video_token] result_placeholders = dict(placeholders) audio_placeholders = [] + video_placeholders = [] - # Each video is paired with one audio + # Each video using audio is paired with one audio + audio_idx = 0 for video_idx, video_placeholder in enumerate(placeholders["video"]): # Create is_embed mask selecting only audio tokens - audio_is_embed = torch.tensor(video_placeholder.tokens) == audio_token_id + placeholder_tokens = torch.tensor(video_placeholder.tokens) + audio_is_embed = placeholder_tokens == audio_token_id + video_is_embed = placeholder_tokens == video_token_id - audio_placeholder = PlaceholderFeaturesInfo( - modality="audio", + if video_use_audio_in_video[video_idx]: + audio_placeholder = PlaceholderFeaturesInfo( + modality="audio", + item_idx=audio_idx, + start_idx=video_placeholder.start_idx, + tokens=video_placeholder.tokens, + is_embed=audio_is_embed, + ) + audio_placeholders.append(audio_placeholder) + audio_idx += 1 + + video_placeholder_with_mask = PlaceholderFeaturesInfo( + modality="video", item_idx=video_idx, start_idx=video_placeholder.start_idx, tokens=video_placeholder.tokens, - is_embed=audio_is_embed, + is_embed=video_is_embed, ) - audio_placeholders.append(audio_placeholder) + video_placeholders.append(video_placeholder_with_mask) result_placeholders["audio"] = audio_placeholders + result_placeholders["video"] = video_placeholders return result_placeholders def _get_raw_input_ids( @@ -1529,10 +1544,12 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( Iterate over multimodal features sorted by position offset. Yields: (offset, modality, feature_data) where feature_data contains: - - image: {"grid_t", "grid_h", "grid_w", "t_factor"} + - image: {"grid_t", "grid_h", "grid_w", "t_factor", + "placeholder_len"} - video: {"grid_t", "grid_h", "grid_w", "t_factor", - "use_audio_in_video", "audio_feature_length"} - - audio: {"audio_feature_length"} + "use_audio_in_video", "audio_feature_length", + "placeholder_len"} + - audio: {"audio_feature_length", "placeholder_len"} """ config = self.config spatial_merge_size = config.vision_config.spatial_merge_size @@ -1555,6 +1572,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( "grid_h": h // spatial_merge_size, "grid_w": w // spatial_merge_size, "t_factor": position_id_per_seconds, + "placeholder_len": mm_feature.mm_position.length, }, ) elif modality == "video": @@ -1576,12 +1594,20 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( "t_factor": second_per_grid_ts * position_id_per_seconds, "use_audio_in_video": use_audio_in_video, "audio_feature_length": audio_for_video.get(offset), + "placeholder_len": mm_feature.mm_position.length, }, ) elif modality == "audio": if offset not in paired_audio_offsets: audio_len = mm_feature.data["audio_feature_lengths"].data.item() - yield offset, "audio", {"audio_feature_length": audio_len} + yield ( + offset, + "audio", + { + "audio_feature_length": audio_len, + "placeholder_len": mm_feature.mm_position.length, + }, + ) def _compute_interleaved_positions(self, start_idx: int, data: dict[str, Any]) -> tuple[np.ndarray, int]: """ @@ -1594,9 +1620,13 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( grid_h = data["grid_h"] grid_w = data["grid_w"] t_factor = data["t_factor"] - audio_feature_length = data["audio_feature_length"] - - audio_len = self._compute_audio_token_count(audio_feature_length) + num_video = grid_t * grid_h * grid_w + placeholder_len = data.get("placeholder_len") + if placeholder_len is None: + audio_feature_length = data["audio_feature_length"] + audio_len = self._compute_audio_token_count(audio_feature_length) + else: + audio_len = max(0, int(placeholder_len) - num_video - 2) h_index = np.tile(np.arange(grid_h).reshape(1, -1, 1), (grid_t, 1, grid_w)).flatten() w_index = np.tile(np.arange(grid_w).reshape(1, 1, -1), (grid_t, grid_h, 1)).flatten() @@ -1612,8 +1642,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( pos_ids_list: list[np.ndarray] = [] video_idx, audio_idx = 0, 0 - num_video = grid_t * grid_h * grid_w - while video_idx < num_video and audio_idx < audio_len: if video_t_values[video_idx] <= audio_t_values[audio_idx]: pos_ids_list.append(video_pos[:, video_idx : video_idx + 1]) @@ -1703,9 +1731,18 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( seq_len = len(input_tokens) llm_pos_ids_list: list[np.ndarray] = [] + feature_summaries: list[dict[str, Any]] = [] st = 0 for offset, modality, data in self.iter_mm_features(mm_features): + feature_summaries.append( + { + "offset": offset, + "modality": modality, + "placeholder_len": data.get("placeholder_len"), + "use_audio_in_video": data.get("use_audio_in_video"), + } + ) text_len = offset - st st_idx = int(llm_pos_ids_list[-1].max()) + 1 if llm_pos_ids_list else 0 @@ -1713,19 +1750,11 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( llm_pos_ids_list.append(np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx) st_idx += text_len - bos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) - llm_pos_ids_list.append(bos_pos) - st_idx += 1 - if modality == "audio": audio_tokens = self._compute_audio_token_count(data["audio_feature_length"]) audio_pos = np.broadcast_to(np.arange(audio_tokens), (3, audio_tokens)) + st_idx llm_pos_ids_list.append(audio_pos) - st_idx = int(audio_pos.max()) + 1 - - eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) - llm_pos_ids_list.append(eos_pos) - st = offset + 1 + audio_tokens + 1 + st = offset + int(data.get("placeholder_len", audio_tokens)) elif modality == "image": grid_t = data["grid_t"] @@ -1739,11 +1768,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( llm_pos_ids_list.append(grid_indices.reshape(3, -1) + st_idx) image_len = grid_t * grid_h * grid_w - st_idx = int(llm_pos_ids_list[-1].max()) + 1 - - eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) - llm_pos_ids_list.append(eos_pos) - st = offset + 1 + image_len + 1 + st = offset + int(data.get("placeholder_len", image_len)) elif modality == "video": grid_t = data["grid_t"] @@ -1758,26 +1783,20 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( llm_pos_ids_list.append(grid_indices.reshape(3, -1) + st_idx) video_len = grid_t * grid_h * grid_w - st_idx = int(llm_pos_ids_list[-1].max()) + 1 - - eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) - llm_pos_ids_list.append(eos_pos) - st = offset + 1 + video_len + 1 + st = offset + int(data.get("placeholder_len", video_len)) else: - audio_bos_pos = np.broadcast_to(np.array([st_idx - 1]), (3, 1)) + audio_bos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) llm_pos_ids_list.append(audio_bos_pos) + st_idx += 1 - pos_ids, _ = self._compute_interleaved_positions(st_idx, data) + pos_ids, token_count = self._compute_interleaved_positions(st_idx, data) llm_pos_ids_list.append(pos_ids) st_idx = int(pos_ids.max()) + 1 eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) llm_pos_ids_list.append(eos_pos) - llm_pos_ids_list.append(eos_pos) - video_len = grid_t * grid_h * grid_w - audio_len = self._compute_audio_token_count(data["audio_feature_length"]) - st = offset + 2 + video_len + audio_len + 2 + st = offset + int(data.get("placeholder_len", token_count + 2)) if st < seq_len: st_idx = int(llm_pos_ids_list[-1].max()) + 1 if llm_pos_ids_list else 0 @@ -1786,7 +1805,11 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1) if llm_positions.shape[1] != seq_len: - raise RuntimeError("Position ids length mismatch with input ids length") + raise RuntimeError( + "Position ids length mismatch with input ids length: " + f"positions={llm_positions.shape[1]}, input_ids={seq_len}, " + f"features={feature_summaries}" + ) mrope_position_delta = int(llm_positions.max()) + 1 - seq_len return torch.from_numpy(llm_positions), mrope_position_delta diff --git a/vllm_omni/outputs/__init__.py b/vllm_omni/outputs/__init__.py index 2f2a6f2ee..34731eeca 100644 --- a/vllm_omni/outputs/__init__.py +++ b/vllm_omni/outputs/__init__.py @@ -58,6 +58,18 @@ class OmniModelRunnerOutput(ModelRunnerOutput): kv_extracted_req_ids: list[str] | None = None omni_connector_output: OmniConnectorOutput | None = None + @classmethod + def with_kv_conn_output_only(cls, kv_connector_output: Any) -> "OmniModelRunnerOutput": + return cls( + req_ids=[], + req_id_to_index={}, + sampled_token_ids=[], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + kv_connector_output=kv_connector_output, + ) + @dataclass class OmniRequestOutput: