""" ============================================================================= CS349E · 第 1 讲 · 03_measure_your_machine.py —— 在你自己的电脑上量出一条 roofline ============================================================================= 【这个脚本要回答的问题】 前两个脚本的 roofline 都是"纸上的"：用厂商标的峰值算力和带宽算出来的。真实硬件上它还成立吗？ 我们没有 GPU，但 roofline 是一个**通用**模型：任何"会算、会搬"的机器都适用，包括你的笔记本 CPU。 本脚本用 numpy 测三件事，然后把它们放在一张图上： 演示 1：主存带宽——复制一个大数组、对一个大矩阵做矩阵-向量乘，每秒能搬多少字节。 演示 2：峰值算力——一个大方阵乘法每秒能做多少次运算。 演示 3：扫描 X[B, D]·W[D, F] 的 B（D = F = 8192，fp32），看实测 FLOP/s 怎样随强度变化， 和"实测带宽 × 强度、封顶于实测峰值"的 roofline 预测对比。 【怎么读输出】这是**墙钟时间测量**：数字取决于你的 CPU、内存、numpy 用的 BLAS 库、当时后台的负载， 每次运行都会略有不同，和讲义里的数字也不会完全一样。**只看趋势**：小 B 时 FLOP/s 随 B 线性上涨 （memory-bound），大 B 时趋于平坦（compute-bound）。 编写讲义时用的是一台 Apple M4 Pro 笔记本（numpy 2.5，底层是 Apple Accelerate）。 【运行】python3 03_measure_your_machine.py （约 10–30 秒） python3 03_measure_your_machine.py --quick （几秒，数组更小，只检查能不能跑通） 【依赖】numpy。 """ import platform import sys import time import numpy as np from console import print_header, print_observations, print_table, fmt_num, fmt_bytes, fmt_time from roofline import matmul_counts, attainable_flops, ascii_roofline QUICK = "--quick" in sys.argv def best_time(fn, repeats=5, warmup=1): """ 多次运行 fn()，返回最短的一次耗时（秒）。 【为什么取最短而不是平均】测量里的"噪声"几乎都是让程序**变慢**的东西（操作系统调度别的进程、 缓存被别人冲掉、CPU 降频……），不会让它变快。所以最短时间最接近"机器本来的能力"。 前 warmup 次不计时：第一次运行会碰到"内存页第一次被访问"、"库第一次初始化"之类的一次性开销。 【参数】fn：不带参数的函数；repeats：计时次数；warmup：预热次数 【返回】float，秒 【Python 小知识】time.perf_counter() 是高精度计时器，专门用来测"一段代码花了多久"； 它的绝对值没有意义，只能做差。 """ for _ in range(warmup): fn() best = float("inf") for _ in range(repeats): t0 = time.perf_counter() fn() best = min(best, time.perf_counter() - t0) return best # --------------------------------------------------------------------------- # 演示 1：带宽 # --------------------------------------------------------------------------- def demo1_bandwidth(): """ 用两种方式测"主存 → CPU"的带宽，返回较大的那个（作为 roofline 的斜坡）。 【方式 A：大数组复制】np.copyto(dst, src)：读 src、写 dst。我们按"读 + 写 = 2 × 数组大小"记流量。 （很多 CPU 写一块内存前会先把它读进 cache，真实流量可能是 3 倍，所以这个数偏保守。） numpy 的复制只用一个核。 【方式 B：矩阵-向量乘】W[D, F]·x[F]：W 要完整读一遍，每个数只用一次（强度 ≈ 0.5 FLOP/byte，fp32）， 所以它一定是 memory-bound，实测的"字节/秒"就是 BLAS 库（可能多线程）能拿到的带宽。 这正是 LLM decode 在做的事：batch = 1 时每一步就是一串矩阵-向量乘。 【数组大小】要远大于 CPU 的各级 cache（几 MB 到几十 MB），否则测到的是 cache 带宽。 【返回】(带宽 byte/s, 结果表格行) """ print_header("演示 1：主存带宽（每秒能搬多少字节）") n = (16 if QUICK else 64) * 2 ** 20 # fp32 元素个数：64 Mi 个 = 256 MiB src = np.ones(n, dtype=np.float32) dst = np.empty_like(src) t_copy = best_time(lambda: np.copyto(dst, src), repeats=3 if QUICK else 5) bw_copy = 2 * src.nbytes / t_copy D = 4096 if QUICK else 8192 W = np.ones((D, D), dtype=np.float32) # 8192² × 4 B = 268 MB x = np.ones(D, dtype=np.float32) t_mv = best_time(lambda: W @ x, repeats=3 if QUICK else 5) bytes_mv = W.nbytes + x.nbytes + D * 4 bw_mv = bytes_mv / t_mv rows = [ [f"复制 {fmt_bytes(src.nbytes)} 数组（单核）", fmt_bytes(2 * src.nbytes), fmt_time(t_copy), fmt_bytes(bw_copy) + "/s"], [f"矩阵-向量乘 {D}×{D} fp32（BLAS）", fmt_bytes(bytes_mv), fmt_time(t_mv), fmt_bytes(bw_mv) + "/s"], ] print_table(["测法", "搬运字节", "最短耗时", "带宽"], rows) bw = max(bw_copy, bw_mv) print_observations([ f"取两者较大的 {fmt_bytes(bw)}/s 作为这台机器 roofline 的'带宽屋顶'。", "单核复制通常拿不满整个内存系统的带宽：一个核同时在路上的内存请求数有限（Little 定律：带宽 = 并发请求数 × 每个请求的大小 / 延迟）。" "多线程的 BLAS 能发出更多并发请求，所以矩阵-向量乘往往测得更高。", "作为对比：H100 的 HBM 是 3.35 TB/s，Blackwell 是 8 TB/s——通常比笔记本内存快十几到几十倍。", ]) return bw # --------------------------------------------------------------------------- # 演示 2：峰值算力 # --------------------------------------------------------------------------- def demo2_peak(): """ 用大方阵乘法 A[n, n]·B[n, n]（fp32）测峰值算力。 【为什么用大方阵】强度 ≈ n/6（fp32，三个矩阵各 4n² 字节，运算 2n³），n = 4096 时约 680 FLOP/byte， 远高于任何 CPU 的屋脊点，一定是 compute-bound，测到的就是"能拿到的最大算力"。 【返回】(峰值 FLOP/s) """ print_header("演示 2：峰值算力（大方阵乘法，fp32）") rows = [] best = 0.0 for n in ((1024, 2048) if QUICK else (2048, 4096)): A = np.ones((n, n), dtype=np.float32) B = np.ones((n, n), dtype=np.float32) t = best_time(lambda: A @ B, repeats=3) flops = 2.0 * n ** 3 best = max(best, flops / t) rows.append([f"{n}×{n} · {n}×{n}", fmt_num(flops), fmt_time(t), fmt_num(flops / t) + "FLOP/s"]) print_table(["矩阵", "FLOPs", "最短耗时", "实测 FLOP/s"], rows) print_observations([ f"取 {fmt_num(best)}FLOP/s 作为这台机器 roofline 的'算力屋顶'（fp32）。", "作为对比：H100 的 bf16 Tensor Core 峰值约 989 TFLOP/s。CPU 和 GPU 的差距主要在算力，带宽的差距小一些——" "所以 CPU 的屋脊点通常比 GPU 低，更'容易' compute-bound。", ]) return best # --------------------------------------------------------------------------- # 演示 3：扫描 batch，画出实测 roofline # --------------------------------------------------------------------------- def demo3_sweep(bw, peak): """ 扫描 X[B, D]·W[D, F]（fp32）的 B，测实际 FLOP/s，和 roofline 预测 min(peak, I·bw) 对比。 【参数】bw、peak：演示 1、2 测出的带宽和峰值算力 【返回】[(B, 强度, 实测 FLOP/s), ...] 【注意】 - 字节数用"每个矩阵读/写一遍"的下界（roofline.matmul_counts，fp32 每个数 4 字节）。 - B 很小时（2–16 行）有的 BLAS 库会走一条对"瘦矩阵"不太优化的路径，可能比 B = 1 的矩阵-向量乘还慢； "等效带宽"一列（= 下界字节数 / 耗时）能直接看出这一点。这是内核实现的问题，不是 roofline 的问题。 """ print_header("演示 3：扫描 B——实测的 FLOP/s 怎样随算术强度变化（D = F = 8192，fp32）") D = F = 4096 if QUICK else 8192 W = np.random.default_rng(0).standard_normal((D, F), dtype=np.float32) Bs = (1, 4, 16, 64, 256) if QUICK else (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048) rows, pts = [], [] for B in Bs: X = np.ones((B, D), dtype=np.float32) t = best_time(lambda: X @ W, repeats=3 if B >= 512 else 5) flops, nbytes = matmul_counts(B, D, F, 4, 4, 4) I = flops / nbytes achieved = flops / t pred = attainable_flops(I, peak, bw) rows.append([B, f"{I:.1f}", fmt_time(t), fmt_num(achieved) + "FLOP/s", fmt_num(pred) + "FLOP/s", f"{achieved / pred * 100:.0f}%", f"{achieved / peak * 100:.0f}%", fmt_bytes(nbytes / t) + "/s"]) pts.append((B, I, achieved, nbytes / t)) print_table(["B", "强度", "最短耗时", "实测", "roofline 预测", "实测/预测", "实测/峰值", "等效带宽"], rows) labels = "123456789ABCDEFGH" plot_pts = [(labels[i], I, y) for i, (_, I, y, _) in enumerate(pts)] print() print("\n".join(ascii_roofline(plot_pts, peak, bw, xmin=0.1, xmax=3e3, title="实测 roofline（线 = 用实测峰值和带宽画的屋顶；数字 = 实测点）"))) print(" 图例：" + "，".join(f"{labels[i]}: B={B}" for i, (B, _, _, _) in enumerate(pts))) ridge = peak / bw first_half = next((B for B, _, y, _ in pts if y >= 0.5 * peak), None) small = [y / attainable_flops(I, peak, bw) for B, I, y, _ in pts if I < ridge / 4] big = [y / peak for B, I, y, _ in pts if I > ridge * 2] # B = 1（矩阵-向量乘，GEMV）和 B = 2..16（通用矩阵乘法 GEMM 处理"瘦"矩阵）的等效带宽对比 bw_gemv = pts[0][3] skinny = [e for B, _, _, e in pts if 2 <= B <= 16] obs = [ f"这台机器实测的屋脊点 = {fmt_num(peak)}FLOP/s ÷ {fmt_bytes(bw)}/s ≈ {ridge:.0f} FLOP/byte。", ] if small: obs.append(f"强度远低于屋脊点的那些 B（memory-bound 区），实测达到 roofline 预测的 " f"{min(small) * 100:.0f}%–{max(small) * 100:.0f}%：FLOP/s 大致随 B 线性上涨，时间几乎不变——" "多处理几行输入'几乎免费'，因为时间花在读 W 上。") if big: obs.append(f"强度远高于屋脊点的那些 B（compute-bound 区），实测达到峰值的 " f"{min(big) * 100:.0f}%–{max(big) * 100:.0f}%：曲线变平，再加大 B 时间开始正比增长。") if skinny and bw_gemv > 1.5 * max(skinny): obs.append(f"B = 1 时等效带宽 {fmt_bytes(bw_gemv)}/s，B = 2–16 却只有 {fmt_bytes(min(skinny))}/s–{fmt_bytes(max(skinny))}/s：" "numpy 对 B = 1 调用专门的矩阵-向量乘（GEMV）代码，B ≥ 2 改走通用矩阵乘法（GEMM），而这台机器的 GEMM " "对'很瘦'的输入没有把带宽用满。roofline 只给出**上界**，能贴多近取决于内核怎么写——这正是推理引擎要专门" "为 decode 的小批量写内核的原因（第 10–11 讲）。") if first_half is not None: obs.append(f"实测 FLOP/s 第一次超过峰值一半是在 B = {first_half}。GPU 上这个转折点要到 B ≈ 几百（见 02 脚本的演示 2），" "因为 GPU 的算力/带宽比更大。") obs.append("同一条 roofline 能解释从笔记本 CPU 到 H100 的现象，这就是它好用的原因；差别只在两个数（峰值、带宽）。") print_observations(obs) return pts if __name__ == "__main__": print(f"机器：{platform.machine()} / {platform.processor() or '未知处理器'}；Python {platform.python_version()}；" f"numpy {np.__version__}{'；--quick 模式' if QUICK else ''}") t0 = time.perf_counter() bw = demo1_bandwidth() peak = demo2_peak() demo3_sweep(bw, peak) print(f"\n⏱ 总耗时 {time.perf_counter() - t0:.1f} 秒")