第 6 讲：核函数与 Triton

对应：Stanford CS336 Spring 2026 Lecture 6（Kernels, Triton）。Spring 2025 的 Lecture 6 内容相近，有公开录像。 配套代码：code/lecture06/ 预计阅读时间：3–4 小时 硬件说明：Triton 只能在 NVIDIA GPU 上运行。配套代码附带了一个用 NumPy 写的”迷你 Triton”模拟器（mini_triton.py），同一份核函数代码在有 NVIDIA GPU 的机器上由真正的 Triton 编译运行，在其他电脑上由模拟器在 CPU 上执行——你可以在任何电脑上读懂、运行、验证每一个核函数。

0. 本讲导航

0.1 学完本讲，你应该能够

说出几个影响性能的硬件细节：warp、占用率、共享内存的 bank 冲突、合并访存、波次量化。
正确地做基准测试（热身、同步），并解释不这样做会错得多离谱。
用性能分析器找出时间花在哪些运算 / 核函数上。
对比 PyTorch、CUDA、Triton 三种编程抽象；理解 Triton “以线程块为单位编程”的模型。
用 Triton 写出：向量加法、融合 GeLU、融合 softmax、分块行求和、分块矩阵乘法 + ReLU 融合、融合 RMSNorm（前向 + 反向）。
理解 mask、tl.constexpr、指针与步长、tl.dot、在程序内循环（分块）这些 Triton 的基本要素。
说出 torch.compile 能自动做什么、不能做什么，以及图中断和重新编译是怎么回事。

0.2 先修知识

知识点	用在哪里	最低要求
第 5 讲	全文	roofline、融合、分块、合并访存、FlashAttention
第 2 讲	第 8 节	张量的步长（stride）
第 3 讲	第 9 节	RMSNorm 的公式
求导	第 9 节	链式法则（看不懂推导也可以直接用结论）

0.3 本讲地图

第1节 硬件回顾：影响性能的几个细节
第2节 基准测试：测"多快"
第3节 性能分析：测"时间花在哪"
第4节 三种编程抽象：PyTorch / CUDA / Triton
第5节 第一个 Triton 核函数：向量加法、GeLU（逐元素）
第6节 融合 softmax（每个程序一行）
第7节 分块行求和（程序内循环）
第8节 分块矩阵乘法 + ReLU 融合（二维网格、tl.dot）
第9节 融合 RMSNorm：前向、反向、接入自动求导
第10节 torch.compile：编译器能做什么，不能做什么

上一讲是”GPU 为什么快、什么时候慢”的理论；这一讲是动手实践：先测量，再优化。

1. 硬件回顾：影响性能的几个细节

1.1 存储层次的带宽

CS336 课件对比了 A100、H100、B200（SM 数分别为 108、132、148）的各级存储带宽，数量级大致是：

\underbrace{\text{寄存器}}_{\text{每秒上百 TB}} \gg \underbrace{\text{L1 / 共享内存}}_{\text{每秒几十 TB}} \gg \underbrace{\text{L2}}_{\text{每秒 5~17 TB}} \gg \underbrace{\text{HBM}}_{\text{每秒 2~8 TB}}

每往下一级，带宽都低一个数量级左右。

1.2 编程模型与硬件的对应

线程（thread）执行代码，处理一部分数据；
线程块（block / CTA）是一组线程，共享一个 SM 的资源（共享内存等）；
网格（grid）是一次启动的所有线程块。
H100 / B200 还支持线程块簇（thread block cluster）：一组线程块可以访问彼此的共享内存（分布式共享内存）。

这些抽象决定了程序的语义（算得对不对）；下面几个硬件细节决定性能（算得快不快）：

warp：32 个线程一组，同步执行同一条指令。warp 内分支发散会让不同路径串行执行（第 5 讲 6.1 节）。SM 在 warp 之间切换是零开销的——一个 warp 等内存时，另一个 warp 马上顶上。

