Source-linked AI summary
SeerAttention-R: Sparse Attention Adaptation for Long Reasoning
Yizhao Gao, Shuming Guo, Shijie Cao, Yuqing Xia, Yu Cheng, Lei Wang, Lingxiao Ma, Yutao Sun, Tianzhu Ye, Li Dong, Hayden Kwok-Hay So, Yu Hua, Ting Cao, Fan Yang, Mao Yang
TL;DR
Long reasoning sequences make decoding increasingly costly, and sparse attention in prolonged reasoning-model decoding remains underexplored. SeerAttention-R adapts self-distilled gating to autoregressive decoding and integrates as a lightweight plug-in without changing pretrained parameters. It maintains near-lossless accuracy under a 4K token budget, while its TileLang kernel achieves up to 8.6× speedup over FlashAttention-3 at 90% sparsity.
Problem
Autoregressive decoding makes later tokens attend to longer contexts, raising compute and KV-cache demands, while sparse attention for prolonged reasoning-model decoding remains underexplored.
Source: §1 · p. 1Method
SeerAttention-R adapts self-distilled attention gating for autoregressive decoding and adds a plug-in gate without modifying pretrained model parameters.
Source: §1 · p. 2Results
SeerAttention-R consistently outperforms Quest and maintains near-lossless accuracy with a 4K token budget, while TileLang achieves up to 8.6× speedup over FlashAttention-3 at 90% sparsity.
Source: §1 · p. 2Takeaways & Limitations
The framework supports coarse-grained attention blocks, including sizes 64 and 128, while retaining near-lossless reasoning accuracy in the reported evaluation.
Source: §1 · p. 2, §8 · p. 13The work focuses on sparse-decoding accuracy and kernel-level speedup; end-to-end system support and optimization are left for future work.
Source: §6.1 · p. 11
Abstract
from arXiv · showhide
We introduce SeerAttention-R, a sparse attention framework specifically tailored for the long decoding of reasoning models. Extended from SeerAttention, SeerAttention-R retains the design of learning attention sparsity through a self-distilled gating mechanism, while removing query pooling to accommodate auto-regressive decoding. With a lightweight plug-in gating, SeerAttention-R is flexible and can be easily integrated into existing pretrained model without modifying the original parameters. We demonstrate that SeerAttention-R, trained on just 0.4B tokens, maintains near-lossless reasoning accuracy with 4K token budget in AIME benchmark under large sparse attention block sizes (64/128). Using TileLang, we develop a highly optimized sparse decoding kernel that achieves near-theoretical speedups of up to 9x over FlashAttention-3 on H100 GPU at 90% sparsity. Code is available at: https://github.com/microsoft/SeerAttention.
1 Introduction
Longer reasoning generations increase attention costs, while the use of sparse attention in prolonged reasoning-model decoding remains underexplored. SeerAttention-R adapts self-distilled gating for decoding, maintains near-lossless accuracy under a 4K token budget, and its TileLang kernel reaches up to 8.6× speedup over FlashAttention-3 at 90% sparsity.
1 Introduction
- Longer reasoning generations raise efficiency demands because per-token generation cost grows linearly with context length while total generation cost grows quadratically.The authors identify attention sparsity as underexplored in reasoning models and report that oracle sparsity can preserve reasoning capability by activating only a subset of important tokens.Source: §1 · p. 1
- SeerAttention-R adapts self-distilled attention gating for autoregressive decoding, removing query sequence pooling and sharing sparsity selection across GQA groups.The plug-in gate can be added to pretrained transformer models without modifying their original parameters.Source: §1 · p. 2
- Across reasoning models and benchmarks, SeerAttention-R consistently outperforms Quest and maintains near-lossless accuracy with a 4K token budget.The reported accuracy gap further diminishes as model size increases; training the gating module uses 0.4B tokens.Source: §1 · p. 2
- At batch size 16 and sequence length ≥32k, TileLang achieves up to 8.6× speedup over FlashAttention-3 and 1.7× over Triton at 90% sparsity.The gains are especially pronounced at large sequence lengths and batch sizes.Source: §1 · p. 2
2 SeerAttention-R
SeerAttention-R adapts SeerAttention’s self-distilled attention gate to autoregressive decoding, with changes to query handling, GQA sparsity sharing, and training targets. The gate predicts block-level attention scores while the pretrained model’s weights remain frozen.
2.1 A Recap of SeerAttention
- SeerAttention uses self-distilled AttnGate scores to select important attention blocks, with pooled Q and K representations supporting long-context prefilling.The original model’s attention provides the ground truth for training the gate while its weights remain frozen.Source: §2.1 · pp. 2–3
2.2 SeerAttention-R: AttnGate for Sparse Decoding
- For decoding, SeerAttention-R leaves Q unpooled along the sequence dimension and projects grouped query heads to KV heads for shared GQA sparsity decisions.Its K branch instead uses non-overlapping block-level max, min, and average pooling before the linear layer; positional embedding in AttnGate consistently improves accuracy in the reported experiments.Source: §2.2 · p. 3
2.3 Distillation/Training
- Training distills AttnGate toward normalized block-level attention targets formed by column-wise 1D max pooling and max pooling within each query-head subgroup.A modified FlashAttention-2 kernel generates the target alongside attention output, reusing intermediate results to improve distillation efficiency.Source: §2.3 · p. 4, Figure 2 · p. 4
3 Inference of SeerAttention-R
At inference, SeerAttention-R turns AttnGate scores into selected block indices, caches compressed keys, and processes only selected blocks with a sparse decoding kernel. The inference design includes a fallback for the newest block while the cache awaits an update.
3.1 Sparsify Methods: Token Budget vs Threshold
- Token-budget sparsification converts the budget to a block count and selects the top-scoring blocks, whereas threshold sparsification selects scores above a threshold and allows head-specific sparsity ratios.Token budgets facilitate direct comparisons with other methods; thresholding is more self-adaptive, with efficiency-accuracy trade-offs between the approaches.Source: §3.1 · p. 5
3.2 K Compression Cache
- The K Compression Cache stores pooled-and-projected key representations and updates after each block of newly generated tokens; the last block remains active until the cache is current.With block size 64, the cache requires 1/128 (<1%) of the original KV-cache memory.Source: Figure 3 · p. 4, §3.2 · p. 5
3.3 Block Sparse Flash Decoding Kernel
- The sparse decoding kernel traverses only AttnGate-selected block indices and partitions work using the maximum selected-block count to balance irregular workloads.Its three-dimensional launch over batch, KV heads, and splits supports parallel work across query groups and KV shards.Source: §3.3 · p. 5
- The H100 TileLang implementation uses wgmma and architecture-targeted tiling, warp specialization, pipelining, and memory-layout optimizations.A Triton implementation uses the same scheduling strategy for comparison.Source: §3.3 · p. 5
4 Experiments
The experiments assess oracle sparsity, compare SeerAttention-R with full attention and Quest, and benchmark a sparse decoding kernel. SeerAttention-R maintains near-lossless accuracy at lower token budgets than Quest, while the TileLang kernel approaches theoretical speedups under high sparsity.
4.1 Experiments Setup
- Evaluation covers four open-source reasoning models and AIME24, AIME25, MATH-500, and GPQA-Diamond; only the AttnGate is trained, leaving original model weights unchanged.The training uses OpenR1-MATH-220K and packs inputs into sequences up to 32k tokens.Source: §4.1 · p. 6
4.2 Oracle Sparse Accuracy: How Sparse is Attention in Reasoning Models?
- Oracle sparse selection achieves lossless performance on all evaluated tasks at 2k tokens; at 1k, degradation is negligible with block sizes 32 or 64 but is observed on AIME24 and AIME25, particularly at 128.The oracle experiment evaluates Qwen3-14B with block sizes 32, 64, and 128; its results motivate block size 64 as the default.Source: Figure 4 · p. 6, §4.2 · p. 7
4.3 Results of SeerAttention-R and Quest
- SeerAttention-R typically reaches lossless AIME24 accuracy at 4k tokens, while Quest fails to reach lossless accuracy even at 8k; on MATH-500 and GPQA-Diamond, the respective budgets are 2k and around 8k.The comparison uses block size 64 and sparse attention in all layers for both methods.Source: Figure 5 · p. 7, §4.3 · p. 8
4.4 Kernel Speedup
- Up to 9× speedup at 0.9 sparsity is reported for the TileLang kernel at batch size 16 and sequence length ≥32k, approaching the theoretical speedup.TileLang also reports up to 6× at 0.9 sparsity for batch size 4 and sequence length 32k.Source: Figure 6 · p. 8, §4.4 · p. 9
5 Ablation Studies
The ablations examine block-size robustness, hybrid dense attention, sparsity selection, reasoning length, and training cost. SeerAttention-R's results include stable accuracy across tested block sizes, marginal benefit from hybridization, and lightweight gating training.
5.1 Block Size for Sparse Attention
- SeerAttention-R maintains accurate sparse block selection across tested block sizes, while Quest accuracy declines as blocks grow.Block size 16 was excluded because it was inefficient and often caused out-of-memory errors during training.Source: §5.1 · p. 9
5.2 Hybrid Dense Attention in the First Two Layers
- Hybrid dense attention significantly improves Quest accuracy but yields only marginal benefits for SeerAttention-R.The authors suggest SeerAttention-R's already accurate first-two-layer sparse predictions may limit hybridization gains.Source: §5.2 · p. 10
5.3 Threshold VS Token Budgets
- At a 4e-3 threshold and 4K token budget, threshold selection produces smoother activation patterns and slightly better accuracy in the high-sparsity region than token budgeting.Token budgeting instead produces a strict piecewise-linear activation pattern; the comparison uses Qwen3-4B on AIME24.Source: §5.3 · p. 10
5.4 Impact of Sparse Attention on Generate Length
- Quest and SeerAttention-R with a 2K budget generate much longer reasoning paths than full attention, whose Qwen3-8B AIME24 accuracy is 74.5 with a generated length of 15.1k.The authors warn that extra reasoning steps can undermine efficiency and emphasize accurate sparse selection.Source: §5.4 · pp. 10–11, Table 1 · p. 11
5.5 Training Budget
- Training only AttnGate parameters for 800 steps makes distillation lightweight; training an 8B model requires 12 GPU hours.The experiments used global batch size 16 and 32k-token packed sequences.Source: §5.5 · p. 11, Table 2 · p. 11
6 Limitation and Future Work
The discussion identifies end-to-end integration, adaptive sparsity selection, and unified sparse prefill and decoding as open directions. The current work focuses on sparse-decoding accuracy and kernel-level speedup.
6.1 End-to-end Speedup
- End-to-end system support and optimization remain future work beyond the paper's focus on sparse-decoding accuracy and kernel-level speedup.The authors identify integration with inference frameworks and support for sparse kernels with PagedAttention as requirements for significant end-to-end speedup.Source: §6.1 · p. 11
6.2 Adaptive Sparsity Ratio
- Choosing an optimal sparsity ratio depends on input length, task difficulty, and reasoning length, motivating precise selection and automatic adaptation to task complexity.The discussion presents Top-p-based selection as a promising direction.Source: §6.2 · p. 11
6.3 Unify Sparse Prefill and Decoding
- Sparse prefill and decoding remain separately trained, and jointly supporting both while preserving accuracy and efficiency remains an active challenge.The discussion notes that decoding-oriented methods lack prefill's query-level parallelism and points to multi-token prediction or speculative decoding as possible unifying directions.Source: §6.3 · p. 12
7 Related Works
The related work distinguishes training-free from training-based sparse attention, describes KV-cache selection strategies, and surveys other efficient-attention approaches. These include architectural changes, cross-layer cache sharing, and linear or hybrid attention.
7.1 Training-free vs. Training-based Sparse Attention
- Sparse-attention methods use either training-free fixed or heuristic patterns, or trainable approaches that reduce accuracy degradation.SeerAttention is a post-training approach that learns sparsity patterns without modifying or fine-tuning original model weights.Source: §7.1 · p. 12
7.2 KV Cache Compression: A Sparse Attention Perspective
- KV-cache methods either permanently evict selected tokens, risking accuracy loss, or dynamically select tokens for each attention step without discarding them.Reducing KV-cache size can save I/O bandwidth and GPU memory.Source: §7.2 · p. 12
7.3 Other Efficient Attention Algorithms
- Other efficiency directions modify attention structure, share KV caches across layers, or use linear attention; hybrid linear/full-attention models have shown performance comparable to full-attention models.The cited discussion notes that pure linear-attention models still struggle on some challenging long-context retrieval or reasoning tasks.Source: §7.3 · pp. 12–13
8 Conclusion
SeerAttention-R is a lightweight, flexible sparse-attention framework for long decoding in reasoning models. It preserves near-lossless reasoning accuracy with coarse attention blocks, while its TileLang decoding kernel achieves near-theoretical speedups at high sparsity.
8 Conclusion
- SeerAttention-R accelerates long decoding in reasoning models through a lightweight, flexible sparse-attention framework.It plugs into pretrained models without modifying original parameters and trains only new gating parameters.Source: §8 · p. 13
- SeerAttention-R maintains near-lossless reasoning accuracy in post-training even with coarse-grained attention blocks.Source: §8 · p. 13
- The TileLang sparse-decoding kernel achieves near-theoretical speedup at high sparsity ratios.Source: §8 · p. 13