Recursive Sparse Attention: Million-Token Length Extrapolation at Logarithmic Cost

Summary (Overview)

• RecSA (Recursive Sparse Attention) is a novel sparse attention mechanism that builds a tree over the KV cache, enabling context length extrapolation to 2M tokens with only O(log L) selection cost and read complexity per query.
• The method uses candidate-local RoPE, where each entry's position is its index within its level's candidate list, ensuring all relative distances stay within the training range (below 4,096 with default k=256) regardless of context length.
• RecSA is trained only at 32K tokens (0.62B tokens total) using a two-stage approach: pooler distillation from dense attention, followed by backbone adaptation—no training at target lengths is needed.
• On RULER benchmark, RecSA outperforms all baselines at every length from 131K to 2M, scoring 60.5 at 1M vs. 40.0 for the best baseline (Dense + DCA), while using only 1.16% of the context as token budget.
• Custom Triton/CUDA kernels provide 3.2× prefill and 3.8× decoding speedup at 524K on one H200 GPU compared to dense FlashAttention, with decoding latency staying nearly constant from 131K to 524K.

Introduction and Theoretical Foundation

Background and Motivation

State-of-the-art language models now support context windows of ~1M tokens, but face two fundamental challenges:

1. Computational cost: Dense attention grows quadratically in prefill and reads the entire KV cache per token in decoding.
2. Data scarcity: Training at target lengths requires sufficiently long documents, which become scarcer as length increases.

Limitations of Existing Approaches

Most sparse attention methods (MoBA, NSA, DSA) score all preceding blocks/tokens, making selection cost O(L) per query. While hierarchical selection (HSA, HiLS, InfLLM, InfiniteHiP) reduces this to O(log L), prior trainable hierarchical methods report accuracy only within their training length. Methods like HiP and PISA that keep original key positions suffer from out-of-distribution relative distances beyond training length.

Key Theoretical Insight

The paper introduces candidate-local RoPE: by assigning each entry its index within its level's candidate list (rather than absolute position), every relative distance to the query stays below the maximum candidate list size (C·k = 4,096 with defaults). This eliminates the need for context extension methods like YaRN or DCA.

Methodology

Tree Construction

Each attention layer maintains a KV tree per KV head. Level 0 holds exact token keys/values; groups of $C = 16$ consecutive entries form parent nodes at higher levels. Each node's key is computed by a learnable pooler, and its value is the mean of children's values:

$$\begin{array}{l} \mathbf {K} _ {i} ^ {(\ell)} = \mathrm{Pool} \big (\mathbf {K} _ {C i} ^ {(\ell - 1)}, \ldots , \mathbf {K} _ {C i + C - 1} ^ {(\ell - 1)} \big), \\ \mathbf {V} _ {i} ^ {(\ell)} = \frac {1}{C} \sum_ {c = 0} ^ {C - 1} \mathbf {V} _ {C i + c} ^ {(\ell - 1)}, \end{array}\tag{1}$$

Hierarchical Selection

Selection proceeds top-down. At each level $\ell \geq 1$, the query scores all candidates and selects the current node plus the $k-1$ highest-scoring others (k=256 default). The children of selected nodes form the next candidate list:

$$\begin{array}{r l} & {\mathcal {A} _ {t} ^ {(D - 1)} = [ 0, \ldots , a _ {t} ^ {(D - 1)} ],} \\ & {\qquad \mathcal {S} _ {t} ^ {(\ell)} = \{a _ {t} ^ {(\ell)} \} \cup \mathrm{TopK} _ {k - 1} \big (\{I _ {t, i} ^ {(\ell)} \} _ {i \in \mathcal {A} _ {t} ^ {(\ell)} \setminus \{a _ {t} ^ {(\ell)} \}} \big),} \\ & {\qquad \mathcal {A} _ {t} ^ {(\ell - 1)} = \underset {i \in \mathcal {S} _ {t} ^ {(\ell)}} {\mathrm{concat}} [ C i, \ldots , \min (C i + C - 1, a _ {t} ^ {(\ell - 1)}) ],} \end{array}\tag{3}$$

Unified Softmax

Summaries of unselected candidates and exact tokens share one softmax:

$$\begin{array}{c} \mathcal {E} _ {t} = (\{0 \} \times \mathcal {A} _ {t} ^ {(0)}) \cup \bigcup_ {\ell = 1} ^ {D - 1} (\{\ell \} \times (\mathcal {A} _ {t} ^ {(\ell)} \setminus \mathcal {S} _ {t} ^ {(\ell)})), \\ \mathbf {o} _ {t, h} = \frac {\sum_ {e \in \mathcal {E} _ {t}} \exp (s _ {t , h , e}) \mathbf {V} _ {e}}{\sum_ {e \in \mathcal {E} _ {t}} \exp (s _ {t , h , e})}, \end{array}\tag{4}$$

