跳转至

CS336 2026 Lecture 6:Benchmarking、Profiling 与 Triton Kernels

LaTeX 源码 · 观看视频

字段 内容
作者/整理 基于 Stanford CS336 Spring 2026 官方可执行讲义整理
来源 Stanford CS336
日期 2026 年春季

CS336 2026 Lecture 6:Benchmarking、Profiling 与 Triton Kernels

\setcounter{tocdepth}{1}

本讲主线:正确性靠编程模型,性能靠硬件模型

Lecture 6 接在 GPU 课之后,目标是把“知道 GPU 很快”推进到“知道怎么测、怎么定位、怎么写 kernel”。核心循环是:benchmark 看到端到端时间,profile 看到时间花在哪里,然后用更贴近硬件的数据组织和 kernel 写法修正瓶颈。

本讲的闭环

编程模型给正确性,硬件模型给性能,benchmark/profiling 给证据,Triton 给可控的 block-level 实现。没有测量就改 kernel,基本是在猜。

课堂提示:正确性与性能需要两套模型

老师在总结中把职责分得很清楚:PyTorch、Triton 与 PTX 等 programming model 让程序语义可表达、可验证;SM、warp、occupancy、bank conflict 与 memory hierarchy 决定实现能否高效。Benchmark 用于观察 scaling,profiler 用于看真实执行了什么、花了多久;Triton 则把思考单位提升到 thread block,形成“从 HBM 读入、片上处理或 fusion、再写回”的统一路径。

术语消化:GPU kernel 基础词汇

  • Kernel:在 GPU 上执行的函数。
  • HBM:High Bandwidth Memory,GPU 的大容量全局显存,带宽高但仍远慢于片上 register/shared memory。
  • Thread block / CTA:一组线程,被调度到同一个 SM,共享 shared memory。
  • Warp:通常 32 个 lockstep 执行的 threads;分支不同会产生 control divergence。
  • Occupancy:SM 上活跃 warps/blocks 的比例,受 registers、shared memory、block size 限制。
  • Bank conflict:多个 threads 同时访问 shared memory 同一个 bank,访问被串行化。
  • Memory coalescing:warp 中线程访问连续 HBM cache line,合并成高效 memory transaction。
  • Tiling:把大矩阵分块搬到 shared memory 中复用。
  • Fusion:把多个操作合并进一个 kernel,减少 HBM 往返。

符号说明:kernel 性能公式

本讲出现的 \(N\) 表示向量或矩阵维度,\(M,N,K\) 在 matmul 中分别表示 \(A\in\mathbb{R}^{M\times K}\)\(B\in\mathbb{R}^{K\times N}\)\(C\in\mathbb{R}^{M\times N}\) 的维度;\(\mathrm{FLOPs}\) 是浮点操作数,bytes moved 是 HBM 读写量。Arithmetic intensity \(I=\mathrm{FLOPs}/\mathrm{bytes}\) 越高,越容易让 compute units 忙起来。

GPU recap:从硬件到编程模型

Source coverage: review_of_gpus. 本节覆盖的核心概念包括:SM/warp/block/shared memory/bank conflict/coalescing/occupancy。

Hardware:容量、带宽与可见层级

GPU 不是一块均匀的“快内存”。本节先对照 SM 数量、register、shared memory/L1、L2 与 HBM 的容量和带宽,建立 kernel 能实际使用的硬件账本;离计算单元越近的存储通常越快但越小,优化的本质是让高复用数据尽量停留在片上。

gpu hardware.png
gpu hardware.png 查看原图

读图:gpu hardware.png

这张图用于把抽象的 kernel 代码连接到硬件执行。读图时先看数据位于 HBM、shared memory 还是 registers,再看线程/blocks 如何映射到 SM。性能优化的关键是让数据在更靠近计算单元的位置被复用,减少慢速 HBM 往返。

Programming model:grid、CTA 与 thread

硬件层级需要一个可编程抽象。CUDA/PTX 把一次 kernel launch 组织为 grid,grid 由多个 thread blocks(也叫 CTA)组成,每个 block 再包含 threads;block 内线程共享同一块 shared memory,并作为整体被调度到一个 SM。

cuda grid.png
cuda grid.png 查看原图

读图:cuda grid.png

这张图用于把抽象的 kernel 代码连接到硬件执行。读图时先看数据位于 HBM、shared memory 还是 registers,再看线程/blocks 如何映射到 SM。性能优化的关键是让数据在更靠近计算单元的位置被复用,减少慢速 HBM 往返。

老师强调:为什么非逐元素算子需要 thread block

课程先用 GeLU 说明 elementwise 工作可自然地“一线程处理一个元素”,随后转向 softmax 与 matmul:这些算子中的 threads 必须交换局部结果,而 HBM 读写太慢。Thread block 的意义正是让一组 threads 共享同一块片上 shared memory;因此一个 block 会被调度到单个 SM,Triton 也选择把 block/program 作为原生思考单位。

