Skip to content

vllm_omni.diffusion.models.auk.alias_free_triton

Triton kernels for the AuK codec's alias-free SnakeBeta activation.

The eager module oversamples with a grouped transposed convolution, applies SnakeBeta, and decimates with a grouped strided convolution. Both grouped convolutions carry one 12-tap filter shared by every channel, and cuDNN and ATen serve them with generic depthwise kernels that dominate the decode (more than half of a 12 s decode on H200). Here the oversampling is written as its polyphase form fused with the activation, and the decimation as a direct strided FIR, each one pass over memory per row. The arithmetic is the same as the eager module in fp32; only the summation order of the taps differs.

The ops are registered with torch.library.triton_op so the compiled decode buckets trace them instead of breaking the graph.

alias_free_snake

alias_free_snake(*args, **kwargs)