--- title: "FlashAttention — 闪电注意力" source: "https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B" wiki: "systems-analysis.info/int" article: "FlashAttention_—_闪电注意力" language: "zh" categories: - "Category:Chinese" - "Category:Large language models" - "Category:Machine learning" revision_id: 2272 wiki_created_at: 2026-09-06T23:01:19Z wiki_modified_at: 2026-09-06T23:01:19Z downloaded_at: 2026-09-07T22:50:16Z --- # FlashAttention — 闪电注意力 **FlashAttention** 是一种革命性的注意力（attention）机制计算算法，旨在显著加速大型语言模型 (LLM)的训练和推理，同时保持完整的计算精度。该算法于2022年由斯坦福大学以 **Tri Dao** 为首的研究团队首次提出^{[\[1\]](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_note-arxiv_main-1)}。 FlashAttention 的核心思想是重组计算过程，充分考虑 GPU 的内存层次结构，从而最大限度地减少对慢速内存的访问次数，消除标准注意力机制的主要瓶颈。 ## 标准注意力机制的问题 Transformer 中的标准自注意力机制通过以下公式计算： $\text{Attention}(Q,K,V) = \text{softmax}\left( \frac{QK^{T}}{\sqrt{d_{k}}} \right)V$ 其中 Q、K、V 分别是查询（Query）、键（Key）和值（Value）矩阵。该方法的主要问题是相对于序列长度 N 的**二次方时间和内存复杂度** (O(N²))^{[\[1\]](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_note-arxiv_main-1)}。在朴素实现中，需要计算并存储大小为 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**）^{[\[2\]](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_note-openreview_fa1-2)}。这种不对称性使得标准注意力算法**受内存带宽限制** (memory-bound)，因为它需要不断地从慢速 HBM 中读取和写入大型矩阵，这正是延迟的主要来源。 ## FlashAttention 的关键创新 FlashAttention 是一种**IO 感知** (IO-aware) 算法，通过最大限度地减少对 HBM 的访问来解决此问题。这主要通过三种核心技术实现。 ### 分块（Tiling）与块处理 FlashAttention 并不一次性处理整个矩阵，而是将输入的 Q、K、V 矩阵分割成能装入高速 SRAM 的小块（**tiles**）。算法会依次加载这些块，在块上执行所有注意力计算并更新最终结果，而**无需在慢速 HBM 中存储完整的注意力矩阵**^{[\[1\]](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_note-arxiv_main-1)}。 ### 在线 Softmax 计算关键的技术突破是“在线”Softmax 计算。标准 Softmax 需要输入向量的所有元素才能进行归一化。FlashAttention 采用了一种改进算法，可以分块计算 Softmax。它维护两个中间值（当前最大值和指数和），这些值会随着新块的处理而更新，从而在不访问整个矩阵的情况下获得精确结果^{[\[2\]](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_note-openreview_fa1-2)}。 ### 将操作融合到单个 CUDA 核心中所有注意力操作（矩阵乘法 QKᵀ、掩码、Softmax、与 V 相乘）都被合并到**一个单一融合的 CUDA 核心** (fused kernel) 中。这极大地减少了对 HBM 的读/写操作次数：算法不再需要多次遍历整个矩阵，而是一次性将一个块加载到 SRAM 中，执行所有计算，然后只写回最终结果。 ## 理论与实践效率 ### 复杂度与最优性 FlashAttention 将内存消耗从 O(N²) 降低到 **O(N)**，实现了线性扩展。事实证明，该算法的 IO 复杂度对于在两级内存层次结构中计算注意力是**理论上最优**的，也就是说，在不改变硬件的情况下，无法更快地执行精确注意力计算^{[\[3\]](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_note-ieee_spectrum_2022-3)}。 ### 实证结果第一版 FlashAttention 取得了显著的改进： - **加速效果**： - BERT-large (序列长度 512)：训练速度提升 **15%**。 - GPT-2 (序列长度 1K)：速度提升 **3倍**。 - Long-Range Arena (1K-4K) 任务：速度提升 **2.4倍**^{[\[1\]](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_note-arxiv_main-1)}。 - **内存节省**：与精确的基线实现相比，内存节省高达 **20倍**。 - **模型质量提升**：由于能够处理更长的上下文，FlashAttention 不仅没有损失模型质量，反而有所提升。例如，GPT-2 的困惑度改善了 0.7，长文档分类任务的准确率提高了 6.4 个百分点^{[\[1\]](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_note-arxiv_main-1)}。 ## 演进与未来发展 FlashAttention 的成功催生了一系列面向硬件的算法。 ### FlashAttention-2 (2023) - 第二代 FlashAttention 第二版旨在更充分地利用 GPU 资源。在最初的 FlashAttention 中，NVIDIA A100 上的效率仅为最大理论值的 25–40%。FlashAttention-2 改进了计算并行化，从而实现了^{[\[4\]](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_note-tridao_flash2-4)}： - 与第一版相比，速度提升了**两倍**。 - 将 GPU 利用率提高到理论最大值的 **50–73%**。 - 扩展了对大小为 256 的注意力头以及多查询注意力 (Multi-Query Attention, MQA) 架构的支持。 ### FlashAttention-3 (2024) - 第三代 FlashAttention 第三版专门针对 **NVIDIA Hopper (H100)** GPU 架构进行了优化^{[\[5\]](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_note-pytorch_blog_fa3-5)}。它利用了新的硬件功能，如**张量核心 (Tensor Cores) 的异步性**和对**FP8**的支持，从而实现了： - 与 FlashAttention-2 相比，速度进一步提升了 **1.5–2倍**。 - 在 FP16 上实现了高达 **740 TFLOPS** 的性能，在 FP8 上接近 **1.2 PFLOPS**。 ### 专用解决方案 FlashAttention 的思想在其他项目中得到了进一步发展： - **FlashInfer** (2025)：一个可定制的注意力引擎，专为 LLM 推理任务优化。它专注于在流式生成模式下高效处理 KV 缓存^{[\[6\]](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_note-arxiv_flashinfer-6)}。 - **FlashMLA** (2024)：一种带有上下文缓存压缩（*latent attention*）的注意力实现，能够在极长序列上节省内存，且信息损失极小^{[\[7\]](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_note-github_flashmla-7)}。 ## 对行业和生态系统的影响 FlashAttention 成为一项基础性突破，并迅速发展为高效训练和推理 LLM 的**行业标准**。它已被集成到 PyTorch 和 Hugging Face 等核心库中，并被大多数大型语言模型（如 LLaMA、MPT、Falcon、Claude 等）所采用。正是 FlashAttention 及其后续版本，在扩大语言模型的**上下文窗口**方面发挥了决定性作用：从 2-4 千个 token (GPT-3) 增加到 12.8 万个 token (GPT-4)，甚至在实验性模型中达到数百万个 token^{[\[8\]](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_note-medium_evolution-8)}。该算法消除了扩展 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. ## 注释 1. ↑ ^{[1.0](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_ref-arxiv_main_1-0)} ^{[1.1](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_ref-arxiv_main_1-1)} ^{[1.2](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_ref-arxiv_main_1-2)} ^{[1.3](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_ref-arxiv_main_1-3)} ^{[1.4](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_ref-arxiv_main_1-4)} Dao, Tri, 等人。“FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness”。*arXiv:2205.14135* \[cs.LG\]，2022年5月28日。[1] 2. ↑ ^{[2.0](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_ref-openreview_fa1_2-0)} ^{[2.1](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_ref-openreview_fa1_2-1)} Dao, Tri, 等人。“FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness”。*OpenReview*。[2] 3. [↑](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_ref-ieee_spectrum_2022_3-0) “We're Training AI Twice as Fast This Year as Last”. *IEEE Spectrum*. [3] 4. [↑](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_ref-tridao_flash2_4-0) Dao, Tri. “FlashAttention-2”. *tridao.me*. [4] 5. [↑](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_ref-pytorch_blog_fa3_5-0) “FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision”. *PyTorch Blog*. [5] 6. [↑](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_ref-arxiv_flashinfer_6-0) “\[2501.01005\] FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving”. *arXiv*. [6] 7. [↑](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_ref-github_flashmla_7-0) “GitHub - deepseek-ai/FlashMLA: FlashMLA: Efficient MLA decoding kernels”. *GitHub*. [7] 8. [↑](https://systems-analysis.info/int/FlashAttention_%E2%80%94_%E9%97%AA%E7%94%B5%E6%B3%A8%E6%84%8F%E5%8A%9B#cite_ref-medium_evolution_8-0) “The Evolution of Flash Attention: Revolutionizing Transformer Efficiency”. *Medium*. [8]