2 min read#flash-attention#attention#gpu#memory-efficiency#io-awareness

Flash Attention

Apr 7, 2026

FlashAttention 是一种 IO-aware 的精确注意力算法，通过减少 HBM（高带宽显存）和 SRAM（片上缓存）之间的数据搬运次数来加速 Transformer 的 self-attention 计算。它不改变注意力的数学结果（是精确计算，不是近似），而是通过算法重排避免将中间矩阵 $S$ 、 $P$ 写回 HBM。

标准 Self Attention 的瓶颈

标准实现的计算流程：

$S = QK^\top$ （ $N \times N$ 矩阵，写回 HBM）
$P = \text{softmax}(S)$ （ $N \times N$ 矩阵，写回 HBM）
$O = PV$

HBM 访问次数为 $O(Nd + N^2)$ ，其中 $N$ 为序列长度(seq_len)， $d$ 为头维度(head_dim)。

核心问题： $S$ 和 $P$ 都是 $N \times N$ 的大矩阵，写完之后立马又读——增加了大量 HBM $\leftrightarrow$ SRAM 的通信开销。能否避免将 $S$ 显式写回 HBM（即避免 materialize），是优化的关键。 $P$ 本身就不需要单独保存。

改进 Softmax

Safe Softmax

在计算 softmax 时将所有 $x_i$ 减去 $\max_{i=1}^{N} x_i$ ，避免指数上溢的数值风险：

$\frac{e^{x_i}}{\sum_{j=1}^{N} e^{x_j}} = \frac{e^{x_i - m}}{\sum_{j=1}^{N} e^{x_j - m}}, \quad m = \max_{j} x_j$

朴素 softmax 需要 3-pass（求 max → 求 sum → 归一化），safe softmax 同样是 3-pass 但数值稳定。

Online Softmax

刚才我们使用替代序列的思想优化并得到了 online softmax 算法，这个思想可以进一步应用到 self attention 的计算上（因为我们并不是为了计算 softmax，而是计算 softmax 乘以矩阵 $V$ 的值）。——虽然 softmax 无法 1-pass 解决，但是 self attention 可以 1-pass 解决。

将 safe softmax 的前两个循环（求 max 和求 sum）合并为一个循环：创建替代序列 $d'$ 代替 $d$ 进行计算，满足 $d'_N = d_N$ 。核心思想是在遍历新元素时，同时更新当前最大值 $m$ 并对已有的部分和做 rescale：

$d'_j = d'_{j-1} \cdot e^{m_{j-1} - m_j} + e^{x_j - m_j}$

这样只需 2-pass（合并的 max+sum → 归一化）。

替代序列思想应用到 Self Attention

实际上不需要把 softmax 的结果 $a_i$ 显式算出来，只需要最终乘以 $V$ 的结果 $\mathbf{o}_i$ 。利用和 online softmax 相同的替代序列思想，可以在遍历 K/V 的同时增量维护输出向量：

$\mathbf{o}'_j = \mathbf{o}'_{j-1} \cdot \frac{d'_{j-1}}{d'_j} + \frac{e^{s_j - m_j}}{d'_j} \cdot v_j$

其中 $s_j = q^\top k_j$ 是当前 query 和第 $j$ 个 key 的点积。这两个循环（online softmax + 增量累加 $O$ ）也可以 fuse 为 1-pass。

Tiling

刚才我们提到需要计算的是 softmax 乘以矩阵 $V$ 的值，既然有矩阵乘法，就可以再进一步利用 tiling 技术进行优化。

将 Q、K、V 矩阵分成若干小块（tile），每次只把一个 tile 加载到 SRAM 中完成局部计算。为确保数值稳定性，每个 tile 单独维护自己的指数偏移量 $m$ 和归一化因子 $d'$ ，最终合并时 rescale 到统一的值。

Tiling 使得 HBM 访问次数降为 $O(N^2 d^2 M^{-1})$ ，其中 $M$ 为 SRAM 大小。当 $M = \Theta(Nd)$ 时，复杂度为 $O(Nd)$ ，是关于序列长度的线性 IO 复杂度。

Flash Attention V2 的改进

改进一：调整内外循环

FlashAttention V1 采用外层循环 K/V、内层循环 Q/O 的方式，导致矩阵 $O$ 被额外加载。V2 调换为外层循环 Q/O、内层循环 K/V，减少了 $O$ 的 HBM 读写。

改进二：减少非矩阵乘法运算

现代 GPU 对矩阵乘法有专门硬件加速（Tensor Core），矩阵乘法的 FLOP 吞吐一般是非矩阵乘法的 16 倍左右。

V1 的每轮迭代末尾都有一个 rescale 操作（除以 $d'_i$ ）。V2 将这个 rescale 推迟到循环结束后统一执行——不再在每轮迭代中除以 $d'_i$ ，而是等循环结束后对 $\mathbf{o}'_N$ 统一除以 $d'_N$ 。减少了循环体内的非矩阵乘法运算。

改进三：Warp 级并行度

一个 thread block 在 GPU 上会被划分为若干 warp（通常 4 个）。V1 的天然划分方式会在 warp 之间产生额外通信。V2 手动调整 warp 的工作分配，让同一个 warp 内完成更多独立计算，减少跨 warp 的同步开销。

Note: 我感觉左图中的划分应该也不是刻意为之，只是在送到 GPU 时被天然划分成了若干个 warp，手动划分的话能减少通信次数。