动画图解 LLM 系统 · Attention kernel

FlashAttention 1、2、3

矩阵乘法很快，数据来回搬却很慢。
把 attention 留在 GPU 的小工作台上，算完再搬出去。

1 · 少搬2 · 分好工3 · 重叠执行

→ 推进每一步 · 拖动时间轴回看 · F 全屏

形状 · 先把维度写清楚

attention 的输出，究竟是哪张矩阵？

量	单个 batch / head 的形状	含义
Q、K	N × d	一行对应一个 token 的 query / key
V、O	N × dᵥ	value 与加权后的输出
S、P	N × N	分数与按行归一化的概率

S = QKᵀ / √d　　P = softmax_rows(S + mask)　　O = PV

以下先取 dᵥ = d、无 dropout、等长 self-attention。总工作量再乘 batch B 和 query head 数 H；GQA 让多个 Q head 共享 K、V。

问题 · HBM 往返

中间那张 N × N 大表，搬了好几遍

每个 query 和每个 key 打分：S = QKᵀ / √d。N 个 token，就得到一张 N × N 的分数表。

朴素实现先把 S 写进 HBM，再读回来做 softmax，得到概率 P，又写回 HBM。

最后读 P 和 V，算加权平均 O。算完的中间结果，马上又得搬回来。

N = 16,384 时，一张 FP16 的 N × N 矩阵就有 512 MiB / head。序列翻倍，这张表变成四倍。

算账 · 显存不只算容量

两次写、两次读，就有 2 GiB / head

N = 16,384，S 和 P 都按 FP16 存储。单张中间矩阵：N² × 2 = 536,870,912 B = 512 MiB。

动作	中间矩阵流量
QKᵀ 写 S	512 MiB
softmax 读 S、写 P	1,024 MiB
PV 读 P	512 MiB
合计	2,048 MiB = 2 GiB / head

这是朴素三步实现的中间张量流量；还没加 Q/K/V/O，也没计缓存命中。FA 消除完整 S/P 的 HBM 往返，但 K/V 仍可能被不同 Q tile 多次读取。

FlashAttention-1 · 分块

大仓库放原料，小工作台算一块

HBM 容量大；SM 里的 shared memory 和寄存器容量小、访问快。把 Q、K、V 切成能放上工作台的小块。

固定一块 Q，依次读入 K、V 块：算局部分数、更新 softmax 状态、累加输出。分数块用完就丢。

所有 K、V 块都走完，才得到这几行最终的 O。完整的 S 和 P 从来不落 HBM。

推导 · 为什么要追踪最大值

重新定标不会改变加权平均

O = Σ exp(sⱼ − m) vⱼ / Σ exp(sⱼ − m)

把基准从 m 换成更大的 m′，每一个旧权重都乘同一个 α = exp(m − m′)。分子、分母一起乘，比例就没变。

单独归一化不能直接相加

两块的局部输出 O₁、O₂ 各自总权重为 1；直接平均它们，会丢掉两块的相对概率质量。

保留还没除的分子

携带 (m, ℓ, o)，合并后只在结尾做 o / ℓ。FA2 还借此减少过程中反复归一化的非矩阵乘工作。

Online softmax · 一边读，一边合并

没见过后面的分数，怎么做 softmax？

每行只随身带三个状态：最大值 m、指数和 ℓ、加权分子 o。新块来了，先统一尺度，再相加。

m′ = max(m, max s_j)

α = exp(m − m′)

p_j = exp(s_j − m′)

s 已包含 1/√d 缩放与 mask。减去最大值，让指数保持在安全范围。

ℓ′ = αℓ + Σ p_j

o′ = αo + Σ p_jv_j

最后：O = o / ℓ

旧的分母和分子必须一起乘 α。o 是向量，每个输出维度都这样更新。

初始 m = −∞、ℓ = 0、o = 0；第一个含有效 key 的块把 α 视为 0。全被 mask 的块跳过。以上写法把最终除法留到最后。

Online softmax · 数字跟着走

最大值变大了，旧结果一起缩小

看一行、一个输出维度。第一块的分数是 [1, 2]，value 是 [2, 8]；m = 2，指数权重是 [e⁻¹, 1]。

第二块出现更大的分数 3。把基准从 2 改成 3：旧的 ℓ 和 o 都乘 e⁻¹，旧的相对贡献仍然保留。

新块的分数是 [3, 0]、value 是 [4, 10]，加上它的权重 [1, e⁻³]。最终 O ≈ 4.966，与一次性算完整 softmax 相同。

实现 · causal 与尾块

mask 要在取最大值之前生效