Interaction between programming model and hardware

Programming model 足以描述正确结果,却不会自动保证高性能。同一份正确 kernel 会因为 warp divergence、register pressure、shared-memory bank conflict、HBM coalescing 与 block wave 切分而出现数量级差异;因此接下来要把抽象对象逐一映射回硬件限制。

block occupancy.png
block occupancy.png 查看原图

读图:block occupancy.png

这张图用于把抽象的 kernel 代码连接到硬件执行。读图时先看数据位于 HBM、shared memory 还是 registers,再看线程/blocks 如何映射到 SM。性能优化的关键是让数据在更靠近计算单元的位置被复用,减少慢速 HBM 往返。

课堂提示:occupancy 不是越高越好,tail 也要单独看

老师一方面提醒 register 用得越多,可同时驻留的 threads/warps 越少,但 low occupancy 不一定坏:thread coarsening 可能让每个 thread 做更多有效工作。另一方面,block occupancy 还受 wave quantization 影响;例如 148 个 SM 执行 160 个 blocks 时,第二波只有 12 个 blocks,大量 SM 空闲。性能判断必须同时看单 block 资源与跨 SM 的 waves,不能只追一个 occupancy 百分比。

本章小结

本节说明了一个 kernel optimization 的共同模式:先确认语义正确,再测时间,再看 profiler,最后用 block、shared memory、tiling、fusion 或 layout 调整降低数据移动。

Benchmarking 与 Profiling:先量再改

Source coverage: benchmarking/profiling. 本节覆盖的核心概念包括:benchmarking, profiling, CUDA synchronize, CUDA events, torch profiler。

老师强调:成功 recipe 是测量闭环

源码把流程写成三步:先 benchmark 与 profile,做出修改,再重新 benchmark 与 profile。Benchmark 只回答端到端 wall-clock 与 scaling,profiling 才回答时间花在哪里、底层实际调用了什么;没有第二轮测量,就无法证明改动解决了原先的瓶颈,也无法排除只是缓存、编译或输入 shape 改变造成的假提升。

方法 回答的问题 常见陷阱
Benchmark 这个操作端到端多快?随 shape 如何变化? 未同步 CUDA、warmup 不足、只测一次。
Profiler 时间花在哪些 kernels 上?实际调用了什么? 只看总时间,不看 kernel 名和 shape。
Nsight 更底层地看 occupancy、memory、warp stalls 信息量大,需要明确假设再查。
CUDA event benchmark 的关键结构
for _ in range(num_warmups):
    run()
torch.cuda.synchronize()
start_event.record()
run()
end_event.record()
torch.cuda.synchronize()
time_ms = start_event.elapsed_time(end_event)

benchmark 必须同步

CUDA kernel launch 是异步的。如果不调用 torch.cuda.synchronize() 或 CUDA events,测到的可能只是 CPU 发起 kernel 的时间,而不是 GPU 真正执行时间。

Profiler case matrix:三种 shape

课程把同一 profiler 放到三个 case 上,不是为了背 kernel 名,而是训练“高层算子 + shape 会选择不同底层实现”的直觉。逐元素加法主要暴露 launch 与 memory traffic;大矩阵乘法进入高吞吐 GEMM;小矩阵乘法的工作量不足以摊薄 launch、调度与 tile 边界开销。

Profiler case 先看什么 应形成的判断
add(dim=2048) kernel 数、CUDA 时间、HBM 读写 Pointwise 算子 FLOPs 少,常受 memory bandwidth 与 launch overhead 主导。
matmul(dim=2048) GEMM kernel、tile、dtype、总 CUDA 时间 尺寸足够大时,矩阵乘法可形成高 arithmetic intensity,接近 compute-bound。
matmul(dim=128) 是否换 kernel、每次 launch 时间占比 小 shape 即使算法仍是 \(O(n^3)\),实际延迟也可能由固定开销主导。

课堂提示:先看 scaling,再读 kernel 名

老师指出,小 dimension 时 timing 大致恒定,尺寸足够大后才显现 cubic scaling。Profiler 进一步显示 tensor dimensions 不同会调用不同 CUDA kernels;像 cutlass3x_sm100_..._64x64x16 这样的名字会暴露库、架构、dtype 与 tile shape。读 kernel 名的目的,是把测到的曲线连接到实际实现,而不是把长字符串当作噪声忽略。

本章小结

本节说明了一个 kernel optimization 的共同模式:先确认语义正确,再测时间,再看 profiler,最后用 block、shared memory、tiling、fusion 或 layout 调整降低数据移动。

GeLU case study:naive、builtin、compiled

Source coverage: naive_vs_builtin_vs_compiled_gelu. 本节覆盖的核心概念包括:kernel fusion, torch.compile, HBM reads/writes。

naive GeLU 的问题:多个小操作导致多次 HBM 往返
def naive_gelu(x):
    return 0.5 * x * (1 + torch.tanh(0.79788456 * (x + 0.044715 * x * x * x)))

