SageAttention2 Technical Report: Accurate 4 Bit Attention for
  Plug-and-play Inference Acceleration
Authors: name: Jintao Zhang; name: Haofeng Huang; name: Pengle Zhang; name: Jia Wei; name: Jun Zhu; name: Jianfei Chen
Published: 2024-11-17 04:35:49

Abstract: Although quantization for linear layers has been widely used, its application
to accelerate the attention process remains limited. SageAttention utilizes
8-bit matrix multiplication, 16-bit matrix multiplication with 16-bit
accumulator, and precision-enhancing methods, implementing an accurate and 2x
speedup kernel compared to FlashAttention2. To further enhance the efficiency
of attention computation while maintaining precision, we propose
SageAttention2, which utilizes significantly faster 4-bit matrix multiplication
(Matmul) alongside additional precision-enhancing techniques. First, we propose
to quantize matrixes (Q, K) to INT4 in a warp-level granularity and quantize
matrixes (widetilde P, V) to FP8. Second, we propose a method to smooth Q
and V, enhancing the accuracy of attention with INT4 QK and FP8 PV.
Third, we analyze the quantization accuracy across timesteps and layers, then
propose an adaptive quantization method to ensure the end-to-end metrics over
various models. The operations per second (OPS) of SageAttention2 surpass
FlashAttention2 and xformers by about 3x and 5x on RTX4090, respectively.
Comprehensive experiments confirm that our approach incurs negligible
end-to-end metrics loss across diverse models, including those for large
language processing, image generation, and video generation. The codes are
available at https://github.com/thu-ml/SageAttention.