arXiv preprint ·
Recursive Sparse Attention: Million-Token Length Extrapolation at Logarithmic Cost
- Published
- Source
- arXiv preprint, posted before peer review
- Original
- Read it at arXiv
- Topics
Abstract
At context lengths of about 1M tokens, dense attention is costly in prefill and decoding, and training typically relies on long documents, which become scarcer as the target length increases. Most sparse attention methods score all preceding blocks or tokens, so their selection cost per query is $O(L)$ at context length $L$; hierarchical selection reduces it to $O(\log L)$, but prior trainable hierarchical methods report accuracy only within their training length. We propose RecSA (Recursive Sparse Attention), which builds a tree over the KV cache in which each node summarizes its subtree with one key and one mean-pooled value. The keys come from one pooler per layer, shared by all levels and distilled from dense attention. Each level above the tokens selects at most $k$ candidates, so each query has $O(\log L)$ selection cost and reads $O(\log L)$ entries; the summaries of unselected candidates share one softmax with the exact tokens. Each entry's RoPE position is its index among the candidates of its level, so at any length every relative distance stays within the range seen during training. Initialized from Qwen3-4B and trained on 0.62B tokens at the original model's 32K training length, RecSA outperforms all baselines on RULER at every length from 131K to 2M; at 1M, it scores 60.5, compared with 40.0 for the best baseline, Dense + DCA. With our optimized kernels, RecSA is faster than dense attention in end-to-end prefill and decoding from 64K to 2M; at 524K on one H200 GPU, it accelerates prefill by 3.2x and decoding by 3.8x.
The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.
Read it on ReadSearch
Each opens after you create an account, and brings you back to this paper.
More in Computation & Language / NLP
How Well Can Large Language Models Impersonate Individual Authors? A Comparative Evaluation
HyFlash: A Prefix-Memory Prior for On-Device Speculative Decoding
Phonological Knowledge of Chinese Characters in Language Models
Encoding and Propagation of Grammatical Number in Large Language Models
Matching the Training-Time Weight Rule in Inference-Time Expert Deactivation
Get the few papers worth your time in Computation & Language / NLP every morning.
ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.