Lecture 12 · ML 编译器（前端）：抓图、中间表示与图级优化 · ML Compilers: Frontends

对应 Maps to：Stanford CS349E / EE339 Efficient ML Inference at Scale（Autumn 2026），原课第 14 次课 · Tue Nov 10 · ML Compilers – Frontends（Fredrik Kjolstad）。 依据说明：这次课没有指定阅读材料，课程网站也没有 slides。本讲以公开文献为主线：J. Ansel et al., PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation（ASPLOS 2024，下文简称 PyTorch 2 论文；TorchDynamo、TorchInductor、动态形状、消融实验）；JAX 的追踪与 jaxpr；XLA / TVM 的图级优化；TFLite 的内存规划。内容的取舍是我定的，不代表课堂上实际讲了什么。 配套代码 Code：code/lecture12/（numpy；一个共享模块 + 5 个脚本，都在几秒内跑完） 先修：前置讲 P4（编译器流水线、SSA、融合与内存规划的入门）、第 2 讲（Transformer 的计算图）、第 11 讲（融合、FlashAttention、内核启动开销）。 预计阅读时间：3 小时（含运行代码）

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

0.1 学习目标 Learning goals

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

1. 说清 ML 编译器为什么要先”抓图”，比较追踪（tracing）、解析源码、分析字节码三种前端，以及它们遇到 Python 控制流、副作用、依赖数据值的分支时的不同行为。
2. 读懂一张 SSA 形式的计算图，说出图级 IR 的设计要点：算子的粒度、形状与类型、函数式（无副作用）。
3. 实现并解释常见的图级 pass：分解、公共子表达式消除、常量折叠、代数化简、权重折叠、模式匹配；解释为什么每个 pass 都必须保持语义、为什么 pass 的顺序重要。
4. 用”每个内核读写多少字节 + 内核个数”的模型分析算子融合：为什么”先分解、再融合”比手写的大算子更好，为什么编译器自己发明不出 FlashAttention。
5. 用生存期和贪心分配做内存规划，说出峰值由什么决定、运算顺序为什么会影响峰值。
6. 解释动态形状的三种处理方式（重新编译、按桶补齐、符号形状）、graph break 的代价，以及 CUDA Graphs 解决的是什么问题。

0.2 本讲地图 Map

§1  为什么要抓图：eager 的灵活与编译的性能
§2  抓图的三种前端：追踪、解析源码、分析字节码             ← 代码 01
§3  中间表示：SSA 的计算图
§4  图级 pass：CSE、常量折叠、化简、权重折叠、模式匹配      ← 代码 02
§5  分解：先拆小，再合大
§6  算子融合                                             ← 代码 03
§7  内存规划                                             ← 代码 04
§8  动态形状与 graph break                               ← 代码 05 演示 1–3
§9  框架开销与 CUDA Graphs                               ← 代码 05 演示 4
§10 ML 编译器的全景（2025–2026）
§11–15  小结、常见误区、练习（含答案）、延伸阅读、下一讲预告

0.3 本讲的代码

graph_lab.py 是一个几百行的迷你 ML 编译器前端：

Python 函数 ──追踪（运算符重载）──▶ 图 IR（SSA）──pass──▶ 更好的图 ──▶ 融合分区 + 代价 / 内存规划 ──▶ 生成 numpy 代码

用它追踪的一层 Llama 风格 Transformer（prefill 和带 KV cache 的 decode），和第 2 讲 transformer_lab 的实现逐位比较，误差在 $10^{-15}$ 量级；每个 pass 之后都重新比较一次。分析 Llama-3.1-8B 这样的大模型时，只追踪形状、不分配权重。

1. 为什么要抓图 · Why capture a graph

PyTorch 默认是 eager 模式：Python 每执行一行 y = torch.matmul(x, w)，就立刻发射一个 GPU 内核。好处是灵活、好调试——你写的就是普通的 Python。坏处是编译器看不到”全局”：

1. 融合不了：gelu(x @ W + b) 在 eager 模式下是 3 个内核，中间结果写回显存再读出来（第 10、11 讲）。要融合，就得知道 + b 后面跟着 gelu。
2. 每个算子都付一次开销：Python 的调用、PyTorch 的分派、CUDA 的内核启动，每个几微秒（第 11 讲 §8）。批 1 的 decode 里这些开销可以超过计算本身。
3. 做不了全局的规划：内存复用、算子的重新排序、把常量预先算好，都需要知道整个计算。

所以 ML 编译器的第一步是抓图：把一段 Python 程序变成一张计算图（中间表示，IR），然后在图上优化、生成代码。前置讲 P4 用一个表达式语言走过了编译器的全流程；这一讲把它放到真实的 ML 框架里：

