wenyi-li Grid 配置策略
Grid 配置是 Triton Kernel 启动的关键。本文档提供 Triton CUDA 的 Grid 配置策略和大 shape 处理方案。
1. Grid 设置规范
维度格式
- Grid 必须是 tuple 类型，最多 3 维
- 支持的格式：
(x,),(x, y),(x, y, z)
# 正确
grid = (100,)
grid = (100, 200)
grid = (100, 200, 50)
# 错误
grid = 100 # 必须是 tuple
grid = [100, 200] # 必须是 tuple，不能是 list
使用 lambda（autotune 场景）
当使用 autotune 时，grid 必须使用 lambda：
# autotune 时必须使用 lambda
grid = lambda meta: (triton.cdiv(M, meta['BLOCK_SIZE_M']) * triton.cdiv(N, meta['BLOCK_SIZE_N']),)
# 非 autotune 时可以直接计算
grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
2. 1D Grid 配置
Element-wise 算子
最常见的配置方式：每个 block 处理 BLOCK_SIZE 个元素。
n_elements = input_tensor.numel()
BLOCK_SIZE = 1024
grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
kernel[grid](input_tensor, output_tensor, n_elements, BLOCK_SIZE=BLOCK_SIZE)
逐行处理（Reduce 类算子）
每个 block 处理一行或多行：
n_rows, n_cols = x.shape
BLOCK_SIZE = triton.next_power_of_2(n_cols)
# 方式 1：每行一个 block
grid = (n_rows,)
# 方式 2：限制并行度（grid stride loop）
num_programs = min(n_rows, 65535)
grid = (num_programs,)
3. 2D Grid 配置
MatMul 类算子
使用 2D Grid 进行行列双向并行：
BLOCK_M, BLOCK_N = 128, 256
grid_m = triton.cdiv(M, BLOCK_M)
grid_n = triton.cdiv(N, BLOCK_N)
# 方式 1：2D Grid
grid = (grid_m, grid_n)
# 方式 2：1D Grid（更灵活，支持 Grouped Ordering）
grid = (grid_m * grid_n,)
1D vs 2D Grid
| 特性 | 1D Grid | 2D Grid |
|---|---|---|
| 灵活性 | 高（支持 Grouped Ordering） | 低 |
| 代码复杂度 | 需要手动计算 pid_m, pid_n | 直接获取 |
| L2 缓存优化 | 容易实现 | 不易实现 |
| 推荐场景 | MatMul（需要缓存优化） | 简单 2D 算子 |
推荐: 对于 MatMul 类算子，使用 1D Grid + Grouped Ordering。
4. 大 Shape 处理：Grid Stride Loop
问题描述
CUDA GPU 对 grid 大小也有限制（通常 2^31 - 1 per dimension），但更重要的是，过大的 grid 会导致：
- 启动开销增加
- 资源浪费（每个 block 只处理少量数据）
Grid Stride Loop 方案
每个 block 通过循环处理多个数据块：
@triton.jit
def grid_stride_kernel(
input_ptr, output_ptr, n_elements,
BLOCK_SIZE: tl.constexpr,
):
pid = tl.program_id(0)
num_pids = tl.num_programs(0)
# Grid stride loop
for block_start in range(pid * BLOCK_SIZE, n_elements, num_pids * BLOCK_SIZE):
offsets = block_start + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
data = tl.load(input_ptr + offsets, mask=mask)
result = compute(data)
tl.store(output_ptr + offsets, result, mask=mask)
# 限制 grid 大小
MAX_GRID = 65535
num_blocks = min(triton.cdiv(n_elements, BLOCK_SIZE), MAX_GRID)
grid = (num_blocks,)
grid_stride_kernel[grid](input_tensor, output_tensor, n_elements, BLOCK_SIZE=1024)
Softmax 的 Grid Stride Loop
@triton.jit
def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride,
n_rows, n_cols, BLOCK_SIZE: tl.constexpr):
row_start = tl.program_id(0)
row_step = tl.num_programs(0)
# 每个 block 处理多行
for row_idx in tl.range(row_start, n_rows, row_step):
# 处理第 row_idx 行...
pass
# 限制程序数量
num_programs = min(32, n_rows)
softmax_kernel[(num_programs, 1, 1)](...)
5. Batch 维度处理
3D Grid
对于 batch 操作，可以使用第三个维度：
# batch matmul
batch_size = Q.shape[0]
grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N), batch_size)
@triton.jit
def batch_kernel(...):
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
batch_idx = tl.program_id(2)
Flatten Batch
将 batch 维度展平为 1D：
# 将 (batch, M, N) 展平
x_flat = x.view(-1, N) # (batch * M, N)
total_rows = batch_size * M
grid = (total_rows,)
6. 最佳实践总结
Element-wise 算子
- 1D Grid：
(triton.cdiv(n_elements, BLOCK_SIZE),) - BLOCK_SIZE = 1024
- 大 shape 使用 grid stride loop
Reduce 算子
- 1D Grid：
(n_rows,)或(min(n_rows, max_programs),) - BLOCK_SIZE = triton.next_power_of_2(n_cols)
- Grid stride loop 处理多行
MatMul 算子
- 1D Grid + Grouped Ordering（推荐）
- 或 2D Grid（简单场景）
- 使用 autotune 搜索最优配置
Attention 算子
- 1D Grid：按查询位置并行
- 或 3D Grid：(seq_len / BLOCK, heads, batch)
7. 常见错误和解决方案
错误 1: Grid 参数类型错误
# 错误
grid = 1024 # 必须是 tuple
grid = [1024] # 必须是 tuple
# 正确
grid = (1024,)
错误 2: Autotune 时未使用 lambda
# 错误：autotune 时直接计算 grid
grid = (triton.cdiv(M, BLOCK_M),) # BLOCK_M 未知
# 正确：使用 lambda
grid = lambda meta: (triton.cdiv(M, meta['BLOCK_SIZE_M']),)
错误 3: Grid 过大导致性能下降
# 不推荐：每个元素一个 block
grid = (n_elements,)
# 推荐：合理的 block 大小
grid = (triton.cdiv(n_elements, 1024),)
8. 性能调优建议
Grid 大小选择
- 足够大: 充分利用所有 SM（A100 有 108 个 SM）
- 不过大: 避免不必要的启动开销
- 经验值: Grid 大小通常为 SM 数量的 2-4 倍以上
负载均衡
- 确保每个 block 处理相近的工作量
- Grid stride loop 天然负载均衡
- MatMul 的 Grouped Ordering 需要注意边界组的大小
与 Autotune 配合
- Grid 使用 lambda 引用 autotune 参数
- 不同的 BLOCK_SIZE 会导致不同的 Grid 大小
- Autotune 会自动搜索最优组合