[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

tests/worker/test_native_duplex_hooks.py

@@ -415,7 +415,8 @@ def test_minicpmo_non_duplex_sample_skips_duplex_row_scan():
     )
 
     assert runner._sample(torch.zeros((1, 4)), spec_decode_metadata=None) == "model-sampler"
-    assert calls == []
+    # The empty row set is still published: the model then skips its per-row prompt scan.
+    assert calls == ["prepare"]
 
 
 def test_minicpmo_duplex_sample_clears_stale_rows_once_without_scanning():
@@ -442,7 +443,7 @@ def test_minicpmo_duplex_sample_clears_stale_rows_once_without_scanning():
 
     assert runner._sample(torch.zeros((1, 4)), spec_decode_metadata=None) == "model-sampler"
     assert runner._sample(torch.zeros((1, 4)), spec_decode_metadata=None) == "model-sampler"
-    assert calls == [()]
+    assert calls == [(), ()]
 
 
 def test_generic_ar_runner_has_no_minicpmo_sampler_state_or_typeerror_probe():
@@ -2890,6 +2891,7 @@ def test_minicpmo_stage0_native_sampler_forced_listen_yields_floor():
     )
 
     sampled = model.sample(logits, sampling_metadata)
+    model._commit_minicpmo45_duplex_pending_samples()  # the runner commits before the next step
 
     assert sampled is not None
     assert sampled.sampled_token_ids.tolist() == [[151705]]
@@ -3155,6 +3157,7 @@ def test_minicpmo_stage0_turn_eos_is_forwarded_not_pending():
     logits[0, 151717] = 30.0
 
     sampled = model.sample(logits, _native_duplex_sampling_metadata([151706, 200, 201]))
+    model._commit_minicpmo45_duplex_pending_samples()  # the runner commits before the next step
 
     assert sampled.sampled_token_ids.tolist() == [[151717]]
     assert state.pending_terminator_token is None
@@ -3166,6 +3169,7 @@ def test_minicpmo_stage0_turn_eos_is_forwarded_not_pending():
     logits = torch.full((1, 151723), -100.0)
     logits[0, 151718] = 30.0
     sampled = model.sample(logits, _native_duplex_sampling_metadata([151706, 200, 201, 151717]))
+    model._commit_minicpmo45_duplex_pending_samples()
 
     assert sampled.sampled_token_ids.tolist() == [[151718]]
     assert state.pending_terminator_token == 151718

vllm_omni/deploy/minicpmo_4_5_duplex_h200.yaml

@@ -1,6 +1,6 @@
 # Opt-in single-H200 native full-duplex codec profile (V1).
 # Inherits the native session/stage capacities, sampling and chunk sizes.
-# Uses the mainline Whole-Euler Flow and shared prompt-state implementation.
+# Enables all Stage 0 (multimodal duplex prefill) and Stage 2 (code2wav DiT) optimizations.
 # Ordinary TF32 for CFM DiT GEMMs (cuBLAS allow_tf32 + Triton tf32 dots).
 # HiFT stays IEEE FP32. Validate audio quality for your workload.
 base_config: minicpmo_4_5.yaml
@@ -9,3 +9,28 @@ connectors:
   connector_of_shared_memory:
     extra:
       token2wav_allow_tf32: true
+      # Stage 2: Fused DiT transformer body (LayerNorm + QKV + Residual)
+      cfm_fused_body: true
+      # Stage 2: Flow condition encoder CUDA graph
+      cfm_encoder_cuda_graph: true
+      # Stage 2: Pre-allocated persistent slot pool and row offset merging
+      cfm_slot_pool: true
+      cfm_row_offset_merge: true
+      # Stage 2: Compress DiT attention cache to bf16
+      code2wav_bfloat16_attention_cache: true
+
+stages:
+  - stage_id: 0
+    enforce_eager: false
+    hf_overrides:
+      # Stage 0: Incremental sliding window Fbank feature extraction
+      duplex_incremental_fbank: true
+      # Stage 0: Fused packed SigLIP encoder layers
+      vision_fused_layers: true
+      # Stage 0: Packed SigLIP tower & resampler CUDA graph
+      vision_cuda_graph: true
+      # Stage 0: Steady Whisper streaming audio encoder CUDA graph
+      duplex_audio_encoder_cuda_graph: true
+      duplex_audio_encoder_cuda_graph_batch_sizes_from_sessions: true
+      # Stage 0: Pinned memory async H2D copy for streaming audio chunks
+      duplex_audio_encoder_pinned_h2d: true

vllm_omni/model_executor/duplex_sampling.py

@@ -22,6 +22,20 @@ class DuplexSamplingRow:
     seq: int | None
     payload: dict[str, object] | None
     max_tokens: int | None
+    # Host copies of the row's sampling parameters (None when the runner has none).
+    temperature: float | None = None
+    top_k: int | None = None
+    top_p: float | None = None
+
+
+def _host_value(values: Any, row_idx: int, cast: type) -> Any:
+    """``values[row_idx]`` of a host array as ``cast``, or None when there is none."""
+    if values is None:
+        return None
+    try:
+        return cast(values[row_idx])
+    except (IndexError, TypeError, ValueError):
+        return None
 
 
 class DuplexSamplingHelper:
@@ -53,15 +67,19 @@ class DuplexSamplingHelper:
         else:
             self.active_request_ids.discard(req_id)
 
-    def update_states(self, runner: object, scheduler_output: object) -> None:
+    def update_states(self, runner: Any, scheduler_output: Any) -> None:
         self.active_request_ids.difference_update(str(req_id) for req_id in scheduler_output.finished_req_ids)
         for request in scheduler_output.scheduled_new_reqs:
             self.refresh_active_request(runner, str(request.req_id))
 
-    def rows(self, runner: object) -> tuple[DuplexSamplingRow, ...]:
+    def rows(self, runner: Any) -> tuple[DuplexSamplingRow, ...]:
         rows: list[DuplexSamplingRow] = []
-        req_ids = [str(req_id) for req_id in getattr(runner.input_batch, "req_ids", [])]
+        input_batch = runner.input_batch
+        req_ids = [str(req_id) for req_id in getattr(input_batch, "req_ids", [])]
         requests = getattr(runner, "requests", {})
+        temperature_cpu = getattr(input_batch, "temperature_cpu", None)
+        top_k_cpu = getattr(input_batch, "top_k_cpu", None)
+        top_p_cpu = getattr(input_batch, "top_p_cpu", None)
         for row_idx, req_id in enumerate(req_ids):
             if req_id not in self.active_request_ids:
                 continue
@@ -73,7 +91,8 @@ class DuplexSamplingHelper:
             if not isinstance(session_id, str) or not session_id:
                 session_id = None
             try:
-                seq = int(duplex.get("seq"))
+                seq_raw = duplex.get("seq")
+                seq = int(seq_raw) if seq_raw is not None else None
             except (TypeError, ValueError):
                 seq = None
             payload = duplex.get("payload")
@@ -93,6 +112,9 @@ class DuplexSamplingHelper:
                     seq=seq,
                     payload=payload,
                     max_tokens=max_tokens if max_tokens > 0 else None,
+                    temperature=_host_value(temperature_cpu, row_idx, float),
+                    top_k=_host_value(top_k_cpu, row_idx, int),
+                    top_p=_host_value(top_p_cpu, row_idx, float),
                 )
             )
         return tuple(rows)
@@ -142,8 +164,9 @@ class DuplexSamplingRunnerMixin:
             return
         helper = getattr(self, "_duplex_sampling_helper", None)
         rows = helper.rows(self) if helper is not None and helper.active_request_ids else ()
-        if rows or (helper is not None and helper.hook_active):
-            prepare_duplex_sampling(logits, prepared_sampling_metadata, rows)
+        # Publish the row set even when empty, so the model skips its per-row
+        # prompt-token scan (a host sync per row) in turn-based serving.
+        prepare_duplex_sampling(logits, prepared_sampling_metadata, rows)
         if helper is not None:
             helper.hook_active = bool(rows)
 

vllm_omni/model_executor/models/minicpmo_4_5/duplex/incremental_fbank.py

