Researchers from MIT, NVIDIA, and Zhejiang University Propose TriAttention: A KV Cache Compression Method That Matches Full Attention at 2.5× Higher Throughput

Long-chain reasoning is without doubt one of the most compute-intensive duties in fashionable massive language fashions. When a mannequin like DeepSeek-R1 or Qwen3 works by way of a posh math drawback, it may generate tens of 1000’s of tokens earlier than arriving at a solution. Every a type of tokens have to be saved in what known as the KV cache — a reminiscence construction that holds the Key and Value vectors the mannequin must attend again to throughout era. The longer the reasoning chain, the bigger the KV cache grows, and for a lot of deployment eventualities, particularly on shopper {hardware}, this development ultimately exhausts GPU reminiscence totally.

A workforce of researchers from MIT, NVIDIA, and Zhejiang University proposed a way known as TriAttention that instantly addresses this drawback. On the AIME25 mathematical reasoning benchmark with 32K-token era, TriAttention matches Full Attention accuracy whereas reaching 2.5× greater throughput or 10.7× KV reminiscence discount. Leading baselines obtain solely about half the accuracy at the identical effectivity degree.

The Problem with Existing KV Cache Compression

To perceive why TriAttention is necessary, it helps to know the usual method to KV cache compression. Most present strategies — together with SnapKV, H2O, and R-KV — work by estimating which tokens within the KV cache are necessary and evicting the remaining. Importance is usually estimated by trying at consideration scores: if a key receives excessive consideration from latest queries, it’s thought of necessary and saved.

The catch is that these strategies function in what the analysis workforce calls post-RoPE area. RoPE, or Rotary Position Embedding, is the positional encoding scheme utilized by most fashionable LLMs together with Llama, Qwen, and Mistral. RoPE encodes place by rotating the Query and Key vectors in a frequency-dependent approach. As a end result, a question vector at place 10,000 seems to be very completely different from the identical semantic question at place 100, as a result of its route has been rotated by the place encoding.

This rotation signifies that solely essentially the most just lately generated queries have orientations which are ‘updated’ for estimating which keys are necessary proper now. Prior work has confirmed this empirically: growing the statement window for significance estimation doesn’t assist — efficiency peaks at round 25 queries and declines after that. With such a tiny window, some keys that may turn into necessary later get completely evicted.

This drawback is very acute for what the analysis workforce calls retrieval heads — consideration heads whose perform is to retrieve particular factual tokens from lengthy contexts. The related tokens for a retrieval head can stay dormant for 1000’s of tokens earlier than all of a sudden turning into important to the reasoning chain. Post-RoPE strategies, working over a slender statement window, see low consideration on these tokens through the dormant interval and completely evict them. When the mannequin later must recall that data, it’s already gone, and the chain of thought breaks.

The Pre-RoPE Observation: Q/Okay Concentration

The key perception in TriAttention comes from trying at Query and Key vectors earlier than RoPE rotation is utilized — the pre-RoPE area. When the analysis workforce visualized Q and Okay vectors on this area, they discovered one thing constant and hanging: throughout the overwhelming majority of consideration heads and throughout a number of mannequin architectures, each Q and Okay vectors cluster tightly round fastened, non-zero heart factors. The analysis workforce phrases this property Q/Okay focus, and measures it utilizing the Mean Resultant Length R — a normal directional statistics measure the place R → 1 means tight clustering and R → 0 means dispersion in all instructions.

On Qwen3-8B, roughly 90% of consideration heads exhibit R > 0.95, which means their pre-RoPE Q/Okay vectors are almost completely concentrated round their respective facilities. Critically, these facilities are steady throughout completely different token positions and throughout completely different enter sequences — they’re an intrinsic property of the mannequin’s discovered weights, not a property of any specific enter. The analysis workforce additional verify that Q/Okay focus is domain-agnostic: measuring Mean Resultant Length throughout Math, Coding, and Chat domains on Qwen3-8B yields almost similar values of 0.977–0.980.

This stability is what post-RoPE strategies can not exploit. RoPE rotation disperses these concentrated vectors into arc patterns that adjust with place. But in pre-RoPE area, the facilities stay fastened.

From Concentration to a Trigonometric Series

The analysis workforce then present mathematically that when Q and Okay vectors are concentrated round their facilities, the eye logit — the uncooked rating earlier than softmax that determines how a lot a question attends to a key — simplifies dramatically. Substituting the Q/Okay facilities into the RoPE consideration components, the logit reduces to a perform that relies upon solely on the Q-Okay distance (the relative positional hole between question and key), expressed as a trigonometric sequence:

$textual content{logit}(Delta) approx sum_{f} underbrace{|bar{q}_f| |bar{okay}_f|}_{textual content{amplitude}} cos(omega_f Delta + underbrace{bar{phi}_f}_{textual content{part}}) = sum_{f} [a_f cos(omega_f Delta) + b_f sin(omega_f Delta)]$

Here, Δ is the positional distance, ω_f are the RoPE rotation frequencies for every frequency band f, and the coefficients a_f and b_f are decided by the Q/Okay facilities. This sequence produces a attribute attention-vs-distance curve for every head. Some heads desire close by keys (native consideration), others desire very distant keys (consideration sinks). The facilities, computed offline from calibration knowledge, absolutely decide which distances are most well-liked.

The analysis workforce validated this experimentally throughout 1,152 consideration heads in Qwen3-8B and throughout Qwen2.5 and Llama3 architectures. The Pearson correlation between the anticipated trigonometric curve and the precise consideration logits has a imply above 0.5 throughout all heads, with many heads reaching correlations of 0.6–0.9. The analysis workforce additional validates this on GLM-4.7-Flash, which makes use of Multi-head Latent Attention (MLA) slightly than customary Grouped-Query Attention — a meaningfully completely different consideration structure. On MLA, 96.6% of heads exhibit R > 0.95, in comparison with 84.7% for GQA, confirming that Q/Okay focus is just not particular to at least one consideration design however is a basic property of recent LLMs.

