T5 讲解：FlashAttention forward

① 先看全貌：省的不是 FLOPs，是显存读写

attention：O = softmax(Q Kᵀ · scale) V，Q、K、V、O 的形状都是 [B, H, S, D]。

naive 实现分三步，每一步都把中间结果写回显存：

算一下题目页的测试规模（B = 4、H = 16、S = 4096、D = 64，fp16）：

naive 实现要写一次 S、读一次 S、写一次 P、读一次 P：分数矩阵和 softmax 结果加起来写 4 GB、读 4 GB，是输入输出本身的 64 倍。S 再翻倍，这一项就变成 4 倍。

② program 怎么分

drills27–35 行 · host 端 drillsgrid（93 行）

grid = (triton.cdiv(S, block_m), B * H)        # 二维 grid

start_m = tl.program_id(0)                     # 第几块 query
bh      = tl.program_id(1)                     # 第几个 (batch, head)
off_b, off_h = bh // H, bh % H
base = off_b.to(tl.int64) * stride_b + off_h.to(tl.int64) * stride_h

• grid 的第 0 维是 query 块，第 1 维把 batch 和 head 合并成一维。S = 4096、BLOCK_M = 64、B·H = 64 时一共 64 × 64 = 4096 个 program。
• 不同的 (batch, head) 之间完全独立；同一个 (batch, head) 里不同的 query 块也独立，它们读同一份 K、V，但各算各的输出。
• base 是这个 (batch, head) 在张量里的起点，用 int64 算（B·H·S·D 很容易超过 2³¹）。之后 q、k、v、o 的地址都是 base + 位置 × stride_s + d，因为 host 端已经 .contiguous()，四个张量的 stride 相同，最后一维 stride 是 1。

③ 循环外：q 只读一次，三个状态初始化

drills38–43 行

q   = tl.load(Q + base + offs_m[:, None] * stride_s + offs_d[None, :], mask=offs_m[:, None] < S, other=0.0)
m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)   # 每行到目前为止的最大分数
l_i = tl.zeros([BLOCK_M], dtype=tl.float32)                 # 每行的 Σ exp(score − m_i)
acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)       # 每行的 Σ exp(score − m_i) · v

q 是 [BLOCK_M, D] = 64 × 64 个 fp16，整个循环里一直留在寄存器里，每一块 K 都和它做一次 tl.dot。这就是 query 块内的复用。

④ 主循环：一块 K、V 进来时发生什么

drills52–76 行

读 K：直接读成转置的

drills55–57 行

k = tl.load(K + base + cols[None, :] * stride_s + offs_d[:, None], mask=cols[None, :] < S, other=0.0)
s = tl.dot(q, k) * sm_scale

K 在内存里是 [S, D]。这里把位置放在列（cols[None, :]）、维度放在行（offs_d[:, None]），读出来就是 [D, BLOCK_N]，也就是 Kᵀ 的一块。于是 tl.dot(q, k) 是普通的 [BLOCK_M, D] @ [D, BLOCK_N]，不需要额外转置。这是 T4 讲过的“用 stride 表达转置”。

mask

drills59–62 行：key 位置超出 S 的设成 −inf；causal 时还要求 query 位置 ≥ key 位置。−inf 经过 exp 变成 0，相当于这个 key 不存在。

online softmax 更新

drills67–76 行

m_new = tl.maximum(m_i, tl.max(s, axis=1))     # 每行的新最大值
p     = tl.exp(s - m_new[:, None])             # 以新最大值为基准的权重
alpha = tl.exp(m_i - m_new)                    # 旧状态的换算系数
l_i   = l_i * alpha + tl.sum(p, axis=1)
acc   = acc * alpha[:, None]
acc   = tl.dot(p.to(v.dtype), v, acc)          # acc += p @ v
m_i   = m_new

和 T2 的 online softmax 一样，只多了一个 acc：分母 l 和分子 acc 都是“以 m 为基准的加权和”，基准变了，两者都要乘 alpha。

手算一个例子：1 个 query，4 个 key，BLOCK_N = 2，D = 1。分数 s = [1, 2, 3, 0]，v = [10, 20, 30, 40]。

