SpargeAttn: Präzise sparsame Aufmerksamkeit zur Beschleunigung jeglicher Modellinferenz

SpargeAttn: Accurate Sparse Attention Accelerating Any Model Inference

February 25, 2025

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

cs.AI

Zusammenfassung

Eine effiziente Implementierung von Attention ist für große Modelle aufgrund ihrer quadratischen Zeitkomplexität entscheidend. Glücklicherweise zeigt Attention häufig Sparsity, d.h., viele Werte in der Attention-Map sind nahe null, was das Weglassen entsprechender Berechnungen ermöglicht. Viele Studien haben dieses spärliche Muster genutzt, um Attention zu beschleunigen. Die meisten bestehenden Arbeiten konzentrieren sich jedoch darauf, Attention innerhalb spezifischer Modelle zu optimieren, indem bestimmte spärliche Muster der Attention-Map ausgenutzt werden. Eine universelle spärliche Attention, die sowohl die Beschleunigung als auch die End-to-End-Leistung verschiedener Modelle gewährleistet, bleibt jedoch schwer fassbar. In diesem Artikel schlagen wir SpargeAttn vor, eine universelle spärliche und quantisierte Attention für jedes Modell. Unsere Methode verwendet einen zweistufigen Online-Filter: In der ersten Stufe sagen wir die Attention-Map schnell und genau voraus, wodurch einige Matrixmultiplikationen in der Attention übersprungen werden können. In der zweiten Stufe entwerfen wir einen Online-Softmax-Filter, der keinen zusätzlichen Overhead verursacht und weitere Matrixmultiplikationen überspringt. Experimente zeigen, dass unsere Methode verschiedene Modelle, einschließlich Sprach-, Bild- und Videogenerierung, erheblich beschleunigt, ohne die End-to-End-Metriken zu beeinträchtigen. Die Codes sind unter https://github.com/thu-ml/SpargeAttn verfügbar.

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.