第 8 讲 · ML 编译器（一）：基于分块的 DSL · ML Compilers Part 1: Tile-based DSLs (Triton, TVM)

对应 Maps to：CMU 15-779（Fall 2025，Zhihao Jia），Lecture 8 · Wed 9/24 · ML Compilers Part 1: Tile-based DSL。 必读论文：P. Tillet, H. T. Kung, D. Cox, Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations（MAPL 2019）；T. Chen, T. Moreau, Z. Jiang, L. Zheng, E. Yan, M. Cowan, H. Shen, L. Wang, Y. Hu, L. Ceze, C. Guestrin, A. Krishnamurthy, TVM: An Automated End-to-End Optimizing Compiler for Deep Learning（OSDI 2018）。选读：C. Lattner et al., MLIR: Scaling Compiler Infrastructure for Domain Specific Computation（CGO 2021）。本讲还引用：Triton 官方教程（向量加法、融合 softmax、矩阵乘法、融合注意力）；TVM 官方教程 How to optimize GEMM on CPU。 依据说明：本讲以原课 slides 08-ML-compilers-part-1.pdf 的主线组织：Triton 的向量加法与 CUDA 对比；Triton 矩阵乘法、指针运算、L2 缓存复用（分组排序）；与 cuBLAS 相当的性能；TensorDescriptor 与 persistent kernel；编译栈（JIT 解析 → Triton-IR → PTX → 执行）；TVM 的张量表达式与调度原语。Triton、TVM 的性能数字引自官方教程和论文（只做转述）。minitriton 解释器、L2 / 缓存模拟器、TVM 风格的迷你调度系统和全部实验都是本讲义自己设计的，不代表课堂上讲了什么。 配套代码 Code：code/lecture08/（4 个脚本 + 共享模块 minitriton.py；只需 numpy；每个脚本 ≤ 2 秒） 预计阅读时间：3–4 小时（含运行代码）

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

0.1 学习目标

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

1. 说清楚”按线程编程”（CUDA）和”按块编程”（Triton）的分工：程序员负责什么，编译器负责什么。
2. 读写 Triton 风格的内核：program_id、arange、指针运算、掩码、tl.load / tl.store / tl.dot；写出向量加法、融合 softmax、矩阵乘法。
3. 解释矩阵乘法的分组排序（GROUP_M）为什么能提高 L2 命中率，并能算出教程里”90 个块 vs 54 个块”的例子。
4. 描述 Triton 的编译栈（JIT → Triton-IR → TritonGPU-IR → LLVM-IR → PTX），说出编译器替你做的几个关键决定（数据布局、合并访存、共享内存、软件流水线）。
5. 解释 TVM 的”计算定义与调度分离”，会用 split / reorder / tile / 改变数据布局 描述循环变换，并能说出每种变换消除了哪一类缓存未命中。
6. 比较 CUDA、Triton、TVM 三种抽象层次的取舍，理解为什么”调度空间太大”引出了下一讲的自动调优。

0.2 先修与本讲范围

0.3 本讲地图

§1  为什么需要 DSL：手写 CUDA 的代价
§2  Triton 的编程模型：程序、块、指针、掩码                  ← 代码 01（演示 1）
§3  融合的 softmax：一个程序一行                             ← 代码 01（演示 2）
§4  矩阵乘法：指针运算、K 循环、分组排序与 L2                  ← 代码 02、04
§5  新特性：persistent kernel 与张量描述符（TMA）             ← 代码 02（演示 4）
§6  编译器做了什么：编译栈与布局                              ← 代码 01（演示 3–4）
§7  TVM：计算定义与调度分离                                   ← 代码 03
§8  三种抽象的比较
§9  论文精读：Triton；TVM
§10 小结  §11 常见误区  §12 练习  §13 延伸阅读  §14 下一讲预告

0.4 关于 minitriton

本机没有 GPU，也没有 Triton。代码里的 minitriton.py 是一个用 NumPy 解释执行的”类 Triton”语言：内核代码几乎可以逐字照搬 Triton 官方教程（去掉 @triton.jit 和 tl.constexpr），在 CPU 上运行、和 NumPy 对照结果，并记录每个程序读写了显存的哪些位置。它让你检查算法（块怎么切、指针怎么算、掩码对不对、读写多少字节），不能告诉你在 GPU 上跑多快——那部分要靠真 Triton 的编译器。

