SageAttention3: Atenção FP4 em Microescala para Inferência e uma Exploração de Treinamento em 8 Bits

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

May 16, 2025

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

cs.AI

Resumo

A eficiência do mecanismo de atenção é crucial devido à sua complexidade de tempo quadrática. Melhoramos a eficiência da atenção por meio de duas contribuições principais: Primeiro, aproveitamos os novos Tensor Cores FP4 nas GPUs Blackwell para acelerar o cálculo da atenção. Nossa implementação alcança 1038 TOPS na RTX5090, o que representa uma aceleração de 5x em relação ao FlashAttention mais rápido na RTX5090. Experimentos mostram que nossa atenção FP4 pode acelerar a inferência de vários modelos de forma plug-and-play. Segundo, somos pioneiros na aplicação de atenção de baixa precisão em tarefas de treinamento. Trabalhos existentes com atenção de baixa precisão, como FlashAttention3 e SageAttention, focam apenas na inferência. No entanto, a eficiência no treinamento de modelos grandes também é importante. Para explorar se a atenção de baixa precisão pode ser efetivamente aplicada a tarefas de treinamento, projetamos uma atenção de 8 bits precisa e eficiente tanto para a propagação direta quanto para a retropropagação. Experimentos indicam que a atenção de 8 bits alcança desempenho sem perdas em tarefas de ajuste fino, mas apresenta convergência mais lenta em tarefas de pré-treinamento. O código estará disponível em 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.