@@ -0,0 +1,220 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Incremental Stage-0 streaming fbank (HF override ``duplex_incremental_fbank``, default off).
+
+The remote ``StreamingMelProcessorExact._extract_full`` recomputes the CPU Whisper log-mel over the whole
+streaming buffer (up to its 30 s slide) on every unit. This keeps the ``|STFT|^2`` rows already computed and
+recomputes only the right edge (frames whose window reaches new samples or the reflect pad) and, after a slide,
+the new left edge. The mel projection, log, dynamic range and scaling still run over the whole buffer with the
+original's shapes and memory layout, so the features stay bitwise equal: the mel GEMM is not row independent
+(MKL's remainder kernel), and an STFT batch's last frames may take a remainder path, so recomputes start on
+``_ALIGN_FRAMES`` boundaries and their last ``_ALIGN_FRAMES`` frames are not kept.
+"""
+
+from __future__ import annotations
+
+import copy
+from typing import Any
+
+import numpy as np
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+# Recompute start granularity, trailing recompute frames not kept, and the
+# left-edge recompute after a slide (>= the 2 frames reading the pad).
+_ALIGN_FRAMES = 16
+_FEATURE_EXTRACTOR_ATTRS = ("n_fft", "hop_length", "n_samples", "sampling_rate", "mel_filters", "set_spac_log_norm")
+_FEATURE_EXTRACTOR_ATTRS += ("dynamic_log_norm", "dynamic_range_db", "log_floor_db")
+_INCREMENTAL_CLASSES: dict[type, type] = {}
+_SELF_CHECKED: dict[tuple[type, type], bool] = {}
+
+
+def _stft_power(buffer: np.ndarray, start: int, stop: int, n_fft: int, hop: int, window: torch.Tensor) -> np.ndarray:
+    """``(stop - start, n_fft // 2 + 1)`` rows of ``|STFT|^2``, framed as the original's centered STFT.
+
+    The batch is ``[start, stop]`` minus its last frame, like the original's
+    ``stft(...)[..., :-1].abs() ** 2``; edge reflect padding is built explicitly.
+    """
+    half = n_fft // 2
+    lo, hi = start * hop - half, stop * hop + half
+    length = len(buffer)
+    segment = buffer[max(lo, 0) : min(hi, length)]
+    if lo < 0 or hi > length:
+        segment = np.pad(segment, (max(0, -lo), max(0, hi - length)), mode="reflect")
+    spec = torch.stft(
+        torch.from_numpy(segment)[None], n_fft=n_fft, hop_length=hop, window=window, center=False, return_complex=True
+    )
+    magnitudes = spec[..., :-1].abs() ** 2  # [1, n_freq, stop - start], frame-major in memory
+    return magnitudes[0].T.numpy()
+
+
+def _log_mel(power: torch.Tensor, mel_filters: torch.Tensor, feature_extractor: Any) -> torch.Tensor:
+    """``_torch_extract_fbank_features`` from the ``|STFT|^2`` rows, op for op and layout for layout."""
+    magnitudes = power.unsqueeze(0).transpose(1, 2)  # [1, n_freq, T], strides of stft(...)[..., :-1]
+    log_spec = torch.clamp(mel_filters.T @ magnitudes, min=1e-10).log10()
+    if feature_extractor.dynamic_log_norm:
+        max_val_bt = log_spec.max(dim=2, keepdim=True)[0].max(dim=1, keepdim=True)[0]
+        log_spec = torch.maximum(log_spec, max_val_bt - feature_extractor.dynamic_range_db)
+    else:
+        log_spec = torch.maximum(log_spec, torch.tensor(feature_extractor.log_floor_db, dtype=log_spec.dtype))
+    return (log_spec + 4.0) / 4.0
+
+
+class _IncrementalFbankMixin:
+    """Overrides of the remote ``StreamingMelProcessorExact``, mixed in by ``enable_incremental_fbank``."""
+
+    #: ``|STFT|^2`` rows (frame t at row t); rows ``[0, _ifb_valid)`` are exact for the current buffer.
+    _ifb_power: np.ndarray | None = None
+    _ifb_valid = 0
+    #: ``left_samples_dropped`` and the buffer length when the rows were kept.
+    _ifb_dropped = 0
+    _ifb_length = 0
+    #: The ``n_fft`` samples ending where the last kept frame's window ends.
+    _ifb_guard: np.ndarray | None = None
+    _ifb_window: torch.Tensor | None = None
+    _ifb_mel_filters: torch.Tensor | None = None
+
+    def reset(self) -> None:
+        super().reset()
+        self._ifb_valid = 0
+
+    def restore_snapshot(self, snapshot: dict) -> None:
+        super().restore_snapshot(snapshot)
+        self._ifb_valid = 0
+
+    def _extract_full(self) -> torch.Tensor:
+        fe, buffer = self.feature_extractor, self.buffer
+        if not (
+            isinstance(buffer, np.ndarray)
+            and buffer.ndim == 1
+            and buffer.dtype == np.float32
+            and fe.n_fft <= len(buffer) <= fe.n_samples
+            and getattr(fe, "dither", 0.0) == 0.0
+            and fe.sampling_rate == self.sample_rate
+        ):
+            # Too short (the original raises), truncated, dithered or misconfigured: the original.
+            self._ifb_valid = 0
+            return super()._extract_full()
+        # The original's log-floor switch, side effect on the feature extractor included.
+        if len(buffer) < 5 * self.sample_rate:
+            fe.set_spac_log_norm(log_floor_db=-10)
+        else:
+            fe.set_spac_log_norm(dynamic_range_db=8)
+        return _log_mel(self._ifb_update_power(buffer, fe), self._ifb_mel_filters, fe)
+
+    def _ifb_update_power(self, buffer: np.ndarray, fe: Any) -> torch.Tensor:
+        n_fft, hop = int(fe.n_fft), int(fe.hop_length)
+        if self._ifb_window is None:
+            self._ifb_window = torch.hann_window(n_fft)
+            self._ifb_mel_filters = torch.from_numpy(fe.mel_filters).to(torch.float32)
+        power = self._ifb_power
+        if power is None or power.ctypes.data % 64:
+            # Rows from a torch allocation: 64-byte aligned like the original's fresh
+            # |X|^2 tensor, which the GEMM reads (a deepcopy's numpy array may not be).
+            aligned = torch.empty(int(fe.n_samples) // hop + 1, n_fft // 2 + 1, dtype=torch.float32).numpy()
+            if power is not None:
+                aligned[: self._ifb_valid] = power[: self._ifb_valid]
+            self._ifb_power = power = aligned
+        frames = len(buffer) // hop  # the original drops the STFT's last frame
+        start = self._ifb_rebase(buffer, n_fft, hop) // _ALIGN_FRAMES * _ALIGN_FRAMES
+        power[start:frames] = _stft_power(buffer, start, frames, n_fft, hop, self._ifb_window)
+        # Keep frames whose window lies inside the buffer, minus the recompute's tail.
+        valid = min((len(buffer) - n_fft // 2) // hop + 1, frames - _ALIGN_FRAMES)
+        guard_end = (valid - 1) * hop + n_fft // 2
+        if guard_end < n_fft:
+            valid = 0
+        else:
+            self._ifb_guard = buffer[guard_end - n_fft : guard_end].copy()
+        self._ifb_valid = valid
+        self._ifb_dropped = int(self.left_samples_dropped)
+        self._ifb_length = len(buffer)
+        return torch.from_numpy(power[:frames])
+
+    def _ifb_rebase(self, buffer: np.ndarray, n_fft: int, hop: int) -> int:
+        """Rows still exact for ``buffer``: the kept rows (moved down after a slide), or 0."""
+        valid = self._ifb_valid
+        if valid == 0:
+            return 0
+        shift = int(self.left_samples_dropped) - self._ifb_dropped
+        guard_end = (valid - 1) * hop + n_fft // 2 - shift
+        if (
+            shift < 0
+            or shift % hop
+            or len(buffer) + shift < self._ifb_length
+            or guard_end < n_fft
+            or not np.array_equal(buffer[guard_end - n_fft : guard_end], self._ifb_guard)
+        ):
+            return 0  # not the same stream (or not a hop-aligned slide of it): start over
+        if shift == 0:
+            return valid
+        drop = shift // hop
+        valid -= drop
+        if valid <= _ALIGN_FRAMES:
+            return 0
+        power = self._ifb_power
+        power[:valid] = power[drop : drop + valid]
+        # The new left edge's frames read the reflect pad.
+        power[:_ALIGN_FRAMES] = _stft_power(buffer, 0, _ALIGN_FRAMES, n_fft, hop, self._ifb_window)
+        return valid
+
+
+def _self_check(mel_processor: Any, incremental_cls: type) -> bool:  # Any: remote-code processor
+    """Compare the incremental and the original features on a probe stream (with a slide) once per class."""
+    fe = mel_processor.feature_extractor
+    saved = {name: getattr(fe, name) for name in ("dynamic_log_norm", "dynamic_range_db", "log_floor_db")}
+    probe = copy.copy(mel_processor)
+    probe.__class__ = incremental_cls
+    hop = int(fe.hop_length)
+    audio = (0.1 * np.random.default_rng(0).standard_normal(8 * int(mel_processor.sample_rate))).astype(np.float32)
+    try:
+        # Below and above the 5 s dynamic-range switch, then a slide of 100 hops.
+        for dropped, end in ((0, len(audio) // 3 + 37), (0, len(audio) - 11), (100 * hop, len(audio))):
+            probe.buffer, probe.left_samples_dropped = audio[dropped:end], dropped
+            expected = type(mel_processor)._extract_full(probe)
+            if not torch.equal(incremental_cls._extract_full(probe), expected):
+                return False
+        return True
+    except Exception:
+        logger.warning("MiniCPM-o incremental fbank self-check raised", exc_info=True)
+        return False
+    finally:
+        for name, value in saved.items():
+            setattr(fe, name, value)
+
+
+def enable_incremental_fbank(mel_processor: Any) -> bool:  # Any: remote-code StreamingMelProcessorExact
+    """Switch one session's streaming mel processor to the incremental ``_extract_full`` (idempotent).
+
+    The instance's class becomes a subclass of its remote class with the mixin
+    in front, so copies keep it. Returns False, leaving the full recompute, for
+    a processor of another shape or one that differs on the probe stream.
+    """
+    if isinstance(mel_processor, _IncrementalFbankMixin):
+        return True
+    cls = type(mel_processor)
+    fe = getattr(mel_processor, "feature_extractor", None)
+    if not callable(getattr(cls, "_extract_full", None)) or not all(
+        hasattr(fe, name) for name in _FEATURE_EXTRACTOR_ATTRS
+    ):
+        return False
+    if not all(hasattr(mel_processor, name) for name in ("buffer", "sample_rate", "left_samples_dropped")):
+        return False
+    incremental_cls = _INCREMENTAL_CLASSES.get(cls)
+    if incremental_cls is None:
+        incremental_cls = type(f"Incremental{cls.__name__}", (_IncrementalFbankMixin, cls), {"__module__": __name__})
+        _INCREMENTAL_CLASSES[cls] = incremental_cls
+    key = (cls, type(fe))
+    if key not in _SELF_CHECKED:
+        _SELF_CHECKED[key] = _self_check(mel_processor, incremental_cls)
+        if not _SELF_CHECKED[key]:
+            logger.warning(
+                "MiniCPM-o incremental fbank differs from %s._extract_full on the probe stream; "
+                "the streaming fbank stays on the full recompute",
+                cls.__name__,
+            )
+    if not _SELF_CHECKED[key]:
+        return False
+    mel_processor.__class__ = incremental_cls
+    return True

vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py

@@ -5,15 +5,19 @@ from __future__ import annotations
 
 import base64
 import copy
+import os
 import time
-from collections.abc import Callable
+from collections.abc import Callable, Generator
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import suppress
 from dataclasses import dataclass, field
 from threading import Lock
+from types import SimpleNamespace
 from typing import TYPE_CHECKING, Any, cast
 
 import numpy as np
 from numpy.typing import NDArray
+from vllm.logger import init_logger
 
 from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import MiniCPMO45DuplexPolicy
 
@@ -22,9 +26,32 @@ if TYPE_CHECKING:
     from PIL import Image
     from torch import nn
 
+logger = init_logger(__name__)
+
 _MINICPMO45_SPECIAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS
 _MINICPMO45_OPTIONAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.OPTIONAL_TOKEN_FIELDS
 _MINICPMO45_PROCESSOR_LOAD_LOCK = Lock()
+_MINICPMO45_FRAME_POOL_LOCK = Lock()
+_MINICPMO45_FRAME_POOL: ThreadPoolExecutor | None = None
+
+
+def _optional_int(value: Any) -> int | None:
+    try:
+        return int(value) if value is not None else None
+    except (TypeError, ValueError):
+        return None
+
+
+def _frame_preprocess_pool() -> ThreadPoolExecutor | None:
+    """Shared pool for camera-frame JPEG decode and slicing (both release the GIL); ``None`` on one CPU."""
+    global _MINICPMO45_FRAME_POOL
+    workers = min(4, os.cpu_count() or 1)
+    if workers <= 1:
+        return None
+    with _MINICPMO45_FRAME_POOL_LOCK:
+        if _MINICPMO45_FRAME_POOL is None:
+            _MINICPMO45_FRAME_POOL = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="minicpmo45-frames")
+    return _MINICPMO45_FRAME_POOL
 
 
 @dataclass
@@ -55,8 +82,13 @@ class _MiniCPMO45Stage0SessionState:
     prepared_inputs_embeds: torch.Tensor | None = None
     prepared_input_token_ids: list[int] = field(default_factory=list)
     prepared_result: dict[str, object] = field(default_factory=dict)
-    #: Opaque transformers cache, handed back to the model unchanged.
+    #: Streaming-encoder state: a ``StreamingAudioKVCache`` on the batched
+    #: path, else the opaque transformers cache handed back to the model.
     audio_past_key_values: object | None = None
+    #: Append identity of the last prefill that did work (successful or not).
+    attempted_append_identity: tuple[int | None, int] | None = None
+    #: This step's batched prefill result, taken by the append's own preprocess.
+    staged_prefill: tuple[tuple[int | None, int], dict[str, object]] | None = None
     pending_terminator_token: int | None = None
     last_terminator_token: int | None = None
     pending_speech_context: bool = False
@@ -65,6 +97,18 @@ class _MiniCPMO45Stage0SessionState:
     generated_tokens: list[int] = field(default_factory=list)
 
 
+@dataclass(frozen=True, slots=True)
+class _MiniCPMO45AudioEncodeRequest:
+    """One streaming unit a prefill needs encoded; ``_stage_prefill_steps`` yields it."""
+
+    #: Any: transformers BatchFeature (or the dict fallback of _process_streaming_audio).
+    batch_feature: Any
+    state: _MiniCPMO45Stage0SessionState | None
+
+
+_PrefillSteps = Generator[_MiniCPMO45AudioEncodeRequest, "torch.Tensor | None", dict[str, object]]
+
+
 class MiniCPMO45Stage0DuplexRuntime:
     """Build scheduler-owned MiniCPM-o 4.5 Stage0 duplex inputs."""
 
@@ -87,6 +131,8 @@ class MiniCPMO45Stage0DuplexRuntime:
         self.model_path = model_path
         self.device = device
         self.sessions: dict[tuple[str, int], _MiniCPMO45Stage0SessionState] = {}
+        # (session_id, epoch, seq) -> (frame payload, per-frame embeddings) of this runner step.
+        self._prefetched_vision: dict[tuple[str, int | None, int], tuple[tuple[object, ...], list[Any]]] = {}
         self.thinker = getattr(stage_model, "thinker", None) or getattr(stage_model, "model", None) or stage_model
         self.processor = (
             getattr(stage_model, "processor", None)
@@ -101,6 +147,49 @@ class MiniCPMO45Stage0DuplexRuntime:
         self._init_token_ids()
         if self.tokenizer is not None:
             self._require_special_token_ids()
+        # ``duplex_incremental_fbank`` (HF override, default off): see incremental_fbank.py.
+        self._incremental_fbank = (
+            getattr(getattr(stage_model, "config", None), "duplex_incremental_fbank", False) is True
+        )
+        # Built with the model, before its weights load and it switches to eval mode:
+        # the streaming encoder's CUDA graphs are captured later, by build_audio_cuda_graph().
+        self._audio_graph_attempted = False
+
+    def build_audio_cuda_graph(self) -> bool:
+        """Capture the streaming-encoder CUDA graphs once weights are loaded (idempotent; failures stay eager).
+
+        The steady unit's mel-frame count comes from a probe processor configured like a session's: the
+        shared processor's default chunk is no model unit, and a first unit has its own size.
+        """
+        thinker = getattr(self, "thinker", None)
+        config = getattr(thinker, "config", None)
+        attempted = getattr(self, "_audio_graph_attempted", False)
+        self._audio_graph_attempted = True
+        if attempted or not bool(getattr(config, "duplex_audio_encoder_cuda_graph", True)):
+            return False
+        apm = getattr(self.thinker, "apm", None)
+        if getattr(apm, "training", False) is True and callable(getattr(apm, "eval", None)):
+            apm.eval()  # inference only; the loader's model.eval() covers it as well
+        build = getattr(self.thinker, "build_streaming_audio_graph_encoder", None)
+        if not callable(build) or self.processor is None:
+            return False
+        try:
+            probe = self._configure_streaming_processor(SimpleNamespace(streaming_processor=None))
+            if probe is None:
+                return False
+            for _ in range(2):  # the first unit, then the steady one
+                silence = np.zeros(self._streaming_chunk_size(probe), dtype=np.float32)
+                feature = probe.process_audio_streaming(silence, reset=False)
+            # A BatchFeature-like mapping or a (features, metadata) tuple, by processor version.
+            frames = int((feature["audio_features"] if isinstance(feature, dict) else feature[0]).shape[-1])
+        except Exception:
+            logger.info(
+                "MiniCPM-o could not derive the steady streaming-audio unit shape; "
+                "the streaming audio encoder stays eager",
+                exc_info=True,
+            )
+            return False
+        return bool(build(unit_frames=frames))
 
     def _stage_runtime_ready(self) -> bool:
         return self.processor is not None and self.tokenizer is not None and self.thinker is not None
@@ -138,7 +227,7 @@ class MiniCPMO45Stage0DuplexRuntime:
             reset_streaming = getattr(processor, "reset_streaming", None)
             if callable(reset_streaming):
                 reset_streaming()
-            return processor
+            return self._maybe_enable_incremental_fbank(processor)
         configure_streaming = getattr(processor, "configure_streaming", None)
         if callable(configure_streaming):
             configure_streaming(
@@ -147,6 +236,19 @@ class MiniCPMO45Stage0DuplexRuntime:
                 slide_trigger_seconds=30.0,
                 slide_stride_seconds=10.0,
             )
+        return self._maybe_enable_incremental_fbank(processor)
+
+    def _maybe_enable_incremental_fbank(self, processor: Any) -> Any:  # Any: remote-code processor
+        """Put the processor's streaming mel front end on the incremental path (``duplex_incremental_fbank``)."""
+        mel_processor = getattr(processor, "_streaming_mel_processor", None)
+        if getattr(self, "_incremental_fbank", False) and mel_processor is not None:
+            from vllm_omni.model_executor.models.minicpmo_4_5.duplex.incremental_fbank import enable_incremental_fbank
+
+            if not enable_incremental_fbank(mel_processor):
+                logger.warning_once(
+                    "MiniCPM-o duplex_incremental_fbank: %s is not supported; the streaming fbank recomputes in full",
+                    type(mel_processor).__name__,
+                )
         return processor
 
     def _prepare_session_context(
@@ -208,24 +310,90 @@ class MiniCPMO45Stage0DuplexRuntime:
             state.current_turn_ended = False
             state.pending_speech_response_open = True
 
-    def _stage_prefill_embeddings_only(
+    def _stage_prefill_embeddings_only(self, state, audio_waveform, **kwargs: Any) -> dict[str, object]:
+        """Run :meth:`_stage_prefill_steps` for one append, encoding its units one at a time."""
+        return self._run_prefill_steps([self._stage_prefill_steps(state, audio_waveform, **kwargs)])[0]
+
+    @staticmethod
+    def take_staged_prefill(
+        state: _MiniCPMO45Stage0SessionState, epoch: int | None, seq: int | None
+    ) -> dict[str, object] | None:
+        """Pop the append :meth:`stage_prefill_batch` built for ``(epoch, seq)`` this step, if any."""
+        staged = state.staged_prefill
+        if staged is None:
+            return None
+        state.staged_prefill = None
+        if seq is None or staged[0] != (epoch, seq):
+            return None
+        return staged[1]
+
+    def batches_audio_encoder(self) -> bool:
+        """Whether the thinker can encode several sessions' audio units in one call."""
+        return self._batched_audio_encoder() is not None
+
+    def needs_prefill(self, state: _MiniCPMO45Stage0SessionState | None, epoch: int | None, seq: int | None) -> bool:
+        """Whether an append has not been built (or attempted) for this session yet."""
+        if seq is None or state is None:
+            return seq is not None
+        built = state.prepared_append_identity == (epoch, seq) and state.prepared_inputs_embeds is not None
+        return not built and state.attempted_append_identity != (epoch, seq)
+
+    def stage_prefill_batch(
+        self,
+        appends: list[tuple[_MiniCPMO45Stage0SessionState, NDArray[np.float32] | None, dict[str, Any]]],
+    ) -> None:
+        """Build distinct sessions' appends (each with a ``seq``), batching their ``k``-th unit encodes.
+
+        Results are staged on the session for the append's own ``preprocess`` this step.
+        """
+        if len({id(state) for state, _, _ in appends}) != len(appends):
+            raise ValueError("stage_prefill_batch needs one append per session")
+        steps = [self._stage_prefill_steps(state, audio, **kwargs) for state, audio, kwargs in appends]
+        for (state, _, kwargs), result in zip(appends, self._run_prefill_steps(steps), strict=True):
+            state.staged_prefill = ((kwargs.get("epoch"), kwargs["seq"]), result)
+
+    def _run_prefill_steps(self, steps: list[_PrefillSteps]) -> list[dict[str, object]]:
+        """Drive prefills to completion, encoding each round's pending units together."""
+        results: list[dict[str, object] | None] = [None] * len(steps)
+        pending: list[tuple[int, _MiniCPMO45AudioEncodeRequest]] = []
+
+        def advance(index: int, value: torch.Tensor | None, *, first: bool = False) -> None:
+            try:
+                request = next(steps[index]) if first else steps[index].send(value)
+            except StopIteration as stop:
+                results[index] = stop.value
+                return
+            pending.append((index, request))
+
+        for index in range(len(steps)):
+            advance(index, None, first=True)
+        while pending:
+            round_, pending = pending, []
+            embeds = self._stage_audio_embeddings_batch([request for _, request in round_])
+            for (index, _), embed in zip(round_, embeds, strict=True):
+                advance(index, embed)
+        return cast("list[dict[str, object]]", results)
+
+    def _stage_prefill_steps(
         self,
         state: _MiniCPMO45Stage0SessionState,
         audio_waveform: NDArray[np.float32] | None,
         *,
         video_frames: list[Image.Image] | None = None,
+        encoded_frames: list[list[torch.Tensor]] | None = None,
         epoch: int | None = None,
         seq: int | None = None,
         is_speech: bool = False,
         final: bool = False,
         stage0_window: dict[str, object] | None = None,
         stage0_reanchor: dict[str, object] | None = None,
-    ) -> dict[str, object]:
+    ) -> _PrefillSteps:
         """Build scheduler-owned Stage0 input embeddings for one audio append.
 
         Unlike the legacy worker-control path, this method never calls an eager
         model forward. The normal vLLM runner consumes the returned embeddings
         and owns attention metadata, block tables, KV cache, and sampling.
+        Yields each unit's encoder request; the caller sends back its embeddings.
         """
         start_time = time.time()
         processor = self._configure_streaming_processor(state)
@@ -239,6 +407,7 @@ class MiniCPMO45Stage0DuplexRuntime:
             result["inputs_embeds"] = state.prepared_inputs_embeds
             result["input_token_ids"] = list(state.prepared_input_token_ids)
             return result
+        state.attempted_append_identity = append_identity
         self._require_special_token_ids()
         if isinstance(stage0_window, dict):
             completed_ids = stage0_window.get("completed_token_ids")
@@ -252,11 +421,12 @@ class MiniCPMO45Stage0DuplexRuntime:
                 state.pending_terminator_token = completed_terminator
         if audio_waveform is None or len(audio_waveform) == 0:
             return self._stage_prefill_result(False, start_time, "empty audio")
+        num_frames = len(encoded_frames) if encoded_frames is not None else len(video_frames or ())
         state.audio_buffer = np.concatenate([state.audio_buffer, np.asarray(audio_waveform, dtype=np.float32)])
         chunk_size = self._streaming_chunk_size(processor)
         self._pad_first_audio_chunk_if_needed(state, processor)
         if len(state.audio_buffer) < chunk_size:
-            if video_frames:
+            if num_frames:
                 # Frames belong to the unit this append closes, and the
                 # scheduler already reserved their vision tokens against this
                 # append. A frame arriving before its unit can close has
@@ -265,7 +435,7 @@ class MiniCPMO45Stage0DuplexRuntime:
                 return self._stage_prefill_result(
                     False,
                     start_time,
-                    f"{len(video_frames)} video frame(s) on an append that closes no model unit "
+                    f"{num_frames} video frame(s) on an append that closes no model unit "
                     f"(need {chunk_size} samples, only {len(state.audio_buffer)}); "
                     "send frames on unit boundaries",
                 )
@@ -278,10 +448,11 @@ class MiniCPMO45Stage0DuplexRuntime:
         # carry them, mirroring official streaming_prefill (feed <unit>, then
         # every frame_list entry as its own <image> block, then the audio).
         frame_blocks: list[list[torch.Tensor]] = []
-        if video_frames:
+        if num_frames:
             self._require_vision_token_ids()
-            encoded_frames = self._stage_vision_embeddings(video_frames)
-            if encoded_frames is None or len(encoded_frames) != len(video_frames):
+            if encoded_frames is None:
+                encoded_frames = self._stage_vision_embeddings(video_frames or [])
+            if encoded_frames is None or len(encoded_frames) != num_frames:
                 return self._stage_prefill_result(False, start_time, "streaming vision embedding failed")
             if any(not slices for slices in encoded_frames):
                 return self._stage_prefill_result(False, start_time, "streaming vision embedding failed")
@@ -312,7 +483,7 @@ class MiniCPMO45Stage0DuplexRuntime:
             ):
                 with suppress(Exception):
                     setattr(batch_feature, name, value)
-            audio_embeds = self._stage_audio_embeddings(batch_feature, state=state)
+            audio_embeds = yield _MiniCPMO45AudioEncodeRequest(batch_feature, state)
             if audio_embeds is None:
                 if units_built == 0:
                     return self._stage_prefill_result(False, start_time, "streaming audio embedding returned empty")
@@ -623,7 +794,8 @@ class MiniCPMO45Stage0DuplexRuntime:
     def _embed_token(self, token_id: int) -> torch.Tensor:
         import torch
 
-        token = torch.tensor([int(token_id)], dtype=torch.long, device=self._model_device())
+        # A device-side fill: ``torch.tensor(..., device=...)`` is a blocking copy, 2-3 per unit.
+        token = torch.full((1,), int(token_id), dtype=torch.long, device=self._model_device())
         embedder = self._token_embedder()
         embeds = embedder(token)
         return self._as_2d_tensor(embeds)
@@ -759,7 +931,11 @@ class MiniCPMO45Stage0DuplexRuntime:
         state: _MiniCPMO45Stage0SessionState | None = None,
     ) -> torch.Tensor | None:
         if hasattr(batch_feature, "to"):
-            batch_feature = batch_feature.to(self.device)
+            batch_feature = self._audio_feature_to_device(batch_feature)
+        if state is not None and hasattr(state.audio_past_key_values, "to_legacy_cache"):
+            # A batched-path session whose unit this path has to take: hand the
+            # history over as the transformers cache the per-session path grows.
+            state.audio_past_key_values = state.audio_past_key_values.to_legacy_cache()
         self._ensure_dynamic_cache_compat()
         has_audio_cache = state is not None and hasattr(self.thinker, "audio_past_key_values")
         previous_audio_past_key_values = (
@@ -799,6 +975,94 @@ class MiniCPMO45Stage0DuplexRuntime:
             if has_audio_cache:
                 self.thinker.audio_past_key_values = previous_audio_past_key_values
 
+    def _stage_audio_embeddings_batch(
+        self,
+        requests: list[_MiniCPMO45AudioEncodeRequest],
+    ) -> list[torch.Tensor | None]:
+        """Encode one unit per request (distinct sessions); those the batched encoder takes share one pass."""
+        results: list[torch.Tensor | None] = [None] * len(requests)
+        encoder = self._batched_audio_encoder()
+        batched: list[tuple[int, _MiniCPMO45AudioEncodeRequest, Any]] = []
+        for index, request in enumerate(requests):
+            chunk = self._batched_audio_chunk(request) if encoder is not None else None
+            if chunk is None:
+                results[index] = self._stage_audio_embeddings(request.batch_feature, state=request.state)
+            else:
+                batched.append((index, request, chunk))
+        if batched:
+            assert encoder is not None
+            embeds, caches = encoder([chunk for _, _, chunk in batched])
+            for (index, request, _), embed, cache in zip(batched, embeds, caches, strict=True):
+                assert request.state is not None
+                request.state.audio_past_key_values = cache
+                results[index] = embed
+        return results
+
+    def _batched_audio_encoder(self) -> Callable[[list[Any]], tuple[list[Any], list[Any]]] | None:
+        encoder = getattr(self, "_batched_audio_encoder_fn", False)
+        if encoder is False:
+            if not getattr(self, "_audio_graph_attempted", False):
+                self.build_audio_cuda_graph()
+            thinker = getattr(self, "thinker", None)
+            batch = getattr(thinker, "get_audio_embedding_streaming_batch", None)
+            supported = getattr(thinker, "supports_streaming_audio_batch", None)
+            encoder = batch if callable(batch) and callable(supported) and supported() else None
+            self._batched_audio_encoder_fn = encoder
+        return encoder
+
+    @staticmethod
+    def _batched_audio_chunk(request: _MiniCPMO45AudioEncodeRequest) -> Any | None:  # Any: StreamingAudioChunk
+        import torch
+
+        from vllm_omni.model_executor.models.minicpmo_4_5.streaming_audio_encoder import (
+            StreamingAudioChunk,
+            StreamingAudioKVCache,
+        )
+
+        cache = getattr(request.state, "audio_past_key_values", None)
+        if request.state is None or (cache is not None and not isinstance(cache, StreamingAudioKVCache)):
+            return None
+        feature = request.batch_feature
+        try:
+            features, lengths = feature["audio_features"], feature["audio_feature_lens"]
+        except (KeyError, TypeError):
+            return None
+        if not isinstance(features, torch.Tensor) or features.ndim != 3 or features.shape[0] != 1:
+            return None
+        if isinstance(lengths, (list, tuple)) and len(lengths) == 1:
+            lengths = lengths[0]
+        if not isinstance(lengths, torch.Tensor) or lengths.numel() != 1:
+            return None
+        # Same extra-context frames the per-session call passes.
+        return StreamingAudioChunk(
+            features=features,
+            cache=cache,
+            prefix_extra_frames=0 if int(getattr(feature, "chunk_idx", 0)) == 0 else 2,
+            suffix_extra_frames=2,
+            feature_length=int(lengths.reshape(-1)[0]),
+        )
+
+    def _audio_feature_to_device(self, batch_feature: Any) -> Any:  # Any: transformers BatchFeature
+        """Move one unit's log-mel features to the encoder without a host sync; integer lengths stay on the host."""
+        import torch
+
+        from vllm_omni.utils.device_copy import to_device_nonblocking
+
+        data = getattr(batch_feature, "data", None)
+        streaming = any(
+            callable(getattr(target, "get_audio_embedding_streaming", None))
+            for target in (self.stage_model, self.thinker)
+        )
+        if torch.device(self.device).type != "cuda" or not isinstance(data, dict) or not streaming:
+            return batch_feature.to(self.device)
+        batch_feature.data = {
+            key: to_device_nonblocking(value, self.device)
+            if isinstance(value, torch.Tensor) and torch.is_floating_point(value)
+            else value
+            for key, value in data.items()
+        }
+        return batch_feature
+
     @staticmethod
     def _decode_ref_audio_from_session_config(session_config: dict[str, object]) -> NDArray[np.float32] | None:
         from vllm_omni.model_executor.models.minicpmo_4_5.duplex.input import decode_native_ref_audio_from_config
@@ -1092,7 +1356,8 @@ class MiniCPMO45Stage0DuplexRuntime:
             return [1]
         return [2] + [1] * (frame_count - 1)
 
-    def _encode_processed_vision(self, processed: Any) -> list[torch.Tensor] | None:  # Any: BatchFeature
+    def _encode_processed_vision_batch(self, processed: list[Any]) -> list[list[torch.Tensor]] | None:  # BatchFeature
+        """Encode every slice of the ``processed`` frames in one vision-tower call."""
         targets = (self.stage_model, self.thinker, getattr(self.stage_model, "model", None))
         for target in targets:
             if target is None:
@@ -1106,18 +1371,23 @@ class MiniCPMO45Stage0DuplexRuntime:
 
                 vpm_param = next(vpm.parameters())
                 device, dtype = vpm_param.device, vpm_param.dtype
-                pixel_nested = processed["pixel_values"]
-                tgt_nested = processed["tgt_sizes"]
                 flat_pixels: list[torch.Tensor] = []
                 flat_tgt: list[torch.Tensor] = []
-                for image_slices, image_tgt in zip(pixel_nested, tgt_nested):
-                    for slice_pixels in image_slices:
-                        flat_pixels.append(slice_pixels.to(device=device, dtype=dtype))
-                    tgt_tensor = image_tgt if hasattr(image_tgt, "reshape") else torch.tensor(image_tgt)
-                    flat_tgt.append(tgt_tensor.reshape(-1, 2))
+                slices_per_frame: list[int] = []
+                for frame in processed:
+                    num_slices = 0
+                    for image_slices, image_tgt in zip(frame["pixel_values"], frame["tgt_sizes"]):
+                        for slice_pixels in image_slices:
+                            flat_pixels.append(slice_pixels.to(device=device, dtype=dtype))
+                            num_slices += 1
+                        tgt_tensor = image_tgt if hasattr(image_tgt, "reshape") else torch.tensor(image_tgt)
+                        flat_tgt.append(torch.as_tensor(tgt_tensor).reshape(-1, 2))
+                    if num_slices == 0:
+                        return None
+                    slices_per_frame.append(num_slices)
                 if not flat_pixels:
                     return None
-                tgt_sizes = torch.cat(flat_tgt, dim=0).to(device=device, dtype=torch.int64)
+                tgt_sizes = torch.cat(flat_tgt, dim=0).to(device="cpu", dtype=torch.int64)
                 with torch.no_grad():
                     hidden = get_hidden({"pixel_values": flat_pixels, "tgt_sizes": tgt_sizes})
             except Exception:  # noqa: BLE001 - prefill fails with a reason
@@ -1132,7 +1402,10 @@ class MiniCPMO45Stage0DuplexRuntime:
                 if embed_dtype is not None and block_2d.dtype != embed_dtype:
                     block_2d = block_2d.to(dtype=embed_dtype)
                 out.append(block_2d)
-            return out
+            if len(out) != sum(slices_per_frame):
+                return None
+            blocks = iter(out)
+            return [[next(blocks) for _ in range(num_slices)] for num_slices in slices_per_frame]
         return None
 
     def _stage_vision_embeddings(self, frames: list[Image.Image]) -> list[list[torch.Tensor]] | None:
@@ -1146,17 +1419,96 @@ class MiniCPMO45Stage0DuplexRuntime:
         Frames are processed one at a time because ``process_image([a, b])``
         packs both PILs into one batch item.
         """
+        processed = self._preprocess_frames(frames)
+        return None if processed is None else self._encode_processed_vision_batch(processed)
+
+    def _preprocess_frames(self, frames: list[Image.Image]) -> list[Any] | None:  # Any: BatchFeature
+        """The CPU half of ``_stage_vision_embeddings``: each frame's processed slices; ``None`` on failure."""
         process_image = getattr(self.processor, "process_image", None)
         if not callable(process_image):
             return None
-        out: list[list[torch.Tensor]] = []
+        out: list[Any] = []
         for frame, max_slices in zip(frames, self._official_max_slice_nums(len(frames))):
             try:
-                processed = process_image([frame], max_slice_nums=max_slices)
+                out.append(process_image([frame], max_slice_nums=max_slices))
             except Exception:  # noqa: BLE001 - prefill fails with a reason
                 return None
-            encoded = self._encode_processed_vision(processed)
-            if encoded is None:
-                return None
-            out.append(encoded)
         return out
+
+    def frame_kwargs(self, duplex: dict[str, Any], payload: dict[str, object]) -> dict[str, Any]:
+        """``encoded_frames`` prefetched for this append, else its decoded ``video_frames`` (``ValueError``)."""
+        encoded = self.take_prefetched_vision(duplex)
+        if encoded is not None:
+            return {"encoded_frames": encoded}
+        return {"video_frames": self._decode_video_frames_payload(payload)}
+
+    @staticmethod
+    def _vision_prefetch_key(duplex: dict[str, Any]) -> tuple[str, int | None, int] | None:
+        """``(session_id, epoch, seq)`` of an append, parsed like ``preprocess`` does; ``None`` if unkeyed."""
+        session_id, seq = str(duplex.get("session_id") or ""), _optional_int(duplex.get("seq"))
+        return (session_id, _optional_int(duplex.get("epoch")), seq) if session_id and seq is not None else None
+
+    def prefetch_vision(self, appends: list[dict[str, Any]]) -> None:
+        """Encode this step's appends' camera frames in ``vision_batch_size``-append tower calls.
+
+        Decode and slicing run on the frame pool, the next chunk's overlapping this chunk's encode.
+        Results are keyed by append and checked against its frame payload on use.
+        """
+        self._prefetched_vision.clear()
+        jobs: list[tuple[tuple[str, int | None, int], tuple[object, ...]]] = []
+        for duplex in appends:
+            key = self._vision_prefetch_key(duplex)
+            payload = duplex.get("payload")
+            frames = payload.get("video_frames") if isinstance(payload, dict) else None
+            if key is None or not isinstance(frames, list) or not frames:
+                continue
+            state = self.sessions.get(key[0])
+            if state is None or state.prepared_append_identity != key[1:] or state.prepared_inputs_embeds is None:
+                jobs.append((key, tuple(frames)))  # else already built: preprocess replays the cached append
+        if not jobs or not callable(getattr(self.processor, "process_image", None)):
+            return
+
+        def preprocess(frames: tuple[object, ...]) -> list[Any] | None:
+            try:
+                decoded = self._decode_video_frames_payload({"video_frames": list(frames)})
+            except ValueError:
+                return None
+            return self._preprocess_frames(decoded) if decoded else None
+
+        pool = _frame_preprocess_pool() if len(jobs) > 1 else None
+
+        def start(chunk: list[tuple[Any, tuple[object, ...]]]) -> list[Any]:
+            return [pool.submit(preprocess, frames) if pool is not None else frames for _, frames in chunk]
+
+        def finish(started: list[Any]) -> list[list[Any] | None]:
+            return [item.result() if pool is not None else preprocess(item) for item in started]
+
+        size = self._vision_prefetch_chunk_size()
+        chunks = [jobs[begin : begin + size] for begin in range(0, len(jobs), size)]
+        started = start(chunks[0])
+        for index, chunk in enumerate(chunks):
+            # Submit the next chunk before this chunk's encode so they overlap.
+            following = start(chunks[index + 1]) if index + 1 < len(chunks) else []
+            ready = [(job, processed) for job, processed in zip(chunk, finish(started)) if processed is not None]
+            items = [item for _, processed in ready for item in processed]
+            encoded = self._encode_processed_vision_batch(items) if ready else None
+            blocks = iter(encoded or ())
+            for (key, frames), processed in ready if encoded is not None else ():
+                self._prefetched_vision[key] = (frames, [next(blocks) for _ in processed])
+            started = following
+
+    def _vision_prefetch_chunk_size(self) -> int:
+        """Jobs per ``prefetch_vision`` chunk: the checkpoint's ``vision_batch_size``, else 16."""
+        for target in (self.stage_model, self.thinker, getattr(self.stage_model, "model", None)):
+            size = getattr(getattr(target, "config", None), "vision_batch_size", None)
+            if isinstance(size, int) and size > 0:
+                return size
+        return 16
+
+    def take_prefetched_vision(self, duplex: dict[str, Any]) -> list[list[torch.Tensor]] | None:
+        """Pop this append's prefetched per-frame embeddings if its frame payload still matches."""
+        key = self._vision_prefetch_key(duplex)
+        frames_payload, encoded = self._prefetched_vision.pop(key, (None, None)) if key is not None else (None, None)
+        payload = duplex.get("payload")
+        frames = payload.get("video_frames") if isinstance(payload, dict) else None
+        return encoded if isinstance(frames, list) and tuple(frames) == frames_payload else None

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py

@@ -17,8 +17,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 from contextlib import suppress
+from dataclasses import dataclass
 from functools import cached_property
 from typing import Any
 
@@ -45,10 +46,35 @@ from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
 from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils.device_copy import index_to_device, to_device_nonblocking
 
 logger = init_logger(__name__)
 
 
+@dataclass(slots=True)
+class _MiniCPMO45PendingSamples:
+    """A deferred Stage-0 duplex step awaiting its host commit."""
+
+    host: torch.Tensor
+    event: Any
+    row_idxs: list[int]
+    stage2: list[bool]
+    rewinds: dict[int, Callable[[], None]]
+    token_ids: dict[str, int]
+    row_sessions: dict[int, str] | None
+    row_payloads: dict[int, dict[str, Any]] | None
+
+
+def _generator_rewind(generator: torch.Generator) -> Callable[[], None]:
+    """A call that puts ``generator`` back where it stands now: its Philox offset, else its full state."""
+    try:
+        offset = generator.get_offset()
+        return lambda: generator.set_offset(offset)
+    except RuntimeError:  # not a Philox generator (CPU)
+        state = generator.get_state()
+        return lambda: generator.set_state(state)
+
+
 @MULTIMODAL_REGISTRY.register_processor(
     MiniCPMO45OmniLLMMultiModalProcessor,
     info=MiniCPMO45OmniLLMProcessingInfo,
@@ -244,7 +270,14 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
     ) -> None:
         """Apply MiniCPM duplex policy before the standard model sampler."""
         del sampling_metadata
+        # The previous step's deferred decisions update the session state read below.
+        self._commit_minicpmo45_duplex_pending_samples()
         self._minicpmo45_active_duplex_rows = [row.row_idx for row in rows]
+        self._minicpmo45_duplex_row_sampling_host = {
+            row.row_idx: (row.temperature, row.top_k, row.top_p)
+            for row in rows
+            if row.temperature is not None and row.top_k is not None and row.top_p is not None
+        }
         self._minicpmo45_duplex_row_sessions = {
             row.row_idx: row.session_id for row in rows if row.session_id is not None
         }
@@ -306,8 +339,9 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
             if force_listen and segment_key in force_listen_segments:
                 continue
             if force_listen:
-                logits[row_idx, :] = float("-inf")
-                logits[row_idx, listen_id] = 0.0
+                # fill_ keeps the scalars on the host; an element assignment copies a scalar tensor and syncs.
+                logits[row_idx].fill_(float("-inf"))
+                logits[row_idx].narrow(0, listen_id, 1).fill_(0.0)
                 force_listen_segments.add(segment_key)
 
     # -------------------- Device utilities --------------------
@@ -418,53 +452,20 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
             embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids)
             return input_ids, embeds, {"duplex": {"prefill_success": False, "reason": "bad_duplex_payload"}}
 
-        session_key = session_id
-        state = helper.sessions.get(session_key)
-        if state is None:
-            from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
-                _MiniCPMO45Stage0SessionState,
-            )
-
-            state = _MiniCPMO45Stage0SessionState(session_id=session_id)
-            helper.sessions[session_key] = state
-            session_config = duplex.get("session_config")
-            session_config = dict(session_config) if isinstance(session_config, dict) else {}
-            runtime_config = duplex.get("runtime_config")
-            runtime_config = dict(runtime_config) if isinstance(runtime_config, dict) else {}
-            if hasattr(helper.thinker, "audio_past_key_values"):
-                helper.thinker.audio_past_key_values = None
-            helper._configure_streaming_processor(state)
-            helper._prepare_session_context(state, session_config, runtime_config=runtime_config)
-
-        audio_waveform = helper._decode_audio_payload(payload)
-        try:
-            video_frames = helper._decode_video_frames_payload(payload)
-        except ValueError as exc:
-            embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids)
-            return input_ids, embeds, {"duplex": {"prefill_success": False, "reason": str(exc)}}
-        seq = duplex.get("seq")
-        try:
-            seq = int(seq) if seq is not None else None
-        except (TypeError, ValueError):
-            seq = None
-        epoch = duplex.get("epoch")
-        try:
-            epoch = int(epoch) if epoch is not None else None
-        except (TypeError, ValueError):
-            epoch = None
-        result = helper._stage_prefill_embeddings_only(
-            state,
-            audio_waveform,
-            video_frames=video_frames,
-            epoch=epoch,
-            seq=seq,
-            is_speech=bool(payload.get("is_speech", False)),
-            final=bool(duplex.get("final")),
-            stage0_window=(duplex.get("stage0_window") if isinstance(duplex.get("stage0_window"), dict) else None),
-            stage0_reanchor=(
-                duplex.get("stage0_reanchor") if isinstance(duplex.get("stage0_reanchor"), dict) else None
-            ),
-        )
+        # A deferred sample of this session updates the latches its append reads.
+        self._commit_minicpmo45_duplex_pending_samples(session_ids={session_id})
+        state = self._minicpmo45_duplex_session_state(helper, session_id, duplex)
+        prefill_kwargs = self._minicpmo45_duplex_prefill_kwargs(duplex, payload)
+        seq = prefill_kwargs["seq"]
+        result = helper.take_staged_prefill(state, prefill_kwargs["epoch"], seq)
+        if result is None:  # not built by this step's preprocess_batch
+            audio_waveform = helper._decode_audio_payload(payload)
+            try:
+                prefill_kwargs.update(helper.frame_kwargs(duplex, payload))
+            except ValueError as exc:
+                embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids)
+                return input_ids, embeds, {"duplex": {"prefill_success": False, "reason": str(exc)}}
+            result = helper._stage_prefill_embeddings_only(state, audio_waveform, **prefill_kwargs)
         update_result = dict(result)
         if result.get("stage0_window_replaced") is True:
             window = duplex.get("stage0_window", {})
