Lecture 8 · 硬件引擎：GPU 的 SM 与 Tensor Core、TPU 的脉动阵列 · Hardware Engines

对应 Maps to：Stanford CS349E / EE339 Efficient ML Inference at Scale（Autumn 2026），原课第 8 次课 · Thu Oct 15 · Hardware Engines（Priyanka Raina）。本讲义把它放在”硬件性能建模”（原第 7 次课，本讲义第 9 讲）之前：先认识硬件，再给它建模（见 SYLLABUS §0.1）。阅读材料：CUDA C++ Programming Guide 第 7 节（Hardware Implementation）；NVIDIA 技术博客 NVIDIA Hopper Architecture In-Depth（2022）；N. P. Jouppi et al., In-Datacenter Performance Analysis of a Tensor Processing Unit（ISCA 2017）。 依据说明：课程网站只公开了课表和阅读材料，没有 slides。本讲以三篇阅读材料为主线，另外用到：How to Scale Your Model 第 2 章（今天的 TPU）、FlashAttention-3 的摘要（ping-pong 调度的数字）、各代 GPU 的公开规格。补充材料是我选的，不代表课堂上实际讲了什么。 配套代码 Code：code/lecture08/（numpy；一个共享模块 + 4 个脚本，都在几秒内跑完） 先修：前置讲 P1（存储层次、Little 定律）、前置讲 P2（CUDA 编程模型、SM、warp、占用率）。本讲在 P2 的基础上往里走一层：SM 里面的矩阵引擎，以及它们怎样被喂饱。 预计阅读时间：3–4 小时（含运行代码）

0. 本讲导航 · How to read this lecture

0.1 学习目标 Learning goals

学完这一讲，你应该能够：

1. 用能量的账解释为什么今天的 AI 芯片都有”矩阵引擎”：一次乘加的能量里，取指令、取操作数比运算本身贵得多；矩阵指令和脉动阵列让每个搬进来的数被用很多次。
2. 说清楚权重驻留的脉动阵列怎样工作（数据怎样流、多少周期、利用率和批大小的关系），以及 TPU v1 的整体结构。
3. 复述 Jouppi et al. 的 roofline 分析：六个生产应用里四个受限于内存带宽；推理有严格的响应时间要求；如果换成更快的内存，TPU 还能再快约 3 倍。
4. 从”SM 数 × 每个 SM 的 Tensor Core × 每周期乘加数 × 时钟”算出一块 GPU 的峰值；说出 V100 → A100 → H100 每一代 Tensor Core 的变化。
5. 解释怎样喂饱 Tensor Core：数据块的强度、共享内存的容量、异步拷贝和 TMA、多级流水线、warp 专门化；用 Little 定律算”要有多少字节在路上”。
6. 比较 GPU、TPU 和以 SRAM 为主存的芯片的设计取舍。

0.2 本讲地图 Map

§1  为什么要打开芯片：峰值和带宽从哪来
§2  能量的账：为什么要有矩阵引擎                                        ← 代码 01 演示 4、03 演示 2
§3  脉动阵列：TPU v1 的心脏                                            ← 代码 01
§4  TPU v1：设计、测量与教训（Jouppi et al. 2017）                      ← 代码 02
§5  今天的 TPU
§6  GPU 的 SM：SIMT 与硬件多线程（CUDA 编程指南第 7 节）
§7  Tensor Core：三代的演进（Hopper 博客）                              ← 代码 03
§8  喂饱 Tensor Core：数据块、异步拷贝、TMA、流水线、warp 专门化          ← 代码 03、04
§9  GPU、TPU 与 SRAM 芯片：三种设计取舍
§10 现状（2025–2026）
§11–15  小结、常见误区、练习（含答案）、延伸阅读、下一讲预告

0.3 本讲的代码

hw_engines.py 有四样东西：一个逐周期模拟的权重驻留脉动阵列（真的做乘加、让数据一格一格地流，输出和 X @ W 逐位比较）；TPU v1 论文里三块芯片和六个应用的数据；V100 / A100 / H100 的 SM 级参数；一个”读 → 算”流水线的事件模拟。没有 GPU 也没关系：这一讲的数字都是从公开规格算出来的，或者是论文的测量。

1. 为什么要打开芯片 · Why open the chip

前七讲用到的硬件知识只有两个数：峰值算力和显存带宽（再加上容量）。roofline 告诉我们一个计算受限于哪一个，但没有回答：

• 为什么 H100 做矩阵乘法有 989 TFLOPS，做普通的 FP32 运算只有约 60 TFLOPS？
• 为什么所有 AI 芯片都把面积押在”矩阵乘法”上？脉动阵列和 Tensor Core 是什么？
• “峰值”要满足什么条件才跑得到？数据块要多大、要多少字节同时在路上？
• 为什么推理（尤其是 decode）总是受限于内存？这是芯片设计的失误，还是物理规律？

这一讲打开芯片看。主角是三篇阅读材料：CUDA 编程指南讲 GPU 怎样执行线程；Hopper 博客讲最新的 SM 和 Tensor Core；Jouppi 的 TPU 论文讲一块专为推理设计的芯片，以及它在 Google 数据中心里的真实表现。

