4. 矩阵乘法基础内核

任务说明

本案例实现一个基础的矩阵乘法Triton内核，支持任意尺寸的方阵（M=N=K）。内核采用分块（tiling）策略，每个程序实例负责计算一个BLOCK_SIZE×BLOCK_SIZE的输出子块。为处理不能被块大小整除的矩阵维度，加载输入时使用掩码（mask）进行边界检查，越界位置填充0，确保任意尺寸方阵均可正确计算。

内核实现

import torch
import triton
import triton.language as tl

@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: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE)
    num_pid_in_group = tl.num_programs(axis=0) // num_pid_m
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * num_pid_in_group
    group_size_m = min(num_pid_in_group, num_pid_m - first_pid_m)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    offs_k = tl.arange(0, BLOCK_SIZE)

    a_ptrs = A_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = B_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)

    mask_m = offs_m[:, None] < M
    mask_n = offs_n[None, :] < N
    mask_k = offs_k[:, None] < K

    acc = tl.zeros((BLOCK_SIZE, BLOCK_SIZE), dtype=tl.float32)

    for k in range(0, tl.cdiv(K, BLOCK_SIZE)):
        a = tl.load(a_ptrs, mask=mask_m & mask_k, other=0.0)
        b = tl.load(b_ptrs, mask=mask_n & mask_k, other=0.0)
        acc += tl.dot(a, b)
        a_ptrs += BLOCK_SIZE * stride_ak
        b_ptrs += BLOCK_SIZE * stride_bk

    offs_cm = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    offs_cn = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    c_ptrs = C_ptr + (offs_cm[:, None] * stride_cm + offs_cn[None, :] * stride_cn)
    mask_c = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, acc, mask=mask_c)

def matmul(a: torch.Tensor, b: torch.Tensor, block_size: int = 16):
    assert a.is_cuda and b.is_cuda
    assert a.shape[0] == a.shape[1] == b.shape[0] == b.shape[1], "仅支持方阵"
    M, K = a.shape
    K, N = b.shape
    c = torch.empty((M, N), device=a.device, dtype=torch.float32)

    grid = (triton.cdiv(M, block_size) * triton.cdiv(N, block_size),)
    matmul_kernel[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),
        BLOCK_SIZE=block_size,
    )
    return c

代码说明

内核按输出矩阵的二维分块进行并行化。每个程序实例通过program_id映射到唯一的(pid_m, pid_n)块位置。K维循环中，每次加载一个BLOCK_SIZE×BLOCK_SIZE的A子块和B子块，调用tl.dot累加部分结果。掩码mask_m、mask_n、mask_k分别对M、N、K三个维度进行边界检查，越界加载自动填充0，保证任意尺寸方阵的正确性。累加器使用float32避免精度损失，最终结果按掩码写回C矩阵。

验证方法

def verify():
    torch.manual_seed(42)
    for size in [32, 100, 257]:  # 覆盖对齐与不对齐尺寸
        a = torch.randn(size, size, device='cuda')
        b = torch.randn(size, size, device='cuda')
        c_triton = matmul(a, b)
        c_ref = (a.float() @ b.float())
        max_diff = (c_triton - c_ref).abs().max().item()
        assert max_diff < 1e-3, f"size={size}, max_diff={max_diff}"
        print(f"size={size}: OK, max_diff={max_diff}")

verify()

运行上述验证脚本，程序会对32、100、257三种尺寸的方阵执行内核计算，并与PyTorch的@矩阵乘法结果逐元素比较。判定标准为最大误差小于1e-3（float32累加下合理容差）。若所有尺寸均通过断言且打印OK，则内核实现正确。

运行方式

将内核代码与验证函数保存为matmul_basic.py，在配备NVIDIA GPU且已安装PyTorch、Triton的环境下执行python matmul_basic.py。脚本会自动运行三种尺寸的验证并打印结果，全部通过即表示内核实现正确。