RACE: RoPE-Aligned Attention Compression Based on Query-Key Activation Energies

Irene Tenison; Abdulrahman Alabdulkareem; Daniel Karl I. Weidele; Nima Dehmamy; Dan Gutfreund

cs.AIOct 20, 2026

RACE: RoPE-Aligned Attention Compression Based on Query-Key Activation Energies

Authors: Irene Tenison, Abdulrahman Alabdulkareem, Daniel Karl I. Weidele, Nima Dehmamy, Dan Gutfreund

Organizations: MIT CSAIL · MIT-IBM Computing Research Lab

Abstract

Head-dimension compression can reduce the memory required by the key-value (KV) cache, but existing methods often reconstruct full-dimensional representations during inference, adding significant decoding FLOPs overhead. We formalize the design considerations for converting pretrained attention into a natively reduced-dimensional form that reduces both KV-cache storage and attention computation using a calibration set. Specifically, we consider reductions that (1) apply a shared projection to queries and keys so that they remain in a common reduced feature space, (2) derive this projection jointly from query and key activations to better preserve their interaction, and (3) are compatible with the pretrained RoPE transformation. To obtain a reconstruction-free realization of this reduction, we restrict it to coordinate selection and theoretically show that RoPE compatibility then holds if and only if complete rotary pairs are retained. Building on this characterization, we introduce RACE, a post-training method that selects rotary pairs using their joint query-key activation energy. We theoretically show that this criterion minimizes a separable, calibration-dependent upper bound on post-RoPE query-key product approximation error. The resulting reductions are folded into the model weights and require neither runtime reconstruction nor specialized kernels. Across three model families and different compression levels, RACE outperforms structured pruning and remains competitive with reconstruction-based methods on language-modeling, zero-shot, and long-context tasks. At 50% KV-cache retention and 256K context, RACE reduces total inference FLOPs by 47.8%, yielding 1.79x higher throughput, 43.2% lower time-to-first-token (TTFT), and 28.1% lower allocated GPU memory in comparison to an uncompressed model.

Figures & tables

Panel of Figure 1 — Figure 1: Quality-efficiency trade-off across KV-cache retention levels on Llama-3.2-3B-Instruct. Left: average zero-shot accuracy. Right: relative decoding FLOPs on a logarithmic scale, normalized to full attention ( $1\times$ ). The dashed line represent the uncompressed model. RAP and Thin Keys have decoding FLOPs that overlap those of RACE, while Thin Keys’ accuracy falls below the displayed y-axis. RACE combines accuracy competitive with reconstruction-based methods with the low decoding cost of native dimensional reduction.

	Excecution			Q/K Reduction Criterion			Efficiency
Method	Reconst- ruction free	Custom- Kernel free	Fwd.-only Calibration	Shared Q/K Map	Joint Q/K Signal	Native RoPE	KV Cache Retained	Attn. Params Retained	Prefill FLOPs	Decode FLOPs
Attention ( Vaswani et al., 2017 )	-	-	-	-	-	-	100%	100%	$\times$ 1	$\times$ 1
EigenAttention ( Saxena et al., 2024 )	✗	✓	✓	✗	✓	✗	50%	69.8%	$\times$ 0.75	$\times$ 11.42
PALU ( Chang et al., 2025 )	✗	✗	✓	✗	✗	✗	50%	58.3%	$\times$ 0.75	$\times$ 21.92
Thin Keys ( Yao et al., 2026 )	✓	✓	✓	✗	✗	✗	75%	50%	$\times$ 0.75	$\times$ 0.75
RAP ( Xin et al., 2026 )	✓	✓	✗	✓	✗	✓	50%	50%	$\times$ 0.5	$\times$ 0.5

Table 1: Comparison of head-dimension KV-cache compression methods. RACE combines a shared, jointly selected, RoPE-compatible Q/K reduction with native reduced-dimensional execution. FLOPs are normalized to full attention. See A.4 for derivations of FLOPs.

Figure 2: Native attention compression along the head-dimension. Reconstruction-based methods (right) may cache low-dimensional representations yet return to the original feature dimension during inference. This is often done using $A_{K}^{low}$ and $B_{K}^{low}$ that are decomposition of $W_{K}$ . We instead seek a native reduction (left) that can be applied before RoPE and folded into the pretrained projections, so that attention is executed natively in the reduced dimension. Dashed boxes indicate the corresponding full-dimensional operations.

Panel of Figure 3 — Figure 3: Quality across KV-cache retention ratios on Llama-3.2-3B-Instruct (Top) and Qwen2.5-7B (Bottom). WikiText perplexity (lower is better), and mean accuracy across 15 zero-shot tasks and mean LongBench score on 16 english tasks (higher is better).