2. 能量的账：为什么要有矩阵引擎 · The energy case for matrix engines

第 1 讲 §2 用 Horowitz 的能量表（45 nm）讲过：运算几乎免费，搬运昂贵。一次 FP16 乘加约 2 pJ，从 DRAM 读一个字节约 160 pJ，一条指令在简单处理器上的总开销约 70 pJ。把这个账算到”一次乘加”上：

（代码 01 演示 4。输入从 24 MiB 的缓冲区读，按”1 MB SRAM”计；部分和的传递、累加器的读写都没算，偏乐观。）

在 GPU 里同样的账发生在寄存器这一层（代码 03 演示 2；寄存器和指令的能量是按 45 nm 量级的假设）：

两件事让矩阵引擎省能量（第 1 讲 §2.6 的两条）：

1. 摊薄控制开销：一条矩阵指令做几千到几十万次乘加；脉动阵列里根本没有”每次乘加一条指令”。
2. 数据复用：矩阵乘法 $[M, K] \times [K, N]$ 有 $MKN$ 次乘加，却只有 $MK + KN$ 个输入——每个输入天然可以被用 $N$ 或 $M$ 次。矩阵引擎在硬件上把这种复用做到最近的地方（寄存器、PE 之间的导线）。

所以芯片面积越来越多地给了矩阵运算：H100 一个 SM 里 Tensor Core 的运算能力是 CUDA core 的 16 倍。代价是：能写成矩阵乘法的计算快，不能的慢。注意力里的 softmax、归一化、激活函数都在慢的那一边（第 1 讲 §5.5、本讲 §8.5）。

3. 脉动阵列：TPU v1 的心脏 · The systolic array

代码：01_systolic_array.py

3.1 怎样工作

TPU v1 的矩阵单元是一个 $256 \times 256$ 的脉动阵列（systolic array，H. T. Kung 1980 年代提出）：65,536 个 8 位乘加单元（PE）排成方阵，每个 PE 只和左边、上边的邻居交换数据。TPU 用的是权重驻留（weight-stationary）的方式：

• 计算 $Y = X W$，$X$ 是 $[B, R]$（$B$ 个输入向量），$W$ 是 $[R, C]$。权重先装进阵列：PE$(i, j)$ 存 $W_{ij}$。
• 输入从左边进来：第 $i$ 行在第 $b + i$ 个周期收到 $X_{bi}$（“斜着”送，skew），之后每个周期向右走一格。
• 部分和从上往下走：PE$(i, j)$ 把上面传来的部分和加上 $X_{bi} W_{ij}$，下个周期交给下面。
• 第 $j$ 列底部在第 $b + (R - 1) + j$ 个周期输出 $Y_{bj}$。

代码 01 演示 1 在一个 $3 \times 3$ 的阵列上逐周期跑 4 个输入向量（数字是 PE 这个周期在处理第几个输入，”.” 是空闲）：

周期   0    1    2    3    4    5    6    7    8
行 0  0..  10.  210  321  .32  ..3  ...  ...  ...
行 1  ...  0..  10.  210  321  .32  ..3  ...  ...
行 2  ...  ...  0..  10.  210  321  .32  ..3  ...

同一个输入向量的计算在阵列里形成一条斜着推进的”波前”（论文图 4：“a diagonal wavefront”）。输出和 $XW$ 完全相同；9 个周期、36 次有用的乘加，利用率 44%。软件看不到这些——对它来说，阵列”好像一次读了 256 个输入、瞬间更新了 256 个累加器”（论文 §2）。

为什么这样设计？论文 §2 的原话：“读一个大 SRAM 的能量比运算多得多，所以矩阵单元用脉动执行来减少对统一缓冲区的读写。“每个输入从缓冲区读一次、被一行的 256 个 PE 用；每个权重装进来一次、被 $B$ 个输入用；部分和在 PE 之间直接传。阵列内部只有相邻 PE 之间的短导线。

3.2 利用率和批大小

$B$ 个输入流过 $R \times C$ 的阵列要 $B + R + C - 1$ 个周期（填充 + 流水 + 排空），利用率是 $B / (B + R + C - 1)$（代码 01 演示 2，逐周期模拟验证了公式）：

阵列越大，需要的批越大。这和 decode 的困境是同一件事：每一步只有很少的 token，再大的矩阵引擎也忙不起来。

3.3 权重从 DRAM 来

更大的问题在权重。TPU v1 的权重存在片外 8 GiB 的 DDR3 里，带宽只有 34 GB/s；经过一个 4 块深的权重 FIFO 送进阵列；阵列里有两份权重存储（双缓冲），算一块的时候把下一块移进来（移满 256×256 要 256 个周期）。

代码 01 演示 3 的粗模型（700 MHz，权重双缓冲）：读一块 64 KiB 的权重要 1.93 µs，让 $B$ 个输入流过去要 $B / 700\text{ MHz}$。$B < 1350$ 时读权重比计算慢，阵列在等数据——这个 1350 正是 TPU roofline 的拐点（§4.1）。一个 4096×4096 的全连接层：

