用FP4加速注意力计算，实现推理5倍提速，并探索8位训练可行性。
SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Training
- 利用黑曜石显卡的FP4张量核心，实现高效注意力计算。
- FP4注意力在RTX5090上达1038 TOPS，比FlashAttention快5倍。
- 首次实现8位训练无损微调，适合追求推理效率的开发者。
由于注意力机制具有二次时间复杂度，其计算效率至关重要。本文通过两大贡献提升注意力效率：首先，利用黑曜石架构GPU中的新型FP4张量核心加速注意力计算，实现1038 TOPS性能，较RTX5090上的最快FlashAttention提升5倍；实验表明，该FP4注意力可即插即用地加速多种模型的推理。其次，开创性地将低比特注意力扩展至训练任务。现有工作如FlashAttention3和SageAttention仅关注推理。为探索低比特注意力在训练中的有效性，我们设计了一种前向与反向传播均高效的8位注意力。实验显示，8位注意力在微调任务中实现无损性能，但在预训练任务中收敛速度较慢。代码已开源：https://github.com/thu-ml/SageAttention。
原文摘要 · Abstract (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 is available at https://github.com/thu-ml/SageAttention.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。