Panel of Figure 4 — Figure 4: Long-context inference efficiency on Llama-3.2-3B-Instruct. The top row reports prefill, decoding, and overall throughput speedups over the uncompressed model (higher is better). The bottom row reports time to first token, allocated GPU memory, and total inference FLOPs (lower is better). Results are shown at 75% and 50% KV-cache retention for prompt lengths of 64K, 128K, and 256K tokens, with the output length fixed at 128 tokens. These results are average over 10 runs.

Panel of Figure 7 — Figure 7: Analysis of RoPE-pair selection and compression structure. (a) Computational requirements of RoPE-compatible scoring criteria. (b) Perplexity under alternative pair-selection criteria. All methods retain complete RoPE pairs. (c) Effect of attention-head width at matched model size across compression ratios.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: Token- and head- dimension KV-cache compression. Token-level methods reduce the number of retained cache entries N, whereas head-dimension methods reduce the dimensionality of each cached representation. These approaches act on orthogonal axes of the KV cache and can be combined.

	Total Inference TFLOPs	TTFT (ms)	Prefill T.put (tok/sec)	Decode T.put (tok/sec)	Total T.put (tok/sec)	GPU Memory (GB)
Eager	3.973	41.051	13937	45.186	224.44	5.884
SDPA	3.973	28.854	17755.36	56.864	282.904	5.882
Flash Attention 3	3.973	34.170	15054	46.23	230.127	5.882

Appendix

Table 2: Comparison of methods that compress KV Cache along head dimension. RACE operates natively in the reduced dimensions without reconstruction. See A.4.1 for the derivations.

Figure 9: Representative RACE $P_{QK}$ and $P_{V}$ matrices before folding. Left: Each column corresponds to one compressed Q/K coordinate and each row to an original head coordinate. Nonzero entries indicate the coordinates retained by $P_{QK}$ . Right: Each column corresponds to one compressed V coordinate and each row to an original head coordinate.

Panel of Figure 10 — Figure 10: Effect of optional post-compression recovery. Perplexity and average zero-shot accuracy before and after a short recovery stage for RACE and RAP at different KV-cache retention ratios. “w/o FT” denotes the folded compressed checkpoint evaluated directly, while the solid curves apply post-fold KL distillation from the uncompressed model. Recovery has little effect at mild compression but substantially restores performance at 50% retention, particularly in perplexity.

Panel of Figure 11 — Figure 11: Memory scaling of RACE with context length. Total model parameter memory (left), attention-parameter memory (center), and KV-cache memory (right) for the base model and RACE at 75% and 50% KV retention across 64K–256K-token contexts. Parameter memory is independent of sequence length and decreases after folding the reduced-dimensional projections, while KV-cache memory scales with context length; RACE correspondingly reduces both the fixed attention-parameter footprint and the context-dependent cache footprint.

					Completion	Reasoning				Knowledge QA			Understanding		Truth	Specialized Knowledge
Method	KV Cache	QK Dim.	V Dim.	Wiki (ppl)	lambada	hella	piqa	wino	arc_e	sciq	obqa	arc_c	boolq	race	truth fulqa	college _CS	college_ Phy	prof _law	HS_ psych	Avg
Qwen2.5 7B
Base	100%	128	128	10.976	0.706	0.795	0.784	0.707	0.759	0.960	0.460	0.495	0.845	0.418	0.524	0.700	0.480	0.529	0.897	0.671
	87.50%	112	112	11.438	0.696	0.791	0.771	0.703	0.772	0.954	0.456	0.504	0.839	0.421	0.457	0.710	0.500	0.511	0.881	0.664
	87.50%	96	128	11.080	0.693	0.797	0.780	0.711	0.782	0.956	0.462	0.514	0.844	0.414	0.516	0.700	0.490	0.512	0.883	0.670
	75%	96	96	12.290	0.687	0.782	0.756	0.703	0.777	0.953	0.460	0.508	0.838	0.414	0.411	0.660	0.451	0.486	0.868	0.650

Appendix

Table 3: Detailed quality results for head-dimension compression. We report WikiText perplexity and per-task zero-shot accuracy across the evaluated model families and KV-cache retention ratios. The final column reports the average zero-shot accuracy across tasks. For each KV Cache compression level per model, we highlight the best performing perplexity and average zero-shot accuracy. RACE is comparable to reconstruction-based methods while not requiring the reconstruction or the additional FLOPs.

	KV Cache Retained	Narrative QA	Qasper	Multifield QA	Hotpot QA	Wikimqa	Musique	gov_report	qmsum	multi_news	trec	triviaqa	samsum	passage _count	passage_ retrieval_en	lcc	repobench	Average
