""" ============================================================================= CS349E · 第 10 讲 · 06_triton_style.py —— 块级编程（Triton 风格）、启动顺序与 L2、融合、怎样测试一个内核 ============================================================================= 【这个脚本要回答的问题】 CUDA 里你写"一个线程做什么"，Triton 里你写"一个块做什么"——线程怎么分、要不要经过共享内存、 用不用 Tensor Core、流水几级，都交给编译器。同一个分块矩阵乘法用这种写法只有十几行。 演示 1：Triton 教程式的矩阵乘法（块指针 + mask + tl.dot），任意形状都对；块越大，读的字节越少。 演示 2：程序实例的启动顺序和 L2：Triton 教程的"分组顺序"（9×9 的例子：90 个块 vs 54 个块；4096³ 的波次模型）。 演示 3：融合：softmax 一行一个程序 vs 逐个算子（Triton 教程的 5MN + 2M 读 / 3MN + 2M 写）。 演示 4：怎样测试一个内核：边缘形状、一个"只在边缘出错"的 bug、FP16 累加 vs FP32 累加的误差与容差。 【运行】python3 06_triton_style.py （约 3 秒） 【依赖】numpy；本目录的 kernel_lab（mini-Triton 解释器 TL / triton_launch）；../lecture01 的 console。 """ import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from kernel_lab import dram_bytes_tiled, tile_order, triton_launch # noqa: E402 from simt_lab import cdiv # noqa: E402 from console import fmt_bytes, fmt_time, print_header, print_observations, print_table # noqa: E402 # ----------------------------------------------------------------------------- # Triton 风格的内核（和官方教程 03-matrix-multiplication 的结构一样，只是"指针"写成了"数组 + 下标"） # ----------------------------------------------------------------------------- def matmul_kernel(tl, a, b, c, M, N, K, BLOCK_M, BLOCK_N, BLOCK_K, GROUP_M, mask_k=True): """ 一个程序实例算 C 的一个 BLOCK_M × BLOCK_N 块。 【分组顺序】程序号 pid → (pid_m, pid_n)：每 GROUP_M 个块行为一组，组内先竖着走（演示 2）。 【mask】最后一个块可能超出矩阵：越界的位置不读（读成 0）、不写。mask_k=False 是演示 4 的故意 bug： K 方向不加 mask（最后一个 K 块读到下一行的数据）。 """ pid = tl.program_id(0) num_pid_m, num_pid_n = cdiv(M, BLOCK_M), cdiv(N, BLOCK_N) 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 offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) offs_k = tl.arange(0, BLOCK_K) acc = tl.zeros((BLOCK_M, BLOCK_N)) for k0 in range(0, K, BLOCK_K): kk = k0 + offs_k a_mask = (offs_m[:, None] < M) & ((kk[None, :] < K) if mask_k else True) b_mask = ((kk[:, None] < K) if mask_k else True) & (offs_n[None, :] < N) a_offs = np.clip(offs_m[:, None] * K + kk[None, :], 0, a.size - 1) # 越界的下标夹住（只为了 numpy 不报错；mask 决定读不读） b_offs = np.clip(kk[:, None] * N + offs_n[None, :], 0, b.size - 1) A = tl.load(a, a_offs, mask=a_mask, other=0.0) B = tl.load(b, b_offs, mask=b_mask, other=0.0) acc = tl.dot(A, B, acc) c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) tl.store(c, np.clip(offs_m[:, None] * N + offs_n[None, :], 0, c.size - 1), acc.astype(c.dtype), mask=c_mask) def matmul(a, b, BLOCK_M=64, BLOCK_N=64, BLOCK_K=32, GROUP_M=8, mask_k=True): """启动 matmul_kernel：网格是一维的，cdiv(M, BLOCK_M)·cdiv(N, BLOCK_N) 个程序实例。返回 (C, 计数)。""" M, K = a.shape _, N = b.shape c = np.zeros((M, N), dtype=np.float32) grid = cdiv(M, BLOCK_M) * cdiv(N, BLOCK_N) cnt = triton_launch(matmul_kernel, grid, a, b, c, M, N, K, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, GROUP_M=GROUP_M, mask_k=mask_k) return c, cnt def demo1_matmul(): print_header("演示 1：Triton 风格的矩阵乘法（FP32）：任意形状、块越大读得越少") rng = np.random.default_rng(0) rows = [] for (M, N, K) in ((256, 256, 256), (300, 200, 150), (1, 1000, 64), (129, 65, 33)): a = rng.standard_normal((M, K)).astype(np.float32) b = rng.standard_normal((K, N)).astype(np.float32) c, cnt = matmul(a, b) err = np.abs(c - a.astype(np.float64) @ b).max() rows.append([f"{M}×{N}×{K}", cnt["programs"], f"{err:.1e}", fmt_bytes(cnt["load_bytes"]), f"{cnt['load_bytes'] / ((M * K + K * N) * 4):.1f}"]) print_table(["M×N×K", "程序实例", "最大误差", "读的字节（不含越界）", "相对\"每个数读一遍\""], rows) print() M = N = K = 512 a = rng.standard_normal((M, K)).astype(np.float32) b = rng.standard_normal((K, N)).astype(np.float32) rows = [] for bm, bn in ((16, 16), (32, 32), (64, 64), (128, 128), (128, 256)): _, cnt = matmul(a, b, BLOCK_M=bm, BLOCK_N=bn) ideal = (M * K * (N // bn) + K * N * (M // bm)) * 4 rows.append([f"{bm}×{bn}", cnt["programs"], fmt_bytes(cnt["load_bytes"]), fmt_bytes(ideal), f"{cnt['dot_flops'] / cnt['load_bytes']:.0f}"]) print_table(["BLOCK_M×BLOCK_N（512³）", "程序实例", "读的字节（计数）", "公式 (M·K·N/BN + K·N·M/BM)·4", "FLOP / 读的字节"], rows) print_observations([ "内核只描述\"一个块\"：offs_m、offs_n 是这个块负责的行号、列号向量，tl.load 一次读一整块，tl.dot 一次算一个块的乘积。" "和 k5 / k6 对比：没有 threadIdx、没有 __syncthreads、没有共享内存——这些都由 Triton 编译器决定。", "mask 让同一个内核处理任意形状：越界的位置不读（当作 0）、不写。1×1000×64 这种\"细长\"形状也对，只是每个程序实例的 64 行里只有 1 行有效，其余都被 mask 掉、白算。", "每个程序实例要读 A 的 BLOCK_M 行、B 的 BLOCK_N 列：总读取 = M·K·(N/BN) + K·N·(M/BM)。块越大，每读一个字节做的运算越多" "（512³ 时从 16×16 的 4 FLOP/字节到 128×256 的约 43）——和第 9 讲的分块、本讲 k3–k5 是同一个式子。", "这里数的是程序实例读的字节（从 L2 的角度）；真正到显存的字节还取决于不同的程序实例之间能不能在 L2 里共享（演示 2）。", ]) def demo2_launch_order(): print_header("演示 2：启动顺序与 L2——Triton 教程的\"分组顺序\"") rows = [] for name, group in (("行主序", None), ("分组顺序 GROUP_M = 3", 3)): pm, pn = tile_order(9, 9, group) first = slice(0, 9) a_blocks = len(np.unique(pm[first])) * 9 # 每个 A 块行有 9 个 K 块 b_blocks = len(np.unique(pn[first])) * 9 order = " ".join(f"({m},{n})" for m, n in zip(pm[:9], pn[:9])) rows.append([name, order, a_blocks, b_blocks, a_blocks + b_blocks]) print_table(["顺序（9×9 个输出块，K 方向也是 9 块）", "前 9 个程序实例算的块 (行,列)", "要读的 A 块", "要读的 B 块", "合计"], rows) print() rows = [] for gpu, n_sm in (("A100（108 SM）", 108), ("H100（132 SM）", 132)): for bm, bn in ((128, 128), (128, 256)): for g in (None, 4, 8, 16): byts = dram_bytes_tiled(4096, bm, bn, 1, n_sm, elem=2, group=g) base = dram_bytes_tiled(4096, bm, bn, 1, n_sm, elem=2, group=None) rows.append([gpu, f"{bm}×{bn}", "行主序" if g is None else f"GROUP_M = {g}", fmt_bytes(byts), f"{byts / base:.0%}"]) print_table(["GPU", "块", "顺序", "显存流量（4096³ FP16，波次模型）", "相对行主序"], rows) print_observations([ "9×9 的例子和 Triton 教程的数字一样：行主序的前 9 个程序实例算 C 的第一行块，要读 A 的 1 行（9 块）和 B 的全部（81 块）= 90 块；" "3×3 分组时读 A 的 3 行和 B 的 3 列 = 54 块。同时在跑的程序实例碰到的数据越\"方\"，L2 的复用越好。", "4096³ 上用本讲的波次模型（同一波的块按 K 齐头并进、在 L2 里共享）：128×128 的块、GROUP_M = 8 时降到行主序的 69%–77%；" "128×256 的块只有 GROUP_M = 16 时降到 87%–89%；GROUP_M 选得不好（4）反而比行主序还多——一波的块跨过两个组时，" "碰到的块行和块列都变多。最优值取决于 SM 数、块大小和矩阵形状，所以它也是自动调优的参数之一。", "这个优化完全不碰内核的计算部分，只改了\"程序号 → 块坐标\"的映射——CUDA 里叫 threadblock swizzle，" "CUTLASS 里是一个可选的模板参数。Boehm 在 kernel 10 上试过，因为 L2 命中率已经约 80%，没有收益：是否有用取决于问题大小和 L2。", ]) # ----------------------------------------------------------------------------- # 融合：softmax # ----------------------------------------------------------------------------- def _rowwise(tl, x, out, M, N, op, BLOCK): """逐个算子时的一个"按行归约"内核：每个程序实例处理一行，op = "max" 或 "sum"，输出长度 M。""" r = tl.program_id(0) cols = tl.arange(0, BLOCK) v = tl.load(x, r * N + np.minimum(cols, N - 1), mask=cols < N, other=-np.inf if op == "max" else 0.0) tl.store(out, np.array(r), np.max(v) if op == "max" else np.sum(v)) def _elementwise(tl, x, vec, out, M, N, op, BLOCK): """逐个算子时的逐元素内核：out = x - vec[行]、exp(x)、x / vec[行]。""" r = tl.program_id(0) cols = tl.arange(0, BLOCK) offs = r * N + np.minimum(cols, N - 1) v = tl.load(x, offs, mask=cols < N) if op == "sub": y = v - tl.load(vec, np.array(r)) elif op == "exp": y = np.exp(v) else: y = v / tl.load(vec, np.array(r)) tl.store(out, offs, y, mask=cols < N) def softmax_eager(x): """PyTorch 教程里的 naive_softmax：max、减、exp、sum、除，5 个内核，每个都从显存读、往显存写。""" M, N = x.shape B = 1 << (N - 1).bit_length() total = {"load_bytes": 0, "store_bytes": 0, "launches": 0} mx = np.zeros(M, np.float32) z = np.zeros_like(x) e = np.zeros_like(x) s = np.zeros(M, np.float32) y = np.zeros_like(x) for cnt in (triton_launch(_rowwise, M, x, mx, M, N, op="max", BLOCK=B), triton_launch(_elementwise, M, x, mx, z, M, N, op="sub", BLOCK=B), triton_launch(_elementwise, M, z, None, e, M, N, op="exp", BLOCK=B), triton_launch(_rowwise, M, e, s, M, N, op="sum", BLOCK=B), triton_launch(_elementwise, M, e, s, y, M, N, op="div", BLOCK=B)): total["load_bytes"] += cnt["load_bytes"] total["store_bytes"] += cnt["store_bytes"] total["launches"] += 1 return y, total def softmax_fused_kernel(tl, x, out, M, N, BLOCK): """Triton 教程 02-fused-softmax 的写法：一个程序实例读一整行进寄存器 / 共享内存，算完只写一次。""" r = tl.program_id(0) cols = tl.arange(0, BLOCK) offs = r * N + np.minimum(cols, N - 1) v = tl.load(x, offs, mask=cols < N, other=-np.inf) v = v - np.max(v) num = np.exp(v) tl.store(out, offs, num / np.sum(num), mask=cols < N) def demo3_fusion(): print_header("演示 3：融合——softmax 逐个算子 vs 一行一个程序（M = 4096 行、N = 1024 列，FP32）") rng = np.random.default_rng(2) M, N = 4096, 1024 x = rng.standard_normal((M, N)).astype(np.float32) ref = np.exp(x - x.max(1, keepdims=True)) ref /= ref.sum(1, keepdims=True) y1, c1 = softmax_eager(x) y2 = np.zeros_like(x) c2 = triton_launch(softmax_fused_kernel, M, x, y2, M, N, BLOCK=1 << (N - 1).bit_length()) MN = M * N rows = [] for name, y, c, launches, formula in (("逐个算子（5 个内核）", y1, c1, c1["launches"], "读 5MN + 2M，写 3MN + 2M"), ("融合（1 个内核）", y2, c2, 1, "读 MN，写 MN")): t = (c["load_bytes"] + c["store_bytes"]) / 3.35e12 rows.append([name, f"{np.abs(y - ref).max():.1e}", launches, f"{c['load_bytes'] / 4 / MN:.4f} MN", f"{c['store_bytes'] / 4 / MN:.4f} MN", formula, fmt_time(t)]) print_table(["写法", "最大误差", "内核数", "读（元素）", "写（元素）", "教程的公式", "H100 上按 3.35 TB/s 的时间"], rows) ratio = (c1["load_bytes"] + c1["store_bytes"]) / (c2["load_bytes"] + c2["store_bytes"]) print_observations([ f"逐个算子：每个算子都把整个 M×N 的中间结果写回显存、下一个算子再读回来，计数正好是 5MN + 2M 读、3MN + 2M 写" f"（多出的 2M 是每行的最大值和总和）。融合以后只读一遍 x、写一遍 y，流量少 {ratio:.1f} 倍。", "softmax 每个元素只有几次运算，完全受限于带宽，所以时间几乎正比于流量——Triton 教程正是据此估计融合能带来约 4 倍的理论加速，" "也是第 12 讲图编译器做算子融合的动机。", "融合的前提是一行放得进一个程序实例的片上存储（寄存器 / 共享内存）：N = 1024 的 FP32 一行 4 KiB 没问题；" "N 很大时一行放不下，就得用\"在线 softmax\"分块地算——那正是 FlashAttention（第 11 讲）。", ]) def demo4_testing(): print_header("演示 4：怎样测试一个内核") rng = np.random.default_rng(3) shapes = [(64, 64, 64), (128, 64, 96), (1, 1, 1), (17, 33, 65), (100, 70, 50), (64, 64, 40)] rows = [] for (M, N, K) in shapes: a = rng.standard_normal((M, K)).astype(np.float32) b = rng.standard_normal((K, N)).astype(np.float32) ref = a.astype(np.float64) @ b tol = 1e-5 * np.sqrt(K) * np.abs(a).max() * np.abs(b).max() * 4 res = [] for mk in (True, False): c, _ = matmul(a, b, BLOCK_M=32, BLOCK_N=32, BLOCK_K=32, mask_k=mk) err = np.abs(c - ref).max() res.append(("通过" if err <= tol else "失败") + f"（{err:.1e}）") mult = all(v % 32 == 0 for v in (M, N, K)) rows.append([f"{M}×{N}×{K}", "是" if mult else "否", f"{tol:.1e}", *res]) print_table(["M×N×K", "都是 32 的倍数？", "容差", "正确的内核", "K 方向忘了 mask 的内核"], rows) print() M = N = 64 rows = [] for K in (64, 256, 1024, 4096): a = rng.standard_normal((M, K)).astype(np.float16) b = rng.standard_normal((K, N)).astype(np.float16) ref = a.astype(np.float64) @ b.astype(np.float64) acc32 = np.zeros((M, N), np.float32) acc16 = np.zeros((M, N), np.float16) for k in range(K): prod = a[:, k:k + 1].astype(np.float32) * b[k:k + 1, :].astype(np.float32) acc32 += prod acc16 = (acc16.astype(np.float32) + prod).astype(np.float16) scale = np.sqrt(K) rows.append([K, f"{np.abs(acc32 - ref).max() / scale:.1e}", f"{np.abs(acc16 - ref).max() / scale:.1e}", f"{np.abs(acc16 - ref).max() / np.abs(acc32 - ref).max():.0f}"]) print_table(["K（FP16 输入）", "FP32 累加的最大误差 / √K", "FP16 累加的最大误差 / √K", "FP16 / FP32"], rows) print_observations([ "\"K 方向忘了 mask\"的内核在所有维度都是块大小倍数的形状上**全部通过**——最后一个 K 块正好不越界；只有 K 不是 32 的倍数时才出错" "（它读到了下一行的数，当成了这一行的）。所以测试集一定要有边缘形状：1、质数、比块大小多一点或少一点。", "容差要随 K 放大（上表用的是正比于 √K 的容差）：K 个乘积相加，累加的舍入误差随 K 增长（下表里 FP32 累加的误差 / √K 也在变大，" "说明绝对误差增长得比 √K 还快；前置讲 P3 §4）。用固定的 1e-6 去比，大 K 的正确内核也会\"失败\"；" "用太松的容差，又会放过真正的 bug。KernelBench 用的是 atol = rtol = 1e-2（第 14 讲会讨论这个选择的后果）。", "FP16 累加的误差是 FP32 累加的几千倍（表里约 7000–9000 倍，差不多就是两种格式的机器精度之比 2^13）——FP16 只有 11 位尾数，" "累加到几十以后每加一个数都要舍掉很多位。Tensor Core 的 GEMM 默认 FP32 累加就是这个原因（04 脚本演示 2）。", "测试一个内核的清单：和高精度的参考实现比较；随机形状 + 边缘形状；随 K 调整的容差；多次运行（竞争导致的错误不一定每次出现，02 演示 3）；" "再用 compute-sanitizer 查越界和竞争。", ]) if __name__ == "__main__": demo1_matmul() demo2_launch_order() demo3_fusion() demo4_testing()