FlashAttention-3 — 闪电注意力-3

FlashAttention-3 是一种用于优化 Transformer 神经网络中注意力（attention）机制的算法，旨在最大限度地利用 NVIDIA Hopper (H100) 架构 GPU 的硬件能力。该算法由 Colfax Research、Meta、NVIDIA、佐治亚理工学院（Georgia Tech）、普林斯顿大学和 Together AI 的一组研究人员于 2024 年提出。该研究成果已被 NeurIPS 2024 大会接收，并被评为 spotlight 论文。

FlashAttention-3 是该算法系列的第三次迭代，前两代分别为 FlashAttention (2022) 和 FlashAttention-2 (2023)。其主要目标是在保持计算精度的同时，显著加速大型语言模型 (LLM)的训练和推理。

引言与背景

注意力机制的难题

Transformer 的关键组件是自注意力（self-attention）机制，但其计算复杂度和内存消耗会随着输入序列长度 (n) 的增加呈平方级 (O(n²)) 增长。这构成了一个严重的“瓶颈”，因为现代 GPU 虽然针对快速矩阵乘法进行了优化，但指数函数（如 Softmax 中的函数）的计算速度要慢几个数量级。此外，在朴素实现中，GPU 内存必须存储一个庞大的中间注意力张量，这限制了模型的可扩展性。

FlashAttention 与 FlashAttention-2

为解决这一问题，2022 年提出的 FlashAttention 通过两种技术减少了对慢速全局内存 (HBM) 的访问：

• 分块处理 (tiling): 将计算分解为多个块（tile），这些块在高速的片上内存 (SRAM) 中处理。
• 算子融合 (operator fusion): 所有操作（矩阵乘法、Softmax）都在单个 GPU 内核中执行，无需将中间结果写回全局内存。

这使得内存复杂度从平方级降低到线性级，并将计算速度提升了 2–4 倍。

2023 年，推出了改进版 FlashAttention-2，它优化了计算的并行化。在 NVIDIA Ampere (A100) 架构的 GPU 上，它达到了理论峰值性能的约 70%。然而，在更新的 NVIDIA Hopper (H100) 架构上，其效率显著降低，仅为约 35%。这是因为该算法未能利用 Hopper 的新硬件特性，从而推动了 FlashAttention-3 的诞生。

Hopper GPU (H100) 的新硬件特性

NVIDIA Hopper 架构提供了一系列新功能，FlashAttention-3 利用这些功能来达到最高性能：

• WGMMA (Warpgroup Matrix Multiply-Accumulate): 一种用于张量核心的新指令类型，执行矩阵乘法时的性能几乎是 Ampere 架构的两倍。
• TMA (Tensor Memory Accelerator): 一种硬件模块，可加速全局内存 (HBM) 与共享内存 (shared memory) 之间的数据传输。TMA 自动执行地址计算，从而为计算核心减负。
• FP8 格式: 对 8 位浮点数据格式的硬件支持，其理论性能是 FP16 的两倍，但由于动态范围有限，存在精度损失的风险。

FlashAttention-3 的技术创新

该算法实现了三种专为 Hopper 架构设计的关键优化方法：

1. 异步执行与 Warp 专业化

FlashAttention-3 采用了“warp-specialization”（Warp 专业化）原则，即 GPU 上的不同线程组（warps）专门负责不同的任务：

• 生产者 Warp (Producer warps): 使用 TMA 从全局内存加载数据。
• 消费者 Warp (Consumer warps): 在张量核心上执行矩阵乘法。

得益于 Hopper 的硬件异步性，这些操作在时间上可以重叠。当一组 warp 在执行计算时，另一组 warp 会并行地为下一个块加载数据。这种基于“乒乓调度”（ping-pong scheduling）的流水线（pipeline）方法，能够隐藏慢速操作（如 Softmax）的延迟，并最大限度地利用 GPU 的所有功能模块。

2. 最小化内存操作

该算法保留了前几代版本中的分块 (tiling)思想，但积极利用 TMA 在执行当前计算的同时，异步加载下一个数据块。从慢速 HBM 到快速 SRAM 的数据传输实际上是在主计算的“幕后”完成的，这使得 GPU 因等待数据而空闲的时间更少。

3. 使用低精度 (FP8) 并减少量化误差

切换到 FP8 可将速度提高一倍，但由于量化，可能会导致严重的精度损失。为了解决这个问题，开发人员引入了“incoherent processing”（非相干处理）方法。其核心思想如下：

1. 在计算注意力之前，将特征向量（查询 Q 和键 K）乘以一个随机正交矩阵（例如，阿达玛矩阵）。
2. 这种变换会将绝对值异常大的值（离群值）“涂抹”到所有坐标上，从而使它们的分布更加均匀。
3. 之后再执行到 FP8 的量化，此时量化误差会更小。
4. 由于该变换是正交的，它不会扭曲最终的注意力结果 (QKᵀ)，因为矩阵在相乘过程中其效果被抵消了。

该技术与不进行变换的标准 FP8 应用相比，将 FP8 注意力计算的误差减少了约 2.6 倍。

性能与意义

应用上述技术后，FlashAttention-3 在 H100 GPU 上的性能显著优于前代版本：

• 与 FlashAttention-2 相比，速度提升 1.5–2 倍。
• 高 GPU 利用率: 达到 H100 理论峰值性能的 约 75–85%。
• 吞吐量:
  • 半精度 (FP16/BF16) 下可达 740–840 TFLOPS。
  • 使用 8 位精度 (FP8) 时可达 1.2–1.3 PFLOPS（千万亿次浮点运算）。

FlashAttention-3 的高效率直接影响了 LLM 的开发和应用：

• 缩短训练时间: 将注意力计算速度提升 75–100%，显著减少了模型的训练时间，这些训练原本可能需要数周或数月。
• 扩大上下文窗口: 模型能够高效处理更长的序列（数十万个 token），这对于分析大型文档或代码至关重要。
• 合理利用资源: 允许在更少的 GPU 上实现相同的性能，或在相同硬件上获得更高的速度，从而降低了模型部署的成本。

可用性与集成

作者已在 GitHub 上以开源许可证发布了 FlashAttention-3 的源代码。预计它将被集成到 PyTorch 和 Hugging Face Transformers 等主流深度学习框架中，从而使广大开发者和研究人员都能使用该技术。其前代版本已成为行业事实上的标准，FlashAttention-3 很可能延续这一趋势。

链接

• FlashAttention 官方 GitHub 仓库
• Together AI 发布的 FlashAttention-3 公告博客

参考文献

• Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv:2407.08608.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
• Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.
• Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
• Chen, Y. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
• Liu, Y. et al. (2024). FastAttention: Extending FlashAttention-2 to NPUs and Low-Resource GPUs. OpenReview: 76NYyOrnfk.
• Dege, P. et al. (2025). FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs. arXiv:2506.01969.
• Wang, G. et al. (2024). FlashMask: Efficient and Rich Mask Extension of FlashAttention. arXiv:2410.01359.
• Abbott, V.; Zardini, G. (2025). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. arXiv:2412.03317.