How TriAttention Uses This

TriAttention is a KV cache compression technique that makes use of these findings to attain keys with no need any stay question observations. The scoring perform has two elements:

The Trigonometric Series Score (S_trig) makes use of the Q heart computed offline and the precise cached key illustration to estimate how a lot consideration the important thing will obtain, primarily based on its positional distance from future queries. Because a key could also be attended to by queries at many future positions, TriAttention averages this rating over a set of future offsets utilizing geometric spacing.

$S_{textual content{trig}}(okay, Delta) = sum_{f} |mathbb{E}[q_f]| cdot |k_f| cdot cos(omega_f Delta + phi_f)$

The Norm-Based Score (S_norm) handles the minority of consideration heads the place Q/Okay focus is decrease. It weights every frequency band by the anticipated question norm contribution, offering complementary details about token salience past distance desire alone.

$S_{textual content{norm}}^{(0)}(okay) = sum_{f} mathbb{E}[|q_f|] cdot |k_f|$

The two scores are mixed utilizing the Mean Resultant Length R as an adaptive weight: when focus is excessive, S_trig dominates; when focus is decrease, S_norm contributes extra. Every 128 generated tokens, TriAttention scores all keys within the cache and retains solely the top-B, evicting the remaining.

Results on Mathematical Reasoning

On AIME24 with Qwen3-8B, TriAttention achieves 42.1% accuracy in opposition to Full Attention’s 57.1%, whereas R-KV achieves solely 25.4% at the identical KV funds of two,048 tokens. On AIME25, TriAttention achieves 32.9% versus R-KV’s 17.5% — a 15.4 proportion level hole. On MATH 500 with only one,024 tokens within the KV cache out of a attainable 32,768, TriAttention achieves 68.4% accuracy in opposition to Full Attention’s 69.6%.

The analysis workforce additionally introduces a Recursive State Query benchmark primarily based on recursive simulation utilizing depth-first search. Recursive duties stress reminiscence retention as a result of the mannequin should keep intermediate states throughout lengthy chains and backtrack to them later — if any intermediate state is evicted, the error propagates by way of all subsequent return values, corrupting the ultimate end result. Under reasonable reminiscence stress as much as depth 16, TriAttention performs comparably to Full Attention, whereas R-KV reveals catastrophic accuracy degradation — dropping from roughly 61% at depth 14 to 31% at depth 16. This signifies R-KV incorrectly evicts essential intermediate reasoning states.

On throughput, TriAttention achieves 1,405 tokens per second on MATH 500 in opposition to Full Attention’s 223 tokens per second, a 6.3× speedup. On AIME25, it achieves 563.5 tokens per second in opposition to 222.8, a 2.5× speedup at matched accuracy.

Generalization Beyond Mathematical Reasoning

The outcomes lengthen nicely past math benchmarks. On LongBench — a 16-subtask benchmark protecting query answering, summarization, few-shot classification, retrieval, counting, and code duties — TriAttention achieves the very best common rating of 48.1 amongst all compression strategies at a 50% KV funds on Qwen3-8B, profitable 11 out of 16 subtasks and surpassing the following greatest baseline, Ada-KV+SnapKV, by 2.5 factors. On the RULER retrieval benchmark at a 4K context size, TriAttention achieves 66.1, a ten.5-point hole over SnapKV. These outcomes verify that the strategy is just not tuned to mathematical reasoning alone — the underlying Q/Okay focus phenomenon transfers to basic language duties.

Key Takeaways

• Existing KV cache compression strategies have a basic blind spot: Methods like SnapKV and R-KV estimate token significance utilizing latest post-RoPE queries, however as a result of RoPE rotates question vectors with place, solely a tiny window of queries is usable. This causes necessary tokens — particularly these wanted by retrieval heads — to be completely evicted earlier than they turn into essential.
• Pre-RoPE Query and Key vectors cluster round steady, fastened facilities throughout almost all consideration heads: This property, known as Q/Okay focus, holds no matter enter content material, token place, or area, and is constant throughout Qwen3, Qwen2.5, Llama3, and even Multi-head Latent Attention architectures like GLM-4.7-Flash.
• These steady facilities make consideration patterns mathematically predictable with out observing any stay queries: When Q/Okay vectors are concentrated, the eye rating between any question and key reduces to a perform that relies upon solely on their positional distance — encoded as a trigonometric sequence. TriAttention makes use of this to attain each cached key offline utilizing calibration knowledge alone.
• TriAttention matches Full Attention reasoning accuracy at a fraction of the reminiscence and compute value: On AIME25 with 32K-token era, it achieves 2.5× greater throughput or 10.7× KV reminiscence discount whereas matching Full Attention accuracy — almost doubling R-KV’s accuracy at the identical reminiscence funds throughout each AIME24 and AIME25.
• The technique generalizes past math and works on shopper {hardware}. TriAttention outperforms all baselines on LongBench throughout 16 basic NLP subtasks and on the RULER retrieval benchmark, and permits a 32B reasoning mannequin to run on a single 24GB RTX 4090 through OpenClaw — a process that causes out-of-memory errors underneath Full Attention.

---

Check out the Paper, Repo and Project Page. Also, be at liberty to observe us on Twitter and don’t neglect to hitch our 120k+ ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.

Need to associate with us for selling your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar and so forth.? Connect with us

The submit Researchers from MIT, NVIDIA, and Zhejiang University Propose TriAttention: A KV Cache Compression Method That Matches Full Attention at 2.5× Higher Throughput appeared first on MarkTechPost.