占用率：每个线程可以用 0~255 个寄存器，每个 SM 共有 65536 个寄存器、最多驻留 64 个 warp。每个线程用的寄存器越多，能同时驻留的 warp 越少。例如每块 128 个线程、每线程 160 个寄存器，一个块就要 20480 个寄存器，一个 SM 只能驻留 3 个块（12 个 warp，占用率 19%）。线程粗化（一个线程处理多个元素）可以在低占用率下保持性能。

共享内存的 bank 冲突：共享内存分成 32 个 bank，每个 bank 宽 4 字节，每个周期每个 bank 只能服务一个地址。如果一个 warp 中的多个线程访问同一个 bank 的不同地址，这些访问就会串行化。最坏的情况是 32 路冲突——比如在矩阵乘法中按列读取一个 32 列宽的共享内存块。解决办法是填充（每行多留一个空位）或swizzling（用异或运算重新排列地址）。

合并访存（HBM）：一个 warp 的 32 个线程各读 4 字节，正好是 128 字节的一个缓存行。最理想的情况是它们全都落在同一个缓存行里（第 5 讲 6.5 节）。

波次量化：线程块按”波”调度。如果线程块数不能被 SM 数整除，最后一波会让很多 SM 闲着。尽量让线程块数是 SM 数（B200 是 148）的整数倍（第 5 讲第 7 节）。

2. 基准测试：测”多快”

2.1 做法

基准测试测的是端到端的挂钟时间。它告诉你”多快”，但不告诉你”时间花在哪”。正确的步骤（配套代码 01_benchmarking_profiling.py 中的 benchmark）：

热身：先运行几次，让一次性的开销（编译核函数、分配显存、选择算法）发生在计时之外；
同步：torch.cuda.synchronize()，确保之前的运算都已完成；
多次试验：每次开始计时 → 运行 → 同步 → 停止计时；
取平均值或中位数。

在 CUDA 上更精确的做法是用 torch.cuda.Event(enable_timing=True)：在 GPU 的时间线上打点（record()），用 elapsed_time() 读出 GPU 实际执行的时间，不受 CPU 端开销的影响。

2.2 两个大坑

在一台 Apple M4 Pro 上用 Apple GPU（MPS）实测一个 4096×4096 的矩阵乘法：

做法	测得的每次耗时
第一次运行（没有热身）	131 ms（是稳定值的 5 倍）
热身之后	26.4 ms
不同步就停表	0.019 ms（比真实值小约 1400 倍！）
同步之后停表	26.1 ms

为什么不同步会错得这么离谱？GPU 是异步执行的：Python 调用 a @ b 时，只是把”做这个矩阵乘法”的命令放进 GPU 的队列就立刻返回了。如果在 GPU 真正算完之前就停表，测到的只是”发命令”的时间。

一个推论：.item()、print(张量)、.cpu() 都会隐式同步（CPU 必须等 GPU 算完才能拿到数值）。训练循环里每一步都打印损失，会让 CPU 每步都等 GPU，失去异步执行带来的”CPU 准备下一步、GPU 计算这一步”的重叠。

2.3 耗时随尺寸的变化

$n$	64	256	1024	2048	4096
矩阵乘法耗时	0.15 ms	0.16 ms	0.55 ms	3.4 ms	25.8 ms

小尺寸时耗时几乎不变——固定开销占主导，硬件远没被喂饱；尺寸足够大后，耗时大约按 $n^3$ 增长。CS336 课上在 B200 上测到的是同样的规律。

3. 性能分析：测”时间花在哪”

PyTorch 自带性能分析器：

from torch.profiler import profile, ProfilerActivity
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    run()
print(prof.key_averages().table(sort_by="cuda_time_total"))

它列出每个运算 / 核函数被调用了多少次、花了多少时间。演示 3 分析了用基本运算手写的 GeLU：调用了 6 次 aten::mul、2 次 aten::add、1 次 aten::tanh——9 个核函数，每个都要读写整个张量；而内置的 F.gelu 只有 1 个 aten::gelu。

