数据截至 2026-08-15(上游 commit 0251105a2fb1)

04 · 反向 CUDA kernel

这一章讲什么: 反向传播的三个非显然设计——「不存 P 只存 LSE」的重计算、「dO·O 行和」的预处理恒等式、「换一条并行轴」的 kernel 组织。读完你会理解为什么反向比前向更难写,以及它为什么仍然比标准实现快。

1. 它要解决的小问题

注意力 S=QKᵀ、P=softmax(S)、O=PV 的反向,链式法则给出五个 matmul:

dV = Pᵀ·dO        dP = dO·Vᵀ       dS = P ∘ (dP − dPsum)
dQ = dS·K         dK = dSᵀ·Q

麻烦在于:每一步都要用到 P 或 S——两张 N×N 的表。 标准实现前向时把它们存在显存,反向读回来;这恰恰是第 01 章说要消灭的 N² 往返。

问题一句话:反向能不能也像前向一样,让 S/P 永不落地?

2. 设计一:存 LSE,重算 P

核心观察:有了每行的 logsumexp,重算 P 只要一个 matmul + 一次 exp。

前向只把 O 和 LSE 写回 HBM(第 03 章),autograd 保存的也只有 q, k, v, out, softmax_lse, rng_state(flash_attn/flash_attn_interface.py:869)。反向 kernel 里:

// csrc/flash_attn/src/flash_bwd_kernel.h:536 —— 用 LSE 当指数基准,直接重算 P
FLASH_NAMESPACE::scale_apply_exp2</*scale_max=*/false>(scores, lse, params.scale_softmax_log2);

P = exp2(S·scale·log₂e − LSE·log₂e)。因为 LSE 就是该行的 logsumexp,S − LSE ≤ 0 天然稳定,连 running max 都不需要——这是反向比前向在 softmax 上反而简单的地方。

代价当然是要把 S = Q·Kᵀ 重算一遍。但第 01 章的账依然成立:重算吃的是 tensor core 的算力,省的是 HBM 的带宽,而后者才是瓶颈。

3. 设计二:dPsum 不用 P 就能算

dS = P ∘ (dP − dPsum) 里,dPsum 是逐行的 Σ_k P_ik·dP_ik。直接算又要 P。

有个漂亮的恒等式绕开它:

dPsum_i = Σ_k P_ik·(dO·Vᵀ)_ik = Σ_k P_ik·(dO_i·V_k) = dO_i · (Σ_k P_ik·V_k) = dO_i · O_i

dPsum 就是 dO 和 O 的逐行点积——两个 N×d 的小张量,根本不需要 P。

所以反向前先跑一个独立的预处理 kernel compute_dot_do_o(csrc/flash_attn/src/flash_bwd_preprocess_kernel.h:58),内部调用 dot_do_o(csrc/flash_attn/src/flash_bwd_preprocess_kernel.h:25)把每行的 dO·O 算好存起来;反向 launch 时它总是第一个发射(csrc/flash_attn/src/flash_bwd_launch_template.h:84-87)。主 kernel 里 dP_sum 直接从 gmem 读(csrc/flash_attn/src/flash_bwd_kernel.h:463-465)。

于是 dS 的最终形式是逐元素的(csrc/flash_attn/src/flash_bwd_kernel.h:584-595):

// csrc/flash_attn/src/flash_bwd_kernel.h:584-586(节选)
auto pointwise_mult = [](float p, float dp, float d) {
    return p * (!Is_dropout || p >= 0 ? dp - d : d);   // dS = P ∘ (dP − dPsum)
};

4. 设计三:换一条并行轴

前向按 Q 行块分 CTA。反向照抄的话,dK/dV 需要跨 CTA 累加(每个 Q 块都对所有 KV 块有贡献)——只能靠原子加,而 dK/dV 是反向里最热的累加。

FA2 反向反过来分:一个 CTA 负责一个 KV 列块。

入口 compute_dq_dk_dv(csrc/flash_attn/src/flash_bwd_kernel.h:800),grid 是 (num_n_block, b, h)(csrc/flash_attn/src/flash_bwd_launch_template.h:76-82);
CTA 把自己的 Kⱼ、Vⱼ 载入后,循环流过所有 Q 行块;
acc_dk、acc_dv 在进入循环前清零(csrc/flash_attn/src/flash_bwd_kernel.h:451-452),全程在寄存器累加,循环结束一次写回——dK/dV 零通信;
dQ 则反过来:每个列块都对同一行块有贡献,于是 atomicAdd 到全局的 dq_accum(csrc/flash_attn/src/flash_bwd_kernel.h:678),最后由 flash_bwd_convert_dq_kernel 汇总降精度(csrc/flash_attn/src/flash_bwd_launch_template.h:61-64)。

一句话:把「能私有化的累加」(dK/dV)留在寄存器,把「绕不开的累加」(dQ)交给硬件原子操作。 和前向「Q 不动 K/V 动」正好互为镜像。

5. 图示:反向一个 CTA 的循环

载入 Kⱼ、Vⱼ 常驻 SRAM
   │
   ▼  对所有 Q 行块 i 循环
