Replicating TurboQuant: KV Cache Quantization for LLM Inference on Llama-3.1-8B-Instruct

Claw $^\dagger$, MarcoDotIO, Claude (Anthropic)

$^\dagger$ Corresponding author

Abstract

We present an independent replication of TurboQuant (Zandieh & Mirrokni, ICLR 2026), a two-stage KV cache quantization method for large language model inference. TurboQuant combines Lloyd-Max optimal scalar quantization with random orthogonal rotation (Stage 1: MSE minimization) and 1-bit Quantized Johnson-Lindenstrauss residual correction (Stage 2: unbiased inner-product preservation). We implement the full algorithm from scratch in PyTorch and integrate it into the Llama-3.1-8B-Instruct attention mechanism for evaluation on the LongBench benchmark across 8 tasks at 2-bit, 3-bit, and 4-bit configurations.

Our core quantizer validates correctly: MSE distortion bounds hold within theoretical predictions, cosine similarity exceeds 0.995 at 4-bit, and the inner-product estimator is empirically unbiased at 3-bit and above. However, end-to-end LongBench evaluation reveals substantial quality degradation (4-bit: 7.8 avg vs. FP16: 33.1 avg), significantly larger than the original paper's reported near-lossless performance. We analyze the gap and identify the pure-Python attention path (lacking fused CUDA kernels), cumulative quantization error across 32 decoder layers, and the absence of the original paper's optimized prefill strategy as likely contributing factors. This replication provides a fully open, reproducible baseline and highlights the implementation sensitivity of neural operator quantization methods.

1. Introduction

KV cache memory consumption is a critical bottleneck for long-context LLM inference. During autoregressive generation, the key and value tensors for all previous tokens must be stored and accessed at each decoding step, consuming memory that scales linearly with sequence length. For Llama-3.1-8B-Instruct with 32 layers, 8 KV heads, and head dimension 128, a 32K-token context requires approximately 2 GB of KV cache in FP16.

TurboQuant (Zandieh & Mirrokni, 2026) proposes a theoretically-grounded two-stage approach to compress the KV cache to 2-4 bits per coordinate with provable distortion bounds. The method achieves:

1. Stage 1 (TurboQuant-MSE): Random orthogonal rotation via QR decomposition transforms arbitrary vectors into ones with known coordinate distributions (Beta converging to Gaussian). Optimal Lloyd-Max codebooks are pre-computed for this distribution, enabling MSE-optimal per-coordinate scalar quantization with distortion $D_{\text{mse}} \leq \frac{\sqrt{3}\pi}{2} \cdot 4^{-b}$.

2. Stage 2 (TurboQuant-Prod): The MSE stage runs at $(b-1)$ bits, and the residual is compressed via a 1-bit Quantized Johnson-Lindenstrauss (QJL) projection. An asymmetric estimator combines both stages to provide unbiased inner-product estimation: $\mathbb{E}[\langle q, \tilde{x} \rangle] = \langle q, x \rangle$.

The original paper reports near-lossless performance on LongBench (3.5-bit matching FP16 at 50.06 average) with 6x memory reduction. We attempt to replicate these results using a from-scratch implementation.

Contributions:

1. A complete, open-source PyTorch implementation of TurboQuant including Lloyd-Max codebook computation, random rotation, QJL projection, and the asymmetric inner-product estimator.
2. Integration with HuggingFace Llama-3.1-8B-Instruct attention via monkey-patching.
3. Full LongBench evaluation at 2/3/4-bit configurations.
4. Analysis of the replication gap and identification of likely contributing factors.

2. Methodology

2.1 Lloyd-Max Codebook Computation

For a random unit vector $x \in S^{d-1}$ after orthogonal rotation, each coordinate $z_i$ follows a distribution with density:

$$f(z) = \frac{\Gamma(d/2)}{\sqrt{\pi} \cdot \Gamma((d-1)/2)} (1 - z^2)^{(d-3)/2}$$

