第 5 讲：GPU 与 TPU

对应：Stanford CS336 Spring 2026 Lecture 5（GPUs, TPUs，Tatsunori Hashimoto 主讲）。Spring 2025 的 Lecture 5（GPUs）有公开录像。 配套代码：code/lecture05/ 预计阅读时间：3–4 小时 硬件说明：本讲的配套代码不需要 NVIDIA GPU。GPU 内部的机制（warp、合并访存、分块、脉动阵列）用 Python 模拟；能在普通电脑上观察到的现象（roofline、算子融合、低精度、跨步访存）会在你的电脑上实测。

0. 本讲导航

0.1 学完本讲，你应该能够

说出 GPU 与 CPU 在设计目标上的根本差别，画出 GPU 的结构（SM、Tensor Core、寄存器、共享内存、L2、HBM）。
解释线程、线程块、warp、网格的含义，以及 SIMT 执行模型。
用 Roofline 模型判断一个运算是计算受限还是内存受限；解释为什么推理解码是内存受限的。
说出让 GPU 跑得快的六个技巧：避免分支发散、低精度、算子融合、重计算、合并访存、分块，并能用”数读写次数”的方法估算它们的收益。
解释矩阵乘法性能随尺寸变化的”锯齿”现象（块量化、对齐、波次量化）。
推导在线 softmax，讲清 FlashAttention 前向和反向的算法，以及它为什么又快又省显存。
说出 TPU 的脉动阵列是如何工作的，TPU 与 GPU 的异同。

0.2 先修知识

知识点	用在哪里	最低要求
第 2 讲	全文	FLOPs、算术强度初步、混合精度
第 3 讲	第 9 节	注意力的公式、softmax
矩阵乘法	第 7.6 节	知道 $C[i,j] = \sum_k A[i,k]B[k,j]$

0.3 本讲地图

第一部分：GPU 是怎么工作的
   第1节 为什么是 GPU：从 Dennard 缩放到并行缩放
   第2节 GPU 的结构：执行单元与存储层次
   第3节 执行模型：线程、线程块、warp
   第4节 计算比内存涨得快
第二部分：让 ML 负载在 GPU 上跑得快
   第5节 Roofline 模型
   第6节 六个技巧：分支、低精度、融合、重计算、合并访存、分块
   第7节 "矩阵之谜"：为什么性能曲线是锯齿状的
第三部分：把一切组合起来 —— FlashAttention
   第8节 在线 softmax 与 FlashAttention 的前向、反向
第四部分：TPU
   第9节 脉动阵列；TPU 与 GPU 的异同

CS336 课上对本讲的定位是：让 CUDA 和 GPU 不再神秘。目标是理解 GPU 什么时候会变慢，以及如何设计快的算法。

第一部分：GPU 是怎么工作的

1. 为什么是 GPU

1.1 计算带来可预测的进步

规模定律（第 9 讲）告诉我们：更多的计算通常带来可预测的效果提升。所以更快的硬件、更高的利用率、更好的并行化，本身就能推动进步。问题是：计算从哪里来？

1.2 Dennard 缩放的终结

1980~2000 年代，芯片的进步主要来自 Dennard 缩放：晶体管变小，频率提高，功耗密度不变，单线程程序自动变快。但这条路在 2000 年代中期走到了尽头——频率再也提不上去（功耗和散热的限制）。

1.3 并行缩放

之后的进步主要来自并行：同时做更多的事。NVIDIA 首席科学家 Bill Dally 的报告显示，GPU 在 10 年间的算力增长超过 1000 倍（来自低精度数字格式、Tensor Core 这样的专用矩阵乘法单元、工艺进步等）。CS336 课上的一句话：没有 GPU 的缩放，就没有大模型的缩放。

1.4 GPU 与 CPU 的根本差别

	CPU	GPU
优化目标	延迟：让每个线程尽快完成	吞吐量：让总处理量最大
线程	少量、很快的线程	成千上万个轻量线程
芯片面积花在哪	大量控制逻辑（分支预测、乱序执行）和大缓存	大量简单的计算单元（ALU），很少的控制逻辑
擅长	复杂的、分支多的串行程序	对大量数据做同样的计算

深度学习的计算（尤其是矩阵乘法）恰好是”对大量数据做同样的计算”。

2. GPU 的结构

2.1 执行单元

GPU 由很多个流式多处理器（SM，Streaming Multiprocessor）组成：A100 有 108 个，H100 有 132 个，B200 有 148 个。每个 SM 独立地执行分配给它的工作（线程块）。
每个 SM 里有许多流处理器（SP，也叫 CUDA 核心），可以并行执行线程；还有几个 Tensor Core——专门做小矩阵乘法的电路。
Tensor Core 让矩阵乘法比其他浮点运算快 10 倍以上（H100：BF16 矩阵乘法约 990 TFLOP/s，而普通 FP32 运算约 67 TFLOP/s）。这意味着：凡是不能写成矩阵乘法的运算，相对来说都很”贵”。

