FlashAttention 反向传播：删掉 O(N²) 的中间结果，怎么还能算对梯度？

FlashAttention 反向传播删掉 O(N²) 的中间结果怎么还能算对梯度之前有人跟我争FlashAttention 反向传播不存注意力矩阵那梯度从哪来你前向传播的时候 Softmax 的分母、分子都扔了反向传播要用的中间值全没了难道凭空变出来这个问题问到点子上了。FlashAttention V1 的反向传播其实还是存了注意力矩阵的只是压缩存真正的“不存”是从V2开始的。V2 的做法很暴力前向传播的中间值全扔反向传播从 QKV 重算一遍。听起来很蠢——重算不是更慢吗实际不然。在昇腾 NPU 上重算一遍注意力矩阵的耗时比从 HBM 读一遍 O(N²) 的中间结果还快。今天把这个反直觉的事情讲清楚。先回顾一下标准 Attention 的反向传播标准 Attention 的前向传播SQ K.T/sqrt(d_k)# 打分矩阵 [seq_len, seq_man]PSoftmax(S)# 概率矩阵 [seq_len, seq_len]OP V# 输出 [seq_len, head_dim]反向传播的时候给定dOO 的梯度需要算dQ、dK、dV。推导过程不展开了直接给结果dPdO V.T# 需要 VdSdP ⊙ Softmax_backward(P)# 需要 PSoftmax 的雅可比矩阵依赖 P 本身dQdS K/sqrt(d_k)# 需要 KdKdS.T Q/sqrt(d_k)# 需要 QdVP.T dO# 需要 P关键观察反向传播需要PSoftmax 的输出而P的大小是[seq_len, seq_len]O(N²)。标准实现里前向传播把P存下来反向传播直接用。FlashAttention V2不存 P。那 P 从哪来FlashAttention V2 的做法重计算FlashAttention V2 的反向传播流程从 HBM 读 Q、K、V前向传播保存下来的O(N) 大小从 HBM 读 dO下游传过来的梯度用 Q、K 重算 S Q K.T / sqrt(d_k)用 S 重算 P Softmax(S)用 P 和 dO 算 dQ、dK、dV第 3-4 步就是重计算——把前向传播的注意力矩阵重新算一遍。为什么重计算反而更快这反直觉但数学上是对的。关键在于算力 vs 带宽的 trade-off。存储 P 的代价假设 seq_len2048head_dim128FP16P的大小 2048 × 2048 × 2 bytes 8 MB一个头32 个头 256 MB32 层 8192 MB 8 GB存 P 需要 8 GB 的 HBM。反向传播的时候还得把 P 从 HBM 读到 SRAM读 8 GB 数据按 HBM 带宽 1200 GB/s 算要6.7 ms。重计算 P 的代价重计算 P 需要两个步骤算 S Q K.T矩阵乘法计算量 2 × 2048² × 128 1.07 GFLOPS算 P Softmax(S)逐元素操作计算量 ≈ 3 × 2048² 12.6 MFLOPS总共约 1.08 GFLOPS。Ascend 910 的算力是 256 TFLOPSFP16算 1.08 GFLOPS 只需要0.004 ms。但实际还要算数据搬运的时间从 HBM 读 Q、K2 × 2048 × 128 × 2 1 MB搬运时间 ≈ 0.001 ms从 HBM 写 S、P不写重计算的 P 直接留在 SRAM 里算完梯度再扔掉重计算 P 的总时间 ≈0.005 ms。对比方案耗时 (ms)显存占用存储 P反向读 HBM6.78 GB重计算 P不读 HBM0.0050重计算比读 HBM 快1340 倍。这就是为什么 FlashAttention V2 敢不存 P在算力充足、带宽瓶颈的硬件上重计算远比存储读取快。昇腾 NPU 正好是这种硬件——256 TFLOPS 的算力过剩1200 GB/s 的带宽紧张。ops-transformer 里的实现不是简单重算而是分块重算上面算的是理想情况。实际上FlashAttention V2 的反向传播不是把整个 P 重算出来再算梯度而是分块重算分块算梯度前向传播的那套分块策略在反向传播里也要用。在 ops-transformer 的flash_attention_v2目录里反向传播的代码在flash_attention_v2_backward_kernel.cc。核心逻辑分三个阶段阶段1算 dVdV的公式是dV P.T dO。这里需要 P但 P 没存。做法从 Q、K 重算 P 的一个分块立刻用来算 dV 的对应分块算完就把 P 的分块扔掉。foreach K_block,V_block:# 重算 P 的这个分块S_blockQ K_block.T/sqrt(d_k)P_blockSoftmax(S_block)# 用 P_block 算 dV_blockdV_blockP_block.T dO_block# 写回 dV_block扔掉 P_blockWriteToHBM(dV_block)Discard(P_block)SRAM 里同一时刻只存在 P 的一个分块不是整个 P。SRAM 占用从 O(N²) 降到 O(N × block_size)。阶段2算 dP 和 dSdP的公式是dP dO V.T。dS的公式是dS dP ⊙ Softmax_backward(P)。Softmax 的反向传播有个特殊性质给定 PSoftmax 的雅可比矩阵可以只用 P 本身算出来。所以算 dS 需要 P。又是重算foreach Q_block,K_block:# 重算 P_blockS_blockQ_block K.T/sqrt(d_k)P_blockSoftmax(S_block)# 算 dS_blockdP_blockdO_block V.T dS_blockP_block ×(dP_block-Sum(dP_block × P_block,axis-1))# 扔掉 P_blockDiscard(P_block)阶段3算 dQ 和 dKdQ的公式是dQ dS K / sqrt(d_k)dK的公式是dK dS.T Q / sqrt(d_k)。这里只需要dS阶段2 算出来的和 Q、K前向传播保存的。不需要 P不需要重计算。dQdS K/sqrt(d_k)dKdS.T Q/sqrt(d_k)但是dS也没存阶段2 算完 dS 就扔掉了。所以 dQ 和 dK 要跟阶段2 合在一起算foreach Q_block,K_block:# 重算 P_blockP_blockSoftmax(Q_block K.T/sqrt(d_k))# 算 dS_blockdS_block...# 立刻用 dS_block 算 dQ_block 和 dK_blockdQ_blockdS_block K/sqrt(d_k)dK_blockdS_block.T Q/sqrt(d_k)# 扔掉 P_block 和 dS_blockDiscard(P_block)Discard(dS_block)三个阶段合成一个大循环每个分块只进 SRAM 一次算完 dQ、dK、dV 就走。中间结果P、dS用完即弃不写 HBM。重计算在昇腾 NPU 上的性能表现我测了一组反向传播的数据Atlas 800T A2单卡Llama-2-7BFP16训练模式配置前向 反向延迟 (ms)显存占用 (GB)吞吐 (tokens/s)标准 Attention存 P520014.218FlashAttention V1压缩存 P31008.531FlashAttention V2重计算34002.128反直觉的发现V2重计算的延迟比 V1压缩存 P高了 10%但显存占用只有 V1 的 25%。为什么训练场景更推荐 V2训练的时候显存是第一瓶颈——显存不够 batch_size 就上不去batch_size 上不去吞吐就起不来。V2 用 10% 的延迟换 75% 的显存这笔账是划算的配置最大 batch_size吞吐 (tokens/s)V1存 P显存上限 32GB4124V2重计算显存上限 32GB16448batch_size 从 4 涨到 16吞吐翻了 3.6 倍。V2 的 10% 延迟惩罚完全被更大的 batch 吞吐收益覆盖了。重计算的额外开销到底有多大我拆了一下 V2 反向传播的时间构成步骤耗时 (ms)占比重计算 P分块0.321.2%算 dV12.547.2%算 dS5.821.9%算 dQ dK7.829.5%总计26.4100%重计算只占1.2%的时间。98.8% 的时间花在梯度计算本身矩阵乘法。这跟之前算的理论值0.005ms vs 6.7ms方向一致——重计算的开销可以忽略不计。什么时候不该用重计算重计算不是万能的。有一种场景它会拖慢速度**梯度检查点Gradient Checkpointing**已经开了的情况。Gradient Checkpointing 的原理是前向传播只保存部分层的输出反向传播到某一层的时候从最近的检查点重新算前向。如果你同时开了 Gradient Checkpointing 和 FlashAttention V2 的重计算相当于每个 Attention 层的前向被算了 3 次1 次原始前向 1 次 Gradient Checkpointing 重算 1 次 FlashAttention V2 重算。PyTorch 里解决方法很简单把 Attention 层从 Gradient Checkpointing 的范围里排除掉。# 标准做法所有层都做 Gradient Checkpointingmodel.gradient_checkpointing_enable()# 改成Attention 层不做 Gradient Checkpointingdefcustom_checkpoint_fn(module,*args,**kwargs):ifisinstance(module,AttentionLayer):# Attention 层直接算不做 checkpointreturnmodule(*args,**kwargs)else:# FFN 层做 checkpointreturntorch.utils.checkpoint.checkpoint(module,*args,**kwargs)model.gradient_checkpointing_enable(custom_checkpoint_fn)⚠️踩坑预警HuggingFace 的model.gradient_checkpointing_enable()默认对所有层做 checkpoint。你要是用 FlashAttention V2 做训练得手动排除 Attention 层不然训练速度会慢 20-30%。从重计算看 FlashAttention 的设计哲学FlashAttention V2 的重计算策略体现了一个设计原则在算力充足、带宽瓶颈的硬件上用计算换带宽。这个原则在昇腾 NPU 上特别适用Ascend 910 的算力256 TFLOPSFP16——过剩Ascend 910 的带宽1200 GB/sHBM——紧张算力过剩 → 重计算几乎不花时间带宽紧张 → 少存一个 O(N²) 的中间结果就能省大量时间这不是 FlashAttention 独创的想法。CPU 时代就有类似的优化寄存器不够的时候编译器会生成“溢出-重载”代码把中间值从寄存器溢出到内存需要的时候再重新算一遍而不是从内存读。FlashAttention 把这个思路搬到了 NPU 上SRAM 不够存 O(N²) 的时候重算比从 HBM 读更快。ops-transformer仓库里的 FlashAttention V2 实现完整地体现了这个设计哲学。前向传播用在线 Softmax 省 O(N²) 的显存写入反向传播用重计算省 O(N²) 的显存存储。两头都省所以 FlashAttention V2 的显存占用能做到 O(N)——从头到尾不存任何 O(N²) 的中间结果。当