[Perf][Qwen-Omni] Check audio-in-video interleaving with tensor reductions

tests/model_executor/models/common/test_audio_in_video.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The tensor-reduction interleave check agrees with vLLM's per-token loop."""
+
+import random
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.common.audio_in_video import check_interleaved_audio_video
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+TEXT, VIDEO, AUDIO, IMAGE, BOUNDARY = 0, 1, 2, 3, 4
+
+
+def _reference():
+    from vllm.model_executor.models.qwen2_5_omni_thinker import check_interleaved_audio_video as reference
+
+    return reference
+
+
+def _masks(tokens):
+    ids = torch.tensor(tokens, dtype=torch.long)
+    return ids == VIDEO, ids == AUDIO
+
+
+def _check(fn, tokens):
+    is_video, is_audio = _masks(tokens)
+    return fn(is_video, is_audio, int(is_video.sum()), int(is_audio.sum()))
+
+
+def _random_request(rng):
+    """One prompt: text, images, and videos with separate or interleaved audio, as the processors lay them out."""
+    tokens = [TEXT] * rng.randint(0, 4)
+    for _ in range(rng.randint(1, 4)):
+        kind = rng.choice(["image", "video", "audio", "video_then_audio", "audio_in_video"])
+        if kind == "image":
+            tokens += [IMAGE] * rng.randint(1, 6)
+        elif kind == "video":
+            tokens += [VIDEO] * rng.randint(1, 8)
+        elif kind == "audio":
+            tokens += [AUDIO] * rng.randint(1, 8)
+        elif kind == "video_then_audio":
+            tokens += [VIDEO] * rng.randint(1, 6) + [AUDIO] * rng.randint(1, 6)
+        else:
+            tokens.append(BOUNDARY)
+            for _ in range(rng.randint(1, 4)):
+                tokens += [VIDEO] * rng.randint(0, 5) + [AUDIO] * rng.randint(0, 3)
+            tokens.append(BOUNDARY)
+        tokens += [rng.choice([TEXT, BOUNDARY])] * rng.randint(0, 2)
+    return tokens
+
+
+@pytest.mark.parametrize("seed", range(20))
+def test_matches_vllm_on_random_batches(seed):
+    rng = random.Random(seed)
+    reference = _reference()
+    outcomes = set()
+    for _ in range(50):
+        # The thinker checks a whole scheduled batch at once, so concatenate several prompts.
+        tokens: list[int] = sum((_random_request(rng) for _ in range(rng.randint(1, 4))), [])
+        expected = _check(reference, tokens)
+        assert _check(check_interleaved_audio_video, tokens) == expected, tokens
+        outcomes.add(expected)
+    assert outcomes == {True, False}
+
+
+def _audio_in_video(num_video, num_audio, chunk_video=600, chunk_audio=50):
+    tokens = [BOUNDARY]
+    while num_video or num_audio:
+        step = min(num_video, chunk_video)
+        tokens += [VIDEO] * step
+        num_video -= step
+        step = min(num_audio, chunk_audio)
+        tokens += [AUDIO] * step
+        num_audio -= step
+    return tokens + [BOUNDARY]
+
+
+@pytest.mark.parametrize(
+    "tokens, expected",
+    [
+        # Daily-Omni sizes: one video with audio in it, and one with its audio as a separate item.
+        ([TEXT] * 50 + _audio_in_video(6600, 750) + [TEXT] * 20, True),
+        ([TEXT] * 50 + [BOUNDARY] + [VIDEO] * 6600 + [BOUNDARY, BOUNDARY] + [AUDIO] * 750 + [BOUNDARY], False),
+    ],
+)
+def test_request_sized_inputs(tokens, expected):
+    assert _check(check_interleaved_audio_video, tokens) is expected
+    assert _check(_reference(), tokens) is expected
+
+
+@pytest.mark.parametrize(
+    "tokens, expected",
+    [
+        ([TEXT, VIDEO, VIDEO, AUDIO, AUDIO, TEXT], False),  # one span, video before audio
+        ([VIDEO, AUDIO, VIDEO], True),
+        ([AUDIO, VIDEO, AUDIO], True),
+        ([BOUNDARY, VIDEO, AUDIO, BOUNDARY, VIDEO, AUDIO, BOUNDARY], False),  # two videos, each with audio after
+        ([VIDEO, VIDEO, BOUNDARY, AUDIO, VIDEO, AUDIO], True),  # only the second span interleaves
+        ([VIDEO, IMAGE, AUDIO, IMAGE, VIDEO], False),  # images split the spans
+        ([TEXT, VIDEO, VIDEO, AUDIO, VIDEO, AUDIO, TEXT], True),  # two prompts touching, no boundary token between
+        ([VIDEO] * 5, False),
+        ([AUDIO] * 5, False),
+        ([TEXT] * 5, False),
+    ],
+)
+def test_span_rules(tokens, expected):
+    assert _check(check_interleaved_audio_video, tokens) is expected
+    assert _check(_reference(), tokens) is expected
+
+
+@pytest.mark.parametrize(
+    "module",
+    [
+        "vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker",
+        "vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker",
+    ],
+)
+def test_thinkers_use_this_check(module):
+    import importlib
+
+    assert importlib.import_module(module).check_interleaved_audio_video is check_interleaved_audio_video
+
+
+def _host_reads(monkeypatch, tokens):
+    is_video, is_audio = _masks(tokens)
+    num_video, num_audio = int(is_video.sum()), int(is_audio.sum())
+    calls = []
+
+    def counting(name):
+        method = getattr(torch.Tensor, name)
+
+        def wrapper(self, *args, **kwargs):
+            calls.append(name)
+            return method(self, *args, **kwargs)
+
+        return wrapper
+
+    with monkeypatch.context() as patch:
+        for name in ("item", "tolist", "__bool__", "__int__", "__index__", "__float__"):
+            patch.setattr(torch.Tensor, name, counting(name))
+        result = check_interleaved_audio_video(is_video, is_audio, num_video, num_audio)
+    return result, calls
+
+
+def test_host_reads_do_not_grow_with_the_prompt(monkeypatch):
+    # The vLLM loop calls .item() twice per placeholder token; this check reads back a fixed number of values.
+    short, short_calls = _host_reads(monkeypatch, _audio_in_video(3000, 250))
+    long, long_calls = _host_reads(monkeypatch, _audio_in_video(6000, 500) * 2)
+    assert short is True and long is True
+    assert short_calls == long_calls
+    assert len(long_calls) <= 1

