SeerAttention-R: Adaptación de Atención Dispersa para Razonamiento Largo

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

Resumen

Presentamos SeerAttention-R, un marco de atención dispersa específicamente diseñado para la decodificación extensa de modelos de razonamiento. Extendido a partir de SeerAttention, SeerAttention-R conserva el diseño de aprendizaje de la dispersión de atención mediante un mecanismo de compuerta auto-distilado, mientras elimina la agrupación de consultas para adaptarse a la decodificación auto-regresiva. Con una compuerta ligera de complemento, SeerAttention-R es flexible y puede integrarse fácilmente en modelos preentrenados existentes sin modificar los parámetros originales. Demostramos que SeerAttention-R, entrenado con solo 0.4 mil millones de tokens, mantiene una precisión de razonamiento casi sin pérdidas con un presupuesto de 4K tokens en el benchmark AIME bajo tamaños de bloques de atención dispersa grandes (64/128). Utilizando TileLang, desarrollamos un núcleo de decodificación dispersa altamente optimizado que logra aceleraciones cercanas a las teóricas de hasta 9x sobre FlashAttention-3 en la GPU H100 con un 90% de dispersión. El código está disponible en: 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.