系统 第 6 讲 / 共 17 讲

第 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 的默认操作存在以下局限：

Fusion 的收益

假设三个连续的 element-wise 操作各需读写 32MB 数据。分开执行需要 6 次 HBM 访问（3 读 + 3 写 = 192 MB 传输），融合后只需 1 读 + 1 写（64 MB 传输），内存传输减少 3 倍。对于 memory-bound 操作，这意味着约 3 倍的速度提升。

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 的关键特点：

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

3. Triton 编程语言介绍

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

3.1 Triton vs CUDA

方面CUDA C++Triton
语言C++Python-like
编程粒度单个 threadBlock 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)

Triton 的关键抽象

在 Triton 中，tl.arange(0, BLOCK_SIZE) 不是创建一个数组——它代表一组 threads 各自的偏移量。tl.load 和 tl.store 是向量化的内存操作，整个 block 的 threads 协同完成。这就是 "block-level programming" 的含义：你写的每一行代码都在描述整个 block 的行为。

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 中，这需要程序员手动确保。

Row-major vs Column-major

在矩阵操作中，沿着行方向（row-major 存储的最内层维度）访问是 coalesced 的，沿着列方向访问则不是。这就是为什么矩阵转置看似简单，却是一个需要仔细优化的操作。Triton 中处理 2D 数据时，需要注意 stride 参数的设置。

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)
    )

Tile 大小的选择

BLOCK_M, BLOCK_N, BLOCK_K 的选择是性能调优的关键。太小：无法充分利用 Tensor Core 的计算能力；太大：shared memory 放不下，或者 occupancy 降低。典型值：BLOCK_M=128, BLOCK_N=128, BLOCK_K=32。Triton 支持 triton.autotune 来自动搜索最优的 tile 大小。

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 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 × KT / √dk（计算 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 读写成为性能瓶颈。

核心瓶颈不是 FLOPs

Standard attention 的 FLOPs 是 O(T²d)，这一点 Flash Attention 并没有改变。Flash Attention 优化的是HBM 访问量——从 O(T²) 降到 O(T²d/M)，其中 M 是 shared memory 大小。对于 memory-bound 的 attention 操作，这带来了巨大的实际加速。

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 分成 Tr 个 block，K 和 V 分成 Tc 个 block
  2. 外层循环遍历 Q 的每个 block
  3. 内层循环遍历 K, V 的每个 block
  4. 在 shared memory 中计算 Sij = Qi × KjT
  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 访问量分析

Standard Attention HBM 访问: O(Td + T²)
Flash Attention HBM 访问: O(T²d² / M)
其中 M 为 SRAM (shared memory) 大小

当 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=...)

实践建议

在实际项目中，推荐直接使用成熟的 Flash Attention 库（如 flash-attn），而不是自己从头实现。但理解其原理对于调试性能问题、选择合适的 block size、理解 profiling 结果都非常有帮助。在 Assignment 2 中，你将使用 Triton 实现一个简化版的 Flash Attention。

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%

Profile 驱动优化

永远不要"猜测"性能瓶颈在哪里。先 profile，看数据，然后有针对性地优化。常见的发现：你以为 matmul 是瓶颈，但实际上 softmax 和 LayerNorm 占了 40% 的时间；你以为 kernel 计算慢，但实际上 kernel launch overhead 就占了 20%。

10. 本讲总结