1. 为什么需要 DSL · Why domain-specific languages

回顾第 3–7 讲手写的 GPU 内核：一个高性能矩阵乘法要处理线程到数据的映射（合并访存）、共享内存分块和 padding（bank 冲突）、合作加载、__syncthreads、寄存器分块、多级流水、warp 特化……每换一个算子（softmax、LayerNorm、注意力的变体）、每换一代 GPU，这些都要重新想一遍。FlashAttention-3 这样的内核要专家写几个月。

这带来两个问题：

• 生产力：新模型的新算子（新的注意力变体、新的归一化）等不及专家写内核；
• 可移植性：为 Hopper 写的内核在 Blackwell 或 AMD GPU 上未必好。

领域特定语言（DSL）的思路是：找到一个合适的抽象层次，程序员描述足够多的信息让编译器生成好代码，但不必描述所有细节。本讲的两个系统选择了不同的抽象：

• Triton：程序员按块（tile）编程，描述”每个程序实例处理哪一块数据、做什么块级运算”；块内的线程分工、共享内存、同步、流水线交给编译器。
• TVM：程序员写计算定义（数学上算什么），再用调度（schedule）描述循环怎样组织；调度可以手写，也可以自动搜索（第 9 讲）。这个”算法与调度分离”的思想来自图像处理 DSL Halide。

2. Triton 的编程模型 · The Triton programming model

2.1 向量加法：CUDA vs Triton

CUDA（第 3 讲）：

__global__ void add(const float* x, const float* y, float* out, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;     // 我是哪个线程？
    if (i < n) out[i] = x[i] + y[i];                    // 每个线程一个元素
}
add<<<cdiv(n, 256), 256>>>(x, y, out, n);

Triton（官方教程 01）：

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)                         # 我是哪个程序？
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 一整块下标（向量）
    mask = offsets < n_elements                         # 最后一块可能不满
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)
add_kernel[(cdiv(n, 1024),)](x, y, out, n, BLOCK_SIZE=1024)

区别：

• 没有 threadIdx：一个程序（program instance，大致对应 CUDA 的一个线程块）处理一整块数据，offsets 是一个向量；
• 指针运算：x_ptr + offsets 得到一”块”指针，tl.load 一次读出一个同形状的数组；
• 掩码代替 if (i < n)：写在块的层面上；
• 块大小 BLOCK_SIZE 是编译期常量（tl.constexpr），要求 2 的幂。

代码 01 演示 1：$n$ = 98,432（教程的例子）、BLOCK = 1024 → 97 个程序；最后一个程序只有 128 个有效元素，由掩码处理。CUDA 写法则是 98,432 个线程。

2.2 程序之间：和 CUDA 线程块一样独立

网格里的程序没有执行顺序的保证，彼此独立（和第 3 讲的线程块一样）。正确的 Triton 内核不能依赖程序之间的执行顺序，也没有程序之间的同步（需要时用原子操作，或拆成多个内核）。

3. 融合的 softmax：一个程序一行 · Fused softmax

Triton 官方教程 02 用 softmax 展示融合的价值。朴素实现是 5 个算子（求最大值、减去、exp、求和、除），每个都从显存读输入、写回输出：

x_max = x.max(1)            # 读 MN，写 M
z = x - x_max[:, None]      # 读 MN + M，写 MN
num = exp(z)                # 读 MN，写 MN
den = num.sum(1)            # 读 MN，写 M
out = num / den[:, None]    # 读 MN + M，写 MN

合计读 $5MN + 2M$、写 $3MN + 2M$ 个元素。融合的内核让一个程序负责一行：

row = tl.program_id(0)
col = tl.arange(0, BLOCK_SIZE)                          # BLOCK_SIZE = next_power_of_2(n_cols)
x = tl.load(in_ptr + row * stride + col, mask=col < n_cols, other=-float('inf'))
x = x - tl.max(x, axis=0)
num = tl.exp(x)
tl.store(out_ptr + row * stride + col, num / tl.sum(num, axis=0), mask=col < n_cols)

一行读一次、写一次。代码 01 演示 2（1823 × 781，教程的例子）逐元素计数：朴素版读 7,122,461 个元素（恰好 $5MN + 2M$）、写 4,274,935 个（恰好 $3MN + 2M$）；融合版读写各 $MN$ = 1,423,763 个，合计少 4.00 倍。softmax 是纯访存受限的，H100 上的估算是 13.6 µs vs 3.4 µs。