tests/model_executor/models/common/test_audio_in_video_cuda.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The Qwen3-Omni thinker passes CUDA masks; the check gives the same answers there with a fixed number of syncs."""
+
+import random
+import warnings
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_marks
+from tests.model_executor.models.common.test_audio_in_video import AUDIO, VIDEO, _audio_in_video, _random_request
+from vllm_omni.model_executor.models.common.audio_in_video import check_interleaved_audio_video
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+    pytest.mark.core_model,
+    *hardware_marks(res={"cuda": "L4"}),
+]
+
+
+def _run(tokens, device):
+    ids = torch.tensor(tokens, dtype=torch.long, device=device)
+    is_video, is_audio = ids == VIDEO, ids == AUDIO
+    return check_interleaved_audio_video(is_video, is_audio, int(is_video.sum()), int(is_audio.sum()))
+
+
+def test_cuda_matches_cpu():
+    from vllm.model_executor.models.qwen2_5_omni_thinker import check_interleaved_audio_video as reference
+
+    rng = random.Random(0)
+    batches: list[list[int]] = [sum((_random_request(rng) for _ in range(rng.randint(1, 4))), []) for _ in range(200)]
+    batches.append(_audio_in_video(6600, 750))
+    for tokens in batches:
+        ids = torch.tensor(tokens, dtype=torch.long)
+        is_video, is_audio = ids == VIDEO, ids == AUDIO
+        expected = reference(is_video, is_audio, int(is_video.sum()), int(is_audio.sum()))
+        assert _run(tokens, "cpu") is expected
+        assert _run(tokens, "cuda") is expected
+
+
+def _device_syncs(tokens):
+    ids = torch.tensor(tokens, dtype=torch.long, device="cuda")
+    is_video, is_audio = ids == VIDEO, ids == AUDIO
+    num_video, num_audio = int(is_video.sum()), int(is_audio.sum())
+    previous = torch.cuda.get_sync_debug_mode()
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        torch.cuda.set_sync_debug_mode("warn")
+        try:
+            result = check_interleaved_audio_video(is_video, is_audio, num_video, num_audio)
+        finally:
+            torch.cuda.set_sync_debug_mode(previous)
+    return result, sum("called a synchronizing CUDA operation" in str(w.message) for w in caught)
+
+
+def test_device_syncs_do_not_grow_with_the_prompt():
+    # vLLM's loop syncs twice per placeholder token; this check syncs for the nonzero size and the final flag.
+    assert _device_syncs(_audio_in_video(3000, 250)) == (True, 2)
+    assert _device_syncs(_audio_in_video(6000, 500) * 2) == (True, 2)

tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """
 Unit tests for Qwen2.5-Omni embed_input_ids to verify embeddings are
 correctly assigned to audio/image/video token positions.
@@ -18,6 +18,8 @@
 from pytest_mock import MockerFixture
 from vllm.multimodal.utils import set_mm_embedding_modality
 
