SpargeAttn: Nauwkeurige Sparse Attention voor Versnelde Inferentie van Elk Model

SpargeAttn: Accurate Sparse Attention Accelerating Any Model Inference

February 25, 2025

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

cs.AI

Samenvatting

Een efficiënte implementatie van aandacht (attention) is essentieel voor grote modellen vanwege de kwadratische tijdcomplexiteit. Gelukkig vertoont aandacht vaak sparsity, wat betekent dat veel waarden in de aandachtmatrix (attention map) bijna nul zijn, waardoor de bijbehorende berekeningen kunnen worden overgeslagen. Veel studies hebben dit sparse patroon benut om aandacht te versnellen. De meeste bestaande werken richten zich echter op het optimaliseren van aandacht binnen specifieke modellen door bepaalde sparse patronen van de aandachtmatrix te benutten. Een universele sparse aandacht die zowel de versnelling als de end-to-end prestaties van diverse modellen garandeert, blijft nog ongrijpbaar. In dit artikel stellen we SpargeAttn voor, een universele sparse en gekwantiseerde aandacht voor elk model. Onze methode gebruikt een tweestaps online filter: in de eerste stap voorspellen we snel en nauwkeurig de aandachtmatrix, waardoor sommige matrixvermenigvuldigingen in de aandacht kunnen worden overgeslagen. In de tweede stap ontwerpen we een online softmax-aware filter dat geen extra overhead veroorzaakt en nog meer matrixvermenigvuldigingen overslaat. Experimenten tonen aan dat onze methode diverse modellen aanzienlijk versnelt, waaronder taal-, beeld- en videogeneratie, zonder in te leveren op end-to-end metrieken. De code is beschikbaar op 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.