[Perf][MiniCPM-o] Stage 0 streaming audio/vision encoder CUDA graphs and candidate sampling (#8430)

Signed-off-by: BeatSeat <wendavid552@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

tests/model_executor/models/minicpmo_4_5/test_fused_residual_layer_norm.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""``residual_layer_norm`` matches the eager gated residual + LayerNorm chain."""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.model_executor.models.minicpmo_4_5.ops import residual_layer_norm
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+@pytest.mark.parametrize("device", ["cpu"] + (["cuda"] if torch.cuda.is_available() else []))
+def test_gated_residual_modulated_layer_norm(device: str) -> None:
+    torch.manual_seed(0)
+    channels = 64
+    x = torch.randn(3, 7, channels, device=device)
+    y = torch.randn(3, 7, channels, device=device)
+    gate, scale, shift = (torch.randn(channels, device=device) for _ in range(3))
+    expected_residual = x + gate * y
+    expected = F.layer_norm(expected_residual, (channels,), eps=1e-6) * (1 + scale) + shift
+    residual = x.clone()
+    out = residual_layer_norm(residual, y, gate=gate, weight=1 + scale, bias=shift, eps=1e-6, residual_out=residual)
+    torch.testing.assert_close(residual, expected_residual, rtol=1e-6, atol=1e-6)
+    torch.testing.assert_close(out, expected, rtol=1e-5, atol=1e-5)

tests/model_executor/models/minicpmo_4_5/test_incremental_fbank.py

@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""``duplex_incremental_fbank`` is bitwise equal to the remote ``StreamingMelProcessorExact`` (skipped without it)."""
+
+import importlib
+from pathlib import Path
+from types import ModuleType, SimpleNamespace
+from typing import Any
+
+import numpy as np
+import pytest
+import torch
+from transformers.dynamic_module_utils import init_hf_modules
+from transformers.utils import HF_MODULES_CACHE
+
+from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import MiniCPMO45Stage0DuplexRuntime
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_SR = 16000
+
+
+@pytest.fixture(scope="module")
+def remote() -> ModuleType:  # openbmb/MiniCPM-o-4_5's remote code, from the transformers module cache
+    candidates = Path(HF_MODULES_CACHE, "transformers_modules").glob("*/*/processing_minicpmo.py")
+    paths = [path for path in candidates if "class StreamingMelProcessorExact" in path.read_text()]
+    if not paths:
+        pytest.skip("MiniCPM-o 4.5 remote processing code is not in the transformers module cache")
+    init_hf_modules()
+    path = max(paths, key=lambda path: path.stat().st_mtime)
+    return importlib.import_module(".".join(path.relative_to(HF_MODULES_CACHE).with_suffix("").parts))
+
+
+def _mel(remote: ModuleType, *, incremental: bool) -> Any:
+    """Stage 0's streaming configuration, switched over by the runtime as a session's processor is."""
+    mel = remote.StreamingMelProcessorExact(
+        feature_extractor=remote.MiniCPMAAudioProcessor(),
+        chunk_ms=1000,
+        first_chunk_ms=1035,
+        cnn_redundancy_ms=20,
+        enable_sliding_window=True,
+        slide_trigger_seconds=30.0,
+        slide_stride_seconds=10.0,
+    )
+    runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
+    runtime._incremental_fbank = incremental
+    runtime._maybe_enable_incremental_fbank(SimpleNamespace(_streaming_mel_processor=mel))
+    assert (type(mel) is not remote.StreamingMelProcessorExact) is incremental
+    return mel
+
+
+@pytest.mark.parametrize("chunks", ["stage0", "random"])
+def test_streams_are_bitwise_equal(remote, chunks: str) -> None:
+    # Stage 0's chunk sizes, or random ones (unaligned buffer ends); a 66 s stream slides the window twice.
+    rng = np.random.default_rng(7)
+    audio = (rng.standard_normal(66 * _SR) * 0.05).astype(np.float32)
+    audio[_SR * 3 : _SR * 7] *= 20  # loud stretch for the dynamic range
+    audio[_SR * 12 : _SR * 14] = 0.0  # digital silence
+    audio[_SR * 45 :] *= 1e-3  # quiet tail
+    ref, inc = _mel(remote, incremental=False), _mel(remote, incremental=True)
+    position, slides = 0, 0
+    while position < len(audio):
+        size = ref.get_chunk_size() if chunks == "stage0" else int(rng.integers(1, 2 * _SR))
+        dropped = ref.left_samples_dropped
+        expected, expected_info = ref.process(audio[position : position + size])
+        got, got_info = inc.process(audio[position : position + size])
+        assert torch.equal(got, expected) and got_info == expected_info, position
+        assert torch.equal(inc._extract_full(), remote.StreamingMelProcessorExact._extract_full(inc)), position
+        position += size
+        slides += ref.left_samples_dropped != dropped
+    assert slides >= 2

tests/model_executor/models/minicpmo_4_5/test_native_duplex_sampling.py

@@ -0,0 +1,187 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Stage-0 duplex sampling: batched rows against a row-by-row reference, the deferred path against both."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch.utils._python_dispatch import TorchDispatchMode
+
+from vllm_omni.model_executor.models.minicpmo_4_5 import minicpmo_4_5_omni
+from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import MiniCPMO45DuplexPolicy
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+Model = minicpmo_4_5_omni.MiniCPMO45OmniForConditionalGeneration
+CHUNK_EOS, LISTEN, TTS_PAD, VOCAB = 5, 6, 10, 40
+_NAMES = ("unit", "chunk_eos", "listen", "tts_bos", "turn_eos", "chunk_tts_eos", "tts_pad")
+TOKEN_IDS = {f"{name}_token_id": token for name, token in zip(_NAMES, (1, CHUNK_EOS, LISTEN, 7, 8, 9, TTS_PAD))}
+
+
+class _CharTokenizer:
+    """Token t decodes to (t % 3) + 1 copies of one CJK character; special ids decode to ""."""
+
+    clean_up_tokenization_spaces, bad_token_ids, all_special_ids = False, list[int](), sorted(TOKEN_IDS.values())
+
+    def __len__(self):
+        return VOCAB
+
+    def decode(self, ids, skip_special_tokens=False):
+        return "".join("" if t in self.all_special_ids else chr(0x4E00 + t) * (t % 3 + 1) for t in map(int, ids))
+
+    def batch_decode(self, batch, skip_special_tokens=False):
+        return [self.decode(ids) for ids in batch]
+
+
+def _setup(seed: int, *, all_greedy: bool = False, max_chars: int = 6, turns_open: bool = True):
+    """Six rows: sampled, listen in an open turn, speak-length cap, boundary chunk_eos, greedy, penalized."""
+    torch.manual_seed(seed)
+    steps = [torch.randn(6, VOCAB) * 2 for _ in range(2)]
+    for logits in steps:
+        logits[3, CHUNK_EOS] += 50.0
+        logits[1, LISTEN] += 40.0
+    states = {row: SimpleNamespace(generated_tokens=[], current_turn_ended=True) for row in range(6)}
+    states[1].current_turn_ended, states[5].generated_tokens = not turns_open, [11, 11, 12]
+    model = Model.__new__(Model)
+    model.model_stage, model.max_new_speak_tokens_per_chunk, model.max_speak_chars_per_chunk = "llm", 5, max_chars
+    model._minicpmo45_native_duplex_token_ids_cache = dict(TOKEN_IDS)
+    model._minicpmo45_tokenizer_cache, model._minicpmo45_duplex_state_for_row = _CharTokenizer(), states.get
+    model._minicpmo45_duplex_payload_for_row = model._minicpmo45_duplex_row_request_max_tokens = lambda row: None
+    md = SimpleNamespace(all_greedy=all_greedy, top_p=torch.tensor([0.9, 1.0, 0.8, 0.8, 0.8, 0.85]))
+    md.generators = {row: torch.Generator().manual_seed(100 + row) for row in range(6)}
+    md.output_token_ids = [[11, 12], [2], [1, 2, 3, 4], [2, 3], [7, 2], [20, 21, 22]]
+    md.temperature, md.top_k = torch.tensor([0.8, 1.1, 0.7, 0.7, 0.0, 0.9]), torch.tensor([10, 0, 100, 100, 100, 20])
+    return model, states, md, steps
+
+
+def _reference(logits, md, states, row_params) -> list[int]:
+    """Row by row: speak-length cap, boundary draw, then the stage-2 draw (no cut, no listen rewrite)."""
+    out = []
+    for row, (temperature, top_k, top_p) in enumerate(row_params):
+        generator, recent = md.generators[row], md.output_token_ids[row]
+        if len(recent) >= 4 or torch.rand((), generator=generator) < torch.softmax(logits[row], -1)[CHUNK_EOS]:
+            out.append(CHUNK_EOS)
+            continue
+        row_logits = logits[row : row + 1].clone()
+        row_logits[0, [TTS_PAD, CHUNK_EOS]] = float("-inf")
+        for token in set((states[row].generated_tokens or recent)[-MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE :]):
+            row_logits[0, token] /= 1.05
+        if temperature <= 0:
+            out.append(int(row_logits.argmax()))
+            continue
+        scaled = row_logits / temperature
+        candidates = Model._duplex_top_k_top_p_candidates(scaled, top_k=top_k, top_p=top_p)
+        # The candidate-space distribution is the dense top-k / top-p filter's.
+        dense = torch.softmax(Model._top_k_top_p_filter(scaled, top_k=int(top_k), top_p=top_p), -1)
+        torch.testing.assert_close(_expand(scaled, candidates), dense)
+        uniform = torch.rand((), generator=generator)
+        out.append(int(Model._duplex_draw_candidates(scaled, candidates, uniform.view(1))))
+    return out
+
+
+def _expand(logits, candidates) -> torch.Tensor:
+    """The candidates as a dense distribution: the tie slot spread over its lowest-index tied tokens."""
+    probs, ids, kth, ties = candidates
+    dense = torch.zeros(logits.shape, dtype=probs.dtype).scatter_(1, ids[:, :-1], probs[:, :-1])
+    tied = logits == kth
+    kept = tied & (tied.cumsum(dim=-1) <= ties)
+    return dense + kept * (probs[:, -1:] / ties.clamp(min=1))
+
+
+def _run(model, md, steps) -> tuple[list[list[int]], list[bool]]:
+    """Two steps through the Stage-0 sampler; also whether each left a deferred commit."""
+    tokens, deferred = [], []
+    for logits in steps:
+        out = model._sample_minicpmo45_native_duplex_stage0(logits.clone(), md, duplex_rows=list(range(6)))
+        tokens.append(out.sampled_token_ids.squeeze(-1).tolist())
+        deferred.append(getattr(model, "_minicpmo45_duplex_pending_samples", None) is not None)
+        model._commit_minicpmo45_duplex_pending_samples()
+        for row, token in enumerate(tokens[-1]):
+            md.output_token_ids[row].append(token)
+    return tokens, deferred
+
+
+@pytest.mark.parametrize("seed", range(4))
+def test_batched_rows_match_the_row_by_row_reference(seed):
+    model, states, md, steps = _setup(seed, max_chars=10**6, turns_open=False)
+    _, ref_states, ref_md, _ = _setup(seed)
+    params = model._minicpmo45_duplex_row_params(md, 6)
+    kwargs = dict(row_idxs=list(range(6)), token_ids=TOKEN_IDS, row_params=params)
+    rows = model._sample_minicpmo45_native_duplex_rows(steps[0], md, **kwargs)
+    assert rows == _reference(steps[0], ref_md, ref_states, params)
+    assert all(torch.equal(md.generators[r].get_state(), ref_md.generators[r].get_state()) for r in range(6))
+
+
+@pytest.mark.parametrize("all_greedy", [False, True])
+@pytest.mark.parametrize("seed", range(3))
+def test_deferred_rows_match_the_synchronous_rows(seed, all_greedy, monkeypatch):
+    sync_model, sync_states, sync_md, steps = _setup(seed, all_greedy=all_greedy)
+    monkeypatch.setattr(sync_model, "_sample_minicpmo45_native_duplex_rows_deferred", lambda *a, **k: None)
+    expected, _ = _run(sync_model, sync_md, steps)
+    model, states, md, _ = _setup(seed, all_greedy=all_greedy)
+    assert _run(model, md, steps) == (expected, [True, True])
+    assert {row: vars(s) for row, s in states.items()} == {row: vars(s) for row, s in sync_states.items()}
+    assert all(torch.equal(md.generators[r].get_state(), sync_md.generators[r].get_state()) for r in range(6))
+
+
+def _tied_logits(device: str) -> torch.Tensor:
+    """Rows whose k-th logit is tied well past k, with the nucleus cutting inside and before the tie."""
+    logits = torch.full((4, 64), -3.0, device=device)
+    logits[0, [3, 9, 17, 30, 41, 50]] = 1.0  # six tied for k=2
+    logits[1, 7], logits[1, [2, 12, 22, 32, 42]] = 4.0, 1.0  # a leader, then five tied for k=3
+    logits[2] = torch.randn(64, generator=torch.Generator().manual_seed(0)).to(device)
+    logits[3, ::4] = 0.5  # sixteen tied for k=4
+    return logits
+
+
+@pytest.mark.parametrize("top_p", [1.0, 0.9, 0.5, 0.2])
+@pytest.mark.parametrize("top_k", [2, 3, 4, 0])
+def test_candidate_draws_follow_the_dense_distribution(top_k, top_p):
+    # A midpoint grid of uniforms integrates the inverse CDF exactly: each token
+    # is drawn in proportion to its dense probability, ties included.
+    logits, n = _tied_logits("cpu"), 4096
+    expected = torch.softmax(Model._top_k_top_p_filter(logits, top_k=top_k, top_p=top_p), dim=-1)
+    rows = logits.repeat(n, 1)
+    grid = ((torch.arange(n, dtype=torch.float32) + 0.5) / n).repeat_interleave(logits.shape[0])
+    draws = Model._duplex_draw_candidates(
+        rows, Model._duplex_top_k_top_p_candidates(rows, top_k=top_k, top_p=top_p), grid
+    )
+    counts = torch.zeros_like(expected).index_put_(
+        (torch.arange(rows.shape[0]) % logits.shape[0], draws), torch.ones(rows.shape[0]), accumulate=True
+    )
+    torch.testing.assert_close(counts / n, expected, atol=2.0 / n, rtol=0)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("top_p", [0.2, 0.25, 0.5, 0.8, 1.0])
+@pytest.mark.parametrize("n", [4, 10, 20, 64, 100])
+def test_tied_nucleus_boundaries_match_the_dense_filter(n, top_p, dtype):
+    # Uniform ties put every top_p on (or next to) a candidate boundary.
+    logits = torch.full((3, n + 8), -30.0)
+    logits[0, :n], logits[1, 3 : 3 + n], logits[2, :2], logits[2, 4 : 4 + n] = 20.0, 0.5, 2.0, 1.0
+    logits = logits.to(dtype)
+    # The candidates compute in float32; bf16 widens exactly, so ties survive.
+    expected = torch.softmax(Model._top_k_top_p_filter(logits.float(), top_k=3, top_p=top_p), dim=-1)
+    actual = _expand(logits, Model._duplex_top_k_top_p_candidates(logits, top_k=3, top_p=top_p))
+    assert torch.equal(actual > 0, expected > 0)
+    torch.testing.assert_close(actual, expected)
+
+
+class _NoHostReads(TorchDispatchMode):
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        if func in (torch.ops.aten._local_scalar_dense.default, torch.ops.aten.is_nonzero.default):
+            raise AssertionError(f"host read: {func}")
+        return func(*args, **(kwargs or {}))
+
+
+@pytest.mark.parametrize(
+    "device",
+    ["cpu", pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"))],
+)
+def test_candidate_sampling_reads_nothing_back_to_the_host(device):
+    logits = _tied_logits(device)
+    with _NoHostReads():
+        for top_k, top_p in [(2, 0.5), (4, 1.0), (0, 0.8)]:
+            candidates = Model._duplex_top_k_top_p_candidates(logits, top_k=top_k, top_p=top_p)
+            Model._duplex_draw_candidates(logits, candidates, torch.rand(logits.shape[0], device=device))

tests/model_executor/models/minicpmo_4_5/test_streaming_audio_batch.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The cross-session batched streaming audio encoder against the per-session path."""
+
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+from torch import nn
+from transformers.models.whisper.modeling_whisper import WhisperConfig
+
+from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import MiniCPMO45OmniForConditionalGeneration
+from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
+    MiniCPMO45OmniLLMForConditionalGeneration,
+    MiniCPMWhisperEncoder,
+    MultiModalProjector,
+)
+from vllm_omni.model_executor.models.minicpmo_4_5.streaming_audio_encoder import StreamingAudioChunk
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+N_MELS, POOL = 16, 5
+# Per round, each session's unit frames (None: sits the round out): 22 / 24 frames mirror 1020 / 1040 ms
+# at 1/10 scale, odd lengths pad conv2, 25 frames leave the pooler a remainder, and 40 positions force resets.
+_SCHEDULE: list[list[int | None]] = [
+    [22, 22, 23, 21],
+    [24, None, 24, 25],
+    [24, 24, None, 24],
+    [24, 23, 24, 24],
+    [24, 24, 24, None],
+    [25, 24, 23, 24],
+]
+
+
+def _thinker() -> Any:
+    """The audio half of the thinker: real modules and methods, toy sizes."""
+    torch.manual_seed(0)
+    config = WhisperConfig(num_mel_bins=N_MELS, d_model=32, encoder_layers=2, encoder_attention_heads=2)
+    config.update(dict(encoder_ffn_dim=128, max_source_positions=40, dropout=0.0, attention_dropout=0.0))
+    config._attn_implementation = "sdpa"
+    thinker = MiniCPMO45OmniLLMForConditionalGeneration.__new__(MiniCPMO45OmniLLMForConditionalGeneration)
+    nn.Module.__init__(thinker)
+    thinker.apm = MiniCPMWhisperEncoder(config).eval()
+    thinker.audio_projection_layer = MultiModalProjector(in_dim=32, out_dim=24)  # takes ffn // 4, as the checkpoint
+    thinker.audio_avg_pooler = nn.AvgPool1d(POOL, stride=POOL)
+    thinker.audio_encoder_layer, thinker.audio_past_key_values = -1, None
+    thinker.config = SimpleNamespace(audio_pool_step=POOL, duplex_audio_kv_page_positions=16)
+    return thinker
+
+
+def test_batched_rounds_match_sequential_sessions() -> None:
+    thinker = _thinker()
+    legacy: list[Any] = [None] * 4
+    batched: list[Any] = [None] * 4
+    resets = 0
+    for round_index, frames_per_session in enumerate(_SCHEDULE):
+        chunks, active, expected = [], [], []
+        for session, frames in enumerate(frames_per_session):
+            if frames is None:
+                continue
+            mel = torch.randn((1, N_MELS, frames), generator=torch.Generator().manual_seed(100 * round_index + session))
+            extra = dict(prefix_extra_frames=0 if legacy[session] is None else 2, suffix_extra_frames=2)
+            features = {"audio_features": mel, "audio_feature_lens": [torch.tensor([frames])]}
+            thinker.audio_past_key_values = legacy[session]
+            with torch.no_grad():
+                nested = thinker.get_audio_embedding_streaming(features, use_extra_context=True, **extra)
+            legacy[session] = thinker.audio_past_key_values
+            expected.append(torch.cat([t for row in nested for t in row]))
+            chunks.append(StreamingAudioChunk(mel, batched[session], **extra))
+            active.append(session)
+        idle = {s: (cache, cache.length if cache else 0) for s, cache in enumerate(batched) if s not in active}
+        outputs, caches = thinker.get_audio_embedding_streaming_batch(chunks)
+        for session, output, cache, reference in zip(active, outputs, caches, expected, strict=True):
+            torch.testing.assert_close(output, reference, rtol=1e-5, atol=1e-5)
+            resets += batched[session] is not None and cache is not batched[session]
+            batched[session] = cache
+            want = legacy[session].self_attention_cache
+            assert cache.length == want.get_seq_length()
+            for got, ref in zip(cache.to_legacy_cache().self_attention_cache.layers, want.layers, strict=True):
+                torch.testing.assert_close((got.keys, got.values), (ref.keys, ref.values))
+        # Not in the batch: the same cache object, the same committed length.
+        assert all(batched[s] is c and (c.length if c else 0) == n for s, (c, n) in idle.items())
+    assert resets >= 2, "the schedule must cross the max_source_positions reset"
+
+
+def test_graph_build_honors_eager_without_engine_config() -> None:
+    # The thinker keeps no vllm_config: the gate reads the flag __init__ stored.
+    thinker = _thinker()
+    thinker._enforce_eager = True
+    assert thinker.build_streaming_audio_graph_encoder(unit_frames=24) is False
+    assert thinker._duplex_audio_cuda_graph_encoder is None
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_graph_build_captures_on_cuda() -> None:
+    thinker = _thinker().cuda()
+    thinker._enforce_eager = False
+    thinker.config.duplex_audio_encoder_cuda_graph_batch_sizes = [1, 2]
+    thinker.config.duplex_audio_encoder_cuda_graph_cache_buckets = [8]
+    assert thinker.build_streaming_audio_graph_encoder(unit_frames=24) is True
+    assert thinker._duplex_audio_cuda_graph_encoder is not None
+
+
+@pytest.mark.parametrize(
+    ("session_mode", "device", "captures"),
+    [("duplex", "cuda", 1), ("turn", "cuda", 0), ("duplex", "cpu", 0)],
+)
+def test_duplex_audio_graphs_capture_while_the_model_loads(monkeypatch, session_mode, device, captures) -> None:
+    # Captured in load_weights, the graphs count against the KV budget instead of past it.
+    model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
+    nn.Module.__init__(model)
+    model.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode=session_mode))
+    model.thinker = nn.Linear(1, 1)
+    calls: list[None] = []
+    model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(build_audio_cuda_graph=lambda: calls.append(None))
+    monkeypatch.setattr(
+        MiniCPMO45OmniForConditionalGeneration, "_module_device", staticmethod(lambda module: torch.device(device))
+    )
+    model._capture_duplex_audio_graphs()
+    assert len(calls) == captures

tests/model_executor/models/minicpmo_4_5/test_vision_fused.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Fused packed SigLIP layers match the eager packed encode; prefetched frames serve only their own append."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.minicpmo_4_5 import vision_fused
+from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import MiniCPMO45Stage0DuplexRuntime
+from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
+    SiglipVisionConfig,
+    SiglipVisionTransformer,
+    _vision_encode_paths,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_fused_layers_match_eager_packed_layers() -> None:
+    torch.manual_seed(0)
+    sizes = dict(hidden_size=32, intermediate_size=72, num_hidden_layers=3, num_attention_heads=4)
+    config = SiglipVisionConfig(**sizes, image_size=28, patch_size=2, attention_dropout=0.0)
+    config._attn_implementation = "sdpa"
+    vpm = SiglipVisionTransformer(config).eval()
+    for norm in [module for module in vpm.modules() if isinstance(module, torch.nn.LayerNorm)]:
+        torch.nn.init.normal_(norm.weight, mean=1.0, std=0.2)
+        torch.nn.init.normal_(norm.bias, std=0.2)
+    seq_groups, hidden = [(0, 3, 16), (48, 2, 15)], torch.randn((78, 32), generator=torch.Generator().manual_seed(1))
+    with torch.inference_mode():
+        expected = hidden.clone()
+        for layer in vpm.encoder.layers:
+            expected = layer.forward_packed(expected, seq_groups)
+        fused = vision_fused.encode_packed_fused(vpm, hidden.clone(), seq_groups)
+        torch.testing.assert_close(fused, vpm.post_layernorm(expected), rtol=1e-5, atol=1e-5)
+
+
+def test_prefetched_frames_are_chunked_and_keyed_by_append() -> None:
+    runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
+    runtime.sessions, runtime._prefetched_vision = {}, {}
+    runtime.processor = SimpleNamespace(process_image=lambda *args, **kwargs: None)
+    runtime.stage_model = runtime.thinker = SimpleNamespace(config=SimpleNamespace(vision_batch_size=2))
+    runtime._decode_video_frames_payload = lambda payload: list(payload["video_frames"])
+    runtime._preprocess_frames = lambda frames: [f"p{frame}" for frame in frames]
+    # Each encoded frame is tagged with its tower call's size: vision_batch_size appends per call.
+    runtime._encode_processed_vision_batch = lambda items: [[f"e{item}/{len(items)}"] for item in items]
+    frames = ["a", "b", "c"]
+    appends = [
+        {"session_id": f"s{i}", "epoch": 0, "seq": 1, "payload": {"video_frames": frames[: i + 1]}} for i in range(3)
+    ]
+    runtime.prefetch_vision(appends)
+    assert runtime.take_prefetched_vision({**appends[1], "payload": {"video_frames": ["x", "y"]}}) is None  # changed
+    assert runtime.take_prefetched_vision(appends[1]) is None  # and the entry is gone
+    assert runtime.take_prefetched_vision({**appends[2], "seq": "2"}) is None  # another append
+    last, payload = appends[2], appends[2]["payload"]
+    assert runtime.frame_kwargs(last, payload) == {"encoded_frames": [["epa/3"], ["epb/3"], ["epc/3"]]}
+    assert runtime.frame_kwargs(last, payload) == {"video_frames": frames}  # taken once
+
+
+def test_packed_vision_adapter_protocol() -> None:
+    shape, out_shape = (1, 3, 16), torch.Size([1, 4, 8])
+    adapter = vision_fused._PackedVisionAdapter(lambda px, runs: px[:, :4, :8], [(2, 2, 1)], shape, out_shape)
+    assert adapter.supports_encoder_cudagraph is True
+    cfg = adapter.get_encoder_cudagraph_config()
+    assert cfg.modalities == ["image"]
+    assert cfg.buffer_keys == ["pixels"]
+    assert cfg.out_hidden_size == 8
+    inputs = adapter.prepare_encoder_cudagraph_capture_inputs(4, 1, 0, torch.device("cpu"), torch.float32)
+    assert "pixels" in inputs.values
+    dest: list[torch.Tensor | None] = [None]
+    adapter.postprocess_encoder_output({"default": torch.ones(1, 4, 8)}, [0], [4], dest, clone=True)
+    assert dest[0] is not None and dest[0].shape == (1, 4, 8)
+
+
+@pytest.mark.parametrize(
+    ("overrides", "encoder_graphs", "expected"),
+    [
+        ({}, True, (False, False)),  # default: the padded batch keeps the vpm encoder graph
+        ({"vision_cuda_graph": True}, True, (True, True)),
+        ({"vision_cuda_graph": True}, False, (False, False)),  # encoder_cuda_graph=False / --enforce-eager
+        ({"vision_packed_encode": True}, True, (True, False)),  # packed eager, for A/B
+        ({"vision_packed_encode": False, "vision_cuda_graph": True}, True, (False, False)),
+    ],
+)
+def test_vision_encode_paths(overrides, encoder_graphs, expected) -> None:
+    assert _vision_encode_paths(SimpleNamespace(**overrides), encoder_graphs=encoder_graphs) == expected