MTP 投机解码与预取

decode 是磁盘受限的——每个 token 都要付 ~11 GB 的专家读。投机解码（speculative decoding）的思路是：用一个便宜的”起草”模型猜出接下来的若干 token，再用主模型在一次批量前向里同时验证——接受的部分等于免费多产出了几个 token。colibrì 用的是 GLM-5.2 自带的 MTP（多 token 预测）头，位于 layer 78。核心是 mtp_draft 和 spec_decode。

1. 投机解码为什么划算 

朴素 decode（每 token 一次前向）:
  token1 → forward → token2 → forward → token3 → ...
  每次 forward 付 ~11 GB 磁盘读

投机 decode（起草 G 个，一次验证）:
  起草：MTP 头猜出 [d1, d2, d3]
  验证：一次批量 forward 同时算 [真, d1, d2, d3] 的 logits
  接受：若 d1、d2 匹配主模型输出 → 一次 forward 产出 3 个 token
  → 有效成本大致减半（缓存暖后）

community 实测（int8 头）：接受率 39–59%，2.2–2.8 tokens/forward。

2. GLM-5.2 原生 MTP 头

MTP 头本身是一个完整的 DeepSeek-V3 式 transformer 层，位于索引 n_layers（GLM-5.2 = 78）。加载在 model_init（755–800 行）：只有当那一层的完整张量集在分片里都存在时才启用（MTP=0 可禁用）。它有自己的注意力投影、共享专家、router，加上 eh_proj（D×2D，790 行）、enorm/hnorm（791 行）、mtp_norm。它的 KV 占用 kv_alloc 里多出的那一行（n_layers+1），kv_start[n_layers]=-1 标记为 decode-only（797 行）。

相关 Model 字段：has_mtp, mtpL, eh_proj, mtp_prop, mtp_acc（133–136 行）——mtp_prop/mtp_acc 是提议数/接受数计数。

2.1 为什么头必须 int8

MTP 头的所有张量都以 dbits（稠密精度，int8）加载——qt_load(..., dbits)（776–790 行）。glm.c 的注释记着 “layer MTP int8 = 2x i layer int4”。原因在 03 详述：int4 头的 draft 几乎总猜错，接受率崩到 0–4%，投机永不启动；int8 头才有 39–59% 的接受率。这不是一个运行时 assert，而是通过把头留在稠密（int8）层而非 int4 专家层来实现的设计。

3. mtp_draft：起草

mtp_draft（1449 行）用 MTP 头自回归地起草最多 G 个 token。每一步走 DeepSeek-V3 的链式结构：

h' = Layer78( eh_proj[ enorm(emb(tok)) ; hnorm(h) ] )
next = argmax( lm_head( mtp_norm(h') ) )

关键行：embedding 上 enorm（1461 行）、final_norm 后 hnorm 作用于携带的隐状态（1462–1463 行）、concat [emb;h]（1465 行，可用 MTP_SWAP 交换顺序）、eh_proj 矩阵乘（1466 行）、MTP 层前向（1472 行）、mtp_norm+lm_head+argmax（1474–1476 行），然后链式 tok=t2; h=hx（1479 行）。

无 MTP 头时的回退是 ngram_draft（1430 行）：prompt-lookup，找最近一次出现的相同 bigram，提议其后续 token，零额外权重。

4. spec_decode：起草 + 验证主循环

spec_decode（1576–1626 行）：

1. 选起草源：has_mtp ? mtp_draft : ngram_draft（1596–1597 行）
2. step_all 一次批量前向，验证 [真 token + G 个 draft]（1602–1603 行）
3. 接受循环（1607–1615 行）
4. mtp_acc += 接受数（1616 行）
5. mtp_absorb 把接受的 token 吸收进 MTP 层 KV（1617 行）

mtp_absorb（1487 行）：验证后，把已接受的 (emb(token@pos+1), h_true@pos) 对在一次批量前向里吸收进 MTP 层的 KV——让下次起草保持因果一致。

4.1 自适应关闭

如果 has_mtp && mtp_prop>=24 && mtp_acc*10 < mtp_prop（接受率 <10%），设 g_draft=0（1589–1593 行）。这对应 README 的诚实警告：冷缓存下每个被验证的 draft 会路由到额外的专家（~660 → ~1100 expert-loads/token），投机可能是净的时间损失——所以有这个自适应 guard 和 DRAFT=0。

旋钮：g_draft（551 行）（DRAFT，有 MTP 时自动解析为 3，否则 0，2445 行）、g_spec（550 行）（SPEC，跨层投机预取）。

5. 拒绝采样：投机在采样下依然无损

投机解码最微妙的部分：在温度采样下也保持逐 token 无损。draft 是确定的（argmax → 点质量分布 q），但主模型是按分布 p 采样的。如果只是”匹配就接受”，输出分布会被 draft 扭曲。colibrì 用 Leviathan 拒绝采样 解决——设计注释在 1512–1516 行。

接受循环（g_temp>0，采样，1607–1615 行）:
  对每个 draft[k]:
    dist_build 出主模型分布 p
    以概率 p(draft[k]) 接受：accept = (rndu() < p[draft[k]])
    若拒绝：carry_ban = draft[k]
      → 下一个 pick_tok 从 p 里把被拒 token 清零、重新归一化后重采
  结果：输出分布【恰好是 p】，投机完全不可见

关键行：accept = (rndu() < g_pbuf[draft[k]])（1610–1611 行）、拒绝时 carry_ban=draft[k]（1612 行）、carry_ban 在下次 pick_tok 消费（1582 行）。贪心路径（g_temp<=0）则用精确 argmax 匹配验证（1609 行），天然无损。

这就是 README 那句”Lossless — and stays lossless under sampling via rejection sampling”的实现。

6. 与专家预取的关系

投机解码和 07 里的 pilot_prefetch（router-lookahead 预取）是互补的两条攻击 decode 磁盘瓶颈的路线：

  • 投机解码：一次前向多产出几个 token，摊薄每 token 的磁盘成本——但冷缓存下会引入额外专家读。
  • router-lookahead 预取：提前把下一层要用的专家读进页缓存，把冷读变成暖读。

两者都依赖同一个观察：MoE 的路由有可预测的结构。README 把它们连同”会学习的缓存”一起，作为把 decode 从”磁盘物理下限”往上推的三个杠杆。

7. 小结

机制 出处 作用
原生 MTP 头（layer 78） model_init:755 GLM-5.2 自带的起草头
头必须 int8 qt_load(dbits):776, 注释:99 int4 接受率崩到 0–4%
自回归起草 mtp_draft:1449 DeepSeek-V3 链式结构
起草+验证主循环 spec_decode:1576 一次批量前向验证
接受吸收 mtp_absorb:1487 保持起草因果一致
拒绝采样 spec_decode:1607 采样下依然无损
自适应关闭 spec_decode:1589 冷缓存下投机反而慢时放弃
This site uses Just the Docs, a documentation theme for Jekyll.