KVQuant: How I Got 2-Bit KV Cache Quality to Match 3-Bit Five Extensions to Near-Optimal Vector Quantization

If you’re running large language models locally, the KV cache is probably your biggest memory problem. At long contexts it dominates VRAM, and the obvious fix quantization comes with a quality cost that gets steep fast below 4 bits.

I spent the last few months building five extensions to make that trade-off much better. Here’s what I built and why it works.

The Starting Point

TurboQuant (Zandieh et al., 2025) gives you a principled baseline: rotate KV vectors into an approximately Gaussian distribution, then apply Lloyd-Max quantization coordinate-by-coordinate. The MSE bound is within 2.7x of the Shannon lower bound as close to optimal as any scalar quantizer gets.

The problem is what it ignores. It treats every token identically. It compresses each vector in isolation. And once the quantization error is committed, there’s no attempt to recover structure from it.

These aren’t minor oversights. They’re properties of how transformers actually behave, and exploiting them gives real gains.

Extension 1: Attention-Weighted Quantization

Not all tokens matter equally. A token receiving 30% of the attention and one receiving 0.001% get identical bit-widths under uniform quantization that’s obviously wrong.

The fix: rank tokens by attention weight, give the top half one extra bit, give the bottom half one fewer. Average bit-width stays the same, bits go where the model actually looks.

Result: 47–70% reduction in attention-weighted distortion per layer on distilgpt2 at the same average bit-width.

Extension 2: Delta Compression

In a streaming KV cache, consecutive tokens are highly correlated. The delta between adjacent vectors is typically much smaller than the vectors themselves compress the delta instead of the absolute vector and you get lower distortion for free.

Store the first vector at full precision as an anchor, then compress each subsequent delta with the quantizer. Errors can accumulate over long sequences, so an adaptive anchor strategy re-anchors when the sequence changes rapidly.

Result: 1.1–2.2x lower MSE across distilgpt2 layers, with earlier layers benefiting most.

Extension 3: Adaptive Bit Allocation

Static allocation has a fundamental problem: you don’t know which tokens will matter when you compress them. A token ignored for the first 40 steps might become the most attended token at step 41.

An EMA-based importance tracker solves this. Each token maintains a running average of the attention it receives. As scores evolve, tokens move between 1/2/3/4-bit tiers and get recompressed. A demoted token can be promoted again if its importance recovers.

Extension 4: Low-Rank Error Correction

Quantization error isn’t random noise it has structure. The top few singular vectors of the residual matrix R = K — K̂ account for a disproportionate share of the total error energy.

A rank-4 SVD of the residual, stored at 7.4% of the full correction budget, recovers most of that structure. Applied at inference, the correction costs O(T·r·d) FLOPs negligible for small r.

Result: 96% of 2-bit PPL degradation recovered on distilgpt2. gpt2-medium drops from +173 dPPL at 2-bit to +5.95 with rank-4 correction essentially 3-bit quality at 2-bit storage.

Extension 5: Product Quantization

Scalar quantization treats each post-rotation coordinate independently. Product Quantization splits each vector into M subspaces and learns a separate k-means codebook per subspace.

At M=16, b=8: 128 bits per vector = 2 bits/dim. Same storage as 2-bit scalar. But instead of 4 centroids per dimension, you get 256 centroids per subspace capturing inter-dimension correlations that scalar quantization misses entirely.

Result: PQ produces coherent, on-topic generation at 2 bits/dim on TinyLlama where 2-bit scalar completely collapses within 30 tokens.

Implementation Fixes

Along the way I fixed several issues in the baseline:

• Codebook fitted to Gaussian approximation instead of true sphere marginal replaced with sampling from the correct distribution
• QR decomposition silently producing reflections (det = -1) fixed with a sign flip ensuring det = +1 - Nearest-centroid search was O(N·d·k) replaced with binary search, 14–22x faster at 4-bit
• k-means++ seeding cuts initial MSE by 75% at 1-bit vs uniform spacing

Results Summary

2-bit + rank-4 correction lands within 0.07 PPL of plain 3-bit 25% storage reduction with no perceptible quality loss.

Code and Paper

Everything is implemented as composable PyTorch modules. The extensions stack use one or all five. 88 tests passing.

Paper: https://doi.org/10.17605/OSF.IO/9WSKZ Code: https://github.com/syedMohib44/kvquant