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

Аннотация

Эффективность механизма внимания имеет важное значение из-за его квадратичной временной сложности. Мы повышаем эффективность внимания за счет двух ключевых вкладов. Во-первых, мы используем новые тензорные ядра FP4 в GPU Blackwell для ускорения вычислений внимания. Наша реализация достигает 1038 TOPS на RTX5090, что в 5 раз быстрее, чем самая быстрая версия FlashAttention на RTX5090. Эксперименты показывают, что наше внимание FP4 может ускорять вывод различных моделей в режиме plug-and-play. Во-вторых, мы впервые применяем низкобитовое внимание к задачам обучения. Существующие работы по низкобитовому вниманию, такие как 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.