SageAttention2 기술 보고서: 플러그 앤 플레이 추론 가속화를 위한 정확한 4비트 어텐션

SageAttention2 Technical Report: Accurate 4 Bit Attention for Plug-and-play Inference Acceleration

November 17, 2024

저자: Jintao Zhang, Haofeng Huang, Pengle Zhang, Jia Wei, Jun Zhu, Jianfei Chen

cs.AI

초록

선형 레이어에 대한 양자화는 널리 사용되어 왔지만, 주의 집중 과정을 가속화하기 위한 적용은 제한되어 있습니다. SageAttention은 8비트 행렬 곱셈, 16비트 행렬 곱셈과 16비트 누산기, 정확하고 2배 속도 향상 커널을 FlashAttention2와 비교하여 구현하는 정밀한 방법을 사용합니다. 정밀도를 유지하면서 주의 계산의 효율성을 더 향상시키기 위해 SageAttention2를 제안합니다. 이는 상당히 빠른 4비트 행렬 곱셈(Matmul)과 추가 정밀도 향상 기술을 사용합니다. 먼저, 행렬(Q, K)을 워프 수준의 INT4로 양자화하고 행렬(widetilde P, V)을 FP8로 양자화하는 것을 제안합니다. 둘째, Q와 V를 부드럽게 만드는 방법을 제안하여 INT4 QK와 FP8 PV로 주의의 정확도를 향상시킵니다. 셋째, 시간 단계와 레이어를 통해 양자화 정확도를 분석한 후, 다양한 모델에서 최종 메트릭을 보장하기 위해 적응형 양자화 방법을 제안합니다. SageAttention2의 초당 연산 횟수(OPS)는 RTX4090에서 FlashAttention2와 xformers를 각각 약 3배와 5배 초과합니다. 포괄적인 실험을 통해 우리의 접근 방식이 대규모 언어 처리, 이미지 생성 및 비디오 생성을 포함한 다양한 모델에서 미미한 최종 메트릭 손실을 초래한다는 것이 확인되었습니다. 코드는 https://github.com/thu-ml/SageAttention에서 사용할 수 있습니다.

English

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.