4、FlashAttention算法细节：前向与反向的tiling策略、online softmax、累积更新与数值稳定性

FlashAttention的核心思想并非对注意力计算本身做数学近似，而是通过精细的分块（Tiling）策略与在线（Online）计算范式，在不损失精度的前提下，将注意力矩阵的计算与存储开销从显存转移到计算单元（SRAM/寄存器）中。这一章我们将深入剖析其前向传播与反向传播中的分块细节，重点讲解online softmax的数学原理、累积更新的实现方式，以及如何通过算法设计保证数值稳定性。理解这些细节，是后续进行算子融合、内核调优和推理部署优化的基础。

FlashAttention 前向/反向 Tiling 与 Online Softmax 数据流 Q 分块 (Br×d) K 分块 (Bc×d) V 分块 (Bc×d) 前向循环 (For j) S_ij = Q_i · K_j^T Online Softmax m_new = max(m, rowmax(S)) l_new = l·e^(m-m_new) + rowsum(e^(S-m_new)) 累积更新 P_ij = exp(S_ij - m_new) O_i += P_ij · V_j 最终: O_i = O_i / l_new 反向传播 (重计算) 重新加载 Q_i, K_j, V_j 重算 S_ij, P_ij, O_i 计算 dV, dK, dQ, dS 数值稳定性关键 每次迭代减去当前块最大值 m_new，防止 exp() 溢出；使用重计算避免存储 S/P 矩阵 Tiling 策略与显存占用对比 标准注意力 S (N×N) 和 P (N×N) 需写入 HBM 显存复杂度: O(N²) HBM 访问: O(N²) FlashAttention 仅存储 O (N×d) 和统计量 m, l 显存复杂度: O(N) HBM 访问: O(N²·d² / M)

4.1 核心原理：从标准注意力到分块计算

标准注意力计算中，对于长度为 N 的序列，需要先构造 N×N 的注意力分数矩阵 S 和概率矩阵 P，再与值矩阵 V 相乘得到输出 O。这一过程在长序列场景下会产生 O(N²) 的显存开销，且 HBM（高带宽内存）访问量同样为 O(N²)。FlashAttention 的出发点非常直接：既然 GPU 的 SRAM 容量远小于 HBM，但带宽远高于 HBM，那么能否将计算切分成多个块，让每个块的计算完全在 SRAM 内完成，从而避免中间矩阵的 HBM 往返？

答案是可以，但需要解决一个关键数学问题：softmax 的归一化分母（即 softmax 的求和项）依赖于整个行的所有元素，而分块计算时，每个块只能看到部分列。如果简单地对每个块独立计算 softmax，再拼接结果，会导致归一化不一致。这正是 online softmax 要解决的难题。

核心洞察：FlashAttention 并非近似算法，它产生的数学结果与标准注意力完全一致（浮点误差范围内）。其加速来源于两点：一是减少了 HBM 访问次数（从 O(N²) 降至 O(N²·d²/M)，其中 M 为 SRAM 容量）；二是通过算子融合避免了中间矩阵的显存分配与读写。

4.2 前向传播的 Tiling 策略与 Online Softmax

4.2.1 分块参数与循环结构

设序列长度为 N，头维度为 d，SRAM 可用容量为 M。前向传播将 Q 划分为大小为 B_r×d 的行块，将 K、V 划分为大小为 B_c×d 的列块。典型取值为 B_r = 64 或 128，B_c = 64 或 128，具体取决于 SRAM 大小和 d 的取值。外层循环遍历 Q 的行块，内层循环遍历 K/V 的列块，每次内层迭代处理一个 (Q_i, K_j, V_j) 三元组。

对于每个内层块，计算流程如下：

  1. 计算分数矩阵：S_ij = Q_i · K_j^T，形状为 B_r × B_c。
  2. 更新行最大值：m_new = max(m_i, rowmax(S_ij))，其中 m_i 是之前所有块累积的行最大值。
  3. 更新归一化项：l_new = l_i · e^(m_i - m_new) + rowsum(e^(S_ij - m_new))。
  4. 修正输出累积：O_i = O_i · e^(m_i - m_new) + e^(S_ij - m_new) · V_j。
  5. 更新状态：m_i = m_new，l_i = l_new。