批小于 1350 时，时间完全由读权重决定，和批大小无关——批处理是”免费”的，这又是 decode 的那条性质（第 1 讲 §4.2）。

3.4 阵列大一点好吗：二维的内部碎片

论文 §7 举了一个例子：LSTM1 有一个 600×600 的权重矩阵。在 256×256 的阵列上要切成 $3 \times 3 = 9$ 块（补零到 768×768），在 512×512 的阵列上是 $2 \times 2 = 4$ 块（补零到 1024×1024）。代码 01 演示 3：

大阵列补的零更多，要从 DRAM 读的字节反而多。论文的性能模型显示：把阵列从 256 扩到 512，所有应用都没有变快（图 11）——“这和大页面的内部碎片类似，只是更糟，因为它是二维的”。后来的 TPU（v2–v5）每个矩阵单元反而是 128×128，靠多放几个矩阵单元提高峰值。

4. TPU v1：设计、测量与教训 · TPU v1 (Jouppi et al. 2017)

代码：02_tpu_rooflines.py

4.1 背景与设计

2013 年，Google 估算：如果每个人每天用 3 分钟语音搜索，数据中心就要翻一倍。于是启动了一个推理专用芯片的项目，目标是性价比比 GPU 高 10 倍——从设计到部署只用了 15 个月（2015 年部署）。

为了快，设计做了很多减法（论文 §2）：

• 作为 PCIe 上的协处理器，像 GPU 一样插进现有服务器；主机通过 PCIe 发 CISC 指令（约十几条，关键的五条：读主机内存、读权重、矩阵乘 / 卷积、激活、写主机内存），每条指令平均 10–20 个周期。
• 矩阵单元：256×256 的 8 位乘加阵列，峰值 92 TOPS；乘积累加到 4 MiB 的 32 位累加器（4096 个 256 元素的累加器——论文说这个数是按”拐点约 1350”向上取到 2048、再翻倍以便双缓冲选的）。
• 24 MiB 的统一缓冲区（Unified Buffer）存中间结果，软件管理，不是 cache；8 GiB 的 DDR3 存权重。
• 确定性的执行：没有 cache、分支预测、乱序执行、多线程、预取。论文的理由：这些机制提高平均性能，却不改善 99 百分位的延迟，而推理服务看重的正是尾延迟。
• 芯片面积（论文图 2）：数据缓冲区 37%、计算 30%、I/O 10%、控制只有 2%（CPU / GPU 的控制逻辑大得多，也难设计得多）。

4.2 roofline：拐点在 1350

论文把 HPC 的 roofline 模型改了两处：运算是 8 位整数的乘加；横轴是”每读 1 字节权重做几次乘加”（权重放不进片上存储，激活可以）。代码 02 演示 1：

TPU 的峰值是 K80 的 33 倍，带宽只有它的约五分之一：拐点在 1350，是 K80 的 150 倍。只有”每个权重字节被用 1350 次以上”的应用才能在 TPU 上跑到峰值。

4.3 六个应用跑得怎样

Google 数据中心 95% 的推理负载由三类网络组成：MLP（61%）、LSTM（29%）、CNN（只有 5%——论文特意指出，当年体系结构会议上的加速器论文几乎全在做 CNN）。代码 02 演示 2：

• MLP 和 LSTM 的强度就是批大小（每个权重被批里的每个输入用一次），64–200，远在拐点左边：受限于内存带宽。性能计数器（论文表 3）显示 58%–79% 的周期在等权重或移权重，阵列真正在算的只有 8%–13%。
• CNN 的权重在空间上被反复使用，强度上千：CNN0 跑到峰值的 93%。
• CNN1 是例外：一半的乘加单元没有有用的权重（有些层的特征深度太浅，填不满 256 宽的阵列），最后 4 个全连接层的强度只有 32、等权重占了 35% 的周期。roofline 按整个应用的平均强度算，看不到”一部分层受限于内存”。

4.4 推理看重的是延迟

论文表 4：MLP0 要求 99 百分位响应时间 ≤ 7 ms。在这个限制下，Haswell 和 K80 只能用 16 的批，吞吐只有它们自己最高值的 42% 和 37%；TPU 能用 200 的批、跑到最高吞吐的 80%。论文的第一个”谬误”（Fallacy）就是”数据中心的推理应用像看重延迟一样看重吞吐”——开发者一旦用上低延迟的硬件，就会把延迟要求定得更紧（LSTM1 的限制从 10 ms 改成了 7 ms）。

结果（论文表 6，相对 Haswell，含主机开销）：K80 几何平均只快 1.1 倍，TPU 快 14.5 倍（按部署比例加权 29.2 倍）；整机性能/瓦是 CPU 的 17–34 倍、GPU 的 14–16 倍（只算加速器的增量功耗时是 CPU 的 41–83 倍）。论文认为 GPU 当时”高吞吐、靠大量线程掩盖延迟”的设计不适合有严格延迟要求的推理。

4.5 如果重新设计：内存带宽是最大的杠杆

