SpargeAttn: Attenzione Sparsa di Precisione per Accelerare l'Inferenza di Qualsiasi Modello

SpargeAttn: Accurate Sparse Attention Accelerating Any Model Inference

February 25, 2025

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

cs.AI

Abstract

Un'implementazione efficiente dell'attenzione è essenziale per i modelli di grandi dimensioni a causa della sua complessità temporale quadratica. Fortunatamente, l'attenzione presenta comunemente una sparsità, ovvero molti valori nella mappa di attenzione sono prossimi allo zero, consentendo l'omissione dei calcoli corrispondenti. Molti studi hanno sfruttato il modello sparso per accelerare l'attenzione. Tuttavia, la maggior parte dei lavori esistenti si concentra sull'ottimizzazione dell'attenzione all'interno di modelli specifici sfruttando determinati schemi sparsi della mappa di attenzione. Un'attenzione sparsa universale che garantisca sia l'accelerazione che le prestazioni end-to-end di modelli diversi rimane elusiva. In questo articolo, proponiamo SpargeAttn, un'attenzione sparsa e quantizzata universale per qualsiasi modello. Il nostro metodo utilizza un filtro online a due stadi: nella prima fase, prevediamo rapidamente e con precisione la mappa di attenzione, consentendo di saltare alcune moltiplicazioni di matrici nell'attenzione. Nella seconda fase, progettiamo un filtro online softmax-aware che non comporta costi aggiuntivi e salta ulteriori moltiplicazioni di matrici. Gli esperimenti dimostrano che il nostro metodo accelera significativamente modelli diversi, inclusi quelli per la generazione di linguaggio, immagini e video, senza sacrificare le metriche end-to-end. I codici sono disponibili all'indirizzo 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.