FlashAttention — 闪电注意力

FlashAttention 是一种革命性的注意力（attention）机制计算算法，旨在显著加速大型语言模型 (LLM)的训练和推理，同时保持完整的计算精度。该算法于2022年由斯坦福大学以 Tri Dao 为首的研究团队首次提出。

FlashAttention 的核心思想是重组计算过程，充分考虑 GPU 的内存层次结构，从而最大限度地减少对慢速内存的访问次数，消除标准注意力机制的主要瓶颈。

标准注意力机制的问题

Transformer 中的标准自注意力机制通过以下公式计算： $\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$ 其中 Q、K、V 分别是查询（Query）、键（Key）和值（Value）矩阵。

该方法的主要问题是相对于序列长度 N 的二次方时间和内存复杂度 (O(N²))。在朴素实现中，需要计算并存储大小为 N×N 的完整注意力矩阵 S，这会导致两个关键问题：

1. 内存消耗大：在处理长上下文时，存储 N×N 的矩阵变得不可行。
2. 输入/输出（IO）操作：主要瓶颈并非浮点运算的数量，而是对 GPU 慢速内存的持续访问。

GPU 内存层次结构

为了理解这一问题，区分 GPU 中的两种内存类型至关重要（以 NVIDIA A100 为例）：

• SRAM (静态随机存取存储器)：高速片上内存，容量小（约 20 MB），但带宽极高（高达 19 TB/s）。
• HBM (高带宽内存)：慢速大容量内存（40–80 GB），带宽远低于 SRAM（约 1.5 TB/s）。

这种不对称性使得标准注意力算法受内存带宽限制 (memory-bound)，因为它需要不断地从慢速 HBM 中读取和写入大型矩阵，这正是延迟的主要来源。

FlashAttention 的关键创新

FlashAttention 是一种IO 感知 (IO-aware) 算法，通过最大限度地减少对 HBM 的访问来解决此问题。这主要通过三种核心技术实现。

分块（Tiling）与块处理

FlashAttention 并不一次性处理整个矩阵，而是将输入的 Q、K、V 矩阵分割成能装入高速 SRAM 的小块（tiles）。算法会依次加载这些块，在块上执行所有注意力计算并更新最终结果，而无需在慢速 HBM 中存储完整的注意力矩阵。

在线 Softmax 计算

关键的技术突破是“在线”Softmax 计算。标准 Softmax 需要输入向量的所有元素才能进行归一化。FlashAttention 采用了一种改进算法，可以分块计算 Softmax。它维护两个中间值（当前最大值和指数和），这些值会随着新块的处理而更新，从而在不访问整个矩阵的情况下获得精确结果。

将操作融合到单个 CUDA 核心中

所有注意力操作（矩阵乘法 QKᵀ、掩码、Softmax、与 V 相乘）都被合并到一个单一融合的 CUDA 核心 (fused kernel) 中。这极大地减少了对 HBM 的读/写操作次数：算法不再需要多次遍历整个矩阵，而是一次性将一个块加载到 SRAM 中，执行所有计算，然后只写回最终结果。

理论与实践效率

复杂度与最优性

FlashAttention 将内存消耗从 O(N²) 降低到 O(N)，实现了线性扩展。事实证明，该算法的 IO 复杂度对于在两级内存层次结构中计算注意力是理论上最优的，也就是说，在不改变硬件的情况下，无法更快地执行精确注意力计算。

实证结果

第一版 FlashAttention 取得了显著的改进：

• 加速效果：
  • BERT-large (序列长度 512)：训练速度提升 15%。
  • GPT-2 (序列长度 1K)：速度提升 3倍。
  • Long-Range Arena (1K-4K) 任务：速度提升 2.4倍。
• 内存节省：与精确的基线实现相比，内存节省高达 20倍。
• 模型质量提升：由于能够处理更长的上下文，FlashAttention 不仅没有损失模型质量，反而有所提升。例如，GPT-2 的困惑度改善了 0.7，长文档分类任务的准确率提高了 6.4 个百分点。

演进与未来发展

FlashAttention 的成功催生了一系列面向硬件的算法。

FlashAttention-2 (2023) - 第二代 FlashAttention

第二版旨在更充分地利用 GPU 资源。在最初的 FlashAttention 中，NVIDIA A100 上的效率仅为最大理论值的 25–40%。FlashAttention-2 改进了计算并行化，从而实现了：

• 与第一版相比，速度提升了两倍。
• 将 GPU 利用率提高到理论最大值的 50–73%。
• 扩展了对大小为 256 的注意力头以及多查询注意力 (Multi-Query Attention, MQA) 架构的支持。

FlashAttention-3 (2024) - 第三代 FlashAttention

第三版专门针对 NVIDIA Hopper (H100) GPU 架构进行了优化。它利用了新的硬件功能，如张量核心 (Tensor Cores) 的异步性和对FP8的支持，从而实现了：

• 与 FlashAttention-2 相比，速度进一步提升了 1.5–2倍。
• 在 FP16 上实现了高达 740 TFLOPS 的性能，在 FP8 上接近 1.2 PFLOPS。

专用解决方案

FlashAttention 的思想在其他项目中得到了进一步发展：

• FlashInfer (2025)：一个可定制的注意力引擎，专为 LLM 推理任务优化。它专注于在流式生成模式下高效处理 KV 缓存。
• FlashMLA (2024)：一种带有上下文缓存压缩（latent attention）的注意力实现，能够在极长序列上节省内存，且信息损失极小。

对行业和生态系统的影响

FlashAttention 成为一项基础性突破，并迅速发展为高效训练和推理 LLM 的行业标准。它已被集成到 PyTorch 和 Hugging Face 等核心库中，并被大多数大型语言模型（如 LLaMA、MPT、Falcon、Claude 等）所采用。

正是 FlashAttention 及其后续版本，在扩大语言模型的上下文窗口方面发挥了决定性作用：从 2-4 千个 token (GPT-3) 增加到 12.8 万个 token (GPT-4)，甚至在实验性模型中达到数百万个 token。该算法消除了扩展 Transformer 的主要障碍之一，为从长文档分析到多模态理解等 AI 应用开辟了新的可能性。

链接

• FlashAttention 官方 GitHub 仓库

参考文献

• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision. arXiv:2407.08608.
• Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.
• Hong, K. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
• Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
• 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. (2025). FlashMask: Efficient and Rich Mask Extension of FlashAttention. OpenReview wUtXB43Chi.
• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (OpenReview version). OpenReview H4DqfPSibmx.
• Gholami, A. et al. (2024). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. OpenReview pF2ukh7HxA.