| --- |
| license: apache-2.0 |
| library_name: kernels |
| tags: |
| - kernel |
| - audio |
| - codec |
| - quantization |
| - cpu |
| --- |
| |
| # 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 |
|
|
| ```python |
| 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 |
|
|
| | Symbol | Purpose | |
| |---|---| |
| | `ResidualVQ(codebooks)` | plan holding the codebook stack and precomputed norms | |
| | `.encode(x, n_stages)` | residual quantization to `[..., Q]` int32 codes | |
| | `.decode(codes, n_stages)` | codes to `[..., D]`; truncate stages for lower bitrate | |
|
|
| ## 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: |
|
|
| | config | rvq-codec | gather + sum | stack + sum | one-hot matmul | |
| |---|---:|---:|---:|---:| |
| | EnCodec Q=8 K=1024 D=128 | 0.33 ms | 0.48 ms | 1.78 ms | 17.5 ms | |
| | SNAC Q=4 K=4096 D=768 | 0.73 ms | 0.76 ms | 1.63 ms | 98.7 ms | |
| | Mimi Q=32 K=2048 D=256 | 3.22 ms | 2.20 ms | 8.45 ms | 141.7 ms | |
|
|
| 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: |
|
|
| | config | rvq-codec | cdist | matmul | dist tensor | |
| |---|---:|---:|---:|---:| |
| | EnCodec Q=8 K=1024 D=128 | 31.6 ms | 24.7 ms | 24.3 ms | 73.7 MB | |
| | DAC Q=9 K=1024 D=8 | 30.3 ms | 16.9 ms | 22.0 ms | 95.1 MB | |
| | SNAC Q=4 K=4096 D=768 | 70.5 ms | 67.9 ms | 66.0 ms | 59.0 MB | |
| | Mimi Q=32 K=2048 D=256 | 224.7 ms | 156.4 ms | 187.1 ms | 393.2 MB | |
|
|
| 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. |
|
|