第一轮 m_i = −inf，alpha = exp(−inf − m_new) = 0，正好把初始的 0 清掉，不需要特判。

为什么不会出现 NaN：如果某一行在第一块里所有 key 都是 −inf，m_new 也会是 −inf，s − m_new 就是 (−inf) − (−inf) = NaN。注释（64–66 行）说明了为什么这里不会发生：每个合法的 query 行至少能看到 key 0（causal 时 query 位置 ≥ 0），而 key 0 在第一块里。L2 的 paged decode 里如果有整行被 mask 掉的情况，就要显式处理。

p 转成 fp16 再乘

tl.dot(p.to(v.dtype), v, acc)：p 是 fp32，Tensor Core 要求两个输入都是 fp16/bf16，所以先转换。p 的值都在 [0, 1] 之间，转成 fp16 的精度损失很小，累加仍在 fp32。这是所有 FlashAttention 实现的标准做法。

⑤ causal：跳过一半的块

drills47–50 行

hi = tl.minimum((start_m + 1) * BLOCK_M, S) if CAUSAL else S
for start_n in range(0, hi, BLOCK_N): ...

causal 时 query i 只能看 key 0..i。第 start_m 块 query 的最后一个位置是 (start_m + 1)·BLOCK_M − 1，所以 key 只需要循环到这里。S = 4 × 64，BLOCK_M = BLOCK_N = 64 时：

16 个块里只算 10 个；S 越大越接近一半。绿色块里所有位置都可见，其实不需要 tl.where；只有橙色的对角块需要逐元素 mask。题目页建议把循环拆成两段来省掉绿色块里的 mask。

CAUSAL 是 constexpr，所以 if CAUSAL: 在编译期就决定了，causal 和非 causal 是两个不同的 kernel，没有运行时分支。

⑥ 结束和 host 端

• drills78–80 行：o = acc / l_i，转回 fp16/bf16 写出。如果要支持 backward 或 L3 的分片合并，还要存 lse = m_i + log(l_i)。
• drillshost 端：sm_scale 默认 1/√D；HEAD_DIM 必须是 2 的幂（tl.arange 的要求），参考解只支持 16–128；固定 num_warps=4, num_stages=2。

⑦ 对照真实实现

• Triton 官方 tutorial 06：triton_attn_fwd_inner 是同一个循环，tritonm_ij、exp2、alpha 一一对应。它用了 exp2 技巧：tritonqk_scale *= 1.44269504（1/ln2），之后用 exp2 代替 exp，硬件的 ex2 指令更快。
• FlashAttention-2（CUDA / CuTe）：主循环在 facompute_attn_1rowblock，causal 跳块在 fan_block_max；online softmax 的重缩放在 fasoftmax_rescale_o，fascores_scale 就是这里的 alpha。
• FlashInfer 的 prefill：flashinferupdate_mdo_states（m、d、o 就是这里的 m、l、acc），kernel 入口 flashinferSinglePrefillWithKVCacheKernel。
• 论文：FlashAttention · FlashAttention-2 · FlashAttention-3

⑧ 常见坑与自测

• acc 忘了乘 alpha：只更新了 l，结果偏向前面的块。
• 用 m_i 而不是 m_new 算 p：max 变大时 exp 可能溢出。
• causal 的 hi 用了 BLOCK_N：(start_m + 1) * BLOCK_N 在 BLOCK_M ≠ BLOCK_N 时会多算或少算；正确的上界是 query 块的结束位置。
• K 按 [BLOCK_N, D] 读进来再 tl.trans：也对，但多一步；直接用 stride 读成转置更省事。

源码与资料索引

本仓库：drillssolutions/t5_flash_attn.py · drillsmine/t5_flash_attn.py · drillstests/test_t5_flash_attn.py · drillscontent/t5.md

实现：tritonTriton tutorial 06 · faFA2 forward · faFA2 online softmax · faFA3 主循环（Hopper） · faFA4（CuTe DSL） · flashinferFlashInfer update_mdo_states

论文：FlashAttention · FlashAttention-2 · Online normalizer calculation for softmax