位置决定可见性

等长 causal attention 的第 i 行只允许 j ≤ i。无效位置先设为 −∞，然后做 max 和 exp。尾块还要同时检查 j < N。

整行或整块都无效

完全无效的块直接跳过，避免 −∞ − (−∞) 产生 NaN。通用 mask 若能遮住整行，要明确返回零等约定，并单独实现这条分支。

N = 5，tile = 4：第 4 行的第二块只有 key 4 有效

带 KV cache 的 decode 还要加 query 的绝对位置偏移；长度 1 的 Q 不等于“只能看 key 0”。

实现 · 一行的可运行逻辑

每轮只更新三个状态

m, l, o = -inf, 0, zeros(dv)
for Kb, Vb, valid in key_tiles:
    s = (q @ Kb.T) / sqrt(d)
    s = where(valid, s, -inf)
    if not any(valid): continue
    m2 = max(m, max(s))
    alpha = exp(m - m2) if l else 0
    p = exp(s - m2)
    o = alpha * o + p @ Vb
    l = alpha * l + sum(p)
    m = m2
if l == 0: return zeros(dv), -inf
return o / l, m + log(l)

这里 p 是未归一化指数权重；全遮蔽行约定输出 0。最后的 m + log(l) 是反向需要的 logsumexp。

反向 · 从 softmax 的导数开始

一行有很多概率，却只需一个归约量

∂Pᵢ / ∂Sⱼ = Pᵢ(δᵢⱼ − Pⱼ)

设传到概率上的梯度为 dP。把链式法则沿第 j 列相加，得到：

dSⱼ = Pⱼ × (dPⱼ − Σᵢ PᵢdPᵢ)

先把公共项算出来

D = Σ PᵢdPᵢ，每一行只要一个标量。dP = dO Vᵀ，所以 D = Σₖ dOₖ Oₖ。

一个方便的自检

Σⱼ dSⱼ = 0：给整行所有分数同时加一个常数，softmax 不变，沿这个方向的导数就应为零。

训练 · 反向重算

反向需要 P？到用它时再算

训练时，保留 Q、K、V、O，并额外存每行一个 L = m + log ℓ。L 把整行 softmax 的归一化信息压成一个数。

反向读一块 Q、K，重新算 S；有了 L，就能恢复 P = exp(S − L)，接着计算这一块的梯度。

临时 P 块用完就释放。多算一些矩阵乘，换掉巨大中间矩阵的存储与搬运。激活重算也在做这笔交换。

反向 · 连回 Q、K、V

缩放 1/√d 也要进入梯度

计算	矩阵形状	含义
dV = PᵀdO	N × dᵥ	按概率把输出梯度分给 values
dP = dOVᵀ	N × N	每个 value 对损失的贡献
dS = P ⊙ (dP − D)	N × N	通过 softmax；D 按行广播
dQ = dSK / √d	N × d	每个 query 的梯度
dK = dSᵀQ / √d	N × d	多个 query 对同一 key 的贡献相加

表里写全矩阵是为了读清形状；kernel 只计算 tile，dK/dV 需要把多个 query tile 的贡献正确归约。mask 位置的 P 与 dS 都为零。

演算 · 一个 query，两个 key

梯度沿着两条 value 路径流回去

取 d = 1，q = 1，k = [0, ln 3]，v = [2, 6]，上游 dO = 2。于是 P = [¼, ¾]，O = 5。

D = dO × O = 10；dS = [¼(4−10), ¾(12−10)] = [−1.5, 1.5]。两项相加为零。

资源 · tile 大小不是越大越好

寄存器、shared memory、任务数一起约束

考虑 Q tile 为 Bᵣ × d，K/V tile 为 B꜀ × d。用 2 字节输入、FP32 分数与输出累加，先列一份教学工作集：

Q/K/V：2d(Bᵣ + 2B꜀) B
局部分数：4BᵣB꜀ B　　输出累加：4Bᵣd B

真实 kernel 会复用缓冲区、用寄存器 fragments 持有中间量；这张表用来识别增长方向，不能直接当成某版本的 shared-memory 分配。

更大的 Bᵣ 复用更多 K/V，却增加输出寄存器与分数 tile；更大的 B꜀ 减少循环次数，却占更多暂存空间。资源超限会降低驻留 block 数，甚至发生 register spill。

演算 · tile 资源预算

Q tile 翻倍，片上账本怎样变化？

教学配置：N = 4,096，d = 128，B꜀ = 64；Q/K/V 为 16 位，S/O 累加为 32 位。图中单位为 KiB。

