TurboQuant KV Cache Compression — Experiments

A research implementation and empirical evaluation of TurboQuant KV cache quantization for transformer inference, based on the paper:

Zandieh et al., "TurboQuant: Post-Training Quantization via Turbocharged Random Hadamard Transforms", ICLR 2026 — arxiv.org/abs/2504.19874

---

Overview

This project measures the real-world impact of TurboQuant KV cache compression across multiple novel evaluation metrics not previously studied in community re-implementations. Experiments run on Qwen2.5 (0.5B / 3B) and Meta-Llama-3.1-8B-Instruct at 2-, 3-, and 4-bit quantization.

What TurboQuant Does

1. Apply a random orthogonal rotation R to each K/V vector (fixes coordinate distribution to Beta)
2. Store the vector norm separately in float32
3. Quantize with a Lloyd-Max codebook (MSE-optimal for that distribution)
4. Decompress via inverse rotation + norm rescaling

---

Project Structure

TurboQuant Google Experiment/
├── experiment.py          # Baseline study — Qwen2.5-0.5B, all core metrics
├── experiment_1.py        # Novel metrics study — Meta-Llama-3.1-8B-Instruct
├── modal_app.py           # Cloud GPU runner (Modal, A100)
├── info.ipynb             # Exploratory notebook (RotationMatrix, QJL classes)
├── results/
│   ├── Llama 3.1 8B/      # Per-seed JSON results (seeds 0–2048)
│   ├── Qwen 2.5 3B/       # Per-seed JSON results
│   └── *.json             # Early-iteration and v1 result snapshots
├── requirements.txt
└── README.md

---

Metrics

Core Metrics (experiment.py)

Metric	Description
Perplexity delta	Baseline vs. compressed NLL on held-out tokens
Cosine similarity	Reconstructed vs. original K/V vectors
Relative MSE	Normalised reconstruction error
Attention correlation	Pearson r between full-precision and compressed attention scores
Top-1 recall	Fraction of top-attended tokens preserved after compression
MSE ratio	Attention score MSE relative to FP16 variance
Memory reduction	FP16 KB vs. compressed KB per layer

Novel Metrics (experiment_1.py)

#	Metric	Why it matters
1	Attention KL Divergence	Detects distribution flattening (not just re-ranking)
2	Per-Layer Sensitivity Curve	Identifies which layers tolerate compression; guides layer-adaptive bit allocation
3	Norm vs. Direction Error	Decomposes error into scalar rescaling vs. geometric rotation
4	Token-Position Degradation	Validates residual-window heuristic; shows older tokens degrade faster
5	GQA Cross-Query Consistency	Unique to Grouped-Query Attention; measures disagreement amplification across query heads sharing one KV head

All five metrics are computed in a single forward pass via registered hooks (~2 min/bit-width on A100).

---

Installation

# Clone the repo
git clone https://github.com/<your-username>/turboquant-experiments.git
cd turboquant-experiments

# Create and activate a virtual environment
python -m venv .venv
source .venv/bin/activate  # Windows: .venv\Scripts\activate

# Install dependencies
pip install -r requirements.txt

GPU required. Experiments default to CUDA; the code falls back to CPU automatically, but Llama-8B is impractical without a GPU.

---

Usage

Local — Qwen2.5-0.5B Baseline

python experiment.py

Results are written to results/turboquant_<timestamp>.json.

Local — Llama-3.1-8B Novel Metrics

python experiment_1.py

Results are written to results/turboquant_novel_<timestamp>.json.

Cloud — Modal (A100)

# One-time setup: store your HuggingFace token
modal secret create huggingface HF_TOKEN=hf_...

# Run with default bit-widths (4, 3, 2)
modal run modal_app.py

# Run with a specific subset
modal run modal_app.py --bits 4
modal run modal_app.py --bits 4,3,2

Jupyter Exploration

jupyter notebook info.ipynb

---

Results

Results are stored as JSON under results/. Each file contains:

{
  "run_id": "<timestamp>",
  "model": "meta-llama/Meta-Llama-3.1-8B-Instruct",
  "device": "cuda",
  "model_cfg": { "n_layers": 32, "n_kv_heads": 8, "head_dim": 128 },
  "ppl_base": 12.34,
  "results": {
    "4": {
      "ppl_base": 12.34, "ppl_tq": 12.51, "ppl_delta": 0.17,
      "attn_corr": 0.998, "top1_recall": 0.94, "mse_ratio": 0.003,
      "cosine_sim": 0.999, "rel_mse": 0.001,
      "mem": { "fp16_kb": 512.0, "tq_kb": 128.0, "reduction_vs_fp16": 4.0 }
    }
  }
}

Pre-computed results across 10 random seeds (0, 7, 42, 100, 148, 256, 512, 999, 1337, 2048) are included for Llama-3.1-8B and Qwen2.5-3B.

---

Dependencies

See requirements.txt. Key packages:

• torch — tensor operations and model inference
• transformers >= 4.43.0 — model loading and DynamicCache API
• scipy — Lloyd-Max codebook construction
• modal — cloud GPU execution (optional)

---

Citation

@inproceedings{zandieh2026turboquant,
  title     = {TurboQuant: Post-Training Quantization via Turbocharged Random Hadamard Transforms},
  author    = {Zandieh, Amir and others},
  booktitle = {International Conference on Learning Representations (ICLR)},
  year      = {2026},
  url       = {https://arxiv.org/abs/2504.19874}
}

---

License

This project is released for research purposes. Model weights are subject to their respective licenses (Meta Llama 3 Community License, Qwen License).