两个细节：

• 越界的位置填 -inf：不影响最大值，exp(-inf) = 0 不影响求和；
• BLOCK_SIZE 要取 2 的幂：781 列时 1024 的块有 24% 被屏蔽。行再长（例如十万列）就放不进一个块，要用第 5 讲的在线 softmax 分块处理。

4. 矩阵乘法 · Matrix multiplication

4.1 指针运算与 K 循环

Triton 官方教程 03 的矩阵乘法（代码 02 逐字照搬）：

pid = tl.program_id(0)
pid_m, pid_n = ...                                      # 程序编号 → 输出块（见 §4.2）
offs_am = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)) % M
offs_bn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)) % N
offs_k = tl.arange(0, BLOCK_K)
a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)   # BM × BK 的一块指针
b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)   # BK × BN
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for k in range(0, tl.cdiv(K, BLOCK_K)):
    a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_K, other=0.0)
    b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0)
    acc = tl.dot(a, b, acc)                             # 块矩阵乘法，FP32 累加（Tensor Core）
    a_ptrs += BLOCK_K * stride_ak                       # 指针块沿 K 前进
    b_ptrs += BLOCK_K * stride_bk
c = acc.to(tl.float16)
tl.store(c_ptrs, c, mask=c_mask)

要点：

• 指针运算：行主序张量 $X[i, j]$ 的地址是 $X + i \cdot \text{stride}_i + j \cdot \text{stride}_j$。用广播 [:, None]、[None, :] 构造二维的指针块。只改 stride 就能处理转置存放的矩阵（代码 02 演示 1 的第三种情况）。
• 边界：M、N 方向越界的行列用 % M、% N 绕回矩阵内部（读到的是有效但无用的数据），写 C 时再用掩码丢掉；K 方向越界用掩码当作 0。代码 02 在 200×300×100（都不整除块大小）上验证，相对误差 3.5e-4（FP16 输出的舍入量级）。
• 和第 4 讲手写的共享内存分块相比：这里没有共享内存、合作加载、__syncthreads、寄存器分块——tl.dot 的输入块放在哪里、怎样分给线程、怎样调用 Tensor Core，由编译器决定。

4.2 程序的顺序与 L2：分组排序

每个程序负责 C 的一个输出块，要读 A 的一个行块和 B 的一个列块（沿 K 的所有小块）。哪些程序同时在跑决定了 L2 能不能复用：同时在跑的程序如果共享 A 的行块或 B 的列块，第二次读就能命中 L2。

行优先（pid 0, 1, 2, … 依次对应 C 的第 0 行的各个块）：同时在跑的程序都在同一行，共享 A 的一个行块，但要读 B 的所有列块。

分组排序（GROUP_M）：把输出块分成每组 GROUP_M 行的组，组内按列编号：

num_pid_in_group = GROUP_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_M)
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m

代码 02 演示 2 用真实的访存记录复现教程的例子：9 × 9 个块，算前 9 个输出块要装入的不同输入块——

同样算 9 个输出块，分组让它们共享更多输入。

4.3 在真实规模上模拟 L2（代码 02 演示 3）

8192 × 8192 × 8192、FP16、块 128 × 128 × 64、H100 的 50 MB L2（能放 3,051 个 16 KB 的块），132 个程序同时在跑（每个 SM 一个），以块为单位模拟 LRU：

A、B 各读一遍只要 0.27 GB，但行优先时一波 132 个程序横跨约 2 行输出块，要用到 B 的全部 64 个列块（8,192 块），L2 放不下，每一波都要重新读 B：显存流量 4.36 GB，访存时间超过了计算时间。GROUP_M = 8 时一波程序落在 8 行 × 约 17 列的小方块里，流量降到 1.74 GB，被计算完全盖住。GROUP_M 太大（= 64）退化成列优先，和行优先对称（变成 A 放不下）。

Triton 教程报告：分组排序在 A100 上把它的矩阵乘法从 220 提到 245 TFLOP/s；整体性能”与 cuBLAS 相当”。这个例子说明了按块编程的抽象边界：块内的事交给编译器，块之间的顺序（影响 L2）仍然由程序员决定——因为它取决于全局的访问模式，编译器很难自动推断。

4.4 FlashAttention 也是一个 Triton 内核（代码 04）