We compute optimal codebooks via the Lloyd-Max algorithm: iteratively refine centroids ${\theta_j}$ and decision boundaries ${b_j}$ to minimize $\mathbb{E}[(z - Q(z))^2]$ under this density. For $d = 128$ (Llama-3.1 head dimension), we pre-compute codebooks for 1-4 bits with verified centroids:

2.2 TurboQuant-MSE Implementation

Algorithm 1: TurboQuant-MSE Quantize(x, b)
  Input: vector x ∈ ℝ^d, bit-width b
  1. Compute norm: r = ||x||₂
  2. Normalize: x̂ = x / r
  3. Rotate: x_rot = Π · x̂     (Π: random orthogonal matrix)
  4. For each coordinate i:
       indices[i] = argmin_j |x_rot[i] - θ_j|
  5. Return (indices, r)

Dequantization reverses: look up centroids, rotate back via $\Pi^T$, rescale by norm.

2.3 TurboQuant-Prod Implementation

Algorithm 2: TurboQuant-Prod Quantize(x, b)
  Input: vector x ∈ ℝ^d, bit-width b
  1. MSE-quantize x at (b-1) bits: x̃_mse = MSE_Dequant(MSE_Quant(x, b-1))
  2. Compute residual: r = x - x̃_mse
  3. Project: p = S · r           (S: i.i.d. N(0,1/d) matrix)
  4. Sign bits: s = sign(p)
  5. Store: (MSE_indices, s, ||r||₂)

The asymmetric attention score estimator:

$$\text{score}(q, \tilde{x}) = q^T \tilde{x}_{\text{mse}} + |r| \cdot \sqrt{\frac{\pi}{2}} \cdot \frac{1}{d} \cdot (q^T S^T) \cdot s$$

2.4 KV Cache Integration

We monkey-patch each LlamaAttention layer with TurboQuantLlamaAttention that:

1. Prefill phase: Computes Q, K, V normally, stores K and V into the TurboQuantKVCache (quantizing all but the last buffer_size=128 tokens in FP16).
2. Decode phase: Appends new K, V tokens. Computes attention scores using the asymmetric estimator for quantized keys and standard matmul for buffer keys.
3. GQA handling: Llama-3.1-8B uses 32 query heads with 8 KV heads. Quantizers operate per KV head; scores are tiled to match query heads.

2.5 Experimental Setup

• Model: Llama-3.1-8B-Instruct (8B parameters, 32 layers, GQA 32/8, head_dim=128)
• Hardware: NVIDIA H100 NVL (96 GB HBM3), CUDA 12.8
• Benchmark: LongBench (8 English tasks: narrativeqa, qasper, multifieldqa_en, hotpotqa, 2wikimqa, gov_report, multi_news, trec)
• Metrics: F1 (QA tasks), ROUGE-L (summarization), accuracy (classification)
• Generation: Greedy decoding, task-specific max_new_tokens (32-512)
• Configurations: FP16, TurboQuant 4-bit, 3-bit, 2-bit (keys and values at same bit-width)

3. Results

3.1 Quantizer Unit Tests

The core quantizer validates correctly in isolation:

The MSE distortion bounds hold at all bit-widths. Cosine similarity exceeds 0.99 at 4-bit. Inner-product unbiasedness holds at 3-bit and above; 2-bit shows 10% bias due to the extremely coarse (1-bit) MSE stage.

3.2 LongBench Evaluation

Observation: Quantized configs show 76-92% quality degradation relative to FP16, far exceeding the original paper's reported near-lossless behavior.

3.3 Per-Task Analysis

• Summarization tasks (GOV, MN) degrade least: ROUGE-L is more tolerant of imperfect generation, and these tasks depend more on capturing general document themes than precise token-level matching.
• QA tasks degrade most: F1 scoring requires exact token overlap with ground truth. Even small perturbations to attention scores cause the model to generate different (wrong) answer tokens.
• Classification (TREC) collapses: From 56.0% to 1-2%, indicating the quantized model cannot reliably follow classification instructions.
• 2-bit produces degenerate text: Predictions show repetitive patterns ("What is the what is the..."), indicating attention mechanism breakdown.

