SageAttention3：微缩FP4注意力机制用于推理及8位训练的探索

注意力机制的效率至关重要，因为其时间复杂度为二次方。我们通过两项关键贡献提高了注意力机制的效率：首先，我们利用Blackwell GPU中的新型FP4张量核心加速注意力计算。我们的实现方案在RTX5090上达到了1038 TOPS的性能，比最快的FlashAttention在RTX5090上的速度提升了5倍。实验表明，我们的FP4注意力机制可以以即插即用的方式加速各种模型的推理过程。其次，我们在训练任务中率先应用低比特注意力机制。现有的低比特注意力研究如FlashAttention3和SageAttention主要集中在推理阶段。然而，大型模型训练的效率同样重要。为了探索低比特注意力是否能有效应用于训练任务，我们设计了一种精确且高效的8比特注意力机制，适用于前向传播和反向传播。实验结果表明，8比特注意力在微调任务中实现了无损性能，但在预训练任务中表现出较慢的收敛速度。代码将在https://github.com/thu-ml/SageAttention发布。