Llama3.2 3B Instruct.
Uncompressed	100%	21.25	39.89	50.06	47.31	36.75	19.15	33.91	23.5	26.64	71	88.41	42.62	8.5	75	52.01	52.5	43.03
	87.5%	24.81	33.42	49.61	47.52	36.52	17.9	33.95	23	26.98	73	90.47	42.91	4	63.5	51.26	53.9	42.05
	75%	22.2	30.79	45.3	40.44	35.29	15.96	32.88	22.73	25.77	71	87.78	41.16	3.5	58.5	44.34	44.45	38.88
RACE	50%	11.45	15.08	25.46	22.86	20.01	15.52	21.44	20.5	20.97	59.5	78.45	35.32	1.18	5.75	34.83	35.91	26.51
	87.5%	23.49	32.55	46.14	41.67	36.7	14.63	30.91	22.6	26.31	68.5	87.58	40.96	4.5	45	49.3	52.24	38.94

Appendix

Table 4: Per-task results of Qwen2.5 7B and Llama3.2 3B Instruct on LongBench. We report zero shot longbench score across the evaluated model families and KV-cache retention ratios. The final column reports the average across tasks. For each KV Cache compression level per model, we highlight the best performing average result.

Tokens /Precision Retained	RACE KV Cache Reatained	Total KV Cache Retained	wikitext	hellaswag	piqa	arc_challenge	boolq	mmlu_ college_cs	Average
H20 (Token Compression)
	100%	50%	12.66	0.66	0.85	0.48	0.82	0.72	0.706
	75%	37.5%	15.8	0.65	0.83	0.45	0.81	0.63	0.674
512 (50%)	50%	25%	17.44	0.62	0.79	0.44	0.74	0.43	0.604
	100%	25%	13.62	0.66	0.85	0.48	0.79	0.66	0.688
	75%	18.75%	16.12	0.65	0.83	0.45	0.78	0.63	0.668

Appendix

Table 5: Per-task results of Qwen2.5 7B on Wikitext (perplexity) and zero-shot accuracy on combining RACE with Token compression and quantization methods.

Figure 12: Layer-wise normalised approximation error as in Eq. 5 with respect to uncompressed model for joint and single-sided energy selectors. Joint Q/K selection yields the lowest average error compared to Q-only or K-only selection.

KV Cache Retained	Selection Criteria	wikitext	piqa	wino	acr_c	sciq	arc_easy	boolq	race	mmlu_cs	mmlu_phy	Avg
	Joint QK Energy (RACE)	13.88	0.74	0.67	0.46	0.93	0.70	0.77	0.42	0.48	0.35	0.61
	Magnitude	14.58	0.74	0.64	0.44	0.89	0.68	0.71	0.41	0.37	0.25	0.57
	Random	14.95	0.71	0.63	0.45	0.89	0.67	0.65	0.39	0.35	0.25	0.55
	Gradient Fisher	14.58	0.76	0.66	0.44	0.87	0.69	0.74	0.41	0.38	0.25	0.58
	Attention Score	14.13	0.74	0.66	0.44	0.92	0.69	0.75	0.40	0.40	0.27	0.59
	Energy Q only	14.02	0.75	0.66	0.45	0.93	0.69	0.75	0.42	0.40	0.32	0.60

Appendix

Table 6: Ablation of selection criteria for RoPE-pair retention. We compare RACE’s joint Q/K activation-energy criterion against magnitude, random selection, gradient Fisher, attention-score selection, and Q-only or K-only energy. Results are reported on WikiText perplexity and zero-shot accuracy across nine tasks at 87.5%, 75%, and 50% KV-cache retention. Joint Q/K energy is consistently strongest overall, with the advantage becoming more pronounced at higher compression.

KV-Cache Compression

Key-value cache compression reduces the memory and bandwidth needed to retain attention states during long-context inference.

See all 269 papers

Jun 23, 2026

RoPE-Aware Bit Allocation for KV-Cache Quantization

Existing low-bit KV-cache quantizers often treat each cached key as a flat vector. Under RoPE, however, a key's contribution to a future attention logit decomposes into a position-dependent sum over two-dimensional frequency blocks. This makes key-cache quantization a block-wise bit-allocation problem: high-energy RoPE blocks are more sensitive to quantization error and should receive more bits. We introduce Block-GTQ, a RoPE-aware bit allocator for key-cache quantization built on TurboQuant-MSE(TQ-MSE). For each layer and KV head, Block-GTQ computes a label-free energy score for each RoPE block and greedily allocates integer bit widths by marginal gain. Under matched K/V bit budgets, Block-GTQ better preserves RoPE query-key logits on a ten-model diagnostic panel, cutting per-layer MAE by 32-80% at 2 and 3 b/dim K-only quantization and winning all 367/367 layer comparisons against uniform TQ-MSE. These fidelity gains translate to stronger downstream long-context retrieval, understanding, and reasoning. At K2V2 on Llama-3.1-8B-Instruct, Block-GTQ raises the six-task NIAH average from 70.6 to 97.4, and the LongBench-EN average from 36.87 to 53.31. On AIME 2024/2025 with DeepSeek-R1-Distill-Qwen-7B, without an fp16 recent-key buffer, Block-GTQ at K3V2 scores 51.7/37.5, close to fp16's 54.2/37.9, whereas uniform TQ-MSE collapses to 0.0/0.0. We further implement a packed-cache serving path. On a single H800 GPU with Qwen2.5-3B-Instruct, packed K3V3 achieves 3.24x KV-cache compression with fp16-comparable quality, runs 1.34x faster than fp16 FlashAttention2 at 128K context, reduces peak memory from 56.31 GB to 19.85 GB, and remains feasible at 256K and 512K where fp16 OOMs. Code is available at https://github.com/JIA-Lab-research/blockgtq.

