∇ Stanford CS336（2026 春季）— 第 6 讲：Kernel 与 Triton

更简洁的张量表达式并不意味着更少的 GPU 内存访问。

以一个构造的 1024 × 1024 float32 softmax 为例，在逻辑读写模型中，五次独立的 pass 大约传输 32.016 MiB。而一个理想的融合行 kernel 只传输 8 MiB：读入输入，将中间值保留在片上，写出结果。

这大约是四倍的传输量比，并非实测的四倍加速。缓存、归约开销、寄存器压力和溢出（spill）等因素仍然不可忽视。

Triton 让这种映射变得显式：一个程序（program）持有一个值块（block），对尾部进行掩码处理，并执行局部计算。BLOCK_SIZE 是逻辑元素的数量，而不是 CUDA 线程的数量。

同样的复用原则从融合激活函数一路延伸到行归约和分块矩阵乘法。先测量基线，检查数值边界情况，再比较实际的 kernel 性能。

我的笔记： https://ickma2311.github.io/ML/CS336/lecture06-kernels-triton.html…

更快的 Kernel 始于更少的内存访问 – ∇ ickma.dev

来源：https://ickma2311.github.io/ML/CS336/lecture06-kernels-triton.html 一个张量表达式可以很简短，但仍然会多次将同一个数组写入 GPU 内存。由一个 kernel 写出的中间结果可能立刻被下一个 kernel 读取。消除这些不必要的内存访问，可能比减少乘法运算次数更重要。

Stanford CS336 第 6 讲（https://cs336.stanford.edu/lectures/?trace=lecture_06），由 Percy Liang 在 2026 春季授课，将这一观察转化为一套实用流程：测量工作量、检查 kernel、融合操作、选择能复用数据的 tile。示例从 GELU 到 softmax、行归约，再到带有激活函数的矩阵乘法。

上一讲（https://ickma2311.github.io/ML/CS336/lecture05-gpus-tpus.html）介绍了硬件预算。本文将这些预算与程序结构联系起来。文中数值示例均为人为构造的计算，并非 GPU 基准测试结果。

官方录像（https://www.youtube.com/watch?v=xnDHaNUvHBg）· 课程代码（https://github.com/stanford-cs336/lectures/blob/main/lecture_06.py）

▶ 在 YouTube 上观看官方课程录像（https://www.youtube.com/watch?v=xnDHaNUvHBg）

修改 Kernel 之前先测量已完成的工作

CUDA 启动是异步的：Python 可能在 GPU 仍在工作时就已返回。因此，围绕启动操作的 CPU 计时器测量的主要是提交开销。一个有用的基准测试必须明确说明它测量的是哪个时间区间。

对于单条 CUDA stream 上的 GPU 区间计时，应记录开始事件、入队操作、记录结束事件、等待完成，然后读取事件经过时间。先进行预热，使编译和初始设置不会主导稳态测量。除非分配输入本身就是你要测量的工作负载的一部分，否则应将输入创建排除在测量区间之外。

同步的主机墙钟测量回答的是另一个问题：从应用视角看，整个调用花费了多长时间。CUDA 事件计时并不是完整的应用延迟测量，已提交的 kernel 之间的间隙仍然会影响事件区间。CUDA 计时指南（https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/index.html#timing）解释了为什么同步边界很重要。

在你关注的各种形状和数据类型上重复测量。非常小的操作可能主要受启动和调度开销支配；较大的矩阵乘法则可能将更多时间花在算术运算上。单一尺寸不能建立普遍的性能排名。

然后进行性能分析。计时器给出一个持续时间；性能分析器（profiler）则显示哪些 kernel 运行了以及时间花在了哪里。两次矩阵乘法调用在不同形状下可能选择不同的 kernel。较长的 kernel 名称可能透露其架构目标、数据类型和 tile 大小，但仅凭名称并不能确定性能。PyTorch 的 profiler 使用教程（https://docs.pytorch.org/tutorials/recipes/recipes/profiler_recipe.html）提供了实用的接口。

硬件细节成为程序的约束

一个 CUDA grid 包含线程块（thread block）；一个线程块在一个流式多处理器（SM）上执行，块内的线程可以通过共享内存进行协作。在 NVIDIA GPU 上，线程被组织为 warp，每个 warp 包含 32 个线程。warp 内的分支发散会降低协同完成的有效工作量。现代设备上的独立线程调度（https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/index.html#branching-and-divergence）使得“所有线程始终同步执行“这一假设在正确性层面已不再成立。

另外三个细节影响映射方式：

• **合并访问（Coalescing）：**一个 warp 访问的相邻地址可以被高效地合并处理。对于现代 CUDA 设备，以 32 字节扇区（sector）为单位进行推理很有用：32 个线程读取连续对齐的 4 字节值，覆盖 128 字节，即四个这样的扇区。跨步（strided）模式可能触及更多扇区。128 字节跨度不应被视为普适的单次事务规则。

• **共享内存 bank：**在常见的 32 bank、32 位字模型中，索引相差 32 的字映射到同一个 bank。warp 内对该 bank 的不同地址访问可能被串行化；读取同一个字则可以广播。填充（padding）或 swizzling 可以改变地址到 bank 的映射。

