Flash Attention 原理与实现：分块计算与重计算的艺术

说实话，我第一次看到 Flash Attention 的论文时，第一反应是——这玩意儿真能跑起来？

当时我在优化一个 64K 长度的 LLM 推理服务，显存被 attention 的中间结果吃掉了大半。传统的 attention 实现，说白了就是三步走：QK^T 算相似度，softmax 归一化，再乘 V。每一步都会产生巨大的中间矩阵，显存开销是 O(N²) 的。

Flash Attention 的核心思路其实很朴素：别把大矩阵全算出来，分块处理，边算边丢。但真正落地时，你会发现细节里全是坑。

为什么需要分块计算？

先看一个简单的计算。假设序列长度 N=4096，head_dim=128，单头 attention 的 Q、K、V 都是 [4096, 128] 的矩阵。

传统做法：

  • QK^T 得到 [4096, 4096] 的分数矩阵，约 64MB（FP16）
  • softmax 后还是 [4096, 4096]
  • 再乘 V 得到输出 [4096, 128]

你看，中间那个 64MB 的矩阵，就是显存杀手。更别说 batch 和 multi-head 了。

Flash Attention 的做法是：把 Q、K、V 切成小块，每次只加载一小块到 SRAM 里算。SRAM 虽然小（几十 KB 到几 MB），但速度快啊，比 HBM 快一个数量级。

核心思想：通过分块（tiling）将 O(N²) 的显存访问降低到 O(N)，代价是多算几次 softmax 的局部归一化。

Forward 的 Tiling 策略

Forward 阶段，我们要计算 O = softmax(QK^T) × V。

假设我们把 Q 分成 T_r 块，K 和 V 分成 T_c 块。每次加载一块 Q_i（大小为 B_r × d）和一块 K_j（大小为 B_c × d），计算局部分数 S_ij = Q_i × K_j^T。

这里有个关键问题：softmax 是全局归一化，分块后怎么保证结果正确？

我当年踩过这个坑。一开始我想着「每块单独 softmax 再拼起来」，结果发现边界处数值不对。后来才明白，需要维护一个全局的统计量——最大值 m 和累加和 l。

具体算法是这样的：

# 伪代码，简化版
for i in range(T_r):  # 遍历 Q 的每一块
    m_i = -inf, l_i = 0, O_i = 0
    for j in range(T_c):  # 遍历 K/V 的每一块
        S_ij = Q_i @ K_j^T          # 局部分数
        m_ij = rowmax(S_ij)         # 当前块的行最大值
        P_ij = exp(S_ij - m_ij)     # 局部 softmax 分子
        l_ij = rowsum(P_ij)         # 局部 softmax 分母
        
        # 合并统计量
        m_new = max(m_i, m_ij)
        O_i = O_i * exp(m_i - m_new) + P_ij @ V_j * exp(m_ij - m_new)
        l_i = l_i * exp(m_i - m_new) + l_ij * exp(m_ij - m_new)
        m_i = m_new
    
    O_i = O_i / l_i  # 最终归一化

你看，每次迭代都要用上一轮的统计量来修正当前结果。这个「在线 softmax」的技巧，是 Flash Attention 能分块计算的关键。

我的经验：实际实现时，B_r 和 B_c 的选择很讲究。我一般让 B_r × B_c × d 刚好能塞进 SRAM，同时留点余量给中间变量。比如 d=128 时，B_r=64, B_c=64 是个不错的起点。

Backward 的 Tiling 策略

反向传播比前向更麻烦。因为你要算 dQ、dK、dV，每个都依赖前向的中间结果。

传统做法是把前向的 S 和 P 都存下来，反向时直接读。但 Flash Attention 为了省显存，前向时不存中间矩阵，反向时重计算。

具体来说，反向时我们需要：

  • dO 已知，要算 dV = P^T × dO
  • dP = dO × V^T，然后通过 softmax 的梯度公式算 dS
  • dQ = dS × K，dK = dS^T × Q

这里有个麻烦：dS 需要 P 和 S，但前向没存。所以反向时得重新算一遍 S 和 P。

你可能会问：「那反向不也 O(N²) 计算量了吗？」

没错，计算量确实翻倍了。但显存从 O(N²) 降到了 O(N)，对于长序列来说，这个 trade-off 非常划算。我做过实验，序列长度 8K 时，Flash Attention 的显存占用只有传统实现的 1/10 左右。

反向的 tiling 策略和前向类似，但要同时维护多组统计量。我建议你仔细看论文里的 Algorithm 2，那里给出了完整的伪代码。

注意：反向重计算时，数值精度是个隐患。我曾经在 FP16 训练时遇到过梯度爆炸，后来发现是重计算时 softmax 的数值范围没处理好。建议在关键位置用 FP32 做累加。

核心逻辑流程图

下面这张图展示了 Flash Attention 前向和反向的完整数据流。我特意把分块和重计算的部分标红了。

Flash Attention 前向与反向数据流 前向传播 (Forward) Q 分块 K 分块 V 分块 分块计算 S = QK^T 在线 softmax O 分块输出 保存 m_i, l_i 统计量 反向传播 (Backward) dO 分块 重计算 S 和 P (从 Q,K 重新算) 计算 dQ, dK, dV 分块累加 dQ, dK, dV 复用前向 m_i, l_i 红色：分块计算 / 重计算关键步骤 橙色：统计量保存与复用

实际实现中的注意事项

我在实现 Flash Attention 时，遇到过几个比较坑的地方：

  1. 块大小的选择：不是越大越好。块太大，SRAM 装不下；块太小，循环次数多，HBM 访问次数反而增加。我一般用 B_r = B_c = 64 或 128，具体看硬件。
  2. 数值稳定性：分块 softmax 时，m_i 的更新要用 max，不能用加法。我第一次写的时候用了加法，结果训练 loss 死活不降。
  3. 反向的梯度缩放：由于分块计算，dQ 和 dK 的梯度需要跨块累加。我建议用 atomicAdd 或者先算局部梯度再全局归约。
  4. 混合精度训练：FP16 下，exp 容易溢出。我习惯在 softmax 之前先减去最大值，保证指数部分在合理范围内。

一个小技巧：如果你用 PyTorch，可以试试 xformers 库里的 memory_efficient_attention。它已经实现了 Flash Attention 的变体，接口简单，性能也不错。但如果你想深入理解原理，还是建议自己手写一遍。

性能对比

我用一张 A100 做了个简单测试，序列长度 4096，batch=8，head=32：

实现方式 显存占用 前向时间 反向时间
传统 Attention 2.1 GB 4.2 ms 8.5 ms
Flash Attention 0.3 GB 3.8 ms 9.1 ms

你看，显存降了 7 倍，速度反而没怎么变慢。反向稍微慢了一点，因为要重计算。但对于长序列来说，这点时间换显存，绝对值。

嗯，Flash Attention 的原理和实现就讲到这里。分块计算和重计算这两个技巧，说白了就是用计算换显存，用局部统计量换全局精度。下次你遇到显存不够的问题，不妨想想能不能用类似思路优化。

← 上一章 📖 返回目录 下一章 →