第 6 讲：GPU Kernels 与 Triton

1. 什么是 Kernel

在 GPU 编程中，Kernel 是在 GPU 上执行的函数。 当你在 PyTorch 中写 y = x + 1（x 是 GPU tensor）时， PyTorch 内部会启动一个 CUDA kernel 来执行这个加法操作。 每个 PyTorch 操作（如 torch.matmul、F.softmax、F.relu） 背后都对应一个或多个 CUDA kernel。

理解和编写自定义 kernel 的能力对于 LLM 系统优化至关重要。 标准的 PyTorch 操作虽然方便，但无法实现某些关键的优化（如 kernel fusion）。

1.1 为什么需要自定义 Kernel

PyTorch 的默认操作存在以下局限：

• 每个操作一个 kernel：y = relu(layer_norm(x + bias)) 会启动 3 个独立的 kernel，每个都需要从 HBM 读取输入、写回输出
• 无法利用操作间的数据局部性：上一个 kernel 的输出写回 HBM 后，下一个 kernel 又得重新读取
• Kernel launch overhead：每次 kernel 启动有 5-10μs 开销

2. CUDA Kernel 基础

传统的 GPU kernel 使用 CUDA C/C++ 编写。以下是一个简单的 vector addition kernel：

// CUDA C++ kernel for vector addition
__global__ void vector_add(float* A, float* B, float* C, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        C[idx] = A[idx] + B[idx];
    }
}

// Launch: vector_add<<>>(A, B, C, N);

CUDA kernel 的关键特点：

• 每个 thread 执行相同的代码，但通过 threadIdx 和 blockIdx 确定自己处理哪部分数据
• 程序员需要手动管理 shared memory、内存合并、同步等底层细节
• 学习曲线陡峭，调试困难

对于复杂操作（如矩阵乘法、attention），手写高效 CUDA kernel 需要数百到数千行代码， 涉及大量的底层优化。这就是 Triton 诞生的背景。

3. Triton 编程语言介绍

Triton 是 OpenAI 开发的 GPU 编程语言，设计目标是让研究者能够用 接近 Python 的语法编写接近手写 CUDA 性能的 GPU kernel。

3.1 Triton vs CUDA

方面	CUDA C++	Triton
语言	C++	Python-like
编程粒度	单个 thread	Block of threads（一次操作一个 tile）
内存管理	手动 shared memory	编译器自动优化
内存合并	手动确保	自动
调试难度	高	中等
性能上限	最高	接近 CUDA（通常 >80%）
开发效率	低	高（代码量减少 5-10 倍）

3.2 Triton 的核心思想：Block-level Programming

CUDA 中，程序员以单个 thread 的视角编写代码。 Triton 中，程序员以一组 threads（一个 block）的视角编写代码， 一次处理一个数据块（tile）。编译器负责将 block-level 操作映射到具体的 threads。

import triton
import triton.language as tl

@triton.jit
def vector_add_kernel(
    A_ptr, B_ptr, C_ptr,
    N,
    BLOCK_SIZE: tl.constexpr,
):
    # Each program instance processes one block of data
    pid = tl.program_id(0)

    # Compute the range of elements this block handles
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

    # Create a mask for boundary handling
    mask = offsets < N

    # Load a block of data from A and B
    a = tl.load(A_ptr + offsets, mask=mask)
    b = tl.load(B_ptr + offsets, mask=mask)

    # Compute element-wise addition
    c = a + b

    # Store the result
    tl.store(C_ptr + offsets, c, mask=mask)

4. 内存合并访问（Memory Coalescing）

GPU 的 HBM 以128 字节为单位进行读写（一个 cache line）。 如果一个 Warp 的 32 个 threads 访问连续的内存地址， GPU 可以将这 32 个请求合并为少量的 cache line 读取——这就是 memory coalescing。

# Coalesced access: threads access consecutive addresses
# Thread 0 reads A[0], Thread 1 reads A[1], ..., Thread 31 reads A[31]
# → 1 memory transaction (32 × 4 bytes = 128 bytes = 1 cache line)

# Strided access: threads access non-consecutive addresses
# Thread 0 reads A[0], Thread 1 reads A[100], ..., Thread 31 reads A[3100]
# → Up to 32 memory transactions (each hits a different cache line)
# → 32× slower!

在 Triton 中，编译器会自动处理 coalescing（前提是你的访问模式是合理的）。在 CUDA 中，这需要程序员手动确保。

