在 Transformer 训练中，注意力层常被概括为矩阵乘法问题，但真正限制性能的往往不只是浮点运算量，还有 GPU 不同层级存储器之间的数据搬运。FlashAttention 的核心价值，就是从 IO 感知的角度重新组织精确注意力计算，通过减少高带宽显存 HBM 的读写、提高片上存储器复用率，让 GPU 将更多时间用于有效计算，进而改善训练吞吐量和长序列处理能力。
传统注意力为何容易受 IO 限制
标准缩放点积注意力需要依次计算 QK 的转置乘积、缩放、掩码、Softmax、Dropout 以及与 V 的乘积。朴素实现通常会把规模为序列长度平方的注意力分数矩阵写入 HBM，后续算子再从 HBM 读取。每经过一个独立算子，中间结果都可能发生一次或多次读写。
现代 GPU 的矩阵计算单元速度很快，但 HBM 的访问速度和能耗成本仍明显高于片上 SRAM。若内核频繁在 HBM 与计算单元之间搬运大型中间张量，即便理论浮点运算量不变，实际执行时间也可能主要消耗在等待数据上。这解释了为什么某些降低 FLOPs 的注意力近似方法，未必能获得同等比例的端到端加速。
FlashAttention 的关键思路
一、用分块计算提高片上数据复用
FlashAttention 不一次性生成完整注意力矩阵，而是把 Q、K、V 划分为能够放入片上 SRAM 的数据块。内核加载一个查询块后，依次遍历键和值的数据块，在片上完成局部分数计算、Softmax 更新和输出累加。数据一旦进入 SRAM，就尽可能参与更多计算，减少重复访问 HBM。
这种 Tiling 策略没有消除注意力随序列长度平方增长的主要计算量，但显著降低了存储层级之间的数据流量。优化目标由“少做多少次乘加”扩展为“每次从 HBM 读取的数据能够被利用多少次”，因此更符合 GPU 的真实执行特征。原始论文将其称为 IO-aware，也就是显式考虑存储层级读写成本的算法设计。相关原理可参考 FlashAttention 原始论文[1]。
二、在线 Softmax 保持精确结果
分块带来的难点是 Softmax：一行注意力分数被拆成多个块后，当前块并不知道整行的最大值和归一化分母。FlashAttention 使用在线 Softmax，在遍历各个 K 块时维护每一行的运行最大值与指数和。当新的分数块到达后，内核更新最大值，并按照新的数值尺度修正此前的累加结果。
这种递推方式既避免保存完整分数矩阵，也能得到与普通 Softmax 数学等价的精确注意力结果。它不是通过稀疏化、低秩近似或缩短上下文来换取速度，因此通常不需要改变模型结构和训练目标。需要注意的是，受浮点运算顺序影响，不同内核之间仍可能出现正常的数值舍入差异。
三、融合算子减少中间张量落盘
传统实现若将矩阵乘法、缩放、掩码、Softmax 和输出乘法交给多个内核，中间张量需要反复写回显存。FlashAttention 把这些步骤组织到更紧凑的融合内核中，在 SRAM 或寄存器里持续处理临时结果，只把必要的输出和少量归一化统计量写回 HBM。
算子融合同时减少了内核启动次数和全局内存同步，但其主要收益仍来自减少大型中间数据的物化。特别是在序列较长时，完整注意力矩阵迅速增大，避免存储该矩阵不仅节省显存，也能缓解显存带宽压力。
四、反向传播选择重计算而非全量保存
训练还需要保存前向激活，以便反向传播计算梯度。若保留完整注意力概率矩阵，显存占用会随序列长度平方增长。FlashAttention 保存输出及必要的 Softmax 统计量，在反向阶段按块重新计算部分中间值，而不是从 HBM 读取此前保存的全部注意力数据。
重计算会增加一定的算术工作，但省去了大量中间张量的写入、保存和读取。当系统受显存容量或带宽约束时，用相对便宜的计算换取昂贵的 IO，整体执行反而可能更快。这也是 FlashAttention 与一般激活检查点思路相通、但优化粒度更深入到注意力内核内部的地方。
IO 优化如何转化为训练吞吐量
- 缩短注意力层耗时：减少 HBM 往返后，计算单元等待数据的时间下降，注意力内核能够获得更高的有效利用率。
- 降低激活显存压力：不物化完整注意力矩阵，让训练可以使用更长序列、更大的批次，或减少激活检查点的使用范围。
- 减少梯度累积负担：如果节省出的显存允许增大单步微批次，完成同一全局批次所需的累积次数可能减少，从而降低部分额外开销。
- 改善端到端吞吐：单个注意力内核加速并不等于整个模型获得相同比例提升，但当注意力占训练时间较高时，收益会更明显。
实际提升幅度取决于 GPU 架构、数据类型、序列长度、头维度、批次大小、掩码形式以及模型中其他算子的占比，因此不宜脱离测试环境引用单一加速数字。短序列任务可能更容易受到内核启动或其他网络层限制；长序列任务通常更能体现降低 IO 和激活显存的价值。
落地使用时应关注什么
- 先检查兼容性：确认 GPU、CUDA 或 ROCm、PyTorch 版本、数据类型和头维度是否符合所使用实现的要求，具体条件应以 官方代码仓库[2] 为准。
- 进行端到端基准测试：同时测量每秒处理 Token 数、单步时间、峰值显存和可用批次，而不是只比较注意力微基准。
- 保持测试条件一致：固定模型结构、序列长度、精度模式、梯度累积、编译选项和随机种子，排除数据加载或首次编译造成的干扰。
- 检查数值稳定性：比较损失曲线、梯度范数和关键评测指标。混合精度训练还应观察是否出现溢出、NaN 或异常收敛。
- 区分训练与推理场景：训练重视前向与反向性能及激活显存，增量推理则更容易受到 KV Cache、批处理策略和解码长度影响，不能直接套用训练结论。
总结
FlashAttention 的突破并不是改变注意力公式，而是改变数据在 GPU 存储层级中的流动方式。它通过分块、在线 Softmax、算子融合和反向重计算，避免在 HBM 中物化完整注意力矩阵，以更少的数据搬运完成精确注意力计算。由此带来的显存节省和内核效率提升，可以为更长上下文、更大批次以及更高训练吞吐量创造条件。评估其价值时，应把理论 IO 优势与具体硬件、模型配置和端到端测试结合起来，而不是只关注孤立的峰值加速结果。
