vocoder-ops

Fused anti-aliased Snake activation for BigVGAN-class neural vocoders, loadable through kernels. The resampling pair that keeps the nonlinearity from aliasing is folded into the activation, so the oversampled signal never reaches memory. The reference baseline is the framework formulation, matched to 1.9e-7.

A neural vocoder's Snake activation must be anti-aliased, and the standard way is to upsample, apply the nonlinearity, and downsample: two depthwise convolutions with the oversampled tensor materialized twice, three times per generator block. None of it is a GEMM, so the whole chain is memory-bound and fuses. This kernel fuses it, and a whole BigVGAN block's activations run in about a third less time with output identical to seven decimal places.

BigVGAN generator stages at ratio 2 and 25 taps on 14 x86 threads: 1.10x, 1.79x, 2.07x, and 1.11x per stage, 1.60x for a whole block's three activations (51.7 ms against 32.3), with output matching the framework to 1.9e-7.

Usage

import kernels

vo = kernels.get_kernel("phanerozoic/vocoder-ops", version=1, trust_remote_code=True)

act = vo.AntiAliasedSnake(channels=512, ratio=2, taps=25, alpha=a, beta=b)
y = act(x)                       # [B, C, L] -> [B, C, L]

y = vo.snake(x, alpha, beta)     # plain Snake, no resampling

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

Method

Three things carry the speedup.

Polyphase decomposition: both filters decompose by the residue of the oversampled index modulo ratio, and within a phase the taps land on consecutive samples. Walking the oversampled grid directly makes both filters strided and forces the chain scalar; walking phases makes every inner loop a unit-stride convolution. This alone was worth 6x.

One transcendental instead of two operations: Snake is x + sin^2(ax)/b, and sin^2(u) = (1 - cos 2u)/2, so the activation is x + (1 - cos(2ax))/(2b). The cosine is evaluated with a Cody-Waite reduction and a minimax polynomial, with quadrant selection done branchlessly; a per-lane branch there spills the vector and runs scalar, which costs more than the polynomial itself.

Multiversioned leaves: on x86-64 the hot loops carry both an AVX2 and a baseline lowering behind an IFUNC resolver. The attribute has to sit on the functions containing the loops rather than their callers, because at::parallel_for takes a type-erased callable and a clone attribute never reaches the lambda body.

Measured

BigVGAN generator stages, ratio 2, 25 taps, against the framework formulation; both produce identical output to 1.9e-7 relative:

Plain snake without resampling is 2.41x on aarch64 and 0.69x on x86-64, where the framework's vectorized sine is already well optimized. The anti-aliased path is the one BigVGAN actually uses.

Correctness

Output matches the framework formulation to 1.9e-7 relative across every stage measured. Expressed as two separate convolutions, the reference pads the activation, so the oversampled signal is zero outside the span rather than the activation of the filter tail; fusing removes that boundary, so this kernel reimposes it and matches the framework result exactly. Without it the interior agrees to floating-point rounding and only the edges differ.

Requirements and limits

• CPU, x86_64-linux and aarch64-linux, for PyTorch 2.11, 2.12, and 2.13 (six prebuilt variants), each architecture built on its own host and validated against the prebuilt binary.
• Plain snake is slower than the framework on x86-64; use the anti-aliased path, which is what BigVGAN calls.

References

Lee et al., BigVGAN (2023); Ziyin et al., "Neural Networks Fail to Learn Periodic Functions and How to Fix It" (the Snake activation, 2020); Cody-Waite range reduction.

License

Apache-2.0.