Fengfeng Liang, Yuechen Zhang, Jiaya Jia

1Hong Kong University of Science and Technology · 2The Chinese University of Hong Kong · 3MiMo, Xiaomi Corporation

KV Caching LLM Inference Acceleration Attention Quantization KV-Cache Quantization Low-Bit Quantization Neural Network Quantization

Ask PDF

Aug 3, 2026

ATFlash: Per-RoPE-Wavelength Attention Windows for Compute/Memory-Efficient LLM Inference

The attention score with rotary position embeddings (RoPE) decomposes exactly into a sum over its 2D-rotation frequency pairs, and each pair's wavelength limits how far it can discriminate position. Aligned with this structure, we propose the per-RoPE-wavelength distance window: it prunes the query--key inner-product terms beyond a wavelength-proportional distance. Unlike a sliding window, every key remains reachable, at least through the low-frequency pairs. The reduction rate is input-independent, with a closed form logarithmic in the sequence length

N

, in contrast to dynamic-sparse methods like MInference. Such token-level selection is orthogonal to our frequency-level pruning. The window can therefore be applied on top of those methods. On Qwen2.5-0.5B and Llama-3.2-3B, the window prunes 37--48% of the query--key inner-product terms within each model's native context length. Relative to full attention, the top-1 match rate stays at 96--98% and the mean output-distribution KL at the

10^{-3}

-nat level on LongBench-v2 contexts. We examine absolute scores on long-context benchmarks such as RULER, OpenAI-MRCR, LongCodeQA, and

\infty

Bench: they are broadly preserved. We implement the window as a slice of the query--key contraction axis, leaving the online-softmax recurrences untouched, and port it with minimal diffs into the released FlashAttention-4 prefill and FlashInfer decode. On RTX PRO 6000 with Llama, both ports outpace stock with gains growing with context length, up to

1.29\times

at 128K. End to end on Qwen2.5-7B-1M, with 57% of the inner-product terms pruned, the speedup reaches

1.31\times

at a 1M-token context.

Shun-ichiro Hayashi, Daichi Mukunoki, Tetsuya Hoshino +1

Graduate School of Informatics, Nagoya University · Information Technology Center, Nagoya University

Long-Context Language Model Inference LLM Inference Acceleration Sparse Attention

Ask PDF

Oct 6, 2026

A Self-Pruning Transformer: Extreme KV-Cache Compression with Universal Attention

The large KV-cache size of modern LLMs creates a barrier to efficient deployment. Recent work has explored replacing attention layers' RoPE positional embeddings with alternative decay-based mechanisms, which can then be used to prune KV-cache during inference. However, these decay functions have limited expressivity, and in practice devolve into sliding-window-like eviction patterns. In this work, we propose a unifying framework for complementary and novel decay mechanisms, capturing complex key statistics and interactions while preserving expressive RoPE embeddings and Softmax attention. The resulting Universal Attention is a highly expressive and end-to-end trainable architecture, whose composite decay mechanism acts as a natural,

\textit{adaptive}

pruning criterion, removing tokens that contribute least to attention computation. Experimentally, Universal Attention achieves state-of-the-art

10\times

compression on natural language and synthetic task data, while

\textit{improving}

downstream performance compared to both state-of-the-art baselines and unpruned oracles. It further demonstrates superior long-context generalization with unprecedented

25\times

compression at length 16k.

Davis Wertheimer, Haochen Shen, Ahan Gupta +6

IBM Research, USA · SSAIL Lab, University of Illinois Urbana-Champaign

KV-Cache Compression Transformer Attention LLM Compression KV-Cache Pruning Transformer Pruning LLM Pruning

Ask PDF

Abstract

KV-Cache Compression

Explore similar work

RoPE-Aware Bit Allocation for KV-Cache Quantization

ATFlash: Per-RoPE-Wavelength Attention Windows for Compute/Memory-Efficient LLM Inference

A Self-Pruning Transformer: Extreme KV-Cache Compression with Universal Attention