Flash Attention

Flash Attention 是一种面向 GPU 显存层次结构的 IO 感知精确注意力算法，由 Tri Dao 等人于 2022 年提出，通过分块（tiling）与在线 softmax 重计算，将中间结果保留在片上 SRAM 中，从而大幅减少对高带宽显存（HBM）的读写次数。它在不牺牲数值精度的前提下显著加速 Transformer 训练与推理，并将注意力层显存占用从序列长度的二次方降至线性。

标准 Transformer 的自注意力在长序列上存在严重效率瓶颈。

• 时间与显存复杂度均为序列长度的 二次方，制约长上下文应用
• 真正的瓶颈不在算术运算量，而在 GPU HBM（高带宽显存） 与片上 SRAM 之间频繁的数据搬运
• 已有的近似注意力方法（如 Linformer）虽降低了理论 FLOPs，却常因 IO 开销大而无法带来实际墙时间（wall-clock time）加速
• IO 感知设计的核心思路：减少 HBM 访问次数，而非仅减少浮点运算量

Flash Attention 通过三项技术组合将注意力计算「锁定」在快速片上缓存内完成。

• 分块（Tiling）：将 Q、K、V 矩阵切分为小块，每次只将一块加载进 SRAM，在片上完成局部 QKᵀ 乘法与 softmax，避免完整注意力矩阵落盘 HBM
• 在线 Softmax（Online Softmax）：利用 log-sum-exp 增量更新技巧，逐块动态修正归一化因子，无需预先得到完整行即可精确计算 softmax
• 反向传播重计算（Recomputation）：反向传播时重新计算注意力权重，而非存储，以算力换显存；最终结果与标准注意力数学等价（精确算法，非近似）
• 上述组合使 HBM 访问量相对标准实现降低数倍，实测在 A100 上端到端加速约 2–4 倍

Flash Attention 自首版发布后持续迭代，形成系列版本。

• 2022：Flash Attention v1 发布（arXiv 2205.14135），在 GPT-2 训练上实现最高 7.6 倍加速，收录于 NeurIPS 2022；作者为 Tri Dao、Daniel Y. Fu 等斯坦福大学团队
• 2022–2023：HuggingFace Transformers、PyTorch 2.0（SDPA 接口）、Megatron-LM 等主流框架陆续集成
• 2023：Flash Attention v2 发布，新增沿序列长度维度的并行化，减少非矩阵乘法 FLOPs，实测再提速约 2 倍
• 2023：Flash Decoding 变体发布，专针对推理阶段小 batch 长 KV 序列场景优化
• 2024：Flash Attention v3 发布，深度适配 NVIDIA Hopper（H100）架构，利用 TMA 异步拷贝、Warp Specialization 流水线，并支持 FP8 低精度，进一步逼近硬件上限

Flash Attention 在实际工程中带来可量化的显著收益。

• 训练速度：GPT-2 训练加速最高 7.6 倍（v1 官方基准）
• 显存占用：注意力层由 O(N²) 降至 O(N)（N 为序列长度），使超长上下文训练成为可能
• 长序列能力：首次使 Transformer 在 Path-X（16K 序列长度）任务上达到高于随机基线的准确率（61.4%）
• 质量无损：精确算法，无精度-速度权衡，模型困惑度（perplexity）不受影响

Flash Attention 已成为主流 LLM 训练与推理栈的标配底层组件。

• 官方实现以 CUDA/C++ 扩展形式开源于 Dao-AILab/flash-attention，提供 PyTorch 接口
• HuggingFace Transformers、vLLM、SGLang、TGI 等主流框架均内置支持
• PyTorch 2.0 的 scaled_dot_product_attention API 已原生集成 Flash Attention 后端
• 与 KV Cache、投机解码（Speculative Decoding）、PagedAttention 等技术协同，共同提升推理吞吐

Flash Attention 常与几个相关技术混淆，需加以区分。

• vs 近似注意力：线性 Attention（Mamba、RetNet）通过改变公式达到 O(N) 理论复杂度，但牺牲精度；Flash Attention 保持精确计算，仅优化实现层面的 IO
• vs PagedAttention：Flash Attention 优化注意力计算内核；PagedAttention（vLLM 提出）优化推理时的 KV Cache 显存管理，两者正交互补，推理引擎通常同时使用
• vs Sparse Attention：稀疏注意力跳过部分 token 对，改变了注意力模式；Flash Attention 对全量 token 对精确计算，仅优化访存效率

Flash Attention 并非在所有场景均能带来显著收益，有其适用边界。

• 硬件依赖：针对 NVIDIA GPU（A100/H100）深度优化，在 AMD GPU、TPU 或 CPU 上需要专项适配
• 短序列收益有限：序列较短时 IO 优化带来的收益相对不明显
• 变体适配成本：滑动窗口注意力、自定义稀疏掩码等非标准变体需要额外工程适配
• 常见误区：Flash Attention 显著降低注意力矩阵的显存，但模型权重、优化器状态等显存不受影响，整体显存节省需结合实际场景评估