feat(cosyvoice3): GPU-resident F0 predictor & pure-PyTorch vocoder (Task C2) (#7518)

Signed-off-by: BeatSeat <wendavid552@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/config/test_environment_variables.py

@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         EnvironmentVariableCategory.PUBLIC_OMNI: 31,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 58,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 59,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
         EnvironmentVariableCategory.INTERNAL: 2,
     }
@@ -181,7 +181,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         ModelEnvironmentVariableDisposition.PROMOTE: 31,
         ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
         ModelEnvironmentVariableDisposition.EXTERNAL: 0,
-        ModelEnvironmentVariableDisposition.INTERNALIZE: 15,
+        ModelEnvironmentVariableDisposition.INTERNALIZE: 16,
         ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 7,
     }
 

tests/model_executor/models/cosyvoice3/test_cosyvoice3_components_cuda.py

@@ -125,3 +125,96 @@ def test_code2wav_streaming_batch_matches_ragged_flow_numerics(monkeypatch):
         assert batched_mel.shape == individual_mel.shape
         rel_mean = (batched_mel - individual_mel).abs().mean() / individual_mel.abs().mean().clamp_min(1e-6)
         assert rel_mean.item() < 0.05
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_causal_hift_f0_predictor_cuda_parity_and_fallback(monkeypatch):
+    """Test GPU-resident F0 predictor on CUDA with streaming/finalize, TF32 restoration,
+    parity against CPU reference, and COSYVOICE3_F0_ON_CPU fallback."""
+    from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
+        CausalConvRNNF0Predictor,
+        CausalHiFTGenerator,
+    )
+
+    torch.manual_seed(42)
+    f0_predictor = CausalConvRNNF0Predictor(num_class=1, in_channels=80, cond_channels=16)
+    hift = CausalHiFTGenerator(
+        in_channels=80,
+        base_channels=32,
+        nb_harmonics=8,
+        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], [1, 3, 5], [1, 3, 5]],
+        resblock_kernel_sizes=[3, 7, 11],
+        resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+        f0_predictor=f0_predictor,
+    ).eval()
+
+    with torch.no_grad():
+        # Force voiced F0 above nsf_voiced_threshold (10.0)
+        hift.f0_predictor.classifier.bias.fill_(150.0)
+
+    # 1. Compute CPU reference for both streaming (finalize=False) and finalize=True
+    speech_feat_cpu = torch.randn(1, 80, 50, dtype=torch.float32)
+    with torch.no_grad():
+        f0_ref_streaming = hift.f0_predictor(speech_feat_cpu, finalize=False)
+        f0_ref_finalize = hift.f0_predictor(speech_feat_cpu, finalize=True)
+
+    # Move HiFT to CUDA
+    hift = hift.to("cuda")
+    speech_feat_cuda = speech_feat_cpu.to("cuda")
+
+    # Set TF32 to True initially to verify `finally` restoration
+    monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", True)
+    monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", True)
+    monkeypatch.delenv("COSYVOICE3_F0_ON_CPU", raising=False)
+
+    # 2. Test CUDA streaming execution (finalize=False)
+    wav_stream, s_stream, phase_acc = hift.inference(speech_feat_cuda, finalize=False)
+    assert wav_stream.device.type == "cuda"
+    assert s_stream.device.type == "cuda"
+    assert next(hift.f0_predictor.parameters()).device.type == "cuda"
+    assert torch.backends.cuda.matmul.allow_tf32 is True
+    assert torch.backends.cudnn.allow_tf32 is True
+
+    # Check streaming F0 parity against CPU reference under explicit tolerances
+    with torch.no_grad():
+        f0_cuda_streaming = hift.f0_predictor(speech_feat_cuda, finalize=False)
+    torch.testing.assert_close(f0_cuda_streaming.cpu(), f0_ref_streaming, atol=1e-3, rtol=1e-3)
+    cos_sim_stream = torch.nn.functional.cosine_similarity(
+        f0_cuda_streaming.cpu().flatten(1), f0_ref_streaming.flatten(1), dim=1
+    ).item()
+    assert cos_sim_stream > 0.9999
+
+    # 3. Test CUDA finalize execution (finalize=True)
+    wav_final, s_final, _ = hift.inference(speech_feat_cuda, finalize=True, phase_acc=phase_acc)
+    assert wav_final.device.type == "cuda"
+    assert s_final.device.type == "cuda"
+    assert next(hift.f0_predictor.parameters()).device.type == "cuda"
+    assert torch.backends.cuda.matmul.allow_tf32 is True
+    assert torch.backends.cudnn.allow_tf32 is True
+
+    # Check finalize F0 parity against CPU reference under explicit tolerances
+    with torch.no_grad():
+        f0_cuda_finalize = hift.f0_predictor(speech_feat_cuda, finalize=True)
+    torch.testing.assert_close(f0_cuda_finalize.cpu(), f0_ref_finalize, atol=1e-3, rtol=1e-3)
+    cos_sim_final = torch.nn.functional.cosine_similarity(
+        f0_cuda_finalize.cpu().flatten(1), f0_ref_finalize.flatten(1), dim=1
+    ).item()
+    assert cos_sim_final > 0.9999
+
+    # 4. Test COSYVOICE3_F0_ON_CPU=1 fallback path
+    monkeypatch.setenv("COSYVOICE3_F0_ON_CPU", "1")
+    wav_fallback, s_fallback, _ = hift.inference(speech_feat_cuda, finalize=True, phase_acc=phase_acc)
+    assert wav_fallback.device.type == "cuda"
+    assert s_fallback.device.type == "cuda"
+    # Predictor should have been moved to CPU
+    assert next(hift.f0_predictor.parameters()).device.type == "cpu"
+    assert torch.backends.cuda.matmul.allow_tf32 is True
+    assert torch.backends.cudnn.allow_tf32 is True
+
+    # Fallback output should match within numerical tolerance
+    torch.testing.assert_close(s_final, s_fallback, atol=1e-2, rtol=1e-2)
+    torch.testing.assert_close(wav_final, wav_fallback, atol=1e-2, rtol=1e-2)

