TriAttention: Efficient Long Reasoning with Trigonometric KV Compression
Abstract
Extended reasoning in large language models (LLMs) requires long and accurate decoding and creates severe KV cache memory bottlenecks. Leading KV cache compression methods estimate KV importance using attention scores from recent post-RoPE queries. However, queries rotate with position during RoPE, making representative queries very few, leading to poor top-key selection and unstable reasoning. To avoid this issue, we turn to the pre-RoPE space, where we observe that Q and K vectors are highly concentrated around fixed non-zero centers and remain stable across positions—Q/K concentration. We show that this concentration causes queries to preferentially attend to keys at specific distances (e.g., nearest keys), with the centers determining which distances are preferred via a trigonometric series. Based on this, we propose TriAttention to estimate key importance by leveraging these centers. Via the trigonometric series, we use the distance preference characterized by these centers to score keys according to their positions, and also leverage Q/K norms as an additional signal for importance estimation. On AIME25 with 32K-token generation, TriAttention matches Full Attention reasoning accuracy while achieving 2.5× higher throughput or 10.7× KV memory reduction, whereas leading baselines achieve only about half the accuracy at the same efficiency.
Lay Summary
Modern AI systems like ChatGPT can reason through complex problems step by step, but this requires storing a large amount of intermediate information in memory—a component called the "KV cache." As reasoning chains get longer, this memory requirement grows dramatically, making such systems expensive and slow. Existing methods try to solve this by discarding less important parts of the memory. However, they rely on clues that shift and become unreliable as the AI generates more text, leading to poor decisions about what to keep. We discovered a surprisingly stable pattern in how these AI models process information: the internal representations that guide attention remain consistent throughout the reasoning process. By exploiting this stability, our method—TriAttention—can predict which parts of memory are important without observing the actual computation, using a mathematical structure similar to a trigonometric series. The result: TriAttention can cut memory usage by over 10× or increase processing speed by 2.5×, while maintaining the same reasoning accuracy as using the full memory. This makes long, complex reasoning tasks significantly more efficient and accessible.