GPU 与 TPU：让硬件不再是黑魔法

0. 本讲导读

前面几讲讲的是"模型该长什么样"：Transformer 的结构、注意力的变体、混合专家的路由。那些讨论里有一个被默认接受的前提——只要 FLOPs 数一样，跑起来就一样快。这个前提是错的，而且错得很离谱。同样是 $2N^3$ 次浮点运算的方阵乘法，在同一块 A100 上，$N=1792$ 能跑到接近峰值，$N=1793$ 只有一半；同样是把一个张量过一遍，写成 x.relu() 和写进一个融合 kernel 里差 5 倍；同样是标准注意力，materialize 出 $N \times N$ 的分数矩阵和不 materialize，差的是能不能训 32K 上下文。

这一讲的目标，用 Tatsu 自己的话说，是 "make CUDA and GPUs less magic"。具体拆成三部分：

• Part 1 — GPU 的解剖学：它由什么组成，为什么长成这样，存储层次是怎么排的，执行模型（thread / warp / block）和硬件（SP / SM）怎么对应。顺带看一眼 TPU，理解"加速器"这个类别的共性与分歧。
• Part 2 — 性能模型：引入 算术强度（arithmetic intensity） 和 roofline 模型，这是全课最重要的分析工具。然后是让 GPU 跑快的六件套：控制分歧、低精度、算子融合、重计算、合并访存、分块。
• Part 3 — 综合应用：用前两部分的全部工具，把 FlashAttention 从头拆开。你会发现它没有任何新数学，就是 tiling + fusion + recomputation 三板斧，加上一个 2018 年就有的在线 softmax 技巧。

这一讲之后紧接着是 Triton 与自定义算子——那一讲教你"怎么写"，这一讲教你"为什么这么写"。没有这一讲的模型，写 kernel 就是在瞎调参数。

1. 算力从哪来：Dennard scaling 的终结与并行 scaling

为什么先讲硬件

Kaplan 等人的 scaling law 告诉我们，语言模型的损失随算力呈幂律下降，而且这个关系惊人地平滑可预测。这件事有一个直接的工程推论：只要你能持续拿到更多的有效算力——通过更快的硬件、更高的利用率、更好的并行策略——你就能持续拿到更好的模型，甚至不需要任何算法创新。Tatsu 在课上补了一句限定："for now"。但至少到目前为止，没有 GPU 的 scaling 就没有 LLM 的 scaling。

所以问题变成：算力本身是怎么 scale 的？

Dennard scaling：免费的午餐及其结束

1974 年 Robert Dennard 观察到：当晶体管的线宽按比例缩小 $\kappa$ 倍时，如果同时把电压和电流也按 $\kappa$ 缩小，那么单位面积的功耗密度保持不变，而开关速度提升 $\kappa$ 倍。这就是 Dennard scaling。它和 Moore 定律（晶体管密度每 18–24 个月翻倍）合起来，意味着 1980–2005 年间：

• 晶体管数量指数增长（Moore）；
• 时钟频率也跟着指数增长，而且不用多花电（Dennard）；
• 单线程性能自动变快，程序员什么都不用做。

2005 年前后这条路断了。原因是亚阈值漏电流（leakage current）：晶体管尺寸小到一定程度后，即使"关断"也会漏电，而且漏电随电压降低反而变严重，所以电压不能再往下降了（所谓 voltage scaling wall）。电压降不下去，功耗密度就随密度增长，芯片会烧掉。于是频率被钉死在 3–5 GHz 这个量级，二十年没怎么动过。

并行 scaling 还在继续

Bill Dally 在 HotChips 的 keynote 里给出过一个著名的统计：NVIDIA GPU 的单芯片推理性能在十年里涨了超过 1000 倍，远远快过同期工艺进步能解释的幅度。他把这 1000× 大致拆成几个来源：数值表示的降低（FP32 → INT8/FP8 这一路，约 16×）、专用复杂指令（DP4A、HMMA 这类张量核指令，约 12.5×）、制程进步（约 2.5×）、以及结构化稀疏（约 2×）。

这个拆解本身就是这一讲的提纲：今天的算力增长几乎全部来自"专门化"而不是"通用变快"。专门化意味着——如果你的工作负载正好落在硬件专门优化的那条路径上（稠密矩阵乘、低精度、规整形状），你能拿到这 1000×；否则你只能拿到那 2.5× 的制程红利。

2. CPU 与 GPU：延迟优化 vs 吞吐优化

CPU 和 GPU 都由晶体管构成，工艺也差不多，差别全在晶体管花在哪。

两种完全不同的目标函数

一个容易被忽略的关键点：GPU 的寄存器文件比 cache 还大

A100 每个 SM 有 65536 个 32-bit 寄存器 = 256 KB 寄存器，而每个 SM 的 L1/shared memory 加起来是 192 KB。整块 A100 的寄存器总量约 27 MB，H100 约 32 MB——比 L2 cache 还大。

这个设计不是偶然的。CPU 线程切换要保存/恢复寄存器上下文，代价是几十到几百个周期；GPU 把所有驻留线程的寄存器全部物理地存在寄存器文件里，切换 warp 只是换一个索引，零开销。这就是 Tatsu 说的 "threads are lightweight and can be stopped and started"。代价是：每个线程能用的寄存器数量有限（最多 255 个），一旦超了就会 spill 到 local memory（其实位于 HBM），性能断崖式下跌。

SIMT：GPU 的编程模型

GPU 采用 SIMT（Single Instruction, Multiple Threads，单指令多线程） 模型：一条指令被广播给一组线程，每个线程用不同的数据执行它。这和 CPU 的 SIMD（AVX-512 之类）不同——SIMD 是"一条指令操作一个宽向量寄存器"，程序员要显式写向量代码；SIMT 是"一条指令驱动 32 个独立线程"，程序员写的是标量代码，每个线程看起来像个普通的顺序程序。这是 CUDA 编程模型好用的根本原因，也是它偶尔坑人的根本原因（见后面的控制分歧）。