Variant comparison:naive_gelu、builtin_gelu 与 compiled_gelu

三个版本计算同一个近似 GeLU,区别不在数学公式,而在最终生成多少 kernels、产生多少中间 tensors,以及每个中间结果是否必须往返 HBM。比较时必须固定输入 shape、dtype、device 与 steady-state 测量方法,再用 profiler 验证 kernel 数。

Executable source 中对应的两个基线标签是 \verb|builtin_gelu| 与 \verb|compiled_gelu|;保留这些原名,便于从 profiler 结果回查源码函数与 profile 节点。

Variant 典型执行形态 性能含义
naive_gelu 多个 primitive pointwise kernels 中间值多次写回和读回 HBM,没有 fusion。
builtin_gelu 框架提供的专用实现 通常合并为更少 kernels,减少 memory traffic。
compiled_gelu torch.compile 生成 Triton kernel 常把整段表达式 fusion 成一个 kernel,一次读、一次写。

为什么 builtin/compiled GeLU 更快

Naive PyTorch 写法会拆成多个 primitive kernels,中间结果反复写回 HBM。Builtin GeLU 或 torch.compile 后的 Triton kernel 能把多个 pointwise 操作融合成一次读、一次写,减少 HBM traffic。

老师强调:快的是执行图,不是 Python 写法更短

课上先观测 builtin 与 compiled 显著更快,再回到 profiler 解释原因:naive 版本产生多个 kernels 和多次 HBM 读写,而 builtin/compiled 版本只有一个 fused kernel;compiled 版本具体生成 Triton kernel。这个顺序很重要——先有 timing 证据,再用执行图归因,不能看到 torch.compile 就预设它一定更快。

本章小结

本节说明了一个 kernel optimization 的共同模式:先确认语义正确,再测时间,再看 profiler,最后用 block、shared memory、tiling、fusion 或 layout 调整降低数据移动。

Triton mental model

Source coverage: triton_introduction / triton_gelu_example. 本节覆盖的核心概念包括:Triton program, block, mask, pointer offsets, PTX。

课堂提示:Triton 把思考单位从 thread 提升到 block

老师对比 CUDA 与 Triton:CUDA 让程序员细粒度指定每个 thread 做什么,控制强但也要手动管理更多 shared-memory 与调度细节;Triton 让程序员描述一个 thread block/program 如何加载一块数据、在片上操作、再写回 global memory。对深度学习 kernel,这个抽象通常已经足够强,同时更接近 tensor 运算的数学结构。

cuda grid.png
cuda grid.png 查看原图

读图:cuda grid.png

这张图用于把抽象的 kernel 代码连接到硬件执行。读图时先看数据位于 HBM、shared memory 还是 registers,再看线程/blocks 如何映射到 SM。性能优化的关键是让数据在更靠近计算单元的位置被复用,减少慢速 HBM 往返。

Triton block-level 思维
pid = tl.program_id(axis=0)
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offsets < num_elements
x = tl.load(x_ptr + offsets, mask=mask)
tl.store(y_ptr + offsets, y, mask=mask)

Triton 与 CUDA 的差异

CUDA 通常让你描述每个 thread 做什么;Triton 更像描述一个 program/block 如何处理一块数据。它牺牲部分底层控制,换来更高层、更适合深度学习 kernel 的表达。

本章小结

本节说明了一个 kernel optimization 的共同模式:先确认语义正确,再测时间,再看 profiler,最后用 block、shared memory、tiling、fusion 或 layout 调整降低数据移动。

Triton softmax:row fits in one block

Source coverage: triton_softmax_example. 本节覆盖的核心概念包括:row-wise reduction, max subtraction, exp, normalization。

triton softmax.png
triton softmax.png 查看原图

读图:triton softmax.png

这张图用于把抽象的 kernel 代码连接到硬件执行。读图时先看数据位于 HBM、shared memory 还是 registers,再看线程/blocks 如何映射到 SM。性能优化的关键是让数据在更靠近计算单元的位置被复用,减少慢速 HBM 往返。

\[ \mathrm{softmax}(x_i)=\frac{e^{x_i-m}}{\sum_j e^{x_j-m}},\qquad m=\max_j x_j. \]

fused softmax 的资源意义

Naive softmax 需要 max、subtract、exp、sum、divide 多个 pass。若一整行能放进一个 Triton block,就可以在 shared memory/register 中完成 reduction,只读写 HBM 各一次左右。

本章小结

本节说明了一个 kernel optimization 的共同模式:先确认语义正确,再测时间,再看 profiler,最后用 block、shared memory、tiling、fusion 或 layout 调整降低数据移动。

Triton row sum:row larger than block

前面的 fused softmax 假设一整行能放进一个 program,本节故意放宽这个条件。Source coverage: triton_row_sum_example;当 row larger than block 时,需要沿 columns 做 tiling,把多个 tile 的局部和累积到 registers,最后再执行一次 reduction,同时处理尾块 mask 与数值类型。