vllm_omni/config/environment_variable_inventory.py

@@ -191,6 +191,7 @@ _MODEL_EXTERNAL: tuple[str, ...] = ()
 
 _MODEL_INTERNALIZE = (
     "COSYVOICE3_BATCH_FLOW_DEBUG",
+    "COSYVOICE3_F0_ON_CPU",
     "DZ_PHASE_TIMING",
     "MAGI2_ALLOW_UNSUPPORTED_TOPOLOGY",
     "MAGI2_ROUTER_BIAS_SOURCE",

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

@@ -6,11 +6,13 @@
 
 """HIFI-GAN"""
 
+import math
+import os
+
 import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from scipy.signal import get_window
 from torch.nn import Conv1d, ConvTranspose1d
 from torch.nn.utils import remove_weight_norm
 
@@ -191,7 +193,7 @@ class SineGen(torch.nn.Module):
         before its own cumsum -- callers must scale both by samples-per-mel.
         """
         f0 = f0.transpose(1, 2)
-        F_mat = torch.zeros((f0.size(0), self.harmonic_num + 1, f0.size(-1))).to(f0.device)
+        F_mat = f0.new_zeros((f0.size(0), self.harmonic_num + 1, f0.size(-1)))
         for i in range(self.harmonic_num + 1):
             F_mat[:, i : i + 1, :] = f0 * (i + 1) / self.sampling_rate
 
@@ -507,9 +509,10 @@ class HiFTGenerator(nn.Module):
         self.num_kernels = len(resblock_kernel_sizes)
         self.num_upsamples = len(upsample_rates)
         # NOTE in CosyVoice2, we use the original SineGen implementation
+        scale = int(math.prod(upsample_rates) * istft_params["hop_len"])
         self.m_source = SourceModuleHnNSF(
             sampling_rate=sampling_rate,
-            upsample_scale=np.prod(upsample_rates) * istft_params["hop_len"],
+            upsample_scale=scale,
             harmonic_num=nb_harmonics,
             sine_amp=nsf_alpha,
             add_noise_std=nsf_sigma,
@@ -517,7 +520,7 @@ class HiFTGenerator(nn.Module):
             sinegen_type="1" if self.sampling_rate == 22050 else "2",
             causal=False,
         )
-        self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates) * istft_params["hop_len"])
+        self.f0_upsamp = torch.nn.Upsample(scale_factor=scale)
 
         self.conv_pre = weight_norm(Conv1d(in_channels, base_channels, 7, 1, padding=3))
 
@@ -540,7 +543,11 @@ class HiFTGenerator(nn.Module):
         self.source_downs = nn.ModuleList()
         self.source_resblocks = nn.ModuleList()
         downsample_rates = [1] + upsample_rates[::-1][:-1]
-        downsample_cum_rates = np.cumprod(downsample_rates)
+        downsample_cum_rates = []
+        cur = 1
+        for r in downsample_rates:
+            cur *= r
+            downsample_cum_rates.append(cur)
         for i, (u, k, d) in enumerate(
             zip(downsample_cum_rates[::-1], source_resblock_kernel_sizes, source_resblock_dilation_sizes)
         ):
@@ -565,7 +572,7 @@ class HiFTGenerator(nn.Module):
         self.reflection_pad = nn.ReflectionPad1d((1, 0))
         self.register_buffer(
             "stft_window",
-            torch.from_numpy(get_window("hann", istft_params["n_fft"], fftbins=True).astype(np.float32)),
+            torch.hann_window(istft_params["n_fft"], periodic=True, dtype=torch.float32),
             persistent=False,
         )
         self.f0_predictor = f0_predictor
@@ -770,9 +777,10 @@ class CausalHiFTGenerator(HiFTGenerator):
 
         self.num_kernels = len(resblock_kernel_sizes)
         self.num_upsamples = len(upsample_rates)
+        scale = int(math.prod(upsample_rates) * istft_params["hop_len"])
         self.m_source = SourceModuleHnNSF(
             sampling_rate=sampling_rate,
-            upsample_scale=np.prod(upsample_rates) * istft_params["hop_len"],
+            upsample_scale=scale,
             harmonic_num=nb_harmonics,
             sine_amp=nsf_alpha,
             add_noise_std=nsf_sigma,
@@ -781,7 +789,7 @@ class CausalHiFTGenerator(HiFTGenerator):
             causal=True,
         )
         self.upsample_rates = upsample_rates
-        self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates) * istft_params["hop_len"])
+        self.f0_upsamp = torch.nn.Upsample(scale_factor=scale)
 
         self.conv_pre = weight_norm(
             CausalConv1d(in_channels, base_channels, conv_pre_look_right + 1, 1, causal_type="right")
@@ -805,7 +813,11 @@ class CausalHiFTGenerator(HiFTGenerator):
         self.source_downs = nn.ModuleList()
         self.source_resblocks = nn.ModuleList()
         downsample_rates = [1] + upsample_rates[::-1][:-1]
-        downsample_cum_rates = np.cumprod(downsample_rates)
+        downsample_cum_rates = []
+        cur = 1
+        for r in downsample_rates:
+            cur *= r
+            downsample_cum_rates.append(cur)
         for i, (u, k, d) in enumerate(
             zip(downsample_cum_rates[::-1], source_resblock_kernel_sizes, source_resblock_dilation_sizes)
         ):
@@ -832,7 +844,7 @@ class CausalHiFTGenerator(HiFTGenerator):
         self.reflection_pad = nn.ReflectionPad1d((1, 0))
         self.register_buffer(
             "stft_window",
-            torch.from_numpy(get_window("hann", istft_params["n_fft"], fftbins=True).astype(np.float32)),
+            torch.hann_window(istft_params["n_fft"], periodic=True, dtype=torch.float32),
             persistent=False,
         )
         self.conv_pre_look_right = conv_pre_look_right
@@ -844,8 +856,8 @@ class CausalHiFTGenerator(HiFTGenerator):
             x = self.conv_pre(x)
         else:
             x = self.conv_pre(x[:, :, : -self.conv_pre_look_right], x[:, :, -self.conv_pre_look_right :])
-            s_stft_real = s_stft_real[:, :, : -int(np.prod(self.upsample_rates) * self.conv_pre_look_right)]
-            s_stft_imag = s_stft_imag[:, :, : -int(np.prod(self.upsample_rates) * self.conv_pre_look_right)]
+            s_stft_real = s_stft_real[:, :, : -int(math.prod(self.upsample_rates) * self.conv_pre_look_right)]
+            s_stft_imag = s_stft_imag[:, :, : -int(math.prod(self.upsample_rates) * self.conv_pre_look_right)]
         s_stft = torch.cat([s_stft_real, s_stft_imag], dim=1)
 
         for i in range(self.num_upsamples):
@@ -875,7 +887,7 @@ class CausalHiFTGenerator(HiFTGenerator):
 
         x = self._istft(magnitude, phase)
         if finalize is False:
-            x = x[:, : -int(np.prod(self.upsample_rates) * self.istft_params["hop_len"])]
+            x = x[:, : -int(math.prod(self.upsample_rates) * self.istft_params["hop_len"])]
         x = torch.clamp(x, -self.audio_limit, self.audio_limit)
         return x
 
@@ -890,10 +902,30 @@ class CausalHiFTGenerator(HiFTGenerator):
         trim: int = 0,
         f0_margin: int = 0,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
-        # mel->f0 NOTE f0_predictor precision is crucial for causal inference, move
-        # self.f0_predictor to cpu if necessary
-        self.f0_predictor.to("cpu")
-        f0 = self.f0_predictor(speech_feat.cpu(), finalize=finalize).to(speech_feat)
+        # mel->f0
+        # By default keep f0_predictor on the same device as speech_feat (GPU) in FP32
+        # with TF32 disabled to prevent causal drift while avoiding CPU syncs and transfers.
+        # COSYVOICE3_F0_ON_CPU=1 forces CPU reference execution.
+        f0_on_cpu = os.getenv("COSYVOICE3_F0_ON_CPU", "0") == "1"
+        if speech_feat.device.type == "cuda" and not f0_on_cpu:
+            p = next(self.f0_predictor.parameters(), None)
+            if p is not None and p.device != speech_feat.device:
+                self.f0_predictor.to(device=speech_feat.device, dtype=torch.float32)
+            orig_matmul_tf32 = torch.backends.cuda.matmul.allow_tf32
+            orig_cudnn_tf32 = torch.backends.cudnn.allow_tf32
+            try:
+                torch.backends.cuda.matmul.allow_tf32 = False
+                torch.backends.cudnn.allow_tf32 = False
+                f0 = self.f0_predictor(speech_feat.float(), finalize=finalize).to(dtype=speech_feat.dtype)
+            finally:
+                torch.backends.cuda.matmul.allow_tf32 = orig_matmul_tf32
+                torch.backends.cudnn.allow_tf32 = orig_cudnn_tf32
+        else:
+            p = next(self.f0_predictor.parameters(), None)
+            if p is not None and p.device.type != "cpu":
+                self.f0_predictor.to("cpu")
+            f0 = self.f0_predictor(speech_feat.cpu().float(), finalize=finalize).to(speech_feat)
+
         if f0_margin > 0:
             f0 = f0[:, f0_margin:]
             speech_feat = speech_feat[:, :, f0_margin:]
@@ -1031,10 +1063,10 @@ class CausalConv1d(torch.nn.Conv1d):
         assert causal_type in ["left", "right"]
         self.causal_type = causal_type
 
-    def forward(self, x: torch.Tensor, cache: torch.Tensor = torch.zeros(0, 0, 0)) -> tuple[torch.Tensor]:
+    def forward(self, x: torch.Tensor, cache: torch.Tensor | None = None) -> torch.Tensor:
         input_timestep = x.shape[2]
-        if cache.size(2) == 0:
-            cache = torch.zeros(x.shape[0], x.shape[1], self.causal_padding).to(x)
+        if cache is None or cache.size(2) == 0:
+            cache = x.new_zeros(x.shape[0], x.shape[1], self.causal_padding)
         assert cache.size(2) == self.causal_padding
         if self.causal_type == "left":
             x = torch.concat([cache, x], dim=2)