From fc8946fcf542cf108fa8b2fdc69de737bbfd52d8 Mon Sep 17 00:00:00 2001 From: boatman <1930807094@qq.com> Date: Sun, 9 Aug 2026 22:29:49 +0800 Subject: [PATCH] [Perf][MiniCPM-O-4.5]Add CUDA Graph For HiFTGenerator (#5869) Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com> --- .../minicpmo_4_5/test_code2wav_batching.py | 3 + .../minicpmo_4_5/test_cuda_graph_wrapper.py | 144 ++++++++++++++++++ .../models/minicpmo_4_5/test_pipeline.py | 1 + .../models/step_audio2/test_hift_parity.py | 54 +++++++ vllm_omni/deploy/minicpmo_4_5.yaml | 5 + vllm_omni/deploy/minicpmo_4_5_2gpu.yaml | 5 + vllm_omni/deploy/minicpmo_4_5_3gpu.yaml | 5 + vllm_omni/deploy/minicpmo_4_5_8x4090.yaml | 5 + .../cosyvoice3/code2wav_core/hifigan.py | 43 +++++- .../models/minicpmo_4_5/batched_token2wav.py | 52 ++++++- .../models/minicpmo_4_5/cuda_graph_wrapper.py | 135 ++++++++++++++++ .../minicpmo_4_5/minicpmo_4_5_code2wav.py | 19 ++- .../step_audio2/step_audio2_token2wav.py | 55 ++++++- 13 files changed, 509 insertions(+), 17 deletions(-) create mode 100644 tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py create mode 100644 tests/model_executor/models/step_audio2/test_hift_parity.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py index c47aa8bbf..ba3c7292b 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py +++ b/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): 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 new file mode 100644 index 000000000..16e04fc09 --- /dev/null +++ b/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 diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index 5236ea883..6b2ead7d1 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/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" diff --git a/tests/model_executor/models/step_audio2/test_hift_parity.py b/tests/model_executor/models/step_audio2/test_hift_parity.py new file mode 100644 index 000000000..d256f62df --- /dev/null +++ b/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) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 66a5f3ac9..c98e85f3a 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/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 diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml index 75693feb7..e44140b2d 100644 --- a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml +++ b/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 diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml index 9254d3c70..71d4833a0 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml +++ b/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 diff --git a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml index 1b38059a3..dfc6499c3 100644 --- a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml +++ b/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 diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py index 63434d57f..d8d99aa21 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py +++ b/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 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 5ffde9e15..651284fe0 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 @@ -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) 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 new file mode 100644 index 000000000..5f6c965d5 --- /dev/null +++ b/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 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 019ab0594..d6d3e3d47 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 @@ -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, + ) diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py index 3923933f6..b116b6bcd 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py +++ b/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)