""" ============================================================================= CS349E · 第 10 讲 · 04_tensor_cores.py —— Tensor Core 内核：warp 级的 mma、swizzle、喂数据的账 ============================================================================= 【这个脚本要回答的问题】 k1–k6 用的是 CUDA core（FFMA，一条指令 32 次乘加）。Tensor Core 的一条 mma 指令做 2048 次乘加， 算力高一个数量级——写内核的难点于是从"算"彻底变成了"喂"。 演示 1：mma.m16n8k16 的片段布局：一个 warp 的 32 个线程各持有矩阵的哪几个元素；模拟一条 mma，和 numpy 比较。 演示 2：一个块级的 Tensor Core 矩阵乘法（FP16 输入、FP32 累加）：数据先进共享内存，再用 ldmatrix 读成片段； 数 ldmatrix 的 bank 冲突，比较行主序和 XOR swizzle 两种布局。 演示 3：喂饱 Tensor Core：每次乘加能从共享内存拿几个字节？warp 的分块要多大？为什么 Hopper 要有 WGMMA。 演示 4：多级流水线：一个阶段的数据要多久才能到、共享内存放得下几级（Little 定律）。 【运行】python3 04_tensor_cores.py （约 2 秒） 【依赖】numpy；本目录的 kernel_lab；../lecture08 的 hw_engines（各代 Tensor Core 的参数）；../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")) sys.path.insert(0, os.path.join(HERE, "..", "lecture08")) from kernel_lab import dram_bytes_tiled, ldmatrix_wavefronts, mma_frag_coords, swizzle_offset, warp_mma # noqa: E402 from hw_engines import GPU_GENS # noqa: E402 from console import fmt_bytes, fmt_time, print_header, print_observations, print_table # noqa: E402 def demo1_fragments(): print_header("演示 1：mma.m16n8k16 的片段布局（D[16×8] = A[16×16]·B[16×8] + C，一个 warp 一条指令）") co = mma_frag_coords() rows = [] for lane in (0, 1, 2, 3, 4, 31): a = " ".join(f"({r},{c})" for r, c in co["A"][lane]) b = " ".join(f"({r},{c})" for r, c in co["B"][lane]) c = " ".join(f"({r},{c})" for r, c in co["C"][lane]) rows.append([lane, lane >> 2, lane % 4, a, b, c]) print_table(["lane", "groupID", "t", "A 片段 a0..a7 的 (行,列)", "B 片段 b0..b3", "C/D 片段 c0..c3"], rows) rng = np.random.default_rng(0) errs = [] for _ in range(100): Am = rng.standard_normal((16, 16)).astype(np.float16) Bm = rng.standard_normal((16, 8)).astype(np.float16) Cm = rng.standard_normal((16, 8)).astype(np.float32) d = warp_mma(Am[co["A"][..., 0], co["A"][..., 1]], Bm[co["B"][..., 0], co["B"][..., 1]], Cm[co["C"][..., 0], co["C"][..., 1]]) D = np.zeros((16, 8), np.float32) D[co["C"][..., 0], co["C"][..., 1]] = d errs.append(np.abs(D - (Am.astype(np.float64) @ Bm.astype(np.float64) + Cm)).max()) print(f"\n100 条随机的 mma：拼回的 D 和 numpy（FP64）的最大误差 {max(errs):.1e}") print(f"一条 mma = 16×8×16 = {16 * 8 * 16} 次乘加；一条 FFMA = 32 次。每个 lane 持有 A 的 8 个、B 的 4 个、C 的 4 个元素，" f"32 个 lane 合起来正好是 256 + 128 + 128 个，不重不漏。") print_observations([ "片段布局是**指令规定的**：lane 0 必须拿着 A 的 (0,0)(0,1)(8,0)(8,1)(0,8)(0,9)(8,8)(8,9)。同一行的相邻两个 FP16 放在一个 32 位寄存器里" "（所以列号总是成对出现）。写 Tensor Core 内核的第一步，就是把共享内存里的数按这个布局搬进寄存器。", "A 的元素在 4 个相邻 lane（t = 0..3）之间按列分，在 8 组（groupID）之间按行分；C 的布局和 A 的左半边一样——" "这让上一条 mma 的输出可以直接当下一条的输入（FlashAttention 里 P = softmax(S) 接着乘 V 就利用了这一点，第 11 讲）。", "误差约 2e-6：FP16 × FP16 的乘积在 FP32 里是精确的（两个 11 位尾数相乘最多 22 位），误差只来自 16 项和 C 在 FP32 里相加的舍入。", ]) def tc_matmul(A16, B16, BM=64, BN=64, BK=32, WM=32, WN=32, swizzle=False): """ 一个块级的 Tensor Core 矩阵乘法（功能模拟 + 计数）：C = A·B，A、B 是 FP16，累加用 FP32。 【结构】每个块算 BM×BN，4 个 warp 排成 2×2，每个 warp 算 WM×WN = 32×32 = 2×4 个 mma 的输出块。 每个 K 步（BK）： 1. 把 A 的 BM×BK 块、B 的 BK×BN 块写进共享内存（B 按"列"存成 Bt[BN][BK]，这样 A、B 的片段都是 "每行 16 字节"的 8×8 小矩阵，都能用 ldmatrix 读）；物理地址由 swizzle_offset 决定； 2. 每 16 个 k：每个 warp 用 ldmatrix.x4 读 A 的 2 个 16×16 块、B 的 2 个 16×16 块，做 8 条 mma。 【计数】mma 条数；ldmatrix 条数、实际轮数、理想轮数。 【返回】(C, stats) 【边界】M、N、K 是 BM、BN、BK 的倍数；BK 是 16 的倍数；WM、WN 是 16 的倍数。 """ M, K = A16.shape _, N = B16.shape co = mma_frag_coords() C = np.zeros((M, N), dtype=np.float32) st = {"mma": 0, "ldm": 0, "ldm_wf": 0, "ldm_ideal": 0} warps = [(wm, wn) for wm in range(BM // WM) for wn in range(BN // WN)] rr, cc = np.meshgrid(np.arange(BM), np.arange(BK), indexing="ij") rb, cb = np.meshgrid(np.arange(BN), np.arange(BK), indexing="ij") offA = swizzle_offset(rr, cc, BK, swizzle=swizzle) // 2 # 元素下标（FP16 = 2 字节） offB = swizzle_offset(rb, cb, BK, swizzle=swizzle) // 2 for bm in range(0, M, BM): for bn in range(0, N, BN): acc = {(w, mi, ni): np.zeros((32, 4), np.float32) for w in warps for mi in range(WM // 16) for ni in range(WN // 8)} for bk in range(0, K, BK): smA = np.zeros(BM * BK, np.float16) smB = np.zeros(BN * BK, np.float16) smA[offA] = A16[bm:bm + BM, bk:bk + BK] # 写共享内存（经过 swizzle） smB[offB] = B16[bk:bk + BK, bn:bn + BN].T # B 转成 [n][k] 存 assert len(np.unique(offA)) == offA.size # swizzle 是一一映射，不会互相覆盖 for ks in range(0, BK, 16): for (wm, wn) in warps: afr, bfr = {}, {} for mi in range(WM // 16): r0 = wm * WM + mi * 16 wf, ideal = ldmatrix_wavefronts(r0, ks, BK, swizzle) st["ldm"] += 1 st["ldm_wf"] += wf st["ldm_ideal"] += ideal rows = r0 + co["A"][..., 0] cols = ks + co["A"][..., 1] afr[mi] = smA[swizzle_offset(rows, cols, BK, swizzle=swizzle) // 2] for nj in range(WN // 16): # 一条 x4 读两个 n8 块 n0 = wn * WN + nj * 16 wf, ideal = ldmatrix_wavefronts(n0, ks, BK, swizzle) st["ldm"] += 1 st["ldm_wf"] += wf st["ldm_ideal"] += ideal for h in range(2): nrow = n0 + h * 8 + co["B"][..., 1] # B 片段的列 = Bt 的行 kcol = ks + co["B"][..., 0] bfr[2 * nj + h] = smB[swizzle_offset(nrow, kcol, BK, swizzle=swizzle) // 2] for mi in range(WM // 16): for ni in range(WN // 8): key = ((wm, wn), mi, ni) acc[key] = warp_mma(afr[mi], bfr[ni], acc[key]) st["mma"] += 1 for ((wm, wn), mi, ni), frag in acc.items(): r = bm + wm * WM + mi * 16 + co["C"][..., 0] c = bn + wn * WN + ni * 8 + co["C"][..., 1] C[r, c] = frag return C, st def demo2_block_kernel(): print_header("演示 2：块级 Tensor Core 矩阵乘法（128³，FP16 输入 / FP32 累加），ldmatrix 的 bank 冲突") rng = np.random.default_rng(1) n = 128 A16 = rng.standard_normal((n, n)).astype(np.float16) B16 = rng.standard_normal((n, n)).astype(np.float16) ref = A16.astype(np.float64) @ B16.astype(np.float64) rows = [] for BK in (32, 64): for sw in (False, True): C, st = tc_matmul(A16, B16, BK=BK, swizzle=sw) rows.append([f"BK = {BK}（每行 {BK * 2} 字节）", "XOR swizzle" if sw else "行主序", f"{np.abs(C - ref).max():.1e}", st["mma"], st["ldm"], st["ldm_wf"], st["ldm_ideal"], f"{st['ldm_wf'] / st['ldm_ideal']:.0f} 路"]) print_table(["共享内存块", "布局", "最大误差（vs FP64）", "mma 条数", "ldmatrix 条数", "实际轮数", "理想轮数", "冲突"], rows) print(f"\nmma 条数 = 128³ / (16·8·16) = {128 ** 3 // (16 * 8 * 16)}；每条 ldmatrix.x4 读 512 字节，理想是 4 轮。") lines = [] for r in range(8): o_plain = int(swizzle_offset(r, 0, 64, swizzle=False)) o_sw = int(swizzle_offset(r, 0, 64, swizzle=True)) lines.append(f"第 {r} 行第 0 列：行主序字节 {o_plain:4d}（bank {o_plain // 4 % 32:2d}） → swizzle 后 {o_sw:4d}（bank {o_sw // 4 % 32:2d}）") print("BK = 64 时，同一列的 8 行在共享内存里的位置：\n " + "\n ".join(lines)) print_observations([ "结果和 FP64 的参考相差约 5e-6：输入本来就是 FP16、乘积在 FP32 里精确，误差全部来自 128 项在 FP32 里累加的舍入（前置讲 P3 §4）。" "如果累加器也用 FP16，误差会大几个数量级——所以 Tensor Core 的 GEMM 几乎总用 FP32 累加（06 脚本演示 4）。", "行主序存放时，ldmatrix 读一个 8×8 小矩阵的 8 行，每行 16 字节：BK = 64 时行宽 128 字节，8 行都从同一组 4 个 bank 开始——8 路冲突；" "BK = 32 时行宽 64 字节，每两行错开半圈——4 路冲突。", "XOR swizzle 把\"16 字节的块号\"和\"行号 mod 8\"做异或：同一列的 8 行落到 8 个不同的 16 字节块上，正好铺满 32 个 bank，" "冲突消失，而且不浪费一个字节（padding 也能解冲突，但会让 TMA / cp.async 的整块拷贝变复杂）。", "swizzle 只改变**物理地址**：写进去和读出来都经过同一个映射，结果完全一样（表里两种布局的误差相同）。" "CUTLASS / CuTe、Triton、FlashAttention 都这样做；Hopper 的 TMA 硬件直接支持几种 swizzle 模式。", ]) def demo3_feeding(): print_header("演示 3：喂饱 Tensor Core——每次乘加能从共享内存拿几个字节？（每 SM 每周期 128 字节）") rows = [] for g in GPU_GENS: fma = g.tc_per_sm * g.tc_fma_per_clk rows.append([g.name, fma, f"{g.fp32_fma_per_sm}", f"{128 / fma:.3f}", f"{fma / 128:.0f}"]) print_table(["GPU", "Tensor Core FP16 乘加 / SM / 周期", "FP32 乘加（CUDA core）", "每次乘加能分到的共享内存字节", "需要的复用（乘加 / 共享内存字节）"], rows) print() rows = [] tiles = [("一条 mma（16×8）", 16, 8, False), ("warp 32×32", 32, 32, False), ("warp 64×32", 64, 32, False), ("warp 64×64", 64, 64, False), ("WGMMA 64×256（A、B 都从共享内存）", 64, 256, False), ("WGMMA 64×256（A 在寄存器）", 64, 256, True)] need = {g.name: g.tc_per_sm * g.tc_fma_per_clk / 128 for g in GPU_GENS} for name, wm, wn, a_in_reg in tiles: fma = wm * wn * 16 byts = (0 if a_in_reg else wm * 16 * 2) + wn * 16 * 2 ratio = fma / byts acc_regs = wm * wn // (128 if "WGMMA" in name else 32) rows.append([name, f"{ratio:.1f}", acc_regs, *[("够" if ratio >= need[g.name] else f"差 {need[g.name] / ratio:.1f} 倍") for g in GPU_GENS[1:]]]) print_table(["每个 k = 16 步、一个 warp（或 warpgroup）算的输出块", "乘加 / 共享内存字节", "每线程的 FP32 累加器个数", "A100（需 8）", "H100（需 16）"], rows) n = 4096 t_c = 2 * n ** 3 / 989e12 rows = [] for bm, bn in ((64, 64), (128, 128), (128, 256)): byts = dram_bytes_tiled(n, bm, bn, 1, 132, elem=2) rows.append([f"{bm}×{bn}", fmt_bytes(byts), fmt_time(byts / 3.35e12), fmt_time(t_c)]) print(f"\n再往外一层：4096³ FP16 矩阵乘法在 H100 上（132 个 SM，每 SM 一个块；波次模型，见 kernel_lab.dram_bytes_tiled）：") print_table(["块的大小 BM×BN", "显存流量", "按 3.35 TB/s 的时间", "按 989 TFLOPS 的计算时间"], rows) print_observations([ "FP32 的 CUDA core 每 SM 每周期 64–128 次乘加，共享内存每次乘加能给 1–2 字节；Tensor Core 是 512–2048 次，" "每次乘加只能分到 0.06–0.25 字节——数据在寄存器里被复用的次数要再高一个数量级。", "一个 warp 做 k = 16 的一步，复用 = WM·WN/(2·(WM + WN))。H100 要 16，至少需要 64×64 的 warp 块，" "可那意味着每个线程 128 个 FP32 累加器，寄存器就快用完了。", "Hopper 的回答是 WGMMA：4 个 warp（一个 warpgroup）合作一条指令，B 直接从共享内存读、不经过寄存器，" "A 可以在寄存器里；64×256 的块复用达到 26–32，累加器分摊到 128 个线程上（第 8 讲 §7.2）。", "块越大，显存流量越少：64×64 的块读显存要约 340 µs，比计算的约 140 µs 还长，整个内核会受限于带宽；" "128×128 降到约 100 µs，128×256 约 60 µs，才落到计算时间以下——所以 Hopper 上的 GEMM 和 FlashAttention-3 都用很大的块，" "再配上下面的多级流水线。", ]) def demo4_stages(): print_header("演示 4：多级流水线要几级？（H100：每 SM 每周期 2048 次 FP16 乘加，1.83 GHz，共享内存每块最多 227 KB）") clock = 1.83e9 rows = [] for bm, bn, bk in ((128, 128, 32), (128, 128, 64), (128, 256, 64), (256, 128, 64)): stage_bytes = (bm + bn) * bk * 2 t_stage = bm * bn * bk / 2048 / clock max_st = (227 * 1024) // stage_bytes need = {lat: 1 + int(np.ceil(lat / t_stage)) for lat in (0.5e-6, 1.0e-6)} bw = stage_bytes / t_stage rows.append([f"{bm}×{bn}×{bk}", fmt_bytes(stage_bytes, binary=True), fmt_time(t_stage), f"{bw / 1e9:.0f} GB/s", need[0.5e-6], need[1.0e-6], max_st]) print_table(["块 BM×BN×BK", "每级的数据", "每级的计算时间", "每 SM 需要的输入带宽", "需要几级（延迟 0.5 µs）", "需要几级（延迟 1 µs）", "共享内存放得下几级"], rows) print_observations([ "Little 定律：要让 Tensor Core 不停，在途的数据 = 需要的带宽 × 延迟。一级的数据算完只要 0.14–0.56 µs，" "而从 L2 / 显存取一块要零点几到一微秒，所以除了正在算的那一级，大块（BK = 64）还要 1–4 级在路上，" "小块（128×128×32）要 4–8 级。", "共享内存是上限：128×256×64 的块每级 48 KiB，227 KiB 只放得下 4 级。块越大、复用越好，但能放的级数越少——又是一个折中。", "\"每 SM 需要的输入带宽\"约 90 GB/s，乘以 132 个 SM 约 11 TB/s，远超显存的 3.35 TB/s：这些数据大部分来自 L2" "（同一个 A 块被同一行的多个块共用），Hopper 的 TMA multicast 还让一个 cluster 里的几个 SM 共享一次 L2 读取。", "写法上：Ampere 用 cp.async（每个线程发异步拷贝），Hopper 用 TMA（一个线程发一个描述符，硬件搬整块）+ mbarrier 计数到达，" "再配合 warp 专门化——一个 producer warp 只管发 TMA，consumer warpgroup 只管 WGMMA（第 8 讲 §8.3–§8.5、本讲 §8）。", ]) if __name__ == "__main__": demo1_fragments() demo2_block_kernel() demo3_feeding() demo4_stages()