前端（本讲）：Python ──抓图──▶ 图 IR ──图级 pass（融合、常量折叠、内存规划……）──▶ 一组内核
后端（第 13 讲）：每个内核 ──调度、自动调优──▶ 高效的机器码（CUDA / Triton / LLVM）

2. 抓图的三种前端 · Three ways to capture a graph

代码：01_tracing.py

2.1 追踪：用假张量跑一遍

最简单的办法：用一个假张量（tracer）代替真数据调用函数。假张量重载了 +、@、exp 等运算——每次运算不真算，而是往图里记一个节点、返回一个新的假张量。JAX 的 jax.jit、PyTorch 的 torch.fx.symbolic_trace 和 torch.jit.trace、PyTorch/XLA 的 Lazy Tensor 都是追踪的变体。

代码 01 演示 1 追踪一个小 MLP（x @ W1 → GELU → @ W2），得到 17 个节点（11 个运算 + 输入、参数、常量）：

%3   = matmul(%0, %1)  : [4×16]
%4   = const()  : [标量]                 # 0.5
%5   = mul(%4, %3)  : [4×16]
...
%13  = tanh(%12)  : [4×16]
%16  = mul(%5, %15)  : [4×16]
%17  = matmul(%16, %2)  : [4×8]
返回 %17

每个 %编号 只定义一次（SSA，前置讲 P4 §3），按执行顺序排列——和 JAX 的 jaxpr、PyTorch 的 FX graph 是同一种东西。注意 math.sqrt(2/π) 在追踪时已经被 Python 算成了一个数，进图的只是一个常量：追踪只记录张量运算，其余的 Python 代码照常执行。

这带来了追踪的全部性质（代码 01 演示 2）：

2.2 解析源码

torch.jit.script（TorchScript）走的是传统编译器的路：解析 Python 源码的语法树，做静态分析，把 Python 的一个子集翻译成 IR。它能正确处理控制流（if、while 都有对应的 IR 节点），但要求代码只用它支持的那部分 Python。PyTorch 2 论文说它在 TorchBench 上”只有约一半的时候能用”，把 HuggingFace 的模型全军覆没——因为它们返回一个 TorchScript 不认识的自定义输出类。

2.3 分析字节码：TorchDynamo

PyTorch 2 的 TorchDynamo 走了第三条路。CPython 在执行每个函数之前，会先把它编译成字节码（代码 01 演示 4：那个小 MLP 在 Python 3.13 里是 40 条字节码）。PEP 523 允许在”执行一个函数的字节码之前”插入一个钩子。TorchDynamo 用这个钩子：

1. 拿到函数的字节码，符号地执行它：遇到张量运算就记进 FX 图，遇到普通的 Python 运算（取属性、算术、列表、字典）就模拟它的效果；
2. 遇到处理不了的东西（依赖数据值的分支、不认识的 C 扩展、某些副作用），就在那里切断图（graph break），生成新的字节码：“调用编译好的图 1 → 执行原来的 Python → 调用图 2”；
3. 为每张图记下一组 guard：这张图是在什么假设下抓到的（张量的形状、dtype、设备，Python 常量的值，全局状态……论文说约 30 种）。下次调用先检查 guard，全部通过就直接运行缓存的编译结果，否则重新抓图、编译。

PyTorch 2 论文的统计（表 1、表 2）：

抓图本身的开销：TorchDynamo 占 eager 时间的 5%（推理）/ 1%（训练）；Lazy Tensors 是 38% / 90%——它每次迭代都要在 C++ 里重新记录、哈希整张图，而 TorchDynamo 只在第一次分析字节码，之后只检查 guard。

2.4 依赖数据值的分支

if x.sum() > 0: 需要算出 x.sum() 的值才知道走哪条路——图里的节点只描述”怎么算”，不能描述”看算出来的值决定下一步算什么”。前端只能三选一（代码 01 演示 3）：

第四条路是让 IR 本身有结构化的控制流：JAX 的 lax.cond / lax.while_loop / lax.scan，PyTorch 的 torch.cond / torch.while_loop。整个分支（包括判断条件）都在图里，编译器甚至可以把判断也放到 GPU 上——代价是用户要改写代码。

3. 中间表示：SSA 的计算图 · The IR

一张图 = 一组节点，每个节点 = (运算, 输入, 属性, 输出的形状和类型)，按拓扑顺序排列、每个值只定义一次。设计 IR 时的几个关键选择：

• 算子的粒度。PyTorch 的 ATen 有两千多个算子，不可能让每个后端都实现一遍。PyTorch 2 定义了更小的子集（Core ATen IR 约两百多个、Prims IR 约几十个），编译器先把高层算子分解成这些基本算子（§5）。JAX 的 jaxpr 一开始就只有几十种基本运算（lax 原语）。XLA 的 HLO / StableHLO 也是一个小而稳定的算子集。
• 形状和类型。每个节点都带输出的形状和 dtype，pass 才能做形状相关的决定（融合、内存规划）。形状可以是符号（§8）。
• 函数式（无副作用）。PyTorch 程序里有大量原地修改（x.add_(y)、x[i] = …），TorchDynamo / AOTAutograd 先做 functionalization：把原地修改改写成”产生一个新值”，图里就只剩纯函数，CSE、重排这些 pass 才安全。我们的 kv_write 就是这样一个”函数式地写一行缓存”的节点；真实系统在生成代码时再把它变回原地写（buffer donation）。
• 视图。reshape、transpose 在 PyTorch 里通常不复制数据，只是换一种下标方式。编译器把它们当成”读的时候换个下标”，不单独生成内核（§6）。

代码 01 / 02 的 IR 就是这样一个最小的设计：二十来种基本运算（逐元素、归约、矩阵乘法、视图），加三个复合运算（softmax、RMSNorm、SiLU）和三个特殊运算（kv_write、flash_attn、因果掩码生成器）。

4. 图级 pass · Graph-level passes

代码：02_graph_passes.py

pass 是一个”图 → 等价的图”的函数。代码 02 先追踪一层 Llama 风格的 Transformer（和第 2 讲的 TINY 模型同尺寸：$d = 128$、8 个 query 头、2 组 KV、头维度 16、FFN 344），和第 2 讲 transformer_lab 的实现比较：最大差 $2.2 \times 10^{-15}$。然后依次跑一条 pass 流水线：

4.1 公共子表达式消除（CSE）

值编号（value numbering）：运算、输入、属性都相同的两个节点只保留一个。这张图上它合并了两个常量（Q 和 K 的 RoPE 各自引用了同一张 cos 表、同一张 sin 表）。在自动生成的图里（分解、自动微分之后）重复的子表达式要多得多。

4.2 常量折叠与代数化简

输入全部是常量或参数的节点，可以在编译时算好。推理时权重是固定的，这类计算每次前向都做一遍是浪费。代码 02 演示 3 的小例子：

W = p["W"].transpose(1, 0) * 2.0           # 参数的变换 → 常量折叠
bias = p["b"] * 0.5                        # → 常量折叠
y = (x @ W) * 1.0 + 0.0 + bias             # ·1、+0 → 化简
z = y.transpose(1, 0).transpose(1, 0)      # 两次转置 → 化简
return -(-z)                               # 两次取负 → 化简

11 个运算节点 → 常量折叠掉 3 个 → 化简 4 处 + 死代码删除 → 只剩 matmul 和 add 两个。真实模型里常见的常量折叠：为某种内核布局预先重排权重、预计算量化参数、把 BatchNorm 折进卷积。

为什么在 Llama 层上常量折叠和化简什么也没做？因为追踪时 Python 已经把 $1/\sqrt{d}$ 这类算术算掉了，模型代码里也没有”参数的转置”。化简 pass 的主要客户是其他 pass：分解、自动微分、模式匹配之后，图里会出现大量 ·1、+0、成对的转置。

只做逐位不变的改写。我们没有把 $x / c$ 改成 $x \cdot (1/c)$：100 万个随机数里有 33% 的结果在最后一位不同（$1/3$ 本身就要舍入）。重新结合、乘倒数、近似的 exp 都要用户显式允许（前置讲 P4 §8.3 的 -ffast-math）。

4.3 权重折叠

RMSNorm 的最后一步是乘一个缩放向量 $w$；它的输出只被几个矩阵乘法使用时，可以折进权重：

$W' = \mathrm{diag}(w) W$ 在编译时算好（常量折叠）。代码 02 演示 4：注意力前的 RMSNorm 折进 Q、K、V 三个投影，MLP 前的折进 gate、up 两个，逐元素乘法少 2 个，结果只差最后几位（$2.4 \times 10^{-15}$）。量化里同样的技巧很常见：SmoothQuant 把激活的逐通道缩放转移到权重上、QuaRot 把旋转矩阵折进相邻的权重（第 3 讲）——都是”图级的代数改写 + 常量折叠”。

4.4 模式匹配与 pass 的顺序

模式匹配 pass 找到”matmul → 乘常数 → 加因果掩码 → softmax → matmul”这一串，换成一个 flash_attn 节点——像 TorchInductor 把注意力的子图换成调用 FlashAttention（第 11 讲）。这一个 pass 让图少了 2 个 matmul、1 个 softmax 和掩码。

顺序很重要：如果先分解（softmax 变成 max、sub、exp、sum、div），再做模式匹配，就一处也认不出来了。真实的编译器（XLA、Inductor）里，pass 的排序、什么时候重复运行哪些 pass，是大量工程经验的结晶。

4.5 每个 pass 都必须保持语义

上表的最后一列在每一步之后都重新和第 2 讲的实现比较。编译器测试的核心就是这种差分测试：同一个输入，优化前后（或编译与 eager）的结果比较。第 14 讲 KernelBench 检查 LLM 写的内核是否正确，用的也是这个办法——以及它的局限（容差怎么定、输入覆盖不全，第 10 讲 §11）。

5. 分解：先拆小，再合大 · Decomposition

为什么要把 softmax、RMSNorm 这种已经有手写高效内核的大算子拆成基本运算？

• 后端只要实现几十种基本运算：新硬件、新后端的门槛降低了。
• 融合的边界不再受算子定义的限制：手写的 RMSNorm 内核读 $x$、写 $\hat x$；拆开以后，“上一个算子的残差加法 + RMSNorm 的平方”可以融进同一个内核，RMSNorm 的最后一步乘法可以折进下一个矩阵乘法（§4.3）。
• 自动微分、量化、形状推断只需要处理基本运算。

代价是：拆开以后如果不能再融合回去，性能比手写的大算子还差。代码 03 演示 2 量化了这一点（§6.2）；PyTorch 2 论文的消融实验是同一个结论（§6.4）。所以”先拆小”和”再合大”必须一起出现：分解是给融合准备原料。

代码 02 的 decompose 把三个复合运算展开：

softmax(x)      → m = max(x)；e = exp(x − m)；e / sum(e)
rms_norm(x, w)  → x · rsqrt(sum(x·x)/d + ε) · w
silu(x)         → x / (1 + exp(−x))

6. 算子融合 · Operator fusion

代码：03_fusion.py

6.1 融合的规则

融合 pass 把图分成内核：每个内核读它的外部输入、写被外部用到的结果，内部的中间值留在寄存器 / 共享内存里。代码里的规则是 TorchInductor 风格的简化版（graph_lab.fuse）：

1. 视图和掩码生成器不单独成内核：消费者按新的下标读（零字节）。
2. 矩阵乘法、注意力各自是一个内核；对矩阵乘法输出（或它的视图）做的逐元素运算可以并进来作为 epilogue（偏置、残差、激活、RoPE），在累加器写回之前做。
3. 逐元素运算并进生产者所在的内核，只要迭代空间相同（同一个形状）。
4. 沿最后一维的归约并进生产者的逐元素内核（输入融合）；之后用到”每行一个数”的逐元素运算也可以并进来，整行留在片上——这样的内核叫 rowwise（Triton 的融合 softmax、RMSNorm 都是这样，第 10 讲 §10.3）。
5. 合并前检查依赖：如果节点的另一个输入来自一个依赖于目标内核的内核，合并会形成环，不能合并。

代码 03 演示 1：一个小层（$d = 128$、$T = 16$）不融合时 45 个内核，融合后 10 个——一个 rowwise 的 RMSNorm、Q / K 投影（带 RoPE 的 epilogue）、V 投影、FlashAttention、O 投影 + 残差、第二个 RMSNorm、gate + SiLU、up + 逐元素乘、down + 残差。

6.2 Llama-3.1-8B 的一层：prefill

用只看形状的追踪分析真实尺寸（$d = 4096$、32 / 8 个头、FFN 14336），prefill 2048 个 token，H100 的 roofline（BF16 矩阵乘法 989 TFLOPS、其他运算 67 TFLOPS、3.35 TB/s）加每个内核 4 µs 的固定开销：

• 只分解不融合，比 eager 还慢：手写的 softmax、RMSNorm 被拆成五六个内核，每个都读写一遍整张量。
• 融合把读写从 5.8 GB 降到 2.1 GB，省在逐元素和 RMSNorm 的中间结果上。
• 但注意力的分数矩阵（32 个头 × 2048 × 2048，268 MB）仍要写出、读回、再写出、再读回：softmax 是按行归约，规则没法把它和前后两个矩阵乘法融成一个内核。编译器自己发明不出 FlashAttention——它需要在线 softmax 这种算法上的改写（第 11 讲 §3）。所以 Inductor、XLA 都是把注意力匹配出来，交给手写的内核（§4.4）。换上以后读写降到 1.0 GB，剩下的基本是矩阵乘法本身。

6.3 同一层的 decode

KV 长度 4096、1 个新 token：

decode 的读写几乎全是权重（一层 436 MB）和 KV cache（17 MB）：融合几乎不减少字节。省下的是内核个数——批 1 的 decode 里每个内核只有几微秒的活，内核个数直接决定了延迟。第 11 讲的 megakernel 把这一点推到极致；§9 的 CUDA Graphs 把每个内核的 CPU 开销降下来。

6.4 PyTorch 2 论文的消融实验

TorchInductor 在 HuggingFace 模型上（float16、A100）的几何平均加速比（论文表 4）：

（其余几行：去掉循环 / 布局重排 1.91× / 1.28×、去掉矩阵乘法模板 1.85× / 1.41×、去掉参数冻结 1.85× / 1.45×、去掉模式匹配 1.83× / 1.45×。）

去掉融合和内联以后，编译后的代码比 eager 还慢——和我们的 §6.2 是同一个现象。论文的解释：Inductor 的分解把大算子拆成了很多小的基本算子，全靠融合把它们组合回来。融合（调度时合并内核，包括”横向”合并两个互不依赖的内核）和内联（降低时把逐元素运算的计算直接复制进每个使用者）功能重叠，所以单独去掉一个只掉一部分。

7. 内存规划 · Memory planning

代码：04_memory_planning.py

有了内核的顺序，编译器就知道每个写回显存的中间结果的生存期：从产生它的内核到最后一个读它的内核。生存期不重叠的结果可以共用同一块显存。

贪心分配（greedy-by-size，TFLite、XLA 等用的一类方法）：按大小从大到小，每个结果放在”和已放好的、生存期重叠的块都不冲突”的最低偏移。下界：任何时刻同时活着的字节之和的最大值——任何分配都不可能低于它。代码 04 演示 1：一个小层融合后有 11 个中间结果，不复用 52.7 KB，贪心 30.2 KB，正好等于下界。

Llama-3.1-8B 一层的 prefill 激活（不含权重和输入）：

• 复用很重要：不融合时 T = 2048 的中间结果加起来 2.3 GB，复用后 541 MB。
• 峰值由同时活着的最大的几个决定：融合减少了要分配的结果，但 T ≥ 2048 时峰值不变——峰值是注意力的分数矩阵和 softmax 的输出。FlashAttention 让分数矩阵根本不出现在显存里（T = 8192 时 8.6 GB → 604 MB），峰值从随 $T^2$ 增长变成随 $T$ 增长。
• 运算顺序也影响峰值（代码 04 演示 3）：4 个分支各算一个 [2048, 16384] 的大矩阵再按行求和。“先把 4 个大矩阵都算出来，再求和”（广度优先）的峰值 268 MB；“一个分支算完、求和，再算下一个”（深度优先）67 MB。追踪把 Python 的执行顺序原样记成了图的顺序；编译器可以在不违反依赖的前提下重排运算来降低峰值（最优排序是 NP 难的，实际用启发式）。

一层的显存账（代码 04 演示 4）：权重每层 436 MB，和 T 无关；KV cache 随 T 线性增长，32K 时每层 134 MB；prefill 的激活是临时的（一层算完就能释放，下一层复用），有 FlashAttention 时 T = 32K 约 2.4 GB，没有时 137 GB。推理引擎在启动时按”最大的 T × 最大的批”预留激活和 KV 的显存；静态的内存规划也是 CUDA Graphs 的前提（§9）。

8. 动态形状与 graph break · Dynamic shapes and graph breaks

代码：05_dynamic_shapes.py 演示 1–3

8.1 形状在变

LLM 推理里形状天天变：decode 时 KV 的长度每一步加 1；连续批处理（第 6 讲）让批大小随请求进出变化；prefill 的长度取决于提示。为一个形状编译的图，换一个形状不能直接用。三种办法（代码 05 演示 1：decode 512 步，KV 长度从 101 到 612）：

• 精确形状：我们的小编译器每次约 0.4 ms（本机），真实的编译器（生成并编译 Triton 内核、可能还要自动调优）每次几秒到几分钟——每一步都重新编译不可行。
• 按桶补齐：长度向上取整到几个固定的档，补出来的位置用掩码屏蔽。桶越细，编译越多、浪费越少。vLLM、TensorRT-LLM 为一组批大小 / 长度预先编译（或捕获 CUDA Graph），实际的请求补齐到最近的那一档。
• 符号形状：用符号 $S$ 追踪一次，所有 $S \ge 2$ 共用一张图（代码 05 演示 2：在 $S$ = 2、7、64、500 上和专门编译的图逐位相同）。形状推断对符号照样成立：$[2, 4, 16] @ [2, 16, S] \to [2, 4, S]$，softmax 沿 $S$ 归约，再 $@ [2, S, 16] \to [2, 4, 16]$，$S$ 消掉了。代价是编译器不知道 $S$ 的具体值，按 $S$ 选分块、完全展开这类优化做不了。

TorchDynamo 默认先按具体形状编译；第二次看到同一个维度变了，才把它标成符号（“automatic dynamic shapes”）。0/1 特化（PyTorch 2 论文 §5.1）：大小为 0 或 1 的维度总是按常数处理（1 决定了广播的语义）并记一个 guard，于是符号变量可以假定 ≠ 0、≠ 1，省掉很多 guard。论文还说它们不支持维度个数（rank）变化的程序——这种模式在真实模型里很少见，支持它要对”符号的形状列表”做归纳推理，不划算。

8.2 graph break 的代价

代码 05 演示 3：迭代 $y \leftarrow \tanh(yW)$ 直到变化小于 $10^{-6}$。每一轮都要看变化的值决定是否继续，于是惰性执行每一轮切断一次图：迭代 25 次，切成 25 张各有 8 个运算的小图（外加取结果时的一张空图）。每次 graph break 都要把结果写回、回到 Python、再作为下一张图的输入读进来，融合只能在一张小图内部做。

推理里典型的数据依赖：投机解码”接受到第几个 token”、检查是否生成了结束符、MoE 里每个专家分到几个 token。它们是推理引擎和编译器之间的边界：要么在图外用 Python 处理，要么写成专门的算子（第 6 讲的向量化接受、MoE 的分组矩阵乘法），要么用结构化的控制流。

9. 框架开销与 CUDA Graphs · Framework overhead and CUDA Graphs

代码：05_dynamic_shapes.py 演示 4（本机实测）

在 CPU 上跑一层 decode（$d = 128$、KV 长度 64），矩阵乘法只有约 0.35 MFLOP，按 numpy 每秒约 50 GFLOP 估计只要约 7 µs。实测（作者机器上一次运行）：

时间几乎全是开销：每个运算一次 Python 调用 + numpy 的分派 + 分配输出，约 0.8 µs。解释执行比生成的代码慢（每个节点多一次字典查找和分派）——这就是编译器最后要”生成代码”的原因。

GPU 上的对应物是每个内核几微秒的 CPU 端开销：PyTorch eager 的算子分派、CUDA 的内核启动。CUDA Graphs 把整串内核的启动在第一次运行时”录下来”，以后一次 API 调用重放整张图，CPU 不再逐个发射内核。前提是形状和显存地址都不变——所以：

• vLLM 为一组固定的批大小各捕获一张 CUDA Graph，实际的批向上补齐到最近的那一档（§8.1 的桶）；
• 显存必须事先规划好、地址固定（§7）；
• 动态形状或数据依赖的分支一出现，就退回 eager。PyTorch 2 论文说 Inductor 只在满足安全条件时才用 CUDA Graphs（遇到动态形状、非 CUDA 张量等自动关闭），消融实验里它在推理上贡献约 0.1×。

CUDA Graphs 去掉的是 CPU 端的开销；GPU 上内核之间的依赖等待、序幕和尾部（第 11 讲 §8.1）还在——那是 megakernel 要解决的问题。

10. ML 编译器的全景（2025–2026）· The landscape

推理引擎里的编译器用法：vLLM 和 SGLang 用 torch.compile 处理模型的大部分（融合逐元素、RMSNorm 等），注意力、MoE、量化矩阵乘法用手写的内核（FlashAttention / FlashInfer、Triton / CUTLASS 内核），整体再用 CUDA Graphs 捕获——编译器和手写内核分工合作，正是 §6.2 的结论。

11. 小结 · Summary

1. ML 编译器先抓图：eager 模式看不到全局，融合不了、每个算子都付一次开销、做不了全局规划。
2. 三种前端：追踪（简单通用；展开循环、特化 Python 分支、写死形状、丢掉副作用）、解析源码（只支持 Python 的子集）、分析字节码（TorchDynamo：符号执行字节码 + guard + graph break，抓图成功率最高、开销最小）。依赖张量值的分支只能报错、悄悄特化、切断图，或改写成结构化控制流。
3. 图 IR：SSA、带形状和类型、函数式；算子的粒度是关键的设计选择（大算子分解成几十到几百个基本算子）。
4. 图级 pass：CSE、常量折叠、代数化简（只做逐位不变的改写）、权重折叠、模式匹配。每个 pass 都必须保持语义（差分测试）；pass 的顺序重要（先分解就认不出注意力）。
5. 先分解、再融合：只分解不融合比 eager 慢（我们的模型 0.72×；PyTorch 2 论文 0.80×）；融合把 Llama-8B 一层 prefill 的读写从 5.8 GB 降到 2.1 GB；FlashAttention 编译器发明不出来，要靠模式匹配换成手写内核（再降到 1.0 GB）。decode 时融合省的是内核个数。
6. 内存规划：生存期 + 贪心分配，峰值由同时活着的最大结果决定；运算顺序可以让峰值相差几倍。
7. 动态形状：重新编译、按桶补齐（编译次数 vs 浪费）、符号形状（一张图，少一些优化）；graph break 让图变小、变多、每次回到 Python；CUDA Graphs 去掉 CPU 的发射开销，但要求形状和地址固定。