• **占用率（Occupancy）：**寄存器和共享内存限制了 SM 上可以驻留的 warp 数量。更多驻留 warp 有助于隐藏延迟，但最大化占用率并不等同于最小化运行时间。

以一个简化的寄存器预算为例，假设一个 SM 有 65,536 个寄存器，支持 64 个驻留 warp。一个包含 128 个线程、每个线程使用 160 个寄存器的线程块需要 20,480 个寄存器。忽略分配粒度和其他限制，只有

\[ \left\lfloor\frac{65{,}536}{128\times160}\right\rfloor = 3 \]

个这样的块可以容纳。这给出 12 个 warp，即 \(12/64 = 18.75\%\) 的占用率。实际的驻留限制还取决于架构、共享内存和分配规则。

还有一个独立的调度效应：部分填充的最后一个 wave 可能在启动结束时让大多数 SM 空闲。在简化的每 SM 一个块的理想模型中，148 个 SM 需要 160 个等时长的块来完成两轮 wave。第二轮 wave 只包含 12 个块。更一般地说，同时可用的块槽位取决于驻留限制；恰好是 SM 数量的整数倍并不是普适的调优规则。

这些细节说明了为什么要测量你提出的映射方案，而不是选择尽可能大的 tile。

融合将临时值保留在算术运算附近

考虑 tanh 近似的 GELU：

\[ \operatorname{GELU}_{\tanh}(x) = \frac{x}{2}\left[1 + \tanh\!\left(\sqrt{\frac{2}{\pi}} \left(x + 0.044715x^3\right)\right)\right]. \]

这是 \(x\Phi(x)\) 的近似，而非完全相同的数学函数。比较时应在两边使用相同的近似方式——例如 PyTorch 的 approximate="tanh" 模式。参见 GELU 定义（https://docs.pytorch.org/docs/stable/generated/torch.nn.GELU.html）。

由多个独立张量操作拼接的 eager 实现可能启动多个 kernel 并物化临时数组。融合实现可以加载每个输入、计算激活值、写出最终输出，而无需将每个中间张量存储到全局内存中。内置算子或 torch.compile 可能已经实现了这一点；在编写自定义 kernel 之前，先检查生成的执行情况。

融合也会改变资源使用。更多的活跃中间值会增加寄存器压力并导致溢出。关键问题是：节省的内存传输和 kernel 启动次数是否在目标工作负载上抵消了这些代价。

数值行为仍然需要注意。恒等式

\[ \tanh(a) = \frac{e^{2a}-1}{e^{2a}+1} \]

在实数范围内成立，但直接的指数运算可能在浮点中溢出。对于 GELU 内层参数在 \(x=11\) 时，\(a \approx 56.263\)，因此 \(e^{2a} \approx 7.41 \times 10^{48}\) 超出了 float32 的有限范围。由此产生的无穷大除以无穷大可能产生 NaN。课程中的示例 GELU kernel 使用了这个指数恒等式。请使用适当的数值稳定实现并测试输入范围；在小型随机输入上的一致性并不能证明在所有情况下的数值安全性。

Triton 程序持有一个值块

Triton 表达的是对值数组的操作，而不需要程序员为每个单独的 CUDA 线程分配操作。对于这里的 kernel，一个程序实例映射到一个线程块。编译器将该程序的值分配到线程和寄存器上，并插入归约所需的通信操作。课程使用共享内存作为直观模型；逻辑值与其物理存储之间的区别在解读代码时很重要。

对于逐元素操作，设 \(N\) 为元素数量，\(B\) 为一个程序处理的数量。启动 \(\lceil N/B\rceil\) 个程序。程序 \(p\) 构造索引

\[ i = pB + [0, 1, \ldots, B-1] \]

并仅加载或存储满足 \(i < N\) 的索引。

**BLOCK_SIZE 计算的是逻辑值，不是 CUDA 线程。**一个线程可以处理多个值；warp 数量和编译器布局决定了具体的映射。同样，tl.load 并不意味着“每个值必然存储在 CUDA 共享内存中“。物理存储可能涉及寄存器、共享内存、缓存和溢出，具体取决于生成的 kernel。

以构造的 \(N=10\)、\(B=4\) 为例，第三个程序持有候选索引 8、9、10、11。只有前两个是有效的。掩码（mask）防止越界访问。这个概念出现在 Triton 的向量加法教程中（https://triton-lang.org/main/getting-started/tutorials/01-vector-add.html）。

构造示例：十个元素，块大小为四。程序零处理索引零到三；程序一处理四到七；程序二处理八和九，并对索引十和十一进行掩码。逻辑元素不是 CUDA 线程的计数。

逻辑块与掩码尾部。每个程序持有四个候选值；最后一个程序只有两个有效值。

此处 \(N\) 计算的是扁平化输入中的元素数量；在下面的矩阵示例中，\(N\) 表示列数。

掩码加载还需要一个适合后续操作的填充值。加法归约使用零。先取最大值再做 softmax 时，负无穷是有用的中性填充值。

