[Perf][MiniCPM-O-4.5]Add CUDA Graph For CFM DiT estimator (#6082)

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

vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py

@@ -14,7 +14,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from vllm.logger import init_logger
 
-from .cuda_graph_wrapper import HiFTGraphWrapper
+from .cuda_graph_wrapper import CFMGraphWrapper, HiFTGraphWrapper
 
 logger = init_logger(__name__)
 
@@ -73,6 +73,7 @@ class BatchedToken2Wav(nn.Module):
         *,
         connector_config: Mapping[str, int] | None = None,
         hift_graph_config: Mapping[str, Any] | None = None,
+        cfm_graph_config: Mapping[str, Any] | None = None,
     ):
         super().__init__()
         self._token2wav = token2wav
@@ -135,6 +136,22 @@ class BatchedToken2Wav(nn.Module):
                 with torch.inference_mode(), _autocast_disabled(hift_parameter.device):
                     self.hift_graph_wrapper.capture()
                 logger.info("HiFT CUDA Graph captured successfully")
+        self._cfm_graph_wrapper: CFMGraphWrapper | None = None
+        cfm_graph_cfg = dict(cfm_graph_config or {})
+        if bool(cfm_graph_cfg.get("enabled", False)):
+            flow_parameter = next(self.flow.parameters(), None)
+            if flow_parameter is not None and flow_parameter.device.type == "cuda":
+                estimator = self.flow.decoder.estimator
+                self._cfm_graph_wrapper = CFMGraphWrapper(
+                    graph_fn=estimator.blocks_forward_chunk,
+                    max_graphs=int(cfm_graph_cfg.get("max_graphs", 32)),
+                )
+                logger.info("CFM CUDA Graph enabled (max_graphs=%d)", int(cfm_graph_cfg.get("max_graphs", 32)))
+            else:
+                logger.info(
+                    "CFM CUDA Graph is disabled on device type %s",
+                    flow_parameter.device.type if flow_parameter is not None else "unknown",
+                )
         self._prompt_features: dict[tuple[str, str], PromptFeatures] = {}
 
     def _hift_inference(
@@ -265,6 +282,19 @@ class BatchedToken2Wav(nn.Module):
         speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width)
         estimator_input = torch.cat((x, mu, speaker_features, cond), dim=1)
         cnn_out, att_out = self._estimator_buffers(estimator, estimator_input, att_cache)
