""" ============================================================================= CS349E · 第 1 讲 · 04_llm_inference_on_rooflines.py —— 把 roofline 用到 LLM 推理上 ============================================================================= 【这个脚本要回答的问题】 02 脚本说明了"矩阵乘法要几百个 token 一起算才 compute-bound"。LLM 推理里情况如何？ 本脚本用 llm_math（公式化的推理账本）和 hw_specs（芯片规格）回答： 演示 1：模型账本——三个 Llama 3.1 模型的参数、权重字节、每 token 的 KV cache 和运算量。 演示 2：一步 decode 的时间拆解——读权重、读 KV、计算，各占多少？批大小 B 怎样改变它们？ 演示 3：上下文越长，decode 越难 compute-bound——每个序列的 KV 读取无法被批处理摊薄。 演示 4：显存容量是批大小的天花板——H100（80 GB）vs Blackwell Ultra（288 GB）。 演示 5：延迟–吞吐–能耗的前沿：每个用户的速度 vs 每块芯片的产出 vs 每个 token 的能耗。 演示 6：prefill 是 compute-bound 的：首 token 延迟（TTFT）的下界。 演示 7：近存 vs 远存——SRAM 芯片把瓶颈从"带宽"换成了"容量"和"互连"。 【所有数字都是"理想下界"】峰值算力和带宽按 100% 计、多芯片按"理想切分"计（不算通信，第 4 讲再加）， 不算激活值读写和每步的固定开销。所以真实系统会更慢，但**谁是瓶颈、怎样随 B 和 ctx 变化**的结论是可靠的。 每次运行完全相同。 【运行】python3 04_llm_inference_on_rooflines.py （不到 1 秒） 【依赖】只用 Python 标准库。 """ import math from console import print_header, print_observations, print_table, fmt_num, fmt_bytes, fmt_time from hw_specs import H100, B300, GROQ3_LPU, CORSAIR from llm_math import (LLAMA31_8B, LLAMA31_70B, LLAMA31_405B, MODELS, decode_step, prefill_time, max_batch, serving_metrics) # --------------------------------------------------------------------------- # 演示 1：模型账本 # --------------------------------------------------------------------------- def demo1_ledger(): print_header("演示 1：模型账本（Llama 3.1 系列）") rows = [] for m in MODELS.values(): rows.append([m.name, fmt_num(m.n_params(), 2), fmt_bytes(m.weight_bytes("bf16")), fmt_bytes(m.weight_bytes("fp8")), fmt_bytes(m.weight_bytes("nvfp4")), fmt_bytes(m.kv_bytes_per_token(), binary=True), fmt_num(m.flops_per_token(0)), fmt_num(m.flops_per_token(8192))]) print_table(["模型", "参数", "BF16", "FP8", "NVFP4", "KV/token (BF16)", "FLOPs/token (ctx=0)", "FLOPs/token (ctx=8K)"], rows) m = LLAMA31_8B print_observations([ f"每个 token 的运算量 ≈ 2 × 参数量：{m.name} 约 {fmt_num(m.flops_per_token(0))} FLOPs（比 2 × {fmt_num(m.n_params(), 2)} " "略小，因为词嵌入只是查表、不做乘法）。", f"上下文 8K 时，注意力打分额外增加 {fmt_num(m.attn_flops_per_token(8192))} FLOPs/token" f"（{m.attn_flops_per_token(8192) / m.linear_flops_per_token() * 100:.0f}%）：长上下文时注意力的运算不能忽略。", "KV cache 每 token 只有几百 KiB，看起来很小；但一个 8K 上下文的序列就要 1–4 GB，几十个序列就和权重一样大了（演示 4）。", "NVFP4 的权重约是 BF16 的 28%（4.5 bit / 16 bit）。权重越小，decode 每一步要搬的字节越少（演示 5）。", ]) # --------------------------------------------------------------------------- # 演示 2：一步 decode 的时间拆解 # --------------------------------------------------------------------------- def demo2_decode_breakdown(model=LLAMA31_8B, chip=H100, ctx=4096): print_header(f"演示 2：{model.name} 在 {chip.name} 上的一步 decode（BF16，每个序列上下文 {ctx}）") rows = [] for B in (1, 4, 16, 64, 128, 256, 512): s = decode_step(model, chip, B, ctx) t_w = s["weight_bytes"] / chip.mem_bw t_kv = s["kv_bytes"] / chip.mem_bw m = serving_metrics(s, B, chip) rows.append([B, fmt_time(t_w), fmt_time(t_kv), fmt_time(s["t_math"]), fmt_time(s["t_step"]), s["bound"], f"{s['intensity']:.1f}", f"{m['tok_s_user']:.0f}", f"{m['tok_s_chip']:,.0f}"]) print_table(["B", "读权重", "读 KV", "计算", "一步下界", "瓶颈", "强度", "tok/s/用户", "tok/s/芯片"], rows) s1 = decode_step(model, chip, 1, ctx) s16 = decode_step(model, chip, 16, ctx) s256 = decode_step(model, chip, 256, ctx) m1, m16 = serving_metrics(s1, 1, chip), serving_metrics(s16, 16, chip) print_observations([ f"B = 1 时一步至少 {fmt_time(s1['t_step'])}，几乎全花在读 {fmt_bytes(s1['weight_bytes'])} 的权重上；计算只要 " f"{fmt_time(s1['t_math'])}，H100 的算力用了不到 {s1['t_math'] / s1['t_step'] * 100:.1f}%。", f"'读权重'这一项不随 B 变化（所有序列共享同一份权重）：B 从 1 涨到 16，每块芯片的产出涨了 " f"{m16['tok_s_chip'] / m1['tok_s_chip']:.1f} 倍，每个用户的速度只降了 {(1 - m16['tok_s_user'] / m1['tok_s_user']) * 100:.0f}%" "——批处理是推理的第一杠杆。", f"但'读 KV'随 B 线性增长：每个序列都有自己的 KV cache。B = 256 时读 KV 是读权重的 " f"{s256['kv_bytes'] / s256['weight_bytes']:.1f} 倍，再加大 B，芯片产出几乎不涨，每个用户却越来越慢。", f"整张表里瓶颈始终是 memory：在 ctx = {ctx} 时，这一步永远到不了 compute-bound（演示 3 解释为什么）。", ]) # --------------------------------------------------------------------------- # 演示 3：上下文长度与"永远 memory-bound" # --------------------------------------------------------------------------- def demo3_context(model=LLAMA31_8B, chip=H100): """ 求出 t_math = t_mem 的批大小 B*（decode 刚好变成 compute-bound 的点）。 【推导】一步 decode：t_math = B·f / P，t_mem = (W + B·k) / β，其中 f = 每 token 运算量（随 ctx 变大），W = 权重字节，k = 每个序列的 KV 字节 = ctx × 每 token KV 字节， P = 峰值算力，β = 带宽。令二者相等：B·(f/P − k/β) = W/β。 若 f/P > k/β：B* = (W/β) / (f/P − k/β)；否则 B 再大也到不了 compute-bound（返回"到不了"）。 直觉：每多一个序列，计算多 f/P 秒，读 KV 多 k/β 秒；只有前者大于后者，增加 B 才能追上固定的读权重时间。 """ print_header(f"演示 3：上下文越长，decode 越难 compute-bound（{model.name}，{chip.name}，BF16）") P, beta = chip.peak_flops("bf16"), chip.mem_bw W = model.matmul_params() * 2 rows = [] for ctx in (0, 128, 512, 1024, 4096, 32768): f = model.flops_per_token(ctx) k = ctx * model.kv_bytes_per_token() per_seq_math, per_seq_kv = f / P, k / beta if per_seq_math > per_seq_kv: Bstar = math.ceil((W / beta) / (per_seq_math - per_seq_kv)) Bs = f"{Bstar:,}" else: Bs = "到不了" limit = f / k if k > 0 else math.inf rows.append([ctx, fmt_time(per_seq_math), fmt_time(per_seq_kv), Bs, "∞" if math.isinf(limit) else f"{limit:.0f}"]) print_table(["ctx", "每序列计算时间", "每序列读 KV 时间", "compute-bound 需要的 B*", "B→∞ 时强度上限"], rows) print_observations([ f"ctx = 0（只看权重）时，B* 约等于 H100 的屋脊点 {chip.critical_intensity('bf16'):.0f}——和 02 脚本的矩阵乘法结论一致。", "上下文一长，每个序列每一步要读的 KV 就超过它带来的计算：B→∞ 时强度上限 = 每 token 运算量 / 每序列 KV 字节，" "随 ctx 下降，低于屋脊点之后批再大也到不了 compute-bound。", f"对 {model.name}（GQA，8 组 KV），这个分界点在 ctx 几百左右。对话、推理、agent 的上下文通常是几千到几十万——" "所以长上下文 decode 本质上是 KV cache 的'带宽 + 容量'问题，不是算力问题。", "这正是 GQA/MQA/MLA、KV 量化、稀疏注意力、prefill/decode 解耦这些技术的共同动机（第 2、3、5、7 讲）。", ]) # --------------------------------------------------------------------------- # 演示 4：显存容量 # --------------------------------------------------------------------------- def demo4_capacity(): print_header("演示 4：显存容量是批大小的天花板（预留 10% 给激活和工作区）") cases = [ (LLAMA31_8B, H100, "bf16", 1), (LLAMA31_8B, B300, "bf16", 1), (LLAMA31_70B, H100, "bf16", 1), (LLAMA31_70B, H100, "fp8", 1), (LLAMA31_70B, H100, "bf16", 2), (LLAMA31_70B, B300, "fp8", 1), (LLAMA31_70B, B300, "nvfp4", 1), (LLAMA31_405B, B300, "fp8", 4), ] rows = [] for m, chip, prec, n in cases: cells = [m.name, f"{n}×{chip.name}", prec, fmt_bytes(m.weight_bytes(prec)), fmt_bytes(n * chip.mem_cap)] for ctx in (1024, 8192, 32768, 131072): cells.append(f"{max_batch(m, chip, ctx, prec, n_chips=n, reserve_frac=0.10):,}") rows.append(cells) print_table(["模型", "芯片", "权重精度", "权重", "总显存", "B @1K", "B @8K", "B @32K", "B @128K"], rows) print_observations([ "70B 模型的 BF16 权重（141 GB）放不进一块 80 GB 的 H100：要么量化到 FP8，要么切到 2 块以上的芯片（第 4 讲）。", "上下文 128K 时，即使是 8B 模型、一块 H100 也只能放下个位数的序列——远低于演示 3 需要的几百。" "容量直接把批大小压在了'低效区'。", "Blackwell Ultra 把显存加到 288 GB（H100 的 3.6 倍），同样的模型能放下多得多的序列；再配合 NVFP4 把权重压小，" "腾出的空间全部给 KV cache。'更大的显存'是推理芯片最重要的卖点之一，不亚于算力。", "KV cache 用 FP8 存储（本表没算）还能再翻一倍；第 3 讲（量化）和第 7 讲（KV cache 管理）会继续展开。", ]) # --------------------------------------------------------------------------- # 演示 5：延迟–吞吐–能耗前沿 # --------------------------------------------------------------------------- def demo5_frontier(ctx=4096): """ 扫描批大小 B，得到每种部署的 (每用户 tok/s, 每芯片 tok/s, J/token, tok/s/MW) 曲线。 【为什么叫"前沿"（Pareto frontier）】每个 B 是一个工作点：B 小 → 每个用户快，但芯片产出低； B 大 → 芯片产出高，但每个用户慢。没有哪个点在两方面都最好，所有"最优折中"连起来就是前沿。 NVIDIA 的 Blackwell Ultra 博客图 6 就是这种图（横轴 TPS/user，纵轴 TPS/MW）。 【简化】8 块芯片按理想张量并行（算力、带宽 ×8，**不计通信**），能耗按最大功耗（偏保守）。 "芯片·时 / 百万 token"乘以你的每卡每小时价格，就是每百万 token 的芯片成本。 """ print_header(f"演示 5：延迟–吞吐–能耗前沿（Llama-3.1-70B，8 块芯片理想切分，ctx = {ctx}）") deployments = [ (H100, "fp8", "8×H100 FP8"), (B300, "fp8", "8×Blackwell Ultra FP8"), (B300, "nvfp4", "8×Blackwell Ultra NVFP4"), ] m, n = LLAMA31_70B, 8 j_b1, j_b256 = [], [] for chip, prec, label in deployments: cap = max_batch(m, chip, ctx, prec, n_chips=n, reserve_frac=0.10) rows = [] for B in (1, 8, 32, 64, 128, 256, 512, 1024): if B > cap: rows.append([B, "（显存放不下）", "", "", "", "", ""]) continue s = decode_step(m, chip, B, ctx, w_prec=prec, n_chips=n) met = serving_metrics(s, B, chip, n_chips=n) chip_hours = n / 3600.0 * s["t_step"] / B * 1e6 rows.append([B, f"{met['tok_s_user']:.0f}", f"{met['tok_s_chip']:,.0f}", fmt_time(s["t_step"]), f"{met['j_per_tok']:.3f}", fmt_num(met["tok_s_mw"]), f"{chip_hours:.3f}"]) if B == 1: j_b1.append(met["j_per_tok"]) if B == 256: j_b256.append(met["j_per_tok"]) print(f"\n{label}（显存最多放 {cap:,} 个序列）") print_table(["B", "tok/s/用户", "tok/s/芯片", "TPOT 下界", "J/token", "tok/s/MW", "芯片·时/百万 token"], rows) print_observations([ "每张表从上到下就是沿前沿移动：用户速度（tok/s/用户）一路下降，芯片产出、每兆瓦产出一路上升，每 token 能耗一路下降。" "选哪个点取决于产品要什么：聊天要 TPOT ≤ 几十 ms，离线批处理只看成本。", f"B = 1 时每个 token 要花 {min(j_b1):.0f}–{max(j_b1):.0f} 焦耳（8 块芯片满功率等一个 token），B = 256 时降到 " f"{min(j_b256):.2f}–{max(j_b256):.2f} 焦耳——批处理同时是'省钱'和'省电'的杠杆。再往后降得很慢：KV 读取不能被摊薄。", "同样 8 块芯片，Blackwell Ultra 在每个 B 上都更快（带宽 8 vs 3.35 TB/s），显存还多放了几倍的序列；" "NVFP4 让权重更小，在小 B（权重主导）时收益最大，大 B（KV 主导）时收益变小——因为本表的 KV cache 仍是 BF16。", "注意这些是不计通信的理想下界：8 卡张量并行每层要做两次 all-reduce，小 B 时通信延迟占比很高（第 4 讲）。", ]) # --------------------------------------------------------------------------- # 演示 6：prefill 与 TTFT # --------------------------------------------------------------------------- def demo6_prefill(model=LLAMA31_8B, chip=H100): print_header(f"演示 6：prefill 的下界（{model.name}，1×{chip.name}，BF16）") rows = [] longest = None for n in (128, 1024, 8192, 32768, 131072): s = prefill_time(model, chip, n) longest = (n, s) attn_share = model.attn_flops_per_token(1) * n * (n + 1) / 2 / s["flops"] rows.append([n, fmt_num(s["flops"]), fmt_time(s["t_math"]), fmt_time(s["t_mem"]), s["bound"], f"{s['intensity']:,.0f}", f"{attn_share * 100:.0f}%", fmt_time(s["t_step"])]) print_table(["提示长度", "FLOPs", "计算", "读写", "瓶颈", "强度", "注意力占比", "TTFT 下界"], rows) print_observations([ "prefill 一次处理整段提示，n 个 token 共享同一次权重读取，强度 ≈ n——几百个 token 就越过了屋脊点，是 compute-bound 的。" "和 decode 正好相反：同一个模型的两个阶段，'性格'完全不同。", f"提示很长时，注意力的运算按 n² 增长：{longest[0]:,} 个 token 的提示里注意力占了大部分运算，" f"TTFT 下界从毫秒级涨到 {fmt_time(longest[1]['t_step'])}（而且这还是 100% 峰值的理想值）。", "prefill 吃算力、decode 吃带宽和容量——把它们放到不同的机器（甚至不同种类的芯片）上，是第 5 讲'解耦式推理'的出发点。", ]) # --------------------------------------------------------------------------- # 演示 7：近存 vs 远存 # --------------------------------------------------------------------------- def demo7_near_vs_far(model=LLAMA31_70B, prec="fp8", hop_latency_s=1e-6): """ 比较"HBM 芯片"和"SRAM 芯片"服务同一个模型时：要几块芯片才放得下权重？batch = 1 的一步至少多久？ 【模型】 芯片数 = ⌈权重字节 / 单芯片容量⌉（只放权重，不留 KV 空间——SRAM 芯片的现实是 KV 往往还得放在别处） 读权重时间下界 = 权重字节 / (芯片数 × 单芯片带宽) 通信时间下界（**假设**）：每一层的激活要从持有上一层的芯片传到下一层，每次至少一个"跳"的延迟。 取 hop_latency_s = 1 µs（How to Scale Your Model 第 2 章给 TPU 的每跳延迟估计）。 层数 L 时每个 token 至少 L 次跳。这是为了说明"延迟"会成为新瓶颈而做的粗略假设。 【边界】芯片数 = 1 时不计通信。 """ print_header(f"演示 7：近存（SRAM）vs 远存（HBM）——batch = 1 服务 {model.name}（{prec.upper()} 权重）") w = model.weight_bytes(prec) rows = [] for chip in (H100, B300, GROQ3_LPU, CORSAIR): n = math.ceil(w / chip.mem_cap) t_w = w / (n * chip.mem_bw) t_hop = 0.0 if n == 1 else model.n_layers * hop_latency_s t = max(t_w, t_hop) if n > 1 else t_w rows.append([chip.name, chip.mem_kind, fmt_bytes(chip.mem_cap), f"{chip.memory_sweeps_per_second():,.0f}", n, fmt_time(t_w), fmt_time(t_hop) if n > 1 else "-", fmt_time(t), f"{1 / t:,.0f}"]) print_table(["芯片", "主存", "单芯片容量", "主存读遍/秒", "放下权重要几块", "读权重下界", f"通信下界（{model.n_layers} 层×1 µs）", "一步下界", "tok/s/用户上限"], rows) print_observations([ "'主存读遍/秒' = 带宽 / 容量：HBM 芯片每秒只能把自己的显存读几十遍，SRAM 芯片能读几十万遍。" "一个塞满显存的模型，batch = 1 时每个用户每秒最多生成这么多个 token。", "SRAM 又快又近，但小：一块 Groq 3 LPU 只有 500 MB，放一个 70B 的 FP8 模型就要一百多块芯片。" "瓶颈从'带宽'变成了'容量'（要多少芯片、多少钱）和'互连'（跨这么多芯片的延迟）。", "按每层 1 µs 的跨芯片延迟粗算，SRAM 方案的一步下界由通信决定，而不是由读权重决定——这和 Ma & Patterson（2026）" "'LLM 推理的瓶颈是内存和互连延迟，而不是算力'的判断一致。", "所以没有哪种芯片在所有阶段都最好：prefill 适合算力强的芯片，decode 适合带宽高的芯片，" "小的草稿模型（投机解码）适合 SRAM 芯片——这就是'异构、解耦'推理（第 5 讲，Gimlet Labs 的做法）的逻辑。", ]) if __name__ == "__main__": demo1_ledger() demo2_decode_breakdown() demo3_context() demo4_capacity() demo5_frontier() demo6_prefill() demo7_near_vs_far()