在 NVIDIA GPU 上，分析器还会显示实际执行的 GPU 核函数名字。CS336 课上的例子：cutlass3x_sm100_simt_sgemm_f32_f32_f32_f32_f32_64x64x16——从名字能看出用的是 NVIDIA 的 CUTLASS 库、B200（sm100）架构、FP32、64×64×16 的分块。另一个有趣的观察：不同的矩阵尺寸会触发不同的核函数（库会根据形状选择最合适的实现）。

更深入的分析工具是 NVIDIA Nsight Systems（作业 2 会用到），它能显示 CPU 和 GPU 的完整时间线，看清楚哪里在等待、哪里在重叠。

4. 三种编程抽象：PyTorch / CUDA / Triton

	你描述的是	优点	缺点
PyTorch	整个张量上的运算	最简单	每个运算一个核函数，无法控制融合
CUDA	每一个线程做什么	完全掌控硬件	要自己管理线程下标、共享内存、同步、合并访存……
Triton	每一个线程块（程序）做什么，操作的是”一块数据”	简单得多，对大多数负载也足够强	对最底层细节的控制不如 CUDA

4.1 CUDA 的样子（示意）

// 示意代码（未在本课程环境中编译运行）：每个线程处理一个元素
__global__ void gelu_kernel(const float* in, float* out, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;      // 我是第几个线程？
    if (i < n) {                                         // 边界检查
        float x = in[i];
        out[i] = 0.5f * x * (1.0f + tanhf(0.79788456f * (x + 0.044715f * x * x * x)));
    }
}
// 主机端启动：gelu_kernel<<<(n + 1023) / 1024, 1024>>>(in, out, n);

CS336 课上用 torch.utils.cpp_extension.load_inline 在 Python 里直接编译这样的 CUDA 代码，并与 PyTorch 版本比较。

4.2 Triton 的模型

在 Triton 里，你写的核函数描述一个程序实例（对应一个线程块）要做的事，它操作的是一整块数据：

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(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)        # 一次写一块

指针：传入的张量在核函数里是指向第一个元素的指针；x_ptr + offsets 是”一块指针”。
mask：最后一个程序的部分位置可能越界，用 mask 掩掉，不读不写（越界访问在 GPU 上会出错）。
tl.constexpr：编译期常量。Triton 会为每个不同的 BLOCK_SIZE 编译一个专门的版本，所以它必须是固定值（通常是 2 的幂）。
启动：add_kernel[grid](x, y, out, n, BLOCK_SIZE=256)，grid = (cdiv(n, 256),) 表示启动多少个程序。

块内的线程如何划分、如何合并访存、何时使用共享内存，大部分交给 Triton 编译器。思维模式：以线程块为单位思考——“从 HBM 加载一块数据 → 在片上计算 → 写回 HBM”。

4.3 关于配套的模拟器

mini_triton.py 用 NumPy 实现了上面用到的 API（@triton.jit、kernel[grid](...)、tl.program_id、tl.arange、tl.load、tl.store、tl.sum、tl.max、tl.exp、tl.dot 等），逐个程序实例地顺序执行核函数。它还能统计每次启动实际读写了多少元素，帮助验证”融合核函数只读写一次”。

它的局限（使用前务必了解）：只实现了一小部分 API；程序实例是顺序执行的，不能用来测速，也可能掩盖”不同程序之间的数据竞争”这类并行 bug；没有 BF16。每个核函数文件开头都会检查：有 NVIDIA GPU 并安装了 triton 时用真正的 Triton（并自动测速），否则用模拟器。

5. 第一个 Triton 核函数：向量加法与 GeLU

配套代码 02_triton_elementwise.py。

向量加法： $n = 1000$ 、每块 256 个元素时，网格有 4 个程序；第 4 个程序的后 24 个位置被 mask 掩掉。模拟器统计到恰好读取 2000 个元素、写入 1000 个——被掩掉的位置不读不写。

