SpargeAttn : Une Attention Sparse Précise Accélérant l'Inférence de Tout Modèle

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

Résumé

Une implémentation efficace de l'attention est essentielle pour les grands modèles en raison de sa complexité temporelle quadratique. Heureusement, l'attention présente souvent de la sparsité, c'est-à-dire que de nombreuses valeurs dans la carte d'attention sont proches de zéro, permettant ainsi d'omettre les calculs correspondants. De nombreuses études ont exploité ce motif sparse pour accélérer l'attention. Cependant, la plupart des travaux existants se concentrent sur l'optimisation de l'attention au sein de modèles spécifiques en exploitant certains motifs sparses de la carte d'attention. Une attention sparse universelle garantissant à la fois l'accélération et les performances end-to-end de divers modèles reste insaisissable. Dans cet article, nous proposons SpargeAttn, une attention sparse et quantifiée universelle pour tout modèle. Notre méthode utilise un filtre en ligne en deux étapes : dans la première étape, nous prédisons rapidement et précisément la carte d'attention, permettant de sauter certaines multiplications matricielles dans l'attention. Dans la deuxième étape, nous concevons un filtre en ligne sensible au softmax qui n'entraîne aucun surcoût supplémentaire et permet de sauter davantage de multiplications matricielles. Les expériences montrent que notre méthode accélère considérablement divers modèles, y compris la génération de langage, d'images et de vidéos, sans sacrifier les métriques end-to-end. Les codes sont disponibles à l'adresse 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.