3. GPU 解剖：执行单元与存储金字塔

执行单元的三层结构

从大到小：

1. GPU 芯片：包含数十到上百个 SM。A100 有 108 个可用 SM（物理 128 个，屏蔽一部分保良率），H100 SXM 有 132 个。
2. SM（streaming multiprocessor，流式多处理器）：独立调度和执行一个 block（CUDA 里叫 CTA，cooperative thread array）。它有自己的 warp scheduler、寄存器文件、L1 cache / shared memory。SM 之间不共享除 L2 和 HBM 之外的任何东西。
3. SP（streaming processor）/ CUDA core：SM 内部的标量 ALU，执行一个线程的一条浮点或整数指令。A100 每 SM 有 64 个 FP32 core。
4. Tensor Core：专用的矩阵乘累加电路，一条指令完成一个小矩阵块的 $D = A \times B + C$。A100 每 SM 有 4 个（每 processing block 一个），全芯片 432 个；H100 有 528 个。

Tensor Core：为什么矩阵乘是"特权操作"

一个 Tensor Core 指令（如 HMMA.16816）在一条指令里完成 $16 \times 8 \times 16$ 的矩阵乘累加，即 $16 \times 8 \times 16 \times 2 = 4096$ FLOPs。相比之下一条 FFMA 指令只有 2 FLOPs。这就是 Tatsu 说的 "matmuls are >10x faster than other floating point ops" 的硬件来源。

历史上这是个反转。早期 NVIDIA GPU 只有可编程着色器（programmable shader），做的是图形管线里的顶点/像素变换。研究者发现这些着色器本质上在做 $4\times4$ 矩阵和向量的乘法，于是把神经网络"伪装"成图形渲染任务，硬塞给 GPU 跑——这就是 GPGPU 的起点，也是 CUDA（2007）出现的动机。二十年后，硬件反过来专门为矩阵乘设计了电路。

存储金字塔

把这些数字放在一起看，跨度是惊人的：从寄存器到 HBM，带宽差约 100 倍，延迟差约 300 倍；从 HBM 到主机内存，带宽再差 100 倍以上。 FlashAttention 论文里用的那张金字塔图给出的具体数字是：SRAM 20 MB @ 19 TB/s，HBM 40 GB @ 1.5 TB/s，CPU DRAM > 1 TB @ 12.8 GB/s。

4. 执行模型与内存模型：软件抽象如何映射到硬件

三个角色：thread / warp / block

这里有两条极其重要的推论，很多性能问题都源于没想清楚它们：

推论一：block 之间要通信，只能走 global memory（HBM）。 因为 block 可能落在不同的 SM 上，甚至可能不同时执行（block 数远多于 SM 数时会排队）。所以任何需要"全局归约"的操作（比如 softmax 的分母、LayerNorm 的均值）如果跨 block，就必须写回 HBM 再读一次——这正是 naive softmax 慢的根源，也正是 FlashAttention 要把一整行的计算塞进一个 block 的原因。

推论二：warp 内的 32 个线程共享一条指令流。 它们的访存也是一起发出的，硬件会尝试把这 32 个地址合并成尽可能少的内存事务（见第 9 节的合并访存）。同样，它们的分支也必须"一起走"（见第 7 节的控制分歧）。

内存模型：谁能看见什么

GPU 模型的优点

• 可扩展性：因为 block 之间完全独立，同一份代码在 108 个 SM 和 132 个 SM 的芯片上都能跑，硬件自动填满。加 SM 就能加性能，不需要改代码。
• 可编程性（相对而言）：SIMT 让你写标量代码，编译器和硬件负责向量化。比手写 AVX intrinsics 友好得多。
• 轻量线程：寄存器物理常驻，warp 切换零开销，因此可以用超额订阅（oversubscription）来掩盖延迟——同时驻留几千个线程，总有一些是 ready 的。

5. TPU：脉动阵列与另一条技术路线

Tatsu 把 TPU 作为一个"side thread"插在这里，理由很好：看过第二种加速器，你才知道 GPU 的哪些设计是本质的、哪些只是历史包袱。

逐项对照

脉动阵列（systolic array）：MXU 的工作原理

TPU 的 MXU 是一个 $128 \times 128$ 的 脉动阵列。这个名字来自"心脏搏动"——数据像血液一样有节奏地在处理单元（PE）阵列中流动。

工作方式（以 weight-stationary 为例）：把权重矩阵 $W \in \R^{128 \times 128}$ 预先加载到 $128 \times 128$ 个 PE 里，每个 PE 存一个权重 $w_{ij}$。然后把激活矩阵按对角线错峰地从左边推入。每个时钟周期，每个 PE 做一件事：

于是数据从左往右流，部分和从上往下累加，$128$ 个周期后底部就吐出完整的结果列。整个过程中，每个操作数只从存储器读一次，之后就在 PE 之间"手手相传"。

关键差异：TPU 没有 warp

这是 Tatsu 特别点出的一条。GPU 的 warp 机制带来的是动态调度：硬件在运行时挑选 ready 的 warp，用超额订阅掩盖不可预测的访存延迟。TPU 走的是另一条路——VLIW + 编译器静态调度：没有 warp，没有运行时的线程切换，所有的指令发射时机、数据搬运时机都由 XLA 编译器在编译期算好。

Tatsu 的原话是"no warps (just blocks) — tradeoffs in matmul vs non-matmul"：TPU 把用于动态调度的晶体管全部换成了矩阵乘单元和片上 SRAM（VMEM 有 128 MB！），代价是非矩阵乘操作相对更弱，以及对不规则计算的容忍度更低。

