CS336 · 第 6 课

Kernels 与 Triton 编程

上一讲说明了 FLOPs 和 HBM bytes 为什么共同决定性能；这一讲进一步回答程序员能做什么：怎样测出真实瓶颈，并通过 fusion、tiling 和更合适的 kernel 编程层次减少数据搬运。

这一讲到底解决什么问题？

硬件规格只能给出上限，无法告诉我们当前 workload 卡在 launch overhead、HBM traffic、occupancy 还是计算单元。优化必须从可重复测量开始，否则 fusion 或更复杂的 kernel 可能只是增加代码量。

PyTorch ops → kernel graph → launches + HBM intermediates
                           ↓ profile
fusion → fewer launches / fewer HBM round trips
tiling → more on-chip reuse
masking → boundary correctness

correctness baseline → warm-up → synchronized timing → profiler → hypothesis

Benchmarking 与 Profiling

为什么要 benchmark？

Benchmark 的正确姿势

import time

def benchmark(func, num_warmups=1, num_trials=5):
    """正确的 benchmark 方法。"""
    # 1. 预热: 让 JIT 编译、缓存预热
    for _ in range(num_warmups):
        func()
    torch.cuda.synchronize()  # 等待 GPU 完成！

    # 2. 正式计时
    times = []
    for _ in range(num_trials):
        start = time.time()
        func()
        torch.cuda.synchronize()  # 关键！GPU 是异步的
        times.append(time.time() - start)

    return sum(times) / len(times)

Profiling：看时间花在哪

from torch.profiler import ProfilerActivity

with torch.profiler.profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    with_stack=True,
) as prof:
    loss = model(x).mean()
    loss.backward()

# 打印时间排名
print(prof.key_averages().table(sort_by="cuda_time_total"))

# 导出 flame graph（可视化调用栈）
prof.export_stacks("stacks.txt", "self_cuda_time_total")

Benchmarking 能回答什么？

• 哪个实现更快？比较不同实现
• 性能如何随规模变化？理解 scaling behavior
• 是否线性加速？batch size 翻倍，时间是否减半？

Kernel Fusion

问题：每个操作都读/写 HBM

# 假设计算 GeLU(A @ B)
# 不 fused 的版本:

# Kernel 1: matmul
tmp = A @ B          # 读 A, B from HBM; 写 tmp to HBM

# Kernel 2: gelu
output = gelu(tmp)   # 读 tmp from HBM; 写 output to HBM

# 问题: tmp 被写入 HBM，然后立即读回——浪费！

解决：融合成一个 kernel

Benchmark 对比

PyTorch 自动 fusion

# torch.compile 可以自动融合
@torch.compile
def fused_operation(x, w):
    return gelu(x @ w)  # 自动融合 matmul + gelu

# 但手动写 kernel 仍然有价值:
# 1. 更激进的融合（跨多个 matmul）
# 2. 自定义操作（FlashAttention）
# 3. 精确控制 shared memory

# Grid: 所有 thread blocks 的集合
# Thread Block: 一组共享 shared memory 的 threads
# Thread: 基本执行单位

# 启动 kernel
kernel_name<<<num_blocks, block_size>>>(args)

# 例如: 1024 个 thread, block_size=256
# → num_blocks = 4
# → 4 个 block, 每个 256 个 thread

// gelu.cu
#include <cuda_fp16.h>

__global__ void gelu_kernel(
    const half* __restrict__ x,
    half* __restrict__ y,
    int num_elements
) {
    // 每个 thread 处理一个元素
    int i = blockIdx.x * blockDim.x + threadIdx.x;

    if (i < num_elements) {
        float xi = __half2float(x[i]);
        // tanh approximation of GELU
        float a = 0.79788456f * (xi + 0.044715f * xi * xi * xi);
        float exp_2a = expf(2.0f * a);
        float tanh_a = (exp_2a - 1.0f) / (exp_2a + 1.0f);
        float result = 0.5f * xi * (1.0f + tanh_a);
        y[i] = __float2half(result);
    }
}

from torch.utils.cpp_extension import load_inline

cuda_source = open("gelu.cu").read()
cpp_source = "torch::Tensor gelu(torch::Tensor x);"

module = load_inline(
    cuda_sources=[cuda_source],
    cpp_sources=[cpp_source],
    functions=["gelu"],
    extra_cflags=["-O2"],
    name="inline_gelu",
    build_directory="var/cuda_gelu",
)

# 使用
y = module.gelu(x)

Triton 入门

Triton 是什么？

Triton 是 OpenAI 2021 年发布的 GPU 编程语言——用 Python 语法写 GPU kernel。

Triton 编程模型

import triton
import triton.language as tl