课堂提示:先问一行能否放进一个 block

课程从 fused softmax 的前提反推限制:若整行能放进一个 block,Triton 可直接完成 block 内 reduction;若 4096 列而 block size 只有 1024,就必须把 row 拆成 4 个 tiles,让 threads 迭代 tiles、累积 partial sums,最后再对各 thread accumulator 做 reduction。这个容量问题决定了算法结构,而不是一个事后调参选项。

讲义提醒:先写正确的分块不变量

优化长 row 时先声明 accumulator 表示“已经处理过的所有 tiles 之和”,再移动 pointer。若循环边界、mask 或 dtype 不清楚,增加 warps 和 block size 只会更快地产生错误结果;Triton 调优必须从语义不变量开始。

triton row sum.png
triton row sum.png 查看原图

读图:triton row sum.png

这张图用于把抽象的 kernel 代码连接到硬件执行。读图时先看数据位于 HBM、shared memory 还是 registers,再看线程/blocks 如何映射到 SM。性能优化的关键是让数据在更靠近计算单元的位置被复用,减少慢速 HBM 往返。

row 不适合一个 block 时怎么办

如果一行太长,不能一次放进一个 block,就把行切成 tiles。每个 program 遍历多个 tiles,局部累加到 accumulator,最后再做 block 内 reduction。这是 tiling 思想的最小例子。

本章小结

本节说明了一个 kernel optimization 的共同模式:先确认语义正确,再测时间,再看 profiler,最后用 block、shared memory、tiling、fusion 或 layout 调整降低数据移动。

Triton matmul + ReLU:tiling and fusion

Row reduction 只沿一个轴累加,本节进一步处理二维矩阵乘。源码中的 matmul + ReLU 示例用 \(M\times N\) 输出 tile、沿 \(K\) 轴循环的 A/B tiles、tl.dot accumulator 和 fused ReLU 展示 shared-memory reuse;核心目标是每次从 HBM 取入的数据参与尽可能多的乘加。

讲义提醒:tile shape 同时受语义和硬件约束

BLOCK_M/N/K 不是越大越好。更大 tile 提高复用,却消耗更多 registers/shared memory,可能降低 occupancy;维度还要对齐 tensor-core 支持的 shape。正确流程是先过 correctness,再 sweep 配置并用 profiler 解释结果。

gemm tiled.png
gemm tiled.png 查看原图

读图:gemm tiled.png

这张图用于把抽象的 kernel 代码连接到硬件执行。读图时先看数据位于 HBM、shared memory 还是 registers,再看线程/blocks 如何映射到 SM。性能优化的关键是让数据在更靠近计算单元的位置被复用,减少慢速 HBM 往返。

Tiled matmul kernel 的核心循环
for k in range(0, K, BLOCK_K):
    a = tl.load(a_ptrs, mask=...)
    b = tl.load(b_ptrs, mask=...)
    acc += tl.dot(a, b)
    a_ptrs += BLOCK_K * stride_ak
    b_ptrs += BLOCK_K * stride_bk
acc = tl.maximum(acc, 0.0)  # fused ReLU

Matmul + ReLU fusion

如果下一步马上做 ReLU,就没有必要先把 matmul 输出写回 HBM,再读出来做 ReLU。把 ReLU 放在 matmul kernel 尾部,就是典型 fusion。

课堂提示:fusion 的收益要用省掉的 HBM traffic 解释

老师强调,融合 ReLU 的价值不是少写一行 PyTorch,而是避免把 matmul 输出完整写回 HBM、再由第二个 kernel 读回。若中间结果本来就在 registers,逐元素 activation 几乎是顺手完成;这也是判断哪些算子值得 fusion 的通用方法。

本章小结

本节说明了一个 kernel optimization 的共同模式:先确认语义正确,再测时间,再看 profiler,最后用 block、shared memory、tiling、fusion 或 layout 调整降低数据移动。

补充推导:从 profiling 到 Triton kernel 的完整工作流

GPU hardware ledger

硬件层级 典型容量/特征 优化含义
Registers 每个 thread 私有,极快 register 用太多会降低 occupancy。
Shared memory / L1 每个 SM 内,快但小 适合 tile、row reduction、softmax 等 block 内复用。
L2 cache 全 GPU 共享缓存 对跨 block 的重复访问有帮助,但不可完全依赖。
HBM GPU global memory,大容量高带宽 仍远慢于片上存储;kernel 优化要减少 HBM 往返。
Tensor cores 专用矩阵硬件 需要 shape/dtype/layout 匹配,才能达到高吞吐。

occupancy 不是越高越好

高 occupancy 能隐藏 HBM latency,但如果每个 thread 做更多 work、复用更多数据,较低 occupancy 也可能更快。正确问题不是“occupancy 是否最大”,而是瓶颈是不是因为可运行 warps 不够。

Benchmarking 的最小正确性条件

