[Hardware][Ascend] Add opt-in chunked anti-alias FIR upsampling (#7896)

Signed-off-by: Jiaqian Liu <591998922@qq.com>

tests/model_executor/models/common/test_alias_free_activation.py

@@ -0,0 +1,201 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU correctness and dispatch tests; no checkpoints or serving processes."""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.model_executor.models.common import alias_free_activation as fir
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+ENV = "VLLM_OMNI_NPU_ANTIALIAS_MAX_CONV_LENGTH"
+
+
+@pytest.fixture(autouse=True)
+def cpu_platform(monkeypatch):
+    monkeypatch.setattr(fir.current_omni_platform, "is_npu", lambda: False)
+    monkeypatch.setattr(fir.current_omni_platform, "is_xpu", lambda: False)
+    monkeypatch.delenv(ENV, raising=False)
+    with torch.device("cpu"):
+        yield
+
+
+def native_resample(module, x):
+    """Independent whole-tensor reference, including the existing cast order."""
+    dtype = x.dtype
+    channels = x.shape[1]
+    weight = module.filter.to(x.device).expand(channels, -1, -1)
+    if isinstance(module, fir.UpSample1d):
+        padded = F.pad(x.float(), (module.pad, module.pad), mode="replicate").to(weight.dtype)
+        out = module.ratio * F.conv_transpose1d(padded, weight, stride=module.stride, groups=channels).to(dtype)
+        return out[..., module.pad_left : -module.pad_right]
+    padded = F.pad(x.float(), (module.pad_left, module.pad_right), mode="replicate").to(weight.dtype)
+    return F.conv1d(padded, weight, stride=module.stride, groups=channels).to(dtype)
+
+
+@pytest.mark.parametrize("kernel,stride", [(1, 1), (7, 1), (11, 2), (12, 2), (12, 3), (12, 4), (15, 5)])
+@pytest.mark.parametrize("length", [31, 32, 33, 63, 64, 65, 127])
+def test_chunked_matches_native(monkeypatch, kernel, stride, length):
+    generator = torch.Generator().manual_seed(123)
+    x = torch.randn(2, 3, length * 2, generator=generator)[..., ::2]
+    weight = torch.randn(3, 1, kernel, generator=generator)
+    native = F.conv_transpose1d
+    expected = native(x, weight, stride=stride, groups=3)
+    lengths = []
+
+    def capture(segment, *args, **kwargs):
+        lengths.append(segment.shape[-1])
+        return native(segment, *args, **kwargs)
+
+    monkeypatch.setattr(F, "conv_transpose1d", capture)
+    actual = fir._chunked_conv_transpose1d(x, weight, stride, 32)
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+    assert max(lengths) <= 32
+    assert (len(lengths) > 1) == (length > 32)
+    assert actual.dtype == x.dtype and actual.device == x.device
+    assert torch.isfinite(actual).all()
+
+
+@pytest.mark.parametrize("length", [6143, 6144, 6145, 8160, 8192, 12288, 12289])
+def test_large_boundary_and_strict_cap(monkeypatch, length):
+    x = torch.randn(1, 2, length, generator=torch.Generator().manual_seed(1))
+    weight = torch.randn(2, 1, 12, generator=torch.Generator().manual_seed(2))
+    native = F.conv_transpose1d
+    lengths = []
+
+    def capture(segment, *args, **kwargs):
+        lengths.append(segment.shape[-1])
+        return native(segment, *args, **kwargs)
+
+    expected = native(x, weight, stride=2, groups=2)
+    monkeypatch.setattr(F, "conv_transpose1d", capture)
+    torch.testing.assert_close(fir._chunked_conv_transpose1d(x, weight, 2, 6144), expected, rtol=1e-5, atol=1e-5)
+    assert max(lengths) <= 6144
+
+
+@pytest.mark.parametrize("kind", ["zero", "constant", "impulse"])
+def test_boundary_signals(kind):
+    x = torch.zeros(1, 2, 105)
+    if kind == "constant":
+        x.fill_(1)
+    elif kind == "impulse":
+        x[..., [0, 25, 26, 27, 51, 52, 53, 104]] = 1
+    weight = fir.kaiser_sinc_filter1d(0.25, 0.3, 12).expand(2, -1, -1)
+    torch.testing.assert_close(
+        fir._chunked_conv_transpose1d(x, weight, 2, 32), F.conv_transpose1d(x, weight, stride=2, groups=2)
+    )
+
+
+@pytest.mark.parametrize("cap,length", [(0, 129), (32, 16), (32, 32)])
+def test_native_fast_path_is_one_call(monkeypatch, cap, length):
+    x = torch.ones(1, 2, length)
+    weight = torch.ones(2, 1, 12)
+    calls = []
+    sentinel = torch.empty(0)
+
+    def capture(segment, *args, **kwargs):
+        calls.append(segment)
+        return sentinel
+
+    monkeypatch.setattr(F, "conv_transpose1d", capture)
+    assert fir._chunked_conv_transpose1d(x, weight, 2, cap) is sentinel
+    assert len(calls) == 1 and calls[0] is x
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("ratio,kernel", [(2, 12), (3, 15)])
+def test_public_npu_branch_on_cpu(monkeypatch, dtype, ratio, kernel):
+    module = fir.UpSample1d(ratio, kernel)
+    x = torch.randn(2, 3, 129, generator=torch.Generator().manual_seed(9)).to(dtype)
+    expected = native_resample(module, x)
+    monkeypatch.setattr(fir.current_omni_platform, "is_npu", lambda: True)
+    module._max_conv_length = 32
+    native = F.conv_transpose1d
+    lengths = []
+
+    def capture(segment, *args, **kwargs):
+        lengths.append(segment.shape[-1])
+        return native(segment, *args, **kwargs)
+
+    monkeypatch.setattr(F, "conv_transpose1d", capture)
+    actual = module(x)
+    torch.testing.assert_close(actual, expected)
+    assert len(lengths) > 1 and max(lengths) <= 32
+    assert actual.dtype == dtype
+    assert module.state_dict() == {}
+
+
+@pytest.mark.parametrize("module_type", [fir.UpSample1d, fir.DownSample1d])
+def test_non_npu_does_not_chunk(monkeypatch, module_type):
+    monkeypatch.setenv(ENV, "32")
+    module = module_type()
+    assert getattr(module, "_max_conv_length", 0) == 0
+
+    def unexpected(*args, **kwargs):
+        pytest.fail("non-NPU path called a chunking helper")
+
+    monkeypatch.setattr(fir, "_chunked_conv_transpose1d", unexpected)
+    x = torch.randn(1, 2, 129)
+    torch.testing.assert_close(module(x), native_resample(module, x))
+
+
+def test_downsampling_is_unchanged(monkeypatch):
+    module = fir.DownSample1d()
+    monkeypatch.setattr(fir.current_omni_platform, "is_npu", lambda: True)
+    monkeypatch.setenv(ENV, "32")
+    x = torch.randn(1, 2, 129)
+    expected = native_resample(module, x)
+    native = F.conv1d
+    lengths = []
+
+    def capture(segment, *args, **kwargs):
+        lengths.append(segment.shape[-1])
+        return native(segment, *args, **kwargs)
+
+    monkeypatch.setattr(F, "conv1d", capture)
+    torch.testing.assert_close(module(x), expected)
+    assert lengths == [129 + module.pad_left + module.pad_right]
+    assert not hasattr(module, "_max_conv_length")
+
+
+@pytest.mark.parametrize("value,expected", [(None, 0), ("0", 0), ("12", 12), ("6144", 6144)])
+def test_npu_configuration(monkeypatch, value, expected):
+    monkeypatch.setattr(fir.current_omni_platform, "is_npu", lambda: True)
+    if value is not None:
+        monkeypatch.setenv(ENV, value)
+    assert fir._npu_conv_max_length(12, 2) == expected
+
+
+@pytest.mark.parametrize("value", ["-1", "1", "11", "invalid"])
+def test_invalid_npu_configuration(monkeypatch, value):
+    monkeypatch.setattr(fir.current_omni_platform, "is_npu", lambda: True)
+    monkeypatch.setenv(ENV, value)
+    with pytest.raises(ValueError):
+        fir._npu_conv_max_length(12, 2)
+
+
+def test_unsupported_transpose_kernel_keeps_native(monkeypatch):
+    monkeypatch.setattr(fir.current_omni_platform, "is_npu", lambda: True)
+    monkeypatch.setenv(ENV, "32")
+    assert fir._npu_conv_max_length(2, 3) == 0
+    x, weight = torch.randn(1, 2, 65), torch.randn(2, 1, 2)
+    torch.testing.assert_close(
+        fir._chunked_conv_transpose1d(x, weight, 3, 32), F.conv_transpose1d(x, weight, stride=3, groups=2)
+    )
+
+
+def test_cap_smaller_than_kernel_rejected():
+    with pytest.raises(ValueError, match="kernel"):
+        fir._chunked_conv_transpose1d(torch.ones(1, 2, 65), torch.ones(2, 1, 12), 2, 11)
+
+
+def test_alias_free_activation_state_dict_roundtrip(monkeypatch):
+    module = fir.AliasFreeActivation1d(torch.nn.PReLU(3))
+    restored = fir.AliasFreeActivation1d(torch.nn.PReLU(3))
+    restored.load_state_dict(module.state_dict(), strict=True)
+    assert set(module.state_dict()) == {"act.weight"}
+    monkeypatch.setattr(fir.current_omni_platform, "is_npu", lambda: True)
+    restored.upsample._max_conv_length = 32
+    x = torch.randn(2, 3, 129)
+    torch.testing.assert_close(restored(x), module(x), rtol=1e-5, atol=1e-6)

tests/model_executor/models/common/test_alias_free_activation_npu.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Real-device public FIR tests without model checkpoints."""
+
+import pytest
+import torch
+
+from tests.model_executor.models.common.test_alias_free_activation import ENV, native_resample
+from vllm_omni.model_executor.models.common import alias_free_activation as fir
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.npu,
+    pytest.mark.skipif(not fir.current_omni_platform.is_npu(), reason="requires Ascend NPU"),
+]
+
+
+@pytest.fixture(autouse=True)
+def full_precision_convolution(monkeypatch):
+    pytest.importorskip("torch_npu")
+    monkeypatch.setattr(torch.npu.conv, "allow_hf32", False)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("batch,channels,length", [(1, 32, 65), (2, 3, 6145), (1, 32, 12288)])
+def test_public_npu_resampler(monkeypatch, dtype, batch, channels, length):
+    monkeypatch.setenv(ENV, "6144")
+    with torch.device("npu"):
+        module = fir.UpSample1d().to(dtype=dtype)
+    assert module._max_conv_length == 6144
+    generator = torch.Generator(device="cpu").manual_seed(42)
+    x_cpu = torch.randn(batch, channels, length, generator=generator, device="cpu").to(dtype)
+    x = x_cpu.to("npu")
+    expected = native_resample(module, x_cpu)
+    native = fir.F.conv_transpose1d
+    lengths = []
+
+    def capture(segment, *args, **kwargs):
+        lengths.append(segment.shape[-1])
+        return native(segment, *args, **kwargs)
+
+    monkeypatch.setattr(fir.F, "conv_transpose1d", capture)
+    with torch.inference_mode():
+        actual = module(x)
+    assert actual.device.type == "npu" and actual.dtype == dtype
+    assert max(lengths) <= 6144
+    assert (len(lengths) > 1) == (length > 6144)
+    assert torch.isfinite(actual).all()
+    tolerance = {torch.float32: (1e-4, 1e-5), torch.float16: (2e-3, 2e-3), torch.bfloat16: (2e-2, 2e-2)}
+    rtol, atol = tolerance[dtype]
+    torch.testing.assert_close(actual.cpu(), expected, rtol=rtol, atol=atol)
+    monkeypatch.setenv(ENV, "0")
+    assert module._max_conv_length == 6144
+
+
+@pytest.mark.parametrize("allow_hf32", [False, True])
+def test_npu_whole_vs_chunked(monkeypatch, allow_hf32):
+    monkeypatch.setattr(torch.npu.conv, "allow_hf32", allow_hf32)
+    monkeypatch.setenv(ENV, "32")
+    with torch.device("npu"):
+        module = fir.UpSample1d()
+    x = torch.randn(2, 3, 129, device="cpu").to("npu")
+    with torch.inference_mode():
+        chunked = module(x)
+        module._max_conv_length = 0
+        whole = module(x)
+    torch.testing.assert_close(chunked, whole, rtol=1e-5, atol=1e-6)

