SageAttention3: Mikroskalierung der FP4-Attention für Inferenz und eine Untersuchung des 8-Bit-Trainings

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

May 16, 2025

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

cs.AI

Zusammenfassung

Die Effizienz der Aufmerksamkeit ist aufgrund ihrer quadratischen Zeitkomplexität von großer Bedeutung. Wir verbessern die Effizienz der Aufmerksamkeit durch zwei wesentliche Beiträge: Erstens nutzen wir die neuen FP4-Tensor-Cores in Blackwell-GPUs, um die Berechnung der Aufmerksamkeit zu beschleunigen. Unsere Implementierung erreicht 1038 TOPS auf der RTX5090, was einer 5-fachen Beschleunigung gegenüber der schnellsten FlashAttention auf der RTX5090 entspricht. Experimente zeigen, dass unsere FP4-Aufmerksamkeit die Inferenz verschiedener Modelle auf Plug-and-Play-Weise beschleunigen kann. Zweitens führen wir Low-Bit-Aufmerksamkeit in Trainingsaufgaben ein. Bisherige Arbeiten zu Low-Bit-Aufmerksamkeit wie FlashAttention3 und SageAttention konzentrieren sich nur auf die Inferenz. Die Effizienz beim Training großer Modelle ist jedoch ebenfalls wichtig. Um zu untersuchen, ob Low-Bit-Aufmerksamkeit effektiv auf Trainingsaufgaben angewendet werden kann, entwickeln wir eine präzise und effiziente 8-Bit-Aufmerksamkeit für sowohl die Vorwärts- als auch die Rückwärtspropagation. Experimente zeigen, dass 8-Bit-Aufmerksamkeit in Feinabstimmungsaufgaben verlustfreie Leistung erzielt, jedoch in Vorabtrainingsaufgaben eine langsamere Konvergenz aufweist. Der Code wird unter https://github.com/thu-ml/SageAttention verfügbar sein.

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.