第 6 讲：Kernel 与 Triton

目录 · ← l5 · l7 →

第 6 讲：Kernel 与 Triton

日期：4 月 15 日（周三，Spring 2026） | 讲师：Percy Liang | 材料：lecture_06.py | 截止：作业 1 到期、作业 2 发布

概览

这是一讲”动手写 kernel”的课：先用 benchmark 与 profiling 找出瓶颈，再用 Triton 写自定义 kernel 消除瓶颈。讲义由易到难开发四个 kernel——GeLU（逐元素）、softmax（行内归约）、row-sum（行超过一个 block 的归约）、matmul+ReLU（用共享内存分块）——并把第 5 讲的 GPU 编程模型（thread → thread block → grid）、占用率、bank conflict、合并访问与波量化落到具体代码上。

核心概念与定义

• Kernel：在 GPU 上运行的函数。用 PyTorch 时，每个基础算子都会启动一个标准 kernel；写自定义 kernel（CUDA/Triton/CUTLASS/ThunderKittens）可以融合与分块，让”GPU 起飞（go brrr）”。
• GPU 硬件表（讲义数据）：A100：108 SM、192KB L1+共享、40MB L2、80GB HBM、2TB/s；H100：132 SM、256KB、50MB L2、80GB、3.35TB/s；B200：148 SM、256KB、96–126MB L2、192GB、8TB/s。寄存器带宽是 HBM 带宽的 4–20 倍——所以”把数据留在寄存器里”。
• 编程模型：thread（在小片数据上执行）→ thread block / CTA（共享共享内存的一组线程，被调度到一个 SM 上）→ grid（线程块的集合）。逐元素算子天然映射到线程；归约/矩阵乘需要线程块，因为线程之间必须通过共享内存通信。
• Triton 的模型：你描述的是每个线程块要做什么（而 CUDA 描述每个线程）：把 tile 从 global memory 载入共享内存、计算、写回。Triton 编译到 PTX（GPU 汇编）。
• Warp：32 个线程锁步执行；控制发散（warp 内 if/else）会串行执行；某个 warp 因访存阻塞时，SM 可以零成本切换到其他 warp。
• 占用率（occupancy）：SM 上可同时驻留的 warp 数，受寄存器（每线程 0–255）、共享内存等限制。低占用率不一定坏，如果每个线程做更多事（thread coarsening）。例：128 线程 × 160 寄存器 = 每块 20480 寄存器 → 65536/20480 = 3 个块。
• Bank conflict（共享内存）：32 个 bank × 4 字节，每周期每 bank 一次访问。32 个线程撞同一个 bank（例如读矩阵的列）→ 32 路串行。Swizzling（行列异或）重排地址可避免冲突。
• 显存合并（HBM）：warp 的 32 次访问若连续，则合并为一个 128 字节事务；完全合并 = 32 线程 × 4 字节一次事务。
• 波量化：线程块按波次填入 SM；148 个 SM 上跑 160 个块 → 148 + 12（第二波大部分闲置）。解法：让块数整除 SM 数。
• Benchmark 与 profiling：benchmark 测端到端墙钟（用于比较实现、研究扩展性）；profiling 显示哪些 kernel 在执行、各花多久（PyTorch profiler、nsight）。kernel 名字本身就泄露实现细节：cutlass3x_sm100_simt_sgemm_f32_..._64x64x16 = CUTLASS 库、Blackwell（sm100）、float32、64×64×16 tile。
• Kernel 融合：朴素 GeLU 会启动多个 kernel（多次 HBM 往返）；融合/builtin/torch.compile 版本只跑一个 kernel（一次读、一次写）——对 memory-bound 的逐元素运算收益巨大。

代码示例：benchmark 与 profiling 框架

代码（Python）：

def benchmark(run: Callable, num_warmups: int = 1, num_trials: int = 3) -> float:
    for _ in range(num_warmups):
        run()
    torch.cuda.synchronize()          # 关键：冲刷异步 CUDA 任务

    times: list[float] = []
    for trial in range(num_trials):
        start_event = torch.cuda.Event(enable_timing=True)
        end_event = torch.cuda.Event(enable_timing=True)
        start_event.record()          # GPU 侧时间戳
        run()
        end_event.record()
        torch.cuda.synchronize()
        times.append(start_event.elapsed_time(end_event))
    return mean(times)

def profile(run: Callable, num_warmups: int = 1):
    for _ in range(num_warmups):
        run()
    torch.cuda.synchronize()
    with torch.profiler.profile(activities=[ProfilerActivity.CUDA]) as prof:
        run()
        torch.cuda.synchronize()
    return prof.key_averages().table(sort_by="cuda_time_total", row_limit=10)