2.2 存储层次

离计算单元越近的存储越快，但也越小、越贵：

层次	位置	大小（H100 量级）	速度
寄存器	每个线程私有	每个 SM 共 256 KB（65536 个 32 位寄存器）	最快
L1 缓存 / 共享内存	SM 内部	每个 SM 最多约 228 KB 共享内存	很快（每秒几十 TB）
L2 缓存	芯片上（所有 SM 共享）	约 50 MB	较快
全局内存（HBM）	芯片旁边的内存颗粒	80 GB	最慢（3.35 TB/s）

片上的 SRAM（共享内存、缓存）比 HBM（DRAM）贵约 100 倍，但快约 8 倍（CS336 课件的数字）。写快的 GPU 程序，核心就是尊重这个存储层次：尽量让数据待在离计算近的地方，少跑去 HBM。

3. 执行模型：线程、线程块、warp

3.1 三个角色

线程（thread）：做实际工作的最小单位。所有线程执行同一段代码，但处理不同的数据（SIMT：单指令多线程）。
线程块（block）：一组线程。一个线程块在一个 SM 上运行，块内的线程共享这个 SM 的共享内存，可以互相通信和同步。
warp：硬件实际调度的单位，由 32 个编号连续的线程组成，它们在同一时刻执行同一条指令。
网格（grid）：一次核函数调用启动的所有线程块。

每个线程通过自己的编号知道该处理哪个数据：i = blockIdx × blockDim + threadIdx。因为数据量不一定是块大小的整数倍，代码里必须有边界检查 if (i < n)（配套代码 02_gpu_execution_model.py 演示 1：处理 1000 个元素、每块 256 个线程，要启动 1024 个线程，其中 24 个什么也不做）。

3.2 存储模型

每个线程可以访问自己的寄存器，以及所在线程块的共享内存。不同线程块之间的信息交换必须经过全局内存（HBM）——这很慢。所以好的算法会让每个线程块尽量独立地完成自己的工作。

3.3 这套模型的优点

容易扩展：更大的 GPU 只是多加一些 SM；
容易（相对而言）编程：写一个线程的代码，自动变成成千上万个线程；
线程很”轻”：可以零开销地暂停和切换（一个 warp 等待内存时，SM 立刻切换去执行另一个 warp）。

4. 计算比内存涨得快

一个贯穿本讲的关键事实：硬件的算力增长远快于内存带宽的增长。配套代码 01_roofline.py 演示 1：

GPU	SM 数	BF16 算力	HBM 带宽	平衡点（算力 / 带宽）
A100	108	312 TFLOP/s	2.0 TB/s	156 FLOP/字节
H100	132	990 TFLOP/s	3.35 TB/s	295 FLOP/字节
B200	148	2250 TFLOP/s	8.0 TB/s	281 FLOP/字节

从 A100 到 B200，算力涨了约 7 倍，带宽只涨了约 4 倍。平衡点在 150~300 之间，意思是：每从 HBM 读一个字节，要做一两百次运算，计算单元才不会闲着。这就是所谓的”内存墙”：我们越来越难以把计算单元”喂饱”。

第二部分：让 ML 负载在 GPU 上跑得快

5. Roofline 模型

5.1 模型

一个运算的耗时由两者中较慢的决定：计算（FLOPs ÷ 峰值算力）和搬运数据（字节数 ÷ 带宽）。定义算术强度 $I = \text{FLOPs} / \text{字节数}$ ，则可达到的算力为：

\text{可达算力} = \min(\text{峰值算力},\ I \times \text{带宽})

画成对数坐标图，就是一个”屋顶”：强度低时沿斜线上升（内存受限），过了平衡点后变成水平线（计算受限）。本节的核心问题就是：怎么避免内存受限？

5.2 常见运算的算术强度

$n = 4096$ 、BF16、H100（演示 2）：

运算	算术强度	Roofline 预测可达算力	占峰值
ReLU（逐元素）	0.25	0.8 TFLOP/s	0.1%
点积	0.5	1.7 TFLOP/s	0.2%
矩阵 × 向量	1.0	3.3 TFLOP/s	0.3%
矩阵 × 矩阵	1365	990 TFLOP/s	100%

矩阵乘法的强度约为 $n/3$ （ $2n^3$ 次运算，只读写 $3n^2$ 个数）， $n$ 足够大时远超平衡点。GPU 真正擅长的只有大矩阵乘法，其他运算都只能用到峰值的零头。

5.3 为什么推理解码是内存受限的