互连拓扑：Pod

另一个真正的差异是芯片之间怎么连（详细内容留给后面的并行化讲）。简单说：

• GPU：节点内 8 张卡用 NVLink / NVSwitch 全互连（H100 上每卡 900 GB/s 双向），节点之间走 InfiniBand（400 Gb/s 量级）。带宽在节点边界处有一个陡降，这是所有分布式训练策略要绕开的核心约束。
• TPU：用 ICI（Inter-Chip Interconnect） 把芯片连成 2D 或 3D 环面（torus）。TPU v4 是 $4096$ 芯片的 3D torus（$16\times16\times16$），v5p 规模更大。环面拓扑的好处是没有"节点边界"——每个芯片只和 6 个邻居直连，带宽均匀，特别适合 all-reduce 这种邻居通信模式。TPU v4 还引入了光路交换（OCS），可以在物理上重新配置 torus 的切片形状。

6. Roofline 模型与算术强度：这一讲最重要的工具

现在进入 Part 2。先看一个 Tatsu 用来"制造困惑"的现象：

算术强度的定义

对任意一段计算，定义它的算术强度（arithmetic intensity，也叫 operational intensity / compute intensity）：

分母是关键：算的是与慢速存储之间的流量，不算片上的 shared memory / 寄存器访问。这个量描述了"每搬一个字节，能榨出多少计算"。

Roofline 模型

数学上极简单。设峰值算力 $P$（FLOP/s），带宽 $B$（byte/s），一段计算有 $F$ FLOPs 和 $Q$ 字节访存，则：

（取 max 是因为计算和访存可以重叠，理想情况下慢的那个决定总时间。）可达吞吐为

两个约束相等处就是 ridge point（脊点、拐点）：

$I < I^\star$ 时 memory-bound，$I > I^\star$ 时 compute-bound。ridge point 是一块硬件的"性格"：

算一遍：矩阵乘

$C = AB$，$A \in \R^{M\times K}$，$B \in \R^{K \times N}$，元素 $b$ 字节。

• FLOPs：$F = 2MKN$（每个输出元素做 $K$ 次乘加 = $2K$ FLOPs）
• 最少访存：$Q = b(MK + KN + MN)$（每个输入读一次、输出写一次；这是理论下界，前提是能全部装进片上内存）

情形 A：方阵 $M=K=N$。

强度线性增长于 $N$。这就是"大矩阵才快"的第一层原因。在 H100 BF16（$b=2$，$I^\star=295$）上，要达到 compute-bound 需要

也就是说 $N$ 小于约 900 的方阵乘法，在 H100 上根本不可能跑满，无论 kernel 写得多好。A100 BF16 上这个门槛是 $N \ge 602$。回看那张散点图，曲线在 512–1024 之间开始"起飞"，正好对上。

情形 B：训练时的典型形状。 一个线性层，batch × seq 展平成 $B = 8 \times 4096 = 32768$ 个 token，隐藏维 $d = 4096$，即 $(B, d) \times (d, d)$：

远超 295——妥妥的 compute-bound。这是训练能跑到 40–50% MFU 的根本原因。

情形 C：自回归解码（batch = 1）。 同一个线性层，但 $B = 1$，变成矩阵–向量乘（GEMV）：

只有 1！离 295 差了 295 倍。解码时你的 H100 实际只能发挥出约 0.3% 的算力，时间 100% 花在把权重从 HBM 搬进来上。

算一遍：element-wise 算子

ReLU：$x \leftarrow \max(0, x)$，长度 $n$ 的向量。

也就是说，一个 ReLU 只能发挥出 H100 算力的 0.08%。具体感受一下：对 $n = 10^9$ 个 BF16 元素做 ReLU，搬运 4 GB，在 3.35 TB/s 下需要 1.2 ms；而 $10^9$ FLOPs 在 989 TFLOP/s 下只需要 0.001 ms。计算单元 99.9% 的时间在发呆。

算一遍：为什么 softmax 和 LayerNorm 跑不满

这是本讲反复出现的主题，值得算细一点。以 naive（未融合）实现的 LayerNorm 为例，输入 $(B, T, d)$ 共 $n$ 个 BF16 元素：

mean = x.mean(-1, keepdim=True)      # 第 1 趟：读 x
var  = (x - mean).pow(2).mean(-1)    # 第 2 趟：读 x（还可能写临时张量）
y    = (x - mean) / (var + eps).sqrt()   # 第 3 趟：读 x，写 y
y    = y * gamma + beta                  # 第 4 趟：读 y，写 y

如果每一行都作为独立的 PyTorch 算子执行（eager 模式就是这样），$x$ 被反复从 HBM 读取，中间结果被反复写回 HBM。粗略数：约 4 次读 + 3 次写 ≈ $7 \times 2n = 14n$ 字节；FLOPs 约 $10n$。

取 $n = 8 \times 4096 \times 4096 = 1.34\times 10^8$ 个元素：访存 $1.9$ GB，在 H100 上需 $0.56$ ms；计算 $1.34\times10^9$ FLOPs，只需 $0.0014$ ms。400 倍的差距。

softmax 同理，naive 实现要走 4 趟（求 max、求 exp 与和、相除；每趟都是全量读写），强度同样在 $0.3$–$1$ 量级。

一个可以直接用的 roofline 计算器

