""" ============================================================================= CS349E · 第 3 讲 · 03_precision_on_the_roofline.py —— 量化省的是什么：字节、FLOPs，还是容量？ ============================================================================= 【这个脚本要回答的问题】 01、02 脚本问的是"量化之后输出偏多少"。这里问另一半："量化之后能快多少、能多服务几个人"——用 roofline 回答。 几种常见方案改变的东西不一样： W8A16 / W4A16（仅权重量化）：权重字节变少；计算仍用 BF16（权重读进来先反量化）。 W8A8（INT8 或 FP8）：权重和激活都是 8 位，用 8 位的 Tensor Core（算力是 BF16 的 2 倍）。 W4A4（NVFP4，Blackwell 才有）：权重和激活都是 4 位（块缩放），算力是 FP8 的 2–3 倍。 KV cache 量化：KV 字节变少（长上下文、大批量时它才是大头）。 演示 1：decode：Llama-3.1-8B、H100，不同批大小下每种方案的一步时间。 演示 2：prefill：同一个模型，2048 个 token 的提示——仅权重量化为什么没用。 演示 3：临界批大小：仅权重量化让 decode 更早变成 compute-bound（第 1 讲练习 3 的推广）。 演示 4：容量：Llama-3.1-70B 放在一块 H100 / Blackwell Ultra 上，32K 上下文时最多服务几个序列、每步多久。 【模型的假设】全部是 roofline 下界（峰值按 100% 算）；不计激活读写、反量化指令和内核启动开销； lm_head 和其他矩阵用同样的比特；INT4 按"每 128 个一组、16 位缩放因子"算 4.125 位，NVFP4 按 4.5 位。 所有输出每次运行完全相同。 【运行】python3 03_precision_on_the_roofline.py （< 1 秒） 【依赖】同目录的 quant_lab；../lecture01 的 console、hw_specs、llm_math。 """ import math import os import sys 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 B300, H100 # noqa: E402 from llm_math import LLAMA31_8B, LLAMA31_70B # noqa: E402 from quant_lab import step_time # noqa: E402 # (名字, 权重比特, 计算精度, KV 比特) SCHEMES_H100 = [ ("W16A16（BF16）", 16, "bf16", 16), ("W8A16（INT8 仅权重）", 8, "bf16", 16), ("W4A16（INT4 g128 仅权重）", 4.125, "bf16", 16), ("W8A8（INT8）", 8, "int8", 16), ("W8A8（FP8）", 8, "fp8", 16), ("W8A8（FP8）+ KV FP8", 8, "fp8", 8), ] SCHEMES_B300 = SCHEMES_H100[:1] + SCHEMES_H100[4:] + [ ("W4A4（NVFP4）+ KV FP8", 4.5, "nvfp4", 8), ] # --------------------------------------------------------------------------- # 演示 1：decode # --------------------------------------------------------------------------- def demo1_decode(): print_header("演示 1：decode 的一步——Llama-3.1-8B，H100，上下文 4096") m, chip, ctx = LLAMA31_8B, H100, 4096 batches = (1, 16, 64, 256) rows = [] base = {B: step_time(m, chip, B, 1, ctx, 16, 16, "bf16")["t_step"] for B in batches} for name, wb, cp, kb in SCHEMES_H100: cells = [name] for B in batches: st = step_time(m, chip, B, 1, ctx, wb, kb, cp) cells.append(f"{fmt_time(st['t_step'])}（{base[B] / st['t_step']:.2f}×）") rows.append(cells) print_table(["方案"] + [f"B = {B}" for B in batches], rows) st = step_time(m, chip, 256, 1, ctx, 4.125, 16, "bf16") print(f"\nW4A16、B = 256 时：读权重 {fmt_bytes(st['weight_bytes'])}，读 KV {fmt_bytes(st['kv_bytes'])}，" f"计算 {fmt_time(st['t_math'])}，搬运 {fmt_time(st['t_mem'])}") print_observations([ "B = 1 时一步几乎全是读权重：W8 快约 1.9 倍，W4 快约 3.5 倍——比特数减半，时间就减半。" "这时计算用 BF16 还是 INT8 / FP8 完全无所谓（计算只占 0.4%）：**小批量 decode 只看权重的字节**。", "B 变大以后，KV 的字节越来越多，权重量化的收益随之缩小：B = 256 时 W4A16 只快 1.1 倍——" "这时一步里 137 GB 是 KV、权重只有 4 GB。要继续加速，就得量化 KV（最后一行：KV 也换成 FP8，每个 B 都正好快 2 倍——权重和 KV 的字节都减半了）。", "INT8 和 FP8 的 W8A8 在 decode 里和 W8A16 一模一样：它们多出来的'算力翻倍'只在 compute-bound 时才用得上（演示 2）。", ]) # --------------------------------------------------------------------------- # 演示 2：prefill # --------------------------------------------------------------------------- def demo2_prefill(): print_header("演示 2：prefill——Llama-3.1-8B，一个 2048 token 的提示") m, n = LLAMA31_8B, 2048 rows = [] for chip, schemes in ((H100, SCHEMES_H100), (B300, SCHEMES_B300)): base = step_time(m, chip, 1, n, (n + 1) / 2, 16, 16, "bf16")["t_step"] for name, wb, cp, kb in schemes: st = step_time(m, chip, 1, n, (n + 1) / 2, wb, kb, cp) rows.append([chip.name, name, fmt_time(st["t_math"]), fmt_time(st["t_mem"]), st["bound"], fmt_time(st["t_step"]), f"{base / st['t_step']:.2f}×"]) print_table(["芯片", "方案", "计算", "搬运", "瓶颈", "TTFT 下界", "加速"], rows) print_observations([ "prefill 是 compute-bound 的（第 1 讲 §4.6）：仅权重量化（W8A16、W4A16）**一点都不快**——字节省了，" "可计算还是 BF16。真实内核里反量化还要额外的指令，可能反而更慢。", "W8A8 让计算也用 8 位，算力翻倍，prefill 快约 2 倍。所以'哪种量化'取决于你要加速哪个阶段：" "decode 主要看字节，prefill 主要看计算精度。", "Blackwell Ultra 上 NVFP4 的算力是 FP8 的 3 倍，prefill 再快 3 倍（相对 BF16 6 倍）。前提是 4 位的激活不伤精度——" "这就是 NVFP4 要用 16 个数一块、E4M3 缩放因子的原因（前置讲 P3 §6.5、本讲 01 脚本）。", ]) # --------------------------------------------------------------------------- # 演示 3：临界批大小 # --------------------------------------------------------------------------- def demo3_critical_batch(): print_header("演示 3：临界批大小——decode 在多大的批时变成 compute-bound（只看权重，ctx → 0）") m = LLAMA31_8B rows = [] for chip, schemes in ((H100, SCHEMES_H100[:5]), (B300, SCHEMES_B300[:2] + SCHEMES_B300[3:])): for name, wb, cp, _ in schemes: w_bytes = m.matmul_params() * wb / 8 b_star = w_bytes * chip.peak_flops(cp) / (chip.mem_bw * m.linear_flops_per_token()) rows.append([chip.name, name, f"{chip.critical_intensity(cp):.0f}", f"{b_star:.0f}"]) print_table(["芯片", "方案", "屋脊点（计算精度）", "临界批大小 B*"], rows) print_observations([ "B* = (读权重的时间) / (每个 token 的计算时间) = 屋脊点 × (每个权重的字节 / 2)。" "全 BF16 时 H100 的 B* ≈ 295（第 1 讲 §3.6）；W8A8 时字节减半、算力翻倍，B* 不变——两边同比例变化。", "仅权重量化是**不对称**的：字节变少、算力不变，B* 按比特数成比例下降（W4A16 约 76）。" "这意味着 decode 更早'吃饱'算力：B 超过约 76 之后，W4A16 的优势就只剩 KV 以外的那点字节了。", "Blackwell Ultra 的 NVFP4 是另一个不对称的例子（第 1 讲 §5.4）：字节只省到 FP8 的 56%，算力却是 3 倍，B* 反而从 312 " "升到约 527。（第 1 讲 §3.6 用单个 8192×8192 的矩阵乘法、把激活的读写也算进去，得到 747；差别来自激活字节。）", ]) # --------------------------------------------------------------------------- # 演示 4：容量 # --------------------------------------------------------------------------- def max_batch_bits(model, chip, ctx, w_bits, kv_bits, reserve=0.1): """显存能放下几个 ctx 长的序列：(容量 × (1 − 预留) − 全部参数的字节) / (ctx × 每 token KV 字节)。放不下权重时返回 0。""" free = chip.mem_cap * (1 - reserve) - model.n_params() * w_bits / 8 if free <= 0: return 0 return int(math.floor(free / (ctx * model.kv_bytes_per_token(kv_bits / 8)))) def demo4_capacity(): print_header("演示 4：容量——Llama-3.1-70B 放在一块芯片上，上下文 32K（预留 10%）") m, ctx = LLAMA31_70B, 32768 combos = [("W16 / KV16", 16, 16, "bf16"), ("W8 / KV16", 8, 16, "fp8"), ("W8 / KV8", 8, 8, "fp8"), ("W4 / KV8", 4.125, 8, "fp8"), ("W4 / KV4", 4.125, 4.5, "fp8")] rows = [] for chip in (H100, B300): for name, wb, kb, cp in combos: B = max_batch_bits(m, chip, ctx, wb, kb) if B == 0: fits = m.n_params() * wb / 8 < chip.mem_cap * 0.9 rows.append([chip.name, name, fmt_bytes(m.n_params() * wb / 8), 0, "放不下一个序列" if fits else "权重放不下", "-", "-"]) continue st = step_time(m, chip, B, 1, ctx, wb, kb, cp) rows.append([chip.name, name, fmt_bytes(m.n_params() * wb / 8), B, fmt_time(st["t_step"]), f"{B / st['t_step']:,.0f}", f"{1 / st['t_step']:.0f}"]) print_table(["芯片", "权重 / KV 比特", "权重", "最多几个序列", "一步下界", "tok/s/芯片", "tok/s/用户"], rows) print_observations([ "70B 的 BF16 权重（141 GB）放不进 H100；8 位权重（71 GB）放得下，但剩下的约 1.5 GB 连一个 32K 序列的 KV 都放不下" "（8 位 KV 也要 5.4 GB）。4 位权重 + 8 位 KV，一块 H100 才能同时服务 6 个 32K 的序列，KV 再降到 4 位是 11 个。", "在 Blackwell Ultra（288 GB）上，每降一档比特，能服务的序列数都增加：KV 从 16 位降到 8 位，序列数翻倍" "（17 → 35）；权重再从 8 位降到 4 位，只多出 6 个（权重已经不是显存的大头）。**量化的第一个好处常常是容量，而不是速度**。", "批变大以后每芯片的吞吐（tok/s/芯片）随之上升，每个用户的速度却下降——量化让你可以沿着第 1 讲 §4.5 的前沿走得更远。", "KV 降到 4 位（每个数 4.5 位）在这里又多出约 80% 的序列（41 → 73）；01 脚本演示 4 说明了它的精度代价比 8 位大一个数量级。", ]) if __name__ == "__main__": demo1_decode() demo2_prefill() demo3_critical_batch() demo4_capacity()