一个 $d\times d$ 的权重矩阵乘以 $B$ 个 token（ $d = 8192$ ，BF16，H100，演示 3）：

批次中的 token 数 $B$	算术强度	可达算力占峰值
1	1.0	0.3%
8	8.0	2.7%
64	63	21%
256	241	82%
1024	819	100%

算术强度约等于 $B$ ：权重读一次，被 $B$ 个 token 重复使用。训练和预填充一次处理成千上万个 token，轻松计算受限；逐个生成 token 的解码，每个请求每步只有 1 个 token，必须把许多请求凑成大批次，否则大部分时间都在搬权重和 KV 缓存。这是第 10 讲推理优化的出发点。

5.4 你电脑上的 roofline

演示 4 在你的电脑上实测带宽和矩阵乘法算力，画出属于你的 roofline。在一台 Apple M4 Pro 上（Apple GPU / MPS）：带宽约 215 GB/s，FP32 矩阵乘法约 5.5 TFLOP/s，平衡点约 26 FLOP/字节。

运算	算术强度	实测	Roofline 预测
`x*2`	0.125	28 GFLOP/s	27 GFLOP/s
`x*2+1`	0.25	30 GFLOP/s	54 GFLOP/s
矩阵乘法 64×64	11	9 GFLOP/s	2295 GFLOP/s
矩阵乘法 1024×1024	171	4640 GFLOP/s	5519 GFLOP/s
矩阵乘法 4096×4096	683	5383 GFLOP/s	5519 GFLOP/s

x*2 与预测吻合：完全被带宽卡住。
x*2+1 只达到预测的一半：PyTorch 把它分成两个核函数（先算 x*2 写回内存，再读出来加 1），搬运量翻了一倍——这正是第 6.3 节的算子融合要解决的问题。
小矩阵远低于预测：roofline 只是上界，还有启动开销、调度等固定成本。

6. 让 GPU 跑得快的六个技巧

6.1 技巧 1：避免分支发散

同一个 warp 的 32 个线程必须执行同一条指令。如果它们遇到 if/else 时走了不同的路，硬件只能先执行 if 分支（else 的线程闲着），再执行 else 分支（if 的线程闲着）——这叫分支发散。演示 2（两个分支各需 10 步，4 个 warp）：

条件	总步数	利用率
`threadIdx % 2 == 0`（每个 warp 内都一半一半）	80	50%
`(threadIdx / 32) % 2 == 0`（按 warp 划分）	40	100%
`threadIdx < 100`（只有跨越边界的 warp 发散）	50	88%

条件判断本身不是问题，问题是同一个 warp 内走了不同的路。（这一条与内存无关，其余五条都是在对付内存瓶颈。）

6.2 技巧 2：低精度

位数少了，要搬的字节就少了。 以 ReLU 为例：FP32 时每次运算要搬 8 字节（读 4、写 4），FP16/BF16 只要 4 字节——算术强度翻倍。配套代码 03_fusion_precision_recompute.py 演示 1：复制 3355 万个元素，FP32 用 1.31 ms，BF16 用 0.59 ms——带宽一样，字节减半，时间减半。

对矩阵乘法，Tensor Core 在低精度下的算力也高得多。前沿方向是更低的精度：Blackwell（B200）支持 MXFP8（E4M3 数值 + 每 32 个元素一个 E8M0 格式的缩放因子）和 MXFP4 / NVFP4（每 16 个元素一个缩放因子）。这些格式带来了新的麻烦，比如转置变得不再简单（分块缩放是按行分块的，转置后分块方向变了，需要重新量化）。实践中也不是所有权重都用 MXFP8，需要仔细选择（第 2 讲第 3.6 节）。

6.3 技巧 3：算子融合

CS336 课上（引用 Horace He 的博客）的比喻：把 GPU 想象成一座工厂，输入存放在仓库（显存）里，在工厂（计算单元）里加工。计算能力（工厂）一直在扩建，但仓库和工厂之间的运输（带宽）跟不上。如果要做很多道工序，每做完一道就把半成品运回仓库、再运回来做下一道，就太傻了。

融合：把多个运算合成一个核函数，中间结果留在寄存器里，不写回显存。以 $\sin^2 x + \cos^2 x$ 为例：

	核函数数	显存读	显存写	合计
未融合	5	$6n$	$5n$	$11n$
融合	1	$n$	$n$	$2n$

理论上可以快 $11/2 = 5.5$ 倍。在 Apple GPU 上实测（1677 万个元素）：逐个运算 3.34 ms，torch.compile 融合后 0.62 ms，加速 5.4 倍——几乎恰好等于理论值。手写的 GELU（多个基本运算）融合后加速 9.1 倍，与 PyTorch 内置的融合版 F.gelu 速度相当。

