SeerAttention-R: Sparse Attention-Adaption für langes logisches Schließen

SeerAttention-R: Sparse Attention Adaptation for Long Reasoning

June 10, 2025

Autoren: 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

cs.AI

Zusammenfassung

Wir stellen SeerAttention-R vor, ein Framework für spärliche Aufmerksamkeit, das speziell für das lange Decodieren von Reasoning-Modellen entwickelt wurde. Als Erweiterung von SeerAttention behält SeerAttention-R das Design bei, die Aufmerksamkeitssparsität durch einen selbst-distillierten Gating-Mechanismus zu erlernen, während das Query-Pooling entfernt wird, um das autoregressive Decodieren zu ermöglichen. Mit einem leichtgewichtigen Plug-in-Gating ist SeerAttention-R flexibel und kann einfach in bestehende vortrainierte Modelle integriert werden, ohne die ursprünglichen Parameter zu verändern. Wir zeigen, dass SeerAttention-R, das mit nur 0,4B Tokens trainiert wurde, eine nahezu verlustfreie Reasoning-Genauigkeit mit einem 4K-Token-Budget im AIME-Benchmark bei großen spärlichen Aufmerksamkeitsblockgrößen (64/128) beibehält. Mit TileLang entwickeln wir einen hochoptimierten spärlichen Decodier-Kernel, der nahezu theoretische Beschleunigungen von bis zu 9x gegenüber FlashAttention-3 auf der H100-GPU bei 90% Sparsität erreicht. Der Code ist verfügbar unter: 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.