anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
This commit is contained in:
amy-why-3459 authored and GitHub committed 2026-10-05 15:01:07 +00:00
1 parent a0e626caaf
commit 9146284c16
5 files changed
+789 -15
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/model_executor/models/minicpmo_4_5/encoder_cuda_graph.py
+262
View File
@@ -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
+82 -15
View File
@@ -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,