Bᵣ 从 32 → 64 → 128，Q tile 数从 128 → 64 → 32。总工作集从 64 → 96 → 160 KiB；缓冲区更大，独立任务也更少。

并行 · 小 batch 的诊断

长序列也可能没有足够多的 block

B = 1，H = 8，N = 4,096，Bᵣ = 128
按 B×H 派发：8 个任务
再切 Q：1 × 8 × 32 = 256 个任务

若 GPU 有 132 个 SM，8 个任务最多让 8 个 SM 同时处理这种任务。256 个 Q block 才有机会覆盖所有 SM。

任务够多只是第一步

一个 block 的寄存器与 shared memory 仍会影响驻留数；因果 mask 还使前后 Q block 的工作量不均匀。

看到 SM 空闲时

先看 grid 与 tile，再看资源占用和工作分配。提高 Tensor Core 理论峰值，无法补上缺失的并行任务。

FlashAttention-2 · SM 之间

把长序列切成更多独立任务

GPU 有很多 SM。只按 batch 和 head 派任务，小 batch、少 head 时，长序列也可能只让少数 SM 忙起来。

不同 query 行的最终输出彼此独立。FA2 把 query block 也放进 grid，每块自己遍历 K、V。

一条长序列变成更多任务，交给更多 SM 同时处理。片上仍然分块，显存里仍然不放 N × N 的中间矩阵。

FlashAttention-2 · 一个 block 内

按 Q 分工：每个 warp 管自己的输出行

FA1 的典型 warp 划分把 K、V 切给不同 warp。它们贡献到同一批输出行，最后还得经 shared memory 合并。

FA2 的 split-Q：把 Q 的行切给不同 warp，让它们读取共同的 K、V tile。每个 warp 都能完成自己负责的输出行。

少一次跨 warp 的部分输出合并，就少一次 shared-memory 往返与同步。再减少重复缩放、归一化等非矩阵乘操作，让 Tensor Core 少等。

FlashAttention-3 · Hopper 异步流水线

下一块在搬，这一块在算

Hopper 上，TMA 负责异步搬数据，WGMMA 发起异步矩阵乘。发出指令后，warp 不必一直守着等待。

producer warpgroup 提前搬下一块 K、V；consumer warpgroup 计算当前块。缓冲区与 barrier 保证读到完整数据、不会提前覆盖。

consumer 还把部分 softmax 与后续矩阵乘交错安排。搬运、Tensor Core、普通计算单元更少轮流空等。

流水线 · 先算理想上限

重叠以后，加法变成较大值

教学时序：每块搬运 2 μs，计算 5 μs，共 4 块；计算阶段已包含该示例里的 softmax，不再单独重叠它。

串行：4 × (2 + 5) = 28 μs
理想双缓冲：2 + 4 × 5 = 22 μs
这部分的加速比：28 / 22 ≈ 1.27×

开始和结束仍要付费

第一块必须先搬好；最后一块算完才能结束。短序列容易被填充、排空与 barrier 开销主导。

双缓冲有资源成本

同时保留当前与下一块，可能挤占 shared memory、降低 occupancy。用 profiler 确认拷贝与计算真的重叠。

FlashAttention-3 · 精度与硬件

FP8 是一条额外的加速路径

异步流水线对 FP16 / BF16 也有用。想继续提高矩阵乘吞吐，可以再选 FP8。

先解决误差

每块设置缩放因子，减少少数大值对整块精度的挤压。量化 Q、K、V 时，尺度也要带进计算。

再分散异常值

对 Q、K 使用匹配的正交变换，把突出的通道打散，同时保持 QKᵀ。论文称为 incoherent processing。

保留关键精度

softmax 的统计量与缩放仍使用 FP32。官方 Hopper 实现提供 FP16 / BF16 前后向与 FP8 前向路径。

FA3 围绕 H100 / H800 等 Hopper GPU 设计。FP8 会引入额外量化误差；是否启用取决于模型精度要求、硬件与具体 kernel 支持。

三代演进

同一个 attention，三层效率问题

版本	最关心什么	关键做法
FA1	数据搬得太多	分块 + online softmax + 反向重算
FA2	任务分得不够好	更多 query block 并行 + split-Q
FA3	不同单元轮流等待	TMA / WGMMA + 异步流水线

密集 attention 仍需计算 O(N²d) 的匹配；省掉的是 N × N 中间矩阵的落盘。FP16 / BF16 路径保持同一数学结果，浮点运算顺序带来正常舍入差异。

这里的大片矩阵主要对应训练与 prefill。decode 通常只有一行新 query，瓶颈更多来自读 KV cache，需要结合 KV cache 管理与专门的 decode kernel。