@@ -542,12 +543,105 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
         if len(input_token_ids) < span_len:
             input_token_ids.extend([pad_token_id] * (span_len - len(input_token_ids)))
         if input_token_ids:
-            req_input_ids = torch.tensor(input_token_ids, dtype=input_ids.dtype, device=input_ids.device)
+            # Pinned + non-blocking: a pageable copy would wait for all queued GPU work.
+            req_input_ids = index_to_device(input_token_ids, input_ids.device, input_ids.dtype)
             update_result["duplex_prompt_token_ids"] = full_input_token_ids
         else:
             req_input_ids = torch.full_like(input_ids, helper._required_token_id("unit_token_id"))
         return req_input_ids, req_embeds, {"duplex": update_result}
 
+    def preprocess_batch(
+        self,
+        *,
+        req_ids: list[str],
+        model_intermediate_buffer: dict[str, dict[str, Any]],
+        device: torch.device,
+    ) -> None:
+        """Runner hook before ``preprocess``: batch this step's duplex appends across sessions.
+
+        With two or more appends, camera frames share vision-tower calls (``prefetch_vision``) and
+        unbuilt appends share streaming-encoder passes (``stage_prefill_batch``); ``preprocess``
+        takes the staged results and builds or reports anything skipped here.
+        """
+        del device
+        if self.model_stage != "llm":
+            return
+        infos = [model_intermediate_buffer.get(req_id) for req_id in req_ids]
+        duplexes = [info.get("duplex") for info in infos if isinstance(info, dict)]
+        appends = [duplex for duplex in duplexes if isinstance(duplex, dict) and duplex.get("data_plane") is True]
+        helper = getattr(self, "_minicpmo45_duplex_data_plane_helper", None)
+        frame_appends = [d for d in appends if isinstance(d.get("payload"), dict) and d["payload"].get("video_frames")]
+        if len(frame_appends) >= 2:
+            try:
+                helper = self._duplex_data_plane_helper()
+                helper.prefetch_vision(frame_appends)
+            except Exception:  # noqa: BLE001 - preprocess encodes each append itself instead
+                logger.warning("MiniCPM-o duplex frame prefetch failed; encoding per request", exc_info=True)
+        elif helper is not None:
+            helper.prefetch_vision([])  # drop entries the previous step did not consume
+        if helper is None or len(req_ids) < 2 or not helper.batches_audio_encoder():
+            return
+        candidates: list[tuple[str, dict[str, Any], dict[str, Any], dict[str, Any]]] = []
+        seen: set[str] = set()
+        for duplex in appends:
+            session_id = str(duplex.get("session_id") or "")
+            payload = duplex.get("payload")
+            repeated = session_id in seen
+            seen.add(session_id)
+            if not session_id or not isinstance(payload, dict) or repeated:
+                continue  # a second append of one session keeps its order: preprocess
+            kwargs = self._minicpmo45_duplex_prefill_kwargs(duplex, payload)
+            if helper.needs_prefill(helper.sessions.get(session_id), kwargs["epoch"], kwargs["seq"]):
+                candidates.append((session_id, duplex, payload, kwargs))
+        if len(candidates) < 2:
+            return
+        self._commit_minicpmo45_duplex_pending_samples(session_ids={candidate[0] for candidate in candidates})
+        staged = []
+        for session_id, duplex, payload, kwargs in candidates:
+            try:
+                audio_waveform = helper._decode_audio_payload(payload)
+                kwargs.update(helper.frame_kwargs(duplex, payload))
+            except ValueError:
+                continue
+            staged.append((self._minicpmo45_duplex_session_state(helper, session_id, duplex), audio_waveform, kwargs))
+        if len(staged) >= 2:
+            helper.stage_prefill_batch(staged)
+
+    def _minicpmo45_duplex_session_state(self, helper, session_id: str, duplex: dict[str, Any]):
+        """The Stage-0 state of ``session_id``, created with its session context on first use."""
+        state = helper.sessions.get(session_id)
+        if state is None:
+            from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import (
+                _MiniCPMO45Stage0SessionState,
+            )
+
+            state = _MiniCPMO45Stage0SessionState(session_id=session_id)
+            helper.sessions[session_id] = state
+            session_config = duplex.get("session_config")
+            session_config = dict(session_config) if isinstance(session_config, dict) else {}
+            runtime_config = duplex.get("runtime_config")
+            runtime_config = dict(runtime_config) if isinstance(runtime_config, dict) else {}
+            if hasattr(helper.thinker, "audio_past_key_values"):
+                helper.thinker.audio_past_key_values = None
+            helper._configure_streaming_processor(state)
+            helper._prepare_session_context(state, session_config, runtime_config=runtime_config)
+        return state
+
+    @staticmethod
+    def _minicpmo45_duplex_prefill_kwargs(duplex: dict[str, Any], payload: dict[str, Any]) -> dict[str, Any]:
+        """Keyword arguments of ``_stage_prefill_embeddings_only`` for one append, frames aside."""
+        from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import _optional_int
+
+        window, reanchor = duplex.get("stage0_window"), duplex.get("stage0_reanchor")
+        return {
+            "epoch": _optional_int(duplex.get("epoch")),
+            "seq": _optional_int(duplex.get("seq")),
+            "is_speech": bool(payload.get("is_speech", False)),
+            "final": bool(duplex.get("final")),
+            "stage0_window": window if isinstance(window, dict) else None,
+            "stage0_reanchor": reanchor if isinstance(reanchor, dict) else None,
+        }
+
     def _duplex_data_plane_helper(self):
         helper = getattr(self, "_minicpmo45_duplex_data_plane_helper", None)
         if helper is not None:
@@ -694,13 +788,11 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
                     if isinstance(key, str) and isinstance(value, int) and value >= 0
                 }
                 if special_keys:
