TriAttention: Efficient Long Reasoning with Trigonometric KV Compression

Weian Mao; Xi Lin; Wei Huang; Yuxin Xie; Tianfu Fu; Bohan Zhuang; Song Han; Yukang Chen

cs.CLApr 6, 2026

TriAttention: Efficient Long Reasoning with Trigonometric KV Compression

Authors: Weian Mao, Xi Lin, Wei Huang, Yuxin Xie, Tianfu Fu, Bohan Zhuang, Song Han, Yukang Chen

Organizations: MIT · ZJU · NVIDIA

Abstract

Extended reasoning in large language models (LLMs) 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.5x higher throughput or 10.7x KV memory reduction, whereas leading baselines achieve only about half the accuracy at the same efficiency. TriAttention enables OpenClaw deployment on a single consumer GPU, where long context would otherwise cause out-of-memory with Full Attention.

Figures & tables

Figure 1: Performance trade-offs on AIME25 (Qwen3-8B). (A) At equivalent accuracy (40.8%), TriAttention achieves 2.5 $\times$ higher throughput than Full Attention. (B) TriAttention reduces KV cache memory by 10.7 $\times$ while matching Full Attention accuracy.

Figure 2: Q/K concentration and its implications for attention. (A) Pre-RoPE Q/K vectors at the dominant frequency band are highly concentrated (high Mean Resultant Length $R$ ). (B) RoPE rotation disperses these vectors into arc patterns. In (A-B), three distinct input sequences are overlayed, showing this structure is stable across content. (C) This concentration holds across nearly all heads. (D) When Q/K are concentrated, attention logits can be accurately reconstructed using a trigonometric series (Pearson $r=0.72$ ).

Figure 3: Attention reconstruction correlation across three DeepSeek-R1 distilled LLMs, including Qwen3 ( Qwen Team, 2025 ) , Qwen2.5 ( Qwen Team, 2024 ) , and Llama3 ( Dubey et al., 2024 ) . Distribution of per-head reconstruction Pearson correlation ( $\bar{r}$ ) across all attention heads. The red dashed line indicates the mean. All models show right-skewed distributions with means above 0.5.

Figure 4: Method overview. From left to right: offline calibration computes Q distribution centers; then during inference, original attention is scored by combining $S_{\text{trig}}$ and norm-based components; the rightmost panel shows the attention map after pruning. We observe that some heads exhibit distance preference—distant keys tend to receive higher attention. However, we also find that certain keys, despite being far from the query, receive little attention due to their low norms. This motivates our two scoring components: $S_{\text{trig}}$ captures distance preference, while the norm-based score identifies low-norm keys. In this example, $S_{\text{trig}}$ correctly assigns low scores to nearby keys, while the norm-based score identifies the earliest token (leftmost) as unimportant due to its low norm, despite its maximal distance. Together, they accurately identify tokens that will not be attended to and prune them. See Appendix D for visualizations with real attention maps.

Method	AIME24				AIME25
Method	Qwen3-8B	DS-Llama	DS-Qwen	GPT-OSS	Qwen3-8B	DS-Llama	DS-Qwen	GPT-OSS
Full Attention	57.1	50.4	43.8	69.2	40.8	31.4	34.2	60.0
SnapKV	34.6	5.0	34.6	48.3	20.0	6.7	25.0	36.7
R-KV	25.4	25.8	34.6	49.6	17.5	11.2	23.3	39.2
TriAttention	42.1	33.8	42.5	59.2	32.9	19.6	30.0	49.2

Table 1: Reasoning performance on AIME24 and AIME25. Best results are in bold , second best are underlined . All methods are compared under the same KV cache budget.

Figure 5: Performance comparison on Qwen3-8B. (A–C) Accuracy vs. KV cache budget on three mathematical reasoning benchmarks. TriAttention consistently outperforms R-KV across all budget levels. (D) Memory retention on Recursive State Query benchmark. Depth refers to DFS recursion depth; deeper recursion requires retaining more intermediate states, increasing memory pressure.

Method	Qwen3-8B	DS-Llama	DS-Qwen	GPT-OSS
Full Attention	69.6	82.4	87.0	91.4
SnapKV	49.2	65.5	66.4	68.2
R-KV	46.4	76.9	71.6	77.4
TriAttention	56.0	80.6	79.6	81.2

Table 2: Reasoning performance on MATH 500 with KV budget of 512. Best results are in bold , second best are underlined .

Table 3: Ablation studies on Qwen3-8B with KV budget of 2048. (A) Effect of removing trigonometric series score

S_{\text{trig}}

