跳到主要内容

数据截至 (上游 commit 0251105a2fb1)

03 · 前向 CUDA kernel

这一章讲什么: 把第 01 章的 tiling 骨架 + 第 02 章的 online softmax,落到 csrc/flash_attn/src/flash_fwd_kernel.h 的真实 CUDA kernel 上:grid 怎么分、主循环每一步干什么、mask 怎么省、以及推理期的 split-KV 是怎么回事。

1. 它要解决的小问题​

数学正确只是入场券。kernel 还要回答三个工程问题:

  • 并行度:上百万个 tile,怎么映射到上百个 SM?
  • 吞吐:搬运(cp.async)和计算(tensor core)怎么互相掩护,谁都不等谁?
  • 特判:causal mask、越界、变长序列,怎么不让边界检查拖慢主循环?

2. 网格划分:一个 CTA 一个 Q 行块​

前向的 launch 配置在 run_flash_fwd(csrc/flash_attn/src/flash_fwd_launch_template.h:55),grid 就一行:

// csrc/flash_attn/src/flash_fwd_launch_template.h:63-64
const int num_m_block = (params.seqlen_q + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM;
dim3 grid(num_m_block, params.b, params.h);

每个 CTA(thread block)负责一个 (Q 行块 × batch × head),在函数 compute_attn_1rowblock(csrc/flash_attn/src/flash_fwd_kernel.h:55)里完成:

  1. 把自己的 Q tile 一次载入 SRAM(kBlockM 常见 128/64 行);
  2. 循环:K/V 按 kBlockN 大小的列块依次流过 SRAM;
  3. 输出 O tile 在寄存器里累加,循环结束才写回。

Q 不动、K/V 流动——这样每个 K/V 元素从 HBM 只读一次。

3. 图示:一个 CTA 的主循环​

Q tile 载入 SRAM(一次)
   │
   ▼  对 K/V 第 j 块循环(j 从右往左)
cp.async 预取 Kⱼ₋₁ ──► SRAM   ←── 与计算重叠
   │
   ▼
Sⱼ = Q·Kⱼᵀ        (tensor core,寄存器)
   │
   ▼
mask + online softmax(m, l 更新,acc_o 折算)
   │
   ▼
acc_o += Pⱼ·Vⱼ    (tensor core)
   │
   ▼  循环结束
O = acc_o / l ──► 写回 gO;LSE ──► 写回 gLSE

怎么读这张图: 预取和计算是并行的两条线,靠 cp.async 的 fence/wait 咬合;S/P 全程只在寄存器和 SRAM 里。

4. 主循环逐步对照源码​

主循环体在 csrc/flash_attn/src/flash_fwd_kernel.h:309-436,一轮迭代的次序:

步骤干什么代码锚点
预取算当前块时,把下一块 K 用 cp.async 发向 SRAMcsrc/flash_attn/src/flash_fwd_kernel.h:341-345
GEMM0acc_s = Q·Kⱼᵀ,tensor core,结果在寄存器csrc/flash_attn/src/flash_fwd_kernel.h:326
maskcausal/越界位置写 −∞csrc/flash_attn/src/flash_fwd_kernel.h:335
online softmax更新 m、l,折算 acc_o(第 02 章)csrc/flash_attn/src/flash_fwd_kernel.h:350
降精度acc_s fp32 → fp16/bf16 得 P,喂给下一个 GEMMcsrc/flash_attn/src/flash_fwd_kernel.h:355
GEMM1acc_o += Pⱼ·Vⱼcsrc/flash_attn/src/flash_fwd_kernel.h:374

两个值得停留的细节:

  • fence 的位置是踩过坑的。 预取下一块 K 的 cp_async_fence 必须包在 if (n_block > n_block_min) 里,注释明说否则「synchronization isn't right and we get race conditions」(csrc/flash_attn/src/flash_fwd_kernel.h:343-344)。
  • V 的 smem 布局是转置的。 SmemLayoutVtransposed(csrc/flash_attn/src/kernel_traits.h:93)让第二个 GEMM 能直接用 ldmatrix 读 V 的转置,不必真的做一次转置。

5. mask 的工程:把「要带 mask 的块」隔离出来​

causal 注意力里只有对角线附近的块真正需要逐元素 mask;其余块全开。kernel 把这个观察写成了两段循环:

// csrc/flash_attn/src/flash_fwd_kernel.h:305-309(结构示意,非逐字)
constexpr int n_masking_steps = (!Is_causal && !Is_local) ? 1 : ceil_div(kBlockM, kBlockN) + 1;
for (int masking_step = 0; masking_step < n_masking_steps; ++masking_step, --n_block) { ... }
// 之后才是无 mask 的快速循环(flash_fwd_kernel.h:385)
for (; n_block >= n_block_min; --n_block) { ... }
  • 第一段从最右侧(对角线所在)的块开始,逐块做完整 mask;
  • 第二段处理剩下所有块,apply_mask</*Causal_mask=*/false> 只做越界防护(csrc/flash_attn/src/flash_fwd_kernel.h:410)。

配套地,causal 时每个 Q 行块要看的 K 块上界直接被裁掉:n_block_max 取到对角线为止(csrc/flash_attn/src/flash_fwd_kernel.h:91-94)——因果掩码省掉的那一半计算,是靠「根本不看右上三角的块」实现的,不是靠把分数乘 0。

6. Epilogue:一次归一化,两个输出​

循环结束后(csrc/flash_attn/src/flash_fwd_kernel.h:438 起):

  1. normalize_softmax_lse 做 O /= l 并算出 LSE(csrc/flash_attn/src/flash_fwd_kernel.h:440,数学见第 02 章);
  2. O 降精度后经 SRAM 中转写回 gO——sO 直接复用 sQ 的 smem(csrc/flash_attn/src/flash_fwd_kernel.h:443-444),因为 Q 已经不需要了;
  3. LSE 写入 gLSE,供反向重算 P 用。

输出就两个:O(N×d)和 LSE(N)。这是「反向重计算」设计的前向半边,见第 04 章。

7. 推理特化:split-KV 与一行查询的 tricks​

7.1 问题:解码期 grid 太小,SM 填不满​

自回归解码时 seqlen_q = 1,grid 只剩 (1, batch, head)——几十上百个 CTA,填不满 108 个 SM,而且每个 CTA 要串行读完整个 KV cache。

7.2 对策:沿 K 维再切一刀​

flash_fwd_splitkv_kernel 把 KV 也分段,多个 CTA 各算一段的「部分注意力」,再由 flash_fwd_splitkv_combine_kernel 合并(csrc/flash_attn/src/flash_fwd_launch_template.h:107 的 grid 变成 (num_m_block, num_splits, b*h))。

切几份由启发式决定:num_splits_heuristic(csrc/flash_attn/flash_api.cpp:281)挑「占用率 ≥ 最优值 85% 的最小 split 数」——注释里算了笔账:48 个 block 对 108 个 SM,2 split 效率 0.89,优于 3 split 的 0.67(csrc/flash_attn/flash_api.cpp:275-280 注释)。

7.3 合并:用 LSE 做 logsumexp 加权​

各 split 产出自己的 O_partial 和 LSE_partial。combine_attn_seqk_parallel(csrc/flash_attn/src/flash_fwd_kernel.h:1118)先对 LSE 沿 split 维再做一次 logsumexp,然后按 e^(LSE_s − LSE_total) 加权求和——数学上这就是把 online softmax 的「折算」在 split 维上再做一遍。第 02 章提过的「行和为 0 时写 −∞」(csrc/flash_attn/src/softmax.h:179)在这里用作「该 split 无贡献」的标记。

7.4 白捡的 reshape:GQA 解码把「头」折进序列​

seqlen_q == 1 且 Q 头数 > KV 头数时,host 侧把 Q reshape 成 (b, ngroups, nheads_kv, d)——把「一个 KV 头服务的 Q 头组」变成序列维,CTA 形状立刻变好(csrc/flash_attn/flash_api.cpp:427-434,注释署名 Daniel Haziza)。零 kernel 改动,纯张量视图游戏。

8. 变长序列(varlen)与 KV cache​

  • varlen:batch 内序列不等长时,不打 padding,而是把所有序列拼接成一条,用 cu_seqlens_q/cu_seqlens_k(累积长度数组)标出边界。kernel 里 BlockInfo</*Varlen=*/!Is_even_MN>(csrc/flash_attn/src/flash_fwd_kernel.h:87)按 bidb 查出本序列的真实长度与偏移,主循环逻辑零改动。Python 入口 flash_attn_varlen_func(flash_attn/flash_attn_interface.py:1391)。
  • KV cache:flash_attn_with_kvcache(flash_attn/flash_attn_interface.py:1485)支持原地追加新 K/V、可选 rotary、paged KV(block table),走 mha_fwd_kvcache(csrc/flash_attn/flash_api.cpp:1244)并强制 split-KV 路径(csrc/flash_attn/flash_api.cpp:1514)。
  • causal 对齐约定:v2.1 起 seqlen_q ≠ seqlen_k 时 causal mask 对齐右下角而非左上角(README.md:420),做 KV cache 解码时这正是你期待的语义;seqlen_q == 1 且无 ALiBi 时 causal 直接被关掉(csrc/flash_attn/flash_api.cpp:426)。

9. 关键细节 / 坑​

  1. 迭代方向是从右往左(n_block 从大到小)。这正好让「需要 mask 的对角线块」先被处理,后面的快速循环一个 if 都不用(inferred——从 flash_fwd_kernel.h:305-309 的步进方向与 causal 裁剪位置可读出来)。
  2. smem 顶格用。 kSmemSize 是 Q、K、V 三块布局的总和,Q/K 可共享时取 max(csrc/flash_attn/src/kernel_traits.h:109);tile 尺寸随 head_dim 和架构变化,Python 侧 _get_block_size_n 要与 CUDA 侧保持一致(flash_attn/flash_attn_interface.py:31,注释明说 "should match the block sizes in the CUDA kernel")。
  3. bank conflict 是被逐个打掉的。 gmem→smem 的拷贝线程布局特意用 kBlockKSmem 而不是 kBlockKGmem,注释记录了 d=128 时快 6–10% 的原因(csrc/flash_attn/src/kernel_traits.h:114-119);smem 内部则用 Swizzle 异或打乱列地址(csrc/flash_attn/src/kernel_traits.h:72)。
  4. dropout 的 RNG 状态由 0 号 CTA 的 0 号线程抢在所有提前退出之前保存(csrc/flash_attn/src/flash_fwd_kernel.h:81-85 注释)——否则短序列 CTA 早退,反向就没有 RNG 状态可重放。

10. 本章小结​

  • grid = Q 行块 × batch × head;CTA 内 Q 常驻、K/V 流动,每个 K/V 元素只读一次。
  • 双缓冲(cp.async)+ 两段 mask 循环 + 对角线裁剪,把「正确性特判」全部挤出主循环。
  • 推理期靠 split-KV(加 LSE 合并)和 GQA reshape 把解码也喂饱。
  • 前向只留 O 和 LSE 给反向——下一章看反向怎么凭这两样东西把梯度算出来。