5. Tiling 策略

Tiling（分块）是 GPU 编程中最核心的优化策略之一。 基本思想是：将大的计算任务分解为小的"块"（tiles）， 每次将一个 tile 加载到高速的 shared memory / registers 中进行计算， 然后写回结果。这样可以最大化数据复用，减少 HBM 访问。

5.1 矩阵乘法的 Tiling

以矩阵乘法 C = A × B（A: M×K, B: K×N）为例。 朴素实现中，计算 C 的每个元素需要读取 A 的一行和 B 的一列——大量重复读取。

Tiled 实现：

1. 将 C 划分为 BLOCK_M × BLOCK_N 的小块
2. 对于每个 C 的小块，沿 K 维度将 A 和 B 分成 BLOCK_K 大小的 tiles
3. 每次加载 A 的一个 tile (BLOCK_M × BLOCK_K) 和 B 的一个 tile (BLOCK_K × BLOCK_N) 到 shared memory
4. 在 shared memory 中计算部分矩阵乘积，累加到结果中
5. 重复直到遍历完所有 K 维度的 tiles

@triton.jit
def matmul_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,
):
    # Which block of C are we computing?
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    # Offsets for this block
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)

    # Initialize accumulator
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    # Iterate over K dimension in tiles
    for k_start in range(0, K, BLOCK_K):
        offs_k = k_start + tl.arange(0, BLOCK_K)

        # Load tiles of A and B
        a = tl.load(
            A_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak,
            mask=(offs_m[:, None] < M) & (offs_k[None, :] < K),
            other=0.0
        )
        b = tl.load(
            B_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn,
            mask=(offs_k[:, None] < K) & (offs_n[None, :] < N),
            other=0.0
        )

        # Accumulate: (BLOCK_M, BLOCK_K) @ (BLOCK_K, BLOCK_N)
        acc += tl.dot(a, b)

    # Store result tile
    c = acc.to(tl.float16)
    tl.store(
        C_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn,
        c,
        mask=(offs_m[:, None] < M) & (offs_n[None, :] < N)
    )

6. Fused Kernels 的优势

Kernel Fusion（算子融合）是将多个连续的操作合并为一个 GPU kernel。 这是 GPU 性能优化中最有效的技术之一，尤其对于 memory-bound 操作。

6.1 Fusion 前后对比

# Without fusion: 3 separate kernels
# Each reads from and writes to HBM

# Kernel 1: LayerNorm
# READ x from HBM → compute → WRITE y to HBM
y = layer_norm(x)         # HBM read: x, HBM write: y

# Kernel 2: Linear
# READ y, W from HBM → compute → WRITE z to HBM
z = linear(y)             # HBM read: y + W, HBM write: z

# Kernel 3: Activation
# READ z from HBM → compute → WRITE out to HBM
out = gelu(z)             # HBM read: z, HBM write: out

# Total HBM accesses: 3 reads + 3 writes of activation tensors
# (plus weights)


# With fusion: 1 kernel
# READ x, W from HBM → compute all three → WRITE out to HBM
out = fused_layernorm_linear_gelu(x, W)
# Total HBM accesses: 1 read + 1 write of activation tensors
# Intermediate values (y, z) stay in registers/shared memory!

6.2 PyTorch 中的 Fusion 方法

• torch.compile：PyTorch 2.0+ 的自动编译器，可以自动发现和执行 fusion
• Triton 自定义 kernel：手动编写融合 kernel，控制力最强
• 现有的融合 kernel 库：如 Flash Attention、xFormers 等

# torch.compile example
@torch.compile
def fused_operation(x, weight, bias):
    x = F.layer_norm(x, x.shape[-1:])
    x = x @ weight + bias
    x = F.gelu(x)
    return x

# torch.compile will automatically fuse compatible operations

7. Flash Attention 原理

Flash Attention 是近年来 LLM 系统优化中最重要的突破之一。 它通过巧妙的 tiling 和 online softmax 算法， 将 standard attention 的内存复杂度从 O(T²) 降低到 O(T)， 同时实现了 2-4 倍的速度提升。

7.1 Standard Attention 的问题

回顾标准的 attention 计算：

1. S = Q × K^T / √d_k（计算 T×T 的 attention score 矩阵）
2. P = softmax(S)（对每行做 softmax）
3. O = P × V（加权求和）

