SageAttention3: 추론을 위한 FP4 주의 메커니즘의 마이크로스케일링 및 8비트 학습 탐구

SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Training

May 16, 2025

저자: Jintao Zhang, Jia Wei, Pengle Zhang, Xiaoming Xu, Haofeng Huang, Haoxu Wang, Kai Jiang, Jun Zhu, Jianfei Chen

cs.AI

초록

주의력의 효율성은 그 이차 시간 복잡도로 인해 중요합니다. 우리는 두 가지 주요 기여를 통해 주의력의 효율성을 향상시켰습니다: 첫째, Blackwell GPU의 새로운 FP4 Tensor Core를 활용하여 주의력 계산을 가속화했습니다. 우리의 구현은 RTX5090에서 1038 TOPS를 달성하며, 이는 RTX5090에서 가장 빠른 FlashAttention 대비 5배의 속도 향상을 보여줍니다. 실험 결과, 우리의 FP4 주의력은 다양한 모델의 추론을 플러그 앤 플레이 방식으로 가속화할 수 있음을 보여줍니다. 둘째, 우리는 저비트 주의력을 훈련 작업에 처음으로 적용했습니다. FlashAttention3 및 SageAttention과 같은 기존의 저비트 주의력 연구는 추론에만 초점을 맞추고 있습니다. 그러나 대규모 모델 훈련의 효율성 또한 중요합니다. 저비트 주의력이 훈련 작업에 효과적으로 적용될 수 있는지 탐구하기 위해, 우리는 정확하고 효율적인 8비트 주의력을 순전파 및 역전파 모두에 대해 설계했습니다. 실험 결과, 8비트 주의력은 미세 조정 작업에서는 손실 없는 성능을 달성하지만, 사전 훈련 작업에서는 더 느린 수렴을 보여줍니다. 코드는 https://github.com/thu-ml/SageAttention에서 제공될 예정입니다.

English

The efficiency of attention is important due to its quadratic time complexity. We enhance the efficiency of attention through two key contributions: First, we leverage the new FP4 Tensor Cores in Blackwell GPUs to accelerate attention computation. Our implementation achieves 1038 TOPS on RTX5090, which is a 5x speedup over the fastest FlashAttention on RTX5090. Experiments show that our FP4 attention can accelerate inference of various models in a plug-and-play way. Second, we pioneer low-bit attention to training tasks. Existing low-bit attention works like FlashAttention3 and SageAttention focus only on inference. However, the efficiency of training large models is also important. To explore whether low-bit attention can be effectively applied to training tasks, we design an accurate and efficient 8-bit attention for both forward and backward propagation. Experiments indicate that 8-bit attention achieves lossless performance in fine-tuning tasks but exhibits slower convergence in pretraining tasks. The code will be available at https://github.com/thu-ml/SageAttention.