代码做了什么： benchmark 包装器先做 warmup（编译/JIT），用 CUDA events 做 GPU 级精确计时（避开 CPU 启动开销），并对多次试验取平均。profiler 则给出按 CUDA 时间排序的逐 kernel 表格。

实现深挖：

• 为什么要 warmup + synchronize：前几次启动可能触发编译；CUDA 是异步的，不同步就会测到启动延迟而非 kernel 时间。CUDA events 在 GPU 上打时间戳，排除 CPU 开销。
• 为什么要多次试验：kernel 时间有抖动（时钟、显存状态）；平均可降噪。做扩展性研究时按维度扫（256→8192）：小矩阵受启动开销支配（时间近似恒定），大矩阵呈现立方增长。
• 为什么 profiling 关键：朴素 vs builtin vs compiled GeLU 的对比揭示了为什么某个更快：profiler 显示是”许多 kernel 启动（未融合）”还是”一个 kernel”。

与作业的联系：作业 2 第一部分正是这件事：为你的作业 1 模型搭建 benchmark + profiling 框架（含 Nsight Compute 与 NVTX range），报告逐 kernel 运行时间，并回答”前向+反向哪个 kernel 占主导”。讲义中的 run_operation1/2、warmup、CUDA event 模式就是参考实现。

代码示例：Triton GeLU（逐元素）

代码（Python）：

import triton
import triton.language as tl

def triton_gelu(x: torch.Tensor):
    assert x.is_cuda and x.is_contiguous()
    y = torch.empty_like(x)
    num_elements = x.numel()
    BLOCK_SIZE = 1024
    num_blocks = triton.cdiv(num_elements, BLOCK_SIZE)   # 向上取整除法
    triton_gelu_kernel[(num_blocks,)](x, y, num_elements, BLOCK_SIZE=BLOCK_SIZE)
    return y

@triton.jit
def triton_gelu_kernel(x_ptr, y_ptr, num_elements, BLOCK_SIZE: tl.constexpr):
    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)      # 从 HBM 读
    # tanh(a) = (exp(2a) - 1) / (exp(2a) + 1)，因为 tl.tanh 不存在
    a = 0.79788456 * (x + 0.044715 * x * x * x)
    exp = tl.exp(2 * a)
    tanh = (exp - 1) / (exp + 1)
    y = 0.5 * x * (1 + tanh)
    tl.store(y_ptr + offsets, y, mask=mask)      # 写回 HBM

代码做了什么： 每个 1024 元素一个线程块；每个块算出自己的索引范围、带边界掩码地载入、计算 GeLU 的 tanh 近似（因为 Triton 没有 tl.tanh，用 exp 重新实现），再写回——全在一个 kernel 内：一次 HBM 读、一次 HBM 写。

实现深挖：

• 为什么要 mask：num_elements 未必整除 BLOCK_SIZE；掩码避免越界访问（否则是静默的数据损坏 bug）。
• 为什么 tl.constexpr：编译期常量 → Triton 会特化/展开；网格大小是运行期量。triton.cdiv 是向上取整除法，保证最后一个块覆盖尾部。
• 为什么自己实现 tanh：Triton 语言算子集有限；(e^{2a}−1)/(e^{2a}+1) 是标准替代写法。这很好地体现了”用 Triton 写 kernel”的取舍：控制力不如 CUDA，但样板代码少得多。
• 为什么一个线程处理 8 个元素：生成的 PTX 显示 thread coarsening——Triton/编译器做了向量化，让一个线程处理多个元素，提高指令级并行与访存吞吐。

与作业的联系：作业 2 要求实现融合的 RMSNorm Triton kernel，用的正是这个模式（分块逐元素 + mask + 单次读/写）。读 PTX 的练习（ld.global/st.global、%ctaid.x、%tid.x）就是验证 kernel 实际行为的方式。

代码示例：Triton softmax（行内归约）与 row-sum（分片循环）

代码（Python）：

@triton.jit
def triton_softmax_kernel(x_ptr, y_ptr, x_row_stride, y_row_stride, num_cols, BLOCK_SIZE: tl.constexpr):
    assert num_cols <= BLOCK_SIZE
    row_idx = tl.program_id(0)                    # 一行一个块
    col_offsets = tl.arange(0, BLOCK_SIZE)
    x_ptrs = x_ptr + row_idx * x_row_stride + col_offsets
    x_row = tl.load(x_ptrs, mask=col_offsets < num_cols, other=float("-inf"))
    x_row = x_row - tl.max(x_row, axis=0)         # 减去行最大值（数值稳定）
    numerator = tl.exp(x_row)
    denominator = tl.sum(numerator, axis=0)
    y_row = numerator / denominator
    tl.store(y_ptr + row_idx * y_row_stride + col_offsets, y_row, mask=col_offsets < num_cols)

