""" ============================================================================= CS349E · 前置讲 P4 · 02_dataflow_and_fusion.py —— 数据流图、算子融合与内存规划 ============================================================================= 【这个脚本要回答的问题】 LLM 里除了矩阵乘法，还有大量"逐元素"运算：加偏置、激活函数（GELU、SiLU）、残差相加、缩放…… 每一个都很便宜（几次运算），但如果每个运算都单独成为一个内核，每次都要把整个张量从显存读出来、再写回去。 编译器最先做、收益最稳的优化就是**融合**（fusion）：把一串逐元素运算合成一个内核。本脚本： 演示 1：数据流图——一个 GELU 激活（含偏置和残差）的 IR：每个值被谁用、在哪一层、什么时候"死掉"。 演示 2：内存流量账——不融合 vs 融合，在 H100 上处理一个 LLM 尺寸的张量要多少时间、多少次内核启动。 演示 3：内存规划——不融合时需要多少个临时缓冲区？复用之后呢？ 演示 4：在你的 CPU 上实测：numpy 逐个运算 / 加上内存规划 / 融合成一个 C 循环（调用 C 库的 tanhf）/ 融合 + 可向量化的 tanh 逼近。融合一定更快吗？ 【怎么读输出】演示 1–3 是公式和计数（每次相同）；演示 4 是墙钟测量（只看趋势，编写时是 Apple M4 Pro，单线程）。 【运行】python3 02_dataflow_and_fusion.py （约 5 秒） python3 02_dataflow_and_fusion.py --quick （数组小 16 倍） 【依赖】numpy；从 ../lecture01 导入 console、hw_specs；同目录的 mini_compiler、cbuild。 """ import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import fmt_bytes, fmt_time, print_header, print_observations, print_table # noqa: E402 from hw_specs import H100 # noqa: E402 from cbuild import best_time, find_cc # noqa: E402 from mini_compiler import (compile_c_kernel, compile_numpy, compile_numpy_planned, eval_ast, # noqa: E402 graph_depth, ir_to_str, last_use, lower, optimize, parse, plan_memory) QUICK = "--quick" in sys.argv SOURCE = """h = x + b; # 加偏置 y = 0.5 * h * (1 + tanh(0.7978845608 * (h + 0.044715 * h * h * h))); # GELU（tanh 近似） out = y + r; # 残差 """ LAUNCH_S = 4e-6 # 每次内核启动的开销：沿用 P2 的假设值（4 µs） def build(): prog = parse(SOURCE) ins, outs, inputs = lower(prog, ["out"]) opt, oouts = optimize(ins, outs) return prog, opt, oouts, inputs def traffic_arrays(instrs, outs, inputs): """ 以"整个数组"为单位数内存流量（读 + 写）： 不融合：每条指令读它的每个数组参数、写一个结果数组； 融合：只读一次每个输入、写一次每个输出，中间结果都在寄存器里。 返回 (不融合的数组次数, 融合的数组次数)。 """ unfused = sum(sum(1 for a in i.args if not isinstance(a, float)) + 1 for i in instrs) fused = len(inputs) + len(outs) return unfused, fused # --------------------------------------------------------------------------- # 演示 1：数据流图 # --------------------------------------------------------------------------- def demo1_dataflow(opt, outs): """ 对每个值列出：由哪个运算产生、用到哪些值、在图里的"层"（最长依赖链上的位置）、被几个指令用到、最后一次被用到是第几条指令。 """ print_header("演示 1：数据流图——GELU（含偏置、残差）的 IR") print("源代码：") print(" " + SOURCE.strip().replace("\n", "\n ")) print("\n优化后的 IR：") print(ir_to_str(opt, outs)) level, users = {}, {} for ins in opt: lv = 0 for a in ins.args: if not isinstance(a, float): lv = max(lv, level.get(a, 0)) users[a] = users.get(a, 0) + 1 level[ins.dst] = lv + 1 last = last_use(opt, outs) rows = [] for idx, ins in enumerate(opt): lu = last.get(ins.dst) rows.append([idx, ins.dst, ins.op, level[ins.dst], users.get(ins.dst, 0), "输出" if lu == len(opt) else lu]) print() print_table(["#", "值", "运算", "层", "被几条指令用到", "最后一次被用到（指令 #）"], rows) print_observations([ f"{len(opt)} 条指令排成一条深度为 {graph_depth(opt)} 的链：几乎每一步都依赖上一步，图里几乎没有可以并行的分支。" "逐元素运算的并行性在'元素'这一维（每个元素独立），而不在图里。", "%h1 被 5 条指令用到（h 出现在公式的好几个地方），活得最久；大多数临时值只被用一次、马上就'死'了——" "演示 3 的内存规划就利用这一点。", ]) # --------------------------------------------------------------------------- # 演示 2：内存流量账 # --------------------------------------------------------------------------- def demo2_traffic(opt, outs, inputs, tokens=4096, d_ff=14336, elem_bytes=2): """ 张量形状取 Llama-3-8B 的 FFN 中间维度（d_ff = 14336）× 4096 个 token，BF16（2 字节）。 时间 = 内存流量 / H100 的 HBM 带宽（3.35 TB/s），即假设逐元素运算完全受限于带宽； 再加上每个内核一次启动开销（不融合 = 每条指令一个内核）。 """ print_header(f"演示 2：内存流量账——{tokens} 个 token × d_ff = {d_ff}（BF16），在 H100 上") n = tokens * d_ff arr = n * elem_bytes unf, fus = traffic_arrays(opt, outs, inputs) rows = [] for label, k, kernels in (("不融合（每条指令一个内核）", unf, len(opt)), ("融合（一个内核）", fus, 1)): nbytes = k * arr t_mem = nbytes / H100.mem_bw t_launch = kernels * LAUNCH_S rows.append([label, k, fmt_bytes(nbytes), fmt_time(t_mem), kernels, fmt_time(t_launch), fmt_time(t_mem + t_launch)]) print(f"每个张量 {n:,} 个元素 = {fmt_bytes(arr)}；H100 带宽 {H100.mem_bw / 1e12:.2f} TB/s；每次启动 {LAUNCH_S * 1e6:.0f} µs（假设）。") print_table(["方式", "读写了几个整张量", "内存流量", "带宽时间", "内核数", "启动开销", "合计"], rows) flops_per_elem = sum(1 for _ in opt) print(f"\n算术强度：每个元素约 {flops_per_elem} 次运算；融合后每个元素读写 {fus * elem_bytes} 字节 → " f"{flops_per_elem / (fus * elem_bytes):.2f} FLOP/B，远低于 H100 的屋脊点 295 ——融合后仍然受限于带宽。") print_observations([ f"不融合时，这 {len(opt)} 个运算一共读写了 {unf} 个整张量；融合后只有 {fus} 个（读 x、b、r，写 out）：内存流量少了约 {unf / fus:.0f} 倍，" "时间也少了这么多——逐元素运算的时间几乎完全由字节数决定。", "融合还把内核启动次数从 11 次降到 1 次。在这个张量尺寸下启动开销不是主要矛盾，但在 decode（每步的张量很小）时，" "启动开销可能占大头（P2 §8）——融合对 decode 更重要。", "融合之后，算术强度也只有 1 FLOP/B 左右，仍然远低于屋脊点：逐元素运算永远是 memory-bound 的。" "所以最好的做法是把它们融合进**前后的矩阵乘法**里（例如在矩阵乘法写回结果之前顺手做完激活和残差，" "叫 epilogue fusion），让它们的读写完全消失（第 11–13 讲）。", ]) # --------------------------------------------------------------------------- # 演示 3：内存规划 # --------------------------------------------------------------------------- def demo3_memory_planning(opt, outs, tokens=4096, d_ff=14336, elem_bytes=2): """ 不融合时，每个临时值都是一个整张量。朴素的做法：每个临时值一个缓冲区； 内存规划：一个值死了以后，它的缓冲区给后面的值复用（mini_compiler.plan_memory，贪心算法）。 """ print_header("演示 3：内存规划——不融合时，临时缓冲区可以复用") n_buf, assign = plan_memory(opt, outs) arr = tokens * d_ff * elem_bytes rows = [[ins.dst, ins.op, assign[ins.dst]] for ins in opt] print_table(["值", "运算", "分到的缓冲区"], rows) print(f"\n朴素：{len(opt)} 个临时缓冲区 = {fmt_bytes(len(opt) * arr)}；" f"规划后：{n_buf} 个 = {fmt_bytes(n_buf * arr)}（每个 {fmt_bytes(arr)}）。") print_observations([ f"贪心复用把 {len(opt)} 个临时缓冲区压到 {n_buf} 个：任何时刻同时'活着'的临时值最多只有 {n_buf} 个。", "逐元素运算可以原地计算（结果覆盖一个刚用完的输入），所以链上的大部分值都写进同一个缓冲区。", "PyTorch 的 eager 模式靠缓存分配器（caching allocator）在运行时复用显存；编译器（XLA、torch.compile）" "则在编译时就做这种规划。融合比规划更彻底：中间结果根本不进显存（第 12 讲）。", ]) # --------------------------------------------------------------------------- # 演示 4：在 CPU 上实测 # --------------------------------------------------------------------------- def demo4_measure(prog, opt, outs, inputs, seed=0): """ 同一个 IR 的四种实现，在 n 个 float32 上实测（单线程）： 1. numpy 逐个运算（生成的 numpy 代码，每条指令一次数组运算，每次分配新的结果数组）； 2. numpy 逐个运算 + 内存规划（结果写进预先分配好的缓冲区，out=）； 3. 融合成一个 C 循环（-O2），tanh 调用 C 库的 tanhf； 4. 融合成一个 C 循环（-O2），tanh 换成可向量化的有理函数逼近（mini_compiler 的 fast_tanh）。 "等效带宽"按各自的内存流量模型算：1、2 用不融合的流量，3、4 用融合的流量。 """ n = 1 << (20 if QUICK else 24) print_header(f"演示 4：实测——{n:,} 个 float32（每个数组 {fmt_bytes(4 * n)}），单线程") rng = np.random.default_rng(seed) env = {k: rng.standard_normal(n).astype(np.float32) for k in inputs} ref = eval_ast(prog, {k: v.astype(np.float64) for k, v in env.items()})["out"] unf, fus = traffic_arrays(opt, outs, inputs) impls = [("numpy 逐个运算", compile_numpy(opt, outs, inputs), unf)] f_plan, _ = compile_numpy_planned(opt, outs, inputs, n) impls.append(("numpy 逐个运算 + 内存规划", f_plan, unf)) if find_cc() is not None: impls.append(("融合 C 循环（C 库 tanhf）", compile_c_kernel(opt, outs, inputs), fus)) impls.append(("融合 C 循环（可向量化的 tanh 逼近）", compile_c_kernel(opt, outs, inputs, fast_tanh=True), fus)) else: print("（没有找到 C 编译器，只测 numpy 的两种实现）") rows, times = [], {} for label, f, k in impls: t = best_time(lambda: f(**env), repeats=3) err = np.abs(f(**env)["out"] - ref).max() times[label] = t rows.append([label, f"{t * 1e3:.1f} ms", f"{k * 4 * n / t / 1e9:.0f} GB/s", f"{err:.1e}"]) print_table(["实现", "时间", "等效带宽（按流量模型）", "和 FP64 参照的最大差"], rows) base = times["numpy 逐个运算"] lines = [f"内存规划让 numpy 版本从 {base * 1e3:.1f} ms 变成 {times['numpy 逐个运算 + 内存规划'] * 1e3:.1f} ms：" "运算和流量都没变，省下的是反复申请新内存（以及操作系统给新内存清零、建立页表）的开销。"] if len(times) == 4: t_libm = times["融合 C 循环（C 库 tanhf）"] t_fast = times["融合 C 循环（可向量化的 tanh 逼近）"] if t_libm > base: lines.append( f"出乎意料：融合成一个 C 循环、但 tanh 调用 C 库的 tanhf 时，用了 {t_libm * 1e3:.1f} ms，是 numpy 不融合版本的 " f"{t_libm / base:.1f} 倍。内存流量确实少了 {unf / fus:.0f} 倍，但每个元素一次**标量**的 tanhf 调用（几纳秒）" "让内核变成了受限于计算——瓶颈换了地方。numpy 的 tanh 内部是向量化实现，所以不融合反而更快。") else: lines.append( f"融合成一个 C 循环（调用 C 库的 tanhf）用了 {t_libm * 1e3:.1f} ms，比 numpy 不融合快 {base / t_libm:.1f} 倍——" "你的 C 库或编译器可能提供了向量化的 tanhf。编写时的机器（macOS）上这一行反而比 numpy 慢 4 倍：" "标量的 tanhf 调用让融合后的内核受限于计算。") lines += [ f"把 tanh 换成只用乘、加、除的有理逼近后，编译器能把整个循环向量化：{t_fast * 1e3:.1f} ms，比 numpy 不融合快 " f"{base / t_fast:.1f} 倍，误差仍在 float32 的舍入量级。", "教训：融合解决的是'搬'的问题，但融合后的内核里每个运算都必须高效——ML 编译器会为 exp、tanh、erf 等生成" "可向量化的多项式逼近（这也是第 1 讲 §5.5 提到的 FlashAttention-4 用多项式代替部分 exp 的道理）。", ] lines.append("这台 CPU 单核的内存带宽大约 100–200 GB/s（P1 演示），表里的等效带宽说明各版本离'只受带宽限制'有多远。") print_observations(lines) if __name__ == "__main__": program, ir, outputs, inp = build() demo1_dataflow(ir, outputs) demo2_traffic(ir, outputs, inp) demo3_memory_planning(ir, outputs) demo4_measure(program, ir, outputs, inp)