""" ============================================================================= CS349E · 第 6 讲 · 04_spec_decode_speedup.py —— 投机解码能快多少：Leviathan 的公式与实测 ============================================================================= 【这个脚本要回答的问题】 03 证明了投机解码"不改变分布"。这里回答"值不值"： 演示 1：复现 Leviathan 表 1——给定接受率 α 和草稿长度 γ，速度提高多少、总运算量增加多少（假设草稿免费）。 演示 2：草稿不免费时（代价系数 c = 草稿一次 / 目标一次）最优的 γ（论文图 3）和推论 3.9 的下界。 演示 3：在第 3 讲的模型上真的跑投机解码（带 KV cache 回滚）：每次目标模型前向平均产出几个 token， 和式 (1) 用实测 α 算出来的比；被接受个数的分布和"独立同分布"假设比。 演示 4：账本：目标模型的总运算量（定理 3.11）和读权重的次数（§3.4："内存访问可以变少"）。 【怎么读输出】"每次目标前向的 token" 就是论文式 (1) 的 E(# generated tokens)。它直接等于"读目标模型权重的次数 减少了几倍"，在 decode 受带宽限制时也近似等于加速比（草稿的代价另算）。 【运行】python3 04_spec_decode_speedup.py （约 20 秒） 【依赖】numpy；本目录的 decode_lab、03 脚本里同样的模型构造（这里重新写一遍，避免导入以数字开头的文件）； ../lecture02 的 transformer_lab；../lecture03 的 quant_lab（需要 ../prereq3）。 """ import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) for sub in ("lecture03", "prereq3"): sys.path.insert(0, os.path.join(HERE, "..", sub)) from decode_lab import (NGramLM, TransformerLM, adjust, autoregressive_generate, best_gamma, # noqa: E402 expected_tokens, ops_factor, speculative_generate, walltime_speedup) from console import fmt_bytes, fmt_num, print_header, print_observations, print_table # noqa: E402 from quant_lab import QCFG, quantize_params, sample_tokens # noqa: E402 from transformer_lab import MATMUL_OPS, Counter, config_like, init_params # noqa: E402 # ----------------------------------------------------------------------------- def demo1_table1(): print_header("演示 1：复现 Leviathan 表 1（假设草稿免费：c = ĉ = 0）") paper = [(0.6, 2, 1.53, 1.96), (0.7, 3, 1.58, 2.53), (0.8, 2, 1.23, 2.44), (0.8, 5, 1.63, 3.69), (0.9, 2, 1.11, 2.71), (0.9, 10, 1.60, 6.86)] rows = [] for a, g, ops_p, sp_p in paper: rows.append([f"{a}", g, f"{expected_tokens(a, g):.2f}", f"{ops_factor(a, g, 0):.2f}×", f"{ops_p:.2f}×", f"{walltime_speedup(a, g, 0):.2f}×", f"{sp_p:.2f}×"]) print_table(["α", "γ", "每轮 token（式 1）", "运算量 公式", "论文", "速度 公式", "论文"], rows) print_observations([ "公式和论文表 1 逐格一致。速度 = 每轮平均产出的 token 数（草稿免费时，一轮的时间 = 目标模型前向一次）。", "运算量总是增加：每轮目标模型算了 γ+1 个位置，被拒绝之后的都白算。α = 0.9、γ = 10 时快 6.9 倍，但多做 60% 的运算——" "投机解码是用\"闲着的算力\"换延迟，前提是算力真的闲着（decode 受带宽限制时成立，05 脚本看什么时候不成立）。", "γ 越大，速度越高但运算浪费越多；α 越低，浪费越多（α = 0.6、γ = 2 就多做 53%）。", ]) # ----------------------------------------------------------------------------- def demo2_optimal_gamma(): print_header("演示 2：草稿不免费时，最优的 γ（论文图 3）与推论 3.9") cs = (0.01, 0.02, 0.05, 0.1, 0.2, 0.5) alphas = (0.5, 0.6, 0.7, 0.8, 0.9) rows = [] for a in alphas: cells = [f"{a}"] for c in cs: g, s = best_gamma(a, c) cells.append("不划算" if g == 0 else f"γ={g}, {s:.2f}×") rows.append(cells) print("每格：最优草稿长度 γ*、那时的墙钟加速比（定理 3.8：(1 − α^{γ+1}) / ((1 − α)(γc + 1))）") print_table(["α \\ c"] + [f"c = {c}" for c in cs], rows) print("\n推论 3.9：只要 α > c，γ = 1 就一定有加速，而且至少 (1 + α) / (1 + c) 倍：") rows2 = [] for a, c in ((0.6, 0.5), (0.5, 0.5), (0.3, 0.2), (0.8, 0.05)): g, s = best_gamma(a, c) rows2.append([a, c, f"{walltime_speedup(a, 1, c):.3f}", f"{(1 + a) / (1 + c):.3f}", f"γ*={g}, {s:.3f}×"]) print_table(["α", "c", "γ = 1 的加速", "(1+α)/(1+c)", "最优"], rows2) print_observations([ "草稿越便宜（c 越小）、越准（α 越高），最优的 γ 越大：α = 0.9 时，c = 0.01 → γ* = 24（7.5 倍），c = 0.05 → 13，" "c = 0.5 → 只草稿 3 个（1.4 倍）。", "c 不是模型本身的属性，取决于硬件和实现（论文 §3.3）：同一个草稿模型在 GPU 上和 SRAM 芯片上 c 可以差十几倍" "（第 5 讲 §8.3），所以最优 γ 也跟着变。", "α = c 时恰好不赚不赔（γ = 1 的加速正好是 1）；α < c 时任何 γ 都更慢。", ]) # ----------------------------------------------------------------------------- def build_models(): """目标 = 第 3 讲的 QCFG 模型；草稿 = INT4 量化的自己、只用前 3 层 / 前 1 层、二元语法表（同 03 脚本）。""" P = init_params(QCFG, seed=0) target = TransformerLM(P, QCFG, "目标") drafts = {"INT4 量化的自己": TransformerLM(quantize_params(P, "int4-g32", include_lm_head=True)[0], QCFG, "int4")} for k in (3, 1): pk = dict(P) pk["layers"] = P["layers"][:k] drafts[f"只用前 {k} 层"] = TransformerLM(pk, config_like(QCFG, n_layers=k), f"前{k}层") drafts["二元语法表"] = NGramLM.fit(sample_tokens(P, QCFG, n_seq=64, length=128, seed=1), QCFG.vocab) return target, drafts def measured_alpha(target, draft, toks, T): """在目标模型自己采样的文本上逐位置算 Σ min(p, q)，取平均——论文推论 3.6 的 α = E(min(p, q))。""" lp = target.full_logits(toks) lq = draft.table[toks] if isinstance(draft, NGramLM) else draft.full_logits(toks) return float(np.minimum(adjust(lp, T), adjust(lq, T)).sum(-1).mean()) def demo3_real_loop(target, drafts): print_header("演示 3：真的跑投机解码（温度 1、批大小 1、8 个提示 × 96 个新 token）") toks = sample_tokens(target.params, QCFG, n_seq=8, length=64, seed=0) prompts = sample_tokens(target.params, QCFG, n_seq=8, length=16, seed=11) n_new = 96 rows = [] hist_case = None for name, dr in drafts.items(): a = measured_alpha(target, dr, toks, 1.0) for g in (1, 2, 4, 8): calls = 0 acc = [] for i, pr in enumerate(prompts): _, st = speculative_generate(target, dr, pr, n_new, g, 1.0, np.random.default_rng(100 + i)) calls += st["target_calls"] acc += st["accepted"] meas = len(prompts) * n_new / calls rows.append([name, f"{a:.3f}", g, f"{meas:.2f}", f"{expected_tokens(a, g):.2f}", f"{np.mean(acc) / g:.3f}"]) if name == "只用前 3 层" and g == 8: hist_case = (a, g, acc) print_table(["草稿", "α（实测 E[Σmin(p,q)]）", "γ", "每次目标前向的 token（实测）", "式 (1) 的预测", "被接受的草稿比例"], rows) a, g, acc = hist_case counts = np.bincount(acc, minlength=g + 1) / len(acc) pred = [a ** n * (1 - a) for n in range(g)] + [a ** g] print(f"\n只用前 3 层、γ = {g}：每轮被接受的草稿个数的分布（共 {len(acc)} 轮） vs 独立同分布假设下的截断几何分布") print_table(["被接受的个数"] + [str(n) for n in range(g + 1)], [["实测"] + [f"{c:.3f}" for c in counts], ["i.i.d. 预测"] + [f"{x:.3f}" for x in pred]]) print_observations([ "实测和式 (1) 大体对得上：草稿越准、γ 越大，每次目标前向产出的 token 越多；INT4 草稿 γ = 8 时平均 5 个以上。", "实测和式 (1) 在抽样误差以内一致（每格只有一两百到四百多轮），有高有低，没有系统偏差；被接受个数的分布也和截断几何分布" "大体相符。式 (1) 并不要求每个位置的 β 一样，只要求各位置**独立**：独立时 E[β₁β₂…βₙ] = αⁿ。这个随机模型里相邻位置的难度" "几乎不相关，所以公式成立；真实文本里难的位置往往成片出现，公式会高估（见讲义 §8.3）。", "γ 越大，\"被接受的草稿比例\"越低：草稿的后半段越来越难全部猜中，白算的越来越多（下一个演示的运算量）。", ]) # ----------------------------------------------------------------------------- def demo4_ledger(target, drafts): print_header("演示 4：账本——目标模型的总运算量（定理 3.11）与读权重的次数（§3.4）") prompts = sample_tokens(target.params, QCFG, n_seq=4, length=16, seed=11) n_new = 96 toks = sample_tokens(target.params, QCFG, n_seq=8, length=64, seed=0) dr = drafts["只用前 3 层"] per_tok = 2 * QCFG.to_model_config().matmul_params() # 目标模型每个 token 的矩阵乘法 FLOPs c_hat = dr.cfg.to_model_config().matmul_params() / QCFG.to_model_config().matmul_params() prefill_extra = len(prompts) * (prompts.shape[1] - 1) * per_tok # 提示里除最后一个 token 以外的部分（两种做法相同） c_base = Counter() for i, pr in enumerate(prompts): autoregressive_generate(target, pr, n_new, 1.0, np.random.default_rng(i), counter=c_base) base_f = c_base.total("flops", MATMUL_OPS) - prefill_extra base_w = c_base.total("w", MATMUL_OPS) rows = [["普通 decode", fmt_num(base_f), "1.00×", "-", fmt_bytes(base_w), "1.00×"]] a = measured_alpha(target, dr, toks, 1.0) for g in (2, 4, 8): c_t = Counter() for i, pr in enumerate(prompts): speculative_generate(target, dr, pr, n_new, g, 1.0, np.random.default_rng(i), counter_t=c_t) f_t = c_t.total("flops", MATMUL_OPS) - prefill_extra w_t = c_t.total("w", MATMUL_OPS) rows.append([f"投机 γ={g}", fmt_num(f_t), f"{f_t / base_f:.2f}×", f"{ops_factor(a, g, 0):.2f}×", fmt_bytes(w_t), f"{w_t / base_w:.2f}×"]) print(f"草稿 = 只用前 3 层（α = {a:.3f}，ĉ = 草稿 / 目标 每 token 运算 = {c_hat:.2f}）；4 个提示 × 96 个新 token；" "只数矩阵乘法（含 lm_head）、不含提示的预填充，按 BF16 记字节；这一栏只算目标模型，不含草稿自己的运算和读取") print_table(["做法", "目标模型运算量", "相对普通", "定理 3.11（ĉ = 0）", "目标模型读权重", "相对普通"], rows) print_observations([ "目标模型的运算量随 γ 增加，和定理 3.11 的预测接近（γ = 8 时多做约 90%）——被拒绝的位置也算过了。", "读目标模型权重的总量反而减少（γ = 8 时只剩约五分之一）：一次前向读一遍权重，却产出好几个 token（论文 §3.4）。" "decode 受带宽限制，读权重少了就是快了——投机解码和批处理是同一个道理：提高 decode 的算术强度，" "只不过\"批\"是同一个序列里的几个位置。", f"这个跳层草稿的 ĉ = {c_hat:.2f}（4 层里用了 3 层，lm_head 也照算），太贵了：真实系统里它会吃掉大部分收益（演示 2 的表里 c = 0.5 时最多快 1.1–1.4 倍）。" "真实的草稿要么小两个数量级（Leviathan：T5-small 之于 T5-XXL），要么像 EAGLE 那样只有一层。", ]) if __name__ == "__main__": demo1_table1() demo2_optimal_gamma() TARGET, DRAFTS = build_models() demo3_real_loop(TARGET, DRAFTS) demo4_ledger(TARGET, DRAFTS)