--- license: apache-2.0 library_name: kernels tags: - kernel - audio - vocoder - bigvgan - cpu --- # 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. ![Per-stage bars showing the framework against the fused kernel across a BigVGAN generator block](https://huggingface.co/kernels/phanerozoic/vocoder-ops/resolve/main/media/hero.gif) *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 ```python 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 | Symbol | Purpose | |---|---| | `AntiAliasedSnake(channels, ratio, taps, alpha, beta)` | fused upsample, Snake, downsample | | `snake(x, alpha, beta)` | plain Snake activation, no resampling | ## 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: | stage | vocoder-ops | framework | speedup | |---|---:|---:|---:| | **x86-64** (i9-12900H, 14 threads) | | | | | C=256 L=800 | 0.93 ms | 1.02 ms | 1.10x | | C=128 L=6400 | 3.17 ms | 5.66 ms | 1.79x | | C=64 L=12800 | 3.30 ms | 6.83 ms | 2.07x | | C=32 L=25600 | 3.37 ms | 3.72 ms | 1.11x | | **whole block, 3 activations** | **32.30 ms** | **51.67 ms** | **1.60x** | | **aarch64** (Cortex-A76, 4 threads) | | | | | C=256 L=800 | 3.68 ms | 20.68 ms | 5.63x | | C=128 L=6400 | 14.76 ms | 99.15 ms | 6.72x | | C=64 L=12800 | 14.72 ms | 100.05 ms | 6.80x | | C=32 L=25600 | 14.62 ms | 94.72 ms | 6.48x | | **whole block, 3 activations** | **143.4 ms** | **943.8 ms** | **6.58x** | 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.