def roofline(flops, bytes_moved, peak_flops=989e12, bandwidth=3.35e12, name=""):
    """给定一个算子的 FLOPs 与 HBM 字节数，判断它被什么卡住。
    默认参数是 H100 SXM 的 BF16 峰值与 HBM3 带宽。"""
    intensity = flops / bytes_moved            # FLOP/byte
    ridge     = peak_flops / bandwidth         # 该硬件的拐点
    t_compute = flops / peak_flops
    t_memory  = bytes_moved / bandwidth
    t         = max(t_compute, t_memory)       # 理想重叠下的下界
    bound     = "compute-bound" if intensity > ridge else "memory-bound"
    print(f"{name:24s} I={intensity:9.2f}  ridge={ridge:6.1f}  {bound}")
    print(f"{'':24s} 下界耗时 {t*1e3:8.4f} ms   "
          f"可达算力 {flops/t/1e12:7.1f} TFLOP/s "
          f"({100*flops/t/peak_flops:5.2f}% of peak)")
    return t

b = 2                       # BF16
N = 4096
roofline(2*N**3, b*3*N**2,          name="matmul 4096^3")
roofline(2*N**3//32, b*3*(N//16)**2, name="matmul 256^3")
n = 8*4096*4096
roofline(n, 2*b*n,                  name="ReLU")
roofline(10*n, 7*b*n,               name="LayerNorm (naive)")
d = 4096
roofline(2*d*d, b*d*d,              name="GEMV (decode)")

把这十几行代码放进你的工具箱。写任何 kernel 之前先跑一遍它，你会立刻知道自己的上限在哪、该往哪个方向优化。

提速的六件套

Tatsu 把手段列成六条，我们按"它们在 roofline 图上做了什么"重新归类：

7. 提速手段（一）：控制分歧与低精度

7.1 控制分歧：唯一一个和内存无关的坑

SIMT 的代价：一个 warp 里的 32 个线程共享一个程序计数器（Volta 之前是严格如此；Volta 引入 independent thread scheduling 后每线程有自己的 PC，但执行仍然按分支分组串行）。所以遇到条件分支时，硬件只能：

1. 算出每个线程的条件结果，得到一个 32 位的活动掩码（active mask）；
2. 执行 if 分支，掩码外的线程被禁用但仍占着执行槽；
3. 执行 else 分支，掩码反过来；
4. 汇合。

最坏情况是 32 路分歧（比如 switch (threadIdx.x % 32)），此时 warp 的有效吞吐降到 $1/32$。

顺带一提：循环边界的分歧也算。for (i = 0; i < n[tid]; i++) 里如果每个线程的 n[tid] 不同，整个 warp 会跑到最长的那个，短的线程一直空转。变长序列不做 padding 直接算，就会踩这个坑。

7.2 低精度：位数少了，要搬的字节就少了

低精度同时在 roofline 的两个方向上起作用，必须分开看：

课上的 ReLU 例子就是最简版本：FP32 是 8 byte/FLOP，FP16 是 4 byte/FLOP，强度翻倍，时间减半。注意这个加速和"计算变快"完全无关——ReLU 的 FLOPs 一个都没少，纯粹是因为搬的字节少了一半。

混合精度：哪些操作可以降，哪些不能

FP8 的前沿：块级缩放（MXFP8）

MXFP8（Blackwell 原生支持）有三个值得注意的设计：

1. 用 E4M3 而不是 E5M2。既然有了细粒度的缩放因子来处理动态范围，就不需要在数据格式里浪费指数位了，把位数留给尾数换精度。
2. 缩放因子本身是 FP8（E8M0），每 32 个元素一个。E8M0 是纯指数格式，意味着"缩放"就是指数相加，硬件实现极便宜。存储开销只有 $8/(32\times8) = 3.1\%$。
3. 转置变成了非平凡操作。这是最反直觉的一点：缩放因子是沿着某一个轴分组的，转置之后分组方向就错了，必须重新量化。所以实际训练流程里，同一个张量要同时保存 rowwise 和 columnwise 两个量化版本（前向用一个，反向用另一个），或者在需要时重新量化。

8. 提速手段（二）：算子融合与重计算

8.1 工厂与仓库：融合的直觉

把这个比喻量化。假设有 $k$ 个连续的 element-wise 算子作用在 $n$ 个元素上（每个 $b$ 字节）：

注意加速比是 $k$，不是 $k$ 的某个分数——是完整的 $k$ 倍。因为这些算子 100% memory-bound，时间正比于访存量。一个 Transformer block 里，从 LayerNorm 到 GeLU 到残差加法到 dropout，串起来轻松有十几个 pointwise 算子。

一个具体例子：$\sin^2 x + \cos^2 x$

这行 PyTorch 代码在 eager 模式下会启动 5 个 CUDA kernel：sin、pow、cos、pow、add。每个 kernel 都要把整个张量从 HBM 读进来、算一下、写回去。

import torch

x = torch.randn(1 << 26, device="cuda", dtype=torch.bfloat16)

def f(x):
    return torch.sin(x) ** 2 + torch.cos(x) ** 2

# eager：5 次 kernel launch，约 5 次读 + 5 次写 = 10 轮 HBM 往返
y1 = f(x)

# 编译后：TorchInductor 生成 1 个 Triton kernel，1 读 1 写
f_fused = torch.compile(f)
y2 = f_fused(x)          # 首次调用会触发编译

# 访存量估算（BF16，n = 2^26 个元素 = 134 MB）
n, b = x.numel(), x.element_size()
print(f"未融合 ≈ {10 * n * b / 1e9:.2f} GB,  融合后 ≈ {2 * n * b / 1e9:.2f} GB")

8.2 重计算：用 FLOPs 换字节

反向传播的标准做法是：前向时把每一层的激活值存下来，反向时读出来算 Jacobian。这套做法在"算力贵、内存便宜"的年代是对的，今天正好反过来。

两种尺度的重计算

from torch.utils.checkpoint import checkpoint

class Block(torch.nn.Module):
    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x

# 普通：block 内部所有中间激活都留在显存里
h = block(x)

# 检查点：只保留 block 的输入 x，反向时重跑一遍 forward 生成中间激活
h = checkpoint(block, x, use_reentrant=False)

激活检查点的显存账：不用检查点时，$L$ 层每层要存 $O(\text{层内激活})$；用了之后只存 $L$ 个 block 的输入，显存从 $O(L \cdot k)$ 降到 $O(L + k)$（$k$ 是每 block 内部的张量数）。经典的 $O(\sqrt{L})$ 策略（Chen et al. 2016）是每隔 $\sqrt{L}$ 层设一个检查点，显存 $O(\sqrt{L})$、额外计算一次前向。

9. 提速手段（三）：合并访存与 DRAM 的物理特性

DRAM 为什么按"突发"读

什么叫"合并"

矩阵乘里的合并问题

这是初学者写 CUDA 时最常见的错误。看下面两个索引方式，数学上完全等价，性能差 8 倍：

// 错误：相邻线程访问相隔 N 的地址（每个线程负责一整行）
int row = blockIdx.x * blockDim.x + threadIdx.x;
for (int k = 0; k < N; ++k)
    acc += A[row * N + k];        // warp 内 32 个线程地址跨度 = 32*N*4 字节

// 正确：相邻线程访问相邻地址（threadIdx.x 映射到最快变化的维度）
int col = blockIdx.x * blockDim.x + threadIdx.x;
for (int k = 0; k < N; ++k)
    acc += A[k * N + col];        // warp 内 32 个线程地址连续 = 128 字节

10. 提速手段（四）：分块（tiling）与矩阵尺寸之谜

Tatsu 称 tiling 为 "the big one"——它是所有手段里效果最大、也最需要人工设计的一个。

问题：朴素矩阵乘把每个输入读了 $N$ 次

算一下朴素版的算术强度：$F = 2N^3$，$Q = 2N^3 b$（每次乘加读两个元素），

和 element-wise 算子一个量级！ 矩阵乘"天生 compute-bound"的性质，只有在你把数据复用做出来之后才成立。这一点非常重要：$I = 2N/(3b)$ 是理论上限，$I = 1/b$ 是朴素实现的实际值，中间差了 $N$ 倍，全靠 tiling 填补。

解法：把数据搬进 shared memory，反复用

最小可读实现

#define T 32                                 // tile 边长，也是 blockDim

__global__ void matmul_tiled(const float* A, const float* B, float* C, int N) {
    __shared__ float As[T][T];               // 每个 block 私有的 shared memory
    __shared__ float Bs[T][T];

    int tx = threadIdx.x, ty = threadIdx.y;
    int row = blockIdx.y * T + ty;           // 本线程负责的输出元素
    int col = blockIdx.x * T + tx;
    float acc = 0.f;                         // 累加器留在寄存器里

    for (int t = 0; t < N / T; ++t) {        // 沿 K 方向逐 tile 推进
        // ---- 阶段 1：协作载入。注意 tx 是最快变化的维度 ----
        As[ty][tx] = A[row * N + (t * T + tx)];   // warp 内地址连续 -> 合并
        Bs[ty][tx] = B[(t * T + ty) * N + col];   // warp 内地址连续 -> 合并
        __syncthreads();                          // 等所有线程载完

        // ---- 阶段 2：在 shared memory 里算 T 次乘加，不碰 HBM ----
        for (int k = 0; k < T; ++k)
            acc += As[ty][k] * Bs[k][tx];
        __syncthreads();                          // 等所有线程算完再覆盖 tile
    }
    C[row * N + col] = acc;                       // 只写一次
}

三个细节值得注意：

• 两个 __syncthreads() 缺一不可。 第一个保证"数据载完了再算"，第二个保证"算完了再覆盖"。漏掉任何一个都会得到随机结果。
• tx 必须映射到最快变化的维度（As[ty][tx] 而不是 As[tx][ty]），否则载入不合并。
• 累加器 acc 在寄存器里，$T$ 次乘加期间完全不访存。真实 kernel 会让每个线程负责 $4\times4$ 或 $8\times8$ 个输出元素（register tiling），进一步提高复用。

复杂性之一：tile quantization（分块量化损失）

影响 tile 尺寸选择的三个因素（互相冲突）：

• 合并访存：tile 的行长要是 burst 大小的整数倍；
• shared memory 容量：tile 越大越好用，但装不下就得降低 occupancy；
• 矩阵维度的整除性：tile 要能整除矩阵尺寸，否则就是上图 (b) 的情况。

复杂性之二：内存对齐

解开谜题：为什么大矩阵更快

现在可以完整解释开头那张散点图了。三种效应叠加：

效应一：算术强度（$N$ 小于约 600–900 时）。 $I = 2N/(3b)$，小矩阵根本达不到 ridge point，被带宽限制。这解释了曲线左半部分的平滑爬升。

11. FlashAttention：三板斧的集大成

Part 3 的目标是证明：你现在已经有了推导出 FlashAttention 的全部工具。 它没有引入任何新的数学，只是把 tiling + fusion + recomputation 用到了极致，再加上一个 2018 年就发表的在线 softmax 技巧。

11.1 标准注意力为什么慢

注意力的计算是三个矩阵乘中间夹一个 softmax（$Q, K, V \in \R^{N \times d}$）：

标准实现把这三步写成三（或更多）个独立的 kernel，于是 $S$ 和 $P$ 这两个 $N \times N$ 的巨型矩阵必须被 materialize 到 HBM。算一笔账（$N = 4096$，$d = 64$，BF16，单个 head）：

差了 9 倍——标准注意力是彻头彻尾的 memory-bound 算子。 而且这个比值随 $N$ 增大而恶化：矩阵乘 FLOPs 是 $O(N^2 d)$，但 $S/P$ 的流量是 $O(N^2)$，所以强度上限被 $d$ 锁死在 $O(d)$，和序列长度无关。

还有第二个问题，甚至更致命：显存。$S$ 和 $P$ 要为反向传播存下来。$N=4096$、32 个 head、BF16：单层单样本就是 $32 \times 33.6 \text{MB} \approx 1.07$ GB，32 层就是 34 GB——一张 80GB 的 H100 连 batch size = 2 都放不下。这就是 FlashAttention 之前长上下文训练的墙。

11.2 第一步：对 $QK^\top V$ 做 tiling

难点在哪？softmax 是沿着 $N$ 维（key 维）的全局归约。要算 $\softmax(S_{i,:})$，你需要整行 $S_{i,:}$ 的最大值和指数和。可是分块之后，你手上一次只有 $S$ 的一小块 $S^{(j)} \in \R^{B_q \times B_k}$——你不知道这一行剩下部分的最大值是多少。

11.3 第二步：在线 softmax（online softmax）

import math, torch

def online_softmax_stats(x, block_size):
    """分块地算出全局 (max, sum of exp)，只扫一遍数据。"""
    m, l = -math.inf, 0.0
    for xb in x.split(block_size):
        m_tilde = xb.max().item()
        m_new   = max(m, m_tilde)
        # 关键一行：把历史累加量重新基准化到 m_new 上
        l = l * math.exp(m - m_new) + torch.exp(xb - m_new).sum().item()
        m = m_new
    return m, l          # softmax(x)_i == exp(x_i - m) / l

x = torch.randn(10000) * 30          # 故意放大，验证数值稳定性
m, l = online_softmax_stats(x, 128)
ref  = torch.softmax(x, dim=0)
print((torch.exp(x - m) / l - ref).abs().max())   # ~1e-8

11.4 组装：FlashAttention 的前向

import torch

def flash_attention_forward(Q, K, V, Bq=128, Bk=128):
    """用纯 PyTorch 复现 FlashAttention 前向的数学。
    真实 kernel 里 Kj/Vj/Sij/Oi 全部驻留在 SRAM 与寄存器中，
    这里用切片来"模拟"分块，重点看 O/m/l 的更新逻辑。
    Q, K, V: (N, d)"""
    N, d = Q.shape
    scale = d ** -0.5
    O = torch.zeros_like(Q)
    L = torch.zeros(N, 1, device=Q.device)      # 存 logsumexp，反向要用

    for i in range(0, N, Bq):                   # 外层：遍历 query 分块
        Qi = Q[i:i + Bq]                        # (Bq, d)  载入一次，全程驻留
        Oi = torch.zeros(Qi.shape[0], d, device=Q.device)   # 未归一化累加器
        mi = torch.full((Qi.shape[0], 1), float("-inf"), device=Q.device)
        li = torch.zeros(Qi.shape[0], 1, device=Q.device)

        for j in range(0, N, Bk):               # 内层：流式遍历 key/value 分块
            Kj, Vj = K[j:j + Bk], V[j:j + Bk]   # (Bk, d)
            Sij = (Qi @ Kj.T) * scale           # (Bq, Bk) —— 只存在于 SRAM

            m_new = torch.maximum(mi, Sij.max(dim=-1, keepdim=True).values)
            P     = torch.exp(Sij - m_new)      # 融合的 exp，不落 HBM
            alpha = torch.exp(mi - m_new)       # 历史累加量的修正因子 (<= 1)

            li = alpha * li + P.sum(dim=-1, keepdim=True)
            Oi = alpha * Oi + P @ Vj            # 向量版的伸缩求和
            mi = m_new

        O[i:i + Bq] = Oi / li                   # 最后统一归一化
        L[i:i + Bq] = mi + torch.log(li)        # logsumexp，反向重算 P 用
    return O, L


N, d = 1024, 64
Q, K, V = (torch.randn(N, d) for _ in range(3))
O, L = flash_attention_forward(Q, K, V)
ref = torch.softmax(Q @ K.T * d ** -0.5, dim=-1) @ V
print((O - ref).abs().max())          # ~1e-6：数值上完全等价（exact attention）

对照 Tatsu 总结的三个要素，逐条落位：

11.5 收益核算

（1）HBM 流量。 沿用 $N=4096, d=64$，BF16。FlashAttention 的流量：$Q, O$ 各读写一次（1.05 MB），$K, V$ 被每个 query 块重读一次。取 $B_q = 128$，共 32 个 query 块：

相比标准的 136 MB，降到约 1/4，算术强度从 32 提到 126 FLOP/byte。把 $B_q$ 调大到 256（受 SRAM 容量限制），只需 16 趟，流量降到 18 MB、强度升到 240——逼近 ridge point。这解释了 kernel 调优里为什么块大小如此关键。

（2）显存。 这是更重要的收益。标准实现要为反向存 $P \in \R^{N\times N}$；FlashAttention 只存 $O \in \R^{N \times d}$ 和 $L \in \R^{N}$（logsumexp）：

还是那个例子（$N=4096$、32 head、单层）：$32 \times 33.6\,\text{MB} = 1.07$ GB $\to$ $32 \times (4096\times64\times2 + 4096\times4) \approx 17$ MB，省了 63 倍。这才是 FlashAttention 真正改变游戏规则的地方——它让长上下文训练从"不可能"变成"routine"。

（3）反向传播。 论文里没在课上展开，但逻辑是对称的：反向时用存下来的 $Q, K, V, O, L$，逐块重新计算 $S^{(j)}$ 和 $P^{(j)}$（因为有 $L$ 在手，不需要再做在线 softmax，直接 $P = \exp(S - L)$），然后算梯度。这就是 §8 的重计算，只不过粒度是 tile 级别。额外的 FLOPs 约 +20–30%，换来 $O(N^2)$ 显存的消失。

11.6 后续版本

这个演进路线本身就是一堂课：每一代改进解决的都是"上一代之后新暴露出来的瓶颈"——先是 HBM 流量，然后是非矩阵乘算力和并行度，最后是计算与搬运的重叠。这正是性能优化的常态：瓶颈会移动，所以必须反复测量。

12. 实战：判断你的 kernel 到底被什么卡住

前面十一节给的是模型，这一节给的是流程。面对一个跑得慢的算子，按下面的顺序排查。

第 0 步：先算，再测

在打开任何 profiler 之前，先用纸笔（或 §6 那个 roofline() 函数）算出理论下界：

import torch, torch.utils.benchmark as bench

def measure(fn, *args, peak=989e12, bw=3.35e12, flops=None, bytes_moved=None):
    t = bench.Timer("fn(*args)", globals={"fn": fn, "args": args}).blocked_autorange()
    sec = t.median
    print(f"实测 {sec*1e3:.3f} ms")
    if flops:       print(f"  达成算力 {flops/sec/1e12:7.1f} TFLOP/s  ({100*flops/sec/peak:5.1f}% of peak)")
    if bytes_moved: print(f"  达成带宽 {bytes_moved/sec/1e9:7.1f} GB/s    ({100*bytes_moved/sec/bw:5.1f}% of peak)")

把实测时间换算成"达成算力"和"达成带宽"两个百分比，然后看落在哪一档：

第三种情况最常见，也最容易改。原因通常在这四个地方。

第 1 步：是不是根本没在 GPU 上跑？

先用 Nsight Systems（或 torch.profiler 导出 chrome trace）看时间线。要找的是：

• kernel 之间的空隙。空隙 = GPU 在等 CPU 派活。典型元凶：Python 层的循环、每步都调 .item() / .cpu() / print() 造成同步、data loader 太慢、大量微小 kernel 导致 launch 开销（每次 launch 约 3–10 μs，如果 kernel 本身只要 5 μs，那一半时间在 launch）。
• 对策：torch.compile（融合 + 减少 launch）、CUDA Graph（把一串 launch 打包成一次）、异步 data loading、去掉所有不必要的同步点。

第 2 步：occupancy（占用率）够不够

Occupancy = 一个 SM 上实际驻留的 warp 数 / 硬件上限。 它衡量的是"GPU 有多少备用工作可以用来掩盖延迟"。A100 每 SM 最多 64 个 warp（2048 线程）。三个东西会限制它：

例：每线程用 64 个寄存器 → $65536/64 = 1024$ 线程 = 32 warp → occupancy 50%。想提到 100% 就得把寄存器压到 32 个以内，但那样每个线程能缓存的中间值就少了，可能反而更慢。

第 3 步：访存合并了吗

用 Nsight Compute 看 Memory Workload Analysis：

• 关键指标：每次请求的 sector 数（l1tex__average_t_sectors_per_request_pipe_lsu_mem_global_op_ld）。理想值：FP32 连续访问是 4（32 线程 × 4 字节 = 128 字节 = 4 个 sector）。如果看到 32，说明完全没合并，带宽浪费 8 倍。
• 另一个信号：gld_efficiency / Global Load Efficiency，理想 100%，低于 50% 就要查索引方式。
• 对策：调整索引让 threadIdx.x 映射到最快变化的维度；改变张量布局；必要时先做一次转置（一次连续拷贝比 32 倍的低效访问便宜）。

第 4 步：shared memory bank conflict

Shared memory 被分成 32 个 bank，每个 bank 宽 4 字节，地址按 4 字节轮流分配到各 bank。一个 warp 的 32 个线程如果访问了同一个 bank 的不同地址，就会串行化（同一地址是广播，不冲突）。

Nsight Compute 里的指标是 l1tex__data_bank_conflicts_pipe_lsu_*，或者直接看 "Shared Memory" 那栏的 bank conflict 百分比。另一个现代解法是用 swizzle（按 XOR 打乱 tile 内的地址映射），Triton 和 CUTLASS 都是这么做的。

第 5 步：形状对不对

回到 §10 的结论，检查清单：

• 所有参与矩阵乘的维度是否是 8 的倍数（BF16 走 Tensor Core 的最低要求）？最好是 64 或 128 的倍数。
• 词表、隐藏维、FFN 维、head 数有没有 pad 到整齐的数？
• batch × seq 切出来的 tile 数是不是 SM 数（108 / 132）的整数倍附近？如果只比整数倍多一点点，就是典型的 wave quantization。
• 数据类型对了吗？torch.set_float32_matmul_precision("high") 打开 TF32；混合精度用 torch.autocast。
• 张量是 contiguous 的吗？x.is_contiguous()、x.stride()。

速查：症状 → 病因 → 处方

本讲小结

一页速查表

要点清单

1. 硬件驱动 scaling，底层细节决定什么能 scale、什么不能。 Dennard scaling 在 2005 年结束之后，算力增长全部来自并行化和专门化——这意味着只有"长得像硬件喜欢的样子"的计算才能享受红利。
2. 当前的 GPU 计算强烈鼓励你围绕"矩阵乘 + 数据搬运"来思考。 Tensor Core 让矩阵乘比其他浮点运算快 10 倍以上；HBM 带宽的增长又远慢于算力。所以：能写成大矩阵乘的就写成大矩阵乘，写不成的就想办法把它融进矩阵乘的前后。
3. 认真对待 GPU 的细节（合并、分块、融合）能带来实实在在的性能。 这不是"锦上添花"的微优化——本讲展示的例子里，随便一条都是 2–10 倍的差距。
4. 建立"先算算术强度"的肌肉记忆。 看到一个算子，30 秒内估出 FLOPs 和字节数，和 ridge point 一比，就知道优化的天花板在哪、该往哪个方向使劲。
5. FlashAttention 不是魔法。 它是 tiling（把 $S$ 留在 SRAM）+ fusion（exp 不落 HBM）+ 在线 softmax（让分块成为可能）+ 重计算（反向重建 $P$）的组合。你现在有能力自己推导出它——下一讲会用 Triton 真正把它写出来。

附录：延伸阅读

性能模型与 GPU 基础

• Roofline: An Insightful Visual Performance Model for Multicore Architectures (Williams et al., 2009) — roofline 模型的原始论文。全篇只有一个公式，但它定义了此后二十年的性能分析方法论。
• AI and Memory Wall (Gholami et al., 2024) — 系统统计了算力/带宽/互连三条曲线的增速差，是本讲"memory wall"论断的数据来源。
• Making Deep Learning Go Brrrr From First Principles (Horace He) — 本讲"工厂与仓库"比喻的出处。把深度学习性能问题干净地分成 compute / memory / overhead 三类，是入门后最该读的一篇博客。
• What Shapes Do Matrix Multiplications Like? (Horace He) — 第 10 节"矩阵尺寸之谜"的完整版本，逐层拆解 tile quantization、对齐、wave quantization。读完你会对"改个数字就快 30%"这件事免疫。
• NVIDIA Matrix Multiplication Background User's Guide — 官方文档，tile quantization 和 wave quantization 的权威解释，附大量实测曲线。查形状对齐规则时的第一手资料。
• How to Scale Your Model: A Systems View of LLMs (JAX 团队) — 课上重点推荐的 "TPU (and now GPU) book"。把 roofline 的思路一路推广到多芯片并行，是本讲和后面并行化讲之间的最佳桥梁。
• GPU MODE (原 CUDA MODE) 讲座与代码 — 社区驱动的 CUDA/Triton 实战课程，从 hello world 到手写 FlashAttention 都有可运行代码。想真正动手时从这里开始。
• Better Performance at Lower Occupancy (Volkov, GTC 2010) — 反直觉但极重要：说明为什么高 occupancy 不是目标，ILP 同样能掩盖延迟。写高性能 kernel 前必读。

FlashAttention 与算子优化

• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Dao et al., 2022) — 本讲 Part 3 的主角。重点看 Section 3 的算法伪代码和 Theorem 2 的 IO 复杂度分析，正好对应第 11 节的推导。
• Online normalizer calculation for softmax (Milakov & Gimelshein, 2018) — 在线 softmax 的原始论文，比 FlashAttention 早四年。短小精悍，说明"关键数学技巧往往早就躺在那里，等一个知道怎么用它的人"。
• FlashAttention-2 (Dao, 2023) — 讲清楚了为什么 FA1 只有 25–40% 的 FLOPs 利用率，以及怎么通过减少非矩阵乘操作、改变并行维度来修复。是学习"如何做第二轮优化"的范本。
• FlashAttention-3 (Shah et al., 2024) — Hopper 架构的异步特性（TMA、WGMMA、warp specialization）如何进一步把搬运和计算重叠。想理解现代 GPU 编程范式的走向就读这篇。
• Data Movement Is All You Need: A Case Study on Optimizing Transformers (Ivanov et al., 2020) — 系统地量化了 Transformer 里每类算子的数据搬运成本，得出"融合能带来 1.3× 端到端加速"的结论。比 FlashAttention 更早提出同样的世界观。
• Training Deep Nets with Sublinear Memory Cost (Chen et al., 2016) — 激活检查点的原始论文，$O(\sqrt{L})$ 显存策略的来源。第 8 节重计算部分的理论基础。
• Min-cut optimal recomputation with AOTAutograd (PyTorch dev-discuss) — 课上三层 sigmoid 例子的出处。讲 PyTorch 编译器如何把"存哪些激活"建模成图上的最小割问题并自动求解。