像这样的逐元素融合，编译器（torch.compile）可以自动完成；更复杂的融合（如 FlashAttention）则需要手写核函数（第 6 讲）。在 CPU 上加速往往没那么明显：sin、cos、tanh 在 CPU 上计算本身就较慢，不完全是内存受限。

6.4 技巧 4：重计算

反向传播需要前向的中间结果（激活值）。通常的做法是存下来，反向时再读回来——但”存”和”读”都要访问显存。以三个 sigmoid 叠加 $y = \sigma(\sigma(\sigma(x)))$ 为例（每个 Function 内部是融合的）：

	前向	反向	合计
保存中间结果	读 $x$ ，写 $s_1, s_2, s_3$	读上游梯度、 $s_1, s_2, s_3$ ，写输入梯度	9
重计算	读 $x$ ，写 $s_3$	读上游梯度、 $x$ ，写输入梯度（ $s_1, s_2, s_3$ 在寄存器里重算）	5

把计算扔掉、重新算一遍，反而可能更快：多做了 3 次 sigmoid，少搬了近一半数据。配套代码用 torch.autograd.Function 实现了两个版本：为反向保存的数据从 12.6 MB 降到 4.2 MB，梯度完全相同，gradcheck 通过。第 2 讲的激活检查点是同一个思想，只是粒度更粗（整层）。

6.5 技巧 5：合并访存

显存（DRAM）以突发模式（burst）读取：由于 DRAM 的物理结构（先把一整行复制到感应放大器），每次读取都会一次性取出一大段连续的字节。如果一个 warp 的 32 个线程访问的地址都落在同一段里，一次内存事务就够了——这叫合并访存（coalesced）。演示 3（每个线程读一个 FP32，128 字节一段）：

访问模式	内存事务数	有效字节占比
线程 $t$ 读 `x[t]`（连续）	1	100%
线程 $t$ 读 `x[2t]`	2	50%
线程 $t$ 读 `x[32t]`	32	3%
随机位置	32	3%

对按行存储的矩阵：32 个线程读同一行的相邻元素只需 1 次事务，读同一列的相邻元素需要 32 次。所以写矩阵乘法等核函数时，要让同一个 warp 中相邻的线程访问内存中相邻的地址。你电脑上的类似现象：复制一个 8192×8192 矩阵，按顺序复制 2.6 ms，转置后复制（跨步读）5.6 ms——数据量相同，只是访问顺序不同。

6.6 技巧 6（最重要的一个）：分块

分块（tiling）：对线程进行分组和排序，最大限度地减少对全局内存的访问。以矩阵乘法为例：

朴素做法：每个输出元素 $C[i,j]$ 都从显存读 $A$ 的第 $i$ 行和 $B$ 的第 $j$ 列。同一个 $A[i,k]$ 被读了 $n$ 次，访问既重复又不合并。
分块做法：把矩阵切成 $T\times T$ 的小块。计算一个输出块时，分阶段把 $A$ 的一个小块和 $B$ 的一个小块搬进共享内存，在里面完成这两个小块能做的全部乘加，累加到部分和上，再换下一对。

每个输入元素从全局内存读取的次数从 $n$ 次降到 $n/T$ 次（在块内被复用 $T$ 次）——全局内存访问减少到 $1/T$ 。配套代码 04_tiling_matmul.py 模拟并计数（ $n = 256$ ）：

方式	从显存读取的元素	算术强度（FP32）
朴素	33,554,432（ $2n^3$ ）	0.25
分块 $T=4$	8,388,608（ $2n^3/4$ ）	1
分块 $T=16$	2,097,152	4
分块 $T=64$	524,288	16

分块还顺带解决了合并访存：搬一个小块时，warp 读的是连续的一段内存。

分块的复杂之处：

块的大小受共享内存限制：H100 每个 SM 最多约 228 KB，要同时放下两个 $T\times T$ 的块，FP32 时 $T$ 最大约 170。实际的核函数用 128×128、128×256 这样的矩形输出块，沿 K 方向每次只搬 32 或 64 列，并用”双缓冲”让搬运和计算重叠。
块量化：矩阵边长不能被块大小整除时，最后一行/列的块大部分是空的。块大小 128 时， $n = 1024$ 的有效计算是 100%， $n = 1025$ 只有 79%。
对齐：如果一行的字节数不是内存事务大小的整数倍，块就对不齐内存段，无法合并访存，需要填充。

6.7 补充：占用率

一个 SM 能同时驻留多少个线程块，受寄存器、共享内存、最大线程数、最大块数四种资源限制。驻留的 warp 越多，一个 warp 等待内存时，SM 越有可能切换去执行别的 warp，从而掩盖延迟。演示 4（H100 的 SM）：每块 256 线程、每线程 32 个寄存器时占用率 100%；每线程 64 个寄存器时降到 50%（寄存器成了瓶颈）；每块 128 线程、每线程 160 个寄存器时只有 19%。占用率低不一定慢——如果每个线程做很多工作（“线程粗化”），或者算法本身数据复用很好，低占用率也能跑得很快。

