""" ============================================================================= CS349E · 第 1 讲 · 02_roofline_basics.py —— roofline 模型：算术强度、临界 batch、分块、网络 ============================================================================= 【这个脚本要回答的问题】 一段计算在一块芯片上"至少要多久"、瓶颈在"算"还是在"搬"？roofline 模型说：只要比较两个数—— 这段计算的**算术强度**（每搬 1 字节做几次运算）和芯片的**临界强度**（峰值算力 / 带宽）。 本脚本按 How to Scale Your Model 第 1 章的顺序，把每个结论都用数字算一遍： 演示 1：几种基本运算的强度——向量加法、点积、矩阵-向量乘、矩阵乘法（随 B 变化）。 演示 2：各芯片、各精度的临界强度，以及"矩阵乘法要多大的 batch 才 compute-bound"。 演示 3：低精度怎样改变强度（书中练习 1、2），以及"每个样本一个矩阵"为什么救不了（练习 4） ——它正是 decode 阶段注意力的形状。 演示 4：一张字符画的 roofline 图，把上面的运算放上去（加 --png 还会存一张 PNG）。 演示 5：分块（tiling）：片上存储能放多大的块，决定了从 HBM 看到的强度——Blackwell 为什么加了 TMEM、 为什么让两个 SM 合作做一次矩阵乘法。 演示 6：网络 roofline：把矩阵乘法切到两块芯片上，临界条件只和 D 有关、和 B 无关。 【怎么读输出】所有数字都是公式算出来的（不测时间），每次运行完全相同。 【运行】python3 02_roofline_basics.py （不到 1 秒） python3 02_roofline_basics.py --png （另外保存 figures/02_roofline.png，需要 matplotlib） 【依赖】只用 Python 标准库；PNG 可选 matplotlib。 """ import os import sys from console import print_header, print_observations, print_table, fmt_num, fmt_bytes, fmt_time from hw_specs import H100, H200, B200, B300, TPU_V5E, TPU_V6E, BYTES_PER_ELEM from roofline import (time_bounds, attainable_flops, matmul_counts, critical_batch, batched_matvec_counts, tiled_matmul_bytes, sharded_matmul_two_chips, ascii_roofline, save_roofline_png) HERE = os.path.dirname(os.path.abspath(__file__)) # --------------------------------------------------------------------------- # 演示 1：基本运算的强度 # --------------------------------------------------------------------------- def basic_ops(N=2 ** 24, D=8192, F=8192): """ 列出几种基本运算（全部 bf16）的 FLOPs、字节数和强度。 【返回】[(名字, flops, nbytes), ...]，演示 1 和演示 4 共用 【各项的推导】 向量加法 z = x + y ：N 次加法；读 2N 个数、写 N 个数 → 强度 N / (6N) = 1/6 点积 x·y ：N 次乘 + (N-1) 次加；读 2N 个数、写 1 个数 → → 1/2 矩阵-向量乘 x[D]·W[D,F]：2DF 次运算；读 W（DF 个数）为主 → ≈ 1 矩阵乘法 X[B,D]·W[D,F] ：见 roofline.matmul_counts → ≈ B """ ops = [ ("向量加法 (N=16M)", float(N), 2.0 * 3 * N), ("点积 (N=16M)", 2.0 * N - 1, 2.0 * (2 * N + 1)), ] for B in (1, 16, 64, 256, 1024, 4096): f, b = matmul_counts(B, D, F) name = "矩阵-向量乘 (B=1)" if B == 1 else f"矩阵乘法 B={B}" ops.append((name, f, b)) return ops def demo1_basic_ops(): print_header("演示 1：几种基本运算的算术强度（bf16；D = F = 8192）") ridge = H100.critical_intensity("bf16") rows = [] for name, f, b in basic_ops(): I = f / b t = time_bounds(f, b, H100.peak_flops("bf16"), H100.mem_bw) rows.append([name, fmt_num(f), fmt_bytes(b), f"{I:.3f}" if I < 10 else f"{I:.1f}", "compute" if I >= ridge else "memory", f"{attainable_flops(I, H100.peak_flops('bf16'), H100.mem_bw) / H100.peak_flops('bf16') * 100:.1f}%", fmt_time(t["lower"])]) print_table(["运算", "FLOPs", "字节", "强度", f"H100 上（屋脊 {ridge:.0f}）", "可达峰值比例", "时间下界"], rows) print_observations([ "逐元素运算（向量加法）和点积的强度是不随规模变化的常数（1/6、1/2）：数据再多，每个数也只被用一次。" "它们在任何 AI 芯片上都是 memory-bound，只能用到峰值算力的零头。", "矩阵乘法的强度 ≈ B：权重 W 搬进来一次，被 B 行输入复用 B 次。B = 1（矩阵-向量乘）时强度 ≈ 1，" "和点积一样糟；B 涨到几百以上才越过 H100 的屋脊点 295。", "B 很大（4096）时强度不再'≈ B'：分母里 B·D 和 B·F 两项（读 X、写 Z）开始变得重要，强度趋于 " "DF/(D+F) 这样的上限（D = F = 8192 时是 4096）。", "这就是 LLM 推理的核心矛盾的雏形：decode 每步每个序列只有 1 个 token，所以要把很多序列凑成一批（B 大），" "矩阵乘法才不会被 HBM 带宽卡住。", ]) # --------------------------------------------------------------------------- # 演示 2：各芯片的临界强度与临界 batch # --------------------------------------------------------------------------- def demo2_critical(): print_header("演示 2：各芯片的临界强度，以及 D = F = 8192 的矩阵乘法要多大的 B 才 compute-bound") D = F = 8192 cases = [ # (芯片, 算力精度, X 字节, W 字节, 输出字节, 说明) (H100, "bf16", 2, 2, 2, "全 bf16"), (H100, "fp8", 1, 1, 2, "fp8 输入 + fp8 权重"), (H200, "bf16", 2, 2, 2, "全 bf16"), (B200, "bf16", 2, 2, 2, "全 bf16（算力为推算）"), (B200, "fp8", 1, 1, 2, "fp8 输入 + fp8 权重"), (B300, "fp8", 1, 1, 2, "fp8 输入 + fp8 权重"), (B300, "nvfp4", BYTES_PER_ELEM["nvfp4"], BYTES_PER_ELEM["nvfp4"], 2, "nvfp4 输入 + nvfp4 权重"), (TPU_V5E, "bf16", 2, 2, 2, "全 bf16"), (TPU_V5E, "int8", 1, 1, 1, "全 int8（书中练习 1）"), (TPU_V6E, "bf16", 2, 2, 2, "全 bf16"), ] rows = [] for chip, prec, xb, wb, ob, note in cases: R = chip.critical_intensity(prec) Bc = critical_batch(D, F, R, xb, wb, ob) rows.append([chip.name, prec, note, f"{R:.0f}", Bc if Bc is not None else "到不了"]) print_table(["芯片", "算力精度", "数据格式", "临界强度 FLOP/B", "临界 B（精确）"], rows) print_observations([ "主流 AI 芯片的 bf16 临界强度都在 200–600 FLOP/byte：要让矩阵乘法跑满算力，一次要处理几百个 token。", "H200 和 H100 是同一颗计算芯片，只是 HBM 更快（4.8 vs 3.35 TB/s），屋脊点从 295 降到 206：" "带宽变大，'更容易'达到 compute-bound。", "fp8 算力是 bf16 的 2 倍，权重和输入也变成 1 字节，强度同样翻倍——临界 B 和 bf16 几乎一样（332 vs 319）：" "低精度让'算'和'搬'同时变快，不会自动改变瓶颈。", "nvfp4 是例外：Blackwell Ultra 的 nvfp4 算力是 fp8 的 3 倍，而字节只少了 1.78 倍（4.5 bit vs 8 bit），" "屋脊点跳到 1875，临界 B 涨到约 750。芯片在 nvfp4 下'算力相对更富余'，要更大的批才能喂饱。", "临界 B 比'临界强度'略大一点：精确公式里还有读 X、写 Z 的字节（演示 1 最后一条）。", ]) # --------------------------------------------------------------------------- # 演示 3：低精度与"每个样本一个矩阵" # --------------------------------------------------------------------------- def demo3_precision_and_per_example(): print_header("演示 3：书中练习 1、2、4——低精度怎样改变强度；为什么 decode 的注意力救不了") D = F = 8192 v5e = TPU_V5E R_bf16 = v5e.critical_intensity("bf16") R_int8 = v5e.critical_intensity("int8") rows = [ ["练习 1：int8 × int8（int8 算力）", f"{R_int8:.0f}", "≈ 2B", critical_batch(D, F, R_int8, 1, 1, 1)], ["练习 2：int8 权重 × bf16 激活（bf16 算力）", f"{R_bf16:.0f}", "≈ 2B", critical_batch(D, F, R_bf16, 2, 1, 2)], ["对照：全 bf16", f"{R_bf16:.0f}", "≈ B", critical_batch(D, F, R_bf16, 2, 2, 2)], ] print_table(["TPU v5e 上的矩阵乘法", "临界强度", "强度", "临界 B"], rows) print("\n练习 4：'每个样本有自己的矩阵'（x[b]·W[b]，b = 0..B-1），D = F = 8192，bf16：") rows = [] for B in (1, 16, 256, 4096): f, b = batched_matvec_counts(B, D, F) rows.append([B, fmt_num(f), fmt_bytes(b), f"{f / b:.3f}"]) print_table(["B", "FLOPs", "字节", "强度"], rows) print("\ndecode 阶段的注意力就是练习 4 的形状：每个序列读自己的 K、V。按头的分组算强度（bf16，ctx = 4096）：") rows = [] H, dh, ctx = 32, 128, 4096 for G, name in ((32, "MHA（每个 query 头一组 KV）"), (8, "GQA（4 个 query 头共享一组）"), (1, "MQA（所有头共享一组）")): flops = 2 * 2 * H * dh * ctx # QKᵀ 和 PV 两次点积，每层每序列 nbytes = 2 * ctx * G * dh * 2 # 读 K 和 V（bf16） rows.append([name, G, fmt_num(flops), fmt_bytes(nbytes), f"{flops / nbytes:.1f}"]) print_table(["注意力结构", "KV 头数", "FLOPs / 序列 / 层", "读 KV 字节", "强度"], rows) print_observations([ "练习 1：int8 算力是 bf16 的 2 倍，但每个数只有 1 字节，强度 ≈ 2B，临界 B 和 bf16 一样（书中近似 240，精确值 256）。", "练习 2：只把**权重**压成 int8、计算仍用 bf16，读权重的字节减半而算力不变，临界 B 减半（书中近似 120，精确值 128）——" "这就是'仅权重量化'（weight-only quantization）在 decode 时那么有效的原因（第 3 讲）。", "练习 4：每个样本一个矩阵时，强度 ≈ 1（bf16），不管 B 多大都一样——批处理只能复用**共享**的数据。", "decode 的注意力正是这样：批里每个序列的 KV cache 都不同，强度 = query 头数 / KV 头数，和 B、ctx 都无关。" "MHA 只有 1，GQA 提到 4，MQA 提到 32——这是 GQA/MQA 被发明出来的系统动机之一（第 2 讲）。", ]) # --------------------------------------------------------------------------- # 演示 4：字符画 roofline # --------------------------------------------------------------------------- def demo4_plot(save_png=False): print_header("演示 4：H100（bf16）的 roofline 图") peak, bw = H100.peak_flops("bf16"), H100.mem_bw labels = "abcdefgh" pts, legend = [], [] for (name, f, b), lab in zip(basic_ops(), labels): pts.append((lab, f / b, None)) legend.append(f"{lab} = {name}") print("\n".join(ascii_roofline(pts, peak, bw, xmin=0.1, xmax=1e4))) print(" 图例：" + "；".join(legend)) if save_png: os.makedirs(os.path.join(HERE, "figures"), exist_ok=True) path = os.path.join(HERE, "figures", "02_roofline.png") curves = [(c.name + " bf16", c.peak_flops("bf16"), c.mem_bw) for c in (H100, H200, B200, TPU_V5E)] points = [(name, f / b, attainable_flops(f / b, peak, bw)) for name, f, b in basic_ops()] ok = save_roofline_png(path, curves, points, title="Roofline (points placed on H100 bf16)") print(f" PNG：{'已保存到 ' + path if ok else '没装 matplotlib，跳过'}") print_observations([ "左边的斜线是'带宽屋顶'：强度每翻一倍，可达算力翻一倍；右边的水平线是'算力屋顶'。两者在屋脊点（+）相交。", "a、b（逐元素和点积）在最左下角：再怎么优化代码也只能用到峰值的零点几个百分点——只能靠'少搬数据'（融合算子、" "第 11 讲）来提速。", "矩阵乘法随 B 增大沿斜线往右上爬，越过屋脊点之后贴着平顶走——这时再增大 B 只会增加每一步的延迟，不再提高吞吐。", ]) # --------------------------------------------------------------------------- # 演示 5：分块与片上存储 # --------------------------------------------------------------------------- def demo5_tiling(): print_header("演示 5：分块（tiling）——片上能放多大的块，决定了从 HBM 看到的强度") m = k = n = 4096 print(f"先验证公式：({m}×{k})·({k}×{n}) 的 bf16 矩阵乘法，不同块大小下从主存搬运的总字节数") rows = [] for bm, bk, bn in ((32, 32, 32), (64, 64, 64), (128, 64, 128), (128, 64, 256), (256, 64, 256)): f, b = tiled_matmul_bytes(m, k, n, bm, bk, bn) approx = bm * bn / (bm + bn) rows.append([f"{bm}×{bk}×{bn}", fmt_bytes(b), f"{f / b:.1f}", f"{approx:.1f}"]) print_table(["块 bm×bk×bn", "主存流量", "强度（精确）", "近似 bm·bn/(bm+bn)"], rows) print("\n块的大小受什么限制？输出块的累加器（fp32）要一直待在片上，直到 k 方向全部算完：") rows = [] for bm, bn in ((64, 64), (128, 128), (128, 256), (256, 256), (512, 512)): acc = bm * bn * 4 I = bm * bn / (bm + bn) fits = "放得下" if acc <= 256 * 1024 else "放不下" rows.append([f"{bm}×{bn}", fmt_bytes(acc, binary=True), fits, f"{I:.0f}", "是" if I >= H100.critical_intensity("bf16") else "否"]) print_table(["输出块", "fp32 累加器大小", "单个 SM 的 256 KiB 寄存器/TMEM", "强度", "≥ H100 屋脊 295？"], rows) print_observations([ "精确公式和近似 bm·bn/(bm+bn) 只差最后'写回输出'那一项：块越大，每个从主存搬进来的数被复用越多次。", "要让一个 SM 单独从 HBM 取数就达到 295 FLOP/byte，输出块得有约 590×590，fp32 累加器要 1.3 MiB——" "远超一个 SM 的 256 KiB 寄存器。所以单个 SM '自己搬自己算'到不了屋脊点。", "真实 GPU 靠两件事补上差距：(1) 50 MB（H100）/ 126 MB（B200）的 L2 cache 让多个 SM 共享搬进来的块，" "HBM 实际流量比'每个 SM 各搬各的'少得多；(2) Blackwell 给每个 SM 加了 256 KB 的 Tensor Memory（TMEM）" "专门放累加器，并支持两个 SM 合作做一次 MMA、共享操作数（Blackwell Ultra 博客）——都是为了做更大的块。", "这是'硬件为了 roofline 而设计'的一个具体例子；第 8 讲（硬件引擎）和第 10–11 讲（写内核）会再回到这里。", ]) # --------------------------------------------------------------------------- # 演示 6：网络 roofline # --------------------------------------------------------------------------- def demo6_network(): print_header("演示 6：网络 roofline——X[B,D]·W[D,F] 沿 D 切到两块芯片，交换部分和") F = 8192 chips = [(TPU_V5E, "bf16", 4.5e10, "单条 ICI 链路（书中例子）"), (TPU_V6E, "bf16", 9e10, "单条 ICI 链路"), (H100, "bf16", H100.link_bw, "NVLink 4 单向"), (B300, "bf16", B300.link_bw, "NVLink 5 单向")] rows = [] for chip, prec, link, note in chips: peak = chip.peak_flops(prec) R_net = peak / link D_crit = 2 * R_net # 强度 = D / 2（bf16），要求 D/2 > R_net verdicts = [] for D in (4096, 8192, 16384): r = sharded_matmul_two_chips(B=512, D=D, F=F, peak=peak, link_bw=link) verdicts.append("算" if r["bound"] == "compute" else "网") rows.append([chip.name, note, f"{R_net:,.0f}", f"D > {D_crit:,.0f}"] + verdicts) print_table(["芯片", "链路", "网络临界强度", "compute-bound 条件", "D=4096", "D=8192", "D=16384"], rows) print("\n换不同的 B（TPU v5e，D = 8192）：结论不变") rows = [] for B in (8, 128, 2048): r = sharded_matmul_two_chips(B=B, D=8192, F=F, peak=TPU_V5E.peak_flops("bf16"), link_bw=4.5e10) rows.append([B, fmt_time(r["t_math"]), fmt_time(r["t_comms"]), f"{r['intensity']:.0f}", "计算" if r["bound"] == "compute" else "通信"]) print_table(["B", "T_math", "T_comms", "网络强度", "瓶颈"], rows) print_observations([ "切分之后，每块芯片算一半，但要把 [B, F] 的部分和发给对方。算的量和发的量都正比于 B，B 在比值里约掉——" "网络强度 = D/2，只和'被切开的那一维'有多大有关。", "TPU v5e 的单条链路要求 D 大于约 8,755（书中结论）；NVLink 带宽大得多，H100 只要 D > 4,396。" "TPU v6e 的算力涨了 4.7 倍、单条链路只涨了 2 倍，要 D > 20,444——实际系统会同时用多条链路。", "B300 的 bf16 算力（推算值）是 H100 的 2.5 倍、NVLink 只涨了 2 倍，网络临界强度反而变高了：" "算力增长快于通信，是贯穿全课的趋势。", "第 4 讲（分布式推理）会系统地用这套方法决定'模型该怎么切到多块芯片上'。", ]) if __name__ == "__main__": demo1_basic_ops() demo2_critical() demo3_precision_and_per_example() demo4_plot(save_png="--png" in sys.argv) demo5_tiling() demo6_network()