低精度与数值格式

• Mixed Precision Training (Micikevicius et al., 2017) — FP16 训练的奠基之作，确立了"低精度存储 + FP32 master weights + FP32 累加 + loss scaling"的标准配方，至今仍是所有混合精度实现的骨架。
• FP8 Formats for Deep Learning (Micikevicius et al., 2022) — E4M3 与 E5M2 两种格式的设计动机，以及为什么前向用前者、反向用后者。
• Microscaling Data Formats for Deep Learning (Rouhani et al., 2023) — OCP MX 标准（MXFP8/6/4）的定义：块级共享指数缩放因子的来龙去脉。
• Recipes for Pre-training LLMs with MXFP8 (2025) — 课上引用的 MXFP8 实战论文。重点看它讨论的"哪些层不能量化"和"转置为什么需要单独量化"，这是理论到工程的真实距离。

硬件与 TPU

• In-Datacenter Performance Analysis of a Tensor Processing Unit (Jouppi et al., 2017) — TPU v1 的架构论文，脉动阵列设计与 roofline 分析的经典案例。第 5 节脉动阵列部分的原始出处。
• TPU v4: An Optically Reconfigurable Supercomputer (Jouppi et al., 2023) — 3D torus 拓扑和光路交换（OCS）如何让 4096 芯片的 pod 保持均匀带宽。为后面的并行化讲打底。
• Scaling Laws for Neural Language Models (Kaplan et al., 2020) — 本讲开头"算力带来可预测的性能提升"这一前提的来源，也是整门课的动机所在。
• Megatron-LM (Shoeybi et al., 2019) 与 Efficient Large-Scale Language Model Training on GPU Clusters (Narayanan et al., 2021) — 把本讲的单卡性能思维扩展到千卡集群，后者报告了 1T 参数模型 52% MFU 的实现细节。
• CUDA Refresher: Reviewing the Origins of GPU Computing (NVIDIA) — 从可编程着色器到 CUDA 的历史，解释了"为什么 GPU 恰好适合深度学习"这件事有多大成分是历史巧合。