diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 7ad153035..dc87aa7be 100644 --- a/tests/config/test_environment_variables.py +++ b/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, } diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components_cuda.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components_cuda.py index aff89f852..0e936bc0e 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components_cuda.py +++ b/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) diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 7e206ebb2..ccdc2133e 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/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", 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 2b8f3e61e..f3ac19d94 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py +++ b/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)