GeLU：把 9 个逐元素运算融合进一个核函数，只读一次 $x$ 、写一次 $y$ ，中间结果全在寄存器里。一个小技巧：

\tanh(a) = 2\sigma(2a) - 1 \quad\Rightarrow\quad \tfrac12\big(1 + \tanh(a)\big) = \sigma(2a) \quad\Rightarrow\quad \text{GeLU}(x) = x\cdot\sigma(2a)

用 sigmoid 代替 tanh，结果完全相同，而且 tl.sigmoid 在所有 Triton 版本中都可用。与 F.gelu(approximate="tanh") 的差异为 $5\times10^{-7}$ 。CS336 课上的实测：在 GPU 上，手写的 PyTorch 版本明显慢于内置版本和 torch.compile 版本；后两者都只有一个融合的核函数（torch.compile 生成的正是 Triton 核函数），速度相当。

看看生成的汇编（PTX）：Triton 把核函数编译成 PTX（NVIDIA GPU 的汇编语言）。CS336 课上展示了 GeLU 核函数的 PTX，可以看到：ld.global.* / st.global.* 是全局内存的读写，%ctaid.x 和 %tid.x 是线程块编号和线程编号，%f*、%r* 是浮点和整数寄存器；还能看到线程粗化——一个线程一次处理 8 个元素。读 PTX 不是必需的，但它能帮你确认编译器确实做了你期望的事。

6. 融合 softmax：每个程序负责一行

配套代码 03_triton_reductions.py。

逐元素运算中各元素互不相干。更常见的是归约：一行的结果依赖这一行的所有元素，比如 softmax、LayerNorm/RMSNorm。

朴素实现的读写（ $M\times N$ 矩阵，CS336 课上的计数）：

x_max = x.max(dim=1)          读 MN，写 M
x = x - x_max[:, None]        读 MN + M，写 MN
num = torch.exp(x)            读 MN，写 MN
den = num.sum(dim=1)          读 MN，写 M
y = num / den[:, None]        读 MN（+M），写 MN
合计：读约 5MN，写约 3MN      理想：读 MN，写 MN（快约 4 倍）

融合 softmax 核函数：网格有 $M$ 个程序，第 $i$ 个程序负责第 $i$ 行；BLOCK_SIZE 取不小于 $N$ 的 2 的幂，把整行一次读进来：

row = tl.program_id(0)
cols = tl.arange(0, BLOCK_SIZE)
mask = cols < n_cols
x = tl.load(x_ptr + row * row_stride + cols, mask=mask, other=float("-inf"))
x = x - tl.max(x, axis=0)
num = tl.exp(x)
tl.store(y_ptr + row * row_stride + cols, num / tl.sum(num, axis=0), mask=mask)

在 $300\times781$ 的矩阵上，与 torch.softmax 的差异为 $10^{-6}$ ；模拟器统计到读、写都恰好是 $MN$ 。

边界：为什么被掩掉的位置要填 $-\infty$ ？ 如果填 0，当整行都是负数时最大值会错误地变成 0，而且 $e^0 = 1$ 会被加进分母。填 $-\infty$ 则 $e^{-\infty} = 0$ ，对结果毫无影响。另一个边界：一行必须能放进一个块（受寄存器容量限制）；更长的行需要在程序内循环 + 在线 softmax（第 5 讲）。

7. 分块行求和：程序内循环

如果一行有 4096 个元素，而块大小只有 1024 呢？在程序内部循环：用一个长度为 BLOCK_SIZE 的累加器（在寄存器里），每次读入这一行的下一段、逐位置累加，循环结束后再对累加器求和：

acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
for start in range(0, n_cols, BLOCK_SIZE):
    cols = start + tl.arange(0, BLOCK_SIZE)
    acc += tl.load(x_ptr + row * row_stride + cols, mask=cols < n_cols, other=0.0)
tl.store(out_ptr + row, tl.sum(acc, axis=0))