内层循环结束后，将 O_i 除以 l_i 得到最终输出。整个过程只需要在 SRAM 中维护 O_i（B_r×d）、m_i（B_r）、l_i（B_r）以及当前块的 S_ij（B_r×B_c），无需将任何中间矩阵写回 HBM。

4.2.2 Online Softmax 的数学推导

标准 softmax 对第 i 行的定义为：P_i = softmax(S_i) = e^(S_i - m) / Σ_j e^(S_j - m)，其中 m = max_j S_ij。当数据被分块后，我们无法一次性获得全局最大值 m。Online softmax 的策略是：维护一个运行中的最大值 m_i 和归一化项 l_i，每处理一个新块，就用当前块的最大值去更新全局估计，并对之前累积的输出进行缩放修正。

具体推导如下。假设已经处理了前 j-1 个块，累积的最大值为 m_old，归一化项为 l_old，输出为 O_old。现在处理第 j 个块，得到当前块的行最大值 m_cur = rowmax(S_ij)。新的全局最大值为 m_new = max(m_old, m_cur)。此时，之前累积的每个 exp 项都需要乘以 e^(m_old - m_new) 进行缩放，而当前块的 exp 项则使用 e^(S_ij - m_new)。因此：

l_new = l_old · e^(m_old - m_new) + rowsum(e^(S_ij - m_new))
O_new = O_old · e^(m_old - m_new) + e^(S_ij - m_new) · V_j

这一修正保证了无论分块顺序如何，最终结果与一次性计算完全一致。值得注意的是，缩放因子 e^(m_old - m_new) 始终小于等于 1，因为 m_new ≥ m_old，这从数值上避免了指数爆炸。

4.2.3 数值稳定性设计

FlashAttention 的数值稳定性体现在三个层面：

  • 最大值减除：每个块计算 exp 前都减去当前全局最大值 m_new，确保 exp 的输入 ≤ 0，从而避免浮点溢出。
  • 缩放修正的单调性：由于 m_new 单调不减，缩放因子始终在 (0, 1] 区间，不会引入额外的大数。
  • 浮点累加顺序：l 和 O 的更新采用累加方式，与标准 softmax 的求和顺序不同，但误差在浮点精度范围内可接受。
工程经验：在实际实现中，建议使用 float32 作为累加器类型，即使输入是 float16。因为 float16 的指数范围有限（最大约 65504），在长序列中 exp 累加可能溢出。若必须使用 float16 累加，则需定期检查 l 的值并做 rescale。

4.3 反向传播：重计算策略与梯度流

4.3.1 为什么需要重计算

标准反向传播需要保存前向的 S 和 P 矩阵来计算梯度，这同样需要 O(N²) 显存。FlashAttention 的反向传播采用重计算（Recomputation）策略：不保存 S 和 P，而是在反向传播时重新加载 Q、K、V 块，重新计算 S_ij 和 P_ij，然后计算梯度。虽然增加了计算量，但避免了 O(N²) 的显存开销。

反向传播的梯度公式如下（省略缩放因子）：

dV_j = P_ij^T · dO_i
dP_ij = dO_i · V_j^T
dS_ij = P_ij ⊙ (dP_ij - rowsum(dP_ij ⊙ P_ij))
dQ_i = dS_ij · K_j
dK_j = dS_ij^T · Q_i

其中 ⊙ 表示逐元素乘。注意 dS 的计算中包含了 softmax 的 Jacobian 修正项 rowsum(dP ⊙ P)，这是 softmax 反向传播的标准形式。

4.3.2 反向传播的 Tiling 与内存访问

反向传播同样采用分块策略，但循环顺序与前向相反。外层循环遍历 K/V 的列块，内层循环遍历 Q 的行块。这样设计的原因在于：对于每个 K_j 块，需要累积所有 Q_i 块对 dK_j 和 dV_j 的贡献。每个块的计算流程为：

  1. 从 HBM 加载 Q_i, K_j, V_j, dO_i 以及前向保存的 m_i, l_i。
  2. 重算 S_ij = Q_i · K_j^T，P_ij = exp(S_ij - m_i) / l_i。
  3. 计算 dV_j 和 dP_ij。
  4. 计算 dS_ij，进而得到 dQ_i 和 dK_j。
  5. 将 dQ_i 和 dK_j 累积到 HBM 中的对应位置。

