---
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.

![Four-way decode bars per codec configuration: one-hot matmul, stack-and-sum, gather-and-sum, and the kernel](https://huggingface.co/kernels/phanerozoic/rvq-codec/resolve/main/media/hero.gif)

*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.