SeerAttention-R: 長文推論のためのスパースアテンション適応

要旨

本論文では、推論モデルの長いデコードに特化したスパースアテンションフレームワークであるSeerAttention-Rを紹介する。SeerAttentionを拡張したSeerAttention-Rは、自己蒸留ゲーティングメカニズムを通じてアテンションのスパース性を学習する設計を保持しつつ、自己回帰デコードに対応するためにクエリプーリングを除去している。軽量なプラグインゲーティングを備えたSeerAttention-Rは柔軟性が高く、既存の事前学習済みモデルに元のパラメータを変更することなく容易に統合できる。0.4Bトークンのみで学習したSeerAttention-Rが、AIMEベンチマークにおいて4Kトークンの予算内で大規模なスパースアテンションブロックサイズ（64/128）のもとでほぼロスレスな推論精度を維持することを実証する。TileLangを使用して、H100 GPU上で90%のスパース性においてFlashAttention-3に対して最大9倍の理論値に近い高速化を達成する高度に最適化されたスパースデコードカーネルを開発した。コードはhttps://github.com/microsoft/SeerAttentionで公開されている。

English

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.

SeerAttention-R: 長文推論のためのスパースアテンション適応

SeerAttention-R: Sparse Attention Adaptation for Long Reasoning

要旨

Support