第 5 讲在 gpusim 上写 FlashAttention 时要操心线程编号、合作加载、padding、__syncthreads。按块写只要 24 行（代码 04）：一个程序负责一个 Q 块，循环读 K/V 块，tl.dot、tl.max、tl.exp 都是块级操作；在线 softmax 的三行和第 5 讲的伪代码一一对应；因果掩码的”整块跳过”只是把循环的上界改成对角线。

正确性：$N$ = 256、300（不是块大小的倍数），因果 / 非因果，与 float64 参考差 ≤ 6.1e-7。读写量：$N$ = 1024、$d$ = 64、BLOCK_M = 64 时逐字节计数 8,916,992 字节，和第 5 讲的公式 io_flash 逐字节相等（K、V 各读 16 遍）。按块编程没有改变算法的 IO，只改变了写代码的方式。

块大小怎样选？$N$ = 4096、$d$ = 128、FP16：BLOCK_M 从 32 增加到 256，显存读写从 258 MiB 降到 34 MiB，但每个程序片上要放的累加器和分数块从 20 KB 增加到 192 KB，程序数从 128 降到 16。这是一个多目标的权衡——下一讲的主题。

5. 新特性：persistent kernel 与张量描述符 · Newer features

persistent kernel（第 6 讲 §7）：网格只启动和 SM 一样多的程序，每个程序循环处理多个输出块：

start = tl.program_id(0)
for tile_id in range(start, num_tiles, NUM_SMS):
    pid_m, pid_n = tile_of(tile_id, ...)    # 仍然可以用分组排序
    ...                                     # 和普通版本一样的主循环

代码 02 演示 4 在 minitriton 上验证了它的正确性（25 个输出块、4 个程序）。在 GPU 上它的好处是：少了成千上万次线程块的启动与收尾、能把上一个块的收尾和下一个块的主循环重叠、程序自己控制块的处理顺序。

张量描述符：新版 Triton 提供了基于描述符（descriptor）的加载接口：在主机端或内核里描述一个张量的形状、步长和块大小，内核里按块坐标读写。在 Hopper 上它可以直接映射到 TMA（第 6 讲 §6.1）：一条指令搬一整块，地址计算、越界处理由硬件完成。从编程模型上看，它把”指针块”换成了”块坐标”，让编译器更容易生成 TMA 指令。Triton 官方教程里有用它写的 persistent 矩阵乘法。

6. 编译器做了什么 · What the compiler does

6.1 编译栈

Python 函数（@triton.jit）
   │  JIT：解析 Python 语法树，按参数类型、constexpr 的值等特化，结果缓存起来
   ▼
Triton-IR（块级的操作：make_range、load、dot、reduce、store……）
   │  与硬件无关的优化
   ▼
TritonGPU-IR（给每个块张量指定"布局"：哪些元素在哪个线程的寄存器里、是否放在共享内存里）
   │  合并访存、软件流水线（num_stages）、选择 Tensor Core 指令、消除不必要的布局转换……
   ▼
LLVM-IR → PTX → 机器码（cubin）

代码 01 演示 3 把融合 softmax 内核追踪成一份清单——这就是编译器拿到的东西：

%0 = tt.get_program_id axis=0
%1 = tt.make_range {start=0, end=1024} : tensor<1024xi32>
%2 = tt.load x (带掩码) : tensor<1024xfloat32>
%3 = tt.reduce(max) axis=0 : 1024 -> 标量
%4 = math.exp : 1024
%5 = tt.reduce(add) axis=0 : 1024 -> 标量
%6 = tt.store y (带掩码) : tensor<1024xfloat32>

全是块级操作，形状都是 1024。编译器要自己决定：1024 个元素怎样分给线程、归约怎样在 warp 内（shuffle）和 warp 间（共享内存）完成。

6.2 一个例子：布局决定了能否合并访存

一个 1024 元素的块交给 4 个 warp（128 个线程），每个线程读 8 个元素。代码 01 演示 4 用第 3 讲的 gpusim 数了三种分法：

第一种看起来”每个线程读连续的一段”很自然，但同一条指令里相邻线程隔 32 字节，每个事务只用一半；第二种让相邻线程读相邻的 16 字节，既合并、又是向量指令。在 CUDA 里这是程序员的责任（第 4 讲 §2.2 的”threadIdx.x 当行还是当列”）；在 Triton 里由编译器的布局分析决定（TritonGPU-IR 里的”blocked layout”描述的就是每个线程拿几个连续元素、线程怎样排列）。

