SeerAttention-R: Adaptação de Atenção Esparsa para Raciocínio Longo

SeerAttention-R: Sparse Attention Adaptation for Long Reasoning

June 10, 2025

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

Resumo

Apresentamos o SeerAttention-R, uma estrutura de atenção esparsa especialmente projetada para a decodificação longa de modelos de raciocínio. Derivado do SeerAttention, o SeerAttention-R mantém o design de aprender a esparsidade da atenção por meio de um mecanismo de portão auto-distilado, enquanto remove o pooling de consultas para acomodar a decodificação auto-regressiva. Com um mecanismo de portão leve e plugável, o SeerAttention-R é flexível e pode ser facilmente integrado a modelos pré-treinados existentes sem modificar os parâmetros originais. Demonstramos que o SeerAttention-R, treinado com apenas 0,4 bilhão de tokens, mantém uma precisão de raciocínio quase sem perdas com um orçamento de 4K tokens no benchmark AIME sob grandes tamanhos de blocos de atenção esparsa (64/128). Utilizando o TileLang, desenvolvemos um kernel de decodificação esparsa altamente otimizado que alcança acelerações próximas às teóricas de até 9x em relação ao FlashAttention-3 na GPU H100 com 90% de esparsidade. O código está disponível em: 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.