问题是：第 1 步生成的 S 矩阵和第 2 步的 P 矩阵大小为 T×T。 当 T=8192 时，这就是 8192² × 2 bytes ≈ 128 MB per head per sample。 这些巨大的中间矩阵必须存储在 HBM 中（因为 shared memory 太小放不下）， 然后再从 HBM 读取——大量的 HBM 读写成为性能瓶颈。

7.2 Online Softmax

Flash Attention 的关键技术之一是 Online Softmax—— 在不完整看到所有数据的情况下，增量地计算 softmax。

标准 softmax 需要两遍扫描：

1. 第一遍：计算 max（数值稳定）和 sum(exp)
2. 第二遍：计算 exp(x - max) / sum

Online softmax 只需一遍，通过维护 running max 和 running sum：

def online_softmax(x_blocks):
    """
    Compute softmax incrementally over blocks of x.
    """
    m_prev = float('-inf')  # running max
    l_prev = 0.0           # running sum of exp
    acc = 0.0              # running weighted sum (for attention)

    for x_block, v_block in zip(x_blocks, v_blocks):
        # Current block's max
        m_curr = max(m_prev, x_block.max())

        # Correction factor: rescale previous sum
        correction = exp(m_prev - m_curr)

        # Update running sum
        l_curr = correction * l_prev + sum(exp(x_block - m_curr))

        # Update accumulator (attention output)
        # Previous accumulator needs rescaling
        p_block = exp(x_block - m_curr)
        acc = correction * acc + p_block @ v_block

        m_prev = m_curr
        l_prev = l_curr

    # Final normalization
    return acc / l_prev

7.3 Flash Attention 的 Tiling 策略

Flash Attention 将 Q, K, V 矩阵分成小块（tiles），在 shared memory 中逐块计算 attention：

1. 将 Q 分成 T_r 个 block，K 和 V 分成 T_c 个 block
2. 外层循环遍历 Q 的每个 block
3. 内层循环遍历 K, V 的每个 block
4. 在 shared memory 中计算 S_ij = Q_i × K_j^T
5. 用 online softmax 增量更新 attention 输出
6. 不需要将完整的 T×T attention 矩阵存入 HBM！

# Flash Attention pseudocode (simplified)
# Q, K, V each (T, d), split into blocks of size B_r × d and B_c × d

for i in range(0, T, B_r):
    Q_i = Q[i:i+B_r]           # Load Q block to SRAM
    O_i = zeros(B_r, d)        # Output accumulator
    m_i = -inf * ones(B_r)     # Row-wise max
    l_i = zeros(B_r)           # Row-wise sum of exp

    for j in range(0, T, B_c):
        K_j = K[j:j+B_c]       # Load K block to SRAM
        V_j = V[j:j+B_c]       # Load V block to SRAM

        # Compute attention scores for this block
        S_ij = Q_i @ K_j.T / sqrt(d_k)   # (B_r, B_c) in SRAM

        # Online softmax update
        m_new = max(m_i, S_ij.max(dim=-1))
        correction = exp(m_i - m_new)
        P_ij = exp(S_ij - m_new.unsqueeze(-1))

        l_i = correction * l_i + P_ij.sum(dim=-1)
        O_i = correction.unsqueeze(-1) * O_i + P_ij @ V_j

        m_i = m_new

    # Final normalization
    O_i = O_i / l_i.unsqueeze(-1)
    O[i:i+B_r] = O_i           # Write output block to HBM

7.4 Flash Attention 的 HBM 访问量分析

当 M 足够大时（现代 GPU 的 shared memory 约 128-228 KB per SM）， Flash Attention 的 HBM 访问量远小于 standard attention。 特别是当 T 很大（长序列）时，节省尤为显著。

8. Flash Attention 的 Triton 实现

Flash Attention 的 Triton 实现比 CUDA 实现简洁很多（约 100-200 行 vs 1000+ 行）， 同时能达到 CUDA 实现性能的 80-90%。以下是核心结构：