6.3 程序员仍然要选的参数

Triton 把块内的决定交给编译器，但留下了几个”旋钮”给程序员：块大小（BLOCK_M / N / K）、num_warps（每个程序几个 warp）、num_stages（软件流水线几级，第 6 讲）、程序的顺序（GROUP_M）。它们对性能影响很大，最优值随形状和硬件变化——Triton 提供 @triton.autotune 在一组候选配置里实测挑选（第 9 讲）。

7. TVM：计算定义与调度分离 · TVM

7.1 想法

TVM 把一个算子分成两部分：

• 计算定义（张量表达式，tensor expression）：数学上算什么。例如矩阵乘法 C = te.compute((M, N), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k))
• 调度（schedule）：循环怎样组织。一串调度原语改变循环结构，但不改变结果：

这个”算法与调度分离”的设计来自图像处理 DSL Halide。好处是：同一个计算定义可以配无数种调度，调度可以手写、也可以自动搜索；正确性由计算定义保证，调度只影响性能。

7.2 一步步优化矩阵乘法（代码 03）

代码 03 实现了一个极小的”张量表达式 + 调度”系统：调度原语改变循环列表，lower() 生成循环嵌套的代码（相当于 tvm.lower 打印的 IR）。例如 tile(i, j, 16, 16) + split(k, 4) + reorder(io, jo, ko, ki, ii, ji) 降低后是：

for i_o in range(4):
    for j_o in range(4):
        for k_o in range(16):
            for k_i in range(4):
                for i_i in range(16):
                    for j_i in range(16):
                        i = (i_o * 16 + i_i); j = (j_o * 16 + j_i); k = (k_o * 4 + k_i)
                        C[i, j] += A[i, k] * B[k, j]

按 TVM 官方 CPU GEMM 教程的顺序逐步优化 64 × 64 的矩阵乘法，执行生成的代码、和 NumPy 对照（差都是 8.9e-15），并用一个 8 KB、4 路组相联、每行 64 字节的缓存模型数未命中（同时对照同样大小的全相联缓存——两者之差就是冲突未命中）：

（三个矩阵各读一遍的最少未命中是 768 次。）每一步消除的是什么：

• 默认顺序最内层沿 k 走，B 按列跳，每次跨 256 字节（4 个缓存行）。一列 B 的 64 行只会落进 32 组里的 8 组（8 组 × 4 路 = 32 行，放不下 64 行），互相挤掉——绝大多数未命中是冲突未命中。
• 循环重排 i, k, j：最内层沿 B、C 的行连续走，一个缓存行的 16 个数都用上，冲突也消失了。
• 分块：把计算限制在 16 × 16 的小块里，数据在被换出之前被多次复用，未命中又少了 4 倍多。
• 数组打包：B 的一个 16 列小块每行相隔 256 字节，同样只映射到 8 组；把 B 重新排列成 [j/16][k][j%16] 后小块在内存里连续，均匀分布到各组，冲突未命中从 1,204 降到 20。

最后一步说明了一点：改变数据布局也是一种调度。全相联缓存看不到它的效果，真实的组相联缓存却很在意（这是代码 README”计划 vs 实际”第 1 条记录的教训：第一版用全相联缓存，看不出块内重排和打包的差别）。

7.3 搜索空间

代码 03 演示 3：只考虑”i、j、k 各选一个 split 因子”和”6 个循环的顺序”，64 × 64 时就有 $7^3 \times 6! = 246{,}960$ 种调度，1024 时 958,320 种——还没算向量化、并行、展开、数据布局、GPU 上的线程绑定。手写调度需要专家；能不能让机器搜索？这就是 TVM 论文的另一半（基于学习的代价模型）和第 9 讲的主题。

7.4 TVM 的其他部分

TVM 是一个端到端的编译器：在算子级的张量表达式之上，还有图级的优化（算子融合、数据布局变换，第 10 讲会再讨论）、自动调优（第 9 讲）、跨硬件的代码生成与运行时（CPU、GPU、移动端、FPGA 加速器）。

8. 三种抽象的比较 · Comparing abstractions

第 2 讲提到的 torch.compile 在 GPU 上的默认后端 Inductor 正是把融合后的算子生成为 Triton 内核——Triton 已经成为 PyTorch 生态里”自动生成的融合内核”的目标语言。