. (B) Effect of removing concentration-based weighting. (C) Cross-domain generalization: comparing calibration on coding vs. reasoning data, both tested on reasoning.

Metric	MATH 500		AIME24		AIME25
Metric	Full Attn.	TriAttn.	Full Attn.	TriAttn.	Full Attn.	TriAttn.
KV Budget	–	1024	–	4096	–	3072
Accuracy	69.6	68.4	57.1	54.6	40.8	40.8
Throughput	222.8	1405.2	222.8	413.9	222.8	563.5
Speedup	6.3 $\times$		1.9 $\times$		2.5 $\times$

Table 4: Throughput comparison on Qwen3-8B. At KV budgets achieving comparable accuracy to Full Attention, TriAttention shows significant throughput gains.

Metric	Comparable Accuracy		Comparable Memory
Metric	R-KV	TriAttn.	R-KV	TriAttn.
KV Budget	2048	1024	1024	1024
MATH 500	68.2	68.4	60.4	68.4 ↑8.0
AIME24	25.4	25.8	10.4	25.8 ↑15.4
Throughput	760.4	1405.2 ↑644.88	1345.5	1405.2

Table 5: Comparison between R-KV and TriAttention at comparable reasoning accuracy and memory.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Figure A : Evaluating memory via recursive simulation. Left : With complete memory, all intermediate states are retained and correct values propagate upward. Right : When an intermediate state is lost (State 2 ), the error propagates through all subsequent return values, corrupting the final result.

Figure B : Method visualization with real attention maps, corresponding to the schematic in Figure 4 . Top row : The four stages of TriAttention. (1) We compute the Q/K centers $\mathbb{E}[q]$ , $\mathbb{E}[k]$ from pre-RoPE distributions. (2) Using the trigonometric series, we compute $S_{\text{trig}}$ which scores keys based on distance preference. (3) We add the norm-based score $S_{\text{norm}}$ , weighted by concentration, to obtain the final score $\bar{S}(k)$ . (4) We retain top-scoring keys and evict the rest. Bottom row : Real attention maps illustrating the scenario described in Figure 4 . From left to right: original attention pattern showing distance preference; $S_{\text{trig}}$ visualization capturing the diagonal structure; combined score $\bar{S}$ incorporating norm information; attention after KV cache pruning, preserving the essential pattern.

Method	10%	20%	30%
FullKV	46.7
H2O ∗	–	–	33.3
TOVA ∗	–	–	36.7
RaaS ∗	–	–	36.7
R-KV ∗	–	–	43.3
LazyEviction	33.3	40.0	43.3

Appendix

Table A : Comparison on AIME24 (DeepSeek-R1-Distill-Qwen-7B) at varying KV budgets. ∗ Results cited from LazyEviction.

	QA						Summarization			Few-shot			Retr.	Count	Code
Method	NarrQA	Qasp	MFQA	HpQA	2Wik	Musi	GovR	QMSu	MNew	TREC	TriQA	SSum	PaRe	PaCn	LCC	ReBe	Avg
Full Attn.	28.8	43.8	55.3	62.8	48.9	35.5	33.5	24.7	24.7	40.5	90.5	40.3	91.8	9.0	64.9	60.0	47.2
SnapKV	26.9	37.0	45.5	59.2	46.3	33.1	32.2	23.0	23.4	40.5	89.9	41.0	91.1	7.5	66.4	59.9	45.2
PyramidKV	25.9	30.4	39.1	52.2	39.9	29.7	29.9	22.2	21.9	33.5	90.0	40.7	93.4	8.0	65.7	60.2	42.7
StreamingLLM	24.1	30.5	31.2	46.5	41.6	21.6	30.8	21.8	23.9	43.0	85.4	38.2	55.2	10.0	64.7	61.2	39.4
KnormPress	17.6	24.4	40.3	29.2	26.4	14.9	28.8	21.6	20.9	50.0	81.6	41.2	79.1	7.1	31.7	47.5	35.1

Appendix

Table B : LongBench results (Qwen3-8B, 50% KV budget). 16 subtasks spanning QA, summarization, few-shot classification, retrieval, counting, and code. Bold = best among compression methods.

Method	RULER Avg
SnapKV	55.6
PyramidKV	40.7
StreamingLLM	61.1
TriAttention	66.1

Appendix

Table C : RULER retrieval results (Qwen3-8B, 50% KV, 4K context). Bold = best among compression methods.

