Long-chain reasoning is one of the most compute-intensive tasks in modern large language models. When a model like DeepSeek-R1 or Qwen3 works through a complex math problem, it can generate tens of thousands of tokens before arriving at an answer. Every one of those tokens must be stored in what is called the KV cache — a memory structure that holds the Key and Value vectors the model needs to attend back to during generation. The longer the reasoning chain, the larger the KV cache grows, and for many deployment scenarios, especially on consumer hardware, this growth eventually exhausts GPU memory entirely.

A team of researchers from MIT, NVIDIA, and Zhejiang University proposed a method called TriAttention that directly addresses this problem. On the AIME25 mathematical reasoning benchmark with 32K-token generation, TriAttention matches Full Attention accuracy while achieving 2.5× higher throughput or 10.7× KV memory reduction. Leading baselines achieve only about half the accuracy at the same efficiency level.

The Problem with Existing KV Cache Compression

To understand why TriAttention is important, it helps to understand the standard approach to KV cache compression. Most existing methods — including SnapKV, H2O, and R-KV — work by estimating which tokens in the KV cache are important and evicting the rest. Importance is typically estimated by looking at attention scores: if a key receives high attention from recent queries, it is considered important and kept.

The catch is that these methods operate in what the research team calls post-RoPE space. RoPE, or Rotary Position Embedding, is the positional encoding scheme used by most modern LLMs including Llama, Qwen, and Mistral. RoPE encodes position by rotating the Query and Key vectors in a frequency-dependent way. As a result, a query vector at position 10,000 looks very different from the same semantic query at position 100, because its direction has been rotated by the position encoding.

This rotation means that only the most recently generated queries have orientations that are ‘up to date’ for estimating which keys are important right now. Prior work has confirmed this empirically: increasing the observation window for importance estimation does not help — performance peaks at around 25 queries and declines after that. With such a tiny window, some keys that will become important later get permanently evicted.

This problem is especially acute for what the research team calls retrieval heads — attention heads whose function is to retrieve specific factual tokens from long contexts. The relevant tokens for a retrieval head can remain dormant for thousands of tokens before suddenly becoming essential to the reasoning chain. Post-RoPE methods, operating over a narrow observation window, see low attention on those tokens during the dormant period and permanently evict them. When the model later needs to recall that information, it is already gone, and the chain of thought breaks.

The Pre-RoPE Observation: Q/K Concentration

The key insight in TriAttention comes from looking at Query and Key vectors before RoPE rotation is applied — the pre-RoPE space. When the research team visualized Q and K vectors in this space, they found something consistent and striking: across the vast majority of attention heads and across multiple model architectures, both Q and K vectors cluster tightly around fixed, non-zero center points. The research team terms this property Q/K concentration, and measures it using the Mean Resultant Length R — a standard directional statistics measure where R → 1 means tight clustering and R → 0 means dispersion in all directions.

On Qwen3-8B, approximately 90% of attention heads exhibit R > 0.95, meaning their pre-RoPE Q/K vectors are nearly perfectly concentrated around their respective centers. Critically, these centers are stable across different token positions and across different input sequences — they are an intrinsic property of the model’s learned weights, not a property of any particular input. The research team further confirm that Q/K concentration is domain-agnostic: measuring Mean Resultant Length across Math, Coding, and Chat domains on Qwen3-8B yields nearly identical values of 0.977–0.980.

This stability is what post-RoPE methods cannot exploit. RoPE rotation disperses these concentrated vectors into arc patterns that vary with position. But in pre-RoPE space, the centers remain fixed.

From Concentration to a Trigonometric Series

The research team then show mathematically that when Q and K vectors are concentrated around their centers, the attention logit — the raw score before softmax that determines how much a query attends to a key — simplifies dramatically. Substituting the Q/K centers into the RoPE attention formula, the logit reduces to a function that depends only on the Q-K distance (the relative positional gap between query and key), expressed as a trigonometric series:

$\text{logit}(\Delta) \approx \sum_{f} \underbrace{\|\bar{q}_f\| \|\bar{k}_f\|}_{\text{amplitude}} \cos(\omega_f \Delta + \underbrace{\bar{\phi}_f}_{\text{phase}}) = \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 each frequency band f, and the coefficients a_f and b_f are determined by the Q/K centers. This series produces a characteristic attention-vs-distance curve for each head. Some heads prefer nearby keys (local attention), others prefer very distant keys (attention sinks). The centers, computed offline from calibration data, fully determine which distances are preferred.