没有一种抽象在所有场景下都最好：越高层，越容易写、越可移植，但越难表达最新硬件的特殊能力（例如 warp 特化、TMA、乒乓调度这样的细节）。这也是为什么 FlashAttention-3 仍然用 CUTLASS 手写，而第 11 讲的 Mirage 试图用搜索跨越多个抽象层次。

9. 论文精读 · Paper reading

9.1 Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations（Tillet, Kung, Cox, MAPL 2019）

1. 问题：深度学习的新算子如果不能用厂商库（cuBLAS、cuDNN）表达，就只能由专家手写 CUDA，否则 GPU 利用率很差。已有的 DSL（例如基于多面体模型或调度语言的系统）要么表达能力不够，要么性能不及手写库。

2. 想法：以块（tile，静态形状的多维子数组）为基本单位：一门类 C 的语言 Triton-C、一个基于 LLVM 的中间表示 Triton-IR（块级的数据流和控制流），以及一组块级的编译优化——与硬件无关的（例如预取、块级的窥孔优化）和与硬件有关的（层次化分块、合并访存、共享内存的分配与同步）。程序员描述块级的计算，线程级的映射交给编译器；性能相关的参数通过自动调优选择。

3. 验证：用 Triton 写的矩阵乘法、卷积与 cuBLAS、cuDNN 性能相当，并能实现库里没有的新算子（论文以 shift 卷积为例）。

4. 从系统的角度看：Triton 找到了一个”恰好的”抽象层次：块比线程高一层，使编译器能做合并访存、共享内存、同步这些容易出错的事；块又比整个算子低一层，程序员仍然控制算法结构（分块方式、循环、块的顺序）。今天的 Triton 已经是用 Python 写内核（论文里是类 C 的语言），编译栈也迁移到了 MLIR，但”按块编程”的核心没有变。

5. 局限：块内必须是规则的张量运算；块之间的调度（顺序、L2 复用）、块大小、流水线级数仍然要程序员选；对最新硬件特性（warp 特化、TMA 等）的支持依赖编译器跟进。

9.2 TVM: An Automated End-to-End Optimizing Compiler for Deep Learning（Chen et al., OSDI 2018）

1. 问题：深度学习框架依赖厂商为每种硬件手工优化的算子库；新硬件（手机 GPU、FPGA 加速器）、新算子、新的算子组合（融合）都缺少高性能实现。

2. 想法：一个端到端的编译器：图级优化（算子融合、数据布局变换）+ 算子级优化（继承 Halide 的计算与调度分离，并新增面向深度学习硬件的调度原语，例如张量化 tensorize、显式的内存层次 cache_read、为加速器隐藏延迟的原语）+ 基于学习的代价模型（从实测数据学习预测一个调度的性能，引导搜索，减少实测次数）。

3. 验证：论文摘要称在低功耗 CPU、移动 GPU、服务器 GPU 上的性能与最先进的手工优化库相当，并能为一个基于 FPGA 的通用深度学习加速器生成代码。

4. 从系统的角度看：TVM 把”为每种硬件手写算子库”变成了”为每种硬件写一个后端 + 自动搜索调度”。它的三层结构（图、算子、调度搜索）几乎定义了后来 ML 编译器的基本形态；第 9 讲的 AutoTVM / Ansor、第 10 讲的图优化都是在这个框架里继续发展的。

5. 局限：调度模板需要人写（第 9 讲 Ansor 想解决这个问题）；搜索需要大量实测时间；对动态形状、不规则计算的支持较弱。

9.3 两篇论文之间的联系

两者都在回答”怎样让编译器生成接近手写的内核”，但选择的抽象不同：Triton 让程序员写块级的程序（命令式，接近 CUDA 的思维），编译器负责块内；TVM 让程序员写数学定义（声明式）加调度，并尝试让机器搜索调度。代码里的对比很直观：minitriton 的矩阵乘法里程序员仍然写了 K 循环和指针运算；TVM 风格的系统里程序员只写了 C[i, j] += A[i, k] * B[k, j]，所有循环结构都来自调度。前者表达更灵活、更容易写出 FlashAttention 这样的复杂内核；后者更容易自动化。

10. 小结 · Summary