Method	Qasp	HpQA	2Wik	Musi	GovR	QMSu	MNew	TREC	TriQA	SSum	NarrQA	MFQA	Avg
H2O	39.2	50.7	43.9	30.7	32.9	23.4	24.4	56.5	89.1	39.1	21.2	45.4	41.4
TriAttention	43.0	60.2	44.9	36.9	32.9	23.8	24.3	69.0	90.3	39.9	28.1	51.4	45.4

Appendix

Table D : Comparison with H2O on LongBench subtasks where H2O fits in 48GB GPU memory (Qwen3-8B, 50% KV). H2O requires materializing the full

O(n^{2})

attention matrix and cannot use FlashAttention. Bold = best.

Max Dist	#Offsets	Acc
128	8	41.7
4096	13	48.8
8192	14	46.2
65536	17	45.8
Linear spacing		28.7
Geometric spacing		45.8

Appendix

Table E : Future offset ablation on Qwen3-8B (AIME24). Top: effect of offset range. Bottom: spacing strategy comparison (17 offsets, range [1, 65536]).

Calibration	Acc
50k tokens	45.4
200k tokens	45.8
960k tokens	45.8
HTML (low qual.)	46.2
Code (mid qual.)	43.3
Chat (high qual.)	46.7

Appendix

Table F : Calibration data sensitivity on Qwen3-8B (AIME24). Top: effect of calibration data size. Bottom: effect of calibration data quality.

Table G : Cross-architecture comparison of reconstruction quality (Pearson

r

) and Q/K concentration (MRL) between GQA (Qwen3-8B) and MLA (GLM-4.7-Flash) architectures.

Figure C : OpenClaw demo on a single RTX 4090 with Qwen3-32B (INT4). Full attention runs out of memory during multi-turn interaction, while TriAttention completes the task within the GPU memory budget. Full video is available on our GitHub page: https://github.com/WeianMao/triattention .

Long-Context Reasoning in Language Models

See all 41 papers

BeaconKV: Key-Value Cache Compression Guided by Beacon Queries for Efficient Large Reasoning Model Inference
Sep 4, 2026Janghyeon Kim, Minsoo Kim, Kyuhong Shim +1Long-Context Reasoning in Language Models KV-Cache Compression
Hanyang University · Sungkyunkwan University
Random Attention: Rethinking KV Cache Eviction for Efficient Reasoning
Sep 3, 2026Heng Wang, Jielin Qiu, Wenting Zhao +8KV Caching Memory-Efficient Optimization
Salesforce AI Research · University of Illinois Urbana-Champaign
LongFlow: Efficient KV Cache Compression for Reasoning Models
Mar 12, 2026Yi Su, Zhenxu Tian, Dan Qiao +3KV-Cache Compression LLM Inference Acceleration
School of Computer Science and Technology, Soochow University, China · ByteDance

Sep 4, 2026J/K move · Enter open · S save

BeaconKV: Key-Value Cache Compression Guided by Beacon Queries for Efficient Large Reasoning Model Inference

Janghyeon Kim, Minsoo Kim, Kyuhong Shim, Jungwook Choi

Hanyang University · Sungkyunkwan University

Long-Context Reasoning in Language Models KV-Cache Compression LLM Inference Acceleration Large Reasoning Models

Open paper AskPDF arXiv

Large Reasoning Models (LRMs) achieve superior problem-solving through extended Chain-of-Thought (CoT) generation, but the resulting key-value (KV) cache grows linearly with sequence length and creates severe memory bottlenecks, often exceeding GPU capacity for long reasoning traces. Existing KV cache compression methods rely on recent queries to estimate future token importance, implicitly assuming these serve as reliable proxies for future attention patterns. We demonstrate that this assumption fails in long-horizon reasoning: certain decoding steps generate Thought Revisiting Tokens (TRT) that re-attend to distant previous context, such as task-solving plans formulated early in the trace. Through systematic analysis, we discover that queries corresponding to the TRT cluster into a small number of similarity groups in the embedding space. Based on this insight, we propose BeaconKV, a training-free KV cache compression method that maintains beacon queries, compact representatives for each global query cluster, to anticipate which KV pairs will be revisited without storing the entire query history. Across four open-source LRMs and diverse reasoning benchmarks, BeaconKV generally outperforms existing compression methods, achieving up to

5.8\times

memory reduction while nearly preserving full cache accuracy and improving throughput by over

4.3\times

Abstract

Long-Context Reasoning in Language Models

Explore similar work

BeaconKV: Key-Value Cache Compression Guided by Beacon Queries for Efficient Large Reasoning Model Inference

Random Attention: Rethinking KV Cache Eviction for Efficient Reasoning

LongFlow: Efficient KV Cache Compression for Reasoning Models