← Infra · 2026-10-25
📖 DEEP DIVE · INFRA

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

Irene Tenison, Abdulrahman Alabdulkareem, Daniel Karl I. Weidele, Nima Dehmamy, Dan Gutfreund
4 minread4.30score

Head-dimension KV compression often reconstructs full-width tensors at decode and adds FLOPs. RACE calibrates pretrained attention into a natively reduced form: shared Q/K projections that stay RoPE-aligned, compressed by query–key activation energy, cutting KV storage and attention compute together. At 50% KV retention and 256K context: total inference FLOPs −47.8%, throughput 1.79×, TTFT −43.2%, allocated GPU memory −28.1%. Versus EigenAttention/PALU/Thin Keys/RAP, RACE checks RoPE keep, shared projection, storage cut, attention-FLOPs cut, and no decode reconstruct. Head-dim compression must be native reduction, not reconstruct-at-decode.

1. Problem

Growing KV hurts memory/bandwidth. Width compression that rebuilds full K/V at decode trades memory for FLOPs. Calibration should produce reduced attention natively.

Full vs compressed attention.
Full vs compressed attention.

2. Method

Shared projections map Q/K into one low-dim space with RoPE still aligned; activation energy picks directions; V compresses with them. One calibration; no up-project at decode. Composes with token reduction and quantization.

3. Core claim

Compression becomes a reduced attention operator that preserves RoPE geometry. The same projection that saves memory should also save attention FLOPs.

4. Experiments

On Llama/Qwen zero-shot, PPL, LongBench: at 50% / 256K → FLOPs −47.8%, 1.79× throughput, TTFT −43.2%, memory −28.1%.

Accuracy curves.
Accuracy curves.
FLOPs/throughput.
FLOPs/throughput.

5. Limits

Calibration domain shift; retune retention with aggressive quant; other ops still bound end-to-end speed.

6. Practice

When KV memory is tight, try native reduced attention before compress-and-reconstruct; scan 50% retention against TTFT and LongBench.