SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference
  Acceleration
Authors: name: Jintao Zhang; name: Jia wei; name: Pengle Zhang; name: Jun Zhu; name: Jianfei Chen
Published: 2024-10-03 10:25:23

Abstract: The transformer architecture predominates across various models. As the heart
of the transformer, attention has a computational complexity of O(N^2),
compared to O(N) for linear transformations. When handling large sequence
lengths, attention becomes the primary time-consuming component. Although
quantization has proven to be an effective method for accelerating model
inference, existing quantization methods primarily focus on optimizing the
linear layer. In response, we first analyze the feasibility of quantization in
attention detailedly. Following that, we propose SageAttention, a highly
efficient and accurate quantization method for attention. The OPS (operations
per second) of our approach outperforms FlashAttention2 and xformers by about
2.1 times and 2.7 times, respectively. SageAttention also achieves superior
accuracy performance over FlashAttention3. Comprehensive experiments confirm
that our approach incurs almost no end-to-end metrics loss across diverse
models, including those for large language processing, image generation, and
video generation.