[Perf] Add CUDA graphs for MiniCPM-o 4.5 input encoders (#8332)

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>

docs/design/minicpmo_encoder_cuda_graph.md

@@ -0,0 +1,61 @@
+# MiniCPM-o 4.5 input encoder CUDA graphs
+
+The vision transformer and stateless audio encoder reuse vLLM's encoder graph
+manager for exact input shapes. Each encoder retains at most four graphs by
+default, and captures on the second call with the same shape and CUDA stream.
+Startup profiling can consume slots. Graphs are never evicted: once the cache
+is full, other shapes run eagerly for the lifetime of the model.
+
+Use `--hf-overrides` to configure both encoders:
+
+```json
+{
+  "encoder_cuda_graph": true,
+  "encoder_cuda_graph_max_graphs": 4,
+  "encoder_cuda_graph_min_capture_calls": 2,
+  "encoder_cuda_graph_min_free_bytes": 1073741824,
+  "encoder_cuda_graph_share_pools": true
+}
+```
+
+`max_graphs` is a nonnegative integer per encoder; zero disables capture.
+`min_capture_calls` is an integer of at least two. Increasing it reduces capture
+of short-lived shapes. Admission history is bounded to four times `max_graphs`,
+so an evicted history entry must accumulate its calls again. Increasing the
+graph cap can improve coverage but retains more GPU memory. Existing graphs
+are not replaced, avoiding repeated capture costs when traffic changes.
+
+By default, graphs in the same encoder on the same device and replay stream
+share a memory pool and capture stream. Different encoders and replay streams
+remain isolated. Graph inputs and output buffers are allocated outside capture;
+only transient capture allocations share the pool. Returned outputs remain
+cloned. A late capture waits for earlier replay work and output clones before
+reusing the pool. Set `encoder_cuda_graph_share_pools` to `false` for an A/B
+comparison with a separate pool and capture stream per graph.
+
+The graph count is **not a byte limit**. Sharing pools reuses intermediate
+storage but does not eliminate per-shape static input and output buffers.
+Before a new capture, the adapter checks device-free memory against
+`min_free_bytes`, a nonnegative integer (default: 1 GiB; zero disables the check).
+Below this floor the call runs eagerly, and a later call may try admission
+again. Existing graphs can still replay. This check is conservative about
+reusable allocator memory and does not predict capture allocations, reserve
+memory, or guarantee that capture fits. Size the floor for the deployment;
+disable encoder graphs when their additional memory cannot be accommodated.
+
+Each adapter exposes `get_cumulative_stats()`: graph hits, total misses, hit
+rate, graph count, and separate ineligible, capacity, warmup and memory misses.
+These include calls rejected before the upstream manager, so a full cache's
+lost coverage is visible. Vision uses `vpm._encoder_graph`; audio uses
+`_audio_encoder_graph` on the model. The counters are per adapter, not global.
+
+Capture failures propagate the original exception and make that adapter
+unusable. Subsequent calls raise without retrying capture or running CUDA eager
+work. Restart the worker, optionally with encoder graphs disabled. This is a
+fatal contract, not automatic recovery: a failed CUDA capture can leave the
+device context unusable. `capture_failures` records this terminal state.
+
+`--enforce-eager` always disables these graphs. CPU, grad/autocast, stateful
+streaming audio and the other model-specific eager fallbacks remain unchanged.
+The adapter clones graph outputs through upstream postprocessing so retained
+embeddings survive later replays.

tests/model_executor/models/minicpmo_4_5/test_encoder_cuda_graph.py

@@ -0,0 +1,379 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.minicpmo_4_5.encoder_cuda_graph import EncoderCudaGraph
+from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
+    SiglipVisionConfig,
+    SiglipVisionTransformer,
+)
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model]
+
+
+def _make_graph(forward, **kwargs):
+    from vllm.config import ModelConfig, VllmConfig
+    from vllm.config.multimodal import MultiModalConfig
+
+    config = VllmConfig()
+    # The manager only needs multimodal limits; no checkpoint is loaded here.
+    config.model_config = ModelConfig.__new__(ModelConfig)
+    config.model_config.multimodal_config = MultiModalConfig()
+    return EncoderCudaGraph(forward, config, **kwargs)
+
+
+@pytest.mark.cpu
+def test_cpu_and_grad_paths_remain_eager():
+    graph = _make_graph(lambda x, mask: x.sin() if mask is None else x.sin() + mask)
+    x = torch.randn(2, 3, requires_grad=True)
+    graph(x, None).sum().backward()
+    torch.testing.assert_close(x.grad, x.detach().cos())
+    assert not graph.graphs
+    assert not graph._seen
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@torch.inference_mode()
+def test_replay_refreshes_inputs_preserves_outputs_and_bounds_shapes():
+    graph = _make_graph(lambda x, mask: x.sin() if mask is None else x.sin() + mask, max_graphs=2)
+    x = torch.randn(2, 8, device="cuda")
+    mask = torch.randn_like(x)
+    graph(x, mask)
+    old = graph(x, mask)
+    expected_old = old.clone()
+    x.mul_(2)
+    mask.add_(3)
+    torch.testing.assert_close(graph(x, mask), x.sin() + mask)
+    torch.testing.assert_close(old, expected_old)
+    graph(x, None)
+    torch.testing.assert_close(graph(x, None), x.sin())
+    assert len(graph.graphs) == 2
+    from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager
+
+    assert all(isinstance(entry, EncoderCudaGraphManager) for entry in graph.graphs.values())
+    assert sum(entry.graph_hits for entry in graph.graphs.values()) == 3
+    assert graph.vllm_config.compilation_config.encoder_cudagraph_token_budgets == []
+    for size in range(3, 24):
+        other = torch.randn(size, 8, device="cuda")
+        torch.testing.assert_close(graph(other, None), other.sin())
+    assert len(graph.graphs) == 2
+    assert len(graph._seen) <= 8
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@torch.inference_mode()
+def test_vision_graph_handles_changed_mask_and_retained_embeddings():
+    config = SiglipVisionConfig(
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        image_size=28,
+        patch_size=14,
+    )
+    config._attn_implementation = "eager"
+    model = SiglipVisionTransformer(config).eval().cuda()
+    pixels = torch.randn(2, 3, 14, 56, device="cuda")
+    sizes = torch.tensor([[2, 2], [1, 2]], dtype=torch.int32)
+    mask = torch.tensor([[[1, 1, 1, 1]], [[1, 1, 0, 0]]], device="cuda", dtype=torch.bool)
+    reference = model(pixels, mask, sizes).last_hidden_state
+    model._encoder_graph = _make_graph(model._encode_last_hidden_state)
+    model(pixels, mask, sizes)
+    output = model(pixels, mask, sizes).last_hidden_state
+    torch.testing.assert_close(output, reference)
+    assert len(model._encoder_graph.graphs) == 1
+    # Same padded shape, different valid positions and image content.
+    pixels.add_(1)
+    sizes[1] = torch.tensor([1, 3])
+    mask[1, 0, 2] = True
+    actual = model(pixels, mask, sizes).last_hidden_state
+    graph = model._encoder_graph
+    model._encoder_graph = None
+    expected = model(pixels, mask, sizes).last_hidden_state
+    torch.testing.assert_close(actual, expected)
+    torch.testing.assert_close(output, reference)
+    model._encoder_graph = graph
+    detailed = model(pixels, mask, sizes, output_hidden_states=True)
+    assert len(detailed.hidden_states) == 3
+    assert len(graph.graphs) == 1
+
+
+def _audio_model(device):
+    from types import SimpleNamespace
+
+    from transformers import WhisperConfig
+
+    from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
+        MiniCPMO45OmniLLMForConditionalGeneration,
+        MiniCPMWhisperEncoder,
+        MultiModalProjector,
+    )
+
+    config = WhisperConfig(
+        d_model=32,
+        encoder_layers=2,
+        encoder_attention_heads=4,
+        encoder_ffn_dim=64,
+        num_mel_bins=80,
+        max_source_positions=1500,
+    )
+    config._attn_implementation = "sdpa"
+    model = MiniCPMO45OmniLLMForConditionalGeneration.__new__(MiniCPMO45OmniLLMForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.config = SimpleNamespace(audio_chunk_length=1.0, audio_pool_step=5)
+    model.apm = MiniCPMWhisperEncoder(config)
+    model.audio_projection_layer = MultiModalProjector(32, 48)
+    model.audio_avg_pooler = torch.nn.AvgPool1d(5, stride=5)
+    model.audio_encoder_layer = -1
+    model.audio_past_key_values = None
+    return model.eval().to(device)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@torch.inference_mode()
+def test_audio_graph_refreshes_mask_and_matches_eager():
+    model = _audio_model("cuda")
+    data = {
+        "audio_features": torch.randn(2, 80, 100, device="cuda"),
+        "audio_feature_lens": torch.tensor([[100], [80]], device="cuda"),
+    }
+    expected = model.get_audio_hidden_states(data)
+    graph = _make_graph(model._encode_audio_features)
+    model._audio_encoder_graph = graph
+    model.get_audio_hidden_states(data)
+    actual = model.get_audio_hidden_states(data)
+    for a, b in zip(actual, expected, strict=True):
+        torch.testing.assert_close(a, b)
+    assert len(graph.graphs) == 1
+    data["audio_features"].add_(1)
+    data["audio_feature_lens"][1, 0] = 30
+    changed = model.get_audio_hidden_states(data)
+    model._audio_encoder_graph = None
+    for a, b in zip(changed, model.get_audio_hidden_states(data), strict=True):
+        torch.testing.assert_close(a, b)
+    for a, b in zip(actual, expected, strict=True):
+        torch.testing.assert_close(a, b)
+
+
+@pytest.mark.cpu
+@torch.inference_mode()
+def test_streaming_audio_does_not_replay_stateless_graph():
+    model = _audio_model("cpu")
+
+    def forbidden(*args):
+        raise AssertionError("streaming KV must not use the stateless graph")
+
+    model._audio_encoder_graph = forbidden
+    data = {"audio_features": torch.randn(1, 80, 100), "audio_feature_lens": torch.tensor([[100]])}
+    for _ in range(2):
+        output = model.get_audio_embedding_streaming(data)
+        assert output[0][0].shape == (10, 48)
+    assert model.audio_past_key_values.get_seq_length() == 100
+
+
+@pytest.mark.cpu
+@torch.inference_mode()
+def test_fp16_audio_keeps_host_overflow_check_eager():
+    model = _audio_model("cpu").half()
+
+    def forbidden(*args):
+        raise AssertionError("FP16 overflow guard cannot run inside capture")
+
+    model._audio_encoder_graph = forbidden
+    data = {
+        "audio_features": torch.randn(1, 80, 100, dtype=torch.float16),
+        "audio_feature_lens": torch.tensor([[100]]),
+    }
+    assert model.get_audio_hidden_states(data)[0].shape == (10, 48)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@torch.inference_mode()
+def test_one_off_shapes_do_not_exhaust_capture_admission():
+    graph = _make_graph(torch.sin, max_graphs=1)
+    for size in range(1, 20):
+        graph(torch.zeros(size, device="cuda"))
+    assert len(graph._seen) == 4
+    assert not graph.graphs
+    x = torch.randn(32, device="cuda")
+    graph(x)
+    torch.testing.assert_close(graph(x), x.sin())
+    assert len(graph.graphs) == 1
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@torch.inference_mode()
+def test_managers_do_not_share_buffers_between_streams():
+    graph = _make_graph(torch.sin)
+    streams = [torch.cuda.Stream(), torch.cuda.Stream()]
+    outputs = []
+    for i, stream in enumerate(streams):
+        with torch.cuda.stream(stream):
+            value = torch.full((4, 8), float(i + 1), device="cuda")
+            graph(value)
+            output = graph(value)
+            graph(value + 1)
+            outputs.append((output, value.sin()))
+    for stream in streams:
+        torch.cuda.current_stream().wait_stream(stream)
+    assert len(graph.graphs) == 2
+    for actual, expected in outputs:
+        torch.testing.assert_close(actual, expected)
+    assert not graph.vllm_config.compilation_config.encoder_cudagraph_token_budgets
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "options",
+    [
+        {"max_graphs": -1},
+        {"max_graphs": True},
+        {"min_capture_calls": 1},
+        {"min_capture_calls": 2.5},
+        {"min_free_bytes": -1},
+    ],
+)
+def test_invalid_admission_options(options):
+    with pytest.raises(ValueError, match="must be an integer"):
+        _make_graph(torch.sin, **options)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@torch.inference_mode()
+def test_configurable_admission_and_capacity_miss_accounting():
+    graph = _make_graph(torch.sin, max_graphs=1, min_capture_calls=3)
+    x = torch.randn(2, 8, device="cuda")
+    for _ in range(2):
+        torch.testing.assert_close(graph(x), x.sin())
+        assert not graph.graphs
+    torch.testing.assert_close(graph(x), x.sin())
+    torch.testing.assert_close(graph(x + 1), (x + 1).sin())
+    for _ in range(5):
+        other = torch.randn(3, 8, device="cuda")
+        torch.testing.assert_close(graph(other), other.sin())
+    stats = graph.get_cumulative_stats()
+    assert stats["num_graphs"] == 1
+    assert stats["graph_hits"] == 2
+    assert stats["graph_misses"] == 7
+    assert stats["capacity_misses"] == 5
+    assert stats["warmup_misses"] == 2
+    assert stats["hit_rate"] == pytest.approx(2 / 9)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@torch.inference_mode()
+def test_memory_admission_defers_capture_but_preserves_existing_replay(monkeypatch):
+    graph = _make_graph(torch.sin, min_free_bytes=1024)
+    x = torch.randn(2, 8, device="cuda")
+    graph(x)
+    monkeypatch.setattr(current_omni_platform, "get_free_memory", lambda device: 1023)
+    for _ in range(3):
+        torch.testing.assert_close(graph(x), x.sin())
+    assert not graph.graphs
+    assert graph.get_cumulative_stats()["memory_misses"] == 3
+    monkeypatch.setattr(current_omni_platform, "get_free_memory", lambda device: 1024)
+    old = graph(x)
+    assert len(graph.graphs) == 1
+    monkeypatch.setattr(current_omni_platform, "get_free_memory", lambda device: 0)
+    torch.testing.assert_close(graph(x + 1), (x + 1).sin())
+    torch.testing.assert_close(old, x.sin())
+    assert graph.get_cumulative_stats()["graph_hits"] == 2
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@torch.inference_mode()
+def test_capture_failure_is_fatal_without_retry_or_eager_fallback(monkeypatch):
+    from unittest.mock import Mock
+
+    forward = Mock(side_effect=torch.sin)
+    graph = _make_graph(forward)
+    x = torch.randn(2, 8, device="cuda")
+    graph(x)
+    capture = Mock(side_effect=RuntimeError("injected capture failure"))
+    monkeypatch.setattr(graph, "_capture", capture)
+    with pytest.raises(RuntimeError, match="injected capture failure"):
+        graph(x)
+    for value in (x, torch.randn(3, 8, device="cuda"), torch.randn(2, 8)):
+        with pytest.raises(RuntimeError, match="restart the worker"):
+            graph(value)
+    capture.assert_called_once()
+    forward.assert_called_once()
+    assert not graph.graphs
+    assert not graph._seen
+    assert not graph._seen_calls
+    assert graph.get_cumulative_stats()["capture_failures"] == 1
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@torch.inference_mode()
+def test_zero_capacity_disables_capture():
+    graph = _make_graph(torch.sin, max_graphs=0)
+    x = torch.randn(2, 8, device="cuda")
+    for _ in range(3):
+        torch.testing.assert_close(graph(x), x.sin())
+    assert not graph.graphs
+    assert not graph._seen
+    assert graph.get_cumulative_stats()["capacity_misses"] == 3
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("share_pools", [False, True])
+@torch.inference_mode()
+def test_pool_reuse_preserves_outputs_across_late_capture_and_arbitrary_replay(share_pools):
+    graph = _make_graph(lambda x: (x.sin() + 2).cos(), share_pools=share_pools)
+    retained = []
+    for size in (64, 128, 32, 96):
+        x = torch.randn(size, 64, device="cuda")
+        graph(x)
+        retained.append((graph(x), (x.sin() + 2).cos()))
+        # Replay older shapes between captures, including smaller -> larger.
+        for previous in (64, size, 64):
+            value = torch.randn(previous, 64, device="cuda")
+            torch.testing.assert_close(graph(value), (value.sin() + 2).cos())
+        for actual, expected in retained:
+            torch.testing.assert_close(actual, expected)
+    pools = {entry.graph_pool for entry in graph.graphs.values()}
+    assert len(pools) == (1 if share_pools else 4)
+    assert len(graph._capture_pools) == (1 if share_pools else 0)
+    for size in (96, 32, 128, 64, 128, 32):
+        value = torch.randn(size, 64, device="cuda")
+        torch.testing.assert_close(graph(value), (value.sin() + 2).cos())
+    for actual, expected in retained:
+        torch.testing.assert_close(actual, expected)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@torch.inference_mode()
+def test_shared_pools_are_isolated_by_encoder_and_replay_stream():
+    graphs = [_make_graph(torch.sin), _make_graph(torch.cos)]
+    streams = [torch.cuda.Stream(), torch.cuda.Stream()]
+    retained = []
+    for graph in graphs:
+        for stream in streams:
+            with torch.cuda.stream(stream):
+                for size in (16, 32):
+                    x = torch.randn(size, 64, device="cuda")
+                    graph(x)
+                    retained.append((graph(x), graph.forward(x)))
+                    graph(x + 1)
+    for stream in streams:
+        torch.cuda.current_stream().wait_stream(stream)
+    pools = {entry.graph_pool for graph in graphs for entry in graph.graphs.values()}
+    assert len(pools) == 4
+    assert all(len(graph._capture_pools) == 2 for graph in graphs)
+    for actual, expected in retained:
+        torch.testing.assert_close(actual, expected)

