wenyi-li MatMul 算子优化
适用于矩阵乘法及相关运算
CUDA GPU MatMul 优化核心
Tensor Core 利用
- Ampere (A100): 支持 FP16, BF16, TF32, INT8 Tensor Core
- Hopper (H100): 额外支持 FP8, wgmma 指令
- 关键:
tl.dot(a, b, allow_tf32=True)启用 TF32 Tensor Core
分块配置建议
常用配置（2 的幂次）：
| 配置 | BLOCK_M | BLOCK_N | BLOCK_K | num_warps | num_stages | 适用场景 |
|---|---|---|---|---|---|---|
| 小矩阵 | 64 | 64 | 32 | 4 | 4 | M, N < 1024 |
| 中矩阵 | 128 | 128 | 32 | 4 | 3 | M, N < 4096 |
| 大矩阵 | 128 | 256 | 64 | 8 | 3 | M, N >= 4096 |
| 高 K | 64 | 128 | 64 | 4 | 4 | K 很大 |
标准 MatMul Kernel（使用 block_ptr）
@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_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
):
pid = tl.program_id(0)
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
# 2D 索引计算
pid_m = pid // num_pid_n
pid_n = pid % num_pid_n
# 创建 block pointers
a_block_ptr = tl.make_block_ptr(
base=a_ptr,
shape=(M, K),
strides=(stride_am, stride_ak),
offsets=(pid_m * BLOCK_SIZE_M, 0),
block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_K),
order=(1, 0)
)
b_block_ptr = tl.make_block_ptr(
base=b_ptr,
shape=(K, N),
strides=(stride_bk, stride_bn),
offsets=(0, pid_n * BLOCK_SIZE_N),
block_shape=(BLOCK_SIZE_K, BLOCK_SIZE_N),
order=(1, 0)
)
# 使用 float32 累加器
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
# K 维度循环
for k in range(0, K, BLOCK_SIZE_K):
a = tl.load(a_block_ptr, boundary_check=(0, 1))
b = tl.load(b_block_ptr, boundary_check=(0, 1))
accumulator += tl.dot(a, b)
# 移动 block pointers
a_block_ptr = tl.advance(a_block_ptr, (0, BLOCK_SIZE_K))
b_block_ptr = tl.advance(b_block_ptr, (BLOCK_SIZE_K, 0))
# 存储结果（需显式转换类型，匹配输出 dtype）
c = accumulator.to(c_ptr.dtype.element_ty)
c_block_ptr = tl.make_block_ptr(
base=c_ptr,
shape=(M, N),
strides=(stride_cm, stride_cn),
offsets=(pid_m * BLOCK_SIZE_M, pid_n * BLOCK_SIZE_N),
block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N),
order=(1, 0)
)
tl.store(c_block_ptr, c, boundary_check=(0, 1))
使用 Autotune 优化
@triton.autotune(
configs=[
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3, num_warps=8),
triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=4),
triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, num_warps=2),
],
key=['M', 'N', 'K'],
restore_value=['c_ptr'], # 必须：列出所有输出指针参数名
)
@triton.jit
def matmul_kernel_autotune(
a_ptr, b_ptr, c_ptr,
M, N, K,
stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
GROUP_SIZE_M: tl.constexpr,
):
pid = tl.program_id(0)
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
# L2 缓存优化：Grouped ordering
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
# ... 后续与标准 kernel 相同
a_block_ptr = tl.make_block_ptr(
base=a_ptr, shape=(M, K), strides=(stride_am, stride_ak),
offsets=(pid_m * BLOCK_SIZE_M, 0),
block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_K), order=(1, 0)
)
b_block_ptr = tl.make_block_ptr(
base=b_ptr, shape=(K, N), strides=(stride_bk, stride_bn),
offsets=(0, pid_n * BLOCK_SIZE_N),
block_shape=(BLOCK_SIZE_K, BLOCK_SIZE_N), order=(1, 0)
)
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
for k in range(0, K, BLOCK_SIZE_K):
a = tl.load(a_block_ptr, boundary_check=(0, 1))
b = tl.load(b_block_ptr, boundary_check=(0, 1))
accumulator += tl.dot(a, b)
a_block_ptr = tl.advance(a_block_ptr, (0, BLOCK_SIZE_K))
b_block_ptr = tl.advance(b_block_ptr, (BLOCK_SIZE_K, 0))
c = accumulator.to(c_ptr.dtype.element_ty)
c_block_ptr = tl.make_block_ptr(
base=c_ptr, shape=(M, N), strides=(stride_cm, stride_cn),
offsets=(pid_m * BLOCK_SIZE_M, pid_n * BLOCK_SIZE_N),
block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N), order=(1, 0)
)
tl.store(c_block_ptr, c, boundary_check=(0, 1))
Host 侧启动
class ModelNew(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, a, b):
M, K = a.shape
K2, N = b.shape
assert K == K2
c = torch.empty((M, N), device=a.device, dtype=a.dtype)
grid = lambda meta: (triton.cdiv(M, meta['BLOCK_SIZE_M']) * triton.cdiv(N, meta['BLOCK_SIZE_N']),)
matmul_kernel_autotune[grid](
a, b, c,
M, N, K,
a.stride(0), a.stride(1),
b.stride(0), b.stride(1),
c.stride(0), c.stride(1),
)
return c
L2 缓存优化：Grouped Ordering
为什么需要 Grouped Ordering？
标准的行优先或列优先遍历会导致 L2 缓存利用率低。通过将相邻的块分组处理，可以增加数据复用：
# 标准遍历：相邻 pid 访问不同行的 A 块
pid_m = pid // num_pid_n
pid_n = pid % num_pid_n
# Grouped ordering：相邻 pid 访问同一组行的 A 块
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)
优化要点
1. 分块配置
- 使用 autotune 搜索最优配置
- 考虑 Tensor Core 的要求（块大小为 16 的倍数）
- 更大的块 → 更好的数据复用，但更高的寄存器压力
2. 精度控制
- 累加器使用 float32:
tl.zeros(..., dtype=tl.float32) - 即使输入是 fp16/bf16，也用 float32 累加
- 最后存储时自动转回目标精度
3. 内存访问
- 优先使用
tl.make_block_ptr和boundary_check - 使用
tl.advance移动块指针 - 利用 Grouped Ordering 优化 L2 缓存
4. 流水线
num_stages控制软件流水线级数- 更多 stage → 更好地隐藏内存延迟
- 但会占用更多共享内存
性能检查清单
- 是否使用了 autotune 搜索最优配置？
- 累加器是否使用 float32？
- 是否使用了 Grouped Ordering 或 swizzle2d 优化 L2 缓存？
- K 维度循环是否正确实现？
- num_warps 和 num_stages 是否合理？
- block 大小是否为 16 的倍数（Tensor Core 要求）？
常见错误
- 累加用 fp16: 精度损失严重
- 忘记 K 维度循环: 结果错误
- block 大小不对齐 Tensor Core: 性能不佳
- L2 缓存未优化: 大矩阵性能下降
- num_warps 不匹配: block 大小和 warp 数不匹配导致资源浪费