@triton.jit
def row_sum_kernel(x_ptr, out_ptr, N, BLOCK_SIZE: tl.constexpr):
    row = tl.program_id(0)
    acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32)   # 每线程一个累加器
    for start in range(0, N, BLOCK_SIZE):            # 遍历列 tile
        cols = start + tl.arange(0, BLOCK_SIZE)
        x = tl.load(x_ptr + row * N + cols, mask=cols < N, other=0.0)
        acc += x
    result = tl.sum(acc, axis=0)                     # 线程间最终归约
    tl.store(out_ptr + row, result)

代码做了什么： softmax kernel 每行一个线程块：载入该行（带掩码、用 -inf 填充，使填充列的 exp(-inf)=0）、减最大值、取指数、求和、归一化、写回——单 kernel 行 softmax。row-sum kernel 处理比一个块还长的行：每个线程跨列 tile 累加，最后用 tl.sum 归约各线程的局部和。

实现深挖：

• 为什么 softmax 用 other=float("-inf")：填充位置必须在求和中等价于 exp(−∞)=0，且不能影响最大值。row-sum 则用 other=0.0。
• 为什么需要循环：assert num_cols <= BLOCK_SIZE 是”行能放进一个块”的假设；当行有 4096 列而 BLOCK_SIZE 为 1024 时，就要按 tile 迭代累加——这是”初级分块”（对归约做分块），也是矩阵乘分块的铺垫。
• 为什么 tl.sum(acc, axis=0)：循环结束后，BLOCK_SIZE 个线程各持有覆盖自己那部分列的部分和；块级归约（共享内存/ warp shuffle，Triton 内部处理）得到标量行和。
• 成本核算：朴素的 PyTorch softmax 约需 5MN 次读 + 3MN 次写（max、减、exp、sum、除）；融合后的 Triton kernel 只需 MN 读 + MN 写——最多减少约 4 倍显存事务，这对 memory-bound 的 softmax 至关重要。

与作业的联系：这是作业 2 的 FlashAttention-2 的结构模板：分块 + 带掩码载入 + 在线 softmax 累加（第 5 讲的技巧，这里用 other=-inf 掩码实现）。理解这个归约循环，才能理解 attention 中那个 O 累加器（形状 [BLOCK_M, head_dim] 的滑动加权和）为什么这样写。

代码示例：Triton matmul + ReLU（共享内存分块）

代码（Python）：