Benchmarking checklist
# 1. Warm up first to avoid compilation / cache effects.
for _ in range(num_warmups):
    run()
torch.cuda.synchronize()

# 2. Use CUDA events for GPU timing.
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
run()
end.record()
torch.cuda.synchronize()
time_ms = start.elapsed_time(end)

为什么第一次运行不能计入 benchmark

第一次运行可能包含 CUDA context 初始化、kernel compilation、allocator warmup、cache miss 等开销。我们关心 steady-state kernel performance,所以需要 warmup,并且要用多次 trial 看方差。

Profiling 如何读 kernel 名

Profiler 不只是给总时间,还暴露实际调用的 CUDA kernels。例如一个名字里可能出现:

  • cutlass:NVIDIA 线性代数 kernel library。
  • sm100:对应 Blackwell/B200 架构。
  • f32:输入/输出/accumulator dtype 信息。
  • 64x64x16:tile shape,暗示 block 内计算组织。

profile 的核心价值

Benchmark 告诉你慢,profile 告诉你为什么慢。比如同样是 matmul,小 shape 可能调用 SIMT kernel,大 shape 才调用 tensor-core-friendly GEMM;同样是 GeLU,naive 写法可能拆成多个 pointwise kernels,compiled 版本可能融合成 Triton kernel。

GeLU fusion 的 HBM 账本

Naive GeLU 公式是:

\[ \mathrm{GeLU}(x)\approx 0.5x\left(1+\tanh\left(0.79788456(x+0.044715x^3)\right) \right). \]

若每个子表达式都变成单独 kernel,可能产生多次 HBM read/write。Fused kernel 则把中间值保留在 registers 中:读一次 \(x\),计算多个子表达式,写一次 \(y\)

实现 kernel 形态 资源效果
Naive PyTorch 多个 pointwise kernels 多次 HBM 往返,launch overhead 多。
Builtin GeLU 库内 fused kernel 减少中间读写。
torch.compile 编译生成 fused Triton kernel 保留 Python 语义,获得接近手写 kernel 的数据移动模式。

Triton GeLU kernel 逐行解释

Triton GeLU kernel 核心结构
pid = tl.program_id(axis=0)
start = pid * BLOCK_SIZE
offsets = start + tl.arange(0, BLOCK_SIZE)
mask = offsets < num_elements
x = tl.load(x_ptr + offsets, mask=mask)
# compute gelu in registers
...
tl.store(y_ptr + offsets, y, mask=mask)

Triton program 的心智模型

每个 Triton program 类似一个 thread block,负责一段 offsets。Mask 处理尾部不满 block 的元素。tl.load 从 global memory 读入向量,后续计算在 registers 中完成,tl.store 写回。这个模型比 CUDA thread-level 更粗,但足以表达很多 ML kernels。

Softmax 为什么适合 row-wise block

Softmax 对每一行独立:先求最大值 \(m\),再计算指数和归一化。

\[ \mathrm{softmax}(x_i)=\frac{e^{x_i-m}}{\sum_j e^{x_j-m}},\qquad m=\max_j x_j. \]

若一行能放进一个 block,Triton 可以一次读入整行,在 block 内完成 max、exp、sum、divide,再写回输出。这样避免 naive 版本多次对同一行读写 HBM。

softmax 的数值稳定性

先减去行最大值 \(m\) 不是性能技巧,而是数值稳定技巧。否则 \(e^{x_i}\) 可能 overflow。性能优化不能破坏这种稳定性约束。

Row sum:当一行放不进一个 block

Row sum 是 softmax 的简化版,用来讲解“行太长”的情形。如果 \(N=4096\) 而 block size 是 1024,一行需要 4 个 tiles。每个 program 在 tile loop 中累加 partial sum,最后做 block 内 reduction。

Row sum 是 baby tiling

它没有 matmul 那么复杂,但已经包含 tiling 的核心:把大数据拆成可放进 block 的块,循环加载,局部累加,最后归约。这是理解 FlashAttention 和 tiled matmul 的前置台阶。

Tiled matmul + ReLU 的完整账本

Naive matmul 对每个输出元素 \(C_{mn}\) 遍历 \(k\),不断从 HBM 读 \(A_{mk}\)\(B_{kn}\)。相邻输出元素会重复用到同一行 A 或同一列 B。如果直接从 HBM 重复读,arithmetic intensity 低。

Tiling 做法是:

  1. 选择一个 \(C\) 的 output tile。
  2. 依次加载对应的 A tile 和 B tile 到 shared memory。
  3. tl.dot 在 tile 上累积 partial sums。
  4. 在写回 C 之前顺手做 ReLU,完成 fusion。

为什么 tile size 决定性能

Tile 太小,数据复用不够;tile 太大,shared memory/register 不够、occupancy 下降。高性能 GEMM kernel 的核心就是在 tile shape、register pressure、memory coalescing、tensor core layout 之间找平衡。

本章小结

