From 877beeaa4957e8caab0daad31930cb274cb601cb Mon Sep 17 00:00:00 2001 From: liwener Date: Thu, 13 Aug 2026 09:36:17 +0800 Subject: [PATCH] [Perf][MiniCPM-O-4.5]Add CUDA Graph For CFM DiT estimator (#6082) Signed-off-by: liwener Co-authored-by: amy-why-3459 --- .../minicpmo_4_5/test_cuda_graph_wrapper.py | 250 +++++++++++++++++- vllm_omni/deploy/minicpmo_4_5.yaml | 2 + vllm_omni/deploy/minicpmo_4_5_2gpu.yaml | 2 + vllm_omni/deploy/minicpmo_4_5_3gpu.yaml | 2 + vllm_omni/deploy/minicpmo_4_5_8x4090.yaml | 2 + .../models/minicpmo_4_5/batched_token2wav.py | 32 ++- .../models/minicpmo_4_5/cuda_graph_wrapper.py | 125 +++++++++ .../minicpmo_4_5/minicpmo_4_5_code2wav.py | 5 + 8 files changed, 418 insertions(+), 2 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py index 16e04fc09..7ad35b86d 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py +++ b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py @@ -1,21 +1,24 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from functools import lru_cache from types import SimpleNamespace from unittest.mock import Mock import pytest import torch import torch.nn as nn +from vllm.platforms import current_platform from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( HiFTGenerator, ) from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import ( + CFMGraphWrapper, HiFTGraphWrapper, ) -pytestmark = [pytest.mark.core_model] +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] class _F0Predictor(nn.Module): @@ -142,3 +145,248 @@ def test_lazy_capture_limit_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) wrapper._capture.assert_not_called() wrapper.decode_fn.assert_called_once_with(speech_feat, cache_source) assert result is wrapper.decode_fn.return_value + + +# --------------------------------------------------------------------------- +# CFMGraphWrapper tests +# --------------------------------------------------------------------------- + + +class _MiniDiT(nn.Module): + """Minimal DiT-like module with blocks_forward_chunk for CFM graph testing. + + Mimics the upstream cosyvoice2 DiT's cache semantics: + - CausalConv1d.forward_chunk: cat([cnn_cache, x], dim=time) when cache is not None + - Attention.forward_chunk: cat([k, k_cache], dim=seq) when att_cache is not None + """ + + def __init__(self, in_dim: int = 16, hidden: int = 8, depth: int = 2) -> None: + super().__init__() + self.in_proj = nn.Linear(in_dim, hidden) + self.blocks = nn.ModuleList([nn.LayerNorm(hidden) for _ in range(depth)]) + self.final_layer = nn.Linear(hidden, in_dim) + + def blocks_forward_chunk( + self, + x: torch.Tensor, + t: torch.Tensor, + mask: torch.Tensor | None, + cnn_cache: torch.Tensor | None = None, + att_cache: torch.Tensor | None = None, + cnn_cache_buffer: torch.Tensor | None = None, + att_cache_buffer: torch.Tensor | None = None, + ) -> torch.Tensor: + x = x.transpose(1, 2) + x = self.in_proj(x) + for b_idx in range(len(self.blocks)): + # Simulate CausalConv1d: cnn_cache contributes to the first frames. + # Upstream: if cnn_cache[b_idx] is None, creates zeros internally (same as zeros). + cnn_b = cnn_cache[b_idx] if cnn_cache is not None else None + if cnn_b is not None: + x[:, : cnn_b.shape[2], :] += cnn_b.transpose(1, 2) + # Simulate Attention: att_cache contributes a bias when non-empty. + # Upstream: if att_cache[b_idx] is None, skips cat entirely (different path). + att_b = att_cache[b_idx] if att_cache is not None else None + if att_b is not None and att_b.shape[3] > 0: + x += att_b.sum(dim=(1, 2), keepdim=False).unsqueeze(1) + x = self.blocks[b_idx](x) + x = x + t + cnn_cache_buffer[b_idx] = x[:, -2:, :].transpose(1, 2).contiguous() + dt = x.shape[1] + att_cache_buffer[b_idx][:, :, :dt, :] = x.unsqueeze(1) + x = self.final_layer(x) + x = x.transpose(1, 2) + return x + + +def _cfm_inputs( + batch_size: int, chunk_size: int, old_att_len: int, *, device: str = "cuda" +) -> tuple[torch.Tensor, ...]: + depth = 2 + hidden = 8 + estimator_input = torch.randn(batch_size, 16, chunk_size, device=device) + time_emb = torch.randn(batch_size, 1, hidden, device=device) + cnn_cache = torch.randn(depth, batch_size, hidden, 2, device=device) + att_cache = torch.randn(depth, batch_size, 1, old_att_len, hidden, device=device) + cnn_out = torch.empty(depth, batch_size, hidden, 2, device=device) + att_out = torch.empty(depth, batch_size, 1, old_att_len + chunk_size, hidden, device=device) + return estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cfm_graph_replay_matches_eager_for_uncached_and_cached_shapes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _MiniDiT().eval().cuda() + wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=32) + + with torch.inference_mode(): + for _, chunk_size, old_att_len in ((2, 10, 0), (2, 10, 5)): + inputs = _cfm_inputs(2, chunk_size, old_att_len) + + eager_inputs = tuple(v.clone() for v in inputs) + with torch.no_grad(): + eager_result = estimator.blocks_forward_chunk( + eager_inputs[0], + eager_inputs[1], + None, + eager_inputs[2], + eager_inputs[3], + eager_inputs[4], + eager_inputs[5], + ) + + wrapper.replay(*inputs) + replay_inputs = tuple(v.clone() for v in inputs) + graph_result, graph_cnn, graph_att = wrapper.replay(*replay_inputs) + + torch.testing.assert_close(graph_result, eager_result, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_cnn, eager_inputs[4], rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_att, eager_inputs[5], rtol=1e-4, atol=1e-5) + + wrapper._capture_graph.cache_clear() + + +def _cfm_mock_wrapper(monkeypatch: pytest.MonkeyPatch, *, max_graphs: int = 1) -> CFMGraphWrapper: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + wrapper = object.__new__(CFMGraphWrapper) + wrapper.max_graphs = max_graphs + wrapper.graph_fn = Mock(return_value=torch.tensor([42.0])) + wrapper.device = torch.device("cuda") + + @lru_cache(maxsize=max_graphs) + def _capture_graph(key: tuple): + return wrapper._do_capture(key) + + wrapper._capture_graph = _capture_graph + wrapper._do_capture = Mock(return_value=None) + return wrapper + + +def test_cfm_unseen_shape_is_lazily_captured(monkeypatch: pytest.MonkeyPatch) -> None: + wrapper = _cfm_mock_wrapper(monkeypatch) + wrapper._capture_graph = Mock(return_value=None) + + inputs = _cfm_inputs(2, 10, 0) + wrapper.replay(*inputs) + + wrapper._capture_graph.assert_called_once() + wrapper.graph_fn.assert_called_once() + + +def test_cfm_capture_failure_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None: + wrapper = _cfm_mock_wrapper(monkeypatch) + wrapper._capture_graph = lru_cache(maxsize=1)(lambda key: None) + + inputs = _cfm_inputs(2, 10, 0) + result = wrapper.replay(*inputs) + + wrapper.graph_fn.assert_called_once() + assert result[0] is wrapper.graph_fn.return_value + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cfm_lru_eviction(monkeypatch: pytest.MonkeyPatch) -> None: + """Verify that LRU cache evicts oldest entries when full.""" + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _MiniDiT().eval().cuda() + wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=2) + + with torch.inference_mode(): + inputs_a = _cfm_inputs(2, 10, 0) + wrapper.replay(*inputs_a) + assert wrapper._capture_graph.cache_info().currsize == 1 + + inputs_b = _cfm_inputs(2, 12, 0) + wrapper.replay(*inputs_b) + assert wrapper._capture_graph.cache_info().currsize == 2 + + wrapper.replay(*inputs_a) + assert wrapper._capture_graph.cache_info().currsize == 2 + + inputs_c = _cfm_inputs(2, 14, 0) + wrapper.replay(*inputs_c) + assert wrapper._capture_graph.cache_info().currsize == 2 + + wrapper.replay(*inputs_b) + hits_before = wrapper._capture_graph.cache_info().hits + wrapper.replay(*inputs_b) + hits_after = wrapper._capture_graph.cache_info().hits + assert hits_after == hits_before + 1 + + wrapper._capture_graph.cache_clear() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cfm_none_cache_parity_between_graph_and_eager( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Verify that None→zeros (graph path) matches [None]*depth (eager path). + + setup_batch calls _decode_cfm with cnn_cache=None on every request's + prompt pass. The graph path replaces None with zeros; the eager path + passes [None]*depth. Both must produce identical outputs. + """ + torch.accelerator.synchronize() + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _MiniDiT().eval().cuda() + wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=32) + + depth = len(estimator.blocks) + hidden = 8 + batch_size = 2 + chunk_size = 20 + + estimator_input = torch.randn(batch_size, 16, chunk_size, device="cuda") + time_emb = torch.randn(batch_size, 1, hidden, device="cuda") + cnn_out = torch.empty(depth, batch_size, hidden, 2, device="cuda") + att_out = torch.empty(depth, batch_size, 1, chunk_size, hidden, device="cuda") + + # Eager path: cnn_cache=[None]*depth, att_cache=[None]*depth + eager_input = estimator_input.clone() + eager_time = time_emb.clone() + eager_cnn_out = torch.empty_like(cnn_out) + eager_att_out = torch.empty_like(att_out) + with torch.no_grad(): + eager_result = estimator.blocks_forward_chunk( + eager_input, + eager_time, + None, + [None] * depth, + [None] * depth, + eager_cnn_out, + eager_att_out, + ) + + # Graph path: cnn_cache=zeros, att_cache=zero-length (simulating _estimator_step's None→zeros) + zero_cnn = torch.zeros_like(cnn_out) + zero_att = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:]) + wrapper.replay(estimator_input, time_emb, zero_cnn, zero_att, cnn_out, att_out) + graph_input = estimator_input.clone() + graph_time = time_emb.clone() + graph_cnn_out = torch.empty_like(cnn_out) + graph_att_out = torch.empty_like(att_out) + graph_zero_cnn = torch.zeros_like(cnn_out) + graph_zero_att = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:]) + graph_result, graph_cnn, graph_att = wrapper.replay( + graph_input, + graph_time, + graph_zero_cnn, + graph_zero_att, + graph_cnn_out, + graph_att_out, + ) + + torch.testing.assert_close(graph_result, eager_result, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_cnn, eager_cnn_out, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_att, eager_att_out, rtol=1e-4, atol=1e-5) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 668954f5d..bc34bfe9f 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -19,6 +19,8 @@ connectors: # A future initial_codec_chunk_frames shape will use lazy capture unless # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. enable_hift_graph: true + enable_cfm_graph: true + cfm_max_graphs: 32 stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml index e44140b2d..2a4abed42 100644 --- a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml @@ -22,6 +22,8 @@ connectors: # A future initial_codec_chunk_frames shape will use lazy capture unless # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. enable_hift_graph: true + enable_cfm_graph: true + cfm_max_graphs: 32 stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml index 4933d0c74..b34673910 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml @@ -22,6 +22,8 @@ connectors: # A future initial_codec_chunk_frames shape will use lazy capture unless # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. enable_hift_graph: true + enable_cfm_graph: true + cfm_max_graphs: 32 stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml index 69fa6d5b6..813a2e2eb 100644 --- a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml @@ -33,6 +33,8 @@ connectors: # A future initial_codec_chunk_frames shape will use lazy capture unless # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. enable_hift_graph: true + enable_cfm_graph: true + cfm_max_graphs: 32 stages: - stage_id: 0 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py index 651284fe0..41223e864 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py +++ b/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( diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py b/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py index 5f6c965d5..97015855e 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py +++ b/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(), + ) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index d6d3e3d47..686b39c5b 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/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, )