这个循环在生成的代码里是一个真正的循环（不会被完全展开）。块大小是一个权衡：太小则循环次数多、每次读得少；太大则占用太多寄存器、占用率下降。这种”沿一个维度循环，把部分结果累加在寄存器里”的模式，就是分块矩阵乘法和 FlashAttention 的骨架。

8. 分块矩阵乘法 + ReLU 融合

配套代码 04_triton_matmul.py。

8.1 回顾：为什么要分块

朴素矩阵乘法：每个输出元素读 $K$ 个 $A$ 的元素和 $K$ 个 $B$ 的元素，读取量 $\propto MNK$ ，算术强度 $O(1)$ ，内存受限。
把 $A$ 、 $B$ 整个放进共享内存：读取量只有 $MK + KN$ ，但矩阵通常太大，放不下。
分块：把输出 $C$ 切成块，每个线程块负责一块；沿 $K$ 方向每次加载一对 $A$ 、 $B$ 的小块，做小矩阵乘法、累加部分和；最后写回。算术强度 $O(\text{块大小})$ ，可调。

8.2 核函数

pid_m, pid_n = tl.program_id(0), tl.program_id(1)             # 二维网格：负责哪个输出块
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
offs_k = tl.arange(0, BLOCK_K)
a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak   # (BLOCK_M, BLOCK_K) 的地址块
b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn   # (BLOCK_K, BLOCK_N)
acc = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
for k in range(0, K, BLOCK_K):
    a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] + k < K), other=0.0)
    b = tl.load(b_ptrs, mask=(offs_k[:, None] + k < K) & (offs_n[None, :] < N), other=0.0)
    acc += tl.dot(a, b)                                       # Tensor Core 完成的小矩阵乘法
    a_ptrs += BLOCK_K * stride_ak                             # 前进到下一块
    b_ptrs += BLOCK_K * stride_bk
acc = tl.maximum(acc, 0.0)                                    # 融合的 ReLU
tl.store(c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn, acc,
         mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))

要点：

步长导航：矩阵在内存中是一维的，x[row, col] 的地址是 x_base + row*stride_row + col*stride_col。把步长作为参数传入，核函数就能处理转置等不连续的矩阵（配套代码对转置输入也做了验证）。
二维的地址块：offs_m[:, None] 和 offs_k[None, :] 通过广播得到一个二维的地址矩阵，一次 tl.load 读入整个小块。
mask： $M$ 、 $N$ 、 $K$ 不一定是块大小的整数倍，越界的位置读取时填 0（不影响乘积之和），写入时跳过。
融合：ReLU 在累加器上、写回之前完成——不需要先把 $A@B$ 写回显存再启动一个 ReLU 核函数（省下 $2MN$ 的读写和一次启动）。
真实 Triton 中 tl.dot 要求块的每一维至少为 16。

8.3 验证读写量

在 $n = 256$ 的矩阵上，模拟器统计的读取量与第 5 讲的分块公式完全一致：

输出块	程序数	读取的元素	写入
16×16	256	2,097,152	65,536
32×32	64	1,048,576	65,536
64×64	16	524,288	65,536
128×128	4	262,144	65,536

输出块越大，每个输入元素被复用得越多，读取量越少；写入量始终是 $n^2$ 。

8.4 真实世界的调优

真正高性能的 Triton 矩阵乘法还会：用 @triton.autotune 自动在多组块大小、num_warps（每个程序用几个 warp）、num_stages（软件流水线的级数：一边算当前块，一边预取后面的块）中挑最快的；调整程序的执行顺序（“分组”排列），让相邻程序读取的 $A$ 、 $B$ 块更可能留在 L2 缓存里。Triton 官方教程的矩阵乘法一章详细讲了这些技巧。

9. 融合 RMSNorm：前向、反向、接入自动求导

配套代码 05_triton_rmsnorm.py。2026 年版 CS336 的作业 2 包含用 Triton 实现融合 RMSNorm 核函数。

9.1 前向

