triton-cuda-grid-config

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)