Pooler Distillation

The pooler computes node keys from children's keys with a gated MLP architecture:

$$\begin{array}{l} \mathbf {K} = \boldsymbol {\mu} + g (\mathbf {p} - \boldsymbol {\mu}) + \sigma \mathbf {m}, \\ \alpha_ {c} = \operatorname{softmax} _ {c} \big (\tau \mathbf {w} ^ {\top} (\mathbf {x} _ {c} - \boldsymbol {\mu}) / \sigma \big), \mathbf {p} = \sum_ {c} \alpha_ {c} \mathbf {x} _ {c}, \\ \mathbf {m} = W _ {2} \text { SiLU } \big (W _ {1} [ \boldsymbol {\mu} / \rho ; (\mathbf {p} - \boldsymbol {\mu}) / \sigma ] + \mathbf {b} _ {1} \big), \end{array}\tag{5}$$

Training uses KL divergence between student importance distribution and teacher (dense attention mass) distribution:

$$\mathcal {L} _ {\mathrm{KL}} = \sum_ {t} \sum_ {\ell = 1} ^ {D - 1} \mathrm{KL} \big (\mathbf {y} _ {t} ^ {(\ell)} \| \boldsymbol {\pi} _ {t} ^ {(\ell)} \big).\tag{6}$$

A shift penalty prevents the MLP from learning positional biases:

$$\mathcal {L} _ {\text { shift }} = \sum_ {\ell = 1} ^ {D - 1} \left\| \frac {1}{M _ {\ell}} \sum_ {i} \sigma_ {i} \mathbf {m} _ {i} \right\| ^ {2} / \left(\frac {1}{M _ {\ell}} \sum_ {i} \| \boldsymbol {\mu} _ {i} \| ^ {2}\right),\tag{7}$$

Training Stages

• Stage 1 (Pooler): Freeze backbone, train pooler with k=64 for deeper trees. Two parts: (a) train w, τ, g with MLP off; (b) train MLP parameters with KL + shift penalty.
• Stage 2 (Backbone): Freeze pooler, train backbone under RecSA with k=256 using language modeling loss only.

Empirical Validation / Results

RULER Benchmark Results

Table 1: RULER score by context length

Perplexity Results

Table 3: Perplexity on PG-19 books

RecSA maintains nearly constant perplexity (10.87–11.12) from 131K to 1M, while all baselines degrade significantly.

Key Findings from Ablations

• Summaries improve average score from 70.4 to 72.3 (131K–1M average)
• Candidate-local RoPE is critical: standard RoPE with YaRN scores 1.9 at 524K vs. RecSA's 69.7
• Distillation + shift penalty + MLP term each contribute to long-context performance
• Increasing k to 512 improves scores at every length (with 1.7×–1.9× token budget)

Efficiency Results

• Prefill: 3.2× speedup at 524K (80s vs. 257s for Dense)
• Decoding: 3.8× speedup at 524K (6.1ms vs. 23.4ms per token); nearly constant latency from 131K–524K
• At 1M/2M: 8.35/10.7 ms per token via KV offloading to host memory

Theoretical and Practical Implications

Theoretical Contributions

1. Logarithmic complexity: RecSA achieves O(log L) selection cost and read complexity per query, a fundamental improvement over the O(L) cost of most sparse attention methods.

2. Length extrapolation without context extension: Candidate-local RoPE eliminates out-of-distribution relative distances entirely, removing the need for YaRN, DCA, or similar techniques.

3. Training efficiency: Only 0.62B tokens at 32K length are needed—no long-document training data required.

Practical Implications

1. Scalable inference: The fixed launch grid and CUDA graph capture in decoding make latency nearly constant across a wide range of context lengths.

2. Memory efficiency: KV offloading keeps GPU memory usage bounded (only levels above level 0, <7% of full KV cache, plus gathered exact tokens).

3. Model-agnostic: The approach works with pretrained models (demonstrated on Qwen3-4B) and can be extended to larger models.

Conclusion

RecSA demonstrates that hierarchical sparse attention with distilled poolers and candidate-local positional encoding can achieve million-token context extrapolation with logarithmic cost, trained only at short lengths. Key takeaways:

• Tree-based hierarchical selection with summaries sharing one softmax with exact tokens provides both accuracy and efficiency
• Candidate-local RoPE eliminates the need for context extension methods by keeping all relative distances within training range
• Two-stage training (pooler distillation + backbone adaptation) enables effective length extrapolation from 32K training data
• Custom kernels with slot-based rotation splitting enable efficient prefill and decoding implementations

Future Directions

• Applying RecSA to larger models and other architectures
• Exploring adaptive k values and branching factors
• Investigating the approach for other attention variants (e.g., sliding window combinations)
• Extending to even longer contexts (beyond 2M tokens)