每个程序处理一行（与 softmax 同一个模式）：转成 FP32， $r = 1/\sqrt{\text{mean}(x^2) + \epsilon}$ ， $y = x\cdot r\cdot g$ 。顺便把每行的 $r$ 存下来（每行只有一个数），反向时直接使用。

9.2 反向的推导

对一行： $y_i = g_i x_i r$ ， $r = (\frac1N\sum_j x_j^2 + \epsilon)^{-1/2}$ ，所以 $\frac{\partial r}{\partial x_j} = -\frac{r^3 x_j}{N}$ 。于是

\frac{\partial L}{\partial x_j} = dy_j\,g_j\,r - \frac{r^3 x_j}{N}\sum_i dy_i\,g_i\,x_i, \qquad \frac{\partial L}{\partial g_i} = \sum_{\text{所有行}} dy_i\,x_i\,r

每行需要 $x$ 、 $g$ 、 $dy$ 和保存的 $r$ ；输出这一行的 $dx$ 和”这一行对 $dg$ 的贡献”。

9.3 数据竞争

$dg$ 要对所有行求和。如果每个程序都直接往同一个 $dg$ 数组里累加，GPU 上并行运行的程序会同时”读-改-写”同一个位置，结果不确定（这就是数据竞争）。办法之一：每行写到一个部分梯度缓冲区的对应行，最后在外面 sum(dim=0)。另一个办法是原子加（tl.atomic_add），但它较慢，而且浮点加法顺序不固定导致结果不可复现（第 2 讲 10.5 节）。更好的做法是让每个程序处理多行、先在寄存器里累加这些行的 $dg$ ，再写出，这样缓冲区只有（行数 / 每个程序的行数）× $N$ 。

注意：模拟器是顺序执行的，不会暴露数据竞争——如果你写了一个有竞争的核函数，在模拟器里可能”碰巧”是对的，在真实 GPU 上却出错。

9.4 接入自动求导

用 torch.autograd.Function 包装：forward 启动前向核函数并 save_for_backward(x, g, r)；backward 启动反向核函数，返回 $dx$ 和 $dg$ 。之后就能像普通 PyTorch 模块一样用 .backward()。三种形状下， $y$ 、 $dx$ 、 $dg$ 与 PyTorch 自动求导的差异都在 $10^{-6}$ 量级。

10. torch.compile：编译器能做什么，不能做什么

配套代码 06_torch_compile.py。

torch.compile(model) 由两部分组成：TorchDynamo 在 Python 执行时截获 PyTorch 运算，记录成计算图；TorchInductor 对计算图做优化（主要是融合），在 NVIDIA GPU 上生成 Triton 代码，在 CPU 上生成 C++ 代码。

实测：一个”RMSNorm + 前馈层 + 残差”的小模块（ $d = 1024$ ，Apple GPU），普通执行 40.5 ms，编译后 29.5 ms（1.37 倍）。矩阵乘法本来就由库高效完成，编译器省下的是 RMSNorm、GeLU、残差相加等内存受限运算的读写。

图中断：Dynamo 遇到无法放进图的 Python 代码（依赖张量数值的 if、.item()、print、不支持的第三方库）时，只能把图切断：

函数	图的数量	图中断次数
中间有 `if y.sum().item() > 0:`	2	1
同样的逻辑改用 `torch.where`	1	0

每次中断都会打断融合，还可能引入 GPU→CPU 同步。排查工具：torch._dynamo.explain(fn)(x)，或环境变量 TORCH_LOGS=graph_breaks；fullgraph=True 会在出现图中断时直接报错。

重新编译：用 5 种不同长度的输入调用一个编译过的函数，一共只编译了 2 张图——第一次按固定形状编译，第二次发现这个维度会变，编译一个”动态形状”的版本，之后复用。训练时通常把序列填充 / 截断到固定长度，避免反复编译。