论文用一个性能模型（和硬件计数器的误差平均 8%）扫描了设计参数（图 11）。代码 02 演示 4 用两个粗模型重算（“roofline”只看上界；“计数器”按表 3 的周期分解做 Amdahl 式的估计）：

论文的结论：如果当初用 GDDR5，TPU 还能快约 3 倍（他们叫它 TPU’），拐点从 1350 移到约 250。我们的两个粗模型把论文的数字夹在中间，但在”GDDR5 再加 1.5 倍时钟”这一点上和论文不一致（论文说不再提高，我们的计数器模型说提高）——粗模型只能看趋势。

后来的 TPU（v2 起）都用了 HBM。这篇 2017 年的论文已经说出了今天 LLM 推理的核心矛盾：算力涨得快，带宽跟不上（第 1 讲 §5.3）。论文里另外几条值得记住的”谬误与陷阱”：

• 用”每秒推理数”比较硬件是误导：同一块 TPU 上 MLP1 每秒 36 万次、CNN1 每秒 4,700 次，差 75 倍——它只反映模型的大小。
• 性能计数器不能事后再加：TPU 有 106 个计数器，作者还嫌少。（第 9 讲的性能建模需要这些数据。）
• 硬件之外还有软件：同一块芯片，编译器和应用的调整还能带来不小的提升。

5. 今天的 TPU · TPUs today

TPU v1 只做推理、只支持 8 位整数。后来的各代（How to Scale Your Model 第 2 章）：

• 一块芯片有一到两个 TensorCore，每个包括：几个 MXU（矩阵单元；v2–v5 每个是 128×128，v6e Trillium 是 256×256；每 8 个周期做一次 bf16[8,128] × bf16[128,128] → f32[8,128]）、一个 VPU（向量单元，做 ReLU、逐元素运算、归约）、VMEM（片上的向量存储，v5e 有 128 MiB，软件管理的暂存区）。
• 主存换成了 HBM；芯片之间用 ICI 直连成 2D 或 3D torus（第 4 讲、前置讲 P5）。
• 执行仍然是编译器排好的：XLA 把计算切成块，安排”HBM → VMEM → MXU”的拷贝和计算重叠——和 v1 的”软件管理的缓冲区 + 确定性执行”一脉相承。

v5e 每个 MXU 在 1.5 GHz 下约 $5 \times 10^{13}$ bf16 FLOP/s，四个 MXU 合计 $2 \times 10^{14}$。第 1 讲的 hw_specs 收录了 v5e、v6e 和 TPU7x（Ironwood）的规格。

6. GPU 的 SM：SIMT 与硬件多线程 · The GPU’s SM

前置讲 P2 已经讲过 GPU 的编程模型。这里从硬件的角度补两点（CUDA 编程指南第 7 节）。

SIMT（单指令多线程）：SM 以 32 个线程（一个 warp）为单位创建、调度、执行线程。一个 warp 一次执行一条公共指令；warp 里的线程分支走了不同的路时，各条路依次执行、不在路上的线程被屏蔽（分支发散，P2 §4.3）。和 SIMD 的区别：SIMD 把向量宽度暴露给软件；SIMT 的指令描述的是单个线程的行为，程序员可以写”每个线程各做各的”的代码，硬件负责把它们打包。从 Volta 起，每个线程有自己的程序计数器和调用栈（独立线程调度），warp 内的线程可以更灵活地分开、汇合。

硬件多线程：每个 warp 的执行上下文（程序计数器、寄存器等）在它的整个生命周期里一直留在芯片上，所以切换 warp 没有代价：每个发射周期，warp 调度器挑一个”下一条指令的操作数已经就绪”的 warp 发射。这是 GPU 掩盖延迟的方法——一个 warp 在等显存，就换另一个 warp 执行。代价是寄存器堆要很大（H100 每个 SM 256 KB），能同时驻留多少 warp 受寄存器和共享内存的限制（占用率，P2 §7）。

H100 的一个 SM（Hopper 博客）分成 4 个分区，每个分区有一个 warp 调度器、32 个 FP32 单元、一个第四代 Tensor Core 等；整个 SM 有 256 KB 的 L1 / 共享内存（共享内存最多 228 KB）。H100 SXM5 有 132 个 SM、50 MB 的 L2、3.35 TB/s 的 HBM3。

TPU 论文对 GPU 的批评（§4.4）在这里能看得更清楚：GPU 的设计目标是吞吐——用大量线程掩盖延迟，平均性能高，但尾延迟不可控。TPU v1 反其道而行：单线程、确定性、编译器安排一切。今天两者在靠拢：GPU 加入了 Tensor Core、TMA、warp 专门化等更”确定”的机制（§7–8），TPU 加入了向量单元和更灵活的编程。

7. Tensor Core：三代的演进 · Tensor Cores

代码：03_tensor_cores.py 演示 1–2

7.1 峰值从哪来

Tensor Core 是 SM 里专门做小矩阵乘加（MMA）的单元，从 V100（2017）开始出现。一块 GPU 的矩阵峰值可以精确地算出来（代码 03 演示 1，FP16 稠密）：