+                    # Host tensors: every consumer reads them on the host, and a pageable
+                    # host->device copy per row and key would wait for the whole forward.
                     multimodal_outputs["meta"] = {
                         key: [
-                            torch.tensor(
-                                [int(value)],
-                                dtype=torch.long,
-                                device=text_hidden_states.device,
-                            )
+                            torch.tensor([int(value)], dtype=torch.long)
                             if isinstance(value, int) and value >= 0
                             else None
                             for duplex_info in duplex_rows
@@ -808,6 +900,7 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
     ) -> SamplerOutput | None:
         if self.model_stage != "llm" or logits.ndim != 2 or logits.shape[0] == 0:
             return None
+        self._commit_minicpmo45_duplex_pending_samples()
         token_ids = self._minicpmo45_native_duplex_token_ids()
         unit_id = token_ids.get("unit_token_id", -1)
         if unit_id < 0:
@@ -823,120 +916,399 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
 
         chunk_terminators = self._minicpmo45_chunk_terminator_token_ids(token_ids)
         output_token_ids = getattr(sampling_metadata, "output_token_ids", None) or []
-        sampled_ids: list[int] = []
-        for row_idx in range(logits.shape[0]):
+        num_rows = logits.shape[0]
+        row_params = self._minicpmo45_duplex_row_params(sampling_metadata, num_rows)
+
+        # Lookahead after a chunk terminator: re-emit it and skip sampling.
+        sampled_ids: list[int] = [-1] * num_rows
+        pending_rows: list[int] = []
+        for row_idx in range(num_rows):
             accepted = output_token_ids[row_idx] if row_idx < len(output_token_ids) else []
             last_accepted = next((int(t) for t in reversed(accepted) if isinstance(t, int) and t >= 0), None)
             if last_accepted in chunk_terminators:
-                # Async scheduling runs one lookahead frame after the chunk
-                # terminator was sampled but before the scheduler observes
-                # the segment stop. The scheduler discards this frame's
-                # token, so decide nothing here: re-emit the terminator and
-                # leave the model-owned policy state exactly as the accepted
-                # history left it. The next append re-injects that terminator.
-                sampled_ids.append(last_accepted)
-                continue
-            row_logits = logits[row_idx : row_idx + 1].clone()
-            sampled = self._sample_minicpmo45_native_duplex_row(
-                row_logits,
-                sampling_metadata,
-                row_idx=row_idx,
-                token_ids=token_ids,
-            )
-            self._record_minicpmo45_duplex_terminator(row_idx, sampled, token_ids)
-            sampled_ids.append(sampled)
+                sampled_ids[row_idx] = last_accepted
+            else:
+                pending_rows.append(row_idx)
+
+        if pending_rows:
+            kwargs = {"row_idxs": pending_rows, "token_ids": token_ids, "row_params": row_params}
+            deferred = self._sample_minicpmo45_native_duplex_rows_deferred(logits, sampling_metadata, **kwargs)
+            if deferred is not None:
+                sampled = index_to_device(sampled_ids, logits.device, torch.int32)
+                sampled[index_to_device(pending_rows, logits.device)] = deferred.to(torch.int32)
+                return SamplerOutput(sampled_token_ids=sampled.unsqueeze(-1), logprobs_tensors=None)
+            synchronous = self._sample_minicpmo45_native_duplex_rows(logits, sampling_metadata, **kwargs)
+            for row_idx, sampled in zip(pending_rows, synchronous, strict=True):
+                sampled_ids[row_idx] = sampled
+                self._record_minicpmo45_duplex_terminator(row_idx, sampled, token_ids)
+
         return SamplerOutput(
-            sampled_token_ids=torch.tensor(sampled_ids, device=logits.device, dtype=torch.int32).unsqueeze(-1),
+            sampled_token_ids=index_to_device(sampled_ids, logits.device, torch.int32).unsqueeze(-1),
             logprobs_tensors=None,
         )
 
-    def _sample_minicpmo45_native_duplex_row(
+    def _minicpmo45_duplex_row_params(
+        self,
+        sampling_metadata: SamplingMetadata,
+        num_rows: int,
+    ) -> list[tuple[float, float, float]]:
+        """Every row's (temperature, top_k, top_p): host copies when present, else one device read per parameter."""
+        host = getattr(self, "_minicpmo45_duplex_row_sampling_host", None) or {}
+        columns = []
+        for column, (name, default) in enumerate((("temperature", 0.7), ("top_k", 100), ("top_p", 0.8))):
+            value = getattr(sampling_metadata, name, None)
+            if not isinstance(value, torch.Tensor):
+                columns.append([self._sampling_metadata_value(sampling_metadata, name, 0, default)] * num_rows)
+            elif value.numel() == num_rows and host.keys() >= set(range(num_rows)):
+                columns.append([float(host[row][column]) for row in range(num_rows)])
+            elif value.numel() == 0:
+                columns.append([default] * num_rows)
+            else:
+                flat = [float(v) for v in value.reshape(-1).tolist()]
+                columns.append([flat[min(row, len(flat) - 1)] for row in range(num_rows)])
+        return list(zip(*columns, strict=True))
+
+    def _minicpmo45_policy_limit(self, name: str, default: int) -> int:
+        """A per-chunk duplex limit set on the model (``0``/``None`` mean ``default``)."""
+        return int(getattr(self, name, default) or default)
+
+    def _minicpmo45_duplex_forbidden_index(self, token_ids: dict[str, int], vocab: int, device) -> torch.Tensor | None:
+        """``_minicpmo45_native_forbidden_token_ids`` inside the vocab, as a cached device index."""
+        key = (vocab, str(device))
+        cache = getattr(self, "_minicpmo45_duplex_forbidden_index_cache", None)
+        if cache is None or cache[0] != key:
+            forbidden = self._minicpmo45_native_forbidden_token_ids(token_ids)
+            valid = sorted({token_id for token_id in forbidden if 0 <= token_id < vocab})
+            cache = (key, torch.tensor(valid, dtype=torch.long, device=device) if valid else None)
+            self._minicpmo45_duplex_forbidden_index_cache = cache
+        return cache[1]
+
+    def _minicpmo45_duplex_cut_tables(self, token_ids: dict[str, int], vocab: int, device):
+        """Per-token decoded length and "never cut" mask for ``_maybe_cut_minicpmo45_native_duplex_text_chunk``.
+
+        Byte-level decoding is additive after a prefix ending on a character boundary:
+        ``len(decode(chunk + [t])) == len(decode(chunk)) + length[t]``. None (no device cut) without batch
+        decode, or with the whitespace clean-up that rewrites across token boundaries.
+        """
+        key = (vocab, str(device))
+        cache = getattr(self, "_minicpmo45_duplex_cut_tables_cache", None)
+        if cache is not None and cache[0] == key:
+            return cache[1]
+        tables = None
+        tokenizer = self._minicpmo45_tokenizer()
+        batch_decode = getattr(tokenizer, "batch_decode", None)
+        if callable(batch_decode) and not getattr(tokenizer, "clean_up_tokenization_spaces", False):
+            try:
+                num_tokens = min(vocab, len(tokenizer))
+                texts = batch_decode([[token_id] for token_id in range(num_tokens)], skip_special_tokens=True)
+            except Exception:
+                texts = None
+            if texts is not None and len(texts) == num_tokens:
+                lengths = torch.zeros(vocab, dtype=torch.int32)
+                lengths[:num_tokens] = torch.tensor([len(text) for text in texts], dtype=torch.int32)
+                never_cut = torch.ones(vocab, dtype=torch.bool)
+                never_cut[:num_tokens] = False
+                special = [t for t in self._minicpmo45_native_special_token_ids(token_ids) if 0 <= t < vocab]
+                if special:
+                    never_cut[torch.tensor(special, dtype=torch.long)] = True
+                tables = (lengths.to(device), never_cut.to(device))
+        self._minicpmo45_duplex_cut_tables_cache = (key, tables)
+        return tables
+
+    def _minicpmo45_duplex_cut_bases(
+        self,
+        row_idxs: list[int],
+        recent_tokens_by_row: dict[int, list[int]],
+        token_ids: dict[str, int],
+        max_chars: int,
+    ) -> list[int] | None:
+        """Each row's decoded current-chunk length (-1: never cut), or None when a row needs the host decode.
+
+        A chunk ending in U+FFFD may end inside a character the next token's bytes complete: not additive.
+        """
+        decode = getattr(self._minicpmo45_tokenizer(), "decode", None)
+        if max_chars <= 0 or not callable(decode):
+            return [-1] * len(row_idxs)
+        bases = []
+        for row_idx in row_idxs:
+            chunk = self._minicpmo45_current_chunk_tokens(recent_tokens_by_row[row_idx], token_ids)
+            try:
+                text = decode(chunk, skip_special_tokens=True)
+            except Exception:
+                return None
+            if not isinstance(text, str) or text.endswith("\ufffd"):
+                return None
+            bases.append(len(text))
+        return bases
+
+    @staticmethod
+    def _minicpmo45_duplex_recent_tokens(
+        sampling_metadata: SamplingMetadata, row_idxs: list[int]
+    ) -> dict[int, list[int]]:
+        """Each row's accepted output tokens (the current chunk's history), placeholders dropped."""
+        output_token_ids = getattr(sampling_metadata, "output_token_ids", None) or []
+        rows = {row_idx: output_token_ids[row_idx] if row_idx < len(output_token_ids) else () for row_idx in row_idxs}
+        return {row: [int(t) for t in tokens if isinstance(t, int) and t >= 0] for row, tokens in rows.items()}
+
+    def _minicpmo45_duplex_capped_rows(self, row_idxs: list[int], recent: dict[int, list[int]]) -> set[int]:
+        """Rows at their per-chunk speak-length cap: chunk_eos without sampling anything."""
+        limit = self._minicpmo45_policy_limit(
+            "max_new_speak_tokens_per_chunk", MiniCPMO45DuplexPolicy.DEFAULT_MAX_NEW_SPEAK_TOKENS_PER_CHUNK
+        )
+        request_max = {row_idx: self._minicpmo45_duplex_row_request_max_tokens(row_idx) for row_idx in row_idxs}
+        row_limit = {row: limit if request is None else min(limit, request) for row, request in request_max.items()}
+        return {row for row, row_cap in row_limit.items() if len(recent[row]) >= max(1, row_cap - 1)}
+
+    def _minicpmo45_duplex_boundary_hits(
+        self, logits: torch.Tensor, rows: list[int], chunk_eos_id: int, sampling_metadata: SamplingMetadata
+    ) -> torch.Tensor:
+        """Whether each row's boundary draw (the released StreamDecoder's first pass) picks chunk_eos.
+
+        The drawn token only matters when it is chunk_eos, so a Bernoulli at
+        ``P(sample == chunk_eos)`` drawn with the row's generator is the same
+        decision, without a full-vocabulary softmax and multinomial per row.
+        """
+        rows_logits = logits[index_to_device(rows, logits.device)]
+        if getattr(sampling_metadata, "all_greedy", False):
+            return torch.argmax(rows_logits, dim=-1) == chunk_eos_id
+        generators = getattr(sampling_metadata, "generators", {}) or {}
+        draws = [torch.rand((), generator=generators.get(row_idx), device=logits.device) for row_idx in rows]
+        return torch.stack(draws) < self._duplex_boundary_chunk_eos_probs(rows_logits, chunk_eos_id)
+
+    def _minicpmo45_duplex_stage2_draws(
+        self,
+        logits: torch.Tensor,
+        rows: list[int],
+        recent: dict[int, list[int]],
+        token_ids: dict[str, int],
+        row_params: list[tuple[float, float, float]],
+        sampling_metadata: SamplingMetadata,
+        rewinds: dict[int, Callable[[], None]] | None = None,
+    ) -> torch.Tensor:
+        """Each row's stage-2 token: forbidden ids masked, repetition penalty, then greedy or a candidate draw.
+
+        ``rewinds`` receives, per row position, a call undoing that row's draw.
+        """
+        device, vocab = logits.device, logits.shape[-1]
+        all_greedy = bool(getattr(sampling_metadata, "all_greedy", False))
+        generators = getattr(sampling_metadata, "generators", {}) or {}
+        stage2_logits = logits[index_to_device(rows, device)].clone()
+        forbidden_index = self._minicpmo45_duplex_forbidden_index(token_ids, vocab, device)
+        if forbidden_index is not None:
+            stage2_logits.index_fill_(1, forbidden_index, float("-inf"))
+        history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE
+        penalty_rows: list[int] = []
+        penalty_cols: list[int] = []
+        for local_idx, row_idx in enumerate(rows):
+            state = self._minicpmo45_duplex_state_for_row(row_idx)
+            repetition_tokens = getattr(state, "generated_tokens", None) or recent[row_idx]
+            penalized = [token_id for token_id in set(repetition_tokens[-history_size:]) if 0 <= token_id < vocab]
+            penalty_rows.extend([local_idx] * len(penalized))
+            penalty_cols.extend(penalized)
+        if penalty_rows:
+            # Distinct (row, col) pairs, so one indexed division equals dividing each once.
+            rows_t, cols_t = index_to_device(penalty_rows, device), index_to_device(penalty_cols, device)
+            stage2_logits[rows_t, cols_t] = stage2_logits[rows_t, cols_t] / 1.05
+
+        samples: list[torch.Tensor | None] = [None] * len(rows)
+        greedy = [local for local, row_idx in enumerate(rows) if all_greedy or float(row_params[row_idx][0]) <= 0]
+        if greedy:
+            greedy_sample = torch.argmax(stage2_logits[index_to_device(greedy, device)], dim=-1)
+            for pos, local_idx in enumerate(greedy):
+                samples[local_idx] = greedy_sample[pos : pos + 1]
+        drawn = [local_idx for local_idx in range(len(rows)) if samples[local_idx] is None]
+        # One uniform per row from its generator, in row order, then a batched draw per (top_k, top_p) group.
+        uniforms: dict[int, torch.Tensor] = {}
+        groups: dict[tuple[int, float], list[int]] = {}
+        for local_idx in drawn:
+            generator = generators.get(rows[local_idx])
+            if rewinds is not None:
+                rewinds[local_idx] = _generator_rewind(generator)
+            uniforms[local_idx] = torch.rand((), generator=generator, device=device)
+            _temperature, top_k, top_p = row_params[rows[local_idx]]
+            groups.setdefault((int(top_k), float(top_p)), []).append(local_idx)
+        if drawn:
+            temps = torch.tensor([float(row_params[row_idx][0]) for row_idx in rows], dtype=stage2_logits.dtype)
+            temps = to_device_nonblocking(temps, device).view(-1, 1)
+            for (top_k, top_p), members in groups.items():
+                # Device-side row gathers: a python-list index would stage a
+                # pageable host copy and block the host on queued GPU work.
+                members_t = index_to_device(members, device)
+                group_logits = stage2_logits[members_t] / temps[members_t]
+                candidates = self._duplex_top_k_top_p_candidates(group_logits, top_k=top_k, top_p=top_p)
+                group_uniforms = torch.stack([uniforms[local_idx] for local_idx in members])
+                draws = self._duplex_draw_candidates(group_logits, candidates, group_uniforms)
+                for i, local_idx in enumerate(members):
+                    samples[local_idx] = draws[i : i + 1]
+        return torch.cat(samples)  # type: ignore[arg-type]
+
+    def _sample_minicpmo45_native_duplex_rows(
         self,
         logits: torch.Tensor,
         sampling_metadata: SamplingMetadata,
         *,
-        row_idx: int,
+        row_idxs: list[int],
         token_ids: dict[str, int],
-    ) -> int:
+        row_params: list[tuple[float, float, float]],
+    ) -> list[int]:
+        """Sample duplex rows with at most two host reads: the boundary decisions, then the stage-2 tokens.
+
+        Per row: chunk_eos at the speak-length cap, else the boundary draw, else
+        the stage-2 draw, the chunk cut and the listen->tts_bos rewrite. Each
+        row's own generator draws as it would one row at a time; rows sharing
+        the process-wide generator draw every boundary before any stage-2 token.
+        """
         chunk_eos_id = token_ids.get("chunk_eos_token_id", -1)
-        generator = getattr(sampling_metadata, "generators", {}).get(row_idx)
-        output_token_ids = getattr(sampling_metadata, "output_token_ids", None) or []
-        raw_recent_tokens = output_token_ids[row_idx] if row_idx < len(output_token_ids) else []
-        recent_tokens = [int(token_id) for token_id in raw_recent_tokens if isinstance(token_id, int) and token_id >= 0]
-        temperature = float(self._sampling_metadata_value(sampling_metadata, "temperature", row_idx, 0.7))
-        top_k = int(self._sampling_metadata_value(sampling_metadata, "top_k", row_idx, 100))
-        top_p = float(self._sampling_metadata_value(sampling_metadata, "top_p", row_idx, 0.8))
-        state = self._minicpmo45_duplex_state_for_row(row_idx)
-        if chunk_eos_id >= 0 and chunk_eos_id < logits.shape[-1]:
-            max_speak_tokens = int(
-                getattr(
-                    self,
-                    "max_new_speak_tokens_per_chunk",
-                    MiniCPMO45DuplexPolicy.DEFAULT_MAX_NEW_SPEAK_TOKENS_PER_CHUNK,
-                )
-                or MiniCPMO45DuplexPolicy.DEFAULT_MAX_NEW_SPEAK_TOKENS_PER_CHUNK
+        recent = self._minicpmo45_duplex_recent_tokens(sampling_metadata, row_idxs)
+        resolved: dict[int, int] = {}
+        if 0 <= chunk_eos_id < logits.shape[-1]:
+            resolved = dict.fromkeys(self._minicpmo45_duplex_capped_rows(row_idxs, recent), chunk_eos_id)
+            boundary = [row_idx for row_idx in row_idxs if row_idx not in resolved]
+            if boundary:
+                hits = self._minicpmo45_duplex_boundary_hits(logits, boundary, chunk_eos_id, sampling_metadata)
+                resolved.update((row_idx, chunk_eos_id) for row_idx, hit in zip(boundary, hits.tolist()) if hit)
+        stage2_rows = [row_idx for row_idx in row_idxs if row_idx not in resolved]
+        if stage2_rows:
+            draws = self._minicpmo45_duplex_stage2_draws(
+                logits, stage2_rows, recent, token_ids, row_params, sampling_metadata
             )
-            request_max_tokens = self._minicpmo45_duplex_row_request_max_tokens(row_idx)
-            effective_max_speak_tokens = max_speak_tokens
-            if request_max_tokens is not None:
-                effective_max_speak_tokens = min(effective_max_speak_tokens, request_max_tokens)
-            if len(recent_tokens) >= max(1, effective_max_speak_tokens - 1):
-                return int(chunk_eos_id)
+            for row_idx, sampled in zip(stage2_rows, draws.tolist(), strict=True):
+                self._record_minicpmo45_duplex_generation_token(row_idx, sampled)
+                sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk(sampled, recent[row_idx], token_ids)
+                resolved[row_idx] = self._finalize_minicpmo45_native_duplex_sample(row_idx, sampled, token_ids)
+        return [int(resolved[row_idx]) for row_idx in row_idxs]
 
-            # Match the released StreamDecoder: first sample the original
-            # distribution only to preserve the model's own chunk boundary.
-            # If it does not choose chunk_eos, mask that token before the
-            # normal text/listen/turn sampling pass below.
-            if getattr(sampling_metadata, "all_greedy", False):
-                boundary_sample = int(torch.argmax(logits, dim=-1).item())
-            else:
-                boundary_probs = F.softmax(logits, dim=-1)
-                boundary_sample = int(torch.multinomial(boundary_probs, num_samples=1, generator=generator).item())
-            if boundary_sample == chunk_eos_id:
-                return int(chunk_eos_id)
+    def _sample_minicpmo45_native_duplex_rows_deferred(
+        self,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+        *,
+        row_idxs: list[int],
+        token_ids: dict[str, int],
+        row_params: list[tuple[float, float, float]],
+    ) -> torch.Tensor | None:
+        """``_sample_minicpmo45_native_duplex_rows`` decided on the device without a host sync, or ``None``.
 
-        logits = logits.clone()
-        forbidden = self._minicpmo45_native_forbidden_token_ids(token_ids)
-        if forbidden:
-            valid_forbidden = [token_id for token_id in forbidden if 0 <= token_id < logits.shape[-1]]
-            if valid_forbidden:
-                logits[:, valid_forbidden] = float("-inf")
-
-        generated_tokens = getattr(state, "generated_tokens", None)
-        repetition_tokens = generated_tokens if generated_tokens else recent_tokens
-        repetition_penalty = 1.05
-        if repetition_penalty != 1.0 and repetition_tokens:
-            history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE
-            for token_id in set(repetition_tokens[-history_size:]):
-                if token_id < 0 or token_id >= logits.shape[-1]:
-                    continue
-                logits[0, token_id] /= repetition_penalty
-
-        if getattr(sampling_metadata, "all_greedy", False) or temperature <= 0:
-            sampled = int(torch.argmax(logits, dim=-1).item())
-            self._record_minicpmo45_duplex_generation_token(row_idx, sampled)
-            sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk(
-                sampled,
-                recent_tokens,
-                token_ids,
+        Host state is committed a step later (``_commit_minicpmo45_duplex_pending_samples``). Stage 2
+        is drawn for every uncapped row; a row the boundary resolves has its generator rewound at
+        commit, so generators advance as on the synchronous path. Process-wide generator rows can not
+        be rewound alone and keep the step synchronous.
+        """
+        device, vocab = logits.device, logits.shape[-1]
+        chunk_eos_id = token_ids.get("chunk_eos_token_id", -1)
+        has_chunk_eos = 0 <= chunk_eos_id < vocab
+        all_greedy = bool(getattr(sampling_metadata, "all_greedy", False))
+        generators = getattr(sampling_metadata, "generators", {}) or {}
+        recent = self._minicpmo45_duplex_recent_tokens(sampling_metadata, row_idxs)
+        capped = self._minicpmo45_duplex_capped_rows(row_idxs, recent) if has_chunk_eos else set()
+        stage2_rows = [row_idx for row_idx in row_idxs if row_idx not in capped]
+        drawn_rows = [row_idx for row_idx in stage2_rows if not (all_greedy or float(row_params[row_idx][0]) <= 0)]
+        if any(generators.get(row_idx) is None for row_idx in drawn_rows):
+            return None
+        cut = None
+        if has_chunk_eos and stage2_rows:
+            max_chars = self._minicpmo45_policy_limit(
+                "max_speak_chars_per_chunk", MiniCPMO45DuplexPolicy.DEFAULT_MAX_SPEAK_CHARS_PER_CHUNK
             )
-            return self._finalize_minicpmo45_native_duplex_sample(row_idx, sampled, token_ids)
+            bases = self._minicpmo45_duplex_cut_bases(stage2_rows, recent, token_ids, max_chars)
+            if bases is None:
+                return None
+            if any(base >= 0 for base in bases):
+                tables = self._minicpmo45_duplex_cut_tables(token_ids, vocab, device)
+                if tables is None:
+                    return None
+                cut = (tables, bases, max_chars)
 
-        logits = logits / temperature
-        logits = self._top_k_top_p_filter(logits, top_k=top_k, top_p=top_p)
-        probs = F.softmax(logits, dim=-1)
-        sampled = int(torch.multinomial(probs, num_samples=1, generator=generator).item())
-        self._record_minicpmo45_duplex_generation_token(row_idx, sampled)
-        sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk(
-            sampled,
-            recent_tokens,
-            token_ids,
+        final = torch.full((len(row_idxs),), max(chunk_eos_id, 0), dtype=torch.long, device=device)
+        position = {row_idx: pos for pos, row_idx in enumerate(row_idxs)}
+        rewinds: dict[int, Callable[[], None]] = {}
+        if stage2_rows:
+            positions_t = index_to_device([position[row_idx] for row_idx in stage2_rows], device)
+            hit = (
+                self._minicpmo45_duplex_boundary_hits(logits, stage2_rows, chunk_eos_id, sampling_metadata)
+                if has_chunk_eos
+                else torch.zeros(len(stage2_rows), dtype=torch.bool, device=device)
+            )
+            rewinds_into = rewinds if has_chunk_eos else None
+            raw = self._minicpmo45_duplex_stage2_draws(
+                logits, stage2_rows, recent, token_ids, row_params, sampling_metadata, rewinds_into
+            )
+            token = raw
+            if cut is not None:
+                (lengths, never_cut), bases, max_chars = cut
+                base_t = index_to_device(bases, device)
+                cut_t = (base_t >= 0) & ~never_cut[raw] & (base_t + lengths[raw].long() >= max_chars)
+                token = torch.where(cut_t, torch.full_like(raw, chunk_eos_id), raw)
+            listen_id = token_ids.get("listen_token_id", -1)
+            if 0 <= listen_id:
+                # The host rewrite of a sampled listen (to tts_bos inside an open turn), decided per row.
+                rewritten = [
+                    self._finalize_minicpmo45_native_duplex_sample(r, listen_id, token_ids) for r in stage2_rows
+                ]
+                if any(value != listen_id for value in rewritten):
+                    token = torch.where(token == listen_id, index_to_device(rewritten, device), token)
+            if has_chunk_eos:
+                token = torch.where(hit, torch.full_like(token, chunk_eos_id), token)
+            final[positions_t] = token
+
+        # One pinned copy of (final, raw stage-2 sample, boundary hit) per row
+        # for the commit; the host reads it once the runner has waited for it.
+        packed = torch.stack([final, torch.full_like(final, -1), torch.zeros_like(final)], dim=1)
+        if stage2_rows:
+            packed[positions_t, 1] = raw
+            packed[positions_t, 2] = hit.long()
+        host = torch.empty(packed.shape, dtype=packed.dtype, pin_memory=device.type == "cuda")
+        host.copy_(packed, non_blocking=True)
+        event = torch.cuda.Event() if device.type == "cuda" else None
+        if event is not None:
+            event.record()
+        self._minicpmo45_duplex_pending_samples = _MiniCPMO45PendingSamples(
+            host=host,
+            event=event,
+            row_idxs=list(row_idxs),
+            stage2=[row_idx not in capped for row_idx in row_idxs],
+            rewinds={position[stage2_rows[local_idx]]: rewind for local_idx, rewind in rewinds.items()},
+            token_ids=token_ids,
+            row_sessions=getattr(self, "_minicpmo45_duplex_row_sessions", None),
+            row_payloads=getattr(self, "_minicpmo45_duplex_row_payloads", None),
         )
-        return self._finalize_minicpmo45_native_duplex_sample(
-            row_idx,
-            sampled,
-            token_ids,
+        return final
+
+    def _commit_minicpmo45_duplex_pending_samples(self, *, session_ids: set[str] | None = None) -> None:
+        """Apply the host state updates of the last deferred step, in the synchronous path's order.
+
+        With ``session_ids``, only when that step sampled one of them: a
+        caller building those sessions' next appends needs their latches, and
+        must not wait for an unrelated step otherwise.
+        """
+        pending = getattr(self, "_minicpmo45_duplex_pending_samples", None)
+        if pending is None:
+            return
+        if session_ids is not None:
+            sessions = pending.row_sessions or {}
+            if not any(sessions.get(row_idx) in session_ids for row_idx in pending.row_idxs):
+                return
+        self._minicpmo45_duplex_pending_samples = None
+        if pending.event is not None:
+            pending.event.synchronize()
+        values = pending.host.tolist()
+        saved = (
+            getattr(self, "_minicpmo45_duplex_row_sessions", None),
+            getattr(self, "_minicpmo45_duplex_row_payloads", None),
         )
+        self._minicpmo45_duplex_row_sessions = pending.row_sessions
+        self._minicpmo45_duplex_row_payloads = pending.row_payloads
+        try:
+            for pos, row_idx in enumerate(pending.row_idxs):
+                _final, raw, hit = values[pos]
+                if pending.stage2[pos] and hit:
+                    pending.rewinds.get(pos, lambda: None)()
+                elif pending.stage2[pos]:
+                    self._record_minicpmo45_duplex_generation_token(row_idx, int(raw))
+            for pos, row_idx in enumerate(pending.row_idxs):
+                self._record_minicpmo45_duplex_terminator(row_idx, int(values[pos][0]), pending.token_ids)
+        finally:
+            self._minicpmo45_duplex_row_sessions, self._minicpmo45_duplex_row_payloads = saved
 
     def _maybe_cut_minicpmo45_native_duplex_text_chunk(
         self,
@@ -1227,7 +1599,8 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
             kth = torch.topk(logits, top_k, dim=-1).values[..., -1, None]
             logits = logits.masked_fill(logits < kth, float("-inf"))
         if 0.0 < top_p < 1.0:
-            sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+            # Stable: tied logits enter the nucleus lowest index first.
+            sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1, stable=True)
             cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
             sorted_remove = cumulative_probs > top_p
             sorted_remove[..., 1:] = sorted_remove[..., :-1].clone()
@@ -1237,6 +1610,80 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
             logits = logits.masked_fill(remove, float("-inf"))
         return logits
 
+    @staticmethod
+    def _duplex_top_k_top_p_candidates(
+        logits: torch.Tensor, *, top_k: int, top_p: float
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """``softmax(_top_k_top_p_filter(logits))`` on its support as ``(probs, indices, kth, ties)``, no host read.
+
+        Nucleus order (value descending, index ascending), O(V + k log k). All logits tied at the k-th value
+        stay; their count is data-dependent, so they share the last slot (index -1): the ``ties`` lowest-index
+        ones the nucleus keeps, at one probability each.
+        """
+        vocab = logits.shape[-1]
+        k = int(top_k) if top_k else 0
+        if 0 < k < vocab:
+            values, indices = torch.topk(logits, k, dim=-1)
+            kth = values[..., -1:]
+            values = values.masked_fill(values <= kth, float("-inf"))
+            perm = torch.argsort(indices, dim=-1)
+            values, order = torch.sort(values.gather(-1, perm), descending=True, dim=-1, stable=True)
+            indices = indices.gather(-1, perm).gather(-1, order)
+            tied = (logits == kth).sum(dim=-1, keepdim=True)
+        else:
+            values, indices = torch.sort(logits, descending=True, dim=-1, stable=True)
+            kth = torch.full_like(values[..., :1], float("-inf"))
+            tied = torch.zeros_like(indices[..., :1])
+        # In float32, as softmax does: exponentials past the largest logit, then one division by their sum.
+        values, kth_value = values.float(), kth.float()
+        top = torch.maximum(values[..., :1], kth_value)
+        weights, tie_weight = (values - top).exp(), torch.where(tied > 0, (kth_value - top).exp(), 0.0)
+        total = weights.sum(dim=-1, keepdim=True) + tied * tie_weight
+        probs, tie_prob = weights / total, tie_weight / total
+        ties = tied
+        if 0.0 < float(top_p) < 1.0:
+            cumulative = probs.cumsum(dim=-1)
+            remove = cumulative > float(top_p)
+            remove[..., 1:] = remove[..., :-1].clone()
+            remove[..., 0] = False
+            probs = probs.masked_fill(remove, 0.0)
+            # Tied token j stays while the mass before it, ``before + j * tie_prob``, is at most top_p (the
+            # first candidate always does); the floor's estimate of the last j is corrected on that mass.
+            before = cumulative[..., -1:]
+            last = ((float(top_p) - before) / tie_prob.clamp(min=torch.finfo(torch.float32).tiny)).floor()
+            last = torch.where(before + (last + 1) * tie_prob <= float(top_p), last + 1, last)
+            last = torch.where(before + last * tie_prob > float(top_p), last - 1, last)
+            kept = torch.where(tie_prob > 0, last + 1, torch.zeros_like(last))
+            ties = torch.minimum(kept.clamp(min=0), tied.float()).long()
+        probs = torch.cat([probs, ties * tie_prob], dim=-1)
+        probs = probs / probs.sum(dim=-1, keepdim=True)
+        return probs, torch.cat([indices, torch.full_like(indices[..., :1], -1)], dim=-1), kth, ties
+
+    @staticmethod
+    def _duplex_draw_candidates(
+        logits: torch.Tensor,
+        candidates: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
+        uniforms: torch.Tensor,
+    ) -> torch.Tensor:
+        """Per row, the inverse-CDF draw of ``_duplex_top_k_top_p_candidates(logits)`` at a uniform in [0, 1)."""
+        probs, indices, kth, ties = candidates
+        cdf = probs.cumsum(dim=-1)
+        target = uniforms.view(-1, 1).to(cdf.dtype) * cdf[..., -1:]
+        slot = torch.searchsorted(cdf, target, right=True).clamp(max=probs.shape[-1] - 1)
+        token = indices.gather(-1, slot)
+        # In the tie slot, the j-th lowest-index tied token, j uniform over the kept ties.
+        block = probs[..., -1:]
+        share = (target - (cdf[..., -1:] - block)) / block.clamp(min=torch.finfo(cdf.dtype).tiny) * ties
+        j = torch.minimum(share.floor().clamp(min=0), (ties - 1).clamp(min=0).to(share.dtype)).long()
+        tied = logits == kth
+        tie_token = (tied & (tied.cumsum(dim=-1) == j + 1)).int().argmax(dim=-1, keepdim=True)
+        return torch.where(slot == probs.shape[-1] - 1, tie_token, token).squeeze(-1)
+
+    @staticmethod
+    def _duplex_boundary_chunk_eos_probs(logits: torch.Tensor, chunk_eos_id: int) -> torch.Tensor:
+        """Per-row ``softmax(logits)[:, chunk_eos_id]`` as ``exp(l_eos - logsumexp(l))``."""
+        return torch.exp(logits[:, chunk_eos_id] - torch.logsumexp(logits, dim=-1))
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights for all components of the omni model."""
         loaded_weights = set()
@@ -1266,4 +1713,19 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
             talker_loaded = add_prefix_to_loaded_weights(talker_loaded, "talker")
             loaded_weights.update(talker_loaded)
 
+        self._capture_duplex_audio_graphs()
         return loaded_weights
+
+    def _capture_duplex_audio_graphs(self) -> None:
+        """Capture the duplex streaming audio encoder graphs while the model loads.
+
+        Inside load_model their memory counts as model memory, so the KV budget leaves room for it.
+        A lazy capture on the first duplex request allocates past gpu_memory_utilization instead,
+        which OOMs the stages sharing the GPU.
+        """
+        model_config = getattr(self.vllm_config, "model_config", None)
+        if self.thinker is None or getattr(model_config, "session_mode", "turn") != "duplex":
+            return
+        if self._module_device(self.thinker).type != "cuda":
+            return
+        self._duplex_data_plane_helper().build_audio_cuda_graph()

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py

@@ -18,6 +18,7 @@
 # limitations under the License.
 
 import inspect
+import itertools
 import math
 import os
 import warnings
@@ -138,6 +139,13 @@ def _encode_tokens(tokenizer: Any, prompt: str) -> list[int]:
 
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
+from vllm_omni.model_executor.models.minicpmo_4_5 import streaming_audio_encoder as streaming_audio
+from vllm_omni.model_executor.models.minicpmo_4_5 import streaming_audio_encoder_graph as streaming_audio_graph
+from vllm_omni.model_executor.models.minicpmo_4_5.vision_fused import (
+    VisionGraphEncoder,
+    encode_packed_fused,
+    supports_fused_layers,
+)
 from vllm_omni.model_executor.models.model_local_kv import (
     ModelLocalKVScope,
     ModelLocalKVSpec,
@@ -952,6 +960,8 @@ class SiglipVisionEmbeddings(nn.Module):
         self.num_patches = self.num_patches_per_side**2
         self.num_positions = self.num_patches
         self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+        # (h, w, device) -> position ids of that patch grid, see layout_position_ids.
+        self._layout_position_ids: dict[tuple[int, int, str], torch.Tensor] = {}
 
     def _create_grid_position_ids(
         self,
@@ -1018,6 +1028,19 @@ class SiglipVisionEmbeddings(nn.Module):
 
         return position_ids.to(device=device)
 
+    def layout_position_ids(self, patch_grid_height: int, patch_grid_width: int, device: torch.device) -> torch.Tensor:
+        """Position ids of one ``(h, w)`` patch grid on ``device``: ``_create_position_ids``'s, cached."""
+        cache = self._layout_position_ids
+        key = (int(patch_grid_height), int(patch_grid_width), str(device))
+        ids = cache.get(key)
+        if ids is None:
+            if len(cache) >= 1024:
+                cache.clear()
+            boundaries = torch.arange(1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side)
+            ids = self._create_grid_position_ids(key[0], key[1], boundaries).to(device=device)
+            cache[key] = ids
+        return ids
+
     def forward(
         self,
         pixel_values: torch.FloatTensor,
@@ -1112,6 +1135,51 @@ class SiglipAttention(nn.Module):
 
         return attn_output, attn_weights
 
+    def _attend_unpadded(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+        """Attention over ``(n, L, heads, head_dim)`` items with no padding; same layout out."""
+        query, key, value = (x.transpose(1, 2) for x in (query, key, value))
+        attn_weights = torch.matmul(query, key.transpose(2, 3)) * self.scale
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+        return torch.matmul(attn_weights, value).transpose(1, 2)
+
+    def forward_packed(self, hidden_states: torch.Tensor, seq_groups: Sequence[tuple[int, int, int]]) -> torch.Tensor:
+        """Self-attention over packed ``(tokens, embed)`` items, one mask-free call per ``seq_groups`` run.
+
+        ``seq_groups`` holds contiguous ``(start_token, num_items, seq_len)`` runs of unpadded items.
+        """
+        states = [projection(hidden_states) for projection in (self.q_proj, self.k_proj, self.v_proj)]
+        attn_output = torch.empty_like(states[0])
+        for start, num_items, seq_len in seq_groups:
+            end, shape = start + num_items * seq_len, (num_items, seq_len, self.num_heads, self.head_dim)
+            attn_output[start:end].view(shape).copy_(self._attend_unpadded(*(x[start:end].view(shape) for x in states)))
+        return self.out_proj(attn_output)
+
+
+class SiglipSdpaAttention(SiglipAttention):
+    """SigLIP attention through ``F.scaled_dot_product_attention`` (adapted from vllm-omni#5385)."""
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        output_attentions: bool | None = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if output_attentions:
+            return super().forward(hidden_states, attention_mask=attention_mask, output_attentions=output_attentions)
+        batch_size, q_len, _ = hidden_states.size()
+        shape = (batch_size, q_len, self.num_heads, self.head_dim)
+        states = [projection(hidden_states).view(shape) for projection in (self.q_proj, self.k_proj, self.v_proj)]
+        attn_output = self._attend_unpadded(*states, attn_mask=attention_mask)
+        return self.out_proj(attn_output.reshape(batch_size, q_len, self.embed_dim)), None
+
+    def _attend_unpadded(
+        self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        query, key, value = (x.transpose(1, 2) for x in (query, key, value))
+        return F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attn_mask, dropout_p=self.dropout if self.training else 0.0, scale=self.scale
+        ).transpose(1, 2)
+
 
 class SiglipFlashAttention2(SiglipAttention):
     """
@@ -1329,6 +1397,10 @@ class SiglipFlashAttention2(SiglipAttention):
             (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
         )
 
+    def _attend_unpadded(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+        # Flash attention consumes the (n, L, heads, head_dim) layout directly.
+        return self._flash_attention_forward(query, key, value, None, query.shape[1])
+
 
 # Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->Siglip
 class SiglipMLP(nn.Module):
@@ -1352,7 +1424,8 @@ class SiglipEncoderLayer(nn.Module):
         super().__init__()
         self.embed_dim = config.hidden_size
         self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
-        self.self_attn = SiglipAttention(config) if not self._use_flash_attention_2 else SiglipFlashAttention2(config)
+        attention_classes = {"flash_attention_2": SiglipFlashAttention2, "sdpa": SiglipSdpaAttention}
+        self.self_attn = attention_classes.get(config._attn_implementation, SiglipAttention)(config)
         self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
         self.mlp = SiglipMLP(config)
         self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
@@ -1404,6 +1477,15 @@ class SiglipEncoderLayer(nn.Module):
 
         return outputs
 
+    def forward_packed(self, hidden_states: torch.Tensor, seq_groups: Sequence[tuple[int, int, int]]) -> torch.Tensor:
+        """``forward`` over packed ``(tokens, embed)`` items; see ``SiglipAttention.forward_packed``."""
+        residual = hidden_states
+        hidden_states = self.self_attn.forward_packed(self.layer_norm1(hidden_states), seq_groups)
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.mlp(self.layer_norm2(hidden_states))
+        return residual + hidden_states
+
 
 class SiglipPreTrainedModel(PreTrainedModel):
     """
@@ -1570,8 +1652,14 @@ class SiglipEncoder(nn.Module):
 class SiglipVisionTransformer(SiglipPreTrainedModel):
     config_class = SiglipVisionConfig
     main_input_name = "pixel_values"
+    # Transformers 5 spells the flash flag ``_supports_flash_attn``; SDPA is the
+    # default vision kernel (see MiniCPMO45OmniLLMForConditionalGeneration).
+    _supports_flash_attn = True
     _supports_flash_attn_2 = True
+    _supports_sdpa = True
     _no_split_modules = []
+    # ``forward_packed`` through ``vision_fused.encode_packed_fused``; set by the model.
+    fused_layers = False
 
     def __init__(self, config: SiglipVisionConfig):
         super().__init__(config)
@@ -1687,6 +1775,35 @@ class SiglipVisionTransformer(SiglipPreTrainedModel):
             attentions=encoder_outputs.attentions,
         )
 
+    def forward_packed(self, pixel_values: torch.Tensor, layouts: Sequence[tuple[int, int, int]]) -> torch.Tensor:
+        """The padded ``forward`` per item, without padding, over items packed along the patch axis.
+
+        ``pixel_values`` is ``(1, C, patch, tokens * patch)``, every item's strip in order (the stride-``patch``
+        embedding cuts at item boundaries); ``layouts`` lists ``(h, w, count)`` runs. Returns ``(tokens, embed)``.
+        """
+        embeddings = self.embeddings
+        device = embeddings.position_embedding.weight.device
+        position_ids = [
+            embeddings.layout_position_ids(height, width, device).repeat(count) for height, width, count in layouts
+        ]
+        position_ids = position_ids[0] if len(position_ids) == 1 else torch.cat(position_ids)
+        patch_embeds = embeddings.patch_embedding(pixel_values).flatten(2).transpose(1, 2)[0]
+        hidden_states = patch_embeds + embeddings.position_embedding(position_ids)
+
+        seq_groups: list[tuple[int, int, int]] = []
+        start = 0
+        # Adjacent equal-length grids share one attention call.
+        for seq_len, runs in itertools.groupby(layouts, key=lambda layout: layout[0] * layout[1]):
+            count = sum(run[2] for run in runs)
+            seq_groups.append((start, count, seq_len))
+            start += count * seq_len
+
+        if self.fused_layers and hidden_states.is_cuda and supports_fused_layers(self):
+            return encode_packed_fused(self, hidden_states, seq_groups)
+        for encoder_layer in self.encoder.layers:
+            hidden_states = encoder_layer.forward_packed(hidden_states, seq_groups)
+        return self.post_layernorm(hidden_states)
+
 
 # ============== Resampler Classes ==============
 # coding=utf-8
@@ -1873,6 +1990,154 @@ class Resampler(nn.Module):
     def _repeat(self, query, N: int):
         return query.unsqueeze(1).repeat(1, N, 1)
 
+    def supports_packed(self) -> bool:
+        attn = self.attn
+        plain = not attn.batch_first and attn._qkv_same_embed_dim and attn.in_proj_weight is not None
+        return plain and attn.bias_k is None and attn.bias_v is None and not attn.add_zero_attn
+
+    def forward_packed(self, x: torch.Tensor, layouts: Sequence[tuple[int, int, int]]) -> torch.Tensor:
+        """``forward`` without key padding for packed ``(tokens, kv_dim)`` features; returns ``(B, Q, D)``.
+
+        Each ``(h, w, count)`` run of equal grids attends in one mask-free call (host ints: no device sync).
+        """
+        max_h = max(height for height, _, _ in layouts)
+        max_w = max(width for _, width, _ in layouts)
+        if max_h > self.max_size[0] or max_w > self.max_size[1]:
+            self.max_size = [max(max_h, self.max_size[0]), max(max_w, self.max_size[1])]
+            self._set_2d_pos_cache(self.max_size, x.device)
+        if self.pos_embed.device != x.device:
+            self.pos_embed = self.pos_embed.to(x.device)
+
+        x = self.ln_kv(self.kv_proj(x))
+        key_input = torch.empty_like(x)
+        start = 0
+        for height, width, count in layouts:
+            end = start + count * height * width
+            pos_embed = self.pos_embed[:height, :width, :].reshape(height * width, -1).to(x.dtype)
+            shape = (count, height * width, -1)
+            torch.add(x[start:end].view(shape), pos_embed, out=key_input[start:end].view(shape))
+            start = end
+
+        attn = self.attn
+        num_heads, head_dim = attn.num_heads, self.embed_dim // attn.num_heads
+        w_q, w_k, w_v = attn.in_proj_weight.chunk(3)
+        b_q, b_k, b_v = attn.in_proj_bias.chunk(3) if attn.in_proj_bias is not None else (None, None, None)
+        query = F.linear(self.ln_q(self.query), w_q, b_q)
+        keys = F.linear(key_input, w_k, b_k)
+        values = F.linear(x, w_v, b_v)
+        num_queries = query.shape[0]
+        query_heads = query.view(num_queries, num_heads, head_dim).transpose(0, 1)
+        eager = getattr(self, "attn_implementation", "sdpa") == "eager"
+
+        outputs: list[torch.Tensor] = []
+        start = 0
+        for height, width, count in layouts:
+            seq_len = height * width
+            end = start + count * seq_len
+            shape = (count, seq_len, num_heads, head_dim)
+            key_heads, value_heads = (t[start:end].view(shape).transpose(1, 2) for t in (keys, values))
+            q = query_heads.unsqueeze(0).expand(count, -1, -1, -1).contiguous()
+            if eager:
+                # Mirrors nn.MultiheadAttention's need_weights path used by ``forward``.
+                weights = torch.matmul(q / math.sqrt(head_dim), key_heads.transpose(-2, -1))
+                out = torch.matmul(F.softmax(weights, dim=-1), value_heads)
+            else:
+                out = F.scaled_dot_product_attention(q, key_heads, value_heads)
+            outputs.append(out.transpose(1, 2).reshape(count, num_queries, self.embed_dim))
+            start = end
+
+        out = outputs[0] if len(outputs) == 1 else torch.cat(outputs)
+        out = self.ln_post(attn.out_proj(out))
+        return out @ self.proj
+
+
+def _vision_encode_paths(config: Any, *, encoder_graphs: bool) -> tuple[bool, bool]:
+    """``(vision_packed_encode, vision_cuda_graph)`` from the HF config.
+
+    The packed path (_encode_vision_packed, vision_fused.py) serves the opt-in
+    ``vision_fused_layers`` and ``vision_cuda_graph``; otherwise the padded
+    batch runs and replays the SigLIP graph of ``vpm._encoder_graph``.
+    ``encoder_graphs`` (``encoder_cuda_graph`` and not ``--enforce-eager``)
+    gates the packed graphs too. ``vision_packed_encode`` overrides the path,
+    e.g. to measure packed eager encoding.
+    """
+    graph = bool(getattr(config, "vision_cuda_graph", False)) and encoder_graphs
+    packed = getattr(config, "vision_packed_encode", None)
+    if packed is None:
+        return bool(getattr(config, "vision_fused_layers", False)) or graph, graph
+    return bool(packed), graph and bool(packed)
+
+
+def _packed_vision_layouts(
+    pixel_values: Sequence[torch.Tensor],
+    tgt_sizes: torch.Tensor | Sequence[Sequence[int]],
+    patch_size: int,
+) -> list[tuple[int, int]] | None:
+    """Host ``(h, w)`` patch grid per slice, or ``None`` on a mismatch; the packed path's only sync."""
+    sizes = tgt_sizes.tolist() if isinstance(tgt_sizes, torch.Tensor) else [list(size) for size in tgt_sizes]
+    if len(sizes) != len(pixel_values) or not sizes:
+        return None
+    layouts: list[tuple[int, int]] = []
+    for pixels, size in zip(pixel_values, sizes):
+        if len(size) != 2:
+            return None
+        height, width = int(size[0]), int(size[1])
+        if (
+            min(height, width) <= 0
+            or pixels.ndim != 3
+            or pixels.shape[-2:] != (patch_size, height * width * patch_size)
+        ):
+            return None
+        layouts.append((height, width))
+    return layouts
+
+
+def _plan_packed_vision_chunks(
+    layouts: Sequence[tuple[int, int]], max_items: int
+) -> list[tuple[list[int], list[tuple[int, int, int]]]]:
+    """Per ``max_items``-slice chunk, its slice indices and ``(h, w, count)`` runs, equal grids adjacent.
+
+    Sorting by ``h * w`` first also puts equal sequence lengths together, sharing an attention call.
+    """
+    order = sorted(range(len(layouts)), key=lambda i: (layouts[i][0] * layouts[i][1], layouts[i]))
+    chunks: list[tuple[list[int], list[tuple[int, int, int]]]] = []
+    for begin in range(0, len(order), max(1, max_items)):
+        indices = order[begin : begin + max(1, max_items)]
+        grids = itertools.groupby(layouts[index] for index in indices)
+        chunks.append((indices, [(height, width, len(list(run))) for (height, width), run in grids]))
+    return chunks
+
+
+def _encode_vision_packed(
+    vpm: "SiglipVisionTransformer",
+    resampler: Resampler,
+    pixel_values: Sequence[torch.Tensor],
+    layouts: Sequence[tuple[int, int]],
+    max_items: int,
+    graph_encoder: VisionGraphEncoder | None = None,
+) -> torch.Tensor:
+    """SigLIP + resampler over unpadded slices packed by grid; ``(num_slices, query_num, embed)`` in input order.
+
+    No slice is padded, so its output does not depend on the others; single-grid chunks replay graphs.
+    """
+    weight = vpm.embeddings.patch_embedding.weight
+    outputs: list[torch.Tensor] = []
+    packed_order: list[int] = []
+    for indices, runs in _plan_packed_vision_chunks(layouts, max_items):
+        parts = [pixel_values[index] for index in indices]
+        pixels = parts[0] if len(parts) == 1 else torch.cat(parts, dim=-1)
+        pixels = pixels.to(device=weight.device, dtype=weight.dtype).unsqueeze(0)
+        out = graph_encoder.encode(pixels, *runs[0]) if graph_encoder is not None and len(runs) == 1 else None
+        if out is None:
+            out = resampler.forward_packed(vpm.forward_packed(pixels, runs), runs)
+        outputs.append(out)
+        packed_order.extend(indices)
+    out = outputs[0] if len(outputs) == 1 else torch.cat(outputs)
+    if packed_order != list(range(len(packed_order))):
+        inverse = torch.tensor(packed_order, dtype=torch.long).argsort()
+        out = out.index_select(0, inverse.to(device=out.device, non_blocking=True))
+    return out
+
 
 class MultiheadAttention(nn.MultiheadAttention):
     def __init__(
@@ -3944,7 +4209,10 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
         multimodal_config = vllm_config.model_config.multimodal_config
 
         self.config = config
-        self.multimodal_config = multimodal_config
+        # Duplex session cap of this deployment; sizes the audio-encoder graph
+        # grid when ``duplex_audio_encoder_cuda_graph_batch_sizes_from_sessions``.
+        self._duplex_max_sessions = int(getattr(vllm_config.model_config, "duplex_max_sessions", 1) or 1)
+        self._enforce_eager = bool(vllm_config.model_config.enforce_eager)
         # Model-local opt-out for A/B measurements; --enforce-eager always wins.
         encoder_graphs = (
             bool(getattr(config, "encoder_cuda_graph", True)) and not vllm_config.model_config.enforce_eager
@@ -3970,10 +4238,15 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
         # Initialize vision encoder (SigLIP)
         if multimodal_config.get_limit_per_prompt("image"):
             # Set attention implementation
-            if hasattr(vllm_config, "model_config") and hasattr(vllm_config.model_config, "_attn_implementation"):
+            # (an engine override, else the vision config, else SDPA: eager math without the (B, H, L, L) scores)
+            if getattr(getattr(vllm_config, "model_config", None), "_attn_implementation", None) is not None:
                 config.vision_config._attn_implementation = vllm_config.model_config._attn_implementation
             else:
-                config.vision_config._attn_implementation = "eager"
+                vision_implementation = getattr(config.vision_config, "_attn_implementation", None)
+                config.vision_config._attn_implementation = vision_implementation or "sdpa"
+            if config.vision_config._attn_implementation == "flash_attention_2" and not is_flash_attn_2_available():
+                logger.warning("flash_attention_2 requested for the MiniCPM-o vision tower but unavailable; using sdpa")
+                config.vision_config._attn_implementation = "sdpa"
 
             self.vpm = SiglipVisionTransformer(config.vision_config)
             if encoder_graphs:
@@ -4029,8 +4302,13 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
                 kv_dim=vision_dim,  # Vision encoder dimension
                 adaptive=True,  # Enable adaptive resampling
             )
+            self.resampler.attn_implementation = config.vision_config._attn_implementation
         else:
             self.resampler = None
+        if self.vpm is not None:
+            self.vpm.fused_layers = bool(getattr(config, "vision_fused_layers", False))
+        self.vision_packed_encode, self.vision_cuda_graph = _vision_encode_paths(config, encoder_graphs=encoder_graphs)
+        self._vision_graph_encoder: VisionGraphEncoder | None = None
 
         # Initialize audio encoder (APM) and audio projection
         if getattr(config, "init_audio", True) and hasattr(config, "audio_config") and config.audio_config is not None:
@@ -4309,6 +4587,19 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
         pixel_values = data["pixel_values"]
         tgt_sizes = data["tgt_sizes"]
 
+        vpm = self.vpm
+        packed = getattr(self, "vision_packed_encode", False) and hasattr(vpm, "forward_packed")
+        if packed and hasattr(self.resampler, "forward_packed") and self.resampler.supports_packed():
+            layouts = _packed_vision_layouts(pixel_values, tgt_sizes, vpm.embeddings.patch_size)
+            if layouts is not None:
+                graph_encoder = None
+                if getattr(self, "vision_cuda_graph", False) and pixel_values[0].is_cuda:
+                    if self._vision_graph_encoder is None:
+                        self._vision_graph_encoder = VisionGraphEncoder(vpm, self.resampler)
+                    graph_encoder = self._vision_graph_encoder
+                max_items = max(1, int(self.config.vision_batch_size))
+                return _encode_vision_packed(vpm, self.resampler, pixel_values, layouts, max_items, graph_encoder)
+
         B = len(pixel_values)
         P = pixel_values[0].shape[-2]
         L = max(item.shape[-1] for item in pixel_values)
@@ -4488,8 +4779,8 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
             audio_embeds = self._encode_audio_features(wavforms, audio_attention_mask)
 
         _, feature_lens_after_pooling = self._get_feat_extract_output_lengths(audio_feature_lens)
-
-        num_audio_tokens = feature_lens_after_pooling
+        # One host read of every length instead of an implicit .item() per slice bound.
+        num_audio_tokens = feature_lens_after_pooling.tolist()
 
         final_audio_embeds = list[torch.Tensor]()
         idx = 0
@@ -4590,18 +4881,23 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
             device=wavforms.device,
         )
 
+        # As in get_audio_embedding: every layer's output only to pick the last one keeps 25 activations alive.
+        selects_final_layer = self.audio_encoder_layer == -1
         audio_outputs = self.apm(
             wavforms,
             past_key_values=self.audio_past_key_values,
             use_cache=True,
-            output_hidden_states=True,
+            output_hidden_states=not selects_final_layer,
             attention_mask=audio_attention_mask,
             use_extra_context=use_extra_context,
             prefix_extra_frames=prefix_extra_frames,
             suffix_extra_frames=suffix_extra_frames,
             cnn_min_length=cnn_min_length,
         )
-        audio_states = audio_outputs.hidden_states[self.audio_encoder_layer]
+        if selects_final_layer:
+            audio_states = audio_outputs.last_hidden_state
+        else:
+            audio_states = audio_outputs.hidden_states[self.audio_encoder_layer]
         self.audio_past_key_values = audio_outputs.past_key_values
 
         audio_embeds = self.audio_projection_layer(audio_states)
@@ -4610,16 +4906,97 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
         audio_embeds = audio_embeds.transpose(1, 2)
 
         _, feature_lens_after_pooling = self._get_feat_extract_output_lengths(audio_feature_lens)
+        # One host read of the lengths instead of a device sync per slice bound.
+        pooled_lens = feature_lens_after_pooling.tolist()
         final_audio_embeds: list[list[torch.Tensor]] = []
         idx = 0
         for i in range(len(audio_feature_lens_raw)):
             target_audio_embeds = []
             for _ in range(len(audio_feature_lens_raw[i])):
-                target_audio_embeds.append(audio_embeds[idx, : feature_lens_after_pooling[idx], :])
+                target_audio_embeds.append(audio_embeds[idx, : pooled_lens[idx], :])
                 idx += 1
             final_audio_embeds.append(target_audio_embeds)
         return final_audio_embeds
 
+    def supports_streaming_audio_batch(self) -> bool:
+        """Whether ``get_audio_embedding_streaming_batch`` can serve this encoder."""
+        if self.apm is None or self.audio_projection_layer is None or self.audio_avg_pooler is None:
+            return False
+        requested = getattr(self.config, "duplex_audio_encoder_batching", None)
+        if requested is not None and not bool(requested):
+            return False
+        reason = streaming_audio.streaming_batch_unsupported_reason(
+            self.apm, audio_encoder_layer=self.audio_encoder_layer
+        )
+        if reason is not None:
+            logger.info("MiniCPM-o streaming audio encoder stays per-session: %s", reason)
+            return False
+        return requested is not None or self.apm.conv1.weight.device.type in ("cuda", "cpu")
+
+    def build_streaming_audio_graph_encoder(self, *, unit_frames: int) -> bool:
+        """Capture CUDA graphs of the streaming encoder's steady unit. Failures stay eager."""
+        self._duplex_audio_cuda_graph_encoder = None
+        config = self.config
+        enabled = bool(getattr(config, "duplex_audio_encoder_cuda_graph", True)) and not self._enforce_eager
+        if not self.supports_streaming_audio_batch() or not enabled:
+            return False
+        if self.apm.conv1.weight.device.type != "cuda":
+            logger.info("MiniCPM-o streaming audio encoder CUDA graph needs a CUDA device; staying eager")
+            return False
+        batch_sizes = getattr(config, "duplex_audio_encoder_cuda_graph_batch_sizes", None)
+        if getattr(config, "duplex_audio_encoder_cuda_graph_batch_sizes_from_sessions", False) is True:
+            # Overrides the explicit list: the grid follows duplex max_sessions.
+            batch_sizes = streaming_audio_graph.batch_sizes_for_sessions(getattr(self, "_duplex_max_sessions", 1))
+        allocated = torch.accelerator.memory_allocated(self.apm.conv1.weight.device)
+        try:
+            wrapper = streaming_audio_graph.StreamingAudioGraphEncoder(
+                self.apm,
+                self.audio_projection_layer,
+                self.audio_avg_pooler,
+                unit_frames=int(unit_frames),
+                pool_step=int(config.audio_pool_step),
+                batch_sizes=batch_sizes,
+                cache_buckets=getattr(config, "duplex_audio_encoder_cuda_graph_cache_buckets", None),
+                page_positions=self._duplex_audio_kv_page_positions(),
+                pinned_h2d=getattr(config, "duplex_audio_encoder_pinned_h2d", False) is True,
+            )
+            wrapper.capture()
+        except Exception:
+            logger.exception("MiniCPM-o streaming audio encoder CUDA graph capture failed; staying eager")
+            return False
+        self._duplex_audio_cuda_graph_encoder = wrapper
+        logger.info(
+            "Captured MiniCPM-o streaming audio encoder CUDA graphs: unit_frames=%d batch_sizes=%s "
+            "cache_buckets=%s pinned_h2d=%s, took %.2f GiB",
+            unit_frames,
+            wrapper.batch_sizes,
+            wrapper.cache_buckets,
+            wrapper.pinned_h2d,
+            (torch.accelerator.memory_allocated(self.apm.conv1.weight.device) - allocated) / (1 << 30),
+        )
+        return True
+
+    def get_audio_embedding_streaming_batch(
+        self,
+        chunks: Sequence[Any],
+    ) -> tuple[list[torch.Tensor | None], list[Any]]:
+        """Batched ``get_audio_embedding_streaming`` (embeddings, updated caches); graphs take steady rows."""
+        graph = getattr(self, "_duplex_audio_cuda_graph_encoder", None)
+        if graph is not None:
+            return graph.encode(chunks)
+        return streaming_audio.encode_streaming_audio_batch(
+            self.apm,
+            self.audio_projection_layer,
+            self.audio_avg_pooler,
+            chunks,
+            pool_step=int(self.config.audio_pool_step),
+            page_positions=self._duplex_audio_kv_page_positions(),
+        )
+
+    def _duplex_audio_kv_page_positions(self) -> int:
+        page_positions = getattr(self.config, "duplex_audio_kv_page_positions", None)
+        return int(page_positions or streaming_audio.DEFAULT_KV_PAGE_POSITIONS)
+
     def _process_audio_input(
         self,
         audio_input: MiniCPMOAudioInputs,

vllm_omni/model_executor/models/minicpmo_4_5/streaming_audio_encoder.py

@@ -0,0 +1,312 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Batched streaming Whisper encoder for MiniCPM-o duplex sessions."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+
+import torch
+from torch import nn
+
+#: Growth step of a session buffer, in encoder positions (a 1 s unit is 50):
+#: at most five reallocations per 1500-position cycle, ~12 MiB average slack.
+DEFAULT_KV_PAGE_POSITIONS = 250
+
+_SUPPORTED_ATTENTION = ("sdpa", "eager")
+_SUPPORTED_DTYPES = (torch.bfloat16, torch.float32)
+
+
+@dataclass(eq=False, slots=True)
+class StreamingAudioKVCache:
+    """One session's streaming-encoder self-attention cache, written in place.
+
+    Layout ``[layers, 2 (key/value), capacity, embed_dim]``: token-major, so appending a unit is one contiguous
+    copy per layer and attention reads a strided ``[1, heads, length, head_dim]`` view without materializing it.
+    Capacity grows in ``page_positions`` steps up to ``max_positions``.
+    """
+
+    num_layers: int
+    embed_dim: int
+    num_heads: int
+    max_positions: int
+    page_positions: int = DEFAULT_KV_PAGE_POSITIONS
+    length: int = field(default=0, init=False)
+    _buffer: torch.Tensor | None = field(default=None, init=False, repr=False)
+
+    def __post_init__(self) -> None:
+        if min(self.num_layers, self.embed_dim, self.num_heads, self.max_positions, self.page_positions) <= 0:
+            raise ValueError("streaming audio KV cache dimensions must be positive")
+
+    @classmethod
+    def for_encoder(cls, encoder: nn.Module, page_positions: int) -> StreamingAudioKVCache:
+        config, positions = encoder.config, int(encoder.embed_positions.weight.shape[0])
+        return cls(
+            len(encoder.layers), int(config.d_model), int(config.encoder_attention_heads), positions, page_positions
+        )
+
+    @property
+    def capacity(self) -> int:
+        return 0 if self._buffer is None else int(self._buffer.shape[2])
+
+    def reserve(self, positions: int, *, dtype: torch.dtype, device: torch.device) -> None:
+        """Make room for ``positions`` entries, keeping the committed prefix."""
+        if positions <= self.capacity:
+            return
+        pages = math.ceil(positions / self.page_positions) * self.page_positions
+        capacity = max(positions, min(self.max_positions, pages))
+        buffer = torch.empty((self.num_layers, 2, capacity, self.embed_dim), dtype=dtype, device=device)
+        if self._buffer is not None and self.length:
+            buffer[:, :, : self.length].copy_(self._buffer[:, :, : self.length])
+        self._buffer = buffer
+
+    def history(self, length: int) -> torch.Tensor:
+        """``[layers, 2, length, embed_dim]`` rows ``[0, length)``, every layer at once."""
+        assert self._buffer is not None
+        return self._buffer[:, :, :length]
+
+    def commit(self, offset: int, new_kv: torch.Tensor) -> None:
+        """Write ``[layers, 2, positions, embed_dim]`` rows at ``offset``; the caller advances ``length``."""
+        assert self._buffer is not None
+        positions = int(new_kv.shape[2])
+        self._buffer[:, :, offset : offset + positions].copy_(new_kv)
+
+    def layer_views(self, start: int, end: int) -> tuple[tuple[torch.Tensor, ...], ...]:
+        """Per layer, ``(key_rows, value_rows, key_heads, value_heads)``: the ``[end - start, embed_dim]``
+        write targets and the ``[1, heads, end, head_dim]`` attention inputs over ``[0, end)``."""
+        assert self._buffer is not None
+        head_dim = self.embed_dim // self.num_heads
+        heads = (
+            self._buffer[:, :, :end]
+            .view(self.num_layers, 2, end, self.num_heads, head_dim)
+            .transpose(2, 3)
+            .unsqueeze(2)
+        )
+        rows = self._buffer[:, :, start:end]
+        return rows[:, 0].unbind(0), rows[:, 1].unbind(0), heads[:, 0].unbind(0), heads[:, 1].unbind(0)
+
+    def to_legacy_cache(self):
+        """The committed history as the ``EncoderDecoderCache`` the per-session path grows."""
+        from transformers.cache_utils import DynamicCache, EncoderDecoderCache
+
+        self_attention = DynamicCache()
+        if self.length:
+            _, _, keys, values = self.layer_views(0, self.length)
+            for layer in range(self.num_layers):
+                self_attention.update(keys[layer].contiguous(), values[layer].contiguous(), layer)
+        return EncoderDecoderCache(self_attention, DynamicCache())
+
+
+@dataclass(frozen=True, slots=True)
+class StreamingAudioChunk:
+    """One session's ready unit: log-mel frames plus the state it continues."""
+
+    #: ``[n_mels, frames]`` or ``[1, n_mels, frames]`` log-mel features, host or device.
+    features: torch.Tensor
+    #: Session cache; ``None`` starts a new history.
+    cache: StreamingAudioKVCache | None
+    prefix_extra_frames: int = 0
+    suffix_extra_frames: int = 0
+    use_extra_context: bool = True
+    #: ``audio_feature_lens`` of the unit when it differs from ``features.shape[-1]``.
+    feature_length: int | None = None
+
+
+@dataclass(slots=True)
+class _Row:
+    index: int
+    frames: int
+    start: int
+    length: int
+    pooled_length: int
+    cache: StreamingAudioKVCache
+    past: int = 0
+    offset: int = 0
+
+    @property
+    def total(self) -> int:
+        return self.past + self.length
+
+
+def _extra_context_trim(frames: int) -> int:
+    return (int(frames) + 1) // 2 if frames > 0 else 0
+
+
+def streaming_batch_unsupported_reason(encoder: nn.Module, *, audio_encoder_layer: int | None) -> str | None:
+    """Why ``encoder`` cannot take the batched path, or ``None`` when it can."""
+    if audio_encoder_layer != -1:
+        return "only the final encoder layer output is batched"
+    if encoder.training:
+        return "training mode (dropout / layerdrop)"
+    weight = getattr(getattr(encoder, "conv1", None), "weight", None)
+    if weight is None or weight.dtype not in _SUPPORTED_DTYPES:
+        # float16 needs the per-session inf/nan clamp, which is per tensor.
+        return f"unsupported encoder dtype {getattr(weight, 'dtype', None)}"
+    layers = getattr(encoder, "layers", None)
+    if not layers:
+        return "encoder has no layers"
+    for layer in layers:
+        attention = getattr(layer, "self_attn", None)
+        if attention is None or not all(
+            hasattr(attention, name) for name in ("q_proj", "k_proj", "v_proj", "out_proj", "scaling")
+        ):
+            return "encoder layers are not Whisper self-attention layers"
+    implementation = getattr(encoder.config, "_attn_implementation", None) or "eager"
+    if implementation not in _SUPPORTED_ATTENTION:
+        return f"attention implementation {implementation!r} is not batched"
+    return None
+
+
+def _attend(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor, implementation: str):
+    """Whisper self-attention as the per-session path calls it; ``query`` already carries ``scaling``."""
+    if implementation == "sdpa":
+        return torch.nn.functional.scaled_dot_product_attention(
+            query, key, value, attn_mask=mask, dropout_p=0.0, scale=1.0
+        )
+    weights = torch.matmul(query, key.transpose(2, 3)) + mask
+    return torch.matmul(torch.nn.functional.softmax(weights, dim=-1), value)
+
+
+def _stack_features(rows: list[_Row], features: list[torch.Tensor], device: torch.device) -> torch.Tensor:
+    """Zero-padded ``[B, mels, max_frames]`` on ``device`` with one host-to-device copy."""
+    max_frames = max(row.frames for row in rows)
+    on_host = all(feature.device.type == "cpu" for feature in features)
+    pin = on_host and device.type == "cuda"
+    batch = torch.zeros(
+        (len(rows), int(features[0].shape[0]), max_frames),
+        dtype=features[0].dtype,
+        device="cpu" if on_host else device,
+        pin_memory=pin,
+    )
+    for slot, (row, feature) in enumerate(zip(rows, features, strict=True)):
+        batch[slot, :, : row.frames].copy_(feature)
+    # Pinned staging keeps the host from waiting on queued GPU work; the
+    # caching host allocator holds the buffer until the copy completes.
+    return batch.to(device, non_blocking=pin)
+
+
+def _position_rows(encoder: nn.Module, rows: list[_Row]) -> torch.Tensor:
+    """Learned positions ``past .. past + n`` per row; past the table, repeat its last row."""
+    table = encoder.embed_positions.weight
+    capacity = int(table.shape[0])
+    pieces: list[torch.Tensor] = []
+    for row in rows:
+        end = row.past + row.length
+        if end <= capacity:
+            pieces.append(table[row.past : end])
+        else:
+            pieces.append(table[row.past :])
+            pieces.append(table[-1:].expand(end - max(row.past, capacity), -1))
+    return pieces[0] if len(pieces) == 1 else torch.cat(pieces)
+
+
+@torch.inference_mode()
+def encode_streaming_audio_batch(
+    encoder: nn.Module,
+    projection: nn.Module,
+    pooler: nn.Module,
+    chunks: Sequence[StreamingAudioChunk],
+    *,
+    pool_step: int,
+    page_positions: int = DEFAULT_KV_PAGE_POSITIONS,
+) -> tuple[list[torch.Tensor | None], list[StreamingAudioKVCache | None]]:
+    """Encode one streaming unit per session in one encoder pass."""
+    implementation = getattr(encoder.config, "_attn_implementation", None) or "eager"
+    weight = encoder.conv1.weight
+    dtype, device = weight.dtype, weight.device
+    max_positions = int(encoder.embed_positions.weight.shape[0])
+    embed_dim, num_heads = int(encoder.config.d_model), int(encoder.config.encoder_attention_heads)
+
+    outputs: list[torch.Tensor | None] = [None] * len(chunks)
+    caches: list[StreamingAudioKVCache | None] = [chunk.cache for chunk in chunks]
+    rows: list[_Row] = []
+    features: list[torch.Tensor] = []
+    for index, chunk in enumerate(chunks):
+        feature = chunk.features
+        if feature.ndim == 3:
+            if feature.shape[0] != 1:
+                raise ValueError("a streaming chunk carries exactly one audio")
+            feature = feature[0]
+        frames = int(feature.shape[-1])
+        conv_length = (frames - 1) // 2 + 1
+        prefix = _extra_context_trim(chunk.prefix_extra_frames) if chunk.use_extra_context else 0
+        suffix = _extra_context_trim(chunk.suffix_extra_frames) if chunk.use_extra_context else 0
+        length = conv_length - prefix - suffix
+        if frames <= 0 or length <= 0:
+            continue
+        feature_length = frames if chunk.feature_length is None else int(chunk.feature_length)
+        pooled_length = ((feature_length - 1) // 2 + 1 - pool_step) // pool_step + 1
+        cache = chunk.cache
+        if cache is None or cache.length + length >= max_positions:
+            cache = StreamingAudioKVCache.for_encoder(encoder, page_positions)
+        rows.append(_Row(index, frames, prefix, length, pooled_length, cache, past=cache.length))
+        features.append(feature)
+    if not rows:
+        return outputs, caches
+
+    offset = 0
+    for row in rows:
+        row.offset = offset
+        offset += row.length
+        row.cache.reserve(row.total, dtype=dtype, device=device)
+
+    # CNN front end on the zero-padded batch, transformer layers on the packed rows.
+    mel = _stack_features(rows, features, device).to(dtype=dtype)
+    hidden = nn.functional.gelu(encoder.conv1(mel))
+    max_frames = int(mel.shape[-1])
+    for slot, row in enumerate(rows):
+        if row.frames < max_frames:
+            # conv2 must see zeros past each row's end, as its own padding would in the per-session call.
+            hidden[slot, :, row.frames :].zero_()
+    hidden = nn.functional.gelu(encoder.conv2(hidden)).permute(0, 2, 1)
+    pieces = [hidden[slot, row.start : row.start + row.length] for slot, row in enumerate(rows)]
+    hidden = pieces[0] if len(pieces) == 1 else torch.cat(pieces)
+    hidden = hidden + _position_rows(encoder, rows)
+
+    # Every per-row view is built once here, not per layer: at batch sizes
+    # that matter the host cost of view ops, not the GPU, bounds this loop.
+    views = [row.cache.layer_views(row.past, row.total) for row in rows]
+    lengths = [row.length for row in rows]
+    total_rows = sum(lengths)
+    head_dim = embed_dim // num_heads
+    # The per-session path passes an all-zero additive mask; keep it for the same attention kernel.
+    masks = [torch.zeros((1, 1, row.length, row.total), dtype=dtype, device=device) for row in rows]
+    for layer_index, layer in enumerate(encoder.layers):
+        attention = layer.self_attn
+        residual = hidden
+        normed = layer.self_attn_layer_norm(hidden)
+        query = attention.q_proj(normed) * attention.scaling
+        key, value = attention.k_proj(normed), attention.v_proj(normed)
+        destinations = [view[0][layer_index] for view in views] + [view[1][layer_index] for view in views]
+        torch._foreach_copy_(destinations, [*key.split(lengths), *value.split(lengths)])
+        queries = query.view(1, total_rows, num_heads, head_dim).transpose(1, 2).split(lengths, dim=2)
+        attended = [
+            _attend(row_query, view[2][layer_index], view[3][layer_index], mask, implementation)
+            for row_query, view, mask in zip(queries, views, masks, strict=True)
+        ]
+        # [1, heads, rows, head_dim] per row -> packed [rows, embed_dim].
+        attended_rows = (attended[0] if len(attended) == 1 else torch.cat(attended, dim=2)).transpose(1, 2)
+        hidden = residual + attention.out_proj(attended_rows.reshape(total_rows, embed_dim))
+        residual = hidden
+        hidden = layer.activation_fn(layer.fc1(layer.final_layer_norm(hidden)))
+        hidden = residual + layer.fc2(hidden)
+    hidden = encoder.layer_norm(hidden)
+    embeds = projection(hidden)
+    if all(row.length % pool_step == 0 for row in rows):
+        pooled = pooler(embeds.transpose(0, 1).unsqueeze(0)).squeeze(0).transpose(0, 1)
+        for row in rows:
+            start = row.offset // pool_step
+            count = min(row.length // pool_step, row.pooled_length)
+            outputs[row.index] = pooled[start : start + count]
+    else:
+        for row in rows:
+            segment = embeds[row.offset : row.offset + row.length].transpose(0, 1).unsqueeze(0)
+            pooled = pooler(segment).squeeze(0).transpose(0, 1)
+            outputs[row.index] = pooled[: row.pooled_length]
+
+    for row in rows:
+        row.cache.length = row.total
+        caches[row.index] = row.cache
+    return outputs, caches

vllm_omni/model_executor/models/minicpmo_4_5/streaming_audio_encoder_graph.py

@@ -0,0 +1,259 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CUDA graphs of MiniCPM-o's streaming Whisper encoder for the steady unit shape."""
+
+from __future__ import annotations
+
+import functools
+import itertools
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import torch
+from torch import nn
+
+from .streaming_audio_encoder import (
+    DEFAULT_KV_PAGE_POSITIONS,
+    StreamingAudioChunk,
+    StreamingAudioKVCache,
+    _attend,
+    _extra_context_trim,
+    encode_streaming_audio_batch,
+)
+
+if TYPE_CHECKING:
+    from torch.cuda import CUDAGraph
+
+_NEG_INF = float("-inf")
+
+# stage0.py builds every unit but a session's first with these trims.
+STEADY_PREFIX_EXTRA_FRAMES = 2
+STEADY_SUFFIX_EXTRA_FRAMES = 2
+
+DEFAULT_GRAPH_BATCH_SIZES: tuple[int, ...] = (1, 2, 4, 8, 16, 32)
+# Shared storage scales with the largest bucket only. 900 stays under ~3 GiB
+# at batch 32 on MiniCPM-o 4.5 (bf16); history past it falls back to eager.
+DEFAULT_GRAPH_CACHE_BUCKETS: tuple[int, ...] = (250, 500, 750, 900)
+
+
+def batch_sizes_for_sessions(max_sessions: int) -> tuple[int, ...]:
+    """Powers of two below ``max_sessions``, plus ``max_sessions`` itself."""
+    top = max(1, int(max_sessions))
+    sizes = [1 << i for i in range(top.bit_length()) if (1 << i) < top]
+    return normalize_buckets([*sizes, top])
+
+
+def select_bucket(value: int, buckets: Sequence[int]) -> int | None:
+    return min((bucket for bucket in buckets if value <= bucket), default=None)
+
+
+def normalize_buckets(values: Sequence[int]) -> tuple[int, ...]:
+    return tuple(sorted({int(v) for v in values if int(v) > 0}))
+
+
+def steady_unit_length(unit_frames: int) -> int:
+    conv_length = (int(unit_frames) - 1) // 2 + 1
+    trim = _extra_context_trim(STEADY_PREFIX_EXTRA_FRAMES) + _extra_context_trim(STEADY_SUFFIX_EXTRA_FRAMES)
+    return conv_length - trim
+
+
+def row_is_steady(chunk: StreamingAudioChunk, *, expected_frames: int) -> bool:
+    frames = int(chunk.features.shape[-1])
+    return (
+        chunk.cache is not None
+        and chunk.use_extra_context
+        and (int(chunk.prefix_extra_frames), int(chunk.suffix_extra_frames))
+        == (STEADY_PREFIX_EXTRA_FRAMES, STEADY_SUFFIX_EXTRA_FRAMES)
+        and frames == expected_frames
+        and (chunk.feature_length is None or int(chunk.feature_length) == frames)
+    )
+
+
+@dataclass(slots=True)
+class _Group:
+    index: int
+    cache_in: StreamingAudioKVCache
+    cache_out: StreamingAudioKVCache
+    past: int
+
+
+class StreamingAudioGraphEncoder:
+    """Pre-captured CUDA graphs of the streaming encoder's steady unit shape."""
+
+    def __init__(
+        self,
+        encoder: nn.Module,
+        projection: nn.Module,
+        pooler: nn.Module,
+        *,
+        unit_frames: int,
+        pool_step: int,
+        batch_sizes: Sequence[int] | None = None,
+        cache_buckets: Sequence[int] | None = None,
+        page_positions: int = DEFAULT_KV_PAGE_POSITIONS,
+        pinned_h2d: bool = False,
+    ) -> None:
+        weight, config = encoder.conv1.weight, encoder.config
+        if weight.device.type != "cuda":
+            raise ValueError("StreamingAudioGraphEncoder requires a CUDA encoder")
+        self.encoder, self.projection, self.pooler = encoder, projection, pooler
+        self.device, self.dtype = weight.device, weight.dtype
+        self.pool_step, self.page_positions, self.pinned_h2d = int(pool_step), int(page_positions), bool(pinned_h2d)
+        self.embed_dim, self.num_heads = int(config.d_model), int(config.encoder_attention_heads)
+        self.max_positions = int(encoder.embed_positions.weight.shape[0])
+        self.implementation = getattr(config, "_attn_implementation", None) or "eager"
+        self.unit_frames = int(unit_frames)
+        self.unit_length = steady_unit_length(self.unit_frames)
+        nominal_pooled = ((self.unit_frames - 1) // 2 + 1 - self.pool_step) // self.pool_step + 1
+        self.pooled_length = min(nominal_pooled, self.unit_length // self.pool_step)
+        if self.unit_length <= 0 or self.pooled_length <= 0:
+            raise ValueError(f"unit_frames={unit_frames} pool_step={pool_step} leaves no steady unit output")
+        self.batch_sizes = normalize_buckets(batch_sizes or DEFAULT_GRAPH_BATCH_SIZES)
+        self.cache_buckets = normalize_buckets(cache_buckets or DEFAULT_GRAPH_CACHE_BUCKETS)
+        if not self.batch_sizes or not self.cache_buckets:
+            raise ValueError("StreamingAudioGraphEncoder needs at least one batch size and one cache bucket")
+        self.max_batch, self.max_cache_bucket = max(self.batch_sizes), max(self.cache_buckets)
+        # (batch, cache length) -> (graph, pooled output); every graph reads prefix views of one storage.
+        self._graphs: dict[tuple[int, int], tuple[CUDAGraph, torch.Tensor]] = {}
+
+    def _views(self, batch: int, cache_len: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        total = cache_len + self.unit_length
+        return (
+            self._mel[:batch],
+            self._cache[:, :, :batch, :total, :],
+            self._mask[:batch, :, :, :total],
+            self._position_ids[:batch],
+        )
+
+    def capture(self) -> None:
+        if torch.cuda.is_current_stream_capturing():
+            raise RuntimeError("cannot capture the streaming audio encoder graph mid-capture")
+        from vllm.platforms import current_platform
+
+        batch, total = self.max_batch, self.max_cache_bucket + self.unit_length
+        zeros = functools.partial(torch.zeros, dtype=self.dtype, device=self.device)
+        self._mel = zeros((batch, int(self.encoder.config.num_mel_bins), self.unit_frames))
+        self._cache = zeros((len(self.encoder.layers), 2, batch, total, self.embed_dim))
+        self._mask = zeros((batch, 1, self.unit_length, total))
+        self._position_ids = torch.zeros((batch, self.unit_length), dtype=torch.long, device=self.device)
+        pool = current_platform.get_global_graph_pool()
+        current_stream = torch.cuda.current_stream(self.device)
+        for key in itertools.product(self.batch_sizes, self.cache_buckets):
+            views = self._views(*key)
+            warmup_stream = torch.cuda.Stream(device=self.device)
+            warmup_stream.wait_stream(current_stream)
+            with torch.cuda.stream(warmup_stream), torch.inference_mode():
+                for _ in range(3):
+                    self._forward(*views, key[1])
+            current_stream.wait_stream(warmup_stream)
+            graph = torch.cuda.CUDAGraph()
+            with torch.inference_mode(), torch.cuda.graph(graph, pool=pool):
+                self._graphs[key] = (graph, self._forward(*views, key[1]))
+
+    def _forward(
+        self,
+        mel: torch.Tensor,
+        cache: torch.Tensor,
+        mask: torch.Tensor,
+        position_ids: torch.Tensor,
+        cache_len: int,
+    ) -> torch.Tensor:
+        encoder = self.encoder
+        batch = mel.shape[0]
+        hidden = nn.functional.gelu(encoder.conv1(mel))
+        hidden = nn.functional.gelu(encoder.conv2(hidden)).permute(0, 2, 1)
+        prefix = _extra_context_trim(STEADY_PREFIX_EXTRA_FRAMES)
+        hidden = hidden[:, prefix : prefix + self.unit_length, :]
+        hidden = hidden + encoder.embed_positions(position_ids)
+        head_dim = self.embed_dim // self.num_heads
+        for layer_index, layer in enumerate(encoder.layers):
+            attention = layer.self_attn
+            residual = hidden
+            normed = layer.self_attn_layer_norm(hidden)
+            query = attention.q_proj(normed) * attention.scaling
+            key, value = attention.k_proj(normed), attention.v_proj(normed)
+            cache[layer_index, 0, :, cache_len : cache_len + self.unit_length, :].copy_(key)
+            cache[layer_index, 1, :, cache_len : cache_len + self.unit_length, :].copy_(value)
+            q = query.view(batch, self.unit_length, self.num_heads, head_dim).transpose(1, 2)
+            k_all = cache[layer_index, 0].view(batch, -1, self.num_heads, head_dim).transpose(1, 2)
+            v_all = cache[layer_index, 1].view(batch, -1, self.num_heads, head_dim).transpose(1, 2)
+            attended = _attend(q, k_all, v_all, mask, self.implementation)
+            attended = attended.transpose(1, 2).reshape(batch, self.unit_length, self.embed_dim)
+            hidden = residual + attention.out_proj(attended)
+            residual = hidden
+            hidden = layer.activation_fn(layer.fc1(layer.final_layer_norm(hidden)))
+            hidden = residual + layer.fc2(hidden)
+        hidden = encoder.layer_norm(hidden)
+        embeds = self.projection(hidden)
+        return self.pooler(embeds.transpose(1, 2)).transpose(1, 2)
+
+    def encode(
+        self,
+        chunks: Sequence[StreamingAudioChunk],
+    ) -> tuple[list[torch.Tensor | None], list[StreamingAudioKVCache | None]]:
+        outputs: list[torch.Tensor | None] = [None] * len(chunks)
+        caches: list[StreamingAudioKVCache | None] = [chunk.cache for chunk in chunks]
+        steady: list[_Group] = []
+        eager: list[int] = []
+        for index, chunk in enumerate(chunks):
+            cache_in = chunk.cache
+            reset = cache_in is not None and cache_in.length + self.unit_length >= self.max_positions
+            past = 0 if reset or cache_in is None else cache_in.length
+            if not row_is_steady(chunk, expected_frames=self.unit_frames) or past > self.max_cache_bucket:
+                eager.append(index)
+                continue
+            cache_out = StreamingAudioKVCache.for_encoder(self.encoder, self.page_positions) if reset else cache_in
+            steady.append(_Group(index, cache_in, cache_out, past))
+        for start in range(0, len(steady), self.max_batch):
+            self._replay(steady[start : start + self.max_batch], chunks, outputs, caches)
+        if eager:
+            eager_outputs, eager_caches = encode_streaming_audio_batch(
+                self.encoder,
+                self.projection,
+                self.pooler,
+                [chunks[i] for i in eager],
+                pool_step=self.pool_step,
+                page_positions=self.page_positions,
+            )
+            for index, output, cache in zip(eager, eager_outputs, eager_caches, strict=True):
+                outputs[index], caches[index] = output, cache
+        return outputs, caches
+
+    @torch.inference_mode()
+    def _replay(
+        self,
+        group: list[_Group],
+        chunks: Sequence[StreamingAudioChunk],
+        outputs: list[torch.Tensor | None],
+        caches: list[StreamingAudioKVCache | None],
+    ) -> None:
+        batch = select_bucket(len(group), self.batch_sizes)
+        cache_len = select_bucket(max((row.past for row in group), default=0), self.cache_buckets)
+        assert batch is not None and cache_len is not None
+        mel, cache, mask, position_ids = self._views(batch, cache_len)
+        mask.fill_(0.0)
+        if len(group) < batch:
+            mask[len(group) :, :, :, :cache_len].fill_(_NEG_INF)
+            position_ids[len(group) :].copy_(torch.arange(self.unit_length, device=self.device))
+        for slot, row in enumerate(group):
+            feature = chunks[row.index].features
+            feature = (feature[0] if feature.ndim == 3 else feature).to(dtype=self.dtype)
+            if self.pinned_h2d and feature.device.type == "cpu":
+                mel[slot].copy_(feature.pin_memory(), non_blocking=True)
+            else:
+                mel[slot].copy_(feature.to(device=self.device))
+            pad = cache_len - row.past
+            if pad:
+                mask[slot, :, :, :pad].fill_(_NEG_INF)
+            if row.past:
+                cache[:, :, slot, pad:cache_len, :].copy_(row.cache_in.history(row.past))
+            position_ids[slot].copy_(torch.arange(row.past, row.past + self.unit_length, device=self.device))
+        graph, pooled = self._graphs[(batch, cache_len)]
+        graph.replay()
+        for slot, row in enumerate(group):
+            outputs[row.index] = pooled[slot, : self.pooled_length].clone()
+            row.cache_out.reserve(row.past + self.unit_length, dtype=self.dtype, device=self.device)
+            row.cache_out.commit(row.past, cache[:, :, slot, cache_len : cache_len + self.unit_length, :])
+            row.cache_out.length = row.past + self.unit_length
+            caches[row.index] = row.cache_out

vllm_omni/model_executor/models/minicpmo_4_5/vision_fused.py

@@ -0,0 +1,273 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Fused packed SigLIP layers and CUDA graphs for MiniCPM-o 4.5 vision encode."""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from collections.abc import Sequence
+from copy import copy
+from math import prod
+from typing import Any, ClassVar, Literal
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from vllm.config import ModelConfig, VllmConfig
+from vllm.config.multimodal import MultiModalConfig
+from vllm.logger import init_logger
+from vllm.model_executor.models.interfaces import SupportsEncoderCudaGraph
+from vllm.platforms import current_platform
+from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager
+from vllm.v1.worker.encoder_cudagraph_defs import (
+    EncoderCudaGraphCaptureInputs,
+    EncoderCudaGraphConfig,
+    EncoderCudaGraphReplayBuffers,
+    EncoderItemSpec,
+)
+
+from .ops import residual_layer_norm
+
+logger = init_logger(__name__)
+
+# Batch buckets of the vision graphs; the checkpoint's vision_batch_size (16)
+# bounds a chunk, so it is the largest bucket.
+DEFAULT_GRAPH_BATCH_SIZES = (1, 2, 4, 8, 16)
+# Distinct (grid, bucket) graphs kept at once; duplex uses one grid.
+_MAX_GRAPHS = 16
+# A (grid, bucket) is captured on its third encode, so one-off shapes (the
+# startup profile run, a single turn-mode image) stay eager.
+_CAPTURE_AFTER = 2
+
+
+def supports_fused_layers(vpm: nn.Module) -> bool:
+    """Whether ``encode_packed_fused`` covers ``vpm`` (SigLIP with a GELU MLP)."""
+    layers = getattr(getattr(vpm, "encoder", None), "layers", None)
+    if not layers or getattr(vpm, "post_layernorm", None) is None:
+        return False
+    if getattr(getattr(vpm, "config", None), "_attn_implementation", None) == "eager":
+        return False
+    attn, mlp = getattr(layers[0], "self_attn", None), getattr(layers[0], "mlp", None)
+    return all(
+        getattr(attn, name, None) is not None for name in ("q_proj", "k_proj", "v_proj", "out_proj")
+    ) and getattr(getattr(mlp, "config", None), "hidden_act", None) in ("gelu_pytorch_tanh", "gelu")
+
+
+def packed_qkv(attn: nn.Module) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """``(3E, E)`` q/k/v weight (and bias) of ``attn``, stacked once."""
+    cached = attn.__dict__.get("_packed_qkv")
+    projections = (attn.q_proj, attn.k_proj, attn.v_proj)
+    if cached is not None and all(
+        p.weight.data_ptr() == cached[0][i * p.weight.shape[0]].data_ptr() for i, p in enumerate(projections)
+    ):
+        return cached
+    with torch.no_grad():
+        weight = torch.cat([p.weight for p in projections])
+        biases = [p.bias for p in projections]
+        bias = torch.cat(biases) if all(b is not None for b in biases) else None
+        rows = int(projections[0].weight.shape[0])
+        for index, projection in enumerate(projections):
+            projection.weight.data = weight[index * rows : (index + 1) * rows]
+            if bias is not None:
+                projection.bias.data = bias[index * rows : (index + 1) * rows]
+    attn.__dict__["_packed_qkv"] = (weight, bias)
+    return weight, bias
+
+
+def _attend(attn: nn.Module, qkv: torch.Tensor, seq_groups: Sequence[tuple[int, int, int]]) -> torch.Tensor:
+    """``(tokens, E)`` attention output of packed ``(tokens, 3E)`` q/k/v, each item attending to itself."""
+    heads, head_dim = attn.num_heads, attn.head_dim
+    embed = heads * head_dim
+    outputs = []
+    for start, count, seq_len in seq_groups:
+        run = qkv[start : start + count * seq_len].view(count, seq_len, 3, heads, head_dim)
+        query, key, value = (run[:, :, index].transpose(1, 2) for index in range(3))
+        out = F.scaled_dot_product_attention(query, key, value, scale=attn.scale)
+        outputs.append(out.transpose(1, 2).reshape(count * seq_len, embed))
+    return outputs[0] if len(outputs) == 1 else torch.cat(outputs)
+
+
+def _mlp_up(mlp: nn.Module, hidden: torch.Tensor) -> torch.Tensor:
+    fc1 = mlp.fc1
+    if mlp.config.hidden_act == "gelu_pytorch_tanh" and fc1.bias is not None and hidden.is_cuda:
+        # cuBLASLt's GELU epilogue is the tanh approximation.
+        return torch._addmm_activation(fc1.bias, hidden, fc1.weight.t(), use_gelu=True)
+    return mlp.activation_fn(fc1(hidden))
+
+
+def _add_norm(residual: torch.Tensor, norm: nn.Module, y=None, y_bias=None, residual_out=None) -> torch.Tensor:
+    """``norm(residual + y + y_bias)`` for a ``(tokens, E)`` ``y``, the sum also written to ``residual_out``."""
+    y = None if y is None else y.unsqueeze(0)
+    return residual_layer_norm(
+        residual, y, y_bias=y_bias, weight=norm.weight, bias=norm.bias, eps=norm.eps, residual_out=residual_out
+    )
+
+
+def encode_packed_fused(
+    vpm: nn.Module, hidden_states: torch.Tensor, seq_groups: Sequence[tuple[int, int, int]]
+) -> torch.Tensor:
+    """The encoder layers and ``post_layernorm`` of ``vpm.forward_packed`` over ``(tokens, E)`` embeddings.
+
+    The residual stream is updated in place in a contiguous copy of
+    ``hidden_states`` (the patch embedding hands over a transposed view).
+    """
+    layers = vpm.encoder.layers
+    residual = hidden_states.contiguous().unsqueeze(0)
+    normed = _add_norm(residual, layers[0].layer_norm1)
+    for index, layer in enumerate(layers):
+        attn, mlp = layer.self_attn, layer.mlp
+        qkv_weight, qkv_bias = packed_qkv(attn)
+        attended = _attend(attn, F.linear(normed[0], qkv_weight, qkv_bias), seq_groups)
+        y = F.linear(attended, attn.out_proj.weight)
+        normed = _add_norm(residual, layer.layer_norm2, y, attn.out_proj.bias, residual_out=residual)
+        last = index + 1 == len(layers)
+        norm = vpm.post_layernorm if last else layers[index + 1].layer_norm1
+        y = F.linear(_mlp_up(mlp, normed[0]), mlp.fc2.weight)
+        normed = _add_norm(residual, norm, y, mlp.fc2.bias, residual_out=None if last else residual)
+    return normed[0]
+
+
+def _copy_pixels_padded(dst: torch.Tensor, src: torch.Tensor) -> None:
+    width_px = int(src.shape[-1])
+    dst[..., :width_px].copy_(src)
+    if width_px < int(dst.shape[-1]):
+        dst[..., width_px:].zero_()
+
+
+class _PackedVisionAdapter(SupportsEncoderCudaGraph):
+    """Adapter exposing the packed SigLIP tower + resampler to vLLM's EncoderCudaGraphManager."""
+
+    supports_encoder_cudagraph: ClassVar[Literal[True]] = True
+
+    def __init__(self, forward: Any, runs: Any, pixel_shape: Any, output_shape: torch.Size):
+        self.forward, self.runs, self.pixel_shape, self.output_shape = forward, runs, pixel_shape, output_shape
+        self.tokens = prod(output_shape[:-1]) or 1
+
+    def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
+        return EncoderCudaGraphConfig(
+            modalities=["image"],
+            buffer_keys=["pixels"],
+            out_hidden_size=self.output_shape[-1],
+            padding_logics={"pixels": _copy_pixels_padded},
+        )
+
+    def get_encoder_cudagraph_budget_range(self, vllm_config: VllmConfig) -> tuple[int, int]:
+        return self.tokens, self.tokens
+
+    def get_encoder_cudagraph_item_specs(self, mm_kwargs: dict[str, Any]) -> list[EncoderItemSpec]:
+        return [EncoderItemSpec(input_size=self.tokens, output_tokens=self.tokens)]
+
+    def select_encoder_cudagraph_items(self, mm_kwargs: dict[str, Any], indices: list[int]) -> dict[str, Any]:
+        return dict(mm_kwargs)
+
+    def prepare_encoder_cudagraph_capture_inputs(self, *args: Any, **kwargs: Any) -> EncoderCudaGraphCaptureInputs:
+        device = kwargs.get("device") or args[3]
+        dtype = kwargs.get("dtype") or args[4]
+        return EncoderCudaGraphCaptureInputs({"pixels": torch.zeros(self.pixel_shape, dtype=dtype, device=device)})
+
+    def prepare_encoder_cudagraph_replay_buffers(
+        self, mm_kwargs: dict[str, Any], *args: Any, **kwargs: Any
+    ) -> EncoderCudaGraphReplayBuffers:
+        return EncoderCudaGraphReplayBuffers(mm_kwargs)
+
+    def encoder_cudagraph_forward(self, inputs: dict[str, torch.Tensor], path: str = "default") -> torch.Tensor:
+        return self.forward(inputs["pixels"], self.runs)
+
+    def encoder_eager_forward(self, mm_kwargs: dict[str, Any], path: str = "default") -> torch.Tensor:
+        return self.encoder_cudagraph_forward(mm_kwargs, path)
+
+    def postprocess_encoder_output(
+        self,
+        outputs: dict[str, torch.Tensor],
+        indices: list[int],
+        per_item_out_tokens: list[int],
+        dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
+        clone: bool = False,
+        batch_mm_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        dest[0] = outputs["default"].clone()
+
+
+class VisionGraphEncoder:
+    """CUDA graphs of the packed SigLIP tower + resampler for single-grid chunks."""
+
+    def __init__(
+        self,
+        vpm: nn.Module,
+        resampler: nn.Module,
+        *,
+        batch_sizes: Sequence[int] = DEFAULT_GRAPH_BATCH_SIZES,
+        vllm_config: VllmConfig | None = None,
+    ):
+        self.vpm, self.resampler, self.vllm_config = vpm, resampler, vllm_config
+        self.batch_sizes = tuple(sorted({int(b) for b in batch_sizes if int(b) > 0}))
+        # (h, w, bucket) -> EncoderCudaGraphManager, least recently used first.
+        self._managers: OrderedDict[tuple[int, int, int], EncoderCudaGraphManager] = OrderedDict()
+        self._failed: set[tuple[int, int, int]] = set()
+        self._seen: dict[tuple[int, int, int], int] = {}
+
+    def encode(self, pixels: torch.Tensor, height: int, width: int, count: int) -> torch.Tensor | None:
+        if torch.cuda.is_current_stream_capturing():
+            return None
+        bucket = min((size for size in self.batch_sizes if count <= size), default=None)
+        if bucket is None:
+            return None
+        key = (int(height), int(width), bucket)
+        manager = self._managers.get(key)
+        if manager is None:
+            if key in self._failed:
+                return None
+            seen = self._seen.get(key, 0)
+            if seen < _CAPTURE_AFTER:
+                if len(self._seen) >= 4 * _MAX_GRAPHS:
+                    self._seen.clear()
+                self._seen[key] = seen + 1
+                return None
+            manager = self._capture(key, pixels)
+            if manager is None:
+                return None
+        self._managers.move_to_end(key)
+        return manager.execute({"pixels": pixels})[0][:count]
+
+    def _run(self, pixels: torch.Tensor, runs: list[tuple[int, int, int]]) -> torch.Tensor:
+        return self.resampler.forward_packed(self.vpm.forward_packed(pixels, runs), runs)
+
+    def _capture(self, key: tuple[int, int, int], like: torch.Tensor) -> EncoderCudaGraphManager | None:
+        height, width, bucket = key
+        runs, patch = [(height, width, bucket)], int(like.shape[-2])
+        pixel_shape = (*like.shape[:-1], bucket * height * width * patch)
+        dummy_pixels = torch.zeros(pixel_shape, device=like.device, dtype=like.dtype)
+        try:
+            stream = torch.cuda.Stream(device=like.device)
+            stream.wait_stream(torch.cuda.current_stream(like.device))
+            with torch.cuda.stream(stream), torch.inference_mode():
+                out = self._run(dummy_pixels, runs)
+                for _ in range(1):
+                    self._run(dummy_pixels, runs)
+            torch.cuda.current_stream(like.device).wait_stream(stream)
+            adapter = _PackedVisionAdapter(self._run, runs, pixel_shape, out.shape)
+            config = copy(self.vllm_config) if self.vllm_config is not None else VllmConfig()
+            if getattr(config, "model_config", None) is None:
+                config.model_config = ModelConfig.__new__(ModelConfig)
+            if getattr(config.model_config, "multimodal_config", None) is None:
+                config.model_config.multimodal_config = MultiModalConfig()
+            cc = config.compilation_config = copy(config.compilation_config)
+            cc.encoder_cudagraph_token_budgets = [adapter.tokens]
+            cc.encoder_cudagraph_max_vision_items_per_batch, cc.encoder_cudagraph_max_frames_per_batch = 1, 0
+            config.parallel_config = copy(config.parallel_config)
+            config.parallel_config.tensor_parallel_size = 1
+
+            manager = EncoderCudaGraphManager(config, like.device, like.dtype, adapter)
+            with torch.cuda.stream(stream):
+                manager.capture(graph_pool=current_platform.get_global_graph_pool())
+            torch.cuda.current_stream(like.device).wait_stream(stream)
+        except Exception:
+            logger.warning("MiniCPM-o vision CUDA graph capture failed for %s; staying eager", key, exc_info=True)
+            self._failed.add(key)
+            return None
+        if len(self._managers) >= _MAX_GRAPHS:
+            _, evicted = self._managers.popitem(last=False)
+            evicted.clear()
+        self._managers[key] = manager
+        logger.info("Captured MiniCPM-o vision CUDA graph grid=%dx%d batch=%d via manager", height, width, bucket)
+        return manager