大型模型在处理序列长度不断延长的过程中，注意力运算（Attention）所需的时间成本逐渐上升，并演变为主要的消耗。此前，由清华大学陈键飞团队研发的即插即用的SageAttention及其升级版SageAttention2，已广泛应用于众多开源和商业的大规模模型之中。这些模型包括但不限于Vidu、CogvideoX、Mochi、Wan、HunyuanVideo、Flux、Llama3和Qwen等。

近期，清华大学陈键飞的研究团队推出了首个专为 BlackWell 架构设计的全 FP4 量化即时安装式注意力算子，命名为 SageAttention3。在推理加速方面，相较于FlashAttention实现了五倍的即时部署效果（先前的SageAttention V1/V2/V2++分别实现了2.1倍、3倍和3.9倍的加速），以RTX 5090为例，SageAttention3的速度可达1040 TOPS，即便是在配备价值远超RTX 5090的H100芯片并使用Hopper特有的FlashAttention3时，其速度也快出1.65倍！SageAttention3 在众多视频及图像生成的大型模型中，如 HunyuanVideo、CogVideoX、Mochi 以及众多图像生成模型，均展现了卓越的端到端精度。此外，还首次提出了可训练的8比特注意力机制（SageBwd），该机制旨在加速大模型的训练过程（注：即便是FlashAttention3的FP8版本，也仅支持前向传播）。在各项微调任务中，这一机制均能保持与全精度注意力机制相当的效果。

效果预览

SageAttention3 引入了高效的注意力机制算法，其推理速度比 FlashAttention 快了五倍，达到了即插即用的效果。无论是输入何种 Q、K、V 矩阵，都能迅速生成 Attention Output (O)，实现了仅需两行代码即可加速任何模型的推理过程。根据官方仓库所公布的开放源代码计划，SageAttention2++的源代码预计将在6月20日前后对外公布，而SageAttention3的源代码则预计会在7月15日前后对外公开。

以 HunyuanVideo 为例，在搭载 RTX5090 显卡的设备上，SageAttention3 技术能够实现端到端生成过程的3倍加速，同时视频的画质损失极小。

该链接指向的文章中，明确指出了对某些内容的严格限制，不得擅自修改或变更，以确保信息的准确性和一致性。

视频 1（使用 FlashAttention2，490s）

https://mp.weixin.qq.com/s/aVlYM_cMrpTKeH3ao-UJuA

视频 2（使用 SageAttention3,164s）

FlashAttention2 已成为 RTX5090 平台上 FlashAttention 方法的最佳实现。

接下来，我们将从以下几个方面对SageAttention3进行详细介绍：首先是前言部分，随后是所面临的挑战，接着是采用的方法，最后是实验所取得的效果。

SageAttention3 总体流程图

前言

随着处理序列长度的不断增加，对Attention机制的加速优化显得尤为关键。图中呈现了在标准Transformer模型中，不同运算步骤随序列长度增长而变化的时间比例分布。

为了便于在注意力运算中引用矩阵，我们首先对注意力的计算公式进行简要回顾：

研究目的包括：（1）Blackwell 架构搭载了极为高效的 FP4 Tensor Core，以 RTX5090 为例，其运行速度达到了 FP16 Tensor Core 的8倍之快。在训练过程中，注意力的计算成本同样至关重要。此前，鲜有研究尝试对低比特注意力加速模型进行训练，即便是FlashAttention3的FP8版本，也仅限于正向过程。此外，我们亦期望对注意力的正向和反向过程进行量化，以实现训练速度的提升。

FP4 注意力量化有什么问题？

FP4的数值类型只包含15个有效的数值，这一限制导致在以张量（每张量）或标记（每标记）为粒度进行量化时，难以确保量化的精确性。

P矩阵的数值跨度较大，若直接进行FP4量化，则会导致量化缩放因子被限定在一个很小的区间内。但是，硬件设备规定这些量化因子必须使用FP8数据类型来表示。在这种情况下，将缩放因子转换为FP8格式时，将会引起明显的精度下降。

8-Bit 注意力用于训练有什么问题？

P 矩阵的梯度对于量化误差表现出极高的敏感性，同时，在反向传播的过程中，它还会导致 Q 和 K 的梯度沿着序列长度逐渐累积误差。

技术方案

为了解决上述的挑战，研究团队提出了对应的解决办法。

为了增强FP4的量化准确性，研究团队引入了Microscaling FP4量化技术，该技术得到了BlackWell硬件层面的支持。这种量化方式允许使用不同的量化粒度进行矩阵量化，同时NVIDIA的硬件自动实现了反量化功能。Microscaling FP4 拥有两种数据表示方式，其一为MXFP4，其二为NVFP4。这两种格式均采纳了E2M1的FP4数据类型。然而，NVFP4在量化块大小上有所不同，其缩放因子的数据类型为E4M3。相较之下，MXFP4的量化块大小各异，其缩放因子的数据格式为E8M0。研究团队选择了NVFP4数据格式，该格式在量化准确性方面显著优于MXFP4。

针对P的缩放因子区间较窄的问题，研究团队提出了两阶段量化（Two-level Quantization）的解决方案。在FlashAttention中，P矩阵的数值被限定在一个特定的区间内，这进而使得P的缩放因子的变动范围也局限于0至0.167之间。若将缩放因子直接转换成FP8格式，将会造成显著的精度下降。

研究团队决定首先将P值通过Per-token的方式进行量化，限定在特定范围内，随后再对FP4进行量化处理。

本表揭示了 Two-Level Scaling 在提高精度方面的效果：

下图展示了 SageAttention3 的算法流程：

在8-Bit训练Attention机制中，研究团队对查询（Q）、键（K）和值（V）进行了块级INT8量化处理，而对P则巧妙地实施了无额外开销的按token量化。其前向算法的具体步骤为：

在反向传播的过程中总共涉及到 5 个矩阵乘法：

研究团队发现是否量化 dOVT 对精度有着较大的影响：

研究团队决定保留 dOVT 的 FP16 精度，同时将剩余的四个矩阵乘法过程进行了量化处理。接下来，我将介绍反向传播算法的具体内容：，

实验效果

SageAttention3 引入了 GPU 的 CUDA Kernel 技术，这一举措在算子运行速度和各模型从输入到输出的整体准确度方面均展现出优异的性能。

具体来看，算子运行速度相较于FlashAttention2（在5090平台上速度最快的FlashAttention版本）提升了约5倍，与xformers相比，则实现了10倍的加速效果。

各个模型在具体应用场景中的整体端到端准确性方面均有出色表现，无论是视频制作还是图像生成这类大型模型，都维持了高水平的端到端精度。

下图是在 HunyuanVideo 当中的可视化实例：

下图是在 Flux 上的可视化实例：

下图是在 Cogvideo 中的可视化实例：

该表格详细呈现了SageAttention3在多种视频及图像生成模型中的端到端精度数据。

在端到端速度表现方面，SageAttention3的实施方案能够高效地提升长序列模型的运行速度，例如，它能够将HunyuanVideo的端到端速度提升至原来的三倍。

在Base Model进行微调以适应Instruct Model的任务时，8-Bit训练的Attention模块表现出了与BF16注意力机制完全相同的精确度，具体的多任务与模型微调效果详见表格。

并且在训练速度上也能起到较好的加速效果：

研究团队进一步揭示，尽管当前用于训练的8比特Attention在微调环节表现完美无瑕，但在预训练阶段，其与全精度Attention在损失值上仍存在明显差距，这预示着未来研究需要继续深入。