""" ============================================================================= 第 5 讲 · 代码 2：GPU 的执行模型 —— 线程、warp、分支发散、合并访存、占用率、波次 ============================================================================= 对应讲义：lectures/lecture05_gpus_tpus.md 第 1、4 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- GPU 编程模型里有几个抽象的词：网格（grid）、线程块（block）、线程（thread）、warp。 它们本身只决定"谁算哪个元素"（正确性），但它们如何映射到硬件，决定了"算得快不快"（性能）。 我们没有 NVIDIA GPU 也能理解它们：这个文件用 Python 模拟了其中最重要的几件事， 并在你的电脑上实测一个与"合并访存"类似的现象： 1. 线程下标：每个线程怎么知道自己该处理哪个元素？ 2. warp 与分支发散：同一个 warp 的 32 个线程必须执行同一条指令 3. 合并访存：同一个 warp 的访存落在同一段连续内存里才高效 4. 占用率：一个 SM 上能同时驻留多少个线程？ 5. 波次量化：为什么矩阵边长从 1792 变成 1793，速度会突然下降？ ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 02_gpu_execution_model.py """ import math import time import torch def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) WARP_SIZE = 32 # ============================================================================= # 第 1 部分：线程下标 # ============================================================================= def simulate_vector_add(a, b, block_size): """ 模拟一次 GPU 核函数调用：c[i] = a[i] + b[i]。 【GPU 的做法】 启动一个"网格"，里面有 num_blocks 个线程块，每个块有 block_size 个线程。 每个线程执行**同一段代码**，但通过自己的编号知道该处理哪个元素： i = blockIdx × blockDim + threadIdx 因为 n 不一定是 block_size 的整数倍，最后一个块里会有"多余"的线程， 所以代码里必须有边界检查：if i < n。 【参数】 a, b : 一维列表 block_size : 每个块的线程数（常用 128、256 等 32 的倍数） 【返回】 (c, 启动的线程总数) """ n = len(a) num_blocks = math.ceil(n / block_size) # 向上取整，保证覆盖所有元素 c = [0.0] * n launched = 0 for block_idx in range(num_blocks): # 真实 GPU 上，这些块在不同 SM 上并行执行 for thread_idx in range(block_size): # 块内的线程也是并行的 launched += 1 i = block_idx * block_size + thread_idx if i < n: # 边界检查！否则会越界访问内存 c[i] = a[i] + b[i] return c, launched def demo_indexing(): """演示 1：线程下标与边界检查。""" print_header("演示 1：每个线程怎么知道自己该算哪个元素？") n, bs = 1000, 256 a = [float(i) for i in range(n)] b = [1.0] * n c, launched = simulate_vector_add(a, b, bs) assert c == [x + 1.0 for x in a] print(f" n = {n}，每块 {bs} 个线程 → 启动 {math.ceil(n / bs)} 个块、共 {launched} 个线程，" f"其中 {launched - n} 个线程因为 i >= n 什么也不做") print(" CUDA 里的写法：int i = blockIdx.x * blockDim.x + threadIdx.x; if (i < n) c[i] = a[i] + b[i];") print(" （第 6 讲的 Triton 以'块'为单位编程，同样需要用掩码 mask = offsets < n 做边界检查。）") # ============================================================================= # 第 2 部分：warp 与分支发散 # ============================================================================= def warp_execution_steps(condition, cost_if, cost_else): """ 模拟一个 warp（32 个线程）执行 if/else 所需的步数。 【思路】 GPU 的一个 warp 中，32 个线程在同一时刻执行**同一条指令**（SIMT：单指令多线程）。 如果 warp 中有的线程走 if、有的走 else，硬件只能**先执行 if 分支（else 的线程闲着）， 再执行 else 分支（if 的线程闲着）**。这叫"分支发散"（control divergence）。 【参数】 condition : 长度 32 的布尔列表，第 t 个线程的条件是否成立 cost_if, cost_else : 两个分支各需要的步数 【返回】 (执行步数, 计算单元的利用率) """ n_if = sum(condition) n_else = WARP_SIZE - n_if steps = (cost_if if n_if > 0 else 0) + (cost_else if n_else > 0 else 0) useful = n_if * cost_if + n_else * cost_else return steps, useful / (steps * WARP_SIZE) def demo_divergence(): """演示 2：同样一半线程走 if，按线程编号分还是按 warp 分，效率完全不同。""" print_header("演示 2：warp 与分支发散") num_threads = 128 # 4 个 warp patterns = { "if (threadIdx % 2 == 0)": [t % 2 == 0 for t in range(num_threads)], "if ((threadIdx / 32) % 2 == 0)": [(t // 32) % 2 == 0 for t in range(num_threads)], "if (threadIdx < 100)": [t < 100 for t in range(num_threads)], } for name, cond in patterns.items(): total_steps, utils = 0, [] for w in range(num_threads // WARP_SIZE): steps, util = warp_execution_steps(cond[w * 32:(w + 1) * 32], cost_if=10, cost_else=10) total_steps += steps utils.append(util) print(f" {name:<32} 4 个 warp 共 {total_steps} 步，平均利用率 {sum(utils) / len(utils):.0%}") print(""" 前两种写法都让恰好一半线程走 if 分支，但： - 按线程编号的奇偶分：每个 warp 里都有两种线程，每个 warp 都要把两个分支各走一遍，利用率 50%； - 按 warp 分：同一个 warp 内的线程走同一条路，没有发散，利用率 100%。 第三种（threadIdx < 100）：只有"跨越边界"的那个 warp（线程 96~127，一部分 < 100、一部分 >= 100） 发生发散，其余 3 个 warp 不发散。 条件判断本身不是问题，问题是**同一个 warp 内**走了不同的路。""") # ============================================================================= # 第 3 部分：合并访存 # ============================================================================= def memory_transactions(addresses, segment_bytes=128): """ 一个 warp 的 32 次访存，一共触及多少个 128 字节的内存段。 【思路】 显存（DRAM）以"突发"（burst）方式读取：读一次就把一整段连续的字节取出来。 如果一个 warp 的 32 个线程访问的地址都落在同一段里，一次事务就够了（"合并访存"，coalesced）； 如果它们分散在 32 个不同的段里，就要 32 次事务，大部分取回来的字节都浪费了。 【参数】addresses : 32 个字节地址 """ return len({a // segment_bytes for a in addresses}) def demo_coalescing(): """演示 3：访存模式决定了内存事务的次数。""" print_header("演示 3：合并访存") elem = 4 # float32 patterns = { "线程 t 读 x[t]（连续）": [t * elem for t in range(32)], "线程 t 读 x[2t]（步长 2）": [2 * t * elem for t in range(32)], "线程 t 读 x[32t]（步长 32）": [32 * t * elem for t in range(32)], "线程 t 读 x[随机位置]": [((t * 7919) % 100000) * elem for t in range(32)], } for name, addrs in patterns.items(): n = memory_transactions(addrs) print(f" {name:<30} {n:>2} 次 128 字节事务，有效字节占比 {32 * elem / (n * 128):.0%}") print("\n 矩阵按行存储（行优先），1024 × 1024 的 float32 矩阵：") ncol = 1024 along_row = [(0 * ncol + t) * elem for t in range(32)] # 32 个线程读同一行的相邻 32 列 along_col = [(t * ncol + 0) * elem for t in range(32)] # 32 个线程读同一列的相邻 32 行 print(f" 32 个线程读同一行的相邻元素：{memory_transactions(along_row)} 次事务") print(f" 32 个线程读同一列的相邻元素：{memory_transactions(along_col)} 次事务") print(" → 写核函数时，要让同一个 warp 里相邻的线程访问内存中相邻的地址。") # 在你的电脑上的类似现象：CPU 缓存行 / GPU 访存都偏爱连续访问 device = torch.device("mps") if torch.backends.mps.is_available() else ( torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")) m = torch.randn(8192, 8192, device=device) def timed(fn, reps=5): fn() if device.type == "mps": torch.mps.synchronize() elif device.type == "cuda": torch.cuda.synchronize() t0 = time.perf_counter() for _ in range(reps): fn() if device.type == "mps": torch.mps.synchronize() elif device.type == "cuda": torch.cuda.synchronize() return (time.perf_counter() - t0) / reps t_copy = timed(lambda: m.clone()) t_transpose = timed(lambda: m.t().contiguous()) print(f"\n 在你的电脑上（{device}）复制一个 8192×8192 矩阵：") print(f" 按顺序复制（连续读、连续写）：{t_copy * 1e3:.1f} ms") print(f" 转置后复制（跨步读、连续写）：{t_transpose * 1e3:.1f} ms（慢 {t_transpose / t_copy:.1f} 倍）") print(" 数据量完全相同，只是访问顺序不同。高性能的转置核函数会用共享内存分块来避免跨步访问。") # ============================================================================= # 第 4 部分：占用率 # ============================================================================= H100_SM = dict(max_threads=2048, max_blocks=32, registers=65536, shared_mem=228 * 1024) def occupancy(threads_per_block, regs_per_thread, smem_per_block, sm=H100_SM): """ 计算一个 SM 上能同时驻留多少个线程块，以及占用率（驻留的 warp 数 / 最大 warp 数）。 【思路】 每个 SM 的资源是有限的：寄存器总数、共享内存总量、最多线程数、最多线程块数。 一个线程块能否驻留，要看这四样资源是否都够用 —— 取四个上限中最小的那个。 驻留的 warp 越多，一个 warp 等待内存时，SM 就能切换到另一个 warp 继续算（零开销切换）， 从而"隐藏"内存延迟。 【返回】(每个 SM 驻留的线程块数, 占用率, 限制因素) """ limits = { "线程数": sm["max_threads"] // threads_per_block, "线程块数": sm["max_blocks"], "寄存器": sm["registers"] // (threads_per_block * regs_per_thread), "共享内存": sm["shared_mem"] // smem_per_block if smem_per_block else 10 ** 9, } blocks = min(limits.values()) bottleneck = min(limits, key=limits.get) warps = blocks * threads_per_block // WARP_SIZE return blocks, warps / (sm["max_threads"] // WARP_SIZE), bottleneck def demo_occupancy(): """演示 4：寄存器和共享内存如何限制占用率（以 H100 的 SM 为例）。""" print_header("演示 4：占用率（H100：每个 SM 65536 个寄存器、228 KB 共享内存、最多 2048 线程）") cases = [ (256, 32, 0), (256, 64, 0), (128, 160, 0), (256, 32, 64 * 1024), (1024, 32, 0), ] print(f" {'每块线程数':>10}{'每线程寄存器':>12}{'每块共享内存':>14}{'驻留块数':>10}{'占用率':>8} 限制因素") for tpb, regs, smem in cases: blocks, occ, why = occupancy(tpb, regs, smem) print(f" {tpb:>10}{regs:>12}{smem // 1024:>11} KB{blocks:>10}{occ:>8.0%} {why}") print(""" 占用率低不一定慢：如果每个线程做很多工作（"线程粗化"，一个线程处理多个元素）， 或者算法本身能很好地复用数据（比如分块矩阵乘法），低占用率也能跑得很快。 但占用率太低时，GPU 就没有足够的 warp 来掩盖内存延迟。""") # ============================================================================= # 第 5 部分：波次量化 # ============================================================================= def wave_utilization(n, tile_m, tile_n, num_sms): """ n×n 矩阵乘法被切成 tile_m × tile_n 的输出块，每个块交给一个 SM；块数超过 SM 数时要分多"波"执行。 【返回】(块数, 波数, 利用率 = 块数 / (波数 × SM 数)) """ tiles = math.ceil(n / tile_m) * math.ceil(n / tile_n) waves = math.ceil(tiles / num_sms) return tiles, waves, tiles / (waves * num_sms) def demo_wave_quantization(): """演示 5：CS336 课件中的"矩阵之谜" —— 1792 与 1793。""" print_header("演示 5：波次量化（A100，108 个 SM，输出块 256 × 128）") print(f" {'n':>6}{'输出块数':>10}{'波数':>6}{'SM 利用率':>10}{'相对 1792 的计算量':>18}{'相对耗时（估计）':>18}") base_tiles, base_waves, _ = wave_utilization(1792, 256, 128, 108) for n in [1536, 1792, 1793, 2048, 2304, 2560]: tiles, waves, util = wave_utilization(n, 256, 128, 108) print(f" {n:>6}{tiles:>10}{waves:>6}{util:>10.0%}{(n / 1792) ** 3:>18.2f}{waves / base_waves:>18.1f}") print(""" n = 1792：7 × 14 = 98 个块，108 个 SM 一波就能做完。 n = 1793：多出一行一列，块数变成 8 × 15 = 120 个 —— 超过了 108，必须分两波， 第二波只有 12 个块在干活，96 个 SM 闲着。计算量只多了 0.2%，耗时却接近翻倍。 （这里的"相对耗时"假设每个块耗时相同，只是一个粗略的模型；真实的库会根据形状选择不同的分块方案， 所以实测曲线是"锯齿状"的，而不是这么极端。） 启示：模型的维度（隐藏维度、词表大小、批次大小）尽量取 64、128 等 2 的幂次的倍数。 Karpathy 在 nanoGPT 中把词表从 50257 填充到 50304（64 的倍数），就带来了明显的加速。""") def main(): demo_indexing() demo_divergence() demo_coalescing() demo_occupancy() demo_wave_quantization() print("\n全部演示运行完毕。") if __name__ == "__main__": main()