""" ============================================================================= CS349E · 第 5 讲 · 05_heterogeneous_spec_decode.py —— 异构解耦：把投机解码的草稿模型放到 SRAM 芯片上 ============================================================================= 【这个脚本要回答的问题】 本讲第二篇阅读材料（Gimlet Labs 博客，2026）：用 gpt-oss-120b 做目标模型、一个 1.6B 的小模型做草稿模型， 把草稿模型从 GPU 搬到 d-Matrix Corsair（每张卡 2 GB 片上 SRAM、150 TB/s），在同样能耗下端到端快 2–10 倍。 这里用一个**参数化的玩具模型**理解为什么： 演示 1：投机解码的期望收益：每轮草稿 γ 个、每个被接受的概率 α，平均能前进几个 token；复现博客里的接受率表。 演示 2：一轮的时间 = γ × 草稿一个 token + 目标模型验证一次 + 两次跨设备的网络延迟。 草稿在 GPU 上 vs 在 SRAM 芯片上：最优草稿长度和加速比怎样变化。以及 MoE 目标模型的一个陷阱：验证的 token 越多，碰到的专家越多。 演示 3：另一种异构切法——注意力 / 前馈网络（FFN）解耦：decode 时两者的算术强度差一个数量级，适合不同的硬件。 【模型的假设（务必看）】 目标模型按 gpt-oss-120b 的公开结构近似：36 层、128 个专家、每个 token 选 4 个、每个专家约 2490 万参数， 非专家部分（注意力、词表等）约 15 亿参数；权重 MXFP4（4.25 位）。草稿模型：16 亿参数的稠密模型，FP8。 每步时间 = 固定开销 + 读权重的字节 / (带宽 × 0.8)。固定开销（GPU 目标 1.5 ms、GPU 草稿 0.3 ms、SRAM 芯片 0.05 ms） 和跨设备的单程网络延迟（10 µs）都是**假设值**。这些数不对应 Gimlet 的实验条件（他们按每千瓦吞吐归一化、有批处理）， 只用来展示趋势。所有输出每次运行完全相同。 【运行】python3 05_heterogeneous_spec_decode.py （< 1 秒） 【依赖】numpy；../lecture02/transformer_lab（演示 3 的逐运算账本）；../lecture01 的 console、hw_specs、llm_math。 """ import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) for sub in ("lecture01", "lecture02"): sys.path.insert(0, os.path.join(HERE, "..", sub)) 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_70B # noqa: E402 from transformer_lab import ATTN_CORE_OPS, MATMUL_OPS, analytic_ops, sum_rows # noqa: E402 # ---- 玩具模型的参数（见文件开头的说明）---- T_LAYERS, T_EXPERTS, T_TOPK, T_EXPERT_PARAMS, T_DENSE_PARAMS, T_BITS = 36, 128, 4, 24.9e6, 1.5e9, 4.25 D_PARAMS, D_BITS = 1.6e9, 8.0 GPU_BW = H100.mem_bw * 0.8 SRAM_BW = 2 * 150e12 * 0.8 # 两张 Corsair 卡 OVH_TARGET, OVH_DRAFT_GPU, OVH_SRAM = 1.5e-3, 0.3e-3, 0.05e-3 NET = 10e-6 # 跨设备单程延迟 def expected_accepted(alpha, gamma): """ 草稿 γ 个 token、每个独立地以概率 α 被接受（第一个被拒绝之后的都作废）：被接受的草稿 token 数的期望 = Σ_{i=1..γ} α^i。 【推导】第 i 个草稿 token 被接受，当且仅当前 i 个都被接受，概率 α^i。期望 = 各个"被接受"的概率之和。 【边界】α = 1 时等于 γ。 """ return sum(alpha ** i for i in range(1, gamma + 1)) def tokens_per_round(alpha, gamma): """ 一轮"草稿 + 验证"平均前进的 token 数 = 被接受的草稿数 + 1（验证时目标模型总会多给一个 token： 要么是第一个被拒绝位置上的修正，要么是全部接受之后的下一个）= (1 − α^(γ+1)) / (1 − α)（Leviathan et al. 2023）。 """ return expected_accepted(alpha, gamma) + 1 def touched_experts(n_tokens): """n 个 token、每个独立地从 128 个专家里选 4 个时，一层里被碰到的不同专家数的期望（第 3 讲 §7.6 的公式）。""" return T_EXPERTS * (1 - (1 - T_TOPK / T_EXPERTS) ** n_tokens) def t_target(n_tokens): """目标模型一次前向（验证 n 个 token，单个用户）的时间：固定开销 + (非专家权重 + 被碰到的专家权重) / 带宽。""" by = (T_DENSE_PARAMS + T_LAYERS * touched_experts(n_tokens) * T_EXPERT_PARAMS) * T_BITS / 8 return OVH_TARGET + by / GPU_BW def t_draft(where): """草稿模型生成一个 token 的时间：在 GPU 上，或者在 SRAM 芯片上。""" by = D_PARAMS * D_BITS / 8 return OVH_DRAFT_GPU + by / GPU_BW if where == "gpu" else OVH_SRAM + by / SRAM_BW def round_time(gamma, where, moe=True): """ 一轮的时间：γ 次草稿 + 一次验证（γ + 1 个 token）+（草稿和目标不在同一台设备上时）一来一回的网络延迟。 【参数】moe=False 时假设验证的开销和 token 数无关（稠密模型、小批量的近似），用来对比 MoE 的影响。 """ verify = t_target(gamma + 1) if moe else t_target(1) net = 2 * NET if where == "sram" else 0.0 return gamma * t_draft(where) + verify + net # --------------------------------------------------------------------------- # 演示 1：期望收益与接受率表 # --------------------------------------------------------------------------- def demo1_acceptance(): print_header("演示 1：投机解码平均每轮前进几个 token——复现 Gimlet 博客的接受率表") table = [(5, 4, 0.926), (10, 7, 0.934), (20, 12, 0.948), (40, 20, 0.962)] rng = np.random.default_rng(0) rows = [] for gamma, reported, alpha in table: acc = expected_accepted(alpha, gamma) sims = [] for _ in range(20000): k = 0 while k < gamma and rng.random() < alpha: k += 1 sims.append(k) rows.append([gamma, f"{alpha * 100:.1f}%", reported, f"{acc:.2f}", f"{np.mean(sims):.2f}", f"{tokens_per_round(alpha, gamma):.2f}"]) print_table(["草稿长度 γ", "每个 token 的接受率 α（博客）", "被接受的 token（博客）", "公式 Σα^i", "模拟（2 万轮）", "每轮前进（+1 个修正 token）"], rows) print_observations([ "博客表里的'被接受的 token 数'正好等于 Σ_{i=1..γ} α^i：第 i 个草稿被接受要求前 i 个全被接受。模拟和公式一致。", "草稿越长，平均被接受的越多，但效率越低：γ = 40 时只接受一半（20 个），后面的大多白算——一旦有一个被拒绝，后面的全部作废。", "所以草稿的**代价**决定了能草稿多长：草稿越便宜，越可以'多猜一些'（演示 2）。博客说的正是这一点。", "（为什么投机解码不改变输出的分布——拒绝采样的正确性——是第 6 讲的内容。）", ]) # --------------------------------------------------------------------------- # 演示 2：草稿放在哪 # --------------------------------------------------------------------------- def demo2_where_to_draft(): print_header("演示 2：草稿模型放在 GPU 上 vs 放在 SRAM 芯片上（单个用户，α = 0.93）") alpha = 0.93 base = t_target(1) print(f"目标模型普通 decode 一步：{fmt_time(base)}（{1 / base:.0f} tok/s）；草稿一个 token：GPU 上 {fmt_time(t_draft('gpu'))}，" f"SRAM 芯片上 {fmt_time(t_draft('sram'))}；验证 6 / 21 / 41 个 token：{fmt_time(t_target(6))} / {fmt_time(t_target(21))} / " f"{fmt_time(t_target(41))}（碰到的专家：每层 {touched_experts(6):.0f} / {touched_experts(21):.0f} / {touched_experts(41):.0f} 个）") rows = [] best = {} for gamma in (1, 2, 3, 5, 8, 10, 15, 20, 30, 40): cells = [gamma, f"{tokens_per_round(alpha, gamma):.2f}"] for where in ("gpu", "sram"): for moe in (True, False): tpot = round_time(gamma, where, moe) / tokens_per_round(alpha, gamma) sp = base / tpot key = (where, moe) if sp > best.get(key, (0, 0))[0]: best[key] = (sp, gamma) cells.append(f"{sp:.2f}×") rows.append(cells) print_table(["γ", "每轮前进", "GPU 草稿（MoE 验证）", "GPU 草稿（验证代价不变）", "SRAM 草稿（MoE 验证）", "SRAM 草稿（验证代价不变）"], rows) rows = [[("GPU" if w == "gpu" else "SRAM 芯片"), "MoE（碰到的专家随 token 数增加）" if m else "验证代价不变（稠密近似）", g, f"{s:.2f}×"] for (w, m), (s, g) in best.items()] print("\n每种情况下最好的草稿长度和加速比：") print_table(["草稿放在", "目标模型的验证", "最优 γ", "相对普通 decode 的加速"], rows) print_observations([ "草稿模型很小（1.6B），但在 GPU 上生成一个 token 也要约 0.9 ms——它和目标模型一样受限于读权重和固定开销，" "草稿 5 个 token（4.5 ms）就将近普通 decode 两步的时间（5.0 ms）。所以在 GPU 上草稿不能太长，加速比也有限。", "放到 SRAM 芯片上，草稿一个 token 只要几十微秒（150 TB/s 的片上带宽，几乎只剩固定开销）：多猜几个几乎不花钱，" "最优草稿长度变长、加速比变大。这就是博客的结论：'草稿越便宜，越可以激进地草稿；也可以在同样延迟下用更大、更准的草稿模型'。", "MoE 目标模型有个陷阱：验证 γ + 1 个 token 时，每个 token 各选 4 个专家，一层要读的专家从 4 个涨到几十个——验证 21 个 token " "的代价是普通一步的好几倍。所以 MoE 上的最优草稿长度比'验证代价不变'时短得多（批大的时候专家本来就都被碰到了，这个陷阱会变小）。", "博客报告：同样能耗下，GPU 上的投机解码比单纯的 prefill / decode 解耦快 2–5 倍；草稿换到 Corsair 上再快 2–10 倍（例子：25 s 的请求" "先快 5 倍多，再快 2.6 倍）。我们的玩具模型只看单个用户、不算能耗，数字不可比，但方向一致。", ]) # --------------------------------------------------------------------------- # 演示 3：注意力 / FFN 解耦 # --------------------------------------------------------------------------- def demo3_attention_ffn(): print_header("演示 3：另一种异构切法——注意力和 FFN 的算术强度（Llama-3.1-70B decode，上下文 4096，BF16）") print("（时间按一块 H100 的峰值算，忽略 70B 放不进一块卡的问题，只用来比较两部分各自受什么限制）") m = LLAMA31_70B ffn_ops = ("q_proj", "k_proj", "v_proj", "o_proj", "mlp_gate", "mlp_up", "mlp_down", "lm_head") rows = [] for B in (16, 64, 256, 1024): a = analytic_ops(m, B, 1, 4095, last_only=True) fa, ba = sum_rows(a, "useful", ATTN_CORE_OPS), sum_rows(a, "kv", ATTN_CORE_OPS) ff, bf = sum_rows(a, "useful", ffn_ops), sum_rows(a, "w", MATMUL_OPS) rows.append([B, f"{fa / ba:.1f}", fmt_time(ba / H100.mem_bw), f"{ff / bf:.0f}", fmt_time(max(ff / H100.peak_flops('bf16'), bf / H100.mem_bw))]) print_table(["批大小 B", "注意力核心：强度", "注意力：读 KV 的时间", "投影 + FFN：强度", "投影 + FFN：时间下界"], rows) print_observations([ "注意力核心的强度固定在组大小（8，GQA），批再大也不变，读 KV 的时间和批大小成正比；投影和 FFN 的强度等于批大小" "（第 2 讲 §4.5、§5.6），B = 1024 时已经是 compute-bound（计算 144 ms > 读权重 41 ms）。" "两者一个是纯粹的带宽 / 容量问题，一个是算力问题——同一块芯片很难同时对两者都合适。", "所以有一类系统（MegaScale-Infer、StepFun 的 Step-3 等）把每一层的注意力和 FFN 拆到不同的机器上：" "注意力放在显存大、带宽高的机器上（它要存所有序列的 KV），FFN（尤其是 MoE 的专家）放在算力强的机器上，用很大的批喂饱它；" "每一层都要把激活在两组机器之间来回传——用网络换硬件的匹配度。", "这和 prefill / decode 解耦是同一个思路的不同切法：**按算术强度把计算分给最合适的硬件，代价是搬数据和协调**（第 1 讲 §6）。", ]) if __name__ == "__main__": demo1_acceptance() demo2_where_to_draft() demo3_attention_ffn()