Softmax 让内存传输的节省变得显式

对于一行 \(x_1, \ldots, x_N\)，在指数运算前减去其最大值 \(m\)：

\[ m = \max_j x_j, \qquad z_i = e^{x_i - m}, \qquad y_i = \frac{z_i}{\sum_j z_j}. \]

共享因子 \(e^{-m}\) 在分子和分母之间约去，数学上的 softmax 结果不变。对于有限输入，最大的指数变为 1，避免了这一步的溢出。远小于 1 的项仍然可能出现下溢。

假设一个 \(M \times N\) 矩阵通过五次独立的 pass 处理。以下统计计算移动的逻辑数组元素数量，每次消费 pass 中每行统计量被读取一次：

最后一次 pass 包含 \(M\) 个分母读取，课程代码的读取计数注释中省略了这一点。这是一个算法层面的内存传输模型。缓存、广播行为、编译器融合和溢出决定了实际的 HBM 传输量。只有在数组使用假定的常见数据类型时，才乘以每元素字节数。

一个融合的行 kernel 在理想情况下可以读入 \(MN\) 个输入元素并写出 \(MN\) 个输出，同时将中间值保留在片上。在该模型中比值为

\[ \frac{(5MN + 2M) + (3MN + 2M)}{2MN} = 4 + \frac{2}{N}. \]

在 \(M = 1024\)、\(N = 1024\)、每个计数元素四字节的情况下，独立 pass 大约移动 32.016 MiB，而理想的融合路径为 8 MiB。比值约为 4.002。这是一个内存传输比，并非实测的四倍加速，也不是超越原生 softmax 算子的承诺。

Softmax 内存传输对比：构造的 1024 × 1024 float32 矩阵。五次独立 pass 的逻辑读写总计 32.016 MiB。理想的融合行 kernel 读入输入并写出输出，总计 8 MiB。约四倍的传输量比并非实测加速比。

Softmax 可以避免写出中间数组。字节总量基于五次 pass 的逻辑传输模型。

Triton 官方 softmax 教程（https://triton-lang.org/main/getting-started/tutorials/02-fused-softmax.html）展开了这种融合模式。其核心条件是行的工作集（working set）能放入所选的片上执行策略中。大行可能需要另一种算法或布局。

填充是数学的一部分

对于五列的行，一个基本的 Triton 实现可以使用八个值的逻辑块。加载前五个条目，剩余三个填充 \(-\infty\)。减去有限的行最大值后，它们的指数变为零，对归一化没有贡献。

对于行 \([0, 0, \log 2, 0, 0]\)，未归一化的权重与 \([1, 1, 2, 1, 1]\) 成比例。结果为

\[ [1/6,\; 1/6,\; 1/3,\; 1/6,\; 1/6]. \]

如果用零填充三个填充位置，将错误地引入三个额外项。仅对最终存储进行掩码并不能修复分母错误。

该基础 kernel 还假设每行至少有一个有限条目，其余条目为有限值或 \(-\infty\)。全为 \(-\infty\) 的行会产生 \(-\infty - (-\infty)\)；NaN 或正无穷输入需要明确的处理策略。非单位列步长（stride）也必须显式处理，而不能默默假定列是连续的。

长行需要循环，然后归约

行求和说明了如何处理超过一个逻辑 tile 大小的值。课程中简单的 row * N + column 寻址假设行是连续的；其他布局需要显式步长。保持一个长度为 \(B\) 的向量累加器。对于每个 tile，将有效值加到累加器上，超出行边界的位置用零替代。循环结束后，将累加器归约为一个标量。

以构造的行 \([1, 2, \ldots, 10]\) 和 \(B=4\) 为例，更新过程为

\[ \begin{aligned} [0,0,0,0] + [1,2,3,4] &= [1,2,3,4],\\ [1,2,3,4] + [5,6,7,8] &= [6,8,10,12],\\ [6,8,10,12] + [9,10,0,0] &= [15,18,10,12]. \end{aligned} \]

最终和为 \(15 + 18 + 10 + 12 = 55\)。逻辑累加器位置不需要与物理线程一一对应。float32 累加器通常对低精度输入有用，但将结果转换回更小的数据类型仍可能溢出或丢失精度。

这个循环减小了工作集大小。它不会自动将单行并行化到多个线程块。当行数非常少时，每行一个程序的设计可能仍然让 GPU 大部分时间处于空闲状态。跨块的并行归约需要额外的组合策略。

矩阵 tile 将一次加载变为多次使用

设 \(A \in \mathbb{R}^{M \times K}\) 和 \(B \in \mathbb{R}^{K \times N}\)。一个程序计算 \(C = AB\) 的一个 \(B_M \times B_N\) tile。在沿 \(K\) 的每一步中，它加载一个 \(B_M \times B_K\) 形状的 \(A\) tile 和一个 \(B_K \times B_N\) 形状的 \(B\) tile，然后更新同一个输出累加器：

\[ C_{\mathrm{tile}} \mathrel{+}= A_{\mathrm{tile}} B_{\mathrm{tile}}. \]

\(A\) tile 的每个元素可以对 \(B_N\) 个输出元素