diff --git a/tests/platforms/npu/test_step_audio2_token2wav.py b/tests/platforms/npu/test_step_audio2_token2wav.py new file mode 100644 index 000000000..ae86f92d9 --- /dev/null +++ b/tests/platforms/npu/test_step_audio2_token2wav.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU unit tests for the Step-Audio2 Ascend HiFT patch.""" + +from __future__ import annotations + +import importlib.util +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +import torch.nn.functional as F + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _load_patch_module(monkeypatch: pytest.MonkeyPatch): + fake_logger = SimpleNamespace(info=lambda *_args, **_kwargs: None) + fake_vllm = types.ModuleType("vllm") + fake_vllm_logger = types.ModuleType("vllm.logger") + fake_vllm_logger.init_logger = lambda _name: fake_logger + monkeypatch.setitem(sys.modules, "vllm", fake_vllm) + monkeypatch.setitem(sys.modules, "vllm.logger", fake_vllm_logger) + + root = next(parent for parent in Path(__file__).resolve().parents if (parent / "vllm_omni").is_dir()) + path = root / "vllm_omni" / "platforms" / "npu" / "models" / "step_audio2_token2wav.py" + spec = importlib.util.spec_from_file_location("test_step_audio2_token2wav_npu_patch", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _reference_f02sine(sine_gen, f0_values: torch.Tensor) -> torch.Tensor: + rad_values = (f0_values / sine_gen.sampling_rate) % 1 + rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], device=f0_values.device) + rand_ini[:, 0] = 0 + rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini + rad_values = F.interpolate( + rad_values.transpose(1, 2), + scale_factor=1 / sine_gen.upsample_scale, + mode="linear", + ).transpose(1, 2) + phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi + phase = F.interpolate( + phase.transpose(1, 2) * sine_gen.upsample_scale, + scale_factor=sine_gen.upsample_scale, + mode="linear", + ).transpose(1, 2) + return torch.sin(phase) + + +def test_even_scale_downsample_matches_linear_interpolate(monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_patch_module(monkeypatch) + x = torch.randn(2, 3, 960) + + expected = F.interpolate(x, scale_factor=1 / 4, mode="linear") + actual = module._linear_downsample_even_scale(x, 4) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize(("shape", "scale"), [((1, 2, 12), 3), ((1, 2, 10), 4)]) +def test_even_scale_downsample_rejects_unsupported_shapes( + monkeypatch: pytest.MonkeyPatch, + shape: tuple[int, ...], + scale: int, +) -> None: + module = _load_patch_module(monkeypatch) + with pytest.raises(ValueError): + module._linear_downsample_even_scale(torch.randn(shape), scale) + + +def test_hift_patch_is_exact_idempotent_and_stays_on_device(monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_patch_module(monkeypatch) + + class FakeSineGen: + sampling_rate = 24000 + upsample_scale = 4 + flag_for_pulse = False + + def __init__(self): + self.original_calls = 0 + + def _f02sine(self, f0_values): + self.original_calls += 1 + return _reference_f02sine(self, f0_values) + + class FakeHiFT: + def __init__(self): + self.m_source = SimpleNamespace(l_sin_gen=FakeSineGen()) + + def to(self, *_args, **_kwargs): + raise AssertionError("HiFT must stay on its accelerator") + + hift = FakeHiFT() + module.patch_step_audio2_hift_for_npu(hift) + patched_method = hift.m_source.l_sin_gen._f02sine + module.patch_step_audio2_hift_for_npu(hift) + assert hift.m_source.l_sin_gen._f02sine is patched_method + + f0_values = torch.rand(1, 16, 3) + torch.manual_seed(7) + expected = _reference_f02sine(hift.m_source.l_sin_gen, f0_values) + torch.manual_seed(7) + actual = patched_method(f0_values) + + assert hift.m_source.l_sin_gen.original_calls == 0 + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize( + ("flag_for_pulse", "upsample_scale", "input_length"), + [ + (True, 4, 16), + (False, 3, 12), + (False, 4.5, 18), + (False, 4, 10), + ], +) +def test_hift_patch_delegates_unsupported_cases_to_original_f02sine_on_cpu( + monkeypatch: pytest.MonkeyPatch, + flag_for_pulse: bool, + upsample_scale: float, + input_length: int, +) -> None: + module = _load_patch_module(monkeypatch) + calls: list[str] = [] + + class FakeSineGen: + def __init__(self): + self.flag_for_pulse = flag_for_pulse + self.upsample_scale = upsample_scale + + def _f02sine(self, f0_values): + calls.append(f0_values.device.type) + return f0_values + 1 + + hift = SimpleNamespace(m_source=SimpleNamespace(l_sin_gen=FakeSineGen())) + module.patch_step_audio2_hift_for_npu(hift) + + f0_values = torch.randn(1, input_length, 3) + output = hift.m_source.l_sin_gen._f02sine(f0_values) + + assert calls == ["cpu"] + assert output.device == f0_values.device + torch.testing.assert_close(output, f0_values + 1) + + +@pytest.mark.parametrize("upsample_scale", [0, -2]) +def test_hift_patch_rejects_non_positive_scale( + monkeypatch: pytest.MonkeyPatch, + upsample_scale: int, +) -> None: + module = _load_patch_module(monkeypatch) + + class FakeSineGen: + flag_for_pulse = False + + def __init__(self): + self.upsample_scale = upsample_scale + + def _f02sine(self, f0_values): + return f0_values + + hift = SimpleNamespace(m_source=SimpleNamespace(l_sin_gen=FakeSineGen())) + module.patch_step_audio2_hift_for_npu(hift) + + with pytest.raises(ValueError, match="upsample_scale must be positive"): + hift.m_source.l_sin_gen._f02sine(torch.randn(1, 16, 3)) + + +def test_hift_patch_rejects_causal_sinegen(monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_patch_module(monkeypatch) + + class FakeCausalSineGen: + causal = True + + def _f02sine(self, f0_values): + return f0_values + + hift = SimpleNamespace(m_source=SimpleNamespace(l_sin_gen=FakeCausalSineGen())) + + with pytest.raises(ValueError, match="only supports non-causal SineGen2"): + module.patch_step_audio2_hift_for_npu(hift) + + +def test_hift_patch_reports_incompatible_layout(monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_patch_module(monkeypatch) + + with pytest.raises(TypeError, match=r"m_source\.l_sin_gen\._f02sine"): + module.patch_step_audio2_hift_for_npu(SimpleNamespace()) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index d14fe14cf..71681c222 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -514,7 +514,7 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): # memory on the NPU sdpa path. Detach the Token2wav vocoder first # so the cast does not drag it onto the accelerator or downcast # its float32 flow/HiFT weights: it manages its own device - # placement (e.g. CPU HiFT on NPU) and may not be an nn.Module. + # placement and may not be an nn.Module. device = current_omni_platform.get_torch_device() audio_tok = getattr(self.tts_obj, "audio_tokenizer", None) if audio_tok is not None: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py index c069b8d09..da868cd58 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py @@ -6,8 +6,8 @@ ``from stepaudio2 import Token2wav`` entry point (``stepaudio2-minicpmo``). That package hard-codes ``.cuda()`` and duplicates the flow/HiFT stack that vLLM-Omni already vendors for Step-Audio2 (which also carries the Ascend/NPU -fixes: CPU HiFT, DiT mask expand, MATH SDPA, compile disable, PE buffer -extension). +fixes: HiFT linear downsample, DiT mask expand, MATH SDPA, compile disable, +PE buffer extension). This module exposes the same call surface MiniCPM expects (``__call__`` / ``set_stream_cache`` / ``stream``) while delegating the 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 1ebad86ce..026276c52 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 @@ -100,8 +100,9 @@ class StepAudio2Token2WavCore(nn.Module): # and apply the Token2Wav NPU patch here (idempotent) before # ``_ensure_models_loaded`` runs. Covers every # construction site (Step-Audio2 wrapper + MiniCPM-o facade): patched - # ``_ensure_models_loaded`` offloads HiFT to CPU and patched ``forward`` - # expands CosyVoice DiT masks + forces MATH SDPA (avoids FA 161001). + # ``_ensure_models_loaded`` replaces HiFT's failing linear downsample, + # while patched ``forward`` expands CosyVoice DiT masks and forces + # MATH SDPA (avoids FA 161001). from vllm_omni.platforms import current_omni_platform if current_omni_platform.is_npu(): diff --git a/vllm_omni/platforms/npu/models/step_audio2_token2wav.py b/vllm_omni/platforms/npu/models/step_audio2_token2wav.py index 04768d04d..1641fee27 100644 --- a/vllm_omni/platforms/npu/models/step_audio2_token2wav.py +++ b/vllm_omni/platforms/npu/models/step_audio2_token2wav.py @@ -4,9 +4,8 @@ Ascend-specific workarounds that must not live in the shared GPU model file: -1. HiFT vocoder CPU offload — STFT/ISTFT and sine-source generation are - unstable on Ascend (aicore 507015). HiFT is tiny, so running it on CPU - is acceptable; inputs/outputs are moved transparently. +1. HiFT sine-source downsample — replace the failing 480x ``linear1d`` + downsample with its exact midpoint form while keeping HiFT on NPU. 2. CosyVoice2 DiT SDPA — force MATH backend (+ DiT attn mask expand) to avoid fused FA rejecting CosyVoice ``(B,1,1,S)`` masks (error 161001). """ @@ -17,7 +16,9 @@ from collections.abc import Iterator from contextlib import contextmanager, nullcontext from types import MethodType +import numpy as np import torch +import torch.nn.functional as F from vllm.logger import init_logger logger = init_logger(__name__) @@ -28,52 +29,87 @@ _original_forward = None _original_stream_chunk_for = None -def patch_hift_module_for_npu(hift: torch.nn.Module) -> None: - """Run the entire HiFT vocoder on CPU when on Ascend NPU. +def _linear_downsample_even_scale(x: torch.Tensor, scale: int) -> torch.Tensor: + """Match ``F.interpolate(..., mode="linear")`` for an even integer scale. - HiFT is tiny, so keeping it on NPU yields negligible speedup while - several of its operators are unstable on Ascend and trigger aicore - exceptions (runtime error 507015): the STFT/ISTFT as well as the - sine-source generation (``f0_upsamp`` + ``m_source`` cumsum/sin over - the full waveform length). Running the whole module on CPU sidesteps - all of them. Inputs are copied to CPU and outputs are moved back to - the original device transparently, so callers stay unchanged. + With ``align_corners=False``, every output location for an even integer + downsample lies exactly halfway between two source samples. Selecting and + averaging those samples avoids Ascend/pytorch#150's ``linear1d`` kernel. """ - if getattr(hift, "_npu_cpu_offload_patched", False): + if scale <= 0 or scale % 2: + raise ValueError(f"scale must be a positive even integer, got {scale}") + if x.shape[-1] % scale: + raise ValueError(f"input length {x.shape[-1]} must be divisible by scale {scale}") + + left = scale // 2 - 1 + right = scale // 2 + return (x[..., left::scale] + x[..., right::scale]) * 0.5 + + +def _run_original_f02sine_on_cpu(self, f0_values: torch.Tensor) -> torch.Tensor: + """Run the unmodified ``_f02sine`` without invoking NPU ``linear1d``.""" + output_device = f0_values.device + output = self._step_audio2_original_f02sine(f0_values.cpu()) + return output.to(output_device) + + +def _f02sine_with_npu_safe_downsample(self, f0_values: torch.Tensor) -> torch.Tensor: + """Use the exact NPU midpoint path, with a narrow CPU fallback.""" + if getattr(self, "flag_for_pulse", False): + return _run_original_f02sine_on_cpu(self, f0_values) + + upsample_scale = self.upsample_scale + if upsample_scale <= 0: + raise ValueError(f"upsample_scale must be positive, got {upsample_scale}") + + scale = int(upsample_scale) + midpoint_supported = scale == upsample_scale and scale % 2 == 0 and f0_values.shape[1] % scale == 0 + if not midpoint_supported: + return _run_original_f02sine_on_cpu(self, f0_values) + + rad_values = (f0_values / self.sampling_rate) % 1 + rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], device=f0_values.device) + rand_ini[:, 0] = 0 + rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini + + rad_values = _linear_downsample_even_scale(rad_values.transpose(1, 2), scale).transpose(1, 2) + phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi + phase = F.interpolate( + phase.transpose(1, 2) * self.upsample_scale, + scale_factor=self.upsample_scale, + mode="linear", + ).transpose(1, 2) + return torch.sin(phase) + + +def patch_step_audio2_hift_for_npu(hift: torch.nn.Module) -> None: + """Patch the non-causal Step-Audio2 HiFT implementation for Ascend. + + The ``flashcosyvoice.SineGen2`` instantiated by Step-Audio2 1.0.0 is + non-causal and reduces a full-rate phase tensor by ``1 / 480`` before + restoring it to the waveform rate. Ascend's ``upsample_linear1d`` kernel + can raise an AIVector UB-address exception (ACL 507015) for that reduction. + + The exact midpoint form keeps the common path on NPU. Unsupported or pulse + configurations delegate only ``_f02sine`` to CPU, preserving upstream + behavior without restoring the old whole-HiFT CPU offload. + """ + if getattr(hift, "_step_audio2_npu_downsample_patched", False): return - hift.to("cpu") - original_forward = hift.forward + try: + sine_gen = hift.m_source.l_sin_gen + original_f02sine = sine_gen._f02sine + except AttributeError as exc: + raise TypeError("expected a Step-Audio2 flashcosyvoice HiFT with m_source.l_sin_gen._f02sine") from exc - def _forward_on_cpu(_module, *args, **kwargs): - output_device: torch.device | None = None + if getattr(sine_gen, "causal", False): + raise ValueError("the Step-Audio2 NPU HiFT patch only supports non-causal SineGen2") - def to_cpu(value): - nonlocal output_device - if isinstance(value, torch.Tensor): - if output_device is None and value.device.type != "cpu": - output_device = value.device - return value.cpu() - return value - - cpu_args = tuple(to_cpu(a) for a in args) - cpu_kwargs = {k: to_cpu(v) for k, v in kwargs.items()} - output = original_forward(*cpu_args, **cpu_kwargs) - - if output_device is None: - return output - - def to_device(value): - if isinstance(value, torch.Tensor): - return value.to(output_device) - return value - - if isinstance(output, tuple): - return tuple(to_device(v) for v in output) - return to_device(output) - - hift.forward = MethodType(_forward_on_cpu, hift) - hift._npu_cpu_offload_patched = True + sine_gen._step_audio2_original_f02sine = original_f02sine + sine_gen._f02sine = MethodType(_f02sine_with_npu_safe_downsample, sine_gen) + hift._step_audio2_npu_downsample_patched = True + logger.info("Patched Step-Audio2 HiFT linear downsample for Ascend NPU") @contextmanager @@ -99,7 +135,7 @@ def _patched_ensure_models_loaded(self) -> None: _original_ensure_models_loaded(self) if was_loaded or self.device.type != "npu" or self._hift is None: return - patch_hift_module_for_npu(self._hift) + patch_step_audio2_hift_for_npu(self._hift) def _patched_forward(self, generated_speech_tokens, prompt_wav, return_bytes=True):