The research team validated this experimentally across 1,152 attention heads in Qwen3-8B and across Qwen2.5 and Llama3 architectures. The Pearson correlation between the predicted trigonometric curve and the actual attention logits has a mean above 0.5 across all heads, with many heads achieving correlations of 0.6–0.9. The research team further validates this on GLM-4.7-Flash, which uses Multi-head Latent Attention (MLA) rather than standard Grouped-Query Attention — a meaningfully different attention architecture. On MLA, 96.6% of heads exhibit R > 0.95, compared to 84.7% for GQA, confirming that Q/K concentration is not specific to one attention design but is a general property of modern LLMs.

How TriAttention Uses This

TriAttention is a KV cache compression method that uses these findings to score keys without needing any live query observations. The scoring function has two components:

The Trigonometric Series Score (S_trig) uses the Q center computed offline and the actual cached key representation to estimate how much attention the key will receive, based on its positional distance from future queries. Because a key may be attended to by queries at many future positions, TriAttention averages this score over a set of future offsets using geometric spacing.

$S_{\text{trig}}(k, \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 attention heads where Q/K concentration is lower. It weights each frequency band by the expected query norm contribution, providing complementary information about token salience beyond distance preference alone.

$S_{\text{norm}}^{(0)}(k) = \sum_{f} \mathbb{E}[\|q_f\|] \cdot \|k_f\|$

The two scores are combined using the Mean Resultant Length R as an adaptive weight: when concentration is high, S_trig dominates; when concentration is lower, S_norm contributes more. Every 128 generated tokens, TriAttention scores all keys in the cache and retains only the top-B, evicting the rest.

Results on Mathematical Reasoning

On AIME24 with Qwen3-8B, TriAttention achieves 42.1% accuracy against Full Attention’s 57.1%, while R-KV achieves only 25.4% at the same KV budget of 2,048 tokens. On AIME25, TriAttention achieves 32.9% versus R-KV’s 17.5% — a 15.4 percentage point gap. On MATH 500 with only 1,024 tokens in the KV cache out of a possible 32,768, TriAttention achieves 68.4% accuracy against Full Attention’s 69.6%.

The research team also introduces a Recursive State Query benchmark based on recursive simulation using depth-first search. Recursive tasks stress memory retention because the model must maintain intermediate states across long chains and backtrack to them later — if any intermediate state is evicted, the error propagates through all subsequent return values, corrupting the final result. Under moderate memory pressure up to depth 16, TriAttention performs comparably to Full Attention, while R-KV shows catastrophic accuracy degradation — dropping from approximately 61% at depth 14 to 31% at depth 16. This indicates R-KV incorrectly evicts critical intermediate reasoning states.

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

Generalization Beyond Mathematical Reasoning

The results extend well beyond math benchmarks. On LongBench — a 16-subtask benchmark covering question answering, summarization, few-shot classification, retrieval, counting, and code tasks — TriAttention achieves the highest average score of 48.1 among all compression methods at a 50% KV budget on Qwen3-8B, winning 11 out of 16 subtasks and surpassing the next best baseline, Ada-KV+SnapKV, by 2.5 points. On the RULER retrieval benchmark at a 4K context length, TriAttention achieves 66.1, a 10.5-point gap over SnapKV. These results confirm that the method is not tuned to mathematical reasoning alone — the underlying Q/K concentration phenomenon transfers to general language tasks.

Key Takeaways

• Existing KV cache compression methods have a fundamental blind spot: Methods like SnapKV and R-KV estimate token importance using recent post-RoPE queries, but because RoPE rotates query vectors with position, only a tiny window of queries is usable. This causes important tokens — especially those needed by retrieval heads — to be permanently evicted before they become critical.
• Pre-RoPE Query and Key vectors cluster around stable, fixed centers across nearly all attention heads: This property, called Q/K concentration, holds regardless of input content, token position, or domain, and is consistent across Qwen3, Qwen2.5, Llama3, and even Multi-head Latent Attention architectures like GLM-4.7-Flash.
• These stable centers make attention patterns mathematically predictable without observing any live queries: When Q/K vectors are concentrated, the attention score between any query and key reduces to a function that depends only on their positional distance — encoded as a trigonometric series. TriAttention uses this to score every cached key offline using calibration data alone.
• TriAttention matches Full Attention reasoning accuracy at a fraction of the memory and compute cost: On AIME25 with 32K-token generation, it achieves 2.5× higher throughput or 10.7× KV memory reduction while matching Full Attention accuracy — nearly doubling R-KV’s accuracy at the same memory budget across both AIME24 and AIME25.
• The method generalizes beyond math and works on consumer hardware. TriAttention outperforms all baselines on LongBench across 16 general NLP subtasks and on the RULER retrieval benchmark, and enables a 32B reasoning model to run on a single 24GB RTX 4090 via OpenClaw — a task that causes out-of-memory errors under Full Attention.

Check out the Paper, Repo and Project Page. Also, feel free to follow us on Twitter and don’t forget to join 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 partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us