（TFLOPS；TC = Tensor Core。）每一代每个 SM 每周期的矩阵运算能力翻一倍。Hopper 博客：“按每个 SM 算，第四代 Tensor Core 的 MMA 速率是 A100 的 2 倍（同样的数据类型），用新的 FP8 是 4 倍；加上更多的 SM 和更高的时钟，整块芯片快到约 6 倍。”

7.2 各代的变化

• V100（第一代）：每个 Tensor Core 每周期做一个 $4 \times 4 \times 4$ 的 FP16 矩阵乘加，FP32 累加。
• A100（第三代）：更大的 MMA 形状（mma.sync.m16n8k16 等，一个 warp 一条指令）；新增 TF32、BF16；2:4 结构化稀疏让吞吐翻倍（第 3 讲 §6.3）；cp.async 让数据不经过寄存器直接从显存进共享内存。
• H100（第四代）：FP8（E4M3、E5M2，第 3 讲、前置讲 P3），吞吐是 FP16 的两倍；warpgroup 级的异步 MMA 指令 wgmma（4 个 warp 一起，形状最大 $64 \times 256 \times 16$，B 直接从共享内存读）；TMA（§8.3）；线程块簇（thread block cluster）和分布式共享内存（一个簇里的 SM 可以直接读写彼此的共享内存）；Transformer Engine（软件 + 硬件按层选择 FP8 / FP16）。
• Blackwell（第五代）：FP4 / NVFP4（第 1 讲 §5.4）；累加器搬进 Tensor Core 旁边专门的 Tensor Memory（TMEM），连累加器的寄存器读写也省掉了；Blackwell Ultra 把指数运算的吞吐翻倍（第 1 讲 §5.5）。

每一代的方向都一样：一条指令做更多乘加、操作数离运算单元更近、数据搬运和计算更异步。§2 的表里，从 FFMA 到 HMMA 到 WGMMA，运算本身占能量的比例从约两成提高到近九成。

8. 喂饱 Tensor Core · Feeding the Tensor Cores

代码：03_tensor_cores.py 演示 3、04_async_pipelines.py

8.1 数据块要多大

H100 的一个 SM 每周期做 4,096 次 BF16 运算；共享内存每周期最多给 128 字节（32 个 bank × 4 字节）。所以数据块的强度至少要 32 FLOP/字节。一个 SM 计算 $b_m \times b_n$ 的输出块，每读进 $k$ 列 A 和 $k$ 行 B（$(b_m + b_n) k$ 个元素）做 $2 b_m b_n k$ 次运算，强度 $= b_m b_n / (b_m + b_n)$ FLOP/字节（BF16）。代码 03 演示 3：

• 共享内存这一层：64×64 刚好够。A100 每个 SM 的运算能力只有一半，16 FLOP/字节就够——H100 把 Tensor Core 翻倍，逼着内核用更大的块。
• 显存这一层：要 295 FLOP/字节（989 TFLOPS / 3.35 TB/s），单个 SM 的块再大也到不了。剩下的复用要靠 50 MB 的 L2（不同 SM 读同一块 A、B 时命中 L2）和线程块簇里的 TMA 多播（一次读、发给簇里的几个 SM）。
• 块越大，每一级流水线占的共享内存越多，能放的级数越少。块大小和流水线深度在抢同一块 228 KB。

第 10 讲会把这些写成真正的矩阵乘法内核。

8.2 Little 定律：要有多少字节在路上

数据从显存进共享内存要几百纳秒。要让 3.35 TB/s 的显存一直忙，任何时刻要有

（延迟 700 ns 是假设值），平均每个 SM 约 17.8 KB。代码 04 演示 1：用普通的线程读、每个线程 4 字节，每个 SM 要同时挂着约 4,400 个请求——相当于 139 个 warp 各挂一条，超过了一个 SM 最多驻留 64 个 warp 的上限；每个线程读 16 字节也要约 1,100 个请求。读回来的数据还要先进寄存器、再写进共享内存。

8.3 异步拷贝与 TMA

• cp.async（Ampere）：数据直接从显存进共享内存，不经过寄存器；线程发出请求后可以去做别的事，之后再等它完成。
• TMA（Tensor Memory Accelerator，Hopper）：一个线程发一条指令，硬件按张量的形状（多维、带步长、自动处理边界）搬一整块（比如 32 KiB），还能多播给簇里的几个 SM。Hopper 博客：“只需要很少的 CUDA 线程就能管理 H100 的全部内存带宽，其余线程可以专心计算。“配合异步事务屏障（asynchronous transaction barrier），等待”这块数据到了”不再需要所有线程一起同步。

8.4 多级流水线

一个 SM 循环处理很多块：读一块、算一块。代码 04 演示 2 模拟了这个循环（每块 32 KiB、算 128×128×64；算一块 280 ns；从 L2 传一块 280 ns；延迟 700 ns），表里是 Tensor Core 忙的比例：

• 同步地读一块、算一块，Tensor Core 只忙 22%；双缓冲只藏住一块的计算时间；要约 $(\text{延迟} + \text{传输} + \text{计算}) / \text{计算} \approx 4.5$ 级才能完全藏住延迟。
• 延迟有抖动时要更多级来吸收偶尔的慢请求——这是 Hopper 内核常用 4–8 级流水线的原因，而级数受共享内存大小的限制（§8.1）。
• 如果每块都要从显存读（没有复用），传输本身比计算慢 4.6 倍，流水线再深也只有 22%：流水线藏得住延迟，藏不住带宽。

