[Perf][MiniCPM-O-4.5]Add CUDA Graph For HiFTGenerator (#5869)

Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>

tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py

@@ -105,6 +105,9 @@ class _FakeHiFT(nn.Module):
         generated_source = speech[:, None]
         return speech, generated_source
 
+    def inference(self, mel, source):
+        return self.forward(mel, source)
+
 
 class _FakeToken2Wav:
     def __init__(self):

tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+import torch.nn as nn
+
+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 (
+    HiFTGraphWrapper,
+)
+
+pytestmark = [pytest.mark.core_model]
+
+
+class _F0Predictor(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.proj = nn.Conv1d(80, 1, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.proj(x).squeeze(1).abs()
+
+
+class _DeterministicSineGen(nn.Module):
+    """Remove source RNG so eager and replay compare only execution paths."""
+
+    def __init__(self, num_harmonics: int) -> None:
+        super().__init__()
+        self.num_harmonics = num_harmonics
+
+    def forward(self, f0: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        shape = (*f0.shape[:-1], self.num_harmonics)
+        sine = f0.new_zeros(shape)
+        uv = f0.new_ones((*f0.shape[:-1], 1))
+        return sine, uv, sine
+
+
+def _small_hift() -> HiFTGenerator:
+    hift = HiFTGenerator(
+        base_channels=32,
+        sampling_rate=24000,
+        upsample_rates=[8, 5, 3],
+        upsample_kernel_sizes=[16, 11, 7],
+        source_resblock_kernel_sizes=[7, 7, 11],
+        source_resblock_dilation_sizes=[[1, 3, 5]] * 3,
+        f0_predictor=_F0Predictor(),
+    )
+    hift.m_source.l_sin_gen = _DeterministicSineGen(hift.nb_harmonics + 1)
+    return hift.eval().cuda()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+def test_hift_graph_replay_matches_eager_for_uncached_and_cached_shapes() -> None:
+    torch.manual_seed(0)
+    hift = _small_hift()
+    token2wav = SimpleNamespace(
+        hift=hift,
+        flow=SimpleNamespace(
+            encoder=SimpleNamespace(pre_lookahead_layer=SimpleNamespace(pre_lookahead_len=3)),
+            token_mel_ratio=2,
+        ),
+        mel_cache_len=2,
+        source_cache_len=960,
+    )
+    wrapper = HiFTGraphWrapper(
+        token2wav,
+        connector_config={"codec_chunk_frames": 2, "codec_left_context_frames": 3},
+        capture_batch_sizes=[1],
+    )
+    wrapper.capture()
+
+    cases = (
+        (torch.randn(1, 80, 4, device="cuda"), torch.zeros(1, 1, 0, device="cuda")),
+        (torch.randn(1, 80, 6, device="cuda"), torch.randn(1, 1, 960, device="cuda")),
+    )
+    with torch.inference_mode():
+        for speech_feat, cache_source in cases:
+            expected_speech, expected_source = hift.inference(speech_feat, cache_source)
+            actual_speech, actual_source = wrapper.replay(speech_feat, cache_source)
+            torch.testing.assert_close(actual_speech, expected_speech, rtol=1e-4, atol=1e-5)
+            torch.testing.assert_close(actual_source, expected_source, rtol=1e-4, atol=1e-5)
+
+
+class _FakeGraph:
+    def replay(self) -> None:
+        return None
+
+
+def _fake_wrapper(monkeypatch: pytest.MonkeyPatch) -> HiFTGraphWrapper:
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+    wrapper = object.__new__(HiFTGraphWrapper)
+    wrapper.capture_batch_sizes = [1]
+    wrapper.graph = {}
+    wrapper.static_speech_inputs = {}
+    wrapper.static_cache_source_inputs = {}
+    wrapper.static_magnitude_outputs = {}
+    wrapper.static_phase_outputs = {}
+    wrapper.static_cache_source_outputs = {}
+    wrapper.lazy_graph_count = 0
+    wrapper.max_lazy_graphs = 1
+    wrapper.decode_fn = Mock(return_value=(torch.tensor([[99.0]]), torch.tensor([[[98.0]]])))
+    wrapper.finalize_fn = lambda magnitude, phase: magnitude + phase
+
+    def capture(batch_size: int, num_frames: int, cache_len: int) -> None:
+        key = (batch_size, num_frames, cache_len)
+        wrapper.graph[key] = _FakeGraph()
+        wrapper.static_speech_inputs[key] = torch.zeros(batch_size, 80, num_frames)
+        wrapper.static_cache_source_inputs[key] = torch.zeros(batch_size, 1, cache_len)
+        wrapper.static_magnitude_outputs[key] = torch.ones(batch_size, 1, num_frames)
+        wrapper.static_phase_outputs[key] = torch.ones(batch_size, 1, num_frames)
+        wrapper.static_cache_source_outputs[key] = torch.ones(batch_size, 1, num_frames)
+
+    wrapper._capture = Mock(side_effect=capture)
+    return wrapper
+
+
+def test_unseen_shape_is_lazily_captured(monkeypatch: pytest.MonkeyPatch) -> None:
+    wrapper = _fake_wrapper(monkeypatch)
+    speech, source = wrapper.replay(torch.randn(1, 80, 7), torch.zeros(1, 1, 0))
+
+    wrapper._capture.assert_called_once_with(1, 7, 0)
+    assert wrapper.lazy_graph_count == 1
+    assert speech.shape == (1, 1, 7)
+    assert source.shape == (1, 1, 7)
+    wrapper.decode_fn.assert_not_called()
+
+
+def test_lazy_capture_limit_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None:
+    wrapper = _fake_wrapper(monkeypatch)
+    wrapper.lazy_graph_count = wrapper.max_lazy_graphs
+    speech_feat = torch.randn(1, 80, 9)
+    cache_source = torch.zeros(1, 1, 0)
+
+    result = wrapper.replay(speech_feat, cache_source)
+
+    wrapper._capture.assert_not_called()
+    wrapper.decode_fn.assert_called_once_with(speech_feat, cache_source)
+    assert result is wrapper.decode_fn.return_value

tests/model_executor/models/minicpmo_4_5/test_pipeline.py

@@ -157,6 +157,7 @@ class TestDeployTopology:
         assert connector["name"] == "SharedMemoryConnector"
         assert connector["extra"]["codec_chunk_frames"] == 25
         assert connector["extra"]["codec_left_context_frames"] == 3
+        assert connector["extra"]["enable_hift_graph"] is True
         assert connector["extra"]["connector_get_max_wait_first_chunk"] == 3000
         assert connector["extra"]["connector_get_max_wait"] == 300
         expected_processor = "tts2code2wav_async_chunk" if deploy.async_chunk else "tts2code2wav_full_payload"

tests/model_executor/models/step_audio2/test_hift_parity.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.cosyvoice3.utils import mel_spectrogram
+from vllm_omni.model_executor.models.step_audio2.step_audio2_token2wav import _build_hift
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+flashcosyvoice_hifigan = pytest.importorskip("flashcosyvoice.modules.hifigan")
+flashcosyvoice_audio = pytest.importorskip("flashcosyvoice.utils.audio")
+
+
+def test_vendored_hift_matches_flashcosyvoice() -> None:
+    reference = flashcosyvoice_hifigan.HiFTGenerator().eval()
+    vendored = _build_hift().eval()
+
+    # Strict loading verifies that the vendored architecture remains compatible
+    # with the Step-Audio2 checkpoint produced for flashcosyvoice.
+    vendored.load_state_dict(reference.state_dict(), strict=True)
+
+    speech_feat = torch.randn(1, 80, 4)
+    cache_source = torch.zeros(1, 1, 0)
+    with torch.inference_mode():
+        torch.manual_seed(0)
+        expected_speech, expected_source = reference(speech_feat, cache_source)
+        torch.manual_seed(0)
+        actual_speech, actual_source = vendored.inference(speech_feat, cache_source)
+
+    torch.testing.assert_close(actual_source, expected_source)
+    torch.testing.assert_close(actual_speech, expected_speech)
+
+
+def test_vendored_mel_spectrogram_matches_flashcosyvoice() -> None:
+    reference = flashcosyvoice_audio.mel_spectrogram
+    vendored = mel_spectrogram
+
+    audio = torch.randn(1, 48000)
+    kwargs = {
+        "n_fft": 1920,
+        "num_mels": 80,
+        "sampling_rate": 24000,
+        "hop_size": 480,
+        "win_size": 1920,
+        "fmin": 0,
+        "fmax": 8000,
+    }
+    with torch.inference_mode():
+        expected_spec = reference(audio, **kwargs)
+        actual_spec = vendored(audio, **kwargs)
+
+    torch.testing.assert_close(actual_spec, expected_spec)

vllm_omni/deploy/minicpmo_4_5.yaml

@@ -14,6 +14,11 @@ connectors:
       connector_get_max_wait: 300
       codec_chunk_frames: 25
       codec_left_context_frames: 3
+      # HiFT capture buckets are derived from codec_chunk_frames and
+      # codec_left_context_frames. (#5869)
+      # A future initial_codec_chunk_frames shape will use lazy capture unless
+      # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
+      enable_hift_graph: true
 
 stages:
   - stage_id: 0

vllm_omni/deploy/minicpmo_4_5_2gpu.yaml

@@ -17,6 +17,11 @@ connectors:
       connector_get_max_wait: 300
       codec_chunk_frames: 25
       codec_left_context_frames: 3
+      # HiFT capture buckets are derived from codec_chunk_frames and
+      # codec_left_context_frames. (#5869)
+      # A future initial_codec_chunk_frames shape will use lazy capture unless
+      # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
+      enable_hift_graph: true
 
 stages:
   - stage_id: 0

vllm_omni/deploy/minicpmo_4_5_3gpu.yaml

@@ -17,6 +17,11 @@ connectors:
       connector_get_max_wait: 300
       codec_chunk_frames: 25
       codec_left_context_frames: 3
+      # HiFT capture buckets are derived from codec_chunk_frames and
+      # codec_left_context_frames. (#5869)
+      # A future initial_codec_chunk_frames shape will use lazy capture unless
+      # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
+      enable_hift_graph: true
 
 stages:
   - stage_id: 0

vllm_omni/deploy/minicpmo_4_5_8x4090.yaml

@@ -28,6 +28,11 @@ connectors:
       connector_get_max_wait: 300
       codec_chunk_frames: 25
       codec_left_context_frames: 3
+      # HiFT capture buckets are derived from codec_chunk_frames and
+      # codec_left_context_frames. (#5869)
+      # A future initial_codec_chunk_frames shape will use lazy capture unless
+      # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
+      enable_hift_graph: true
 
 stages:
   - stage_id: 0

vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py

@@ -200,6 +200,11 @@ class SineGen2(torch.nn.Module):
         self.flag_for_pulse = flag_for_pulse
         self.upsample_scale = upsample_scale
         self.causal = causal
+        self.register_buffer(
+            "harmonic_ids",
+            torch.arange(1, self.harmonic_num + 2, dtype=torch.float32).view(1, 1, -1),
+            persistent=False,
+        )
         if causal is True:
             self.rand_ini = torch.rand(1, 9)
             self.rand_ini[:, 0] = 0
@@ -277,7 +282,7 @@ class SineGen2(torch.nn.Module):
         output uv: tensor(batchsize=1, length, 1)
         """
         # fundamental component
-        fn = torch.multiply(f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device))
+        fn = torch.multiply(f0, self.harmonic_ids)
 
         # generate sine waveforms
         sine_waves = self._f02sine(fn) * self.sine_amp
@@ -463,7 +468,11 @@ class HiFTGenerator(nn.Module):
         self.ups.apply(init_weights)
         self.conv_post.apply(init_weights)
         self.reflection_pad = nn.ReflectionPad1d((1, 0))
-        self.stft_window = torch.from_numpy(get_window("hann", istft_params["n_fft"], fftbins=True).astype(np.float32))
+        self.register_buffer(
+            "stft_window",
+            torch.from_numpy(get_window("hann", istft_params["n_fft"], fftbins=True).astype(np.float32)),
+            persistent=False,
+        )
         self.f0_predictor = f0_predictor
 
     def remove_weight_norm(self):
@@ -485,7 +494,7 @@ class HiFTGenerator(nn.Module):
             self.istft_params["n_fft"],
             self.istft_params["hop_len"],
             self.istft_params["n_fft"],
-            window=self.stft_window.to(x.device),
+            window=self.stft_window,
             return_complex=True,
         )
         spec = torch.view_as_real(spec)  # [B, F, TT, 2]
@@ -500,11 +509,15 @@ class HiFTGenerator(nn.Module):
             self.istft_params["n_fft"],
             self.istft_params["hop_len"],
             self.istft_params["n_fft"],
-            window=self.stft_window.to(magnitude.device),
+            window=self.stft_window,
         )
         return inverse_transform
 
-    def decode(self, x: torch.Tensor, s: torch.Tensor = torch.zeros(1, 1, 0)) -> torch.Tensor:
+    def _decode_pre_istft(
+        self,
+        x: torch.Tensor,
+        s: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         s_stft_real, s_stft_imag = self._stft(s.squeeze(1))
         s_stft = torch.cat([s_stft_real, s_stft_imag], dim=1)
 
@@ -533,11 +546,17 @@ class HiFTGenerator(nn.Module):
         x = self.conv_post(x)
         magnitude = torch.exp(x[:, : self.istft_params["n_fft"] // 2 + 1, :])
         phase = torch.sin(x[:, self.istft_params["n_fft"] // 2 + 1 :, :])  # actually, sin is redundancy
+        return magnitude, phase
 
+    def _finalize_decode(self, magnitude, phase):
         x = self._istft(magnitude, phase)
         x = torch.clamp(x, -self.audio_limit, self.audio_limit)
         return x
 
+    def decode(self, x: torch.Tensor, s: torch.Tensor = torch.zeros(1, 1, 0)) -> torch.Tensor:
+        magnitude, phase = self._decode_pre_istft(x, s)
+        return self._finalize_decode(magnitude, phase)
+
     def forward(
         self,
         batch: dict,
@@ -555,7 +574,11 @@ class HiFTGenerator(nn.Module):
         return generated_speech, f0
 
     @torch.inference_mode()
-    def inference(self, speech_feat: torch.Tensor, cache_source: torch.Tensor = torch.zeros(1, 1, 0)) -> torch.Tensor:
+    def _inference_pre_istft(
+        self,
+        speech_feat: torch.Tensor,
+        cache_source: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         # mel->f0
         f0 = self.f0_predictor(speech_feat)
         # f0->source
@@ -565,7 +588,13 @@ class HiFTGenerator(nn.Module):
         # use cache_source to avoid glitch
         if cache_source.shape[2] != 0:
             s[:, :, : cache_source.shape[2]] = cache_source
-        generated_speech = self.decode(x=speech_feat, s=s)
+        magnitude, phase = self._decode_pre_istft(speech_feat, s)
+        return magnitude, phase, s
+
+    @torch.inference_mode()
+    def inference(self, speech_feat: torch.Tensor, cache_source: torch.Tensor = torch.zeros(1, 1, 0)) -> torch.Tensor:
+        magnitude, phase, s = self._inference_pre_istft(speech_feat, cache_source)
+        generated_speech = self._finalize_decode(magnitude, phase)
         return generated_speech, s
 
 

vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Mapping
 from contextlib import nullcontext
 from dataclasses import dataclass
 from typing import Any
@@ -11,6 +12,11 @@ from typing import Any
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from vllm.logger import init_logger
+
+from .cuda_graph_wrapper import HiFTGraphWrapper
+
+logger = init_logger(__name__)
 
 _SILENCE_TOKEN = 4218
 
@@ -60,7 +66,14 @@ class BatchedToken2Wav(nn.Module):
     asset loader and prompt feature extractor.
     """
 
-    def __init__(self, token2wav: Any, trt_stepper: Any | None = None):
+    def __init__(
+        self,
+        token2wav: Any,
+        trt_stepper: Any | None = None,
+        *,
+        connector_config: Mapping[str, int] | None = None,
+        hift_graph_config: Mapping[str, Any] | None = None,
+    ):
         super().__init__()
         self._token2wav = token2wav
         # Optional TrtDiTStepper (step_audio2_dit_trt): replaces only the
@@ -82,7 +95,7 @@ class BatchedToken2Wav(nn.Module):
                 _autocast_disabled(device),
             ):
                 # 50 mel frames match the default first streamed vocoder chunk.
-                speech, source = self.hift(
+                speech, source = self.hift.inference(
                     torch.zeros((1, mel_channels, 50), device=device, dtype=dtype),
                     torch.zeros((1, 1, 0), device=device, dtype=dtype),
                 )
@@ -98,8 +111,41 @@ class BatchedToken2Wav(nn.Module):
             token2wav.speech_window.detach().clone(),
             persistent=False,
         )
+        self.hift_graph_wrapper: HiFTGraphWrapper | None = None
+        graph_config = dict(hift_graph_config or {})
+        if bool(graph_config.get("enabled", False)):
+            if hift_parameter is None:
+                raise ValueError("MiniCPM-o HiFT Graph requires a parameterized HiFT module")
+            if hift_parameter.device.type != "cuda":
+                logger.info("HiFT CUDA Graph is disabled on device type %s", hift_parameter.device.type)
+            else:
+                if connector_config is None:
+                    raise ValueError("MiniCPM-o HiFT CUDA Graph requires connector chunk configuration")
+                if self.mel_cache_len <= 0 or self.source_cache_len % self.mel_cache_len != 0:
+                    raise ValueError(
+                        "MiniCPM-o HiFT CUDA Graph requires source_cache_len to be divisible by mel_cache_len"
+                    )
+                capture_batch_sizes = graph_config.get("capture_batch_sizes", [1])
+                logger.info("Enabling HiFT CUDA Graph with batch sizes %s", capture_batch_sizes)
+                self.hift_graph_wrapper = HiFTGraphWrapper(
+                    token2wav=token2wav,
+                    connector_config=dict(connector_config),
+                    capture_batch_sizes=capture_batch_sizes,
+                )
+                with torch.inference_mode(), _autocast_disabled(hift_parameter.device):
+                    self.hift_graph_wrapper.capture()
+                logger.info("HiFT CUDA Graph captured successfully")
         self._prompt_features: dict[tuple[str, str], PromptFeatures] = {}
 
+    def _hift_inference(
+        self,
+        mel: torch.Tensor,
+        source_cache: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.hift_graph_wrapper is None:
+            return self.hift.inference(mel, source_cache)
+        return self.hift_graph_wrapper.replay(mel, source_cache)
+
     def prepare_prompt(self, prompt_cache_id: str, prompt_wav: str) -> PromptFeatures:
         cache_key = (prompt_cache_id, prompt_wav)
         cached = self._prompt_features.get(cache_key)
@@ -463,7 +509,7 @@ class BatchedToken2Wav(nn.Module):
         old_source = torch.cat([state.hift_cache["source"] for state in states], dim=0)
         old_speech = torch.cat([state.hift_cache["speech"] for state in states], dim=0)
         mel = torch.cat((old_mel, chunk_mel), dim=2)
-        speech, source = self.hift(mel, old_source)
+        speech, source = self._hift_inference(mel, old_source)
         if old_speech.shape[-1] > 0:
             window = self.speech_window.to(device=speech.device, dtype=speech.dtype)
             speech = self._fade_in_out(speech, old_speech, window)

vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py

@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from torch.cuda import CUDAGraph
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+
+class HiFTGraphWrapper:
+    def __init__(self, token2wav, connector_config, capture_batch_sizes):
+        self.decode_fn = token2wav.hift.inference
+        self.graph_fn = token2wav.hift._inference_pre_istft
+        self.finalize_fn = token2wav.hift._finalize_decode
+        self.codec_chunk_frames = connector_config["codec_chunk_frames"]
+        self.codec_left_context_frames = connector_config["codec_left_context_frames"]
+        lookahead_layer = getattr(token2wav.flow.encoder, "pre_lookahead_layer", None)
+        pre_lookahead_len = getattr(lookahead_layer, "pre_lookahead_len", None)
+        self.pre_lookahead_len = int(pre_lookahead_len) if pre_lookahead_len is not None else 3
+        self.mel_cache_len = int(token2wav.mel_cache_len)
+        self.source_cache_len = int(token2wav.source_cache_len)
+        self.mel_frames = int(token2wav.hift.conv_pre.in_channels)
+        self.flow_upsample_rate = int(getattr(token2wav.flow, "token_mel_ratio", 2))
+        self.capture_bucket_size, self.capture_source_cache_len = self.derive_capture_bucket_size()
+        self.capture_batch_sizes = capture_batch_sizes
+        self.graph: dict[tuple[int, int, int], torch.cuda.CUDAGraph] = {}
+        self.static_speech_inputs: dict[tuple[int, int, int], torch.Tensor] = {}
+        self.static_magnitude_outputs: dict[tuple[int, int, int], torch.Tensor] = {}
+        self.static_phase_outputs: dict[tuple[int, int, int], torch.Tensor] = {}
+        self.static_cache_source_inputs: dict[tuple[int, int, int], torch.Tensor] = {}
+        self.static_cache_source_outputs: dict[tuple[int, int, int], torch.Tensor] = {}
+        parameter = next(token2wav.hift.parameters())
+        self.device = parameter.device
+        self.dtype = parameter.dtype
+        self.max_lazy_graphs = 8
+        self.lazy_graph_count = 0
+
+    def derive_capture_bucket_size(self):
+        chunk_mel_frames = (
+            self.codec_chunk_frames + self.codec_left_context_frames - self.pre_lookahead_len
+        ) * self.flow_upsample_rate
+
+        return [chunk_mel_frames, chunk_mel_frames + self.mel_cache_len], [
+            0,
+            self.source_cache_len,
+        ]
+
+    def capture(self):
+        for batch_size in self.capture_batch_sizes:
+            for mel_frames, source_cache_len in zip(
+                self.capture_bucket_size,
+                self.capture_source_cache_len,
+                strict=True,
+            ):
+                self._capture(batch_size, mel_frames, source_cache_len)
+
+    def _capture(
+        self,
+        batch_size: int,
+        mel_frames: int,
+        source_cache_len: int,
+    ):
+        if torch.cuda.is_current_stream_capturing():
+            raise RuntimeError("Cannot capture HiFT graph during an active stream capture")
+
+        key = (batch_size, mel_frames, source_cache_len)
+
+        if key in self.graph:
+            return
+
+        static_mel = torch.zeros(batch_size, self.mel_frames, mel_frames, device=self.device, dtype=self.dtype)
+        static_source_cache = torch.zeros(batch_size, 1, source_cache_len, device=self.device, dtype=self.dtype)
+        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):
+                warmup_outputs = self.graph_fn(static_mel, static_source_cache)
+        current_stream.wait_stream(warmup_stream)
+        del warmup_outputs
+
+        graph = CUDAGraph()
+        with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
+            static_magnitude_output, static_phase_output, static_cache_source_output = self.graph_fn(
+                static_mel,
+                static_source_cache,
+            )
+
+        self.graph[key] = graph
+        self.static_speech_inputs[key] = static_mel
+        self.static_cache_source_inputs[key] = static_source_cache
+
+        self.static_magnitude_outputs[key] = static_magnitude_output
+        self.static_phase_outputs[key] = static_phase_output
+        self.static_cache_source_outputs[key] = static_cache_source_output
+        logger.info("Captured HiFT CUDA Graph for shape %s", key)
+
+    def replay(self, speech_feat, cache_source):
+        if torch.cuda.is_current_stream_capturing():
+            logger.info("Falling back to eager HiFT inference during an active stream capture")
+            return self.decode_fn(speech_feat, cache_source)
+
+        batch_size = speech_feat.shape[0]
+        num_frames = speech_feat.shape[2]
+        cache_source_len = cache_source.shape[2]
+        target_b = next((b for b in sorted(self.capture_batch_sizes) if b >= batch_size), None)
+
+        if target_b is None:
+            logger.info("Falling back to eager HiFT inference for unsupported batch size %d", batch_size)
+            return self.decode_fn(speech_feat, cache_source)
+
+        key = (target_b, num_frames, cache_source_len)
+
+        if key not in self.graph:
+            if self.lazy_graph_count >= self.max_lazy_graphs:
+                logger.info("Falling back to eager HiFT inference after reaching the lazy Graph limit")
+                return self.decode_fn(speech_feat, cache_source)
+            logger.info("Lazily capturing HiFT CUDA Graph for shape %s", key)
+            self._capture(*key)
+            self.lazy_graph_count += 1
+
+        static_speech_inputs = self.static_speech_inputs[key].zero_()
+        static_speech_inputs[:batch_size].copy_(speech_feat)
+        static_cache_sources = self.static_cache_source_inputs[key].zero_()
+        static_cache_sources[:batch_size].copy_(cache_source)
+
+        self.graph[key].replay()
+        static_magnitude_output = self.static_magnitude_outputs[key]
+        static_phase_output = self.static_phase_outputs[key]
+        static_cache_source_output = self.static_cache_source_outputs[key]
+        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

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py

@@ -150,6 +150,18 @@ class MiniCPMO45Code2Wav(nn.Module):
             prefix="minicpmo45-runtime-prompts-",
         )
         extra = self._extra_config()
+        self._connector_config = {
+            "codec_chunk_frames": int(extra.get("codec_chunk_frames", 25)),
+            "codec_left_context_frames": int(extra.get("codec_left_context_frames", 3)),
+        }
+        if self._connector_config["codec_chunk_frames"] <= 0 or self._connector_config["codec_left_context_frames"] < 0:
+            raise ValueError(f"Invalid MiniCPM-o connector chunk configuration: {self._connector_config}")
+        raw_capture_batch_sizes = extra.get("hift_graph_capture_batch_sizes")
+        capture_batch_sizes = [1] if raw_capture_batch_sizes is None else raw_capture_batch_sizes
+        self._hift_graph_config = {
+            "enabled": bool(extra.get("enable_hift_graph", False)),
+            "capture_batch_sizes": capture_batch_sizes,
+        }
         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")
@@ -785,4 +797,9 @@ class MiniCPMO45Code2Wav(nn.Module):
             logger.info("MiniCPM-o Code2Wav: DiT estimator running on TensorRT (%s)", dtype_name)
             token2wav.enable_trt_spk_embedding()
             logger.info("MiniCPM-o Code2Wav: campplus speaker embedding running on TensorRT")
-        self.backend = BatchedToken2Wav(token2wav, trt_stepper=trt_stepper)
+        self.backend = BatchedToken2Wav(
+            token2wav,
+            trt_stepper=trt_stepper,
+            connector_config=self._connector_config,
+            hift_graph_config=self._hift_graph_config,
+        )

vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py

@@ -13,9 +13,12 @@ import torch
 import torch.nn as nn
 import torchaudio
 import torchaudio.compliance.kaldi as kaldi
-from flashcosyvoice.modules.hifigan import HiFTGenerator
-from flashcosyvoice.utils.audio import mel_spectrogram
 from hyperpyyaml import load_hyperpyyaml
+
+try:
+    from torch.nn.utils.parametrizations import weight_norm
+except ImportError:
+    from torch.nn.utils import weight_norm
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import SupportsPP
@@ -25,6 +28,8 @@ from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
+from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import HiFTGenerator
+from vllm_omni.model_executor.models.cosyvoice3.utils import mel_spectrogram
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import (
     DEFAULT_STREAM_CONFIG,
@@ -66,6 +71,30 @@ class _StreamState:
     finished: bool = False
 
 
+class _ConvRNNF0Predictor(nn.Module):
+    def __init__(self):
+        super().__init__()
+        layers: list[nn.Module] = []
+        for in_channels in [80, 512, 512, 512, 512]:
+            layers.extend([weight_norm(nn.Conv1d(in_channels, 512, 3, padding=1)), nn.ELU()])
+        self.condnet = nn.Sequential(*layers)
+        self.classifier = nn.Linear(512, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.classifier(self.condnet(x).transpose(1, 2)).squeeze(-1).abs()
+
+
+def _build_hift() -> HiFTGenerator:
+    return HiFTGenerator(
+        sampling_rate=24000,
+        upsample_rates=[8, 5, 3],
+        upsample_kernel_sizes=[16, 11, 7],
+        source_resblock_kernel_sizes=[7, 7, 11],
+        source_resblock_dilation_sizes=[[1, 3, 5]] * 3,
+        f0_predictor=_ConvRNNF0Predictor(),
+    )
+
+
 class StepAudio2Token2WavCore(nn.Module):
     """Core Token2Wav model - handles token to waveform conversion"""
 
@@ -142,8 +171,8 @@ class StepAudio2Token2WavCore(nn.Module):
             torch.load(f"{self.model_path}/flow.pt", map_location="cpu", weights_only=True), strict=True
         )
         self._flow.to(self.device).eval()
+        self._hift = _build_hift()
 
-        self._hift = HiFTGenerator()
         hift_state_dict = {
             k.replace("generator.", ""): v
             for k, v in torch.load(f"{self.model_path}/hift.pt", map_location="cpu", weights_only=True).items()
@@ -232,7 +261,21 @@ class StepAudio2Token2WavCore(nn.Module):
         if sample_rate != 24000:
             audio_24k = librosa.resample(y=audio_24k, orig_sr=sample_rate, target_sr=24000)
         audio_24k_t = torch.from_numpy(audio_24k).unsqueeze(0)  # [1, T]
-        prompt_mel = mel_spectrogram(audio_24k_t).transpose(1, 2).squeeze(0)  # [T, num_mels]
+        prompt_mel = (
+            mel_spectrogram(
+                audio_24k_t,
+                n_fft=1920,
+                num_mels=80,
+                sampling_rate=24000,
+                hop_size=480,
+                win_size=1920,
+                fmin=0,
+                fmax=8000,
+                center=False,
+            )
+            .transpose(1, 2)
+            .squeeze(0)
+        )  # [T, num_mels]
         prompt_mels = prompt_mel.unsqueeze(0).to(self.device)
         prompt_mels_lens = torch.tensor([prompt_mels.shape[1]], dtype=torch.int32, device=self.device)
         prompt_mels = torch.nn.functional.pad(
@@ -278,7 +321,7 @@ class StepAudio2Token2WavCore(nn.Module):
                 self.n_timesteps,
             )
 
-        wav, _ = self.hift(speech_feat=mel)
+        wav, _ = self.hift.inference(speech_feat=mel)
 
         if return_bytes:
             output = io.BytesIO()
@@ -366,7 +409,7 @@ class StepAudio2Token2WavCore(nn.Module):
         hift_cache_speech = state.hift_cache_dict["speech"]
 
         mel = torch.cat([hift_cache_mel, chunk_mel], dim=2)
-        speech, source = self.hift(mel, hift_cache_source)
+        speech, source = self.hift.inference(mel, hift_cache_source)
 
         if hift_cache_speech.shape[-1] > 0:
             speech = fade_in_out(speech, hift_cache_speech, self.speech_window)