@triton.jit
def matmul_relu_kernel(
    a_ptr, b_ptr, c_ptr, M, N, K,
    stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    indices_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    indices_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    indices_k = tl.arange(0, BLOCK_K)

    # A、B tile 的指针网格
    a_ptrs = a_ptr + indices_m[:, None] * stride_am + indices_k[None, :] * stride_ak
    b_ptrs = b_ptr + indices_k[:, None] * stride_bk + indices_n[None, :] * stride_bn

    acc = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
    for k in range(0, K, BLOCK_K):                       # 沿 K 维遍历 tile
        a = tl.load(a_ptrs, mask=(indices_m[:, None] < M) & (indices_k[None, :] + k < K), other=0.0)
        b = tl.load(b_ptrs, mask=(indices_k[:, None] + k < K) & (indices_n[None, :] < N), other=0.0)
        acc += tl.dot(a, b)                              # 对 tile 做 tensor core 矩阵乘
        a_ptrs += BLOCK_K * stride_ak                    # 推进到下一个 tile
        b_ptrs += BLOCK_K * stride_bk

    acc = tl.maximum(acc, 0.0)                           # 融合 ReLU！
    c_ptrs = c_ptr + indices_m[:, None] * stride_cm + indices_n[None, :] * stride_cn
    tl.store(c_ptrs, acc, mask=(indices_m[:, None] < M) & (indices_n[None, :] < N))

代码做了什么： 把 C 划分为 BLOCK_M × BLOCK_N 的输出 tile（每个 tile 一个线程块，二维网格）。每个块沿 K 维以 BLOCK_K 为步长循环：载入 A tile 与 B tile（带边界掩码），把 tl.dot(a, b) 累加进输出累加器，最后在 kernel 内先做 ReLU 再写回。

实现深挖：

• 为什么要分块：朴素矩阵乘对每个 (m,n) 都要从 HBM 读 A[m,k] 和 B[k,n]——O(M·K·N) 次读，算术强度 O(1)。把 tile 载入共享内存、在每个 K 分片上复用给整个输出 tile，强度提升到 O(tile_size)；讲义”理想化”版本（全部放共享内存）可达 O(N)。
• 为什么每次载入/写回都要 mask：网格按 triton.cdiv(M, BLOCK_M) × triton.cdiv(N, BLOCK_N) 大小构造，边缘 tile 会越界 M/N/K；(用 other=0.0) 掩码让越界贡献为 0——这对矩阵乘是数学上正确的填充。
• 为什么用 tl.dot：它会下沉为 tensor core 指令（第 5 讲”矩阵乘比其他浮点运算快 10 倍以上”的硬件）。Triton 会替你选择 MMA 布局；在 CUDA 里你得手动管理 fragment。
• 为什么融合 ReLU：在写回前做 tl.maximum(acc, 0.0) 避免第二次 HBM 往返（读 C、激活、写 C）——融合原则用于矩阵乘。FlashAttention 把 softmax 融进 attention 矩阵乘也是同一手法。
• 为什么用带步长的指针：矩阵在内存中是线性化的；index = row*stride_row + col*stride_col。用 stride（而不是硬编码形状）让 kernel 支持非连续布局，也让同一个 kernel 服务于转置操作数。

与作业的联系：作业 2 的 FlashAttention-2 Triton kernel 就是同一结构（BLOCK_M × BLOCK_N 输出 tile、K 维循环配 tl.dot、掩码），再加上在线 softmax、mask/bias 处理，以及重算注意力分数的反向传播。FlashAttention 的前向”确实就是 KQV 矩阵乘的分块”（第 5 讲）——softmax 是额外的那部分。

关键要点

1. 配方是：benchmark → profile → 改 → 再 benchmark。benchmark 给端到端时间；profiling 告诉你哪些 kernel 占主导以及它们叫什么名字。
2. 理解硬件（SM、warp、占用率、bank conflict、合并访问、波量化）——即使代码正确，它仍决定性能。
3. Triton 的思维单位是线程块而非线程：从 HBM 载入 tile → 在共享内存/寄存器里计算 → 写回 HBM。把逐元素算子融进 kernel 以避免往返。
4. 四类典型 kernel：逐元素（GeLU）、行归约（softmax）、分片归约（row sum）、分块矩阵乘（matmul+ReLU）——真实 kernel 都是它们的组合。
5. 融合与分块是把 memory-bound 代码变成 compute-bound 代码的两大杠杆；HBM 流量下降时 MFU 就上升。

常见陷阱

• 忘记掩码：越界的读/写会造成静默损坏或崩溃；掩码的 other= 填充值必须对当前算子是数学中性的。
• benchmark 里漏掉 torch.cuda.synchronize()：你会测到异步启动开销而不是 kernel 时间。
• 填充值用错：softmax 用 -inf（使 exp→0），matmul/sum 用 0.0——用错会在边界上给出错误结果。
• tile 布局导致 bank conflict：朴素的共享内存布局会在同一 bank 上串行；按真实 FlashAttention 实现那样做 swizzling（行列异或）。
• 网格尺寸不整除 SM 数：波量化会让 SM 闲置；需要调块数。
• 非连续输入：kernel 往往断言 is_contiguous()；直接传入转置张量而不处理 stride 会破坏合并访问或正确性。
• 精度漂移：即使输入是 bf16，也要用 fp32 累加（acc = tl.zeros(..., dtype=tl.float32)）；作业也强调结果要在容差内匹配 PyTorch 参考实现。

复习题

1. 问： 融合后的 Triton GeLU 与朴素 PyTorch 表达式 0.5*x*(1+tanh(...)) 数学上相同，为什么快这么多？
   • 答： 朴素版本会启动很多 kernel（乘、加、tanh、加、乘），每个都要把整个张量读写一遍 HBM。融合 kernel 只做一次读、在寄存器里算完、一次写——对逐元素运算而言 HBM 流量（瓶颈所在）降低约 4–5 倍。
2. 问： 在矩阵乘 kernel 中，为什么用 other=0.0 掩码载入 A 能在 M 不整除 BLOCK_M 时保持正确？
   • 答： 填充行的点积贡献为 0，所以越界行的累加器是 0；最后的带掩码写回根本不写这些行。0 是加法单位元，使填充在数学上无害。
3. 问： 什么是”波量化”？如果 kernel 在 N=1793 时性能骤降，你会怎么修？
   • 答： 线程块按波次调度到 SM；若 tile 数不整除 SM 数，最后一个波次会部分闲置。修法：调整 tile 大小/块数使网格整除 SM 数，或对问题做 padding 使 tile 对齐。