vllm_omni/model_executor/models/common/alias_free_activation.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Alias-free activation for BigVGAN-style speech decoders.
 
 Provides anti-aliased activation (upsample → activation → downsample) with
@@ -12,6 +12,7 @@ Used by: Qwen2.5-Omni, Qwen3-TTS v1, CoVo-Audio, and future BigVGAN vocoders.
 from __future__ import annotations
 
 import math
+import os
 
 import torch
 import torch.nn as nn
@@ -104,6 +105,50 @@ def replication_pad_1d(hidden_states: torch.Tensor, pad_left: int, pad_right: in
     return torch.cat(segments, dim=-1)
 
 
+def _npu_conv_max_length(kernel_size: int, stride: int) -> int:
+    """Read the opt-in input-length cap once, when a resampler is constructed.
+
+    Zero disables chunking. Other platforms and kernels shorter than the stride
+    keep the native path. A positive cap must accommodate at least one kernel.
+    """
+    if not current_omni_platform.is_npu():
+        return 0
+    name = "VLLM_OMNI_NPU_ANTIALIAS_MAX_CONV_LENGTH"
+    value = int(os.environ.get(name, "0"))
+    if value < 0 or (0 < value < kernel_size):
+        raise ValueError(f"{name} must be 0 or >= kernel_size ({kernel_size}), got {value}")
+    return value if kernel_size >= stride else 0
+
+
+def _chunked_conv_transpose1d(
+    x: torch.Tensor, weight: torch.Tensor, stride: int, max_input_length: int
+) -> torch.Tensor:
+    """Depthwise FIR transpose convolution with bounded input chunks.
+
+    Only bias-free, padding=output_padding=0, dilation=1 FIRs with kernel >=
+    stride are supported. Global edge padding must already be present in x.
+    Left context supplies every contribution to each retained output; outputs
+    are concatenated, not overlap-added. This is mathematically equivalent to
+    the native convolution, but backend rounding can depend on the input shape.
+    """
+    length = x.shape[-1]
+    kernel_size = weight.shape[-1]
+    if max_input_length <= 0 or length <= max_input_length or kernel_size < stride:
+        return F.conv_transpose1d(x, weight, stride=stride, groups=x.shape[1])
+    if max_input_length < kernel_size:
+        raise ValueError("max_input_length must accommodate the FIR kernel")
+    overlap = (kernel_size - 1 + stride - 1) // stride
+    block_size = max_input_length - overlap
+    pieces = []
+    for start in range(0, length, block_size):
+        end = min(start + block_size, length)
+        context_start = max(0, start - overlap)
+        out = F.conv_transpose1d(x[..., context_start:end], weight, stride=stride, groups=x.shape[1])
+        trim = (start - context_start) * stride
+        pieces.append(out[..., trim:] if end == length else out[..., trim : trim + (end - start) * stride])
+    return torch.cat(pieces, dim=-1)
+
+
 # ---------------------------------------------------------------------------
 # Up / Down sampling with Kaiser sinc filters
 # ---------------------------------------------------------------------------
@@ -123,17 +168,18 @@ class UpSample1d(nn.Module):
 
         filt = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size)
         self.register_buffer("filter", filt, persistent=False)
+        self._max_conv_length = _npu_conv_max_length(self.kernel_size, self.stride)
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         channels = hidden_states.shape[1]
         if current_omni_platform.is_npu():
             input_dtype = hidden_states.dtype
             hidden_states = replication_pad_1d(hidden_states.to(self.filter.dtype), self.pad, self.pad)
-            hidden_states = self.ratio * F.conv_transpose1d(
+            hidden_states = self.ratio * _chunked_conv_transpose1d(
                 hidden_states,
                 self.filter.expand(channels, -1, -1),
                 stride=self.stride,
-                groups=channels,
+                max_input_length=self._max_conv_length,
             ).to(input_dtype)
         else:
             input_dtype = hidden_states.dtype