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

초록

큰 모델에서는 어텐션의 이차 시간 복잡도로 인해 효율적인 어텐션 구현이 필수적입니다. 다행히도, 어텐션은 흔히 희소성을 보이는데, 이는 어텐션 맵의 많은 값들이 0에 가까워 해당 계산을 생략할 수 있음을 의미합니다. 많은 연구들이 이러한 희소 패턴을 활용하여 어텐션을 가속화해 왔습니다. 그러나 기존 연구 대부분은 어텐션 맵의 특정 희소 패턴을 활용하여 특정 모델 내에서 어텐션을 최적화하는 데 초점을 맞추고 있습니다. 다양한 모델의 속도 향상과 엔드투엔드 성능을 모두 보장하는 보편적인 희소 어텐션은 아직 달성되지 못했습니다. 본 논문에서는 모든 모델에 적용 가능한 보편적인 희소 및 양자화 어텐션인 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.