KV Cache Compression Through the Lens of Transform Coding
Authors: Hannah Laus, Claudio Mayrink Verdun, Hao Wang, Flavio du Pin Calmon, Felix Krahmer
Organizations: Technical University of Darmstadt, the Technical University of Munich and the Munich Center for Machine Learning · Massachusetts Institute of Technology · Red Hat AI Innovation Team and the MIT-IBM Watson AI Lab · John A. Paulson School Of Engineering And Applied Sciences at Harvard University
The key-value (KV) cache stores information from past tokens and is a major memory bottleneck in long-context inference. Existing quantization methods address this bottleneck by representing the KV cache uniformly with lower-precision data types and designing quantization schemes to minimize reconstruction error in the cache itself, without accounting for how that error propagates through attention mechanisms. We prove that, under a white-noise quantization model, the expected attention-aware distortion decomposes into additive key and value contributions that factor across tokens and channels. Building on transform coding and reverse water-filling, which are classical tools from signal processing and rate-distortion theory, we introduce Attention-Aware Transform Coding (AATC), which allocates bits over a calibration set to minimize attention-aware distortion. On Llama-3.1-8B-Instruct and Qwen-2.5-7B-Instruct, evaluated across LongBench, RULER, GSM8K, MMLU-Pro, and MATH-500, our method achieves near-lossless accuracy at approximately 5.8× compression, whereas each baseline degrades in at least some settings.
TABLE I: Distortion factors in D=DK+DV and the design axes they motivate. Each row corresponds to one factor in the per-token, per-channel decomposition of attention-output error under KV quantization. Methods are listed by the factor they primarily target. Here, VQ = Vector quantization, MP = mixed precision and TC = transform coding.
Algorithm 1 Attention-Aware Transform Coding (AATC) for KV cache Compression
Fig. 1: RULER score versus context length for AATC and AATC var-only at 2.5-bit (solid) and 2-bit (dashed), on Llama-3.1-8B-Instruct and Qwen-2.5-7B-Instruct .
LongBench
RULER
Reasoning
Method
Comp.
avg.
4k
8k
16k
32k
GSM8K
MMLU -m
MMLU -cs
MATH-500
Llama-3.1-8B-Instruct
Baseline
1.00
0.570
0.859
0.856
0.828
0.775
0.789
0.438
0.456
0.454
KIVI
5.00
0.564
0.846
0.840
0.810
0.763
0.730
0.390
0.420
0.452
PALU
5.39
0.550
0.851
0.849
0.83
0.769
0.617
0.274
0.320
0.314
KVQuant
5.91
0.549
0.715
0.729
0.676
0.654
0.754
0.376
0.402
0.392
TABLE II: Comp. is the KV cache compression ratio relative to the FP16 baseline. LongBench is the task-averaged score over seven English subsets (per-subtask breakdown in Table III ). Here σ is the standard error of each score, the standard deviation of the per-example scores divided by N over the N evaluated samples. A method is called statistically indistinguishable from FP16 when its score lies within 2σ of the FP16 score.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Method
TriviaQA
Qasper
TREC
SAMSum
LCC
RepoB-P
QMSum
avg.
Llama-3.1-8B-Instruct
Baseline
0.917
0.455
0.725
0.437
0.634
0.567
0.254
0.570
KIVI
0.913
0.440
0.725
0.448
0.619
0.552
0.250
0.564
PALU
0.892
0.449
0.710
0.438
0.539
0.576
0.248
0.550
KVQuant
0.921
0.413
0.675
0.415
0.633
0.546
0.241
0.549
AATC var
0.922
0.452
0.715
0.430
0.634
0.559
0.251
0.566
AATC (ours)
0.919
0.444
0.690
0.434
0.623
0.546
0.252
0.558
Appendix
TABLE III: Per-subtask LongBench scores at the main-table operating point: F1 for QA (TriviaQA, Qasper), ROUGE-L for summarization (SAMSum, QMSum) and query summarization, accuracy for classification (TREC), and edit-similarity for code (LCC, RepoBench-P). Bold = best compressed method per column, and the FP16 baseline is a reference.
LongBench
RULER
Reasoning
Method
Comp.
avg.
4k
8k
16k
32k
GSM8K
MMLU-m
MMLU-cs
MATH
Llama
Baseline (FP16)
1.00
0.570
0.859
0.856
0.828
0.775
0.789
0.438
0.456
0.454
AATC var
7.0
0.543
0.818
0.822
0.775
0.677
0.767
0.402
0.427
0.484
AATC (ours)
7.0
0.568
0.851
0.861
0.809
0.735
0.756
0.416
0.434
0.448
Qwen
Baseline (FP16)
1.00
0.574
0.720
0.716
0.718
0.720
0.818
0.708
0.567
0.680
AATC var
6.9
0.518
0.688
0.653
0.640
0.566
0.795
0.644
0.510
0.664
Appendix
TABLE IV: AATC at the aggressive 2-bit operating point ( ≈ 7 × compression), against the FP16 reference. Baseline compression methods (KIVI, PALU, KVQuant), which run at 2-bit and ≈ 5 to 6 × , are reported in the main results table ( II ).
Fig. 2: LongBench average and GSM8K versus KV averages bits on Llama-3.1-8B-Instruct .
Config
Comp
TriviaQA
Qasper
TREC
SAMSum
LCC
RepoB-P
QMSum
LB avg.
GSM8K
Llama-3.1-8B–Instruct
Baseline FP16
1.0
0.917
0.455
0.725
0.437
0.634
0.567
0.254
0.570
0.789
KIVI 4-bit
3.1
0.917
0.454
0.725
0.440
0.634
0.568
0.250
0.570
0.795
KVQuant 4-bit
3.4
0.914
0.447
0.725
0.426
0.631
0.560
0.256
0.566
0.779
AATC 4-bit
3.8
0.916
0.446
0.725
0.436
0.624
0.565
0.253
0.566
0.783
KIVI 3-bit
3.8
0.914
0.450
0.725
0.435
0.634
0.564
0.254
0.568
0.792
KVQuant 3-bit
4.3
0.921
0.452
0.725
0.433
0.631
0.559
0.250
0.567
0.776
Appendix
TABLE V: LongBench per-subtask scores and GSM8K across KV multiple bit-budgets. Metrics: F1 (TriviaQA, Qasper), accuracy (TREC), ROUGE-L (SAMSum, QMSum), edit-similarity (LCC, RepoBench-P). LB avg. is the seven-task mean.
AATC variant (2-bit, Llama-3.1-8B-Instruct )
TriviaQA
Qasper
TREC
SAMSum
LCC
RepoB-P
QMSum
LB avg.
GSM8K
AATC 2-bit (standard)
0.917
0.456
0.720
0.438
0.630
0.570
0.248
0.568
0.751
Allocation
no qc2 in allocation
0.910
0.410
0.685
0.444
0.545
0.473
0.248
0.531
0.752
no WO in allocation
0.916
0.452
0.720
0.430
0.602
0.551
0.251
0.560
0.763
per-layer allocation
0.913
0.430
0.715
0.445
0.583
0.533
0.249
0.553
0.760
global allocation
0.847
0.337
0.595
0.432
0.475
0.462
0.231
0.483
0.740
Appendix
TABLE VI: Ablation of AATC’s bit-allocation and calibration choices (2-bit, Llama-3.1-8B-Instruct , with LongBench per-subtask, LB avg. = seven-task mean and GSM8K ).
Long-context LLM decoding reads the key-value (KV) cache at every step. Loading it takes longer than computing attention over it, so throughput is bandwidth-bound. Hence, reducing the cache size can raise both decoding speed and serving capacity. The challenge is to reduce cache size while preserving the attention products, keeping reconstruction cheap, and using a fixed per-token bit count. At two bits per element, the most competitive methods rely on orthogonal transforms. However, existing techniques are either data-oblivious or use the query statistics without deriving the transform from a distortion criterion. Moreover, they rely on transforms built on top of random or Hadamard rotations, which equalize variances across entries rather than compacting energy, and fixed-width scalar quantizers, which are suboptimal at low rates. In this paper, we formulate KV cache quantization as a transform coding problem in which distortion is the error in the attention products. We derive closed-form optimal transforms for keys and values from calibration statistics, under a high-resolution model. We show that the optimal key transform is not orthogonal and satisfies a generalized Parseval relation: the attention-aware distortion becomes mean-squared error (MSE) in the transform domain. Thus, we can use MSE-optimal vector quantizers applied directly to the transformed key coefficients. To meet the fixed-width layout requirement, we show that grouping coefficients into equal-volume partitions makes equal-size codebooks attain the variable-rate optimum under the same high-resolution model. At two bits per element, our method, termed NOVA-KV, recovers most of the long-context retrieval accuracy lost by scalar quantization methods at comparable throughput.
Samuel Fernández-Menduiña, Amir Ziashahabi, Eduardo Pavez +2
Department of Electrical and Computer Engineering, University of Southern California
Transformer inference on long sequences is expensive because softmax attention repeatedly reads from a large KV cache. The prevalent approach to this bottleneck is KV cache compression, which replaces the full cache with a compact summary. Despite its practical importance, the design of such summaries is largely driven by empirical experimentation. On the theoretical side, existing results show that KV cache compression can be impossible in the worst case, but offer little systematic guidance for designing algorithms in regimes where accurate compression is possible. We bridge this gap by characterizing the minimax risk of KV cache compression in terms of the intrinsic compressibility of a cache, revealing when and how accurate compression is possible. These results yield novel design principles for KV cache compression under causal masking that map efficiently to prefill and autoregressive decoding while achieving minimax-optimal risk. We instantiate these principles in a practical algorithm and report promising performance on LongBench in targeted experiments. Overall, our results provide a principled avenue for practical KV cache compression with theoretical guarantees.
Lukas Haverbeck, Carmen Amo Alonso, Andres Felipe Posada-Moreno +2
The key-value (KV) cache is the dominant memory bottleneck in long-context language model inference. Existing compression methods apply low-rank factorization or quantization independently, without jointly allocating rank and precision under a shared storage budget. We introduce JoLT, a training-free compressor that treats grouped prefill caches as fourth-order tensors and applies partial Tucker decomposition along the token and feature modes, the two axes that carry low-rank structure, while leaving the head and layer modes intact. A rotated low-bit quantizer captures the truncation residual, and a single Lagrangian dual allocates per-group Tucker ranks and residual bit-widths under a global byte constraint. FlashJoLT replaces the exact token-mode SVD with a randomized approximation that matches JoLT within the free zone at a fraction of the compression cost, and a fused Triton decode kernel evaluates attention directly over the stored factors without materializing dense KV tensors. Across five models from four architecture families, covering multi-head attention, grouped-query attention, and mixture-of-experts architecture, JoLT achieves 2 - 3x compression with less than 0.2% perplexity degradation, without retraining. On RULER at 64K context with LLaMA-3.1-8B, retrieval accuracy remains near-lossless through 3x and declines by only 0.90 and 2.40pp at 4x and 5x, respectively. JoLT demonstrates that tensor-aware low-rank decomposition and quantized residuals, unified under a single storage budget, achieve near-lossless KV-cache compression across diverse model architectures without retraining.
Rahul Krishnan, Volker Schulz
Universität Trier · Fachbereich IV, Mathematik, Universität Trier