SpargeAttn: Точное разреженное внимание, ускоряющее вывод любой модели

SpargeAttn: Accurate Sparse Attention Accelerating Any Model Inference

February 25, 2025

Авторы: Jintao Zhang, Chendong Xiang, Haofeng Huang, Jia Wei, Haocheng Xi, Jun Zhu, Jianfei Chen

cs.AI

Аннотация

Эффективная реализация внимания является важной для крупных моделей из-за квадратичной сложности по времени. К счастью, внимание обычно проявляет разреженность, т.е. многие значения в карте внимания близки к нулю, что позволяет опустить соответствующие вычисления. Многие исследования использовали разреженный шаблон для ускорения внимания. Однако большинство существующих работ сосредотачиваются на оптимизации внимания в пределах конкретных моделей, эксплуатируя определенные разреженные шаблоны карты внимания. Универсальное разреженное внимание, которое гарантирует как ускорение, так и конечную производительность различных моделей, остается недостижимым. В данной статье мы предлагаем SpargeAttn, универсальное разреженное и квантованное внимание для любой модели. Наш метод использует двухэтапный онлайн-фильтр: на первом этапе мы быстро и точно предсказываем карту внимания, что позволяет пропустить некоторые умножения матриц во время внимания. На втором этапе мы разрабатываем онлайн-фильтр, осведомленный о софтмаксе, который не влечет за собой дополнительных накладных расходов и дополнительно пропускает некоторые умножения матриц. Эксперименты показывают, что наш метод значительно ускоряет различные модели, включая генерацию языка, изображений и видео, не жертвуя конечными метриками. Коды доступны по ссылке 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.