""" ============================================================================= CS349E · 第 6 讲 · 05_spec_decode_on_hardware.py —— 投机解码在硬件上：c、批大小、上下文与树形验证 ============================================================================= 【这个脚本要回答的问题】 Leviathan 的公式有两个关于硬件的假设：(1) 目标模型一次验证 γ+1 个 token 和一次只算 1 个一样快； (2) 代价系数 c 是个已知常数。这两个假设什么时候成立、什么时候不成立？用第 1 讲的 roofline 算一算： 演示 1：Llama-3.1-8B（目标）+ Llama-3.2-1B（草稿），一块 H100，批大小 1：c 是多少；验证 γ+1 个 token 要多久； 最优 γ 和加速比。 演示 2：批大小 × 上下文长度：批一大，验证不再"免费"，投机解码的收益缩水；但上下文越长，收益保持得越久。 演示 3：树形验证（Medusa、EAGLE、SpecInfer）：每层草稿给出前 k 名，一次验证一棵树。在第 3 讲的模型上 （贪心解码）算出接受长度，再换算成 H100 上的加速比：批小时树很划算，批大时不划算。 【模型的假设（务必看）】 每次前向的时间 = max(FLOPs / (峰值 × 0.6), 字节 / (带宽 × 0.8)) + 固定开销（目标 1 ms、草稿 0.3 ms，假设值）； 验证时同一个序列的 γ+1 个 query 一起读一遍 KV（融合注意力内核）。草稿模型自己的 KV cache 不计入显存。 接受率 α 是**给定的参数**（真实值见 EAGLE 表 2：每个位置约 0.65–0.8），不是测出来的。 演示 3 的接受长度来自随机初始化的小模型，只说明"树比链多接受多少"的机制。 【运行】python3 05_spec_decode_on_hardware.py （约 3 秒） 【依赖】numpy；本目录的 decode_lab；../lecture01 的 hw_specs、llm_math、console； ../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 (LLAMA32_1B, TransformerLM, autoregressive_generate, best_round, draft_ranks, # noqa: E402 expected_tokens, tree_accept_lengths, verify_time) from console import fmt_time, print_header, print_observations, print_table # noqa: E402 from hw_specs import H100 # noqa: E402 from llm_math import LLAMA31_8B, max_batch # noqa: E402 from quant_lab import QCFG, quantize_params, sample_tokens # noqa: E402 from transformer_lab import config_like, init_params # noqa: E402 TGT, DRF = LLAMA31_8B, LLAMA32_1B # ----------------------------------------------------------------------------- def demo1_c_and_verify(): print_header("演示 1：Llama-3.1-8B + Llama-3.2-1B，一块 H100，批大小 1、上下文 1024") ctx = 1024 t_base = verify_time(TGT, H100, 1, ctx, 1, overhead=1e-3) t_d = verify_time(DRF, H100, 1, ctx, 1, overhead=0.3e-3) print(f"目标模型普通 decode 一步 {fmt_time(t_base)}；草稿模型一步 {fmt_time(t_d)} → c = {t_d / t_base:.3f}" f"（草稿参数 {DRF.n_params() / 1e9:.2f} B，是目标的 {DRF.n_params() / TGT.n_params():.0%}）") rows = [] for n in (1, 2, 3, 5, 9, 17, 33, 65, 129, 257): t = verify_time(TGT, H100, 1, ctx, n, overhead=1e-3) rows.append([n, fmt_time(t), f"{t / t_base:.3f}"]) print("\n目标模型一次前向送进 n 个 token（验证 γ = n − 1 个草稿）要多久：") print_table(["n", "时间", "相对 n = 1"], rows) rows2 = [] for a in (0.6, 0.7, 0.8, 0.9): g, r = best_round(TGT, DRF, H100, 1, ctx, a, overhead_t=1e-3, overhead_d=0.3e-3) rows2.append([a, g, f"{expected_tokens(a, g):.2f}", f"{r['speedup']:.2f}×" if r else "1.00×"]) print_table(["α", "最优 γ", "每轮 token", "加速比（一轮 = γ 次草稿 + 1 次验证）"], rows2) print_observations([ "批大小 1 时，送进 2 个 token 和 33 个 token 几乎一样快（前 129 个 token 都在\"读权重\"的阴影里）：" "decode 的算力利用率极低，验证几乎是免费的——Leviathan 的假设 (1) 在这里成立。", f"但 c 不小：1B 的草稿只读 8B 的约 15% 的字节，加上每步的固定开销，c ≈ {t_d / t_base:.2f}。" "按 04 演示 2 的表，c ≈ 0.2 时最优只草稿几个、加速 2 倍左右（α = 0.8）。", "想更快，要么草稿更便宜（EAGLE 那样的单层草稿、n-gram / 提示查找，或者放到 SRAM 芯片上，第 5 讲 §8），" "要么更准（同系列的模型、在目标模型的输出上蒸馏），要么一次验证更多候选（演示 3 的树）。", ]) # ----------------------------------------------------------------------------- def demo2_batch_ctx(): print_header("演示 2：批大小 × 上下文长度（α = 0.8）——每格：最优 γ、加速比；\"放不下\"= KV 超出 80 GB 显存") ctxs = (512, 4096, 32768) rows = [] for B in (1, 16, 64, 128, 256, 512): cells = [B] for ctx in ctxs: if B > max_batch(TGT, H100, ctx): cells.append("放不下") continue g, r = best_round(TGT, DRF, H100, B, ctx, 0.8, overhead_t=1e-3, overhead_d=0.3e-3) cells.append("不投机" if r is None else f"γ={g}, {r['speedup']:.2f}×（验证 {r['verify_ratio']:.2f}）") rows.append(cells) print_table(["批大小"] + [f"上下文 {c}" for c in ctxs], rows) print("（括号里是\"验证 γ+1 个 token 的时间 / 普通 decode 一步\"。）") B, ctx = 512, 512 base = verify_time(TGT, H100, B, ctx, 1, overhead=1e-3) ver = verify_time(TGT, H100, B, ctx, 5, overhead=1e-3) print(f"\n例：批 {B}、上下文 {ctx}：普通一步 {fmt_time(base)}，一次验证 5 个 token {fmt_time(ver)}（{ver / base:.1f} 倍）——" f"批 × 5 = {B * 5} 个 token 已经让矩阵乘法变成算力受限。") print_observations([ "批小时验证几乎免费，加速约 1.9 倍；批大到\"批 × (γ+1)\"超过临界批大小（H100 上按 60% / 80% 的效率约 220 个 token，" "第 1 讲 §3.6），验证的时间就按 γ+1 增长，最优 γ 变小：批 128 时 1.6 倍、256 时 1.4 倍、512 时只剩 1.1 倍。", "长上下文时，显存先装不下大批（32K 上下文最多十几个序列），decode 一直受带宽限制，验证始终几乎免费——" "投机解码在长上下文、小批的场景里最值钱。同一批大小下长上下文的加速略低，是因为草稿模型也要读自己的 KV，c 变大了。", "所以\"批处理\"和\"投机\"在抢同一种资源——闲着的算力。连续批处理把批做大以后，留给投机解码的空间就少了；" "服务系统要按负载决定用不用、草稿多长（vLLM、SGLang 都支持按批大小调整或关闭投机）。", ]) # ----------------------------------------------------------------------------- def toy_ranks(): """ 在第 3 讲的模型上：目标模型贪心生成 64 条序列（每条 8 个提示 token + 48 个新 token）， 再让每个草稿模型在这些真实序列上算一遍 logits，记下"正确答案在草稿里排第几"。 （只生成 48 个：随机模型贪心生成太长会陷入重复循环，接受率被虚高。） """ P = init_params(QCFG, seed=0) tgt = TransformerLM(P, QCFG, "目标") prompts = sample_tokens(P, QCFG, n_seq=64, length=8, seed=3) seqs = np.array([np.concatenate([pr, autoregressive_generate(tgt, pr, 48, 0, np.random.default_rng(0))[0]]) for pr in prompts]) drafts = {} pk = dict(P) pk["layers"] = P["layers"][:1] drafts["只用前 1 层（弱）"] = TransformerLM(pk, config_like(QCFG, n_layers=1), "d1") drafts["INT4 量化的自己（强）"] = TransformerLM(quantize_params(P, "int4-g32", include_lm_head=True)[0], QCFG, "d4") out = {} for name, dr in drafts.items(): lg = dr.full_logits(seqs) out[name] = np.concatenate([draft_ranks(lg[i, 7:], seqs[i, 7:]) for i in range(len(seqs))]) return out def demo3_trees(): print_header("演示 3：树形验证——每层取草稿的前 k 名，一次验证整棵树（贪心解码）") ranks = toy_ranks() trees = [(1,), (1, 1), (1, 1, 1, 1), (1,) * 6, (2, 2, 2), (3, 3, 3), (4, 2, 2, 1), (8, 4, 2, 1), (4, 4, 4, 4)] for name, rk in ranks.items(): print(f"\n草稿：{name}；正确答案在草稿的前 1/2/4/8 名里的比例：" + " / ".join(f"{(rk < k).mean():.2f}" for k in (1, 2, 4, 8))) rows = [] for br in trees: acc, nodes = tree_accept_lengths(rk, br) tau = 1 + acc.mean() cells = ["链 γ=" + str(len(br)) if max(br) == 1 else "树 " + "×".join(map(str, br)), nodes, f"{tau:.2f}"] for B in (1, 32): ctx = 1024 t_base = verify_time(TGT, H100, B, ctx, 1, overhead=1e-3) t_round = len(br) * 0.05 * t_base + verify_time(TGT, H100, B, ctx, nodes + 1, overhead=1e-3) cells.append(f"{tau * t_base / t_round:.2f}×") rows.append(cells) print_table(["草稿形状", "每轮验证的草稿 token", "τ = 每轮产出的 token", "H100 加速 批 1", "批 32"], rows) print("\n（换算到 H100 时：目标 = Llama-3.1-8B、上下文 1024；草稿每层的代价假设为目标一步的 5%（EAGLE 式的单层草稿），" "树有几层就草稿几次；验证 = 一次送进\"节点数 + 1\"个 token。）") print_observations([ "链式草稿越长，后面的位置越难全部猜中，τ 涨得越来越慢；树在每一层多放几个候选，只要正确答案在前几名里就能继续，" "τ 明显更高——弱草稿受益更多（它的第一名常错，但前 4 名常对）。这是 SpecInfer、Medusa、EAGLE 用树的原因。", "代价是每轮要验证的 token 多得多（几十到几百个）。批 1 时这些 token 几乎免费，大树加速最多；批 32 时，" "32 × 几百个 token 让验证变成算力受限，大树反而比小树、甚至比链更慢。", "EAGLE 报告树形验证让平均接受长度多 0.6–0.8、加速比多 0.3–0.5（表 5、图 7）：接受长度涨得多，加速比涨得少，" "就是因为每次前向的 token 变多了。EAGLE-2 进一步按草稿的置信度动态决定树的形状。", ]) if __name__ == "__main__": demo1_c_and_verify() demo2_batch_ctx() demo3_trees()