tests/model_executor/models/minicpmo_4_5/test_vision_batching.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from dataclasses import dataclass
 
 import pytest
@@ -97,6 +100,8 @@ class _FakeVisionModel:
 
 
 class _FakeAudioModel:
+    _encode_audio_features = MiniCPMO45OmniLLMForConditionalGeneration._encode_audio_features
+
     def __init__(self, audio_encoder: _FakeAudioEncoder, audio_encoder_layer: int) -> None:
         self.config = _AudioConfig()
         self.apm = audio_encoder

vllm_omni/model_executor/models/minicpmo_4_5/encoder_cuda_graph.py

@@ -0,0 +1,262 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exact-shape MiniCPM adapters for vLLM encoder CUDA graph management."""
+
+from collections import Counter
+from collections.abc import Callable, Hashable
+from copy import copy
+from math import prod
+
+import torch
+from vllm.config import VllmConfig
+from vllm.model_executor.models.interfaces import SupportsEncoderCudaGraph
+from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager
+from vllm.v1.worker.encoder_cudagraph_defs import (
+    EncoderCudaGraphCaptureInputs,
+    EncoderCudaGraphConfig,
+    EncoderCudaGraphReplayBuffers,
+    EncoderItemSpec,
+)
+
+from vllm_omni.platforms import current_omni_platform
+
+
+def _copy_exact(target: torch.Tensor, source: torch.Tensor) -> None:
+    target.copy_(source)
+
+
+class _ExactShapeEncoder(SupportsEncoderCudaGraph):
+    """One already-packed local encoder batch is one indivisible manager item.
+
+    Packing and distributed placement belong to the caller. Keeping the batch
+    intact preserves attention masks and avoids introducing padding semantics.
+    """
+
+    def __init__(
+        self, forward: Callable[..., torch.Tensor], inputs: tuple[torch.Tensor | None, ...], output_shape: torch.Size
+    ):
+        self.forward = forward
+        self.templates = tuple(None if x is None else (x.shape, x.dtype, x.device) for x in inputs)
+        self.output_shape = output_shape
+        self.tokens = prod(output_shape[:-1]) or 1
+        self.keys = [str(i) for i, x in enumerate(inputs) if x is not None]
+
+    def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
+        return EncoderCudaGraphConfig(
+            modalities=["image", "audio"],
+            buffer_keys=self.keys,
+            out_hidden_size=self.output_shape[-1],
+            padding_logics=dict.fromkeys(self.keys, _copy_exact),
+        )
+
+    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, torch.Tensor]) -> list[EncoderItemSpec]:
+        return [EncoderItemSpec(input_size=self.tokens, output_tokens=self.tokens)]
+
+    def select_encoder_cudagraph_items(
+        self, mm_kwargs: dict[str, torch.Tensor], indices: list[int]
+    ) -> dict[str, torch.Tensor]:
+        assert indices == [0]
+        return dict(mm_kwargs)
+
+    def prepare_encoder_cudagraph_capture_inputs(
+        self,
+        token_budget: int,
+        max_batch_size: int,
+        max_frames_per_batch: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        path: str = "default",
+        axis_keys: tuple[Hashable, ...] | None = (),
+    ) -> EncoderCudaGraphCaptureInputs:
+        return EncoderCudaGraphCaptureInputs(
+            {
+                str(i): torch.zeros(shape, dtype=input_dtype, device=input_device)
+                for i, template in enumerate(self.templates)
+                if template is not None
+                for shape, input_dtype, input_device in [template]
+            }
+        )
+
+    def prepare_encoder_cudagraph_replay_buffers(
+        self,
+        mm_kwargs: dict[str, torch.Tensor],
+        max_batch_size: int,
+        max_frames_per_batch: int,
+        path: str = "default",
+    ) -> EncoderCudaGraphReplayBuffers:
+        return EncoderCudaGraphReplayBuffers(mm_kwargs)
+
+    def encoder_cudagraph_forward(self, inputs: dict[str, torch.Tensor], path: str = "default") -> torch.Tensor:
+        return self.forward(*(inputs.get(str(i)) for i in range(len(self.templates))))
+
+    def encoder_eager_forward(self, mm_kwargs: dict[str, torch.Tensor], 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, torch.Tensor] | None = None,
+    ) -> None:
+        output = outputs["default"]
+        # Callers retain embeddings across subsequent replays.
+        dest[0] = output.clone() if clone else output
+
+
+class EncoderCudaGraph:
+    """Capture repeated shapes without padding or changing attention semantics.
+
+    The owner must bypass this object for training, stateful KV caches and
+    data-dependent Python branches. Entries on the same device and replay
+    stream share a graph pool and capture stream. Different replay streams
+    and encoder instances remain isolated. Graphs must not run concurrently
+    on a shared pool, and their intermediate tensors must not escape capture;
+    callers get a clone so a later replay cannot overwrite retained embeddings.
+    Unseen shapes after the cap run eagerly instead of growing GPU memory.
+    Admission is first-come, with no eviction or automatic recapture. Startup
+    profiling can consume slots. The free-memory floor is a pre-capture guard,
+    not a bound on pool bytes or a guarantee that capture will fit.
+
+    Capture errors are fatal for this instance: propagate the original error
+    and reject subsequent calls without retrying or falling back to CUDA eager
+    execution, since a failed capture may leave the CUDA context unusable.
+    """
+
+    def __init__(
+        self,
+        forward: Callable[..., torch.Tensor],
+        vllm_config: VllmConfig,
+        *,
+        max_graphs: int = 4,
+        min_capture_calls: int = 2,
+        min_free_bytes: int = 1 << 30,
+        share_pools: bool = True,
+    ):
+        for name, value, minimum in (
+            ("max_graphs", max_graphs, 0),
+            ("min_capture_calls", min_capture_calls, 2),
+            ("min_free_bytes", min_free_bytes, 0),
+        ):
+            if type(value) is not int or value < minimum:
+                raise ValueError(f"{name} must be an integer >= {minimum}")
+        self.forward = forward
+        self.vllm_config = vllm_config
+        self.max_graphs = max_graphs
+        self.min_capture_calls = min_capture_calls
+        self.min_free_bytes = min_free_bytes
+        self.share_pools = share_pools
+        self._capture_pools: dict[tuple, tuple[tuple[int, int], torch.cuda.Stream]] = {}
+        self.graphs: dict[tuple, EncoderCudaGraphManager] = {}
+        self._seen: dict[tuple, torch.Size] = {}
+        self._seen_calls: Counter[tuple] = Counter()
+        self._misses: Counter[str] = Counter()
+        self._capture_failed = False
+
+    def get_cumulative_stats(self) -> dict[str, int | float]:
+        """Include eager admission misses that never reach an upstream manager."""
+        hits = sum(entry.graph_hits for entry in self.graphs.values())
+        misses = sum(self._misses.values()) + sum(entry.graph_misses for entry in self.graphs.values())
+        return {
+            "graph_hits": hits,
+            "graph_misses": misses,
+            "hit_rate": hits / (hits + misses) if hits + misses else 0.0,
+            "num_graphs": len(self.graphs),
+            "ineligible_misses": self._misses["ineligible"],
+            "capacity_misses": self._misses["capacity"],
+            "warmup_misses": self._misses["warmup"],
+            "memory_misses": self._misses["memory"],
+            "capture_failures": int(self._capture_failed),
+        }
+
+    def __call__(self, *inputs: torch.Tensor | None) -> torch.Tensor:
+        if self._capture_failed:
+            raise RuntimeError("Encoder CUDA graph capture previously failed; restart the worker before reuse")
+        tensors = [value for value in inputs if value is not None]
+        if (
+            not tensors
+            or any(value.device.type != "cuda" for value in tensors)
+            or len({value.device for value in tensors}) != 1
+            or torch.is_grad_enabled()
+            or torch.is_autocast_enabled("cuda")
+            or torch.cuda.is_current_stream_capturing()
+        ):
+            self._misses["ineligible"] += 1
+            return self.forward(*inputs)
+        # Do not share mutable replay buffers across independent CUDA streams.
+        key = (
+            torch.cuda.current_stream(tensors[0].device).cuda_stream,
+            tuple(None if x is None else (x.shape, x.dtype, x.device) for x in inputs),
+        )
+        entry = self.graphs.get(key)
+        if entry is None:
+            if len(self.graphs) >= self.max_graphs:
+                self._misses["capacity"] += 1
+                return self.forward(*inputs)
+            if key not in self._seen:
+                # Avoid paying capture latency for one-off shapes. Bound even
+                # the admission metadata for streams with arbitrary resolutions.
+                if len(self._seen) >= self.max_graphs * 4:
+                    oldest = next(iter(self._seen))
+                    self._seen.pop(oldest)
+                    self._seen_calls.pop(oldest)
+                output = self.forward(*inputs)
+                self._seen[key] = output.shape
+                self._seen_calls[key] = 1
+                self._misses["warmup"] += 1
+                return output
+            self._seen_calls[key] = min(self._seen_calls[key] + 1, self.min_capture_calls)
+            if self._seen_calls[key] < self.min_capture_calls:
+                self._misses["warmup"] += 1
+                return self.forward(*inputs)
+            if self.min_free_bytes and current_omni_platform.get_free_memory(tensors[0].device) < self.min_free_bytes:
+                self._misses["memory"] += 1
+                return self.forward(*inputs)
+            try:
+                entry = self._capture(inputs, self._seen[key])
+            except Exception:
+                self._capture_failed = True
+                self._seen.clear()
+                self._seen_calls.clear()
+                raise
+            self.graphs[key] = entry
+            self._seen.pop(key)
+            self._seen_calls.pop(key)
+        return entry.execute({str(i): x for i, x in enumerate(inputs) if x is not None})[0]
+
+    def _capture(self, inputs: tuple[torch.Tensor | None, ...], output_shape: torch.Size) -> EncoderCudaGraphManager:
+        tensor = next(x for x in inputs if x is not None)
+        adapter = _ExactShapeEncoder(self.forward, inputs, output_shape)
+        # Do not mutate the runner configuration. This adapter has already
+        # packed and placed a local batch, so manager-level DP must be disabled.
+        config = copy(self.vllm_config)
+        config.compilation_config = copy(config.compilation_config)
+        config.compilation_config.encoder_cudagraph_token_budgets = [adapter.tokens]
+        config.compilation_config.encoder_cudagraph_max_vision_items_per_batch = 1
+        config.compilation_config.encoder_cudagraph_max_frames_per_batch = 0
+        config.parallel_config = copy(config.parallel_config)
+        config.parallel_config.tensor_parallel_size = 1
+        manager = EncoderCudaGraphManager(config, tensor.device, tensor.dtype, adapter)
+        replay_stream = torch.cuda.current_stream(tensor.device)
+        pool_key = (tensor.device, replay_stream.cuda_stream)
+        if self.share_pools and pool_key in self._capture_pools:
+            pool, stream = self._capture_pools[pool_key]
+        else:
+            pool = torch.cuda.graph_pool_handle()
+            stream = torch.cuda.Stream(device=tensor.device)
+        # A late capture must wait for prior replays and output clones before
+        # reusing their pool. Reusing the capture stream also lets the allocator
+        # reuse freed blocks from earlier captures. Manager input/output buffers
+        # are allocated outside capture; only transient intermediates share it.
+        stream.wait_stream(replay_stream)
+        with torch.cuda.stream(stream):
+            manager.capture(graph_pool=pool)
+        replay_stream.wait_stream(stream)
+        if self.share_pools:
+            self._capture_pools[pool_key] = (pool, stream)
+        return manager

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py

@@ -55,6 +55,8 @@ from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPo
 from transformers.modeling_utils import PreTrainedModel
 from transformers.models.whisper.modeling_whisper import ACT2FN
 
+from vllm_omni.model_executor.models.minicpmo_4_5.encoder_cuda_graph import EncoderCudaGraph
+
 try:
     from transformers.models.whisper.modeling_whisper import WHISPER_ATTENTION_CLASSES
 except ImportError:
@@ -1580,6 +1582,7 @@ class SiglipVisionTransformer(SiglipPreTrainedModel):
         self.encoder = SiglipEncoder(config)
         self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
         self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+        self._encoder_graph: EncoderCudaGraph | None = None
 
         # Initialize weights and apply final processing
         self.post_init()
@@ -1587,6 +1590,18 @@ class SiglipVisionTransformer(SiglipPreTrainedModel):
     def get_input_embeddings(self) -> nn.Module:
         return self.embeddings.patch_embedding
 
+    def _encode_last_hidden_state(
+        self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None
+    ) -> torch.Tensor:
+        output = self.encoder(
+            inputs_embeds=hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=False,
+            output_hidden_states=False,
+            return_dict=True,
+        )
+        return self.post_layernorm(output.last_hidden_state)
+
     @add_start_docstrings_to_model_forward(SIGLIP_VISION_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=SiglipVisionConfig)
     def forward(
@@ -1636,6 +1651,21 @@ class SiglipVisionTransformer(SiglipPreTrainedModel):
                 else patch_attention_mask
             )
 
+        # Position/mask construction above performs host decisions. Capture
+        # only the transformer stack; padded FlashAttention unpadding also
+        # reads dynamic lengths on the host and remains eager.
+        if (
+            self._encoder_graph is not None
+            and not self.training
+            and not output_attentions
+            and not output_hidden_states
+            and not (self._use_flash_attention_2 and attention_mask is not None)
+        ):
+            last_hidden_state = self._encoder_graph(hidden_states, attention_mask)
+            if not return_dict:
+                return (last_hidden_state, None)
+            return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state)
+
         encoder_outputs = self.encoder(
             inputs_embeds=hidden_states,
             attention_mask=attention_mask,
@@ -3915,6 +3945,18 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
 
         self.config = config
         self.multimodal_config = multimodal_config
+        # 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
+        )
+        # Per-encoder limits, configurable via --hf-overrides. No eviction:
+        # increasing the cap trades retained GPU memory for shape coverage.
+        encoder_graph_options = {
+            "max_graphs": getattr(config, "encoder_cuda_graph_max_graphs", 4),
+            "min_capture_calls": getattr(config, "encoder_cuda_graph_min_capture_calls", 2),
+            "min_free_bytes": getattr(config, "encoder_cuda_graph_min_free_bytes", 1 << 30),
+            "share_pools": getattr(config, "encoder_cuda_graph_share_pools", True),
+        }
 
         # Initialize image processor
         self.image_processor = MiniCPMVImageProcessor(
@@ -3934,6 +3976,10 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
                 config.vision_config._attn_implementation = "eager"
 
             self.vpm = SiglipVisionTransformer(config.vision_config)
+            if encoder_graphs:
+                self.vpm._encoder_graph = EncoderCudaGraph(
+                    self.vpm._encode_last_hidden_state, vllm_config, **encoder_graph_options
+                )
             # Drop last layer if configured
             if config.drop_vision_last_layer:
                 self.vpm.encoder.layers = self.vpm.encoder.layers[:-1]
@@ -4001,6 +4047,11 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
             self.audio_encoder_layer = None
             self.audio_past_key_values = None
 
+        self._audio_encoder_graph = (
+            EncoderCudaGraph(self._encode_audio_features, vllm_config, **encoder_graph_options)
+            if encoder_graphs
+            else None
+        )
         self.mm_token_ids = set[int]()
         self.make_empty_intermediate_tensors = self.llm.make_empty_intermediate_tensors
 
@@ -4424,6 +4475,36 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
             audio_attention_mask_ = torch.logical_or(audio_attention_mask_, torch.logical_not(chunk_mask))
 
         audio_attention_mask[audio_attention_mask_] = float("-inf")
+        graph = getattr(self, "_audio_encoder_graph", None)
+        if (
+            graph is not None
+            and not self.training
+            and self.audio_encoder_layer == -1
+            # Whisper's FP16 overflow guard reads tensor booleans on the host.
+            and self.apm.conv1.weight.dtype != torch.float16
+        ):
+            audio_embeds = graph(wavforms, audio_attention_mask)
+        else:
+            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
+
+        final_audio_embeds = list[torch.Tensor]()
+        idx = 0
+        for i in range(len(audio_feature_lens_raw)):
+            target_audio_embeds_lst = list[torch.Tensor]()
+            for _ in range(len(audio_feature_lens_raw[i])):
+                target_audio_embeds_lst.append(audio_embeds[idx, : num_audio_tokens[idx], :])
+                idx += 1
+
+            final_audio_embeds.append(torch.cat(target_audio_embeds_lst))
+
+        return final_audio_embeds
+
+    def _encode_audio_features(self, wavforms: torch.Tensor, audio_attention_mask: torch.Tensor) -> torch.Tensor:
+        """Stateless audio encoder, projection and pooling (no streaming KV)."""
         selects_final_layer = self.audio_encoder_layer == -1
         audio_outputs = self.apm(
             wavforms,
@@ -4443,21 +4524,7 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
         audio_embeds = self.audio_avg_pooler(audio_embeds)
         audio_embeds = audio_embeds.transpose(1, 2)
 
-        _, feature_lens_after_pooling = self._get_feat_extract_output_lengths(audio_feature_lens)
-
-        num_audio_tokens = feature_lens_after_pooling
-
-        final_audio_embeds = list[torch.Tensor]()
-        idx = 0
-        for i in range(len(audio_feature_lens_raw)):
-            target_audio_embeds_lst = list[torch.Tensor]()
-            for _ in range(len(audio_feature_lens_raw[i])):
-                target_audio_embeds_lst.append(audio_embeds[idx, : num_audio_tokens[idx], :])
-                idx += 1
-
-            final_audio_embeds.append(torch.cat(target_audio_embeds_lst))
-
-        return final_audio_embeds
+        return audio_embeds
 
     def get_audio_embedding_streaming(
         self,