使用 · 从 PyTorch 调用融合 attention

先确认形状、mask 和真正选中的 backend

import torch.nn.functional as F
from torch.nn.attention import sdpa_kernel, SDPBackend
# q/k/v: [batch, heads, sequence, head_dim]，CUDA 张量
with sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION]):
    out = F.scaled_dot_product_attention(
        q, k, v, dropout_p=0.0, is_causal=True)

显式限制 backend 有助于发现不支持的 dtype、head dimension 或 mask；正常使用时也可让 PyTorch 自动选择。SDPA 名称本身不能证明选中了哪一代 FA。

推理时把 dropout_p 明确设为 0；布尔 attention mask 中 True 表示允许参与。非方阵 causal 对齐和 GQA 的支持要按所用接口检查。

测量 · 一次 forward 的 GPU 时间

先热身，再用设备事件围住计算

for _ in range(20):
    out = run_attention(q, k, v)
torch.cuda.synchronize()
a = torch.cuda.Event(enable_timing=True)
b = torch.cuda.Event(enable_timing=True)
a.record()
for _ in range(100):
    out = run_attention(q, k, v)
b.record(); b.synchronize()
print(a.elapsed_time(b) / 100, "ms / forward")

run_attention 是待测调用；输入分配、随机初始化和编译不放进计时区。多轮报告中位数；固定 B/H/N/d、dtype、causal、dropout 与训练/推理模式。

forward 与 forward+backward 分开测；另测峰值显存。GPU 计时示例需 CUDA；CPU 实验用于验证数学结果。

排查 · 先定位差异长什么样

错在边界、精度，还是执行路径？

观察	优先检查	下一步
只在 causal / 尾块错误	位置偏移、mask、padding	用 N=5、tile=4 对照逐项概率
换 tile 后误差巨大	m 改变时只缩放了分子	检查 α 同时乘 ℓ 与 o
正确但小 N 更慢	launch、调度、工作量太小	扩大 N 并比较多种形状
显存低了，时间几乎没变	瓶颈可能移到算力或投影	分开测 attention 与整层
强制 backend 后报错	形状 / dtype / mask 不支持	读诊断信息，避免误认 fallback

用最大绝对误差与相对误差一起看；接近零的参考值不适合只看相对误差。先用 FP64 CPU oracle 排逻辑，再评估低精度的舍入。

练习 1 · 独立分析

两块 softmax，为什么不能平均局部输出？

块 A 只有 s = 0、v = 10；块 B 只有 s = ln 3、v = 2。计算正确 O，再算直接平均两个局部输出的错误答案。

先独立计算，再展开解答。

练习 2 · 独立分析

从概率算回分数梯度

已知 P = [0.2, 0.3, 0.5]，V = [1, 3, 5]，dO = 2。求 O、D、dS 与 dV，并给出一个自检。

先独立计算，再展开解答。

练习 3 · 独立分析

为什么吞吐提高，整层却只快了一点？

教学测量：整层 10 ms，其中 attention 4 ms。更换 kernel 后 attention 快了 2 倍，其余耗时不变。整层加速多少？还应该测什么？

先独立计算，再展开解答。

动手 · CPU 参考实现

让密集实现与分块实现互相对答案

python3 lab.py --length 7 --block 3
# 再试：--block 1、--block 8、--length 5

实验逐行合并 online softmax，再利用 L 重建概率并计算 dQ/dK/dV。完整密集 softmax 作前向对照，中心差分作独立梯度对照。

默认：前向最大误差 < 1e−12
Q/K/V 梯度差分误差 < 2e−8；最终打印 PASS

修改题：删掉旧分子的 α、漏掉 1/√d、把 causal 的 ≤ 改成 <，分别预测哪项检查会失败。

下载 lab.py · NumPy / CPU

参考与延伸阅读

论文与实现

Dao et al.（2022），FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness：分块、online softmax、IO 分析与反向重算。
Dao（2023），FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning：query block 并行、split-Q 与减少非矩阵乘工作。
Shah et al.（2024），FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision：Hopper 异步流水线、FP8 分块量化与 incoherent processing。
Dao-AILab / flash-attention：官方实现；FA3 的硬件与 dtype 支持见 README 和 hopper/。
PyTorch 文档：scaled_dot_product_attention：backend 选择、mask 与 dropout 语义。

动画中的块数、warp 数与时间长度用于解释数据依赖，不代表某个 kernel 的固定配置或实测耗时。

继续：RMSNorm 与 GPU 算子优化 · 返回博客