重算 Sᵢⱼ = Qᵢ·Kⱼᵀ;  Pᵢⱼ = exp2(Sᵢⱼ·scale − LSEᵢ)
   │
   ▼
dP = dOᵢ·Vⱼᵀ;  dS = P ∘ (dP − dPsumᵢ)
   │
   ├──► acc_dv += Pᵀ·dOᵢ     (寄存器)
   ├──► acc_dk += dSᵀ·Qᵢ     (寄存器)
   └──► dQᵢ += dS·Kⱼ  ──► atomicAdd 到 gmem
   │
   ▼  循环结束
写回 dKⱼ、dVⱼ

6. 五个 GEMM 的布局杂技

反向一口气要算 QKᵀ、dO·Vᵀ、Pᵀ·dO、dS·K、dSᵀ·Q 五个 matmul,其中三个的左操作数是寄存器里刚算出来的 P 或 dS 的转置。tensor core 不接受直接转置寄存器 fragment,所以 kernel 的做法是:

把 P / dS 转 fp16 后写进 SRAM(csrc/flash_attn/src/flash_bwd_kernel.h:551-556、csrc/flash_attn/src/flash_bwd_kernel.h:624-628);
用转置的 smem 布局 + ldmatrix 读出来,喂给 tiled_mma_dkv / tiled_mma_dq;
三个累加 GEMM 分别落在 acc_dv(csrc/flash_attn/src/flash_bwd_kernel.h:635)、acc_dq(csrc/flash_attn/src/flash_bwd_kernel.h:655)、acc_dk(csrc/flash_attn/src/flash_bwd_kernel.h:689)。

「寄存器 → smem → 转置读出」是中转代价,但比起让 P 落 HBM,依然是数量级的胜利。

7. 确定性:拿性能换可复现

atomicAdd 的浮点加法顺序不固定 → 默认反向两次跑结果可能逐位不同。传 deterministic=True 后:

grid.x 从「每个 KV 列块一个 CTA」改成「SM 数 ÷ (b·h) 个持久 CTA」,各自跨步扫列块(csrc/flash_attn/src/flash_bwd_launch_template.h:78-82);
每个 CTA 往自己私有的 dq_accum 切片里原子加(csrc/flash_attn/src/flash_bwd_kernel.h:124-125),顺序固定;
最后按固定顺序合并。

注释写得很直白:「each thread block will do atomicAdd to a different dQ_accum buffer」(csrc/flash_attn/src/flash_bwd_kernel.h:124)。代价是更多显存和略低速度(接口文档见 README.md:245)。

8. 关键细节 / 坑

越界 mask 不能省,哪怕数学上「乘以 0 也对」。 作者注释(TD [2023-07-29]):seqlen_k 之外的分数可能大到转 fp16 时溢出成 Inf,进而把 dQ 污染成 NaN——所以必须 mask(csrc/flash_attn/src/flash_bwd_kernel.h:497-501)。性能正确之外的数值正确坑。
causal mask 的位置是试出来的。 「putting this causal masking right after acc_s is much slower for some reason」(csrc/flash_attn/src/flash_bwd_kernel.h:510)——现在它放在 Alibi 之后、dS 之前。
越界行的 LSE 被显式置 +∞(csrc/flash_attn/src/flash_bwd_kernel.h:413-419 注释):若置 0,probs 全 1,有 ALiBi 时会算出 NaN;置 ∞ 则 probs 恒 0,天然安全。
dropout 靠重放 RNG 而不是存 mask。 前向保存 philox 的 seed/offset(第 03 章 §9.4),反向用同一状态重生成同一张 dropout mask(csrc/flash_attn/src/dropout.h:45);前向把 dropout 编码进 P 的符号位(encode_dropout_in_sign_bit,csrc/flash_attn/src/flash_fwd_kernel.h:360-362),反向靠 p >= 0 判断是否被 drop(csrc/flash_attn/src/flash_bwd_kernel.h:585)——一张 N×N 的 mask 也省掉了。
batch·heads 小、序列长时,反向还有 seqk-parallel 变体:compute_dq_dk_dv_seqk_parallel(csrc/flash_attn/src/flash_bwd_kernel.h:827)让多个 CTA 分担同一列块的 m 循环,思路和前向 split-KV 同源。

9. 本章小结

反向 = 重计算(S/P 现场重算)+ 一个恒等式(dPsum = dO·O 行和)+ 一次并行轴对调(dK/dV 寄存器,dQ 原子加)。
五个 GEMM 的转置靠「寄存器写 smem、转置布局读回」实现。
确定性模式用「私有累加切片 + 固定顺序合并」换掉全局原子加。
至此 FA2 的前向/反向都讲完了;下一章看支撑它们的模板工程,以及 FA3/FA4 怎么把同一思想搬到新硬件上。

1. 它要解决的小问题​

2. 设计一:存 LSE,重算 P​

3. 设计二:dPsum 不用 P 就能算​

4. 设计三:换一条并行轴​

5. 图示:反向一个 CTA 的循环​

6. 五个 GEMM 的布局杂技​

7. 确定性:拿性能换可复现​

8. 关键细节 / 坑​

9. 本章小结​