6.8 小结

目标	手段
减少内存访问	合并访存、算子融合
把数据搬到共享内存	分块
用计算/精度换内存	量化（低精度）、重计算

7. “矩阵之谜”：为什么性能曲线是锯齿状的

CS336 课件引用了 Horace He 的一个实验：测量方阵乘法在不同尺寸下的吞吐量，得到一条整体上升、但有很多周期性下降的锯齿曲线。

整体上升：矩阵越大，算术强度越高，越接近计算受限（第 5 节）。
与对齐相关的下降：尺寸不是 2 的若干次幂的倍数时，块对不齐，访存无法合并（第 6.6 节）。
波次量化：输出被切成许多块，每个块交给一个 SM；块数超过 SM 数时要分多”波”执行。

演示 5 模拟了课件中的例子（A100，108 个 SM，输出块 256×128）：

$n$	输出块数	波数	SM 利用率	相对 1792 的计算量	相对耗时（模型估计）
1792	$7\times14 = 98$	1	91%	1.00	1.0
1793	$8\times15 = 120$	2	56%	1.00	2.0
2560	200	2	93%	2.92	2.0

$n$ 从 1792 变成 1793，计算量只多了 0.2%，但块数从 98 变成 120，超过了 108 个 SM，必须分两波——第二波只有 12 个块在干活。（这是一个粗略模型：真实的库会根据形状选择不同的分块方案，实测曲线是锯齿状的，而不会这么极端。）

实用启示：让模型的各个维度（隐藏维度、词表大小、批次大小）取 64 或 128 的倍数。Karpathy 在 nanoGPT 中把 GPT-2 的词表从 50257 填充到 50304（64 的倍数），就带来了明显的加速。

第三部分：把一切组合起来 —— FlashAttention

8. FlashAttention

8.1 标准注意力的问题

S = QK^\top/\sqrt d \quad\to\quad P = \text{softmax}(S) \quad\to\quad O = PV

标准实现要把两个 $n\times n$ 矩阵（ $S$ 和 $P$ ）写进显存再读出来。一个头、 $n = 64$ K、BF16：光这两个矩阵就要 17 GB。而且这些运算的算术强度很低（softmax 是逐行的归约 + 逐元素运算），完全内存受限。

FlashAttention（Dao et al. 2022）大幅加速了注意力。它的技巧，拆开看就是本讲学过的东西：

分块：对 $QK^\top$ 和 $PV$ 这两个矩阵乘法做分块（论文的图 1 本质上就是分块矩阵乘法）；
在线 softmax：softmax 需要整行的最大值和指数和，分块后怎么办？——一块一块地增量计算；
融合：指数运算等全部在一个核函数里完成， $n\times n$ 矩阵从不写回显存；
重计算：反向传播时逐块重算 $P$ ，而不是存下来。

8.2 在线 softmax

标准 softmax 需要三趟扫描：求最大值 $m$ 、求 $\ell = \sum_j e^{x_j - m}$ 、归一化。在线 softmax（Milakov & Gimelshein 2018）把前两趟合成一趟：维护运行中的最大值 $m$ 和以它为基准的指数和 $\ell$ 。读入新的一块后，最大值可能变大（ $m \to m'$ ），之前累加的指数和要”换基准”：

m' = \max(m,\ \max_{j\in\text{新块}} x_j),\qquad \ell' = \ell\cdot e^{m - m'} + \sum_{j\in\text{新块}} e^{x_j - m'}

因为 $e^{x_j - m} \cdot e^{m - m'} = e^{x_j - m'}$ ，所有旧项都被正确地换到了新基准上（CS336 课件称之为”伸缩求和”的技巧）。配套代码 05_flash_attention.py 演示 1 验证了在线算法与 torch.softmax 的差异只有 $5\times10^{-10}$ 。

8.3 FlashAttention 前向

对每个 query 块 $Q_i$ （在 GPU 上由一个线程块负责，数据留在 SRAM）：

\begin{aligned} &\text{初始化 } O_i = 0,\ m_i = -\infty,\ \ell_i = 0\\ &\text{对每个 key/value 块 } (K_j, V_j):\\ &\qquad S_{ij} = Q_iK_j^\top/\sqrt d \quad\text{（只在 SRAM 中）}\\ &\qquad m_i^{new} = \max(m_i, \text{rowmax}(S_{ij})),\quad P_{ij} = e^{S_{ij} - m_i^{new}}\\ &\qquad \ell_i \leftarrow e^{m_i - m_i^{new}}\ell_i + \text{rowsum}(P_{ij})\\ &\qquad O_i \leftarrow e^{m_i - m_i^{new}}O_i + P_{ij}V_j \quad\text{（之前的结果也要换基准）}\\ &\qquad m_i \leftarrow m_i^{new}\\ &O_i \leftarrow O_i/\ell_i,\qquad L_i = m_i + \log\ell_i\ \text{（logsumexp，留给反向）} \end{aligned}

