Attention-aware vector quantization for the KV cache
2-bit and 1-bit caches that beat data-oblivious SOTA — a codebook that's data-aware but not data-dependent, calibrated in 15 seconds.
Alara Dirik · PhD, Imperial College London
Laguna-XS.2model-agnostic
01 / The bottleneck
Long context is a memory problem.
The KV cache stores a key and value vector for every past token, and the model re-reads the whole thing on every step. On Laguna-XS.2 only 10 of 40 layers grow with context — but those full-attention layers still push cache size linearly with length, and at 131K it becomes a memory and bandwidth wall.
fp16 · 131K tokens
5.4 GB
YaRN max · 10 full-attn layers
AttnVQ 2-bit
0.73 GB
7.4× smaller · attnvq-32×256
AttnVQ 1-bit
0.40 GB
14× smaller · attnvq-16×256
Only the 10 full-attention layers are compressed; 30 sliding-window layers stay fp16.
AttnVQThe bottleneck
02 / Prior work
Prior methods pick one: deployability or fidelity.
Data-oblivious · KIVI · TurboQuant ’26 · RaBitQ ’24
Cheap, but coarse
Fixed grids or random rotations + a data-independent codebook. No calibration, runs online on GPU — but fidelity drops sharply below ~2 bits.
Learn the codebook from data for higher fidelity. The cost: second-order preprocessing, poor GPU support, and re-tuning when the data shifts.
This work · AttnVQ
Learned, at deploy cost
A learned codebook without the heavy preprocessing — 10–15 traces, ~15 s, on GPU. The fidelity of data-dependent VQ with the deployability of data-oblivious methods.
AttnVQPrior work
03 / Method
Codebooks optimized for attention, not MSE.
AttnVQ learns a small per-subspace codebook with the Linde–Buzo–Gray algorithm, but replaces the usual MSE objective with attention-output error — allocating resolution to the directions attention actually weights. Calibration is light: 10–15 traces, 15–30 s, on GPU.
01
Split
Each K/V vector → M subspaces.
→
02
Learn (attn-LBG)
Codebook fit to minimize attention-output error.
→
03
Encode
Store the 8-bit nearest-centroid index.
→
04
Attend
Look up centroids at decode; near-full fidelity.
A codebook fit on SWE-agent traces quantizes political-debate caches as well as in-domain — and the reverse holds.
The K/V geometry is fixed by the model's frozen projections and RoPE, not the input text; calibration captures the model, not the corpus. Hence data-aware, not data-dependent.
RoPESplit — Laguna rotates only half its RoPE dims; this variant aligns subspaces to that split (strongest at 1-bit)
AttnVQMethod
04 / Results — accuracy vs budget
Best method at every budget below 4 bits.
AttnVQ (ours)
attention-aware product VQ
RoPESplit (ours)
Laguna variant · 1-bit
TurboQuant
data-oblivious SOTA
Scalar · INT2/4 · sign
+ KIVI, ternary (points)
96% of fp16 at 2 bits
vs TurboQuant 83%, INT2 75%. At 1 bit we take the top two spots; at 0.5 bit we're the only method that runs.
Mean LongBench F1 over 4 language+code tasks · Laguna-XS.2 · single 15-SWE-trace codebook · synthetic passage-retrieval shown separately →
AttnVQAccuracy frontier
05 / Results — LongBench cheap metrics
SWE-fit codebooks on LongBench.
AttnVQ (ours)
product VQ · 0.5–4 b
RoPESplit
Laguna 1-bit variant
TurboQuant
MSE rotation baseline
Scalar · KIVI · sign
+ ternary (points)
0.960 key cosine @ 2 b
vs INT2 0.866 at higher bpe · cross-domain calib → LongBench prompts
Mean over 5 LongBench tasks · 50 prompts/task · no generation (cheap proxy)
AttnVQCheap metrics
06 / Results — LongBench, per task
At 1 bit, both our variants beat every baseline.
LongBench v1 · F1×100 · iso ≈1 bpe · single 15-SWE-trace codebook · ✻ passage-retrieval is the lone synthetic task
AttnVQPer-task
07 / Takeaway
For a fixed model, learning the codebook is nearly free.
▹
2-bit: 96% of fp16 across LongBench — TurboQuant 83%, INT2 75%.
▹
1-bit: our two variants beat every iso-budget baseline on every task.
▹
0.5-bit: the only method that still runs at all.
▹
Data-aware, not data-dependent: 15 traces, ~15 s, on GPU; codebooks transfer across domains.
Honest limits & next: synthetic passage-retrieval is the one task we trail at 2 bits; the memory-efficient kernel is real but not yet fused — fusing only adds speed, not memory. Next: the fused decode path.
Builds on: Linde–Buzo–Gray (1980) · Product Quantization (Jégou et al., 2011) · vs. TurboQuant (Zandieh et al., ICLR ’26) · RaBitQ (Gao & Long, SIGMOD ’24) · KIVI (2024)