8.5 warp 专门化

Hopper 的高性能内核（CUTLASS 3、FlashAttention-3）普遍用 warp 专门化（warp specialization）：一个 producer warp 只负责发 TMA 读请求、管理缓冲区；几个 consumer warpgroup 只负责用 wgmma 算。两类工作不在同一条指令流里互相等待，流水线才能稳定地保持很多级。

warp 专门化的另一个用处是让不同的硬件单元同时忙。注意力里每个分数要做 512 次矩阵运算（$d = 128$：$QK^\top$ 和 $PV$ 各 $2d$）和 1 次指数运算（在特殊函数单元 SFU 上）。代码 04 演示 3（整块 H100，每秒 4.5 T 次 exp）：

FlashAttention-3 的 ping-pong 调度：两个 consumer warpgroup 交替，一个做矩阵乘法时另一个做 softmax。FP8 时 Tensor Core 快一倍，exp 几乎和矩阵运算一样长，重叠更关键。FlashAttention-3 报告 BF16 前向约 740 TFLOPS（75% 利用率）、FP8 接近 1.2 PFLOPS。第 11 讲会回到 FlashAttention。

9. GPU、TPU 与 SRAM 芯片：三种设计取舍 · Three design points

三者的共同点比差别更重要：都把面积押在矩阵运算上，都在想办法让每个从远处搬来的字节被多用几次。差别在于”谁来安排数据的搬运”——硬件（GPU 的 cache、warp 调度）还是软件（TPU / SRAM 芯片的编译器）——以及主存放在哪里。第 1 讲 §4.7、第 5 讲 §8 已经算过 SRAM 芯片的账：每秒能把整个主存读几十万遍，所以 decode 极快，但一块芯片只放得下零点几到两 GB。这也是异构推理的硬件基础：prefill 放在算力强的 GPU 上，对延迟敏感的 decode 或草稿模型放在 SRAM 芯片上。

10. 现状（2025–2026）· State of the field

• NVIDIA：Blackwell / Blackwell Ultra（第 1 讲 §5）：双 die、NVFP4、TMEM、更快的 SFU；Rubin（2026）继续提高 FP4 算力和 HBM4 带宽；Rubin CPX 专做长上下文的 prefill。
• Google TPU：v5e、v6e（Trillium，256×256 的 MXU）、TPU7x（Ironwood，面向推理，大容量 HBM）。
• 其他矩阵引擎：AWS Trainium（第 9 讲会读它的 NKI 性能指南：Tensor / Vector / Scalar 几个引擎、SBUF / PSUM 两级片上存储、DMA）、AMD Instinct（Matrix Core）、Meta MTIA、各种 SRAM 芯片（Groq 3 LPX、d-Matrix Corsair、Cerebras）。
• 共同的趋势：更低的精度（FP8 → FP4，带块缩放）、更大的片上存储、更异步的数据搬运、更多的”非矩阵”运算单元（exp、归约），以及芯片之间更快的互连。内存墙依然是主角。

11. 小结 · Summary

1. 一次乘加的能量里，取指令和取操作数比运算本身贵得多；矩阵引擎靠”一条指令很多乘加”和”每个输入被用很多次”把运算的占比提高到七八成以上。
2. 权重驻留的脉动阵列：输入斜着从左边流进来，部分和往下流，$B$ 个输入要 $B + R + C - 1$ 个周期；阵列越大越需要大批。权重从 DRAM 读进来的速度决定了拐点：TPU v1 是每字节 1350 次乘加。
3. TPU v1 在 Google 的推理负载上：MLP / LSTM（占 90%）受限于内存带宽，阵列只有约一成的周期在算；它仍然比同时期的 GPU 快 13–15 倍，因为推理的延迟要求让 CPU / GPU 只能用小批。最好的改进是更快的内存。
4. GPU 的峰值 = SM 数 × 每个 SM 的 Tensor Core × 每周期乘加 × 时钟；每代 Tensor Core 每个 SM 的能力翻倍（V100 1024 → A100 2048 → H100 4096 FLOPs/周期）。
5. 喂饱 Tensor Core 要大的数据块（H100 至少 64×64 才能让共享内存跟上）、L2 和多播的复用、足够多的字节在路上（每个 SM 约 18 KB），以及多级流水线（约 4–5 级才能藏住 700 ns 的延迟）。TMA 和 warp 专门化让一个线程就能发起整块的搬运、让搬运和计算分工。
6. 流水线藏得住延迟，藏不住带宽；不同硬件单元（Tensor Core 和 SFU）可以靠 warp 专门化重叠。
7. GPU、TPU、SRAM 芯片的区别在于谁来安排数据的搬运、主存放在哪里；它们都在和同一堵内存墙作战。

12. 常见误区 · Common pitfalls

