""" ============================================================================= CS349E · 第 12 讲 · 04_memory_planning.py —— 内存规划：中间结果的生存期与显存复用 ============================================================================= 【这个脚本要回答的问题】 有了图和内核的顺序，编译器就知道每个中间结果什么时候产生、什么时候最后一次被用到。 生存期不重叠的结果可以共用同一块显存——提前规划好，运行时就不必再分配。 演示 1：一个小层（融合后）的中间结果和它们的生存期；贪心分配的结果。 演示 2：Llama-3.1-8B 的一层，prefill 的激活显存：不复用 / 贪心 / 下界；融合和 FlashAttention 的影响；随 T 的变化。 演示 3：运算的顺序也影响峰值：同一个计算，"先把所有分支算完"和"一个分支一个分支地算完"。 演示 4：一层的显存账：权重、KV cache、激活。 【运行】python3 04_memory_planning.py （约 1 秒） 【依赖】numpy；本目录的 graph_lab；../lecture01 的 console。 """ import os import sys HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from graph_lab import (F, LLAMA8B, LayerCfg, constant_fold, cse, decompose, fold_norm_weight, fuse, layer_params, # noqa: E402 lifetimes, lower_bound, match_attention, plan_memory, prefill_layer, simplify, trace) from console import fmt_bytes, print_header, print_observations, print_table # noqa: E402 def optimize(g, attention=True): if attention: g, _ = match_attention(g) g = decompose(g) g = cse(g) g, _ = constant_fold(g) g, _ = simplify(g) g, _ = fold_norm_weight(g) return g def param_shapes(cfg): return {"attn_norm": (cfg.d,), "wq": (cfg.d, cfg.H * cfg.dh), "wk": (cfg.d, cfg.G * cfg.dh), "wv": (cfg.d, cfg.G * cfg.dh), "wo": (cfg.H * cfg.dh, cfg.d), "mlp_norm": (cfg.d,), "w_gate": (cfg.d, cfg.f), "w_up": (cfg.d, cfg.f), "w_down": (cfg.f, cfg.d)} def demo1_small(): print_header("演示 1：一个小层（d = 128，T = 16，融合 + FlashAttention）的中间结果") cfg = LayerCfg(128, 8, 2, 16, 344) g = optimize(trace(prefill_layer(cfg, 16), {"x": (16, cfg.d)}, layer_params(cfg))) ks = fuse(g) lts = lifetimes(g, ks) peak, off = plan_memory(lts) rows = [[f"%{nid}", g.nodes[nid].op, "×".join(map(str, g.nodes[nid].shape)), fmt_bytes(sz), s, e if e < len(ks) else "输出", fmt_bytes(off[nid])] for nid, sz, s, e in sorted(lts, key=lambda t: t[2])] print_table(["结果", "产生它的运算", "形状", "大小", "产生于内核", "最后使用于内核", "贪心分配的偏移"], rows) print(f"\n不复用：{fmt_bytes(plan_memory(lts, 'naive')[0])}；贪心：{fmt_bytes(peak)}；下界：{fmt_bytes(lower_bound(lts))}") print_observations([ "只有\"写回显存\"的结果才需要分配（融合在内核内部的中间值在寄存器 / 共享内存里）。每个结果有一个生存期：从产生它的内核，到最后一个读它的内核。", "贪心（greedy-by-size）：按大小从大到小，每个结果放在\"和已放好的、生存期重叠的块都不冲突\"的最低偏移。TFLite、XLA 的缓冲区分配都是这一类。", "下界 = 任何时刻同时活着的字节之和的最大值。这张图上贪心正好达到下界：生存期是区间、彼此嵌套得比较规整。一般情况下（区间着色的变体）贪心不保证最优，但很接近。", ]) def demo2_llama(): cfg = LLAMA8B print_header("演示 2：Llama-3.1-8B 一层的 prefill 激活显存（BF16；不含权重和输入 x）") rows = [] for T in (512, 2048, 8192): g = trace(prefill_layer(cfg, T), {"x": (T, cfg.d)}, param_shapes(cfg)) for name, gg, en in (("不融合", decompose(g), False), ("融合", optimize(g, False), True), ("融合 + FlashAttention", optimize(g, True), True)): ks = fuse(gg, enable=en) lts = lifetimes(gg, ks) rows.append([T, name, len(lts), fmt_bytes(plan_memory(lts, "naive")[0]), fmt_bytes(plan_memory(lts)[0]), fmt_bytes(lower_bound(lts))]) print_table(["T", "分区", "要分配的中间结果", "不复用", "贪心", "下界"], rows) print_observations([ "复用很重要：不融合时 T = 2048 的中间结果加起来 2.3 GB，贪心复用后 541 MB——生存期不重叠的块共用显存。", "融合减少了要分配的结果的个数（45 → 13），\"不复用\"那一列大幅下降；但 T ≥ 2048 时峰值几乎没变：峰值由同时活着的最大的那几个决定——" "注意力的分数矩阵（32 × T × T）和 softmax 的输出。", "FlashAttention 让分数矩阵根本不出现在显存里：T = 8192 时峰值从 8.6 GB 降到 604 MB。没有它，长序列的 prefill 光一层的激活就放不下" "（这也是第 5 讲 chunked prefill 把长提示切成小块算的原因之一）。", "有了 FlashAttention，峰值随 T 线性增长（MLP 的中间结果 T × 14336）；没有它，随 T² 增长。", ]) def branches(depth_first, k=4): """k 个分支：每个分支 y_i = tanh(x @ W_i)（很大），再按行求和（很小），最后相加。depth_first 决定代码里的执行顺序。""" def f(x, p): if depth_first: acc = None for i in range(k): s = F.tanh(x @ p[f"W{i}"]).sum(axis=-1, keepdims=True) acc = s if acc is None else acc + s return acc ys = [F.tanh(x @ p[f"W{i}"]) for i in range(k)] ss = [y.sum(axis=-1, keepdims=True) for y in ys] out = ss[0] for s in ss[1:]: out = out + s return out return f def demo3_order(): print_header("演示 3：运算顺序影响峰值（4 个分支，每个 y_i = tanh(x @ W_i)：[2048, 16384]，再按行求和）") T, d, k = 2048, 4096, 4 ps = {f"W{i}": (d, 4 * d) for i in range(k)} rows = [] for name, dfs in (("先把 4 个分支的大矩阵都算出来，再求和（广度优先）", False), ("一个分支算完、求和，再算下一个（深度优先）", True)): g = trace(branches(dfs, k), {"x": (T, d)}, ps) for fz in (False, True): gg = decompose(g) ks = fuse(gg, enable=fz) lts = lifetimes(gg, ks) rows.append([name, "融合" if fz else "不融合", len(ks), fmt_bytes(plan_memory(lts)[0]), fmt_bytes(lower_bound(lts))]) print_table(["代码里的顺序", "分区", "内核数", "贪心峰值", "下界"], rows) print_observations([ "同一个计算、同样的运算，只是先后不同：广度优先时 4 个大矩阵（每个 67 MB）同时活着；深度优先时每个大矩阵求完和就可以释放，峰值只有约四分之一。", "追踪把 Python 的执行顺序原样记成了图的顺序。编译器可以在不违反依赖的前提下重排（调度）运算来降低峰值——TorchInductor、XLA 都有按内存排序的调度；" "最优的排序是 NP 难的，实际用启发式。", "融合让 tanh 和矩阵乘法合进同一个内核（epilogue），少写一份大矩阵，但大矩阵本身（求和要读它）仍要写出来：融合不能代替好的顺序。", ]) def demo4_budget(): cfg = LLAMA8B print_header("演示 4：一层的显存账（Llama-3.1-8B，BF16，批 1）") w = 2 * (cfg.d * cfg.H * cfg.dh + 2 * cfg.d * cfg.G * cfg.dh + cfg.H * cfg.dh * cfg.d + 3 * cfg.d * cfg.f) rows = [] for T in (2048, 8192, 32768): kv = 2 * T * cfg.G * cfg.dh * 2 g = trace(prefill_layer(cfg, T), {"x": (T, cfg.d)}, param_shapes(cfg)) res = {} for name, att in (("无 FA", False), ("FA", True)): gg = optimize(g, att) lts = lifetimes(gg, fuse(gg)) res[name] = plan_memory(lts)[0] rows.append([T, fmt_bytes(w), fmt_bytes(kv), fmt_bytes(res["无 FA"]), fmt_bytes(res["FA"])]) print_table(["T", "权重（每层）", "KV cache（每层）", "prefill 激活峰值（融合，无 FlashAttention）", "（融合 + FlashAttention）"], rows) print_observations([ "权重每层 436 MB，和 T 无关；KV cache 随 T 线性增长，32K 时每层 134 MB、32 层 4.3 GB——decode 时它们一直占着显存（第 7 讲）。", "prefill 的激活是临时的：一层算完就能释放（下一层复用同一块显存，规划时只要按\"一层的峰值\"留出来）。有 FlashAttention 时它和 T 成正比，" "T = 32K 时约 2.4 GB；没有时随 T² 增长到 137 GB，一块 80 GB 的卡根本放不下。", "推理引擎在启动时就按\"最大的 T × 最大的批\"把激活和 KV 的显存预留好（vLLM 的 gpu_memory_utilization 先跑一次最大规模的前向来测峰值）；" "静态的内存规划也是 CUDA Graphs 的前提：图里的每个内核的地址在捕获时就固定了（05 脚本）。", ]) if __name__ == "__main__": demo1_small() demo2_llama() demo3_order() demo4_budget()