SpargeAttn: Atenção Esparsa Precisas Acelerando a Inferência de Qualquer 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

Resumo

Uma implementação eficiente de atenção é essencial para modelos grandes devido à sua complexidade de tempo quadrática. Felizmente, a atenção comumente exibe esparsidade, ou seja, muitos valores no mapa de atenção estão próximos de zero, permitindo a omissão de cálculos correspondentes. Muitos estudos têm utilizado o padrão esparso para acelerar a atenção. No entanto, a maioria dos trabalhos existentes se concentra em otimizar a atenção dentro de modelos específicos, explorando certos padrões esparsos do mapa de atenção. Uma atenção esparsa universal que garanta tanto a aceleração quanto o desempenho de ponta a ponta de diversos modelos ainda é um desafio. Neste artigo, propomos o SpargeAttn, uma atenção esparsa e quantizada universal para qualquer modelo. Nosso método utiliza um filtro online de dois estágios: no primeiro estágio, prevemos rapidamente e com precisão o mapa de atenção, permitindo a omissão de algumas multiplicações de matrizes na atenção. No segundo estágio, projetamos um filtro online consciente do softmax que não gera sobrecarga adicional e omite ainda mais algumas multiplicações de matrizes. Experimentos mostram que nosso método acelera significativamente diversos modelos, incluindo geração de linguagem, imagem e vídeo, sem sacrificar métricas de ponta a ponta. Os códigos estão disponíveis em 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.