做不到的事：编译器擅长融合逐元素和归约运算，但不会发明新算法。标准注意力 → FlashAttention 不是”把运算合并”，而是用在线 softmax 重新组织计算（PyTorch 会识别注意力模式，替换成现成的融合核函数）。所以大模型的高性能代码通常是”编译器 + 少数手写核函数”：大部分交给 torch.compile，最关键的热点（注意力、归一化、MoE 分发、量化矩阵乘法）用 Triton / CUDA 手写。

11. 本讲小结

理解硬件才能优化：warp（分支发散）、占用率（寄存器用量）、共享内存 bank 冲突（填充、swizzling）、合并访存、波次量化。
基准测试测”多快”：热身 + 同步 + 多次试验。不同步可能错上千倍；.item()、print 会隐式同步。
性能分析测”时间花在哪”：PyTorch profiler 列出每个运算 / 核函数；Nsight Systems 显示完整时间线。
Triton 以线程块为单位编程：program_id 决定负责哪块，arange + 指针得到地址块，mask 做边界检查，load / store 读写，constexpr 是编译期常量。
四种核函数模式：逐元素（GeLU）、每行一个程序的归约（softmax、RMSNorm）、程序内循环的分块（行求和）、二维分块 + tl.dot（矩阵乘法），以及在写回前完成的融合（ReLU）。
反向传播中的”全局求和”（如 RMSNorm 的 $dg$ ）要避免数据竞争：部分和缓冲区或原子操作。
torch.compile 自动完成简单的融合；注意图中断和重新编译；需要改变算法的优化仍要手写。

12. 与作业 2 的关系

作业 2（系统）的前半部分直接对应第 5、6 讲：

基准测试与性能分析：写一个正确的基准测试脚本，测量 Transformer 前向、反向的时间（注意热身和同步）；用 PyTorch profiler / Nsight Systems 找出最耗时的核函数；比较混合精度的效果；分析显存。
Triton 核函数：2025 年版要求用 Triton 实现 FlashAttention-2 的前向（第 5 讲 05 号文件的算法 + 本讲的 Triton 编程），并实现反向；2026 年版的介绍中提到了融合 RMSNorm 核函数（本讲第 9 节）。
与 PyTorch 原生实现、torch.compile 版本比较速度和显存。

建议：先在模拟器里把核函数的逻辑写对（它能帮你检查越界访问和数值结果），再到真正的 GPU 上调试性能。但要记住模拟器不会暴露数据竞争，也不能测速。

13. 练习

练习 1（基准测试）：下面的计时代码有什么问题？怎么改？

t0 = time.time()
y = model(x.cuda())
print(time.time() - t0)

参考答案

(1) 没有热身，第一次运行包含 CUDA 初始化、核函数编译等一次性开销；(2) 没有同步，GPU 可能还没算完就停表了；(3) x.cuda() 的数据搬运也被算进去了；(4) 只测一次，波动大。应该先把 x 放到 GPU、热身几次、每次计时前后 torch.cuda.synchronize()、多次取中位数（或用 CUDA 事件）。

练习 2（占用率）：一个核函数每块 256 个线程、每线程 96 个寄存器、每块使用 48 KB 共享内存。在 H100 的一个 SM 上（65536 个寄存器、228 KB 共享内存、最多 2048 个线程、最多 32 个块）能同时驻留几个块？占用率是多少？瓶颈是什么？

参考答案

寄存器：65536 / (256×96) = 2.67 → 2 块；共享内存：228/48 = 4.75 → 4 块；线程数：2048/256 = 8 块。取最小值 2 块，即 512 个线程 = 16 个 warp，占用率 16/64 = 25%，瓶颈是寄存器。

练习 3（mask 与 other）：在 softmax 核函数里把 other=float("-inf") 改成 other=0.0，对一行全是负数、且 $N$ 不是 2 的幂的输入，结果会怎样？

参考答案

填充位置的值为 0，比所有真实值都大，于是最大值变成 0（不影响正确性，因为减去任何常数 softmax 都不变）；但 $e^{0-0} = 1$ 会被加进分母，每个填充位置都多贡献 1，导致所有概率偏小、每行之和小于 1。