因果掩码下，完全在”未来”的 key 块可以直接跳过，省掉约一半计算。演示 2：与标准注意力、F.scaled_dot_product_attention 的差异在 $10^{-16}$ 量级——FlashAttention 是精确计算，不是近似。

8.4 FlashAttention 反向：重计算

反向需要 $P$ ，但我们没有存它。FlashAttention-2 的写法只需要 $Q, K, V, O, L$ 和 $dO$ ：

P_{ij} = e^{S_{ij} - L_i}\ \text{（由 logsumexp 重算）},\quad dV_j \mathrel{+}= P_{ij}^\top dO_i,\quad dP_{ij} = dO_iV_j^\top

dS_{ij} = P_{ij}\odot(dP_{ij} - D_i),\quad D_i = \text{rowsum}(dO_i\odot O_i),\quad dQ_i \mathrel{+}= dS_{ij}K_j/\sqrt d,\quad dK_j \mathrel{+}= dS_{ij}^\top Q_i/\sqrt d

其中 $D_i$ 这一项来自 softmax 的求导（ $\sum_j P_{ij}dP_{ij}$ ），可以用 $dO_i$ 和 $O_i$ 直接算出，不需要完整的 $P$ 。演示 3：用 torch.autograd.Function 包装后，梯度与 PyTorch 自动求导的差异在 $10^{-15}$ 量级。

8.5 省了多少

一个头、BF16、 $d = 128$ （演示 4）：

$n$	标准： $S$ 和 $P$ 占用	Flash：额外占用
4096	0.07 GB	0.02 MB
16384	1.07 GB	0.07 MB
65536	17.2 GB	0.26 MB

显存：从 $O(n^2)$ 降到 $O(n)$ 。训练时也不再需要为注意力矩阵保存激活（第 3 讲：GPT-2 XL 一半以上的激活就是这两个矩阵）。
显存访问：论文的 IO 复杂度分析给出 $\Theta(n^2d^2/M)$ （ $M$ 为 SRAM 大小），标准实现是 $\Theta(n^2)$ 的若干倍。按演示中的估算约少一个数量级。论文报告 GPT-2 的端到端训练约快 3 倍。
计算量并没有减少（反向时还因为重算增加了）——但注意力是内存受限的，少搬数据才是关键。

第四部分：TPU

9. 脉动阵列与 TPU

9.1 TPU 与 GPU：大同小异

CS336 课上的观点：GPU、TPU 和许多其他加速器，在高层次上是相似的——轻量的控制逻辑 + 很快的（大）矩阵乘法单元 + 很快的片上存储。差别主要在：

核心的粒度：GPU 有很多个 SM（H100 有 132 个）；TPU 每个芯片只有少数几个 TensorCore，但每个里面有很大的矩阵乘法单元（MXU）。两者的矩阵乘法峰值性能相近。
没有 warp：TPU 只有”块”的概念，这让它在非矩阵乘法的运算上相对弱一些（需要权衡）。
联网方式：TPU 芯片之间用专门的高速互连（ICI）直接连成大规模的 pod，这在并行训练时很重要（第 7–8 讲）。

9.2 脉动阵列

TPU 的 MXU 是一个脉动阵列（TPU v2~v5 为 128×128，v6e Trillium 为 256×256）：数据像心脏泵血一样有节奏地流过一个网格状的处理单元（PE）阵列—— $A$ 的行从左边流入、每周期右移一格， $B$ 的列从上边流入、每周期下移一格，每个 PE 把经过它的两个数相乘并累加。 $A$ 的第 $i$ 行要比第 0 行晚 $i$ 个周期进入（“错位”输入），这样 $A[i,k]$ 和 $B[k,j]$ 才能恰好在同一个周期到达 PE $(i,j)$ 。

配套代码 06_systolic_array.py 逐周期模拟了这个过程：

4×4 阵列、 $K = 6$ ：12 个周期算完，结果与 A @ B 一致；开始和结束时只有部分 PE 在工作（流水线的”填充”和”排空”），利用率 50%。
8×8 阵列： $K = 8$ 时利用率 36%， $K = 512$ 时 97%——TPU 擅长又大又规整的矩阵乘法。
每个数据只从内存读一次，在阵列中被 8 个 PE 依次使用。这和 GPU 上”分块 + 共享内存”追求的目标一样，只是用硬件结构直接实现了数据复用，因此只需要很少的控制逻辑和片上缓存。