这一节把源码中的 kernel examples 连接成完整工作流:GeLU 展示 pointwise fusion,softmax 展示 row-wise reduction,row sum 展示跨 tile accumulation,matmul+ReLU 展示 shared-memory tiling 和 activation fusion。四者共同说明:Triton 的价值在于让你以 block 为单位显式组织数据移动。

源码节点全覆盖索引

Lecture 6 是 executable source,真正的“slide-complete”不是逐页 PDF,而是逐个教学函数覆盖。下面这张表说明每个函数在讲义中的位置和处理方式。

源码函数 教学目标 讲义处理
review_of_gpus 复习硬件与编程模型 GPU 图、grid 图、block occupancy 图、术语表。
benchmarking 端到端计时 CUDA events、warmup、synchronize 代码。
profiling 看 kernel 级时间分布 profiler 表解释、kernel name 解读。
GeLU variants fusion 的性能收益 GeLU 公式、naive/builtin/compiled 对比。
Triton GeLU elementwise kernel block offsets、mask、load/store、PTX 观察。
Triton softmax row-wise reduction max-subtract-exp-sum-divide 的 fused row kernel。
Triton row sum row too large for one block tile loop、accumulator、final reduction。
Triton matmul ReLU tiled GEMM + fusion tile pointers、tl.dot、fused ReLU。

这张索引表的目的

它是本讲的 coverage matrix 的正文版本。读者可以看到:本讲不是泛泛讲 Triton,而是沿源码中的每个教学函数建立 benchmark → profile → optimize → write kernels 的完整链条。

更细的硬件解释:为什么这些细节会影响性能

Warps、divergence 与 occupancy

Warp 是 GPU 执行的基本调度粒度。一个 warp 内的 32 个 threads 通常执行同一条指令。如果一半 threads 走 if 分支 A,另一半走分支 B,硬件往往需要串行执行 A 和 B,这叫 control divergence。

Occupancy 则描述一个 SM 上有多少 warps 可以同时驻留。高 occupancy 能隐藏 memory latency,因为一个 warp 等 HBM 时,SM 可以切到另一个 ready warp。但 occupancy 太高不一定最好:如果每个 thread 使用更多 registers 来做更多 work,低 occupancy 也可能换来更高 arithmetic intensity。

occupancy 计算直觉

如果每个 thread 用 160 个 registers,一个 block 有 128 threads,那么一个 block 需要 \(160\times128=20480\) registers。若一个 SM 有 65536 registers,最多只能同时放下 3 个这样的 blocks。register pressure、block size 和 max warps 共同决定 occupancy。

Bank conflicts 与 swizzling

Shared memory 被分成多个 banks。若同一 cycle 内多个 threads 访问同一个 bank 的不同地址,访问会被串行化,形成 bank conflict。矩阵乘法中,读取 A 的行和 B 的列很容易造成不同访问模式,因此高性能 kernel 常通过 swizzling 改变 shared memory layout,减少冲突。

bank conflict 很隐蔽

代码语义完全正确,但性能可能因为 bank conflict 急剧下降。Profiler 或更底层的 Nsight 指标才能确认这类问题。只看 Python 层时间很难定位。

Memory coalescing 与 HBM transaction

HBM 访问通常以 cache line / transaction 为单位。若一个 warp 的 threads 访问连续地址,硬件可以合并成少量 transaction;若访问跨越很多不连续位置,就会浪费带宽。Coalescing 是 global memory 访问优化的第一原则。

coalescing 的一句话判断

同一个 warp 的相邻 threads 最好访问相邻地址。若 thread 0、1、2、... 访问的是矩阵同一行的连续元素,通常更 coalesced;若访问同一列且 row stride 很大,往往更差。

Triton kernel 参数表

参数/概念 出现位置 含义
BLOCK_SIZE GeLU、softmax、row sum 一个 program 处理多少元素或列。太小复用差,太大 register/shared memory 压力高。
tl.program_id 所有 Triton kernels 当前 program/block 的 id,用来决定处理哪一块数据。
offsets GeLU/softmax 当前 block 对应的全局元素下标。
mask 尾部处理 防止读写越界,是 block size 不整除数据长度时的必要保护。
stride softmax/matmul 描述逻辑二维 tensor 在一维内存中的跳步方式。
BLOCK_M/N/K matmul C tile 的行/列和 K 维 chunk 大小,决定 tile shape 和复用。
tl.dot matmul 调用 Triton 的 block-level dot,通常映射到底层矩阵硬件。

为什么这些参数不是随便调

Triton kernel 的参数直接映射到硬件资源:block size 影响并行粒度,mask 影响边界处理,stride 决定内存访问连续性,tile shape 影响 shared memory/register pressure 和 tensor core 使用。调参不是玄学,而是在硬件约束下找平衡。

扩展讲解:从源码到性能判断

为什么 GeLU 是第一个好例子