@triton.jit
def gelu_kernel(
    x_ptr, y_ptr,
    num_elements,
    BLOCK_SIZE: tl.constexpr,
):
    # 每个 program (block) 处理 BLOCK_SIZE 个元素
    pid = tl.program_id(0)
    block_start = pid * BLOCK_SIZE

    # 这个 block 处理的元素索引
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < num_elements  # 边界保护

    # 读取
    x = tl.load(x_ptr + offsets, mask=mask)

    # 计算 (和标准 GeLU 一样)
    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)

def triton_gelu(x):
    y = torch.empty_like(x)
    num_elements = x.numel()
    block_size = 1024
    num_blocks = triton.cdiv(num_elements, block_size)
    gelu_kernel[(num_blocks,)](x, y, num_elements, BLOCK_SIZE=block_size)
    return y

Triton 的优势

• Python 语法：不需要写 C++/CUDA
• 自动优化：memory coalescing、shared memory 管理
• 可调试：可以 step through Python 代码
• 性能：接近甚至超过手写 CUDA

Triton 实战：Fused Softmax

Naive Softmax 的问题

# Naive 实现: 5 次遍历数据
def manual_softmax(x):
    x_max = x.max(dim=1)[0]      # 遍历 1: 求 max
    x = x - x_max[:, None]        # 遍历 2: 减 max
    numerator = torch.exp(x)      # 遍历 3: exp
    denominator = numerator.sum(dim=1)  # 遍历 4: sum
    y = numerator / denominator[:, None]  # 遍历 5: 除法
    return y

# 5MN 次读取 + 3MN 次写入
# 理论上只需要 MN 读 + MN 写（4× 差距！）

Fused Softmax Kernel

@triton.jit
def softmax_kernel(
    x_ptr, y_ptr,
    x_row_stride, y_row_stride,
    num_cols,
    BLOCK_SIZE: tl.constexpr,
):
    # 每个 program 处理一行
    row_idx = tl.program_id(0)
    col_offsets = tl.arange(0, BLOCK_SIZE)

    # 读取一行（从 HBM 到 SRAM）
    x_start_ptr = x_ptr + row_idx * x_row_stride
    x_ptrs = x_start_ptr + col_offsets
    x_row = tl.load(x_ptrs, mask=col_offsets < num_cols, other=float("-inf"))

    # 在 SRAM 中完成所有计算
    x_row = x_row - tl.max(x_row, axis=0)        # 减 max
    numerator = tl.exp(x_row)                     # exp
    denominator = tl.sum(numerator, axis=0)       # sum
    y_row = numerator / denominator                # 除法

    # 写回 HBM
    y_start_ptr = y_ptr + row_idx * y_row_stride
    y_ptrs = y_start_ptr + col_offsets
    tl.store(y_ptrs, y_row, mask=col_offsets < num_cols)

性能对比

Tiling（分块）与 Shared Memory

为什么需要 Tiling？

矩阵乘法 C = A × B，朴素实现每个元素独立计算：

Tiled Matmul

Triton 的自动 Tiling

# Triton 的 matmul kernel 自动处理 tiling:
@triton.jit
def matmul_kernel(A, B, C, M, N, K, BLOCK_M: tl.constexpr, ...):
    # Triton 自动:
    # - 选择 tile 大小
    # - 管理 shared memory
    # - 处理边界条件
    # - 优化内存访问模式

# 你只需要指定 block 大小，编译器处理细节

L2 Cache 优化

不同 tile 处理顺序对 L2 cache 命中率影响巨大：

# Triton 生成的 PTX 示例（GeLU kernel）:
ld.global.v4.f32  {%f0, %f1, %f2, %f3}, [%rd0];  # 读 4 个 float
fma.rn.f32        %f4, %f0, %f5, %f6;            # 融合乘加
ex2.approx.f32    %f7, %f8;                      # 近似 exp
st.global.v4.f32  [%rd1], {%f0, %f1, %f2, %f3};  # 写 4 个 float

本讲知识地图：5 种实现层次

趋势：自动编译器（Triton、torch.compile）在快速进步。手写 kernel 的价值在于：(1) 自定义操作；(2) 极端性能优化；(3) 理解底层。

知识检查

如果只能记住 5 件事

1. 优化从同步正确的 benchmark 和 profiler 时间线开始；没有基线就无法判断改动是否有效。
2. Kernel fusion 的主要收益是消除中间张量的 HBM 写回与再次读取，而不是减少数学 FLOPs。
3. Tiling 把可复用数据搬到 registers / shared memory，让一次 HBM load 服务更多计算。
4. Triton 提供 block-level 编程模型，但 tile、layout、num warps 等选择仍需针对目标 shape 和 GPU 测量。
5. Occupancy、coalescing、bank conflict 与 wave quantization 是解释结果的机制，不应脱离实际 kernel 单独追求某个数字。

延伸阅读