Listen: TurboQuant: From Paper to Triton Kernel in One Session

As large language models process longer prompts, storing their previously computed states—known as the key-value cache—quickly clogs up graphics processing unit, or GPU, memory. Google recently introduced TurboQuant, a compression algorithm designed to squeeze this cache down to just a few bits with virtually zero loss in accuracy.

The secret behind TurboQuant is a beautiful math trick. By applying a random rotation to the key-value vectors, their data distributions become highly predictable. This allows developers to use a precomputed, optimized codebook to quantize the data, bypassing the need for expensive, block-by-block normalization.

To see if the claims held up, I implemented TurboQuant from scratch for the Gemma 3 model, eventually building a custom Triton kernel to compute attention directly on compressed bytes.

The journey was full of roadblocks. An early, faithful implementation of the random rotation using Python loops was painfully slow, making the GPU wait endlessly for the CPU. I fixed this by replacing it with a single, precomputed matrix multiplication. I also struggled with subclassing the Hugging Face cache structure, which broke continuously across different library versions. Patching the update method directly proved to be a much cleaner solution.

Ultimately, the results were striking. The optimized two-bit implementation delivered character-for-character identical text output compared to the uncompressed baseline, running at the exact same speed while using up to six times less video random access memory.