Back to Documentation Index | Algorithm | Architecture | CUDA Kernels | Integration | Testing | Improvements
This page provides a detailed mathematical treatment of the TurboQuant algorithm as implemented in FlashQuant. The algorithm achieves near-optimal rate-distortion performance for KV cache compression through three stages: PolarQuant (MSE-optimal compression), QJL correction (unbiased inner products), and fused attention kernels.
Reference: Ashkboos et al., "TurboQuant: Online Vector Quantization for Efficient KV Cache Compression," arXiv 2504.19874, 2025.
- Why Compress the KV Cache?
- PolarQuant: Stage 1
- Why Rotation Works
- Lloyd-Max Closed-Form via erfinv
- Distortion Bounds (Theorem 1)
- QJL Correction: Stage 2
- Variance Bound for TurboQuant_prod (Theorem 2)
- Comparison with Other Quantization Methods
In autoregressive LLM inference, each generated token requires attending over the entire Key-Value cache. For a model with L layers, H attention heads, head dimension d, and sequence length S:
KV cache memory = 2 * L * H * S * d * sizeof(dtype)
For Llama-3-70B (L=80, H=8 KV heads, d=128) at S=8K with FP16:
= 2 * 80 * 8 * 8192 * 128 * 2 bytes = 26.8 GB
This often exceeds the available VRAM, limiting context length or batch size. TurboQuant compresses each coordinate from 16 bits (FP16) to 4 bits, reducing the KV cache by approximately 7.5x when accounting for the stored norm overhead.
PolarQuant separates each vector into its norm and direction, then exploits the fact that random rotation makes directional coordinates approximately i.i.d. Gaussian to apply optimal scalar quantization independently per coordinate.
Given an input vector x in R^d:
1. Compute norm: gamma = ||x||
2. Normalize: x_hat = x / gamma
3. Rotate: y = Pi^T * x_hat (Pi is Haar-random orthogonal)
4. Scalar quantize: q_i = Q_LM(y_i) (Lloyd-Max per coordinate)
5. Store: {nibble_pack(q), gamma} (4 bits/coord + fp32 norm)
1. Unpack indices: q_i from nibble bytes
2. Centroid lookup: y_hat_i = c[q_i] (16-entry table for 4-bit)
3. Inverse rotation: x_tilde = Pi * y_hat
4. Rescale: x_hat = gamma * x_tilde
| Component | Size | Notes |
|---|---|---|
| Packed indices | 64 bytes | 128 coords * 4 bits / 8 bits/byte |
| Norm | 4 bytes | fp32 |
| Total | 68 bytes | vs. 256 bytes (FP16) = 3.76x compression |
For both K and V combined, the compression ratio is effectively 7.5x compared to storing both in FP16.
Henri Poincare observed that if x is uniformly distributed on the unit sphere S^{d-1}, then any single coordinate x_i is approximately Gaussian with distribution N(0, 1/d) for large d. More precisely, for a uniform random unit vector:
x_i ~ N(0, 1/d) as d -> infinity
This is a consequence of the concentration of measure on the unit sphere: in high dimensions, the projection of a uniformly random point on S^{d-1} onto any axis concentrates around zero with variance 1/d.
For finite d, the exact marginal distribution of each coordinate is:
x_i ~ Beta((d-1)/2, (d-1)/2) rescaled to [-1, 1]
The Gaussian approximation becomes excellent for d >= 64, which covers all practical LLM head dimensions (64, 128, 256).
The key insight is that we do not need x to already be uniform on the sphere. Applying a Haar-distributed random orthogonal matrix Pi^T to the normalized vector x/||x|| makes the rotated coordinates behave as if they were drawn from the Poincare distribution, regardless of the original distribution of x.
This is because:
- x/||x|| lies on S^{d-1} (deterministic point on the sphere)
- Pi^T is a uniformly random rotation (Haar measure on O(d))
- Pi^T applied to any fixed point on S^{d-1} produces a uniformly random point on S^{d-1}
Therefore, each coordinate of y = Pi^T * (x/||x||) is approximately N(0, 1/d), and critically, the coordinates are approximately independent -- enabling independent scalar quantization with minimal loss.
FlashQuant generates Pi via QR decomposition of a Gaussian matrix with sign correction:
gen = torch.Generator(device="cpu").manual_seed(seed)
G = torch.randn(d, d, generator=gen, device="cpu", dtype=torch.float32)
Q, R = torch.linalg.qr(G)
# Sign correction for Haar distribution
diag_sign = torch.sign(torch.diag(R))
diag_sign[diag_sign == 0] = 1.0
Pi = Q * diag_sign.unsqueeze(0)The sign correction ensures that Pi is uniformly distributed over O(d) (the orthogonal group), not just the special orthogonal group SO(d). Without this step, the QR decomposition produces a matrix with a systematic bias in the sign of the diagonal of R.
The Lloyd-Max quantizer minimizes expected distortion E[||x - Q(x)||^2] for a scalar random variable with known distribution. The standard approach is an iterative algorithm:
- Initialize centroids uniformly
- Assign each point to its nearest centroid (Voronoi boundaries)
- Update each centroid to the conditional expectation within its region
- Repeat until convergence
When the distribution is Gaussian N(0, sigma^2) with sigma = 1/sqrt(d), the Lloyd-Max solution has a closed-form expression using the inverse error function:
Boundaries:
b_i = sigma * sqrt(2) * erfinv(2*i/L - 1), i = 1, ..., L-1
where L = 2^bits is the number of quantization levels.
Centroids:
c_i = sigma^2 * [phi(a_i) - phi(b_i)] / [Phi(b_i) - Phi(a_i)]
where:
- phi(x) = (1 / (sigma * sqrt(2pi))) * exp(-x^2 / (2sigma^2)) is the Gaussian PDF
- Phi(x) = 0.5 * (1 + erf(x / (sigma * sqrt(2)))) is the Gaussian CDF
- a_i and b_i are the left and right boundaries of region i (with a_0 = -infinity, b_L = +infinity)
This closed-form solution has two critical advantages:
-
No scipy dependency. The original turboquant-vllm used
scipy.optimizefor iterative Lloyd-Max. FlashQuant computes exact boundaries and centroids analytically, reducing runtime dependencies from 12+ packages to justtorch. -
Exact optimality. The iterative algorithm can converge to local optima or stop early. The closed-form is the globally optimal solution.
The C++ implementation (csrc/core/codebook.cpp) provides two paths:
solve_lloyd_max(dim, bits):
if dim >= 64:
return gaussian_lloyd_max(dim, bits) # Closed-form via erfinv
else:
return beta_lloyd_max(dim, bits) # Iterative for exact Beta dist
For dim >= 64, the Gaussian approximation error is negligible (< 0.1% MSE difference). For smaller dimensions, the exact Beta distribution is used with iterative refinement.
Theorem 1 (MSE Distortion Bound). For the PolarQuant quantizer with b bits per coordinate and dimension d, the normalized MSE satisfies:
MSE(x, x_hat) / ||x||^2 <= (sqrt(3) * pi/2) / 4^b
This means the distortion is:
- Only 2.72x from the rate-distortion theoretical optimum
- Dimension-independent (does not grow with d)
- Exponentially decreasing with bit-width (4^b = 256 for 4-bit)
| Bits | 4^b | Theoretical MSE/||x||^2 bound | Empirical (d=128) | |------|-----|-------------------------------|-------------------| | 2 | 16 | 0.1700 | ~0.08 | | 3 | 64 | 0.0425 | ~0.02 | | 4 | 256 | 0.0106 | ~0.005 |
-
After rotation, each coordinate y_i of the normalized vector is approximately N(0, 1/d).
-
The Lloyd-Max quantizer for N(0, 1/d) with L = 2^b levels achieves:
E[(y_i - Q(y_i))^2] <= (sqrt(3) * pi/2) / (d * 4^b) -
Since the coordinates are approximately independent after rotation, the total MSE on the unit sphere is:
E[||y - Q(y)||^2] = sum_{i=1}^{d} E[(y_i - Q(y_i))^2] <= d * (sqrt(3) * pi/2) / (d * 4^b) = (sqrt(3) * pi/2) / 4^b -
Rescaling by ||x||^2 gives the final bound.
The test suite (tests/test_numerical_bounds.py) empirically verifies this bound with a 5x safety factor to account for finite-dimensional effects:
mse = ((x - reconstructed) ** 2).mean().item()
theoretical_bound = math.sqrt(3) * math.pi / 2 * (1.0 / 4 ** bits)
assert mse < theoretical_bound * 5Stage 1 (PolarQuant) minimizes MSE, which is optimal for value cache reconstruction. However, for the key cache, what matters is not reconstruction quality but inner product accuracy -- attention scores are computed as <q, k>.
A quantizer that minimizes MSE does not necessarily minimize inner product error. The quantization error r = k - k_hat may be systematically correlated with k, introducing bias in <q, k_hat> as an estimator of <q, k>.
TurboQuant adds a Quantized Johnson-Lindenstrauss (QJL) correction that makes the inner product estimate unbiased:
<q, k> ~ <q, k_hat_mse> + gamma * sqrt(pi/2) / m * <S*q, sign(S*r)>
\___ MSE term ___/ \___ QJL correction ___/
Where:
k_hat_mseis the Stage 1 MSE reconstructionr = k - k_hat_mseis the quantization residualgamma = ||r||is the residual normSis a random Gaussian projection matrix (d x d), shared across all keyssign(S*r)is stored as 1 bit per dimensionmis the projection dimension (= d in FlashQuant)
TurboQuantProd allocates the total bit budget as:
- (bits - 1) bits for Lloyd-Max MSE quantization (Stage 1)
- 1 bit for QJL sign correction (Stage 2)
For 4-bit TurboQuantProd: 3 bits for MSE (8 levels) + 1 bit for QJL signs.
Unbiasedness:
E[estimate] = <q, k> (exact, not approximate)
The QJL correction term has zero mean when the true residual r aligns randomly with q, and its expectation equals the missing <q, r> component.
Low variance:
Var(correction) = O(||q||^2 * ||r||^2 / m)
Since ||r|| decreases exponentially with the MSE bit-width, and m = d is typically 64-256, the correction variance is negligible.
Cheap storage: The QJL signs are stored as int8 (one sign per dimension), adding only 1 byte per 8 dimensions. In the original turboquant-vllm, these were wastefully stored as float32 (32x overhead); FlashQuant stores them as int8.
Theorem 2 (Inner Product Variance Bound). For TurboQuantProd with b total bits per coordinate (b-1 for MSE, 1 for QJL), the variance of the inner product estimator satisfies:
Var(<q, k_hat>) <= sqrt(3) * pi^2 * ||y||^2 / (d * 4^b)
where ||y|| = ||q|| is the query norm and d is the dimension.
-
The variance is O(||y||^2 / (d * 4^b)), which decreases with:
- Higher dimension d (more averaging)
- Higher bit-width b (exponential improvement)
- Smaller query norm ||y||
-
For practical LLM settings (d=128, b=4): the standard deviation of the error is approximately 0.003 * ||q||, which is negligible compared to typical attention score magnitudes.
The variance has two components:
- MSE contribution:
<q, k_hat_mse - k>has variance proportional to||q||^2 * E[||r||^2] - QJL contribution: The sign quantization adds variance proportional to
||q||^2 * ||r||^2 * pi/(2m)
The total is dominated by the MSE component for b >= 3.
FlashQuant's test suite verifies Theorem 2 with unbiasedness tests (mean error < 0.01 over 10,000 samples) and variance bound tests:
empirical_variance = (errors ** 2).mean().item()
theoretical_bound = math.sqrt(3) * math.pi ** 2 / (dim * 4 ** b_eff)
assert empirical_variance < theoretical_bound * 100| Method | Compression | Bias | Complexity | Quality |
|---|---|---|---|---|
| TurboQuant (MSE) | 7.5x | Biased IP | O(d) per vector | Excellent MSE |
| TurboQuant (Prod) | ~6x | Unbiased | O(d) per vector | Best for attention |
| Scalar Uniform | 7.5x | Biased | O(d) per vector | Poor (~5x from optimal) |
| Product Quantization | Variable | Biased | O(d/m * K) | Good, but slow |
| Vector Quantization | Variable | -- | O(K * d) | Best MSE, impractical K |
Simple uniform quantization without rotation:
q_i = round((x_i - min) / (max - min) * (2^b - 1))
Problems:
- Does not exploit the Gaussian structure after rotation
- Distortion bound is ~5x from optimal (vs. 2.72x for TurboQuant)
- Non-uniform coordinate distributions lead to wasted quantization levels
Splits the vector into subvectors and quantizes each with a codebook learned via k-means:
x = [x_1 | x_2 | ... | x_m] (m subvectors of dimension d/m)
q_j = argmin_k ||x_j - c_k|| (per-subvector codebook)
Problems:
- Requires training codebooks on representative data
- Online learning is expensive (k-means updates per token)
- Subvector independence assumption may not hold
Quantizes the entire vector with a single codebook:
q = argmin_k ||x - c_k||
Problems:
- Codebook size K must be exponential in the effective bits: K = 2^(b*d)
- For d=128, b=4: K = 2^512 -- completely impractical
- Even with structured codebooks, search complexity is prohibitive
TurboQuant achieves the best of all worlds:
- Near-optimal distortion (2.72x from rate-distortion bound) via rotation + Gaussian quantization
- O(d) complexity -- linear in dimension, same as uniform quantization
- No training data needed -- codebook is computed analytically from the Gaussian distribution
- Unbiased inner products -- the QJL correction is unique to TurboQuant
- Online operation -- each vector is compressed independently, enabling streaming KV cache append
FlashQuant — By Ayi NEDJIMI