练习 4（网格）：用 matmul_relu_kernel 计算 $(1000\times512)@(512\times700)$ ，块大小 64×64×32。网格是多少？每个程序沿 $K$ 循环几次？

参考答案

网格 $(\lceil1000/64\rceil, \lceil700/64\rceil) = (16, 11)$ ，共 176 个程序；每个程序循环 $512/32 = 16$ 次。

练习 5（RMSNorm 反向）★：推导 RMSNorm 对 $x$ 的梯度，并用 torch.autograd.gradcheck 验证（提示：用 float64 的参照实现）。

练习 6（数据竞争）★：如果在 RMSNorm 反向核函数里直接用 tl.store(dg_ptr + cols, 已有值 + 本行贡献) 来累加 $dg$ （先 load 再 store），在真实 GPU 上会出什么问题？为什么模拟器发现不了？

参考答案

多个程序并行运行时，可能同时读到同一个旧值，各自加上自己的贡献后写回，后写的会覆盖先写的，一部分贡献丢失（“丢失更新”）。模拟器按顺序一个一个地执行程序，不存在”同时读”的情况，所以结果碰巧正确。这也提醒我们：模拟器只能验证单个程序的逻辑，不能替代在真实 GPU 上的测试。

练习 7（torch.compile）：写一个包含 for i in range(x.shape[0]): ... Python 循环的函数，用 torch._dynamo.explain 看看它被编译成了几张图。再改用向量化的写法，比较结果。

练习 8（动手）： (a) 仿照 gelu_kernel 写一个 SiLU（ $x\cdot\sigma(x)$ ）核函数，并用 mini_triton 的统计验证只读写一次。 (b) 把 softmax_kernel 改成”每个程序处理多行”（线程粗化），比较程序数。 (c) 改进 rmsnorm_bwd_kernel：每个程序处理 ROWS_PER_PROGRAM 行，在寄存器里累加 $dg$ ，缩小部分梯度缓冲区。 (d) 如果你有 NVIDIA GPU：用 @triton.autotune 为矩阵乘法核函数自动选择块大小，与 torch.matmul 比较速度。

14. 延伸阅读

CS336 Spring 2026 Lecture 6（GitHub stanford-cs336/lectures 中的 lecture_06.py，可执行讲义，包含 CUDA 与 Triton 代码及 PTX 输出）。
Triton 官方教程（triton-lang.org）：向量加法、融合 softmax、矩阵乘法、LayerNorm（含反向）、FlashAttention。本讲的核函数与之风格一致，强烈建议对照阅读。
Tillet et al.，Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations（2019）。
PyTorch 文档：torch.compile tutorial、torch.profiler；NVIDIA Nsight Systems 用户指南。
Kirk & Hwu，Programming Massively Parallel Processors：CUDA、共享内存、bank 冲突的系统讲解。
GPU MODE（原 CUDA MODE）社区的讲座：从入门到 CUTLASS、Triton 的深入内容。

15. 下一讲预告

到目前为止，我们都在一块 GPU 上想办法。但训练前沿模型需要成千上万块 GPU：一个 700 亿参数的模型，光参数、梯度、优化器状态就要 1 TB 以上的显存（第 2 讲）。下一讲（并行 I）开始多卡训练：

多块 GPU 之间如何通信？all-reduce、reduce-scatter、all-gather 分别是什么？各要传输多少数据？
数据并行：每块卡一份完整模型，各算一部分数据，再同步梯度；
ZeRO / FSDP：把优化器状态、梯度、参数切分到各块卡上，突破单卡显存的限制；
用 PyTorch 的 torch.distributed（在 CPU 上用多进程模拟多卡）亲手实现这些算法。

下一讲之前，建议你：跑通 code/lecture06/ 下的 6 个脚本，读懂 mini_triton.py 的实现，完成练习 2、3、6，并尝试练习 8(a)。