1. “峰值算力就是能跑到的速度”——峰值要求数据块够大、在路上的字节够多、流水线够深；decode 这类强度低的计算离峰值很远（§3.3、§8）。
2. “阵列越大越好”——大阵列需要更大的批，切块补零的浪费也更大（600×600 在 512×512 上比在 256×256 上读更多字节，§3.4）。
3. “TPU 赢在算力”——TPU v1 的大部分负载受限于内存带宽；它赢在能在延迟限制下用大批、能效高（§4）。
4. “GPU 的 CUDA core 和 Tensor Core 只是快慢不同”——差 16 倍，而且只有能写成矩阵乘法的计算才能用 Tensor Core；softmax、归一化在慢的那一边。
5. “流水线越深越好”——级数受共享内存限制，而且流水线只能藏延迟，藏不住带宽不足（§8.4）。
6. “用更多线程就能掩盖延迟”——Little 定律要求的在路上的请求数可能超过一个 SM 能驻留的 warp 数；TMA 这类大块搬运才是出路（§8.2）。
7. “用每秒推理数比较芯片”——它主要反映模型大小，TPU 上不同模型能差 75 倍（§4.5）。
8. “roofline 的判断对整个应用都成立”——它按平均强度算；一部分层受限于内存（CNN1 的全连接层）会被平均掉（§4.3）。

13. 练习 · Exercises

练习 1（脉动阵列的周期） 一个 4×4 的权重驻留阵列，送进 6 个输入向量。要多少个周期？利用率多少？如果是 256×256 的阵列、同样 6 个输入呢？

答案：$6 + 4 + 4 - 1 = 13$ 个周期，利用率 $6/13 = 46\%$。256×256：$6 + 511 = 517$ 个周期，利用率 $6/517 = 1.2\%$。

练习 2（拐点） TPU v1：92 TOPS、权重带宽 34 GB/s。(a) 以”每读 1 字节权重做几次乘加”计，拐点是多少？(b) 换成 180 GB/s 的 GDDR5 呢？

答案：(a) $92\text{e}12 / 2 / 34\text{e}9 \approx 1353$（TOPS 把乘和加各算一次，所以除以 2）。(b) $92\text{e}12 / 2 / 180\text{e}9 \approx 256$（论文：“from 1350 to 250”）。

练习 3（读权重 vs 计算） TPU v1 一块权重 $256 \times 256$ 字节、700 MHz、34 GB/s、权重双缓冲。批大小 $B$ 多大时计算比读权重慢？和练习 2 有什么关系？

答案：读一块 $65536 / 34\text{e}9 = 1.93\ \mu s$；算一块 $B / 700\text{e}6$。$B \ge 1.93\text{e-}6 \times 700\text{e}6 \approx 1350$。正是拐点——每块权重的每个字节被 $B$ 个输入各用一次，强度就是 $B$。

练习 4（二维的内部碎片） 一个 1000×1000 的权重矩阵，分别放在 128×128、256×256、512×512 的阵列上，各要多少块？补零后要读多少字节（int8）？

答案：128：$8 \times 8 = 64$ 块，$1024^2 = 1{,}048{,}576$ 字节；256：$4 \times 4 = 16$ 块，$1024^2$ 字节；512：$2 \times 2 = 4$ 块，$1024^2$ 字节。三者恰好一样（1000 向上取到 1024 正好是三种块大小的整数倍）——补零的浪费取决于矩阵大小和块大小的”对齐”，不是块越大越浪费（600 的例子里是）。

练习 5（MLP0 的上界） MLP0 的批大小 200（每个权重字节 200 次乘加）。TPU v1 上它的 roofline 上界是多少 TOPS？实测 12.3 TOPS 是上界的多少？

答案：$2 \times 200 \times 34\text{e}9 = 13.6$ TOPS；$12.3 / 13.6 = 90\%$。

练习 6（GPU 的峰值） H100 SXM5：132 个 SM、每个 4 个 Tensor Core、每个每周期 512 次 FP16 FMA、1.83 GHz。(a) FP16 峰值？(b) FP8 每周期的 FMA 翻倍，FP8 峰值？(c) 每个 SM 128 个 FP32 单元，FP32 峰值？

答案：(a) $132 \times 4 \times 512 \times 2 \times 1.83\text{e}9 \approx 989$ TFLOPS。(b) 约 1979 TFLOPS。(c) $132 \times 128 \times 2 \times 1.83\text{e}9 \approx 61.8$ TFLOPS（官方 67 按更高的时钟）。

练习 7（操作数复用） mma.sync.m16n8k16：A 是 16×16、B 是 16×8，做 $16 \times 8 \times 16$ 次乘加。每次乘加读多少个输入元素？A 的每个元素被用几次？和 FFMA 比少了多少倍？

答案：$(256 + 128) / 2048 = 0.1875$；A 的每个元素被用 $n = 8$ 次（B 的被用 $m = 16$ 次）；FFMA 每次乘加读 3 个，少了 16 倍。

练习 8（数据块的强度） H100 一个 SM 计算 128×128 的输出块（BF16）。(a) 强度多少 FLOP/字节？(b) 共享内存每周期 128 字节、SM 每周期 4096 FLOPs，够吗？(c) 如果所有数据都直接来自显存（295 FLOP/字节才能跑满），能跑到峰值的多少？