10. 本讲小结

硬件的进步来自并行缩放；GPU 为吞吐量而设计：很多 SM、很多线程、Tensor Core 让矩阵乘法快 10 倍以上。
存储层次：寄存器 > 共享内存 > L2 > HBM，越近越快越小。算力比带宽涨得快，平衡点在 150~300 FLOP/字节。
执行模型：线程组成块，块在 SM 上运行；32 个线程组成 warp 同步执行；块之间只能通过 HBM 交流。
Roofline 模型：可达算力 = min(峰值, 强度 × 带宽)。只有大矩阵乘法是计算受限的；逐元素运算、小批次解码都是内存受限的。
六个技巧：避免 warp 内分支发散；低精度减少字节；融合避免中间结果往返显存；重计算用计算换访存；合并访存让 warp 读连续内存；分块让数据在共享内存里复用 $T$ 次。
性能曲线的锯齿来自对齐、块量化和波次量化；维度取 64/128 的倍数是最便宜的优化。
FlashAttention = 分块 + 在线 softmax + 融合 + 重计算：精确，显存 $O(n)$ ，更快。
TPU 用巨大的脉动阵列直接实现数据复用；与 GPU 大同小异，差别在核心粒度和互连方式。

11. 与作业 2 的关系

作业 2（系统）要求你：

基准测试与性能分析：正确地测量 Transformer 前向、反向、优化器步骤的时间（热身、同步），用性能分析工具（PyTorch profiler、NVIDIA Nsight Systems）找出瓶颈（第 6 讲详细讲）；分析混合精度的效果和显存占用。
FlashAttention-2：先用纯 PyTorch 实现分块的前向（本讲 05 号文件的思路），再用 Triton 写成真正的 GPU 核函数（第 6 讲），反向用重计算实现；并与标准注意力比较速度和显存。
（2026 年版的第 1 讲介绍中还提到用 Triton 实现融合的 RMSNorm 核函数，思路与本讲的融合、第 6 讲的 Triton 编程完全一致。）
分布式训练部分见第 7–8 讲。

12. 练习

练习 1（平衡点）：一块 GPU 的 BF16 算力为 400 TFLOP/s，带宽 2 TB/s。它的平衡点是多少？一个算术强度为 50 的运算最多能达到多少算力？

参考答案

平衡点 = 400/2 = 200 FLOP/字节。强度 50 < 200，内存受限，可达算力 = 50 × 2 TB/s = 100 TFLOP/s（峰值的 25%）。

练习 2（算术强度）： $n\times n$ 的 BF16 矩阵乘法（读 $A$ 、 $B$ ，写 $C$ ），算术强度是多少？在 H100 上， $n$ 至少多大才能达到计算受限？

参考答案

FLOPs $= 2n^3$ ，字节 $= 3n^2\times2 = 6n^2$ ，强度 $= n/3$ 。 $n/3 \ge 295$ → $n \ge 885$ 左右。（实际中还要考虑分块、占用率等因素，所以要更大一些才能接近峰值。）

练习 3（解码的算术强度）：推理解码时，批次里有 $B$ 个请求，每个请求这一步只产生 1 个 token。权重矩阵 $d\times d$ 、BF16。忽略激活的读写，算术强度约是多少？在 H100 上， $B$ 至少多大才能不被带宽卡住？

参考答案

FLOPs $= 2Bd^2$ ，字节 $\approx 2d^2$ ，强度 $\approx B$ 。需要 $B \ge 295$ 左右。所以推理服务要把很多请求凑成一批（连续批处理，第 10 讲）。

练习 4（融合）：计算 $y = \text{ReLU}(x\cdot w + b)$ （逐元素， $x, w, b$ 都是长度 $n$ 的向量），未融合时有几个核函数？读写多少？融合后呢？

参考答案

未融合 3 个核函数： $t_1 = x\cdot w$ （读 2n 写 n）、 $t_2 = t_1 + b$ （读 2n 写 n）、 $y = \text{ReLU}(t_2)$ （读 n 写 n），共 $8n$ 。融合后读 $x, w, b$ （3n）写 $y$ （n），共 $4n$ ，理论加速 2 倍。

练习 5（合并访存）：一个 warp 的 32 个线程读取一个行优先存储的 BF16 矩阵，每行 4096 个元素。(a) 32 个线程读同一行的 32 个相邻元素；(b) 读同一列的 32 个相邻元素。各需要几次 128 字节的内存事务？

参考答案

(a) 32 × 2 字节 = 64 字节，1 次事务（如果起始地址对齐）。(b) 相邻两行相距 8192 字节，每个元素在不同的段，需要 32 次事务。

练习 6（分块）★： $n = 4096$ 的 FP32 矩阵乘法，分块大小 $T = 128$ 。朴素方式和分块方式分别从显存读取多少字节？分块后的算术强度是多少？

