From f420cfcfbf5d7dff2ed471305e0ef145dc32f956 Mon Sep 17 00:00:00 2001 From: ChangXiong Deng <101780062+peterDengcx@users.noreply.github.com> Date: Wed, 16 Sep 2026 09:06:41 +0800 Subject: [PATCH] [NPU][CosyVoice3] Add CPU STFT fallback (#7520) Signed-off-by: dengchangxiong Co-authored-by: dengchangxiong --- .../cosyvoice3/test_cosyvoice3_components.py | 68 +++++++++++++++++++ .../cosyvoice3/code2wav_core/hifigan.py | 42 ++++++++++-- 2 files changed, 106 insertions(+), 4 deletions(-) diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py index 1b1df6367..7dfbe6e74 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py @@ -76,6 +76,74 @@ def test_causal_hift_stft_moves_window_to_input_device(causal_hift): assert causal_hift.stft_window.device == waveform.device +def test_causal_hift_routes_npu_transforms_to_cpu_fallback(causal_hift, monkeypatch): + class NpuTensor: + device = SimpleNamespace(type="npu") + + tensor = NpuTensor() + stft_calls = [] + istft_calls = [] + monkeypatch.setattr( + causal_hift, + "_stft_on_cpu", + lambda value: stft_calls.append(value) or ("real", "imag"), + ) + monkeypatch.setattr( + causal_hift, + "_istft_on_cpu", + lambda magnitude, phase: istft_calls.append((magnitude, phase)) or "waveform", + ) + + assert causal_hift._stft(tensor) == ("real", "imag") + assert causal_hift._istft(tensor, tensor) == "waveform" + assert stft_calls == [tensor] + assert istft_calls == [(tensor, tensor)] + + +def test_causal_hift_stft_cpu_fallback_uses_float32(causal_hift, monkeypatch): + original_stft = torch.stft + + def assert_cpu_float32_stft(input, *args, window, **kwargs): + assert input.device.type == "cpu" + assert input.dtype == torch.float32 + assert window.device.type == "cpu" + assert window.dtype == torch.float32 + return original_stft(input, *args, window=window, **kwargs) + + monkeypatch.setattr(torch, "stft", assert_cpu_float32_stft) + + waveform = torch.randn(1, 64, dtype=torch.float64) + real, imag = causal_hift._stft_on_cpu(waveform) + + assert real.device == waveform.device + assert imag.device == waveform.device + assert real.dtype == waveform.dtype + assert imag.dtype == waveform.dtype + + +def test_causal_hift_istft_cpu_fallback_uses_float32(causal_hift, monkeypatch): + original_istft = torch.istft + + def assert_cpu_complex64_istft(input, *args, window, **kwargs): + assert input.device.type == "cpu" + assert input.dtype == torch.complex64 + assert window.device.type == "cpu" + assert window.dtype == torch.float32 + return original_istft(input, *args, window=window, **kwargs) + + monkeypatch.setattr(torch, "istft", assert_cpu_complex64_istft) + + waveform = torch.randn(1, 64, dtype=torch.float32) + spec = torch.stft(waveform, 16, 4, 16, window=causal_hift.stft_window, return_complex=True) + magnitude = torch.abs(spec).to(torch.float64) + phase = torch.angle(spec).to(torch.float64) + reconstructed = causal_hift._istft_on_cpu(magnitude, phase) + + assert reconstructed.device == magnitude.device + assert reconstructed.dtype == magnitude.dtype + torch.testing.assert_close(reconstructed, waveform.to(torch.float64), rtol=1e-5, atol=1e-5) + + class TestPreLookaheadLayer: """Tests for PreLookaheadLayer.""" 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 2a5d472a0..7cbd31e9c 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py @@ -489,6 +489,9 @@ class HiFTGenerator(nn.Module): block.remove_weight_norm() def _stft(self, x): + if x.device.type == "npu": + return self._stft_on_cpu(x) + spec = torch.stft( x, self.istft_params["n_fft"], @@ -500,18 +503,49 @@ class HiFTGenerator(nn.Module): spec = torch.view_as_real(spec) # [B, F, TT, 2] return spec[..., 0], spec[..., 1] + def _stft_on_cpu(self, x): + target_device = x.device + target_dtype = x.dtype + spec = torch.stft( + x.float().cpu(), + self.istft_params["n_fft"], + self.istft_params["hop_len"], + self.istft_params["n_fft"], + window=self.stft_window.float().cpu(), + return_complex=True, + ) + spec = torch.view_as_real(spec).to(device=target_device, dtype=target_dtype) + return spec[..., 0], spec[..., 1] + def _istft(self, magnitude, phase): + if magnitude.device.type == "npu": + return self._istft_on_cpu(magnitude, phase) + magnitude = torch.clip(magnitude, max=1e2) real = magnitude * torch.cos(phase) - img = magnitude * torch.sin(phase) - inverse_transform = torch.istft( - torch.complex(real, img), + imag = magnitude * torch.sin(phase) + return torch.istft( + torch.complex(real, imag), self.istft_params["n_fft"], self.istft_params["hop_len"], self.istft_params["n_fft"], window=self._get_stft_window(magnitude), ) - return inverse_transform + + def _istft_on_cpu(self, magnitude, phase): + target_device = magnitude.device + target_dtype = magnitude.dtype + magnitude = torch.clip(magnitude, max=1e2) + real = magnitude * torch.cos(phase) + imag = magnitude * torch.sin(phase) + waveform = torch.istft( + torch.complex(real.float().cpu(), imag.float().cpu()), + self.istft_params["n_fft"], + self.istft_params["hop_len"], + self.istft_params["n_fft"], + window=self.stft_window.float().cpu(), + ) + return waveform.to(device=target_device, dtype=target_dtype) def _get_stft_window(self, tensor: torch.Tensor) -> torch.Tensor: if self.stft_window.device != tensor.device: