第 11 讲 · 高级 ML 编译：ML 超优化 · Advanced ML Compilation: ML Superoptimization (Mirage)

对应 Maps to：CMU 15-779（Fall 2025，Zhihao Jia），Lecture 11 · Fri 10/3 · Advanced ML Compilation: ML Superoptimization。 必读论文：M. Wu, X. Cheng, S. Liu, C. Shi, J. Ji, M. K. Ao, P. Velliengiri, X. Miao, O. Padon, Z. Jia, Mirage: A Multi-Level Superoptimizer for Tensor Programs（OSDI 2025；arXiv 2405.05751）。 依据说明：本讲以原课 slides 11-kernel-superoptimization.pdf 的主线组织：

• FlashAttention 能否表示成图优化；
• Mirage 的三层图表示（kernel / thread block / thread graph）；
• RMSNorm + MatMul、GQA 的例子；
• 搜索的关键挑战、抽象表达式剪枝、Schwartz–Zippel 与概率等价验证；
• 硬件定制的图（A100 vs H100 的 GPC 级运算）；
• 各种超优化机会（Norm + Linear、LoRA、Gated MLP、注意力变体）；
• Triton、TVM、TASO、Mirage 的比较。

论文数字以 OSDI 版（arXiv v3）为准，只做转述；slides 上个别数字和论文版本不同时会注明。mugraph.py（有限域、抽象表达式的规范形式、块图执行器、RMSNorm + MatMul 的生成器、代价模型）以及全部实验，都是本讲义自己设计的缩小版，不代表课堂上讲了什么。 配套代码 Code：code/lecture11/（4 个脚本 + 共享模块 mugraph.py；只需 numpy；02 号脚本约 35 秒，其余 ≤ 15 秒） 预计阅读时间：3–4 小时（含运行代码）

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

0.1 学习目标

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

1. 解释为什么 FlashAttention、“RMSNorm + MatMul 融合”这样的优化不能表示成计算图改写：它们同时需要代数变换、定制内核和调度变换。
2. 读懂一个 µGraph：
   • 内核图、块图、线程图三层；
   • 网格维、for 循环维；
   • imap / omap / fmap；
   • 输入迭代器、累加器、输出保存器。
3. 说清楚 Mirage 生成器的两个降低搜索量的办法：
   • 内核层和块层穷举，线程层用规则；
   • 抽象表达式剪枝，以及它的理论保证（定理 1）和局限（剪枝只是必要条件）。
4. 解释有限域上的概率等价验证：
   • 为什么要两个素数 q | p − 1，exp 怎样用单位根实现；
   • LAX 片段是什么，出错概率由什么决定。
5. 用代码的数字说明 µGraph 优化器为什么还要选网格大小、循环次数、布局和内存分配。
6. 比较 Triton、TVM、TASO、Mirage 在”图级 / 内核级”两个维度上各覆盖了什么。

0.2 先修与本讲范围

0.3 本讲地图

§1  图级优化够不够？FlashAttention 与 RMSNorm + MatMul      ← 代码 01
§2  µGraph：内核图、块图、线程图
§3  案例：RMSNorm + MatMul 的 µGraph                          ← 代码 01、02
§4  µGraph 生成器：混合搜索 + 抽象表达式剪枝                    ← 代码 02
§5  概率等价验证：有限域上的随机测试                           ← 代码 03
§6  µGraph 优化器：布局、调度、内存规划                         ← 代码 02（演示 4）
§7  案例：GQA、QKNorm、LoRA、Gated MLP、nTrans；硬件定制的图     ← 代码 04
§8  搜索时间与四个系统的比较
§9  论文精读：Mirage
§10 小结  §11 常见误区  §12 练习  §13 延伸阅读  §14 下一讲预告

0.4 关于 mugraph.py

代码实现了 Mirage 最核心的几样东西，但都缩小了：

• 有限域 FF：同一段”程序”既能用浮点数算，也能在有限域里算。
• 块图 BlockGraph：可以执行，也可以算出每个张量的抽象表达式。
• 生成器：为 RMSNorm + MatMul 枚举块图里的算子（循环体 ≤ 4 个、循环后 ≤ 3 个），并固定了”网格沿输出列切、循环沿归约维走”这一种 imap / fmap。
• 抽象表达式的子式判断：用一种规范形式 + 子多重集实现，覆盖了 Mirage 表 2 里和本例相关的公理，而不是像 Mirage 那样交给 Z3。

H100 上的耗时都是 roofline + 启动开销的模型估算。Mirage 论文报告的加速普遍比我们的模型大，我们没有去拟合。

1. 图级优化够不够？· Can FlashAttention be a graph optimization?

1.1 slides 的问题

第 10 讲的 TASO / PET 改写的是”由现成内核组成的计算图”。slides 第 3 页问：FlashAttention 能不能表示成 matmul、exp、add、mul、div 的某种组合？

不能。 计算图里的每个内核都从显存读输入、往显存写输出。FlashAttention 需要同时做三件事：

1. 代数变换：把 softmax 拆成在线的、可以分块合并的形式（第 5 讲）；
2. 定制内核：把 QKᵀ、softmax、PV 放进一个内核；
3. 调度变换：按块计算，中间结果留在共享内存和寄存器里。

1.2 用数字看（代码 01 演示 1）

一个头的注意力，N = 4096，d = 128，BF16：

计算图级的优化再好，N×N 的中间结果也至少要写一次、读一次（67 MB）。FlashAttention 比”图级最好”少 17 倍，于是注意力从访存受限变成了计算受限（两个矩阵乘法 8.6 GFLOP，H100 上至少 8.7 µs）。

1.3 Mirage 的目标

slides 第 4 页：同时考虑代数变换、定制 GPU 内核、调度变换，自动生成高度优化的内核。slides 列出三个好处：

• 工程量少：几千行 CUDA → 几行 Python；
• 性能更好：比现有系统快 1.1–3.5 倍（论文 OSDI 版：最多 3.3 倍）；
• 适配新模型更快：第 0 天就能支持新模型，不需要人工。

论文的一句话概括：Mirage 是第一个多层的张量程序超优化器。现有的自动化方法分两类，各自缺一块：

• 调度优化（Halide、TVM、Ansor）：固定”算法”（算什么），只搜索”调度”（怎么算）；
• 代数变换（TASO、PET、TENSAT）：改写算法，但只能用一组现成的内核。

论文还提到，Triton 写的 FlashAttention 有 700 多行代码。

2. µGraph：内核图、块图、线程图 · Multi-level graph representation

2.1 GPU 层次

µGraph 和 GPU 的执行 / 存储层次一一对应（论文图 2）：

2.2 内核图

每个节点是一个内核，每条边是在内核之间传递的张量，所有张量都在显存里（不同内核之间不能共享寄存器或共享内存）。节点可以是两种之一：

• 预定义的内核算子：cuBLAS 的矩阵乘法、cuDNN 的卷积……；
• 图定义的内核算子（graph-defined operator）：它的行为由一个更低层的块图定义。新的定制内核就是这样表示的。

2.3 块图

描述一个线程块里的计算，块里的中间张量都放在共享内存。它带着以下属性：

• 网格维（grid dims）：最多 3 维（x, y, z），决定启动多少个块。
• imap：每个输入张量怎样分给各个块。网格的每一维映射到张量的一个数据维（沿这一维均匀切开），或者映射到复制维 φ（每个块都拿完整的一份）。
• omap：各个块的输出怎样拼成整个输出。必须映射到数据维，因为不同的块要写显存里互不重叠的部分。
• for 循环维：为了让大张量放得进共享内存，并让加载和计算重叠，块图可以带一个 for 循环，循环体执行若干次。
• 输入迭代器（input iterator）：每次迭代把输入的一部分从显存搬进共享内存。它带一个 fmap：循环维映射到张量的哪一维（沿它切），或映射到 φ（每次迭代都用整个张量）。
• 累加器（for-loop accumulator）：把每次迭代的结果累加起来（求和、最大值等），存在共享内存。
• 输出保存器（output saver）：循环之后的算子在累加结果上执行，最后把结果从共享内存写回显存。

2.4 线程图

把计算范围从一个块缩小到一个线程。输入迭代器把共享内存里的张量读进寄存器，输出保存器写回共享内存。线程图里只有预定义的线程级算子。

2.5 合法性

一个 µGraph 合法，要同时满足三条（论文定义 2.1）：

1. 每个算子的输入输出符合它的规格；
2. 每层的张量放得进对应的存储（显存、共享内存、寄存器）；
3. 带 for 循环的块图或线程图里，从输入到输出的每条路径恰好经过一个输入迭代器、一个累加器、一个输出保存器。

网格维、for 循环维和 imap / omap / fmap 一起构成了”调度”的搜索空间，三层图则让代数变换可以在每一层发生——两类变换就这样统一在了一种表示里。

3. 案例：RMSNorm + MatMul · Case study

3.1 现有系统为什么用两个内核

RMSNorm 是 LLM 里常用的归一化：

$Y_{ij} = \frac{X_{ij} G_j}{\text{RMS}(X)_i}, \qquad \text{RMS}(X)_i = \sqrt{\tfrac{1}{H} \textstyle\sum_j X_{ij}^2},$

后面通常接一个矩阵乘法 O = Y · W。

两个算子都沿隐藏维 h 做归约（求平方和、矩阵乘法）。RMSNorm 必须先看完整一行才能除，所以现有的编译器把它们放进两个内核，中间结果 Y 写回显存。

3.2 Mirage 找到的 µGraph

Mirage 找到的最好的图（论文图 3b）只用一个图定义的内核，改变了三件事：

1. 代数变换：利用矩阵乘法和逐元素除法可以交换，把除以 RMS 挪到矩阵乘法之后：

   $O = \frac{(X \odot G)\, W}{\text{RMS}(X)}.$

   代码 01 演示 2 的数值检查：两种算法最大差 1.1e-14，只是舍入。

2. 调度变换：两个累加（Σ X² 和 Σ X·G·W）在同一个 for 循环里并行进行，部分和放在共享内存，不写回显存。

3. 线程图：循环后的乘常数、开方、除法是一串逐元素运算，融合成一个线程图，在寄存器里算完。

论文报告这个 µGraph 比手写的内核快 1.5 倍（A100）和 1.9 倍（H100）。slides 上写的是 2.2 倍，来自较早的版本。

3.3 代码里的估算（代码 01 演示 3）

LLaMA-2-7B 的形状 H = D = 4096，H100 上：

小 batch 时两种做法都受读 32 MB 权重 W 的限制。在我们的 roofline 模型里，融合省下的是一次内核启动（4 µs）和 Y 的一次往返，所以加速比论文小。一个可能的原因是：真实的内核在小 batch 时离 roofline 更远，而模型对两边都假设跑满。

batch 256 时，我们这个”网格只沿 d 切”的块图把所有行都放在一个块里，共享内存放不下——需要再沿 b 切一维网格。Mirage 会同时枚举网格维的映射，我们的迷你版只固定了一种。

4. µGraph 生成器：混合搜索 + 抽象表达式剪枝 · Expression-guided generator

4.1 搜索空间比 TASO 大得多

TASO / PET 只在内核层用一组现成的内核搜索；Mirage 要同时搜索内核层、块层，还有网格维、循环维、各种映射。它用两个办法控制规模：

1. 混合搜索：内核层和块层穷举不超过一定大小的图（评估时内核图 ≤ 5 个算子，块图 ≤ 11 个），线程层用规则构造（把相邻的逐元素算子融合进线程图）。理由是：显存、共享内存的访问比寄存器贵几个数量级，内核层和块层的优化对性能影响最大。
2. 抽象表达式剪枝（下面 §4.3）。

生成的过程是”逐个算子扩展前缀”（论文 Algorithm 1）：

• 每一步枚举下一个算子的类型和输入；
• 如果是图定义的内核，就枚举网格维和 for 循环维，再递归生成它的块图；
• 每加一个算子都检查形状和内存用量。

为了不重复生成同一个图，Mirage 定义了规范形式：每个算子有一个”秩”（输入张量的编号列表、算子类型），只按秩递增的顺序加算子。任何图都能通过重排算子变成规范形式，所以这不会漏掉合法的图。代码 generate_rms_graphs 也是这样做的。

4.2 鸡生蛋的问题

搜索时我们想剪掉”对目标计算没用”的前缀。例如对输入程序 X·Z + Y·Z：

• 算 X·Y 的前缀应该剪掉；
• 算 X + Y 的前缀不应该剪，因为 (X + Y)·Z 等价于目标。

可是在搜索出结果之前，怎么知道一个前缀”有没有用”？Mirage 的办法是抽象。

4.3 抽象表达式

抽象表达式把每条边上的张量函数抽象成一阶逻辑的项，忽略同一个输入张量里不同元素之间的差别（论文表 1）：

例如 X 是 k×k 的矩阵时，“对行求和”和”对列求和”的抽象表达式都是 sum(k, E(X))。但 k 要保留，它对剪枝很关键。

在抽象表达式上定义两种关系，并用 SMT 求解器（Z3）依据两组公理判断（论文表 2）：

• 等价公理 A_eq：加法、乘法的交换律和结合律；分配律；sum 对加法 / 乘法 / 除法的分配；sum(i, sum(j, x)) = sum(i·j, x)；除法的结合律；exp、sqrt 的乘法性质…… 这组公理不要求可靠（抽象表达式相同的图未必真的等价），而且故意不包含约分（div(mul(x, y), y) = x）——否则”什么都是什么的子式”，剪枝就失效了。
• 子式公理 A_sub：x 是 add(x, y)、mul(x, y)、div(x, y)、exp(x)、sqrt(x)、sum(i, x) 的子式；y 是 div(x, y) 的子式；子式关系满足自反性和传递性。

剪枝规则：先算出输入程序的抽象表达式 E_O；如果前缀 G 的抽象表达式不是 E_O（在 A_eq 下的某个等价形式）的子式，就剪掉 G。

定理 1（剪枝的保证）：如果最优的图 G 和输入程序在 A_eq 下抽象表达式相等，那么 G 一定会被生成。证明的要点是：

• 任何算子的输入的抽象表达式都是输出的子式；
• 加上传递性，G 的任何前缀都是 E(G) 的子式，也就是 E_O 的子式，所以不会被剪掉。

这里有一个折中：公理越强，定理覆盖的图越多，但能剪掉的越少。

4.4 代码里的搜索（代码 02 演示 1、2）

目标的抽象表达式（H = 8 时）：

$E_O = \Sigma_8\Big[\frac{G \cdot W \cdot X}{\sqrt{\Sigma_8[X \cdot X / 8]}}\Big]$

带剪枝（循环体 ≤ 4 个算子，循环后 ≤ 3 个）：生成了 5,260 个前缀，88% 被剪掉，用时约 17 ms，得到 3 个形状正确、没有死代码的完整块图：

• 前两个就是 §3.2 的结构，机器自己找到了：除法挪到循环之后，两个归约在同一个循环里。
• 候选 3 的抽象表达式 Σ8[W·X / sqrt(…)] 仍然是 E_O 的子式（少了一个 G 的乘积是子积），所以剪枝放过了它，验证才把它淘汰。这正好说明定理 1 的方向：剪枝保证”不剪掉有用的”，不保证”留下的都对”。正确性要靠下一节的等价验证。
• 被剪掉的典型前缀是 sqr(W)、sqr(G)：目标里 W、G 都只出现一次，W·W 不可能是子式。

不剪枝会怎样：

最大的规模下，不剪枝多生成 2,454 倍的前缀。而且 50 万个完整图都要做等价验证，抽样估计约需 8 分钟，剪枝时只需验证 3 个。块图每多允许一个算子，不剪枝的搜索就大一个数量级。

论文表 5（RMSNorm）：

• 不用抽象表达式时，块图最多 5 / 6 个算子要 768 / 19,934 秒，7 个以上超过 10 小时；
• 用抽象表达式（加多线程）时，5–11 个算子只要 11–28 秒。

4.5 线程图的构造

Mirage 不在线程层穷举，而是在内核图、块图生成之后，用规则把可以和前一个算子融合的算子融合进线程图（论文 Algorithm 1 第 32–36 行）。RMSNorm 的例子里，循环后的”乘常数 → 开方 → 除法”融合成一个线程图，中间结果不回共享内存。

5. 概率等价验证 · Probabilistic equivalence verifier

5.1 为什么不用定理证明器，也不用浮点数

• 证明（TASO 的做法）：输入输出有上百万个元素，太慢。
• 浮点随机测试：给不出保证，还有舍入误差。代码 03 演示 3：两个等价的程序浮点结果差 7.3e-17，不完全相同，只能用阈值比较；而阈值太松会放过错误的程序，太紧又会冤枉正确的。

Mirage 的做法是在有限域上随机测试，并证明对一类程序（LAX）这能给出强保证。这推广了第 10 讲 PET 的定理 2：PET 只处理线性 / 多线性算子，Mirage 加上了除法和指数。

5.2 LAX 片段

LAX µGraph（论文定义 5.1）满足两条：

1. 只含多线性算子（矩阵乘法、卷积、求和……）、除法和指数；
2. 从输入到输出的每条路径上最多一次指数。

这样输出的每个元素都可以写成

$\frac{\sum_i f_i \exp(g_i / h_i)}{\sum_i f'_i \exp(g'_i / h'_i)},$

其中 f、g、h 都是输入元素的多项式。两个 LAX 图的差也是这个形式，所以判断等价就是判断这样一个式子是不是恒为零。

Mirage 先把输入程序切成 LAX 子程序，再分别优化。RMSNorm（平方根另有处理）、注意力（不减最大值的形式）、Gated MLP（SiLU(z) = z / (1 + e^(−z))）都在 LAX 里；ReLU、max 不在。

5.3 两个有限域

指数怎么在有限域里算？Mirage 取两个素数 p、q，要求 q | p − 1（这样 Z_p 里才有 q 次单位根 ω，即 ω^q ≡ 1 且 ω ≠ 1）。每个值同时有”模 p 的部分”和”模 q 的部分”（论文表 3）：

为什么这样做是对的：

• 指数只需要知道 x mod q：因为 ω^q = 1，ω^x 只取决于 x mod q。
• 保留了 e^a · e^b = e^(a+b)：因为 ω^a · ω^b = ω^(a+b)。

这样”指数里面”的运算在 Z_q 里做，“指数外面”的在 Z_p 里做。exp 之后的值没有模 q 的部分，所以每条路径只能有一次 exp。

代码 03 演示 1（p = 1,861,303，q = 10,007，p − 1 = 186q）：

• exp(a)·exp(b) = exp(a + b)、(a/b)·b = a、(a + b)² = a² + 2ab + b² 在 1000 组随机数上全部成立；
• exp(a + b) = exp(a) + exp(b)、(a + b)² = a² + b² 不成立。

代码里的 sqrt 用 x^((p+1)/4)（p ≡ 3 mod 4 时，x 是平方数时它就是平方根，否则只是一个确定的函数）。它保持 sqrt(a)·sqrt(b) = sqrt(ab)。

5.4 保证

定理 2 把多项式恒等测试（PIT，Schwartz–Zippel 引理）推广到 LAX 的形式：

• 如果那个式子不恒为零，在 Z_p × Z_q × {q 次单位根} 里均匀随机取一点，它碰巧为零的概率至多 8dk⁴/q + q^(−1/k²)；
• d 是多项式的次数，k 是求和的项数。

定理 3：

• 等价的图一定通过验证；
• 对不等价的图，做 Ω((k² / ln q) · ln(1/δ)) 次随机测试，全部通过的概率至多 δ——可以做到任意小。

实现上的折中：Mirage 在 GPU 上做测试，用的是乘积放得进 16 位整数的最大素数 p = 227、q = 113。它们只做一次随机测试，但比较输出张量的所有元素。论文说实践中没有观察到误判，并计划对最终选出的 µGraph 再做一次有理论保证的验证。

代码 03 演示 4 用错误的注意力合并（见下）演示这个折中：

小域上单个元素的误判率和 1/p 同一个量级，比较所有元素就几乎不会误判。

5.5 注意力的几种写法（代码 03 演示 2）

一个查询、32 个 key、头维 8，FlashDecoding 式地沿 KV 切成 4 块：

“减任意 m 都相同”说明了安全 softmax 减最大值的本质：max 本身不是多项式运算，不在 LAX 里，但减去它不改变数学结果，只影响浮点数的数值稳定性。

数值稳定性是另一个问题。代码 03 演示 3 里，把 q、K 放大 40 倍（logits 最大约 1882），不减最大值直接算会出现 8 个 nan / inf，而在有限域里两个程序仍然相同。所以 Mirage 在有限域验证之外，还用浮点测试过滤掉数值误差大的 µGraph。

对于 ReLU 这类非 LAX 的程序，Mirage 另有一个基于定理证明器的验证器（类似 TASO：用户提供算子性质），论文没有展开。

6. µGraph 优化器：布局、调度、内存规划 · µGraph optimizer

验证之后，Mirage 对每个通过的 µGraph 再做三类不影响正确性的优化。放到验证之后做有两个好处：一是生成器不用区分”拓扑相同、只差布局 / 顺序 / 内存分配”的图，搜索空间小得多；二是只需要优化已经验证等价的图。

1. 张量布局：每个中间张量（内核、块、线程三层）怎样线性化。不同算子对布局有约束，例如 cuBLAS 要求输入的最内维在最后两维里；不同布局性能也不同，例如有的布局能用整块拷贝。Mirage 把布局选择写成整数线性规划（每个张量的每种布局一个 0/1 变量），用 Z3 求最优解。
2. 算子调度：块里的算子有多种拓扑序。Mirage 给每个节点标上”深度”（从输入出发的最长路径长度），按深度升序执行，只在不同深度之间插 __syncthreads()，同步次数最少。
3. 内存规划：给各层的中间张量分配偏移（动态存储分配问题），穷举找最优方案。

代码里的调度选择（代码 02 演示 4）：同一个块图结构，网格大小 nx 和循环次数 ni 怎么选。batch 16 时：

• ni 太小，每次迭代的切片太大，共享内存放不下；
• nx 太小，只有 nx 个 SM 在干活：ni = 64 时，nx = 16 要 37.3 µs，nx = 64 只要 14.1 µs（H100 有 132 个 SM）。

两个结构哪个更省共享内存，还取决于 batch：

• (X ⊙ G)·W 多存一个 b × h/ni 的中间结果，X·(G ⊙ W) 多存一个 h/ni × d/nx 的；
• 放不下的格子数，batch 16 时 11 vs 12，batch 64 时 15 vs 14。

所以结构和调度要一起考虑。

7. 案例 · Case studies

论文在 A100、H100 上用 6 个基准比较了 Mirage 和 TASO/PET、PyTorch（torch.compile + FlashAttention）、TensorRT / TensorRT-LLM、FlashAttention / FlashDecoding、Triton，基线都用了 CUDA Graph。Mirage 比最好的现有方法快最多 3.3 倍。代码 04 用有限域验证了代数变换，再用模型估算代价。

7.1 GQA（分组查询注意力）

LLaMA-3-70B 用 4 卡张量并行，每卡 16 个查询头、2 个 KV 头，每组 8 个查询头共享一对 K、V。上下文 8K。Mirage 找到的 µGraph 比 FlashAttention / FlashDecoding 等最多快 2.2 倍，论文给出两个原因：

1. 网格维：现有实现用固定的启发式，例如 TensorRT-LLM 在 batch 1 / 8 时用 (8, 2, 1) / (8, 2, 8) 的网格，用不满 A100（108 个 SM）、H100（132 个 SM）。Mirage 为每个 µGraph 搜索网格维；把最好的 µGraph 换成 TensorRT-LLM 的网格，会慢 18%。
2. 并行的维度：
   • FlashAttention 沿样本、头、查询序列并行；FlashDecoding、TensorRT-LLM 沿样本、头、KV 序列并行。这对头很多的 MHA 高效，对头少的 GQA 不是最优。
   • Mirage 在样本、KV 头、查询序列、KV 序列里自动选，并为不同场景生成不同的 µGraph，比现有的启发式最多少 7 倍显存访问。
   • slides 的说法是”消除冗余的 KV 缓存访问，更高效地使用 Tensor Core”：一组 8 个查询头叠成 8 行一起算，K、V 只读一遍，矩阵乘法有 8 行而不是 1 行。

代码 04 演示 1：一组 8 个查询头叠成 8 行一起算，和逐头算在有限域上完全相同。模型估算：

• batch 1 时 KV 只有 8 MB，放得进 L2，差别在网格：只按 16 个查询头并行时 16 个块，12.3 µs；沿 KV 再切 32 段之后 6.5 µs。
• batch 16 时 KV 有 134 MB，超过 L2：FlashDecoding 220–325 µs，按 KV 头分组的 µGraph 44 µs。这比论文的差距大：我们对 L2 的估算很粗，只能当”流量最多差 8 倍”的上界。

7.2 QKNorm

一些新模型（如 Chameleon）在注意力前对查询和键做 LayerNorm（QKNorm），现有的注意力实现不支持，要单独启动归一化内核。Mirage 把注意力的计算重新组织，使它能和两个归一化融合进一个定制内核，避免中间结果写回显存，内核时间最多降低 1.4 倍。

7.3 LoRA

现有系统为原来的线性层和 LoRA 的两个小线性层分别启动内核。LoRA 的计算量很小，开销主要是启动。Mirage 用代数变换

$W \times X + B \times A \times X = (W \,\|\, B) \times \big(X \,\|\, (A \times X)\big)$

把三个矩阵乘法和一次加法合成一个内核里的两个块级矩阵乘法。拼接只是在共享内存里改偏移，不搬数据。论文报告快 1.1–2.4 倍；slides 写”比 Triton 快 2.3 倍”。

代码 04 演示 2：

• 有限域上验证了这个恒等式；
• 模型估算 4 个内核 29.3 µs → 1 个 14.1 µs（2.07 倍）。

第 10 讲的 TASO 能把 X·W 和 X·A 拼起来，但 B·(A·X) 依赖 A·X 的结果，中间要写回显存。只有在一个块图里同时安排两次矩阵乘法才能完全合并。

7.4 Gated MLP

现有优化器一般把两个矩阵乘法融合成一个内核（X 只读一次），但 SiLU 和逐元素乘法是另一个内核，两个矩阵乘法的输出要写回显存。Mirage 的 µGraph：

• 两个矩阵乘法在同一个块图里并行；
• SiLU 和乘法作为循环后的处理。

论文报告 A100 上快 1.5 倍，H100 上快 2.7–3.3 倍。

代码 04 演示 3：

• 有限域上验证了块图形式（SiLU(z) = z / (1 + e^(−z)) 只用一次 exp，在 LAX 里）；
• 模型估算：三个内核 32.2 µs，拼接权重的一个矩阵乘法 + 一个逐元素内核 28.1 µs，µGraph 24.1 µs。比论文小得多，原因同 §3.3。

7.5 nTrans：Mirage 不如 TensorRT 的情况

nGPT 的归一化 Transformer 把归一化和逐元素运算交错在一起，现有系统用三个内核。Mirage 把它融合成一个内核，比其他基线快，但比 TensorRT 慢（论文图 7 里只有它的 0.3–0.4 倍）。

原因是：Mirage 的图定义内核对每个张量都要从显存搬进共享内存、再写回。这改善了内存效率、方便异步流水，但对计算很轻的内核，这些搬运开销占了主导。论文计划支持”加载时绕过共享内存”。

7.6 硬件定制的图

slides 第 30–33 页展示了 Mirage 针对硬件生成不同的图：

• A100：找到的注意力 µGraph 类似专家手写的实现（FlashDecoding++ 的思路）。
• H100：引入了处理器簇（GPC）这一层——同一个簇里的 SM 可以通过分布式共享内存通信。Mirage 加入 GPC 级的 AllReduce / Broadcast / AllGather 算子，找到的注意力 µGraph 比最好的现有内核快 2.2 倍。

slides 把这叫”无处不在的超优化机会”：Norm + Linear、LoRA、Gated MLP、各种注意力变体。

7.7 端到端

把 Mirage 生成的内核接进 PyTorch（只改几行代码），在 Chameleon-7B、LLaMA-3-8B、GPT-3-7B-LoRA、nGPT-1B 上，端到端加速 0.9–1.9 倍（个别配置比原来的 PyTorch 略慢）。

8. 搜索时间与四个系统的比较 · Search time and comparison

8.1 搜索时间

论文：优化一个 LAX 程序最多要 4 小时，这是部署前的一次性开销。Mirage 实现约 3 万行 C++ / CUDA / Python：

• 内核层用 cuDNN / cuBLAS；
• 块层和线程层用 CUTLASS 和 PTX；
• 生成的 CUDA 代码即时编译，可以直接接进 PyTorch。

8.2 四个系统

slides 最后的讨论题：Triton、TVM、TASO、Mirage 的关键区别是什么？按两个维度看：

这四个系统又是同一条暗线：把专家的手工优化变成搜索空间 + 代价模型 + 正确性保证。

9. 论文精读 · Paper reading

Mirage: A Multi-Level Superoptimizer for Tensor Programs（Wu et al., OSDI 2025）

1. 问题：

• 现有的自动化张量程序优化器只做一半：调度优化器固定算法，代数优化器只能用现成的内核。
• FlashAttention 这类需要”代数变换 + 定制内核 + 调度变换”协同的优化，任何现有方法都搜不到，只能由专家手写（Triton 版的 FlashAttention 超过 700 行）。

2. 想法：

• µGraph：内核、线程块、线程三层统一表示，网格维、for 循环维和 imap / omap / fmap 构成调度空间，三层图让代数变换在每一层都能发生。
• 混合的、表达式引导的生成器：内核层和块层穷举、线程层用规则；用抽象表达式（一阶逻辑项 + 两组公理 + Z3）剪掉不可能有用的前缀，并给出不剪掉最优图的理论保证。
• 概率等价验证：对 LAX 程序，在两个有限域上随机测试；推广了多项式恒等测试，错误概率可以做到任意小。
• µGraph 优化器：验证后再优化布局（ILP）、算子顺序（按深度、最少同步）、内存分配。

3. 验证：

• 6 个 LLM 常用基准（GQA、QKNorm、RMSNorm、LoRA、GatedMLP、nTrans），A100 和 H100 上，比最好的现有方法快最多 3.3 倍；即使是被大量手工优化过的 GQA，也快到 2.2 倍。
• nTrans 上不如 TensorRT。
• 端到端（接进 PyTorch）0.9–1.9 倍。
• 抽象表达式剪枝让块图能扩展到 11 个算子，搜索时间从 >10 小时降到几十秒（RMSNorm）。

4. 从系统的角度看：Mirage 是本单元的集大成者：

• 它把第 8–9 讲的”调度搜索”和第 10 讲的”代数搜索”放进一个表示；
• 它把 PET 的有限域测试推广到除法和指数，使得像注意力、RMSNorm 这样的程序也能被快速验证；
• 它的输出是可以直接部署的 CUDA 内核。

它和第 7 讲的 MPK（Mirage Persistent Kernel）出自同一个项目：MPK 把整个模型编译进一个常驻内核，Mirage 负责生成单个子图的最优内核。

5. 局限：

• 只对 LAX 子程序有概率验证；ReLU、max 等要靠另一个基于定理证明器的验证器。
• 搜索最多要几个小时；动态形状要为每个形状重新搜索。
• 抽象表达式的公理决定了”能保证找到哪些图”：不含约分，所以可能漏掉一些等价的图。
• 图定义的内核总是经过共享内存，计算很轻的内核会吃亏（nTrans）。
• 实现里的验证用很小的域（p = 227、q = 113）和单次测试，理论保证要靠后续的最终验证。

与前几讲的联系

10. 小结 · Summary

1. FlashAttention、RMSNorm + MatMul 这样的优化需要同时做代数变换、定制内核和调度变换，计算图级的改写表达不了：注意力的图级最好仍要 71 MB 显存读写，FlashAttention 只要 4.2 MB。
2. µGraph 用内核图、块图、线程图三层统一表示，对应显存、共享内存、寄存器。块图带网格维、for 循环维、imap / omap / fmap、累加器。
3. RMSNorm + MatMul 的最优 µGraph：除法挪到矩阵乘法之后，两个归约在同一个循环里，一个定制内核（论文：1.5× A100，1.9× H100）。
4. 生成器：内核层和块层穷举、线程层用规则。抽象表达式剪枝去掉不可能有用的前缀（代码：88% 被剪，前缀少 2,454 倍）。它保证不剪掉”抽象表达式等价”的最优图，但留下的不一定对——忘了乘 G 的候选就靠验证淘汰。
5. 验证：LAX 程序在两个有限域（q | p − 1，exp = ω^x）上随机测试，没有舍入、不会溢出，错误概率可以做到任意小。数值稳定性另用浮点测试检查。
6. 优化器在验证后选布局、算子顺序、内存分配。网格大小、循环次数对性能影响很大（代码：nx = 16 vs 64，37.3 vs 14.1 µs）。

11. 常见误区 · Common misconceptions

• “FlashAttention 只是把几个算子融合起来。” 不只是融合：还需要把 softmax 改写成可以分块合并的形式（代数变换），并按块调度、中间结果留在片上（调度变换）。只融合的话，N×N 的中间结果仍要进显存。
• “抽象表达式剪枝保证留下的图都正确。” 剪枝只是必要条件：抽象表达式忽略了下标和很多细节。代码里”忘了乘 G”的候选的抽象表达式仍是目标的子式，没被剪掉，是等价验证把它淘汰的。
• “随机测试不可靠，所以 Mirage 的结果不可信。” 在有限域上、对 LAX 程序，随机测试有严格的概率保证：等价的图一定通过，不等价的图通过 t 次测试的概率随 t 指数下降。这和”用几组浮点数试一试”完全不同。
• “在有限域里验证过，就说明 GPU 上的浮点结果也对。” 有限域验证的是数学等价，不管数值稳定性。不减最大值的 softmax 在数学上和安全 softmax 等价，浮点数上却会溢出。Mirage 另做浮点测试过滤这类问题。
• “超优化器总是比手写的快。” nTrans 上 Mirage 只有 TensorRT 的 0.3–0.4 倍：它的图定义内核总是经过共享内存，对计算很轻的内核反而吃亏。
• “搜出了结构就完事了。” 同一个结构，网格大小、循环次数、布局不同，性能差几倍（代码 02 演示 4）；Mirage 的 GQA 例子里，网格维换成 TensorRT-LLM 的固定值就慢 18%。

12. 练习 · Exercises（含参考答案）

练习 1　按 Mirage 表 1 的规则，写出下面块图里每个张量的抽象表达式（循环 4 次，每次迭代 h 方向 16 个元素）：t0 = sqr(X)；t1 = sum_h(t0)；A = Accum(t1)。它是目标 E_O = Σ64[G·W·X / sqrt(Σ64[X·X / 64])] 的子式吗？

答案：

• t0 = mul(X, X)，t1 = sum(16, mul(X, X))，A = sum(4, sum(16, mul(X, X))) = sum(64, mul(X, X))。
• 它是 sqrt 里面那个表达式 div(sum(64, X·X), 64) 的子式（分子），而 sqrt 的参数又是 sqrt(·) 的子式、sqrt(·) 是除法的分母的子式。所以是 E_O 的子式，不会被剪掉。

练习 2　为什么 Mirage 的等价公理不能包含约分 div(mul(x, y), y) = x？

答案：有了约分，任何表达式 z 都等价于 div(mul(z, w), w)，而 w 是它的子式（除数是子式）。于是任何 w 都是任何 z 的子式，剪枝条件永远满足，一个前缀也剪不掉。公理越强，定理 1 覆盖的图越多，但剪枝越弱——这是论文明确说的折中。

练习 3　（小有限域上的 exp）取 p = 11、q = 5（5 | 10）。验证 ω = 2^((11−1)/5) mod 11 = 4 是 5 次单位根，并计算 exp(2)·exp(3) 和 exp(5)，它们相等吗？exp(2) + exp(3) 呢？

答案：

• 4⁵ = 1024 = 93 × 11 + 1，所以 4⁵ ≡ 1 (mod 11)，ω = 4 ≠ 1 是 5 次单位根。
• exp(2) = 4² = 16 ≡ 5，exp(3) = 4³ = 64 ≡ 9，乘积 45 ≡ 1。exp(5) = 4^(5 mod 5) = 4⁰ = 1，相等。
• exp(2) + exp(3) = 14 ≡ 3 ≠ 1。

练习 4　为什么 LAX 要求每条路径上最多一次 exp？如果程序里有 exp(exp(x))，在 §5.3 的双域方案里会发生什么？

答案：exp(x) 的结果是 Z_p 里的数（ω^(x mod q)），它没有”模 q 的部分”；而再做一次 exp 需要知道它 mod q 的值。ω^x mod p 再取模 q，和真正的 e^x 的”指数代数”没有对应关系，测试就失去了意义。代码里会报 NotLAX。数学上，两层指数让输出不再是论文公式 (2) 的形式，定理 2 不适用。

练习 5　（共享内存）按代码的代价模型，batch 16、H = D = 4096、网格 nx = 64、循环 ni = 16 时，(X ⊙ G)·W 这个块图每次迭代的输入切片有多大？加上中间结果和累加器一共多少？（输入 BF16、双缓冲 ×2；中间结果和累加器 FP32）

答案：

• 每次迭代 h 方向 256 个元素，每个块 d 方向 64 列。
• 输入：X 16 × 256 + G 256 + W 256 × 64 = 20,736 个元素 × 2 字节 × 2 = 82,944 字节 ≈ 81 KB。
• 中间结果（FP32）：sqr(X) 16 × 256、X ⊙ G 16 × 256、Σ_h 的结果 16、矩阵乘法的结果 16 × 64，共 9,232 个 × 4 = 36,928 字节 ≈ 36 KB。
• 累加器：16 × 64 + 16 = 1,040 个 × 4 = 4,160 字节 ≈ 4 KB。
• 合计约 121 KB，和代码 02 演示 3 的输出一致，放得进 H100 的 227 KB。

练习 6　（GQA 的流量）每卡 2 个 KV 头、头维 128、上下文 8192、batch 16、BF16。K、V 缓存一共多少字节？如果每个查询头都从显存把它那组的 K、V 读一遍，流量是多少？按 KV 头分组呢？H100 上各自至少要多久？

答案：

• K、V：16 × 2 × 8192 × 128 × 2（K 和 V）× 2 字节 ≈ 134 MB。
• 逐查询头读：每组 8 个查询头，流量 8 × 134 ≈ 1074 MB，至少 1074 MB / 3.35 TB/s ≈ 321 µs。
• 按 KV 头分组：134 MB，约 40 µs。
• 代码 04 的模型给出 220–325 µs vs 44 µs（多了 4 µs 启动开销，以及 L2 命中的估算）。

练习 7　证明 LoRA 的变换 W·X + B·(A·X) = [W | B] · [X ; A·X]，并说明为什么这样就能在一个内核里算完。

答案：

• 分块矩阵乘法：[W | B] · [X ; A·X] = W·X + B·(A·X)。
• 一个内核里：每个块先在共享内存里算出 A·X（只有 r 行，很小），在共享内存里把它接在 X 下面（只改偏移），然后用一个块级矩阵乘法乘 [W | B]。
• 原来 4 个内核之间的中间结果（A·X、W·X、B·A·X）都不用写回显存，也省掉了 3 次启动。

练习 8　（开放）ReLU 不在 LAX 片段里。如果要让 Mirage 优化一个含 ReLU 的 MLP（Linear → ReLU → Linear），你会怎么做？

参考思路：

1. 在 ReLU 处把程序切开，前后两段各是 LAX 子程序，分别优化——但会失去跨 ReLU 的融合。
2. 用论文提到的基于定理证明器的验证器：用户写出 ReLU 的性质（例如逐元素、对正数是恒等），像 TASO 那样证明等价。
3. 把 ReLU 当成一个”不透明”的逐元素函数，只允许它出现在两个等价程序的同一个位置上（两边对同一个值做 ReLU），这时可以把它当成一个随机函数来测试。

13. 延伸阅读 · Further reading

• Mirage 项目主页与代码：github.com/mirage-project/mirage（含 Mirage Persistent Kernel，第 7 讲）。
• J. T. Schwartz, Fast Probabilistic Algorithms for Verification of Polynomial Identities（JACM 1980）；R. Zippel, Probabilistic Algorithms for Sparse Polynomials（1979）——多项式恒等测试的来源。
• T. Dao et al., FlashAttention；Flash-Decoding for long-context inference（2023 博客）；K. Hong et al., FlashDecoding++（MLSys 2024）——Mirage 自动”重新发现”的注意力内核。
• B. Zhang, R. Sennrich, Root Mean Square Layer Normalization（NeurIPS 2019）——RMSNorm。
• J. Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints（EMNLP 2023）——分组查询注意力。

14. 下一讲预告 · Next lecture

第 3 单元（ML 编译）到此结束：从”让程序员按块写”（Triton），到”让机器搜调度”（AutoTVM / Ansor），再到”让机器改写计算图”（TASO / PET）和”三层一起搜”（Mirage），我们一直在一块 GPU 上想办法。

第 4 单元转向多块 GPU：当一块卡算不完、放不下时怎么办？第 12 讲《数据并行》从最简单的分工开始：每块卡一份完整的模型、一部分数据，算完梯度再求平均。问题马上变成通信：

• 参数服务器为什么会成为瓶颈？
• 朴素、环形、树形、蝶形 AllReduce 各要传多少数据、几步？
• 在 α-β 通信模型下，什么时候哪种算法最快？

我们会在”设备列表”上把这些算法真的实现一遍并核对结果，然后算出不同消息大小、不同卡数下的最优选择。最后看到数据并行的另一个问题：每块卡都要放一整份模型、梯度和优化器状态——这引出第 13 讲的 ZeRO。