1. 手写 CUDA 内核成本高、难移植；DSL 选择一个抽象层次，让程序员描述足够多的信息，其余交给编译器。
2. Triton 按块编程：程序 = 处理一个数据块；指针块、掩码、tl.load / tl.dot / tl.store；块内的线程分工、共享内存、同步由编译器决定。
3. 融合 softmax：一个程序一行，读写从 $8MN + 4M$ 降到 $2MN$（代码 01：少 4.00 倍）；行太长时要用在线 softmax。
4. 矩阵乘法：指针运算 + K 循环 + tl.dot；块之间的顺序由程序员决定，分组排序（GROUP_M）让同时在跑的程序共享输入（9 × 9 块时 90 → 54；8192³ 时显存流量 4.36 → 1.74 GB）。
5. Triton 编译栈：JIT → Triton-IR → TritonGPU-IR（布局）→ LLVM-IR → PTX；布局分析决定了访存能否合并。
6. TVM：计算定义与调度分离；split / reorder / tile / 数据布局变换都是调度；每一步消除一类缓存未命中（容量、空间局部性、冲突）。调度空间巨大，引出自动调优。

11. 常见误区 · Common misconceptions

• “Triton 的一个程序就是一个线程。” 不是。一个程序大致对应一个线程块（默认 4 个 warp），处理一整块数据；块内的元素由编译器分给各个线程。
• “用了 Triton 就不用考虑性能了。” 块大小、num_warps、num_stages、程序的顺序（L2 复用）都要选；选错了性能可以差很多（代码 02：GROUP_M 不同，显存流量差 2.5 倍）。
• “掩码只是为了防止越界。” 它也是让一份代码处理任意形状的关键：块大小必须是 2 的幂，真实的形状不是；被屏蔽的位置还要填合适的值（softmax 填 -inf，矩阵乘法填 0）。
• “TVM 的调度会改变计算结果。” 不会（浮点运算顺序变化带来的舍入差异除外）。调度只改变循环怎样组织，计算定义保证了结果。
• “循环分块只是为了让数据放进缓存。” 那是容量问题。代码 03 显示还有空间局部性（循环顺序）和冲突（数据布局）两类问题，需要不同的变换。
• “Triton 和 TVM 是竞争关系，只能选一个。” 它们处在不同的抽象层次，解决的问题有重叠也有互补；例如 TVM 的后继系统也可以生成 Triton 风格的代码，PyTorch 的编译器则直接生成 Triton。

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

练习 1　用 Triton 写一个”逐行求 L2 范数”的内核：输入 $M \times N$，输出长度 $M$。一个程序负责一行，写出 offsets、mask 和需要的 tl 调用。越界的位置应该填什么？

答案：row = tl.program_id(0); col = tl.arange(0, BLOCK); x = tl.load(ptr + row*stride + col, mask=col < N, other=0.0); tl.store(out + row, tl.sqrt(tl.sum(x*x, axis=0)))。越界填 0：平方和不受影响。（求最大值时才填 -inf。）

练习 2　融合 softmax 的读写量是 $2MN$，朴素版是 $8MN + 4M$。如果朴素版用了”把减最大值和 exp 融合在一起”的两个内核（第一个求最大值，第二个算 exp、求和、除——但第二个内核要先算完和才能除，怎么办？），读写量最少能到多少？

答案：除法需要整行的和，所以第二步要么读两遍 x（一次求和、一次算输出），要么把 exp 的结果写出去再读回。一种做法：内核 1 读 x、写 max（$MN + M$）；内核 2 读 x 和 max、算 exp 的和、写和（$MN + 2M$）；内核 3 读 x、max、和，写输出（$2MN + 2M$）——共约 $4MN$，仍是融合版的 2 倍。只有在线 softmax（第 5 讲）能把”求最大值”和”求和”合进一遍。

练习 3　对一个 $12 \times 12$ 块的矩阵乘法，算前 12 个程序在行优先和 GROUP_M = 4 时要装入的不同输入块数（K 方向也是 12 块）。

答案：行优先：A 的 1 行（12）+ B 的全部（144）= 156。GROUP_M = 4：前 12 个程序是 4 行 × 3 列，A 的 4 行（48）+ B 的 3 列（36）= 84。

练习 4　代码 02 演示 3 中 GROUP_M = 8 和 16 的结果几乎一样，而 64 退化成行优先的对称情形。直观上，最好的 GROUP_M 应该怎样随”同时在跑的程序数”变化？

