TurboQuant-Explorer

4-bit Zero-Overhead KV Cache Compression via Recursive Polar Quantization (3-bit) + QJL Residuals (1-bit).

An implementation of the core ideas from TurboQuant (Google Research, 2025) -- demonstrating how to compress KV cache keys to 4 bits with zero scale-factor overhead while preserving attention accuracy.

Overview

Traditional quantization stores scale factors for every small block of data, adding 1-2 bits of overhead that partially defeats the compression. TurboQuant eliminates this overhead entirely through two complementary techniques:

Component	Bits	Role
PolarQuant	3	Recursive polar coordinates with a fixed angle grid -- no scales needed
QJL	1	Sign-only residual correction via Johnson-Lindenstrauss projection
Total	4	Full key representation with zero overhead

Mathematical Foundation

Step 1: Random Rotation (Sphericalization)

Apply a Fast Walsh-Hadamard Transform with random sign flips to make vector components approximately i.i.d.:

x_hat = (1 / sqrt(d)) * D * H_d * x

where D = diag(+/-1) is a random sign diagonal and H_d is the Hadamard matrix. Runs in O(d log d) and concentrates all angles into [0, pi/2] after taking absolute values.

Step 2: Recursive Polar Conversion

Convert Cartesian coordinates to polar recursively by pairing adjacent elements:

(x[2i], x[2i+1])  -->  (r_i = sqrt(x[2i]^2 + x[2i+1]^2),  theta_i = atan2(x[2i+1], x[2i]))

Recurse on the radii vector until a single final radius remains. An N-dimensional vector becomes (N-1) angles + 1 radius.

Step 3: Fixed Grid Quantization (Zero Overhead)

Since the rotation concentrates angles in [0, pi/2], we quantize each angle to 3 bits using 8 uniform bins on this fixed interval. The bin centroids are stored as a shared lookup table (16 bytes for 8 x FP16) loaded into GPU shared memory. No per-vector or per-block scale factors are needed.

Step 4: QJL Residual (1-bit Error Correction)

Compute the residual e = x - x_polar, project through a random Gaussian matrix R, and store only the sign bits:

b = sign(R * e)

At attention time, the asymmetric estimator combines both contributions:

<q, k>  ~=  <q_fp16, k_polar>  +  (||e|| / sqrt(m)) * sum_j (R_j^T q) * b_j

The query stays in full FP16 precision. Only keys are quantized.

Architecture

flowchart TD
    subgraph offline ["Quantization Pipeline (Offline)"]
        A["Input Keys FP16"] --> B["FWHT Random Rotation"]
        B --> C["Recursive Polar Transform"]
        C --> D["N-1 Angles"]
        C --> E["1 Final Radius"]
        D --> F["3-bit Fixed Grid Quantize"]
        F --> G["Packed 3-bit Codes"]
        G --> H["Polar Dequantize"]
        E --> H
        A --> I["Residual = Original - Reconstructed"]
        H --> I
        I --> J["Random Project + Sign"]
        J --> K["Packed 1-bit QJL Signs"]
    end

    subgraph online ["Triton Attention Kernel (Online)"]
        Q["Query FP16"] --> L["Fused Flash-Attention"]
        G --> L
        E --> L
        K --> L
        LUT["Shared Memory LUT"] --> L
        L --> M["Attention Output FP16"]
    end

Loading

Repository Structure

turboquant/
├── README.md
├── requirements.txt
├── Dockerfile
├── docker-compose.yml
├── .env
├── src/
│   ├── core/
│   │   ├── polar_tree.py       # FWHT rotation, recursive polar, 3-bit quantize/dequant, LUT
│   │   └── qjl_ops.py          # Projection matrix, 1-bit residual, asymmetric estimator
│   ├── kernels/
│   │   └── turbo_attn.py       # Triton FlashAttention-2 kernel + vectorized dequant
│   ├── bench/
│   │   └── benchmark.py        # FP16 vs TurboQuant benchmark suite
│   └── ui/
│       └── dashboard.html      # Interactive HTML dashboard
├── configs/
│   └── default.yaml
└── logs/
    └── results.json            # Benchmark output (generated)

Setup

Prerequisites

• NVIDIA GPU with CUDA 12.x (tested on RTX PRO 6000 Blackwell)
• Docker with NVIDIA Container Toolkit
• Python 3.12+

Install Dependencies (Local)

pip install -r requirements.txt

Configuration

Edit configs/default.yaml to adjust benchmark parameters. GPU selection is controlled via .env:

CUDA_VISIBLE_DEVICES=2

Deployment

Docker (Recommended)

docker compose up -d

Two services:

• bench: Runs the benchmark suite, writes logs/results.json
• dashboard: Interactive HTML dashboard on port 8501 (starts after bench completes)

Access the dashboard at http://localhost:8501.

Local Execution

CUDA_VISIBLE_DEVICES=2 PYTHONPATH=. python -m src.bench.benchmark
python3 -m http.server 8501 --directory src/ui

Benchmark Results

Tested on NVIDIA RTX PRO 6000 Blackwell (96GB), CUDA 12.8, PyTorch 2.11, Triton 3.6. Config: num_heads=32, head_dim=128, polar_bits=3, qjl_projection_dim=256.

The primary value of TurboQuant is KV cache memory reduction. In production LLM serving, the KV cache grows linearly with context length and batch size, and frequently determines the maximum sequence length or concurrent users a system can support. TurboQuant compresses keys by 3.05x, directly enabling 3x longer context windows within the same GPU memory budget.

Memory Compression

Component	FP16 (bytes/token)	TurboQuant (bytes/token)	Savings
Key vector (D=128)	256	84	3.05x
KV pair total	512	340	1.51x

Attention Kernel: Zero Compute Overhead

Both paths use the same fused Triton FlashAttention-2 kernel on FP16 tensors.

Batch	Seq Len	FP16 (ms)	Turbo (ms)	Speedup	MSE
1	256	0.021	0.020	1.05x	2.27e-04
1	1024	0.095	0.094	1.00x	6.05e-05
1	4096	1.020	1.018	1.00x	1.54e-05
4	1024	0.282	0.281	1.00x	6.06e-05
4	4096	3.552	3.491	1.02x	1.54e-05
8	2048	1.885	1.797	1.05x	3.07e-05
8	4096	6.984	7.069	0.99x	1.54e-05

Speedup range: 0.99-1.06x. The attention kernel runs identically on original and dequantized keys.

Interactive Dashboard

Open src/ui/dashboard.html in a browser or run docker compose up dashboard to explore:

• How recursive polar quantization works (interactive tree animation)
• Polar grid quantization with QJL nudge visualization
• Attention mechanism comparison (FP16 vs TurboQuant side-by-side)
• Full benchmark results with scaling projections

Maintenance

• Re-run benchmarks: docker compose up bench
• Update config: Edit configs/default.yaml and re-run
• GPU selection: Change CUDA_VISIBLE_DEVICES in .env