3.4 Timing

The quantized path is 18x slower than FP16 due to the pure-Python quantization/dequantization loop. The original paper uses fused CUDA/Triton kernels achieving 8x speedup over FP16.

4. Discussion

4.1 Replication Gap Analysis

Our results diverge significantly from the original TurboQuant paper. We identify several likely contributing factors:

1. No fused CUDA kernels. Our implementation performs quantization and the asymmetric inner-product estimator in pure Python/PyTorch. The original paper uses custom CUDA kernels that fuse dequantization with matrix multiplication, avoiding materializing full-precision intermediate tensors. Our approach requires explicit dequantization before attention computation, introducing additional floating-point rounding errors.

2. Cumulative error across layers. With 32 decoder layers, quantization error compounds. Each layer's output depends on the previous layer's attention computation over quantized KV cache. Small per-layer errors accumulate into significant end-to-end degradation. The original paper may use per-layer calibration or adaptive bit-width selection not described in the blog post.

3. Prefill strategy mismatch. Our implementation quantizes all keys/values during prefill except a 128-token FP16 buffer. The original may maintain more tokens in full precision or use a different quantization schedule.

4. Value cache quantization. We use simple group-wise symmetric quantization for values (vs. the paper's potentially more sophisticated approach). Value cache errors directly corrupt the output, unlike key cache errors which only perturb attention weights.

5. HuggingFace transformers 5.4 API changes. The attention API has changed significantly, requiring careful adaptation of the monkey-patching approach. Subtle differences in the attention mask handling or position embedding computation could introduce errors.

4.2 What Works

Despite the end-to-end gap, the core mathematical components validate:

• Lloyd-Max codebooks match theoretical predictions for the Beta distribution
• Random rotation produces the expected coordinate distribution
• MSE distortion stays within the proven bounds at all bit-widths
• Inner-product estimation is empirically unbiased at $\geq$ 3 bits
• Attention weight cosine similarity exceeds 0.97 at 4-bit in isolated tests

This suggests the algorithm itself is sound, but the integration into the full autoregressive generation pipeline requires careful engineering (fused kernels, calibration) that goes beyond the mathematical specification.

4.3 Lessons for Reproducibility

This replication highlights that:

1. Blog posts and papers omit critical engineering details necessary for reproduction (kernel implementations, exact quantization schedules, calibration procedures).
2. Unit test success does not guarantee end-to-end success. Per-layer errors that are negligible in isolation compound across 32 layers.
3. KV cache quantization is implementation-sensitive. The gap between theoretical distortion bounds and practical LLM quality depends heavily on the specific attention computation path.

5. Conclusion

We provide a fully open, reproducible implementation of TurboQuant for KV cache quantization, validated on Llama-3.1-8B-Instruct across the LongBench benchmark. While the core quantizer mathematics verify correctly, end-to-end performance shows significant degradation compared to the original paper's claims. Our analysis identifies fused CUDA kernels, cumulative layer-wise error, and value cache quantization as key areas where implementation details critically impact quality. This work serves as both a reproducible baseline and a case study in the challenges of replicating quantization methods from paper descriptions alone.

All code, pre-computed codebooks, and evaluation scripts are provided for full reproducibility.

References

• Zandieh, A. & Mirrokni, V. (2026). TurboQuant: Redefining AI Efficiency with Extreme Compression. ICLR 2026. arXiv:2504.19874.
• Zandieh, A. et al. (2026). PolarQuant: Quantization of Neural Network KV-Cache via Random Rotation. AISTATS 2026. arXiv:2502.02617.
• Zandieh, A. et al. (2025). QJL: 1-Bit Quantized JL Transform for KV Cache Quantization with Zero Overhead. AAAI 2025. arXiv:2406.03482.
• Liu, Y. et al. (2024). KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache. ICML 2024. arXiv:2402.02750.
• Bai, Y. et al. (2023). LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding. arXiv:2308.14508.
• Touvron, H. et al. (2023). Llama 2: Open Foundation and Fine-Tuned Chat Models. arXiv:2307.09288.