266 lines
8.8 KiB
Python
266 lines
8.8 KiB
Python
# 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 contextlib import contextmanager
|
|
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())
|
|
|
|
|
|
def test_sdpa_context_preserves_exceptions_from_token2wav_body(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
module = _load_patch_module(monkeypatch)
|
|
fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn")
|
|
fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = lambda: None
|
|
|
|
@contextmanager
|
|
def fake_math_context(*, require_available=False):
|
|
del require_available
|
|
yield
|
|
|
|
fake_attention_patch.npu_math_sdpa_context = fake_math_context
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"vllm_omni.platforms.npu.models.cosyvoice2_dit_attn",
|
|
fake_attention_patch,
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="original graph failure"):
|
|
with module.npu_token2wav_sdpa_context():
|
|
raise RuntimeError("original graph failure")
|
|
|
|
|
|
def test_sdpa_context_requires_math_for_npu_graph(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
module = _load_patch_module(monkeypatch)
|
|
required = []
|
|
fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn")
|
|
fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = lambda: None
|
|
|
|
@contextmanager
|
|
def fake_math_context(*, require_available=False):
|
|
required.append(require_available)
|
|
yield
|
|
|
|
fake_attention_patch.npu_math_sdpa_context = fake_math_context
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"vllm_omni.platforms.npu.models.cosyvoice2_dit_attn",
|
|
fake_attention_patch,
|
|
)
|
|
|
|
with module.npu_token2wav_sdpa_context(require_math=True):
|
|
pass
|
|
|
|
assert required == [True]
|
|
|
|
|
|
def test_sdpa_context_does_not_hide_patch_setup_failures(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
module = _load_patch_module(monkeypatch)
|
|
fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn")
|
|
|
|
def fail_patch():
|
|
raise RuntimeError("attention patch failed")
|
|
|
|
fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = fail_patch
|
|
fake_attention_patch.npu_math_sdpa_context = contextmanager(lambda: iter((None,)))
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"vllm_omni.platforms.npu.models.cosyvoice2_dit_attn",
|
|
fake_attention_patch,
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="attention patch failed"):
|
|
with module.npu_token2wav_sdpa_context():
|
|
pass
|