SeerAttention-R: Adattamento dell'attenzione sparsa per il ragionamento a lungo termine

SeerAttention-R: Sparse Attention Adaptation for Long Reasoning

June 10, 2025

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

Abstract

Introduciamo SeerAttention-R, un framework di attenzione sparsa specificamente progettato per il lungo decoding di modelli di ragionamento. Estendendo SeerAttention, SeerAttention-R mantiene il design di apprendimento della sparsità dell'attenzione attraverso un meccanismo di gating auto-distillato, rimuovendo però il pooling delle query per adattarsi al decoding auto-regressivo. Con un gating leggero e plug-in, SeerAttention-R è flessibile e può essere facilmente integrato in modelli pre-addestrati esistenti senza modificare i parametri originali. Dimostriamo che SeerAttention-R, addestrato su soli 0,4 miliardi di token, mantiene un'accuratezza di ragionamento quasi senza perdite con un budget di 4K token nel benchmark AIME, utilizzando grandi dimensioni di blocco di attenzione sparsa (64/128). Utilizzando TileLang, abbiamo sviluppato un kernel di decoding sparsa altamente ottimizzato che raggiunge accelerazioni quasi teoriche fino a 9x rispetto a FlashAttention-3 su GPU H100 con una sparsità del 90%. Il codice è disponibile all'indirizzo: 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.