triton-cuda-optimization

Triton CUDA 性能优化指南

1. 性能优化策略

1.1 块大小选择

• 原则: 平衡并行度与资源占用
• 建议: 使用 2 的幂次（256, 512, 1024）
• GPU 考量: 需要足够多的 warp 来隐藏延迟

1.2 Warp 和 Stage 调优

CUDA 后端特有的两个重要参数：

• num_warps: 每个 block 的 warp 数量（每个 warp = 32 个线程）

  • 小 BLOCK_SIZE：使用较少的 warp (2-4)
  • 大 BLOCK_SIZE：使用较多的 warp (4-8)
  • MatMul：通常使用 4-8 个 warp

• num_stages: 软件流水线级数

  • 更多的 stage 可以更好地隐藏内存延迟
  • 但会占用更多共享内存
  • 通常 2-5 之间选择

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

1.3 内存访问优化

• 合并访问 (Coalesced Access): 同一 warp 内的线程应访问连续内存地址
• 2D数据: 优先使用 tl.make_block_ptr 配合 boundary_check
• 步幅设计: 仔细设计 stride 参数，错误设置会严重影响性能
• 数据布局: 保持内存访问的连续性和局部性

1.4 算子拆分策略

• 复杂算子: 拆分为多个简单 kernel，避免单个 kernel 过于复杂
• 融合策略: 适度融合以减少全局内存读写（如 fused attention）
• 平衡: CUDA 后端融合通常比 NPU 更有效，但仍需注意 register pressure

1.5 Occupancy 优化

GPU 利用率（Occupancy）是性能的关键指标：

• 寄存器使用: 减少每个线程的寄存器使用量，增加并发 block 数
• 共享内存: 合理使用共享内存，不超过硬件限制
• Block 大小: 选择能整除 SM 最大线程数的 block 大小

2. 数值稳定性

2.1 防溢出处理

Softmax 数值稳定化:

# 减去最大值防止 exp 溢出
max_val = tl.max(scores, axis=0)
scores = scores - max_val
p = tl.exp(scores)  # CUDA 后端直接使用 tl.exp

2.2 防负值开方

# 方差计算前确保非负
variance = tl.maximum(variance, 0.0)
std = tl.sqrt(variance + eps)

2.3 精度提升

• 使用 float32 进行累加: 即使输入是 float16/bfloat16
• 最后再转换: 计算完成后再转回目标精度
• TF32: Ampere+ GPU 上可使用 TF32 加速 MatMul

accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
# ... 累加计算 ...
result = tl.cast(accumulator, output_dtype)

3. API 使用限制

3.1 禁止使用的语法

禁止使用: return, break, continue, lambda

Triton 内核是一次性执行完整逻辑，不支持提前返回或跳转语句。

# 错误：使用 return
@triton.jit
def kernel(ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    if pid >= n:
        return  # 编译错误！

# 正确：使用 mask
@triton.jit
def kernel(ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offsets < n
    data = tl.load(ptr + offsets, mask=mask, other=0.0)
    # ... 所有代码都在同一层级执行