SageAttention3：用于推理的微缩放 FP4 注意力机制及对 8 位训练的探索
机器学习
2026-01-16 v3 人工智能
硬件体系结构
计算机视觉与模式识别
性能
摘要
由于注意力机制的二次时间复杂度，其效率至关重要。我们通过两个关键贡献来提升注意力机制的效率：首先，我们利用 Blackwell GPU 中新的 FP4 Tensor Core 来加速注意力计算。我们的实现在 RTX5090 上达到了 1038 TOPS，比 RTX5090 上最快的 FlashAttention 快了 5 倍。实验表明，我们的 FP4 注意力机制可以以即插即用的方式加速各种模型的推理。其次，我们率先将低位注意力机制应用于训练任务。现有的低位注意力机制工作，如 FlashAttention3 和 SageAttention，仅关注推理。然而，训练大模型的效率同样重要。为了探索低位注意力机制是否能有效应用于训练任务，我们设计了一种用于前向和后向传播的精确且高效的 8 位注意力机制。实验表明，8 位注意力机制在微调任务中实现了无损性能，但在预训练任务中表现出较慢的收敛速度。代码可在 https://github.com/thu-ml/SageAttention 获取。
引用
@article{arxiv.2505.11594,
title = {SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Training},
author = {Jintao Zhang and Jia Wei and Pengle Zhang and Xiaoming Xu and Haofeng Huang and Haoxu Wang and Kai Jiang and Jianfei Chen and Jun Zhu},
journal= {arXiv preprint arXiv:2505.11594},
year = {2026}
}