triton-cuda-memory

内存访问优化

内存访问是 GPU 性能的关键瓶颈。本文档提供 Triton CUDA 的内存访问优化策略。

---

1. GPU 内存层次

内存带宽和延迟

内存类型	带宽 (A100)	延迟	容量
寄存器	~19 TB/s	1 cycle	256 KB/SM
共享内存	~19 TB/s	~20 cycles	164 KB/SM
L2 缓存	~5 TB/s	~100 cycles	40 MB
全局内存 (HBM)	~2 TB/s	~400 cycles	40/80 GB

优化原则

• 减少全局内存访问: 利用共享内存和寄存器
• 合并访问 (Coalesced Access): 同一 warp 内线程访问连续地址
• 提高 L2 缓存命中率: 通过 Grouped Ordering 等技术

---

2. 合并访问 (Coalesced Access)

什么是合并访问？

当同一 warp 中的 32 个线程访问连续的内存地址时，GPU 可以将这些请求合并为一次或少量内存事务，大幅提高带宽利用率。

# 正确：合并访问（连续地址）
@triton.jit
def coalesced_kernel(input_ptr, output_ptr, n, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)  # 连续偏移
    mask = offsets < n
    data = tl.load(input_ptr + offsets, mask=mask)
    tl.store(output_ptr + offsets, data, mask=mask)

# 错误：非合并访问（跳跃地址）
@triton.jit
def strided_kernel(input_ptr, output_ptr, n, stride, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    # 每个线程跳跃 stride 个元素，导致非合并访问
    strided_offsets = offsets * stride
    mask = strided_offsets < n
    data = tl.load(input_ptr + strided_offsets, mask=mask)

---

3. 块大小选择策略

调优原则

• 平衡并行度与资源占用，避免过大或过小
• BLOCK_SIZE 常用值：128, 256, 512, 1024
• 过小：并行度不足，无法充分利用 warp
• 过大：寄存器/共享内存溢出，occupancy 下降

推荐设置

• Element-wise 算子：BLOCK_SIZE = 1024 或 512
• Reduce 算子：BLOCK_SIZE = triton.next_power_of_2(n_cols)
• MatMul 算子：BLOCK_M = 128, BLOCK_N = 128, BLOCK_K = 32-64

---

4. 2D 数据内存访问优化

优先使用 tl.make_block_ptr

对于 2D 数据（如矩阵），优先使用 tl.make_block_ptr 配合 boundary_check，可自动优化内存合并。

@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,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    
    # 创建 2D Block Pointer
    A_block_ptr = tl.make_block_ptr(
        base=A_ptr,
        shape=(M, K),
        strides=(stride_am, stride_ak),
        offsets=(pid_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, BLOCK_K),
        order=(1, 0),  # Row-major
    )
    
    # 使用 boundary_check 自动处理边界
    a = tl.load(A_block_ptr, boundary_check=(0, 1))