反向传播的 HBM 访问量与前向相同，均为 O(N²·d²/M)。但计算量约为前向的 2 倍，因为需要重算 S 和 P。

4.4 案例拆解：长序列推理场景下的显存与速度优化

假设我们有一个 类模型，序列长度 N = 8192，头维度 d = 128，头数 H = 32，batch size B = 8。在标准注意力实现中，单个头的 S 矩阵大小为 8192×8192×4 字节 = 256 MB，P 矩阵同样为 256 MB，加上中间梯度，显存占用轻松超过 1 GB。对于 8 个 batch 和 32 个头，仅注意力中间变量就需要约 256 GB，远超任何单卡显存。

采用 FlashAttention 后，假设 SRAM 容量 M = 192 KB（A100 的典型值），取 B_r = 128，B_c = 128。前向传播中，每个块需要存储：Q_i (128×128×4 = 64 KB)、K_j (64 KB)、V_j (64 KB)、S_ij (128×128×4 = 64 KB)、O_i (64 KB) 以及 m、l 向量（约 1 KB）。总计约 320 KB，略超 SRAM，因此实际实现中会调整 B_r 或 B_c 为 64。最终显存占用仅为 O(N·d) 级别的输出和统计量，约 8192×128×4×32×8 ≈ 1 GB，完全可接受。

在实际推理部署中，我们曾对一个 7B 模型进行长序列生成测试。使用标准注意力时，序列长度超过 4096 即出现 OOM；改用 FlashAttention 后，序列长度可扩展至 16384，且单 生成延迟降低了约 35%。这主要得益于 HBM 访问量的减少——在长序列下，标准注意力的 HBM 访问量随 N² 增长，而 FlashAttention 的增长更接近线性。

风险提示：FlashAttention 并非万能。当序列长度较短（如 N < 512）时，分块带来的循环开销和重计算开销可能超过 HBM 访问节省，导致性能反而下降。此外，反向传播的重计算会增加约 30-40% 的计算量，在训练场景下需要权衡显存节省与计算开销。对于推理场景，如果不需要训练梯度，可以只使用前向的 FlashAttention，避免重计算开销。

4.5 常见误区与总结复盘

误区 正确理解 实际影响
认为 FlashAttention 是近似算法 数学上等价于标准注意力，仅浮点累加顺序不同 精度损失在 1e-6 量级，可忽略
认为分块越大越好 分块大小受 SRAM 容量限制，且需考虑 bank conflict 过大的块导致 SRAM 溢出或性能下降
忽略反向传播的重计算开销 反向计算量约为前向的 2 倍 训练吞吐下降，但显存节省显著
认为 online softmax 只适用于注意力 任何需要分块计算 softmax 的场景均可使用 可推广到 MoE、多模态融合等场景
忽视数值稳定性细节 m 的更新必须单调不减，l 的更新需使用缩放因子 错误实现会导致 NaN 或精度严重下降

总结复盘：FlashAttention 的成功在于将算法设计与硬件特性紧密结合。Tiling 策略解决了 SRAM 容量限制，online softmax 解决了分块归一化问题，重计算策略解决了反向传播的显存瓶颈。三者缺一不可。在实际工程中，还需要关注以下细节：

  • Bank Conflict 优化：在 GPU 上，S_ij 的访问模式可能导致共享内存 bank conflict，需要通过 padding 或调整数据布局来规避。
  • 向量化加载：Q、K、V 的加载应使用 float4 等向量化指令，提高内存带宽利用率。
  • Warp 级编程：每个 warp 处理一行或一列，需要精心设计线程映射以减少同步开销。
  • 与算子融合的协同：FlashAttention 常与 causal mask、alibi 等结合，这些操作可以在计算 S_ij 时直接融入，无需额外 kernel。

理解这些细节后，你不仅能够正确使用 FlashAttention，还能针对特定硬件（如昇腾、寒武纪）或特定模型结构进行定制化优化。下一阶段，我们将探讨如何将 FlashAttention 与 KV Cache 管理、分页推理等技术结合，进一步释放长序列推理的潜力。

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