diff --git a/docs/design/minicpmo_encoder_cuda_graph.md b/docs/design/minicpmo_encoder_cuda_graph.md new file mode 100644 index 000000000..2f7922618 --- /dev/null +++ b/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. diff --git a/tests/model_executor/models/minicpmo_4_5/test_encoder_cuda_graph.py b/tests/model_executor/models/minicpmo_4_5/test_encoder_cuda_graph.py new file mode 100644 index 000000000..4ab35d30e --- /dev/null +++ b/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) diff --git a/tests/model_executor/models/minicpmo_4_5/test_vision_batching.py b/tests/model_executor/models/minicpmo_4_5/test_vision_batching.py index 250d63b12..3c2482a8b 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_vision_batching.py +++ b/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 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/encoder_cuda_graph.py b/vllm_omni/model_executor/models/minicpmo_4_5/encoder_cuda_graph.py new file mode 100644 index 000000000..dc4286dae --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index 7f1d29345..086bf74e8 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/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,