rvq-codec

Residual vector quantization encode and decode for EnCodec, DAC, and SNAC-class audio codecs and the audio language models built on them, loadable through kernels. Decode output is bit-identical to the framework result (torch.equal), not merely close. The reference baselines are three framework formulations of the same operation, all measured.

Every audio language model ends at an RVQ decode: a stack of codebook lookups summed into a waveform latent, run once per frame of generated audio. Written in framework ops it is Q separate advanced-index gathers, each materializing an [N, D] tensor, followed by Q-1 adds, and every one of those intermediates exists only to be summed away. This kernel is a single traversal that produces exactly the same bytes.

Decode of 30 s of audio on 14 x86 threads, against three framework formulations of the same operation: the one-hot matmul route at 17.5 to 141.7 ms, stack-and-sum at 1.63 to 8.45, a hand-written gather-and-sum at 0.48 to 2.20, and the kernel at 0.33 to 3.22, with output identical to all of them.

Usage

import kernels

rv = kernels.get_kernel("phanerozoic/rvq-codec", version=1, trust_remote_code=True)

q = rv.ResidualVQ(codebooks)      # [Q, K, D], or a list of [K, D]
codes = q.encode(x)               # [..., D] -> [..., Q] int32
recon = q.decode(codes)           # [..., Q] -> [..., D]
recon = q.decode(codes, n_stages=4)   # truncate for a lower bitrate

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

API

Method

The two directions have opposite structure, and the measurements say so.

Decode is a single traversal: Q codebook lookups accumulated in place, with no intermediate per stage. It is exact, so the output is the framework result to the bit.

Encode's cost is a dense [N, D] x [D, K] product per stage, which BLAS already runs at hardware peak. A hand-blocked search measured 0.03-0.09x against it, and transposing the codebook to remove the horizontal reductions did not help, because the real constraint is codebook reuse across vectors, which is exactly what a blocked GEMM provides. So the product stays in the framework and only what follows it is fused: argmin, the gather, and the residual subtract, which a framework implementation materializes twice. That took the worst case from 4641 ms to 272 ms. Distances use ||r - c||^2 = ||r||^2 - 2 r.c + ||c||^2, dropping the term constant in k, with codebook norms precomputed once by the plan.

Measured

Decode of 30 s of audio, 14 threads, x86-64, against three framework formulations of the same operation:

The kernel is 44x to 136x the one-hot matmul formulation and 2.2x to 5.5x the stack-and-sum route, while a carefully hand-written gather-and-sum loop is at parity; the kernel's advantage there is that it needs no intermediates and is one call rather than a loop.

Encode, same host, against the framework's matmul formulation:

Encode runs at 0.73-0.94x of the framework matmul formulation while avoiding the [N, K] distance materialization, which is 59 to 393 MB across the stage loop at these sizes.

Correctness

Decode output is torch.equal to the framework result at every configuration measured, not merely close.

Reconstruction error is non-increasing in the number of stages only when each stage's codebook scale tracks the residual it quantizes, which is what training produces. With every stage at the same scale, stage 2 subtracts a code far larger than stage 1's residual and error goes up. That is a property of the codebooks, not of the algorithm, and the test suite pins both behaviors.

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.
• Encode is at parity with the framework matmul by design; the win there is the removed distance materialization, not throughput.

References

Zeghidour et al., SoundStream (2021); Défossez et al., EnCodec (2022); Kumar et al., DAC (2023); Siuzdak et al., SNAC (2024).

License

Apache-2.0.