GeLU 是 elementwise operation,数学上不复杂,却非常适合说明 kernel fusion。因为 naive GeLU 由乘法、加法、立方、tanh、再乘法组成,如果框架把它拆成多个 kernels,每个中间 tensor 都要写回 HBM 再读出。

版本 执行方式 资源结果
naive PyTorch 多个 pointwise kernels launch 多,HBM 往返多。
builtin GeLU 手写/库内 fused kernel 读一次、算完、写一次。
compiled GeLU torch.compile 生成 Triton kernel 保持 Python 表达,获得 fusion。

elementwise kernel 的核心瓶颈

Elementwise 操作的 FLOPs 很少,bytes moved 却和 tensor 大小成正比。单独执行时通常 memory-bound。Fusion 的目标不是减少数学 FLOPs,而是减少中间结果在 HBM 中来回搬运。

Profiler 表应该怎么看

当 profiler 显示很多小 kernel 时,通常有三种解释:

  1. Python/PyTorch 代码被拆成多个 primitive operations。
  2. 每个 primitive operation 都 launch 一个独立 CUDA kernel。
  3. 中间结果 materialize 到 HBM,形成额外 memory traffic。

若 profiler 显示一个长名字的 CUTLASS/Triton kernel,则要读名字里的线索:架构代号、dtype、tile shape、layout alignment。这些信息帮助判断是不是 tensor core kernel、tile 是否过小、是否有 alignment 问题。

profile 后怎么行动

如果瓶颈是许多小 pointwise kernels,优先考虑 fusion;如果瓶颈是 GEMM 但 MFU 低,检查 shape、dtype、layout、tile;如果瓶颈是 memory copy 或 communication,优化 kernel 本身可能无效。

Triton masks 的必要性

很多 tensor 长度不能被 block size 整除。例如 \(N=10{,}000\)\(\texttt{BLOCK\_SIZE}=1024\),最后一个 block 会越界。Triton 用 mask 防止读写越界:

mask 保护尾部 block
offsets = start + tl.arange(0, BLOCK_SIZE)
mask = offsets < num_elements
x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
tl.store(y_ptr + offsets, y, mask=mask)

mask 是 correctness 条件,不只是性能细节

没有 mask,最后一个 block 可能读到非法地址或写坏输出。很多 Triton kernel 的第一类 bug 就来自 tail block 没处理好。

Softmax 的读写账本

Naive row softmax 至少经历这些阶段:row max、subtract max、exp、row sum、divide。若每步都是独立 kernel,读写次数近似为多次 \(MN\)。Fused row-wise softmax 则把一整行放进 block 中,局部完成 reduction 和 normalize。

阶段 naive 实现 fused Triton 实现
max 读整行,写 max block 内 reduction。
subtract/exp 再读整行,写中间结果 registers/shared memory 中完成。
sum 再读 numerator,写 denominator block 内 reduction。
divide 再读 numerator/denominator,写输出 只写最终输出。

Matmul tiling 的指针结构

Triton matmul kernel 不是直接写三重 for loop,而是构造指针矩阵:

  • indices_m 选择 C tile 的行。
  • indices_n 选择 C tile 的列。
  • indices_k 选择当前 K tile。
  • a_ptrsb_ptrs 指向当前要加载的 A/B tile。

每轮 K tile 做一次 tl.dot(a,b),累加到 acc。循环结束后可以在 acc 上做 ReLU,再写回 C。

Tiled matmul 的统一解释

Tiling 把 HBM 里的大矩阵切成可放入片上存储的小块,让一个 tile 内的数据服务多个输出元素。Fusion 则把 matmul 后的 ReLU 放在写回前完成。两者都在减少 HBM 访问。

本章小结

源码中的每个 kernel example 都在回答同一个问题:如何把高层 tensor expression 改写成更少 HBM traffic、更高片上复用、更适合 SM/warp/block 执行的程序。Triton 的价值在于让这种改写比 CUDA 更接近数学表达,但仍保留 block-level 控制。

Assignment 2 视角:从 correctness 到 performance

Lecture 6 对应的 systems assignment 不只是“写出能跑的 kernel”,而是要求你 benchmark 和 profile 实现。一个 kernel 的评价至少包含三层:

层级 问题 证据
Correctness 输出是否和 PyTorch reference 接近? torch.allclose、单元测试、极端 shape。
Performance 是否比 baseline 快?随 shape 如何 scaling? CUDA event benchmark,多次 trial,warmup。
Attribution 时间花在哪些 kernels 和 memory 操作上? PyTorch profiler、Nsight、kernel 名、occupancy/throughput 指标。

不要跳过 correctness

高性能错误 kernel 没有意义。Triton 的 mask、stride、边界条件、dtype accumulation 都可能造成 subtle bug。Assignment 的正确姿势是先让 reference 对齐,再 benchmark,再 profile,再优化。

性能比较要公平