答案：(a) $128 \times 128 / 256 = 64$。(b) 需要 $4096 / 128 = 32$，64 ≥ 32，够。(c) $64 / 295 \approx 22\%$——要靠 L2 和多播提高复用。

练习 9（Little 定律） 显存带宽 3.35 TB/s、延迟 700 ns、132 个 SM。每个 SM 平均要有多少字节在路上？如果每个线程一次读 16 字节，要多少个请求同时挂着？

答案：$3.35\text{e}12 \times 700\text{e-}9 / 132 \approx 17.8$ KB；$17{,}770 / 16 \approx 1{,}110$ 个请求（约 35 个 warp 每个线程挂一条）。

练习 10（流水线的级数） 每块的延迟 700 ns、传输 280 ns、计算 280 ns。至少要几级流水线才能让 Tensor Core 一直忙？每级 32 KiB，共享内存要多少？

答案：$(700 + 280 + 280) / 280 = 4.5$，要 5 级；$5 \times 32 = 160$ KiB，H100 的 228 KB 放得下（代码 04 演示 2：4 级 88%、6 级 99%）。

练习 11（ping-pong 的收益） 注意力 $d = 128$，每个分数 512 次矩阵运算、1 次 exp。H100：BF16 989 TFLOPS，exp 每秒 4.5 T 次。串行时每个分数要多久？两个 warpgroup 完全重叠呢？如果 exp 的吞吐翻倍（Blackwell Ultra 的做法），串行时快多少？

答案：矩阵 $512 / 989\text{e}12 = 0.517$ ps，exp $1 / 4.5\text{e}12 = 0.222$ ps；串行 0.740 ps，重叠 0.517 ps（1.43 倍）。exp 翻倍后 exp 0.111 ps，串行 0.628 ps，比原来快 1.18 倍——硬件的办法和软件的重叠解决的是同一个问题。

练习 12（为什么 K80 只比 Haswell 快一点） K80 的峰值是 Haswell 的两倍多、带宽是三倍，可论文里推理只快 1.1 倍（几何平均）。用论文表 4 解释。

答案：推理要满足 7 ms 的 99 百分位延迟。CPU 和 GPU 的执行时间不确定（cache、多线程、动态调度），只能用 16 的小批，吞吐只有自己最高值的约四成；小批时强度低、更受限于内存和固定开销，峰值的优势发挥不出来。TPU 的执行是确定的，能用 200 的批。

14. 延伸阅读 · Further reading

阅读材料

• NVIDIA, CUDA C++ Programming Guide, §7 Hardware Implementation（SIMT 架构、硬件多线程）. https://docs.nvidia.com/cuda/cuda-c-programming-guide/
• NVIDIA 技术博客, NVIDIA Hopper Architecture In-Depth, 2022. https://developer.nvidia.com/blog/nvidia-hopper-architecture-in-depth/
• N. P. Jouppi et al., In-Datacenter Performance Analysis of a Tensor Processing Unit, ISCA 2017. https://arxiv.org/abs/1704.04760

更多

• H. T. Kung, Why Systolic Architectures?, IEEE Computer 1982——脉动阵列的经典介绍。
• J. Austin et al., How to Scale Your Model, 第 2 章（TPU）和第 12 章（GPU）. https://jax-ml.github.io/scaling-book/
• J. Shah et al., FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision, 2024. https://arxiv.org/abs/2407.08608
• N. P. Jouppi et al., TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings, ISCA 2023（之后几代 TPU 的设计）。
• NVIDIA, NVIDIA H100 Tensor Core GPU Architecture 白皮书（Hopper 博客的完整版）。

本仓库

• 前置讲 P1：存储层次、Little 定律、功耗。
• 前置讲 P2：CUDA 编程模型、合并访存、bank 冲突、占用率。
• 第 1 讲 §2（能量）、§5（Blackwell Ultra 的设计取舍）。

15. 下一讲预告 · Next lecture

这一讲看到：同一个矩阵乘法，数据怎样在存储层次之间搬、在哪一层被复用多少次，决定了能量和速度。TPU 的脉动阵列是”权重驻留”的；还有”输出驻留""行驻留”等别的方式；切块的大小、循环的顺序、哪个维度在硬件上并行展开……组合起来有成千上万种。哪种最好？要不要每种都造一块芯片试试？

第 9 讲给硬件建模：

• 把 DNN 的计算写成循环嵌套；数据流（dataflow）= 循环顺序 + 哪些循环在硬件上展开；分块 = 把循环映射到存储层次；
• 一个简单的能量模型（每一级存储的访问次数 × 每次的能量）就能比较各种设计——这正是 Interstellar（用 Halide 的调度语言描述加速器）的思路，它的结论是：数据流本身没有那么重要，存储层次和分块才是关键；
• AWS 的 NKI 性能指南：在 Trainium 上，怎样用 Tensor / Vector / Scalar 引擎和 SBUF / PSUM 写出跑满硬件的内核。

→ 第 9 讲：硬件性能建模