Abstract
An efficient attention implementation is essential for large models due toits quadratic time complexity. Fortunately, attention commonly exhibitssparsity, i.e., many values in the attention map are near zero, allowing forthe omission of corresponding computations. Many studies have utilized thesparse pattern to accelerate attention. However, most existing works focus onoptimizing attention within specific models by exploiting certain sparsepatterns of the attention map. A universal sparse attention that guaranteesboth the speedup and end-to-end performance of diverse models remains elusive.In this paper, we propose SpargeAttn, a universal sparse and quantizedattention for any model. Our method uses a two-stage online filter: in thefirst stage, we rapidly and accurately predict the attention map, enabling theskip of some matrix multiplications in attention. In the second stage, wedesign an online softmax-aware filter that incurs no extra overhead and furtherskips some matrix multiplications. Experiments show that our methodsignificantly accelerates diverse models, including language, image, and videogeneration, without sacrificing end-to-end metrics. The codes are available athttps://github.com/thu-ml/SpargeAttn.