比较 naive、builtin、compiled、Triton 版本时,必须使用相同输入 shape、相同 dtype、相同 device,并包含 warmup。否则测到的可能是编译时间、CPU launch overhead、allocator 行为或缓存状态,而不是 kernel 本身。

从本讲走向下一讲

本讲所有优化仍在单 GPU 内部。下一讲进入多 GPU 后,同样的原则会扩大:HBM 往返变成 GPU 间通信,thread blocks 变成 ranks 和 collectives,kernel fusion 的“少搬数据”变成 all-reduce、reduce-scatter、all-gather 的通信账本。

统一心智模型

单 GPU kernel 优化和多 GPU parallelism 都是在问:数据在哪里、谁需要它、什么时候移动、能不能少移动或多复用。Lecture 6 学的是这个问题在一个 GPU 内部的版本。

总结与延伸

本讲把 GPU kernel 工作流压缩成一条路径:理解硬件层级,写出正确 kernel,benchmark 测端到端时间,profile 找真实 kernel 和瓶颈,再用 Triton 调整 block、mask、tiling 和 fusion。

最终 takeaway

性能优化不是玄学。每一次优化都应该能回答:减少了哪次 HBM 读写?增加了多少复用?提高了 occupancy 还是降低了 bank conflict?benchmark 和 profiler 是否支持这个判断?

拓展阅读

下面的材料分别补足 Triton 算法、CUDA 硬件语义、GEMM 实现与测量工具。阅读时不要只抄 API,而要为每份材料记录“它解释了哪一次数据移动、哪一级并行和哪一个 profiler 指标”。

材料 重点问题 读完应能做什么
Triton fused softmax tutorial 一行如何映射到一个 program;mask、reduction 与 num warps 独立写出 row-fit-in-block 的 softmax,并数 HBM reads/writes。
NVIDIA CUDA programming guide warp、shared memory、coalescing、occupancy 用硬件术语解释一个正确 kernel 为什么仍然慢。
CUTLASS GEMM documentation tile hierarchy、tensor core shape、pipeline 从 kernel 名与 tile 配置反推数据复用和 resource pressure。
PyTorch profiler + Nsight guides end-to-end、kernel timeline、memory/warp stalls 把 timing 异常定位到 launch、kernel、memory 或 occupancy。

复现实验:从 shape sweep 到 profiler attribution

选择 GeLU、row sum 或 matmul 中任意一个算子,固定 dtype 与 device,执行一次小型 hold-out 实验:

  1. 预先写下随 shape 增长的延迟趋势,以及何时可能从固定开销区进入 bandwidth/compute 主导区。
  2. 对每个 shape 做 warmup 和多次 CUDA-event 计时,保留均值与方差。
  3. 只对转折点前后两个 shape 做 profiler,记录 kernel 数、kernel 名、CUDA time 与关键 memory/occupancy 指标。
  4. 修改一个变量,例如 fusion、block size 或 tile shape,再重复同一协议;若预测失败,优先修正心智模型而不是挑选对自己有利的样本。

Kernel 优化交付清单

最后把实验结果整理成一份可复核交付,而不是散落的截图和最快数字。每一项都应能指向代码、命令或 profiler 证据:

验收项 必须记录 不合格信号
Correctness reference、容差、边界 shape、mask 与 dtype 只测一个整齐 shape,或用速度掩盖数值错误。
Measurement warmup、trials、CUDA events、均值与方差 混入编译/CPU launch,只报告最快一次。
Scaling 多个 shape、转折点、预先预测 只在单点比 baseline,无法判断适用范围。
Attribution kernel 数/名、CUDA time、memory/occupancy 指标 只有 wall-clock,没有解释时间为何变化。
Mechanism 减少的 HBM bytes、增加的 reuse、改变的 tile/fusion 用“编译器更聪明”代替可检验机制。
Portability 至少一个 hold-out shape 或 dtype 配置只对调参样本有效,换输入立即退化。

实验验收不是“最快一次”

合格证据必须包含 reference correctness、完整 shape sweep、steady-state timing、profile attribution 与失败 case。只报告单个最佳数字,无法证明优化机制,也无法迁移到新 shape。

失败诊断:看到什么先查什么

当结果不符合预测时,先用最便宜的证据缩小范围,再打开更重的 Nsight 指标:

症状 第一证据 优先排查
小 shape 几乎一样慢 CUDA-event 延迟与 kernel 数 launch/dispatch 固定开销,工作量不足以占满设备。
大 shape scaling 突然变差 kernel 名、tile、memory throughput kernel selection 改变、tile 边界、bandwidth 或 workspace。
Fusion 后没有加速 profiler 的中间 kernel 与 HBM traffic 编译器未 fusion、额外 layout copy,或原算子并非 memory-bound。
增大 block/tile 反而变慢 registers、shared memory、active warps resource pressure 降低 occupancy,或 tail/wave quantization 恶化。
结果只在边界 shape 错 reference diff、mask、stride、accumulator dtype 越界 load/store、非连续 layout、reduction 精度或尾 tile。