参考答案

朴素： $2n^3\times4 \approx 5.5\times10^{11}$ 字节；分块： $2n^3/T\times4 \approx 4.3\times10^9$ 字节。FLOPs $= 2n^3 \approx 1.37\times10^{11}$ ，分块后强度 $\approx 1.37\times10^{11}/4.3\times10^9 = 32$ FLOP/字节（即 $T/4$ ）。这还没算上 L2 缓存和寄存器层次的复用，真实实现的有效强度更高。

练习 7（波次量化）：H100 有 132 个 SM。如果输出块是 128×128， $n$ 取多少时恰好一波做完所有块（块数不超过 132 且尽量多）？ $n$ 再增加 1 会怎样？

参考答案

$\lceil n/128\rceil^2 \le 132$ → $\lceil n/128\rceil \le 11$ （ $11^2 = 121$ ），所以 $n \le 1408$ 时一波（121 个块，利用率 92%）。 $n = 1409$ 时变成 $12^2 = 144$ 个块，需要两波，第二波只有 12 个块，利用率降到约 55%。

练习 8（在线 softmax）：用在线 softmax 算法，分两块处理 $x = [1, 3 \mid 2, 5]$ 。写出每块处理后的 $m$ 和 $\ell$ ，并验证最终的 softmax。

参考答案

第一块： $m = 3$ ， $\ell = e^{1-3} + e^{0} = e^{-2} + 1 \approx 1.1353$ 。第二块： $m' = 5$ ， $\ell' = 1.1353\times e^{3-5} + e^{2-5} + e^{0} \approx 0.1537 + 0.0498 + 1 = 1.2035$ 。softmax $(x_i) = e^{x_i - 5}/1.2035$ ，例如 $x_4 = 5$ 的概率 $\approx 0.831$ 。

练习 9（FlashAttention 反向）★★：证明 $D_i = \sum_j P_{ij}dP_{ij}$ 等于 $\sum_k dO_{ik}O_{ik}$ 。为什么这一点对 FlashAttention 很重要？

参考答案

$\sum_j P_{ij}dP_{ij} = \sum_j P_{ij}\sum_k dO_{ik}V_{jk} = \sum_k dO_{ik}\sum_j P_{ij}V_{jk} = \sum_k dO_{ik}O_{ik}$ 。这样 $D_i$ 可以用 $O(nd)$ 的数据（ $O$ 和 $dO$ ）算出来，不需要完整的 $n\times n$ 矩阵 $P$ ，反向传播才能逐块进行。

练习 10（动手）： (a) 在 05_flash_attention.py 中改变 block_q、block_k（包括不能整除 $n$ 的值），验证结果始终正确。 (b) 在 03_fusion_precision_recompute.py 中加入你自己的逐元素函数（比如 SiLU 的手写版本），测量 torch.compile 的加速比，并用”数读写次数”的方法预测它。 (c) 在 06_systolic_array.py 中改用”权重驻留”（weight-stationary）的方式：让 $B$ 预先存放在 PE 中，只让 $A$ 流动。

13. 延伸阅读

CS336 Spring 2026 Lecture 5 课件（GitHub stanford-cs336/lectures 中的 lecture_05.pdf）。
Horace He，Making Deep Learning Go Brrrr From First Principles（博客）：计算、内存、开销三种瓶颈。
Horace He，What Shapes Do Matrix Multiplications Like?（thonking.ai）：“矩阵之谜”的完整分析。
Dao et al.，FlashAttention（2022）；Dao，FlashAttention-2（2023）。
Milakov & Gimelshein，Online normalizer calculation for softmax（2018）。
Kirk & Hwu，Programming Massively Parallel Processors（教材）：CUDA 编程的系统入门；GPU MODE（原 CUDA MODE）社区的讲座。
Google DeepMind，How to Scale Your Model（jax-ml.github.io/scaling-book）：TPU 与 GPU、并行训练的”餐巾纸计算”。
NVIDIA，Matrix Multiplication Background User’s Guide：块量化与波次量化的官方说明。

14. 下一讲预告

本讲讲了”GPU 为什么快、什么时候慢”。下一讲（核函数与 Triton）动手：

怎么正确地测量一个运算的耗时？（热身、同步、CUDA 事件）为什么”不同步”会得到完全错误的结论？
怎么用性能分析器看清时间花在了哪些核函数上？
用 Triton 写 GPU 核函数：向量运算、GeLU、softmax、分块行求和、分块矩阵乘法 + ReLU 融合——把本讲的分块和融合变成真正的代码。
torch.compile 能自动做什么，不能做什么？

下一讲之前，建议你：跑通 code/lecture05/ 下的 6 个脚本，完成练习 2、3、6、8。