[Feature] Support per-video audio masks for Qwen Omni (#4656)

Signed-off-by: hbhflw2000 <417911774@qq.com>

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

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,

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

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: