T4 讲解：Triton GEMM，tl.dot、分组顺序与 autotune

① 先看全貌：每个 program 一个 C tile

和 C4 的 v2 一样：每个 program 负责 C 的一个 BLOCK_M × BLOCK_N 的 tile，沿 K 方向每次取 A 的 BLOCK_M × BLOCK_K 块和 B 的 BLOCK_K × BLOCK_N 块，累加到 fp32 的 acc 里。

② 分组顺序：program 和 tile 怎么对应

drills24–32 行

num_pid_in_group = GROUP_M * num_pid_n
group_id     = pid // num_pid_in_group
first_pid_m  = group_id * GROUP_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_M)   # 最后一组可能不满
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m

这几行只是把一维的 pid 换算成 tile 的 (行, 列)。最直接的换算是行主序（pid_m = pid // num_pid_n，pid_n = pid % num_pid_n），分组顺序则是每 GROUP_M 行 tile 为一组，组内先往下走、再往右走。

C 被切成 4×4 个 tile，GROUP_M = 2，每格里写的是负责它的 pid：

左图里 pid 0–3 是第一行的 4 个 tile；右图里 pid 0–3 是左上角的 2×2。下半部分橙色是第二组（group_id = 1）。

拿 pid = 5 验算右图：num_pid_in_group = 2 × 4 = 8，group_id = 0，first_pid_m = 0，group_size_m = 2；pid_m = 0 + (5 % 8) % 2 = 1，pid_n = (5 % 8) // 2 = 2。所以 pid 5 负责第 1 行第 2 列，和图里一致。

为什么这样更快： 假设 GPU 同时在跑 pid 0–3 这 4 个 program。

同时运行的 program 数越多，差距越大。真实场景里同时运行的是几百个 program，行主序会让它们需要一整行 B 的 panel，L2 装不下，就会反复从显存读。分组之后，同时运行的 program 集中在一个近似正方形的区域里，A、B 的 panel 都能留在 L2 里被重复使用。题目页建议把 GROUP_M 改成 1 来对比 L2 命中率。

③ 二维指针块：广播出来的

drills35–39 行

offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)       # [BLOCK_M]
offs_k = tl.arange(0, BLOCK_K)                         # [BLOCK_K]
a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak   # [BLOCK_M, BLOCK_K]

offs_m[:, None] 是一列（形状 [BLOCK_M, 1]），offs_k[None, :] 是一行（[1, BLOCK_K]），两者相加按 NumPy 的规则广播成 [BLOCK_M, BLOCK_K] 的矩阵。每个格子是 A 里对应元素的地址。

BLOCK_M = BLOCK_K = 4，A 是 [M, 8] 的行优先矩阵（stride_am = 8，stride_ak = 1），pid_m = 0 时：

格子里是 a_ptrs 相对 a_ptr 的偏移（以元素为单位）。左列是 offs_m × stride_am，上行是 offs_k × stride_ak，每格 = 行头 + 列头。

为什么要传 stride，而不是假设连续：传 b.t() 这样的转置视图进来，stride_bk = 1、stride_bn = K，同一套代码照样对，不需要先 .contiguous() 拷贝一份。PyTorch 的 Linear 存的权重就是 [N, K]，乘的时候需要的正是它的转置视图。

④ 主循环：load、tl.dot、指针后移

drills42–50 行

acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for k in range(0, tl.cdiv(K, BLOCK_K)):
    k_remaining = K - k * BLOCK_K
    a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < k_remaining), other=0.0)
    b = tl.load(b_ptrs, mask=(offs_k[:, None] < k_remaining) & (offs_n[None, :] < N), other=0.0)
    acc = tl.dot(a, b, acc)
    a_ptrs += BLOCK_K * stride_ak
    b_ptrs += BLOCK_K * stride_bk

• mask 也是二维广播：(offs_m[:, None] < M) 是一列 bool，(offs_k[None, :] < k_remaining) 是一行 bool，& 之后是 [BLOCK_M, BLOCK_K] 的 mask。越界填 0，和 C4 的道理一样：0 不影响点积。
• k_remaining：最后一轮 K 方向可能不满一个 BLOCK_K。用“还剩多少”来比较，指针只需要往后移，offs_k 不用变。
• tl.dot(a, b, acc) 等于 acc += a @ b。fp16 / bf16 输入时编译成 Tensor Core 指令，在 fp32 里累加。第三个参数让编译器把累加融合进 mma 指令。
• 指针整体后移：A 往右移 BLOCK_K 列，B 往下移 BLOCK_K 行。

num_stages：软件流水线

你写的循环是“读一块、算一块”，但编译器会把 load 提前发出：

每个 stage 要存一块 A 和一块 B（fp16 每个元素 2 字节）：

超过 GPU 每个 SM 的 shared memory 上限，这个 config 就会编译失败或被跳过；占用太多，每个 SM 能同时驻留的 program 就少。

⑤ epilogue

drills53–55 行：acc 从 fp32 转回 fp16/bf16，用 C 的 stride 算出二维指针块，带 mask 写回。题目页的进阶练习是在这里融合 bias 和激活：c = gelu(acc + bias)，这是 GEMM 融合最常见的形式。

⑥ autotune 在背后做了什么

drills60–70 行 · host 端 drills84–93 行

matmul_kernel_autotuned = triton.autotune(configs=AUTOTUNE_CONFIGS, key=["M", "N", "K"])(matmul_kernel)
grid = lambda meta: (triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]),)
matmul_kernel_autotuned[grid](*args)

• grid 必须是函数：autotuner 在尝试不同 config 时 BLOCK_M、BLOCK_N 各不相同，program 数也跟着变，所以 grid 要根据 meta（当前 config 的参数）现算。
• CPU 解释器不调优（85–86 行）：解释器上计时没有意义，直接用一个固定的小 config。
• 源码：tritonAutotuner，查缓存和 benchmark 的逻辑在 tritonrun()。

⑦ 对照与延伸

• 官方 tutorial 03 和参考解几乎逐行相同：tritonmatmul_kernel，L2 分组的讲解在 triton“L2 Cache Optimizations” 一节（带图），tritontl.dot。
• persistent kernel：program 数等于 SM 数，每个 program 循环处理多个 tile，省掉反复启动 program 的开销：tritonmatmul_kernel_persistent。题目页的面试追问提到这是 cuBLAS 的优势之一。
• grouped GEMM（一次算很多个不同 shape 的小 GEMM，M2 的 MoE 会用到）：triton08-grouped-gemm；block-scaled GEMM（M3 的 FP8）：triton10-block-scaled-matmul。

⑧ 常见坑与自测

• grid 用了固定的 BLOCK_M：配合 autotune 时 program 数不对，部分 tile 没人算或者越界。
• 只 mask 了 M、N，没 mask K：K 不是 BLOCK_K 的倍数时读进垃圾值。
• acc 用了 fp16：大 K 时累加误差很大。Tensor Core 本来就支持 fp32 累加，没有性能损失。
• GROUP_M 的最后一组：num_pid_m 不是 GROUP_M 的倍数时，最后一组更矮，必须用 group_size_m 而不是 GROUP_M 取模。

源码与资料索引

本仓库：drillssolutions/t4_matmul.py · drillsmine/t4_matmul.py · drillstests/test_t4_matmul.py · drillscontent/t4.md

Triton：tritontutorial 03 · triton09 persistent · triton08 grouped · tritonAutotuner · API：tl.dot · triton.autotune · triton.Config

延伸阅读：Simon Boehm · How to Optimize a CUDA Matmul Kernel（C4 → cuBLAS 的每一步） · CUTLASS