@triton.jit
def flash_attention_kernel(
    Q_ptr, K_ptr, V_ptr, O_ptr,
    stride_qb, stride_qh, stride_qt, stride_qd,
    stride_kb, stride_kh, stride_kt, stride_kd,
    # ... other strides
    T, d_k,
    BLOCK_Q: tl.constexpr,
    BLOCK_KV: tl.constexpr,
):
    # Identify which (batch, head, q_block) this program handles
    pid_q = tl.program_id(0)
    pid_bh = tl.program_id(1)

    # Load Q block (BLOCK_Q, d_k) into SRAM
    q_offsets = pid_q * BLOCK_Q + tl.arange(0, BLOCK_Q)
    q = tl.load(...)  # (BLOCK_Q, d_k)

    # Initialize accumulators
    m_i = tl.full((BLOCK_Q,), float('-inf'), dtype=tl.float32)
    l_i = tl.zeros((BLOCK_Q,), dtype=tl.float32)
    acc = tl.zeros((BLOCK_Q, d_k), dtype=tl.float32)

    # Inner loop: iterate over K, V blocks
    for kv_start in range(0, T, BLOCK_KV):
        k = tl.load(...)  # (BLOCK_KV, d_k)
        v = tl.load(...)  # (BLOCK_KV, d_k)

        # Compute attention scores
        s = tl.dot(q, tl.trans(k)) * (1.0 / tl.sqrt(d_k.to(tl.float32)))

        # Apply causal mask
        causal_mask = q_offsets[:, None] >= (kv_start + tl.arange(0, BLOCK_KV))[None, :]
        s = tl.where(causal_mask, s, float('-inf'))

        # Online softmax
        m_new = tl.maximum(m_i, tl.max(s, axis=1))
        correction = tl.exp(m_i - m_new)
        p = tl.exp(s - m_new[:, None])

        l_i = correction * l_i + tl.sum(p, axis=1)
        acc = correction[:, None] * acc + tl.dot(p.to(v.dtype), v)
        m_i = m_new

    # Normalize and store
    acc = acc / l_i[:, None]
    tl.store(O_ptr + ..., acc.to(O_ptr.dtype.element_ty), mask=...)

9. Benchmarking 和 Profiling 方法

编写 kernel 之后，准确地测量和分析性能是至关重要的。

9.1 Triton Benchmarking

import triton

@triton.testing.perf_report(
    triton.testing.Benchmark(
        x_names=['T'],  # Argument to vary
        x_vals=[128, 256, 512, 1024, 2048, 4096],
        line_arg='provider',
        line_vals=['triton', 'torch'],
        line_names=['Triton', 'PyTorch'],
        styles=[('blue', '-'), ('red', '-')],
        ylabel='ms',
        plot_name='flash-attention-benchmark',
        args={'B': 4, 'H': 32, 'd': 128},
    )
)
def benchmark(B, H, T, d, provider):
    Q = torch.randn(B, H, T, d, device='cuda', dtype=torch.float16)
    K = torch.randn(B, H, T, d, device='cuda', dtype=torch.float16)
    V = torch.randn(B, H, T, d, device='cuda', dtype=torch.float16)

    if provider == 'triton':
        fn = lambda: flash_attention_triton(Q, K, V)
    else:
        fn = lambda: torch.nn.functional.scaled_dot_product_attention(Q, K, V)

    ms = triton.testing.do_bench(fn)
    return ms

benchmark.run(print_data=True, save_path='./benchmarks/')

9.2 PyTorch Profiler

# Detailed GPU profiling
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    record_shapes=True,
    with_stack=True,
    profile_memory=True,
) as prof:
    output = model(input_ids)
    loss = criterion(output, labels)
    loss.backward()

# Print kernel-level breakdown
print(prof.key_averages().table(
    sort_by="cuda_time_total",
    row_limit=20
))

# Export for visualization in Chrome trace viewer
prof.export_chrome_trace("trace.json")

9.3 关键性能指标

指标	如何计算	目标
Achieved TFLOPS	理论 FLOPs / 实测时间	接近 GPU 峰值
Memory Bandwidth Utilization	实际数据传输量 / 实测时间 / HBM 峰值带宽	memory-bound 操作接近 100%
Occupancy	活跃 Warp 数 / SM 支持的最大 Warp 数	>50%

10. 本讲总结

• GPU Kernel 是在 GPU 上执行的函数，每个 PyTorch 操作背后都对应 kernel
• 自定义 kernel 的核心价值是 kernel fusion——减少 HBM 读写次数
• Triton 提供了接近 Python 的语法来编写高性能 GPU kernel（block-level programming）
• Memory Coalescing 和 Tiling 是两个最基础的 GPU 优化策略
• Flash Attention 通过 tiling + online softmax 将 attention 的内存复杂度从 O(T²) 降到 O(T)
• Flash Attention 不减少 FLOPs，而是减少 HBM 访问，对 memory-bound 的 attention 操作带来巨大加速
• Benchmarking 和 Profiling 是性能优化的必备工具——先测量，再优化