答案：一波程序覆盖的区域是 GROUP_M 行 × (P / GROUP_M) 列（P = 同时在跑的程序数）。需要读的输入块数 ∝ GROUP_M + P / GROUP_M，在 GROUP_M ≈ √P 时最小（P = 132 时约 11）。8 和 16 都接近这个值；太大或太小都会让一边放不进 L2。

练习 5　解释代码 03 里”默认 i, j, k”为什么在 4 路组相联缓存里的未命中是全相联缓存的约 16 倍（266,760 vs 16,896）。如果把矩阵宽度从 64 改成 65（每行 260 字节），结果会怎样？

答案：B 的一列每行相隔 256 字节 = 4 个缓存行，行号模 32 只取 8 个值，一列 64 行挤在 8 组里（每组 4 路），每组要放 8 行，必然互相挤掉。宽度 65 时相邻两行相隔 260 字节，行号不再是 4 的倍数，映射到的组分散开。用代码 03 的 Schedule(65) 实测：默认顺序的未命中从 266,760 降到 17,814，而且和全相联缓存一样多——冲突未命中完全消失。这就是在矩阵的每行末尾加 padding 的经典技巧，和第 4 讲共享内存的 tile[32][33] 是同一个道理。

练习 6　写出代码 03 演示 1 那个调度对应的 TVM 原语序列，并说出它和第 4 讲的共享内存分块矩阵乘法有什么对应关系。

答案：io, jo, ii, ji = s.tile(i, j, 16, 16); ko, ki = s.split(k, 4); s.reorder(io, jo, ko, ki, ii, ji)。第 4 讲里 (io, jo) 对应一个线程块负责的输出块，ko 循环每次把一个 A 块和 B 块搬进共享内存（相当于再加 cache_read(..., "shared") 和 compute_at(..., ko)），(ii, ji) 由块内的线程并行执行（bind 到 threadIdx）。

练习 7　为什么 Triton 要求 tl.arange 的长度是 2 的幂？这给程序员带来什么负担？

答案：块的形状决定了寄存器分配、线程布局、向量化、归约树等，2 的幂让这些映射规整、容易生成高效代码。负担是：真实形状通常不是 2 的幂，程序员要用掩码处理多出来的部分，并接受被屏蔽部分的浪费（代码 01：781 列浪费 24%）。

练习 8　（开放）对比 minitriton 的 FlashAttention（代码 04）和第 5 讲 gpusim 上的 FlashAttention（第 5 讲代码 03），列出 Triton 版本里”消失”的所有实现细节，并说明 Triton 编译器需要为每一项做什么。

参考思路：线程到行的映射（布局分析）、Q/K/V 的合作加载与合并访存（coalesce pass）、共享内存的 padding 或 swizzle（共享内存布局）、__syncthreads（编译器插入同步）、K/V 加载与计算的重叠（num_stages 软件流水线）、tl.dot 映射到 Tensor Core 指令及其操作数布局。

13. 延伸阅读 · Further reading

• Triton 官方教程：01 向量加法、02 融合 softmax、03 矩阵乘法、06 融合注意力，以及 persistent 矩阵乘法的教程。
• TVM 官方教程：How to optimize GEMM on CPU（本讲代码 03 的顺序来自它）。
• J. Ragan-Kelley et al., Halide: A Language and Compiler for Optimizing Parallelism, Locality, and Recomputation in Image Processing Pipelines（PLDI 2013）——算法与调度分离的来源。
• C. Lattner et al., MLIR: Scaling Compiler Infrastructure for Domain Specific Computation（CGO 2021）——Triton 当前的编译栈建立在它上面。
• J. Ansel et al., PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation（ASPLOS 2024）——TorchInductor 怎样生成 Triton 代码。

14. 下一讲预告 · Next lecture

本讲留下了一个问题：Triton 的块大小、num_warps、num_stages、GROUP_M，TVM 的几十万种调度——谁来选？ 第 9 讲《ML 编译器（二）：内核自动调优》从 Triton 的 @triton.autotune（在一小组候选里实测）讲到 AutoTVM（用学习的代价模型预测性能、减少实测次数）和 Ansor（不需要人写调度模板：自动生成程序草图、随机标注、进化搜索微调）。我们会在一个”假想 GPU”上设计一个调优问题，比较网格搜索、随机搜索、代价模型引导、进化搜索四种策略”测多少次能找到多好的配置”。