12. 常见误区 · Common pitfalls

1. “torch.jit.trace 抓到的图就是这个函数”——它只是示例输入走过的那条路；换一个输入可能悄悄算错（§2.4）。
2. “追踪会记下所有的代码”——只记张量运算；Python 的循环被展开、分支被特化、副作用只在追踪时发生一次（§2.1）。
3. “分解成基本算子一定更好”——不能融合回去时，比手写的大算子还慢（§6.2、§6.4）。
4. “融合能自动发现 FlashAttention”——在线 softmax 是算法上的改写，不是融合规则能推出来的；编译器靠模式匹配换成手写内核（§6.2）。
5. “融合总能减少显存读写”——decode 时读写几乎全是权重，融合省的是内核个数（§6.3）。
6. “x / c 和 x · (1/c) 一样，编译器可以随便换”——约三分之一的结果最后一位不同；这类改写要用户允许（§4.2）。
7. “峰值显存由算子的总和决定”——由同时活着的结果决定；复用和运算顺序能让它小好几倍（§7）。
8. “符号形状没有代价”——一些依赖具体大小的优化做不了；0 和 1 还要特化（§8.1）。
9. “用了 CUDA Graphs 就没有空隙了”——它只去掉 CPU 端的发射开销，而且要求形状和地址固定（§9）。

13. 练习 · Exercises

练习 1（追踪展开循环） for _ in range(4): x = tanh(x @ W) 追踪以后图里有几个 matmul？如果改成 while x.sum() > 0: ... 呢？

答案：4 个（循环被展开）。while 看的是张量的值：jax.jit 报错，torch.jit.trace 按示例输入的迭代次数特化，TorchDynamo 每次判断都切断图。要一张图，就得用 lax.while_loop / torch.while_loop。

练习 2（为什么 TorchDynamo 的开销小） Lazy Tensors 和 TorchDynamo 都能在数据依赖处切断图。为什么前者的抓图开销（推理 38%）比后者（5%）大得多？

答案：Lazy Tensors 每一次迭代都在 C++ 里重新记录所有运算、计算图的哈希来查缓存；TorchDynamo 只在第一次分析字节码，之后每次调用只检查一组 guard（形状、类型、常量），然后直接运行编译好的代码。

练习 3（CSE） 下面的图做 CSE 之后剩几个运算节点？a = x·y；b = x·y；c = a + b；d = exp(c)；e = exp(a + b)；返回 d·e

答案：b 合并进 a；a + b 变成 a + a，和 c 相同，所以 e 的输入就是 c，e = exp(c) 和 d 相同。剩下 a = x·y、c = a + a、d = exp(c)、d·d，共 4 个。

练习 4（权重折叠省了什么） Llama-3.1-8B（$d = 4096$）每层有两处 RMSNorm 的缩放可以折进权重。每个 token、每层省掉多少次乘法？和这一层矩阵乘法的乘加次数（约 2.18 亿）比呢？

答案：$2 \times 4096 = 8192$ 次，约是矩阵乘法的 $4 \times 10^{-5}$——运算可以忽略。真正的好处是图里少了逐元素运算（不融合时少两个内核和两次整张量的读写），以及 RMSNorm 的参数不用再读。

练习 5（只分解不融合） eager 的 softmax 是一个手写内核（读 $MN$、写 $MN$）。分解成 max、sub、exp、sum、div 五个内核以后读写多少？（第 10 讲 §10.3）

答案：读 $5MN + 2M$、写 $3MN + 2M$，约是原来的 4 倍。这就是 §6.2 里”只分解不融合”比 eager 慢的原因。

练习 6（decode 的固定开销） 一层 decode 不融合时 46 个内核，每个固定开销 4 µs，roofline 时间 136 µs。32 层一个 token 要多久？融合到 12 个内核呢？

答案：$32 \times (136 + 46 \times 4) = 32 \times 320 = 10.2$ ms；$32 \times (135 + 12 \times 4) = 32 \times 183 = 5.9$ ms。融合几乎不改变字节，但延迟少了约 42%。

练习 7（内存规划的下界） 四个中间结果：A（100 MB，内核 0–2）、B（50 MB，1–3）、C（80 MB，3–4）、D（60 MB，2–4）。下界是多少？按大小从大到小贪心分配，峰值是多少？

答案：各时刻活着的：t0：A=100；t1：A+B=150；t2：A+B+D=210；t3：B+C+D=190；t4：C+D=140。下界 210 MB。贪心：A 放在 0（0–100）；C（80）和 A 不重叠，放在 0（0–80）；D（60）和 A、C 都重叠：A 占 0–100、C 占 0–80，放在 100（100–160）；B（50）和 A、D、C 都重叠：0–100 被 A 占、100–160 被 D 占，放在 160（160–210）。峰值 210 MB = 下界。