+        if self._cfm_graph_wrapper is not None:
+            if cnn_cache is None:
+                cnn_cache = torch.zeros_like(cnn_out)
+            if att_cache is None:
+                att_cache = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:])
+            return self._cfm_graph_wrapper.replay(
+                estimator_input,
+                time_embedding,
+                cnn_cache,
+                att_cache,
+                cnn_out,
+                att_out,
+            )
         old_cnn: Any = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks)
         old_att: Any = att_cache if att_cache is not None else [None] * len(estimator.blocks)
         result = estimator.blocks_forward_chunk(

vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from functools import lru_cache
+
 import torch
 from torch.cuda import CUDAGraph
 from vllm.logger import init_logger
@@ -133,3 +135,126 @@ class HiFTGraphWrapper:
         cache_source = static_cache_source_output[:batch_size].clone()
         speech = self.finalize_fn(static_magnitude_output[:batch_size], static_phase_output[:batch_size]).clone()
         return speech, cache_source
+
+
+def _tensor_signature(value: torch.Tensor) -> tuple:
+    return tuple(value.shape), str(value.dtype), str(value.device)
+
+
+_DTYPE_MAP = {
+    "torch.float32": torch.float32,
+    "torch.float16": torch.float16,
+    "torch.bfloat16": torch.bfloat16,
+    "torch.float64": torch.float64,
+}
+
+
+def _tensors_from_key(key: tuple) -> tuple[torch.Tensor, ...]:
+    """Rebuild zero tensors from a cache key (shape, dtype, device tuples)."""
+    tensors = []
+    for shape, dtype_str, device_str in key[1:]:
+        dtype = _DTYPE_MAP.get(dtype_str, torch.float32)
+        device = torch.device(device_str)
+        tensors.append(torch.zeros(shape, dtype=dtype, device=device))
+    return tuple(tensors)
+
+
+class CFMGraphWrapper:
+    """Per-shape CUDA graph capture/replay for the CFM DiT estimator.
+
+    Captures one blocks_forward_chunk call (in_proj -> DiT blocks -> final_layer)
+    as the graph target. The 10-step Euler loop stays in Python, replaying
+    the graph 10 times per decode.
+
+    Uses functools.lru_cache for automatic LRU eviction. Cache misses
+    trigger capture; capture failures fall back to eager. Outputs are
+    cloned after replay to prevent streaming cache corruption.
+    """
+
+    def __init__(self, graph_fn, *, max_graphs: int = 32) -> None:
+        self.graph_fn = graph_fn
+        self.max_graphs = int(max_graphs)
+        self.device = next(graph_fn.__self__.parameters()).device
+
+        @lru_cache(maxsize=self.max_graphs)
+        def _capture_graph(key: tuple):
+            return self._capture(key)
+
+        self._capture_graph = _capture_graph
+
+    def _call_graph_fn(self, args: tuple[torch.Tensor, ...]) -> torch.Tensor:
+        return self.graph_fn(args[0], args[1], None, args[2], args[3], args[4], args[5])
+
+    def _capture(self, key: tuple) -> tuple | None:
+        """Capture a CUDA graph for the given key. Returns None on failure."""
+        static_inputs = _tensors_from_key(key)
+
+        current_stream = torch.cuda.current_stream(self.device)
+        warmup_stream = torch.cuda.Stream(device=self.device)
+        warmup_stream.wait_stream(current_stream)
+        with torch.cuda.stream(warmup_stream), torch.no_grad():
+            for _ in range(3):
+                self._call_graph_fn(static_inputs)
+        current_stream.wait_stream(warmup_stream)
+
+        try:
+            graph = CUDAGraph()
+            with torch.no_grad(), torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
+                static_output = self._call_graph_fn(static_inputs)
+        except Exception:
+            logger.warning(
+                "CFM graph capture failed for shape=%s; using eager",
+                key,
+                exc_info=True,
+            )
+            return None
+
+        logger.info(
+            "Captured CFM CUDA Graph for shape %s (cache=%d/%d, hits=%d, misses=%d)",
+            key,
+            self._capture_graph.cache_info().currsize,
+            self.max_graphs,
+            self._capture_graph.cache_info().hits,
+            self._capture_graph.cache_info().misses,
+        )
+        return (static_inputs, static_output, graph)
+
+    def replay(
+        self,
+        estimator_input: torch.Tensor,
+        time_emb: torch.Tensor,
+        cnn_cache: torch.Tensor,
+        att_cache: torch.Tensor,
+        cnn_out: torch.Tensor,
+        att_out: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        inputs = (estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out)
+
+        if torch.cuda.is_current_stream_capturing() or estimator_input.device.type != "cuda":
+            with torch.no_grad():
+                result = self._call_graph_fn(inputs)
+            return result, inputs[4], inputs[5]
+
+        key = ("estimator_step",) + tuple(_tensor_signature(v) for v in inputs)
+        entry = self._capture_graph(key)
+
+        if entry is None:
+            logger.debug(
+                "CFM graph eager fallback for shape=%s (hits=%d, misses=%d)",
+                key,
+                self._capture_graph.cache_info().hits,
+                self._capture_graph.cache_info().misses,
+            )
+            with torch.no_grad():
+                result = self._call_graph_fn(inputs)
+            return result, inputs[4], inputs[5]
+
+        static_inputs, static_output, graph = entry
+        for static, current in zip(static_inputs, inputs, strict=True):
+            static.copy_(current)
+        graph.replay()
+        return (
+            static_output.detach().clone(),
+            static_inputs[4].detach().clone(),
+            static_inputs[5].detach().clone(),
+        )

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py

@@ -162,6 +162,10 @@ class MiniCPMO45Code2Wav(nn.Module):
             "enabled": bool(extra.get("enable_hift_graph", False)),
             "capture_batch_sizes": capture_batch_sizes,
         }
+        self._cfm_graph_config = {
+            "enabled": bool(extra.get("enable_cfm_graph", False)),
+            "max_graphs": int(extra.get("cfm_max_graphs", 32)),
+        }
         self._min_batch_size = int(extra.get("code2wav_min_batch_size", 1))
         if self._min_batch_size < 1:
             raise ValueError("MiniCPM-o Code2Wav code2wav_min_batch_size must be >= 1")
@@ -802,4 +806,5 @@ class MiniCPMO45Code2Wav(nn.Module):
             trt_stepper=trt_stepper,
             connector_config=self._connector_config,
             hift_graph_config=self._hift_graph_config,
+            cfm_graph_config=self._cfm_graph_config,
         )