SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Training
Created by
Haebom
저자
Jintao Zhang, Jia Wei, Pengle Zhang, Xiaoming Xu, Haofeng Huang, Haoxu Wang, Kai Jiang, Jun Zhu, Jianfei Chen
개요
본 논문은 어텐션 메커니즘의 계산 복잡도를 해결하기 위해 두 가지 주요 기여를 제시합니다. 첫째, Blackwell GPU의 FP4 Tensor Cores를 활용하여 어텐션 연산을 가속화하는 방법을 제안하며, RTX 5090에서 기존 최고 성능 대비 5배 빠른 1038 TOPS의 성능을 달성했습니다. 다양한 모델의 추론 속도를 향상시키는 플러그 앤 플레이 방식을 제시합니다. 둘째, 저비트 어텐션을 훈련 과정에 적용하는 선구적인 연구를 진행했습니다. 기존의 FlashAttention3이나 SageAttention과 달리 추론뿐 아니라 훈련 효율 향상에도 초점을 맞춰, 정확하고 효율적인 8-bit 어텐션을 전방 및 역방향 전파에 적용했습니다. 실험 결과, 8-bit 어텐션은 미세 조정 작업에서는 손실 없는 성능을 달성하지만, 사전 훈련 작업에서는 수렴 속도가 느린 것으로 나타났습니다. 코드는 https://github.com/thu-ml/SageAttention 에서 공개될 예정입니다.