Mirror Speculative Decoding: Breaking the Serial Barrier in LLM Inference

1 · 动机:为什么 SD 卡在串行墙上

1.1 历史脉络:从 Leviathan 2023 到 EAGLE-3 的"局部最优陷阱"

把时间线拉一下,SD 的演化基本就是沿着同一条 Pareto 曲线在挪点:

• 2023:Leviathan & Chen 的 vanilla SD。draft 是个独立小模型,γ 个 token AR 生成,target 一次 forward 验证。serial pipeline 写在公式里:T_SD = T_draft + T_target。
• 2023 末:Medusa。把 draft 拆成 target 上的 K 个并行 head,γ 步压缩成一步——但 head 之间没自回归依赖,acceptance 上限低。
• 2024:Hydra。在 Medusa 基础上给 head 加上序列耦合,acceptance 回升,但开销也回来了。
• 2024:EAGLE-1/2/3。改用共享 target hidden state 的轻量 draft layer,EAGLE-3 还引入 multi-layer fusion + token-level 预测,把 acceptance 推到当前最强(对应 SpecBench 上 2.5×–3.5×)。
• 2024:tree drafts (SpecInfer / EAGLE-2 dyn-tree)。draft 不再是一条直线,而是个候选树,target 一次验证一棵树。

但所有这些仍然是单一加速器、单一时间轴上的串行:draft 必须先把 γ 步推完,target 才能开 verification(或 verification 必须等 target 完整 forward 才出 logits)。本文用一句话把这条天花板拍死:

"larger speculative modules improve acceptance at the cost of higher draft construction latency, while smaller ones reduce overhead but lower acceptance and limit speedup. ... Eq. 5 exposes the core limitation: improvements in acceptance must compensate for the added draft latency, intrinsically coupling acceptance with latency."

1.2 别的方案为什么不够

本文最尖锐的 framing 是:前面所有方法都在 "acceptance ↔ latency" 这条二维曲线上挪——你想要 acceptance 高就要付 latency。Mirror-SD 不是挪,而是给曲线加一个新的维度("跨设备时间"),让 draft 时间在临界路径上消失。

1.3 为什么"对称速推"不平凡

简单想:既然 target 的中间层 hidden 已经"差不多知道下一个 token 是什么了"(Pal et al. 2023 早就报过 logit lens 现象),那为啥不在中间层就让 draft 启动?三个非显然的难点:

1. Early exit 的分布 ≠ 最终分布。如果你直接拿 ℓ_e 层的 token 提交,会污染输出分布(违反 lossless)。所以 early exit 只能当路标(top-κ 候选 + 概率),最终接受还是要看第 N 层 logits。
2. Draft 收到 top-κ 路标后,要替每一个候选准备好后续 γ 步——这是 branch-complete。不是 draft 一条线 γ 个 token,而是 κ 棵深度 γ 的子树。κ=8、γ=7 时 draft batch 显著膨胀,如何让它仍能在 target suffix 时间内跑完?这逼出了 SS draft + NPU 上的特殊 sharding。
3. "Mirror" 的另一半:target 同时也在为 draft 做 correction。当 draft 选错分支(early-exit top-κ 没命中真正的下一 token),target 在第 N 层算出 correction 后,要么命中树里某条预算分支(reuse,见 §3),要么必须 fallback 重新启 draft。这一对称(draft 在猜 target / target 在校 draft)就是名字里"mirror"的来源。

把这三件事合在一起:Mirror-SD 既不是单纯"提早启 draft"(那只是 early-exit decoding),也不是"换硬件"(那只是异构推理),而是三件事必须同时成立才能保 lossless 又拿到加速。这是它跑 v.s. 看起来"好像谁都能想到"的差距所在。

2 · 背景速查

Vanilla SD 复习一遍。step t:

1. draft 从 (x, y_
2. target 一次 forward 算出 p^(N)(·| y_
3. 取最长前缀 r,使所有 ŷ_{t+j} = argmax p^(N);若 r<γ,补一个 correction token c,从 (p^(N) − p^(draft))₊ 重采(stochastic 模式)。
4. commit y_{t+1:t+r}(+ c),进入下一 step。

关键观察:第 2 步 target forward 必须等第 1 步 draft 完成,这就是"serial barrier"。

3 · 方法 A:Early-exit + branch-complete 速推

3.1 关键动作:在 ℓ_e 层"分叉"

step t 开始,target 像往常一样跑 1..ℓ_e 层。在 ℓ_e 层,做两件事 同时:

• 把 LM head 应用到 h_t^(ℓ_e) 得到 p^(ℓ_e);取 top-κ,打包成 M_t={(v_i, log p̃_i)}_i=1..κ 通过低带宽 token channel 发给 draft 所在的 NPU(O(Bκ) 字节,微秒级)。
• target 自己继续 ℓ_e+1..N 层,生成最终的 p^(N)——这是 verifier 真正要看的分布。

NPU 这边,draft 收到 κ 个候选根之后,同时对每一个 v_i 跑深度 γ 的 rollout(branch-complete):

这个"对每个候选都把后面 γ 步铺开"就是 branch-complete。draft 不再是一条直线,而是 κ 棵深度 γ 的树,合在一起记作 𝓣_t。

3.2 Verification vs. reuse 规则

target 跑完到 N 层得到 p^(N),沿"当前 draft 选定路径"做标准 SD verify,得到 A_t。两种情况:

• A_t = γ:整个 window 全接受,直接 commit;下一步从树根重新出发。
• A_t < γ: 在深度 τ=A_t+1 处出现 mismatch,target 用 c_t+τ ∼ p^(N)(·) 修正。

关键的"reuse"判定:被 commit 的 prefix Π_t⁺ = (accepted...,c) 在不在 𝓣_t 已经预算好的 paths 里?

Reuse 命中 = 下一 step 的 draft 不用重来,直接拿树里那条分支当新 window;fallback = draft 必须从被修正后的 context 重新 rollout,白费一次 NPU 时间。

3.3 Worked example: 一个 token "积" 的命运

设 target = Qwen3-32B,N=64,ℓ_e=32,κ=8,γ=7。当前 prefix = "微积" 之后,要预测下一个字。

1. target 跑完 32 层,LM head(h⁽³²⁾) → p⁽³²⁾。top-8 = ["分", "学", "公", "(", "方", "在", "运", "的"](假设)。把这 8 个 token + log-prob 通过 token channel 发给 NPU(8 项 × bf16,几十字节)。
2. 同一时刻:GPU 继续跑 33..64 层。NPU 上,draft 同时在 batch 维度展开 8 棵 γ=7 的树,例如对 "分":"分→法→是→一→种→重→要" 之类。SS draft 实际只需 J ≤ ⌈γ/η̄⌉ 步(后面会讲),通常 J=2~3 而非 7。
3. GPU 第 64 层算出 p⁽⁶⁴⁾,argmax = "分"——命中 v_1!
4. 用 "分" 这条分支沿深度 verify:draft 给 "分→法→是" 而 target 给 "分→法→中"。则 A_t=2(接受了 "分", "法"),τ=3,correction c="中"。
5. 检查 reuse:Π_t⁺=("分","法","中") 在不在树里?如果 draft 在 "分→法" 的 depth-2 节点下还展开了多个分支(包括 "中"),那就 reuse 命中,下一 step 直接用预算好的 "分→法→中→…" 当 window;否则 fallback。

反向论证: 如果不做 branch-complete(只让 draft 给一条线),那么 step t 的 fallback 概率 = P(target argmax ≠ draft 选的那个根 v*),通常 30%–50%——白费一半。Branch-complete 把它降到 P(c_t+τ ∉ 任何已展开 path),实测 0.1–0.3(见 §8 fallback 曲线)。

4 · 方法 B:Speculative Streaming on draft(只动 draft,不动 target)

问题:branch-complete 让 draft 要在很短时间内吐出 κ·γ 个 token。如果 draft 还是普通 AR,要 γ 个内部步——很可能撑爆 overlap budget Δ。

解法:把作者自己 2024 年的 Speculative Streaming (SS) 装到 draft 上(不是装到 target 上,这个区分很重要)。SS 的核心是 multi-stream attention:每一层除了主 stream M_t^(ℓ),还附带 γ 个 lookahead stream S_t,j^(ℓ)(j=1..γ),每个 lookahead 只看前面已 verified 部分 + 比它早的 lookahead。一次 forward 同时给出:

• 主 stream 的下一 token 分布(用来 verify 上一轮 draft)
• γ 个 lookahead 的下一 token 分布(用来填本轮 window)

每个内部 step 因此能吐 η_j ≥ 1 个 token,平均 η̄ > 1。要覆盖 γ 个 token,只需 J ≤ ⌈γ/η̄⌉ 步而不是 γ 步。实测 J ≈ 2–3 时,T_draft 砍掉 50%+。

5 · 方法 C:GPU↔NPU 异构 sharding

这一节是论文最"systems"的部分,也是 EAGLE-3 没做的(EAGLE-3 把 draft 和 target 全堆在 GPU 上,共享 KV)。Mirror-SD 用 8 张 GPU + 8 颗 NPU(Apple M2 Ultra 系统),分工:

SPD 这件事关键到不能含糊。普通 TP 是每一层都要 all-reduce,因为权重是按 hidden dim 切的。SPD 把 draft 的 N_d 层切成两段,段内每个 NPU 跑全部层(用自己那份权重 shard),段间才同步。代价是每个 NPU 都要装一份权重,但 draft 才 0.6B,16 层 × 256 hidden / shard 完全装得下。

跨设备 token channel:每 step 两次小消息——M_t(early exit)和 final verify 触发——都是 O(Bκ) 个 token IDs + log-probs(bf16),通过 GPU→CPU→NPU 拷贝走。论文说微秒级,几乎可忽略。

6 · Latency 公式 + 灵敏度

把上面三件事合到一个 step 的时间公式里:

记 Δ ≜ T_target^ℓ_e+1:N(overlap budget,target 后半段时间)。

所以本作的全部工程目标可以用一句话总结:把 draft 时间塞进 Δ 这个时间窗里。方法 B (SS) 缩 T_draft;方法 C (SPD) 把 draft sync 从 2N_d 砍到 2;ℓ_e 选在 N/2 让 Δ 最大化但又保证 early-exit 信号还够准。

6.1 灵敏度小表

论文实测 sweet spot:κ=8, ℓ_e = N/2, γ=7。这个组合在 SpecBench 6 个域 (Humanities/Math/Roleplay/STEM/Coding/Math-reasoning) 上 fallback frequency 都落在 0.1–0.3,Ω_κ 接近饱和。

7 · Worked example:一条序列穿过双 pipeline

把 §3.3 的"积分"例子完整跑一遍,把所有计时也填进去。Qwen3-32B,N=64,ℓ_e=32,γ=7,κ=8,M2 Ultra GPU+NPU 集群 (×8)。

时刻 0(target step t 开始): commit 到 prefix "微". 假设 target 单层 ≈ 0.4 ms。

对比 vanilla SD 同 prefix:需要 draft AR 7 步(7×3.5≈24.5 ms)+ target 整 forward(64×0.4=25.6 ms)= 50 ms,假设接受 2 个产出 3 token,16.7 ms/token。≈1.9× 加速,与 Math 域报的 5×–5.8× 相比稍低,因为这只是 single step 的简化估计;实际增益还要乘上"reuse 把下 step T_draft 也省了"那部分。

8 · 实验关键结果 + 30% 拆分

8.1 SpecBench 主表(摘录)

"30% 平均"是几十个 task × temperature × model 的 micro-average。从分布看:长输出、高确定性的任务(数学、翻译、长 reasoning)增益最大(高 acceptance + 长 window 都能稳定 hide draft);高熵任务(多轮对话、QA)增益最低(频繁 fallback 把 NPU 计算白白浪费)。

8.2 30% 增益拆分(我的解读)

论文没给一张正经的 ablation table 把 30% 切成 (cross-device + mirror + early-exit + SS) 四个分量,但从 Fig 4 (OPT/Mistral 上的 4-way) 和 Fig 6 (SS 单独 speedup) 能反推:

8.3 Fallback 曲线(Apdx E.4)

论文给了 6 个域的 FF (fallback frequency) vs κ × ℓ_e 曲线:

• κ 从 1 到 128:FF 单调降,κ=8 时已经在膝盖处,16+ 收益边际递减——支持 κ=8 的默认。
• ℓ_e ∈ {1/4, 1/2, 3/4} N:深一点 fallback 更低,但论文最终选 1/2 是因为要兼顾 Δ(深一点 Δ 太小,draft 钻不进去)。
• Math/Math-reasoning 域 FF 最低(0.1 左右),Multi-turn / Roleplay 最高(0.3+)——直接对应 SpecBench 表里这俩 domain 的加速差异。

8.4 Batching 影响(Apdx D.1)

B 从 8 → 128:vanilla SD 加速从 2.84× 跌到 1.77×;Mirror-SD 从 3.05× 跌到 1.95×。两者都跌(因为 target verify 越来越主导),但 Mirror-SD 始终领先。论文的对策是随 B 增大动态缩 κ(κ=8 → 4 → 2 → 1)以控制 draft 端 compute——这意味着大 B 部署时 mirror 优势会变小,但仍存在。

9 · 与同类工作对比

立场总结: 90% 的 SD 改进论文都在"单一加速器内"挖榨;Mirror-SD 把 axis 加到"跨加速器异构 + 早退路标 + 双向校正"。代价是需要 NPU——这是它的硬件假设,也是它的硬伤。

10 · 局限 / 个人 take / 待验证问题

论文承认的局限

• 大 batch 下 mirror 优势缩小(Apdx D.1);本质是 draft 也变成 compute-bound,overlap 窗口被吃掉。
• 需要异构硬件;论文用 Apple M2 Ultra(GPU+ANE),换到 H100-only 集群没法直接用——你要么虚拟出"两组 GPU",要么完全没法跑 mirror 这部分。
• Apdx F.1 透露,跨 SoC 用 Thunderbolt 5 (120 Gbps)互连,8 节点拓扑是纸糊的:groups of 2,group 间靠 fabric。Inter-group sync 论文说"sub-millisecond",但具体怎么 schedule 没展开,扩展性堪忧。
• Lossless 证明依赖"draft 在 verification 时呈现的分布与 vanilla SD 同分布"——Apdx B Eq.12。这要求 branch 选择策略是 conditional law-preserving 的;论文用"按 v_i 概率加权"暗示了一种实现,但没证全。

我的疑问

1. 真无 NPU 能复现多少?把 NPU 替换成"target 同集群里的另一组 GPU + DGL 切到 PCIe 通道",overlap 还在,但带宽完全不同。希望看 H100-only ablation。
2. RL post-training 的 hidden states drift(见 08 号笔记)对 mirror 多敏感? 早退用的是 target 中间层,如果 RL 让中间层漂移加快(早期 fine-tune 后 logit lens 退化),Ω_κ 会跌得多猛?
3. 30% 拆分中"branch-complete reuse"那 10% 是否过高估?论文没单独 ablate,我从 fallback 0.1–0.3 推测的;实测 reuse rate 想看个真实曲线。
4. τ=0 vs τ=1 的差距并不剧烈,意味着 Mirror-SD 对 stochastic decoding 也稳。但 nucleus / top-p 没测,top-κ 在 ℓ_e 是硬切的——和 nucleus sampler 在 N 层组合可能违反 lossless。
5. γ=7、κ=8 是否对所有 14B–66B 都最优?论文说"是",但 OPT-66B 比 Qwen3-14B 大 5 倍,Δ 也大 5 倍,直觉上 OPT 应该可以承受更激进的 γ。
6. "Mirror" 的真正双向是不是被夸大了?读全文,target 主要在做 verify + 给 correction;draft 在做 branch-complete。论文修辞上说"target speculates correction paths for the draft",我看更像"target 给 single-token correction,draft 自己看在不在树里" — 这并不算 target 在主动 speculate。

记忆点

精读笔记 v1 · 2026-05-07 · 配套 PDF: /data/szhang967/papers/paper-notes/spec-rl/MirrorSD_2510.13161.pdf
相关:01 NeMo-RL × EAGLE-3 · 02 EAGLE evolution · 07 DFlash · 08 Hidden States Drift