练习 8（运算顺序） 代码 04 演示 3 的分支例子，如果有 8 个分支（每个大矩阵 67 MB），广度优先和深度优先的峰值大约是多少？

答案：广度优先时 8 个大矩阵同时活着，约 $8 \times 67 = 537$ MB（融合后）；深度优先只有一个，约 67 MB。差距随分支数线性增长。

练习 9（按桶补齐的浪费） KV 长度从 1 到 256 各出现一次。按 128 的倍数补齐，多读了百分之几？按 2 的幂补齐呢？

答案：真实总长度 $\sum_{S=1}^{256} S = 32{,}896$。128 的倍数：$128 \times 128 + 128 \times 256 = 49{,}152$，多 49.4%。2 的幂：$1 + 2 + 2 \times 4 + 4 \times 8 + 8 \times 16 + 16 \times 32 + 32 \times 64 + 64 \times 128 + 128 \times 256 = 43{,}691$，多 32.8%。短序列时粗的桶浪费大，2 的幂在这里反而更好；代码 05 的场景（101 → 612）结论相反（48% vs 18%）。

练习 10（符号形状的推断） decode 时 $q$：[8, 1, 64]，$K$：[8, S, 64]，$V$：[8, S, 64]。写出 $qK^\top$、softmax、乘 $V$ 之后的形状。

答案：$K^\top$ 是 [8, 64, S]；$qK^\top$：[8, 1, S]；softmax 沿最后一维：[8, 1, S]；乘 $V$：[8, 1, S] @ [8, S, 64] = [8, 1, 64]。输出和 $S$ 无关。

练习 11（0/1 特化） 为什么大小为 1 的维度不能当成普通的符号变量？

答案：广播的语义依赖它：形状 [S, 4] 和 [1, 4] 相加，当 S = 1 时两者形状相同，S > 1 时第二个被广播。如果 S 是符号，编译器就要为”S 是否等于 1”生成两套代码或加运行时检查。把 1 当常数特化（并记 guard），符号变量就可以假定 ≠ 1，推理大大简化。

练习 12（graph break 的次数） 一个循环最多 100 轮，每轮先算一次、再根据张量的值判断是否结束，实际在第 30 轮结束。惰性执行会产生几张图？如果改成 torch.while_loop 呢？

答案：每轮一次 graph break：30 张非空的小图（外加最后取结果时可能的一张）。while_loop 的循环体和条件都在一张图里：1 张图（编译器还可能把判断放在 GPU 上，不必每轮回到 CPU）。

14. 延伸阅读 · Further reading

主线

• J. Ansel et al., PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation, ASPLOS 2024. https://pytorch.org/assets/pytorch2-2.pdf
• PyTorch 文档：torch.compile、Dynamo Overview、Dynamic Shapes、IRs（Core ATen IR 与 Prims IR）。
• R. Frostig, M. J. Johnson, C. Leary, Compiling machine learning programs via high-level tracing（SysML 2018，JAX 的设计）；JAX 文档 Understanding Jaxprs、Control flow。

更多

• T. Chen et al., TVM: An Automated End-to-End Optimizing Compiler for Deep Learning, OSDI 2018（图级优化 + 自动调优）。
• OpenXLA 文档（HLO、StableHLO、融合与缓冲区分配）。
• Y. Pisarchyk, J. Lee, Efficient Memory Management for Deep Neural Net Inference, 2020（TFLite 的内存规划：greedy-by-size 等）。
• C. Lattner et al., MLIR: Scaling Compiler Infrastructure for Domain Specific Computation, CGO 2021.
• NVIDIA 博客 Getting Started with CUDA Graphs（2019）。

本仓库

• 前置讲 P4：编译器流水线、SSA、优化 pass、融合与内存规划的入门（mini_compiler.py）。
• 第 10 讲 §10–§11：融合的 softmax、内核的测试。
• 第 11 讲：FlashAttention（编译器发明不出来的那种融合）、megakernel。

15. 下一讲预告 · Next lecture

前端把程序变成了”一组内核”——每个内核是一个融合好的循环嵌套：逐元素、按行归约、带 epilogue 的矩阵乘法。下一步是为每个内核生成高效的代码：循环怎样分块、怎样映射到线程、数据怎样经过共享内存、哪些循环展开。第 13 讲讲 ML 编译器的后端：Halide 的”算法与调度分离”（以及它最重要的一个调度决策：生产者在消费者的哪一层循环里算——局部性、并行性和冗余计算的三角取舍）、多面体模型（用整数集合描述循环、找合法又高效的变换）、TVM 和 Triton 的做法、自动调优和代价模型，以及稀疏张量的编译（TACO——本课老师 Kjolstad 的工作）。

→ 第 13 讲：ML 编译器（内核生成）