SeerAttention-R : Adaptation de l'attention parcimonieuse pour le raisonnement long

SeerAttention-R: Sparse Attention Adaptation for Long Reasoning

June 10, 2025

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

Résumé

Nous présentons SeerAttention-R, un cadre d'attention parcimonieuse spécialement conçu pour le décodage long des modèles de raisonnement. Développé à partir de SeerAttention, SeerAttention-R conserve la conception d'apprentissage de la parcimonie de l'attention via un mécanisme de gating auto-distillé, tout en supprimant le pooling des requêtes pour s'adapter au décodage auto-régressif. Avec un mécanisme de gating léger et modulaire, SeerAttention-R est flexible et peut être facilement intégré à des modèles pré-entraînés existants sans modifier les paramètres originaux. Nous démontrons que SeerAttention-R, entraîné sur seulement 0,4 milliard de tokens, maintient une précision de raisonnement quasi sans perte avec un budget de 4K tokens dans le benchmark AIME, même avec de grandes tailles de blocs d'attention parcimonieuse (64/128). En utilisant TileLang, nous avons développé un noyau de décodage parcimonieux hautement optimisé qui atteint des accélérations proches des limites théoriques, jusqu'à 9x par rapport à FlashAttention-3 sur un GPU H100 avec une parcimonie de 90%. Le code est disponible à l'adresse suivante : 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.