+from vllm_omni.model_executor.models.common.audio_in_video import check_interleaved_audio_video
+
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
@@ -83,35 +85,31 @@ def make_interleaved_seq(video_chunks: list[int], audio_chunks: list[int], text_
 class TestCheckInterleavedAudioVideo:
     def test_non_interleaved_audio_then_video(self):
         """Audio entirely before video -> not interleaved."""
-        m = _qwen2_5_omni_thinker_mod()
         input_ids, is_multimodal = make_token_seq(5, 0, 4)
         is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
         is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
-        assert not m.check_interleaved_audio_video(is_video, is_audio, is_video.sum().item(), is_audio.sum().item())
+        assert not check_interleaved_audio_video(is_video, is_audio, is_video.sum().item(), is_audio.sum().item())
 
     def test_non_interleaved_with_image(self):
         """Audio + image + video (the mixed_modalities case) -> not interleaved."""
-        m = _qwen2_5_omni_thinker_mod()
         input_ids, is_multimodal = make_token_seq(5, 4, 6)
         is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
         is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
-        assert not m.check_interleaved_audio_video(is_video, is_audio, is_video.sum().item(), is_audio.sum().item())
+        assert not check_interleaved_audio_video(is_video, is_audio, is_video.sum().item(), is_audio.sum().item())
 
     def test_no_audio(self):
         """Video only -> not interleaved."""
-        m = _qwen2_5_omni_thinker_mod()
         input_ids, is_multimodal = make_token_seq(0, 0, 6)
         is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
         is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
-        assert not m.check_interleaved_audio_video(is_video, is_audio, is_video.sum().item(), is_audio.sum().item())
+        assert not check_interleaved_audio_video(is_video, is_audio, is_video.sum().item(), is_audio.sum().item())
 
     def test_interleaved(self):
         """V A V A interleaved -> True."""
-        m = _qwen2_5_omni_thinker_mod()
         input_ids, is_multimodal = make_interleaved_seq([4, 4], [3, 3])
         is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
         is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
-        assert m.check_interleaved_audio_video(is_video, is_audio, is_video.sum().item(), is_audio.sum().item())
+        assert check_interleaved_audio_video(is_video, is_audio, is_video.sum().item(), is_audio.sum().item())
 
 
 # ---------------------------------------------------------------------------

vllm_omni/model_executor/models/common/audio_in_video.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Audio-in-video interleave detection for the Qwen-Omni thinkers."""
+
+import torch
+
+
+def check_interleaved_audio_video(
+    is_video: torch.Tensor,
+    is_audio: torch.Tensor,
+    num_video: int,
+    num_audio: int,
+) -> bool:
+    """Whether any contiguous video/audio placeholder span interleaves the two.
+
+    Same result as vLLM's ``qwen2_5_omni_thinker.check_interleaved_audio_video``
+    (a span is a run of consecutive video or audio positions; it interleaves
+    when its first video position precedes its last audio position and its
+    first audio position precedes its last video position), computed with
+    tensor reductions and two host reads in total.
+    """
+    if num_video == 0 or num_audio == 0:
+        return False
+
+    positions = (is_video | is_audio).nonzero(as_tuple=True)[0]
+    if positions.numel() == 0:
+        return False
+
+    span_starts = torch.ones_like(positions, dtype=torch.bool)
+    span_starts[1:] = positions[1:] != positions[:-1] + 1
+    span_ids = span_starts.cumsum(0) - 1
+
+    video = is_video[positions]
+    audio = is_audio[positions]
+    # One slot per position (at least one per span), so the span count is never read back. Unused slots
+    # and spans without video (or audio) keep the past-the-end and -1 fills, which never compare as interleaved.
+    past_end = torch.full_like(positions, is_video.numel())
+
+    def span_bounds(mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        first = past_end.scatter_reduce(0, span_ids, torch.where(mask, positions, past_end), "amin")
+        last = torch.full_like(positions, -1).scatter_reduce(0, span_ids, torch.where(mask, positions, -1), "amax")
+        return first, last
+
+    video_first, video_last = span_bounds(video)
+    audio_first, audio_last = span_bounds(audio)
+    interleaved = (video_first < audio_last) & (audio_first < video_last)
+    return bool(interleaved.any())

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py

@@ -34,7 +34,6 @@
 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,
 )
@@ -88,6 +87,7 @@
 from vllm.sequence import IntermediateTensors
 from vllm.utils.collection_utils import is_list_of
 
+from vllm_omni.model_executor.models.common.audio_in_video import check_interleaved_audio_video
 from vllm_omni.quantization.component_config import (
     resolve_encoder_quant_config,
 )
@@ -1348,8 +1348,8 @@ def embed_input_ids(
         video_token_id = self.config.video_token_index
         audio_token_id = self.config.audio_token_index
 
-        # The interleaving check reads scalar positions in a Python loop.
-        # Keep its masks on CPU to avoid repeated CUDA synchronization.
+        # Keep the masks on CPU: the counts below and the per-modality counts
+        # in merge_interleaved_embeddings then read back without CUDA syncs.
         input_ids_cpu = input_ids.cpu()
         is_video = is_multimodal & (input_ids_cpu == video_token_id)
         is_audio = is_multimodal & (input_ids_cpu == audio_token_id)

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -71,7 +71,6 @@
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniAudioFeatureInputs,
     Qwen2_5OmniThinkerDummyInputsBuilder,
-    check_interleaved_audio_video,
     merge_interleaved_embeddings,
 )
 from vllm.model_executor.models.qwen2_5_vl import (
@@ -120,6 +119,7 @@
 from vllm.utils.torch_utils import async_tensor_h2d
 
 from vllm_omni.data_entry_keys import OmniPayload
+from vllm_omni.model_executor.models.common.audio_in_video import check_interleaved_audio_video
 from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
     Qwen2_5OmniThinkerMultiModalDataParser,