[Hardware][Ascend] Keep HiFT vocoder on NPU (#5242)

Signed-off-by: fatSheep <tzh2005t@163.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

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())

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:

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

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():

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):