FlashAttention 动效 02

分块 + 在线 softmax:不物化 N×N,也能算出精确结果

小例子:4 个 query、6 个 key(分成 3 块,每块 2 个)。每加载一个 K/V 块,就更新一次行状态 m(运行最大值)、ℓ(指数和)和未归一化输出 Õ。点「下一步」逐块推进。

Q · 4×2(每格是一行)

S_j = Q K_jᵀ · 4×2

K / V · 列 = 第 1/2/3 块,每格是一行

K
V

行状态(每个 query 一行)

rowmℓÕ₀Õ₁
未加载的块 已处理 当前块