Triton Cuda Memory

内存访问优化

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

1. GPU 内存层次

内存带宽和延迟

优化原则

• 减少全局内存访问: 利用共享内存和寄存器
• 合并访问 (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))

Stride 设计要点

• 仔细设计 stride 参数，错误设置会严重影响性能
• 连续访问：确保内存访问的连续性和局部性
• 行主序 (Row-major): PyTorch 默认，stride(0) > stride(1)

5. 连续内存的一维访问优化

推荐方案：转连续后用一维访问

class ModelNew(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, input_tensor):
        # 非连续张量转为连续（一次性开销）
        if not input_tensor.is_contiguous():
            input_tensor = input_tensor.contiguous()
        
        output_tensor = torch.empty_like(input_tensor)
        n_elements = input_tensor.numel()
        grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
        
        elementwise_kernel[grid](
            input_tensor, output_tensor, 
            n_elements, 
            BLOCK_SIZE=1024
        )
        return output_tensor

性能对比

建议：非连续张量先调用 .contiguous() 转换，再用一维访问，整体性能更优。

6. L2 缓存优化

Grouped Ordering

对于 MatMul 等 2D 算子，通过分组遍历提高 L2 缓存命中率：

# 标准遍历：L2 缓存利用率低
pid_m = pid // num_pid_n
pid_n = pid % num_pid_n

# Grouped Ordering：L2 缓存利用率高
GROUP_SIZE_M = 8
num_pid_in_group = GROUP_SIZE_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m

swizzle2d

task_m, task_n = tl.swizzle2d(pid_m, pid_n, num_pid_m, num_pid_n, GROUP_SIZE)

7. 软件流水线 (Software Pipelining)

num_stages 参数

通过 num_stages 控制预取级数，隐藏内存延迟：

• num_stages=2: 最少的共享内存使用
• num_stages=3-4: 通常最优
• num_stages=5+: 可能超出共享内存限制

@triton.autotune(
    configs=[
        triton.Config({...}, num_stages=2, num_warps=4),
        triton.Config({...}, num_stages=3, num_warps=4),
        triton.Config({...}, num_stages=4, num_warps=8),
    ],
    key=[...],
    restore_value=['output_ptr'],  # 必须：列出所有输出指针参数名
)

8. 最佳实践

Element-wise 算子

1. 转连续：input.contiguous()
2. 一维访问：ptr + offsets
3. BLOCK_SIZE = 1024

2D 算子（MatMul、Attention）

1. 使用 tl.make_block_ptr
2. 配合 boundary_check
3. Grouped Ordering 优化 L2 缓存
4. 合理设置 num_stages

避免的陷阱

• 非连续张量直接用 stride 访问
• BLOCK_SIZE 设置过大导致 occupancy 下降
• 忘记边界检查导致越界访问
• 忽略 L2 缓存优化

9. 调试建议

性能问题排查

1. 检查张量是否连续：tensor.is_contiguous()
2. 检查内存访问是否合并
3. 使用 Nsight Compute 分析内存带宽利用率
4. 检查 occupancy 是否合理

常见错误

• 内存访问越界：检查 mask 和 boundary_check
• 性能不佳：检查合并访问和 L2 缓存优化
• 结果错误：检查 stride 计算是否正确