SpargeAttn: Atención Dispersa Precisa que Acelera la Inferencia de Cualquier Modelo

SpargeAttn: Accurate Sparse Attention Accelerating Any Model Inference

February 25, 2025

Autores: Jintao Zhang, Chendong Xiang, Haofeng Huang, Jia Wei, Haocheng Xi, Jun Zhu, Jianfei Chen

cs.AI

Resumen

Una implementación eficiente de la atención es esencial para modelos grandes debido a su complejidad temporal cuadrática. Afortunadamente, la atención comúnmente exhibe dispersión, es decir, muchos valores en el mapa de atención son cercanos a cero, lo que permite omitir los cálculos correspondientes. Muchos estudios han aprovechado este patrón disperso para acelerar la atención. Sin embargo, la mayoría de los trabajos existentes se centran en optimizar la atención dentro de modelos específicos explotando ciertos patrones dispersos del mapa de atención. Una atención dispersa universal que garantice tanto la aceleración como el rendimiento de extremo a extremo en diversos modelos sigue siendo esquiva. En este artículo, proponemos SpargeAttn, una atención dispersa y cuantizada universal para cualquier modelo. Nuestro método utiliza un filtro en línea de dos etapas: en la primera etapa, predecimos rápida y precisamente el mapa de atención, permitiendo omitir algunas multiplicaciones de matrices en la atención. En la segunda etapa, diseñamos un filtro en línea consciente del softmax que no incurre en sobrecarga adicional y omite más multiplicaciones de matrices. Los experimentos muestran que nuestro método acelera significativamente diversos modelos, incluyendo generación de lenguaje, imágenes y videos, sin sacrificar las métricas de extremo a extremo. Los códigos están disponibles en https://github.com/thu-ml/SpargeAttn.

English

An efficient attention implementation is essential for large models due to its quadratic time complexity. Fortunately, attention commonly exhibits sparsity, i.e., many values in the attention map are near zero, allowing for the omission of corresponding computations. Many studies have utilized the sparse pattern to accelerate attention. However, most existing works focus on optimizing attention within specific models by exploiting certain sparse patterns of the attention map. A universal sparse attention that guarantees both the speedup and end-to-end performance of diverse models remains elusive. In this paper, we propose SpargeAttn, a universal sparse and quantized attention for any model. Our method uses a two-stage online filter: in the first stage, we rapidly and accurately predict the attention map, enabling the skip of some matrix multiplications in attention. In the second stage, we design an online softmax-aware filter that incurs no extra overhead and further skips some matrix multiplications. Experiments show that our method significantly accelerates diverse models, including language, image, and video generation, without sacrificing end-to-end metrics. The codes are available at https://github.com/thu-ml/SpargeAttn.