FastMTP: Accelerating LLM Inference with Enhanced Multi-Token Prediction

1 · 动机:已经训好的 MTP 为什么没人在 inference 时用?

1.1 历史脉络:MTP 的两个时代

Multi-Token Prediction 由 Gloeckle et al. 2024 提出(见 14_MTP_Gloeckle 的精读),最初是把它当成一种训练辅助任务:在主头之外挂 n 个独立的 head,各自负责预测 t+1, t+2, ..., t+n。其论点是 "提供更密集的监督信号 + 鼓励 planning ahead",主要好处兑现在训练时——下游任务质量轻微提升、收敛更快。

到 DeepSeek-V3 把它进一步演化成串行 MTP (sequential MTP):n 个 module 之间有 cross attention 形成因果链,每个都吃前一个 module 的 hidden state,而不是 Gloeckle 的并行独立设计。这一改保住了 autoregressive 性质,看起来"顺手就能拿来做 speculative decoding 的 draft"。

但实际部署中,几乎所有开放权重的 MTP 模型(DeepSeek-V3、GLM-4.5、MiMo)都做出了同一选择:

• 开源时只放一个 MTP module(明明训了多个)
• SGLang / vLLM 等 serving 框架的 MTP 路径只支持 K=1
• K≥2 时性能反而崩

这就是 FastMTP 看到的"商机":这些已被付费训练的 MTP 模块,本来应该能加速 inference,但实际拿来一用就崩,作者要回答崩在哪、怎么修。

1.2 别的方案为什么都"够呛"

1.3 为什么这事不平凡 — train-inference mismatch 的根源

把 vanilla MTP 拿来递归用 K 次,本质上是 exposure bias 的极端版本。考虑训练时第 k 个 module 看见什么:

那 train 时为什么不直接做 scheduled sampling / use ĥ_{k-1}?——因为 vanilla MTP 的设计目标是训练辅助,不是为 inference 优化的。每个 module 各管各的位置,就像 Gloeckle 的"n 头并行各预测一个 future token"——根本没想过递归。

FastMTP 的本质修补就两点:

1. 架构层:把 K 个独立 module 改成同一个 head 在 K 个位置上反复用 → 强迫它在权重里融合"对不同 lookahead 距离都 robust"的能力
2. 训练层:显式让训练时的输入分布 = 推理时的递归输入分布,不再单纯 teacher-force

2 · 背景速查

2.1 关键术语

2.2 EAGLE-style verify 流程速记 (本文用的就是这个)

1. 主模型 F 跑一次 prefill / 当前 step,得到 ĥ_i 和下一 token ĥ_{i+1}(这一个总是被接受)
2. MTP head M 递归 K 次,生出 ĥ_{i+2}, ĥ_{i+3}, ..., ĥ_{i+K+1}
3. 主模型 F 一次 forward 处理这 K+1 个位置,得到每个位置的 p_target
4. 从左到右逐位置匹配 draft token vs argmax p_target;greedy 模式下"匹配 ≡ 接受",一旦不匹配就停,补一个 p_target 重采样的 token
5. 本步产出 1..K+1 个 token

本文实验 greedy + temperature=0,所以"接受"就是 draft 与 target argmax 严格相等;非 greedy 情况标准 rejection sampling 一样适用,lossless 性质由 Leviathan 2023 / Chen 2023 证明。

3 · 方法① · Position-shared 单 head 架构

3.1 与 Gloeckle / DeepSeek-V3 的对比图

3.2 为什么 weight share 反而能涨 acceptance(反向论证)

"weight share"在 ML 史上往往是表达力的限制(每个 head 只能学同一个 mapping)。这里为什么能涨 acceptance?

1. 训练 signal 在 K 个位置上叠加:同一份 θ 同时受到"预测 t+1"、"预测 t+2"、"预测 t+3"三种 loss 拉扯,被迫学一个对 lookahead 距离无关的"下一 token mapping",而不是 over-fit 到固定位置。
2. 训练分布与推理分布严格对齐:推理时 M 在第 k 步看到的是 ĥ_{k-1}(自己生成的);训练时也按这个递归方式喂数据(下一节展开)。两边输入分布 i.d.d.。
3. 推理时显存常数:1 个 KV cache、1 份权重。哪怕 K=7,显存与 K=1 相同,这让更长 draft变得可负担。

论文 §3.4 的 ablation 给的硬数据(详见 §8):同样 fine-tune,但 fix-data(用源数据集 response)的 K=3 acceptance length τ=2.54;换成 self-distilled 后 τ=2.73。再加 FR vocab compression τ=2.66 但 token/s 涨。这说明 share weight 是必要但不充分条件 — 配合 self-distill 才能解锁全部潜力。

4 · 方法② · Self-distilled 训练数据 — fine-tune 的核心 trick

4.1 数据流程图

4.2 Self-distill vs Fixed-data 的硬数据 (论文 Table 1, K=3)

从 Vanilla MTP → Fixed-data FT,τ 从 1.83 涨到 2.54(+39%);单单换数据(Fixed → Self),τ 从 2.54 → 2.73(+7%),speedup 从 1.67× → 1.81× (+8%)。看似不大,但乘进 1.81 → 2.03 这一步,self-distill 是必要的踏板。

4.3 反向思考 — 不 self-distill 会怎样?

假设你直接用原数据集的 (x_n, y_n) 来训 MTP head:

• MTP head 学到的"下一 token"是 y_n 的下一 token,即原作者的写法
• 但推理时 main model F 会用自己的分布写 response,例如更喜欢 "Let me think step by step" 而 y_n 是 "First, "
• MTP head 草稿出 "First, ",F verify 时发现自己想写 "Let me",reject
• 哪怕 F 和 y_n 写法在语义上等价,token-level mismatch → acceptance 跳水

Self-distill 的本质就是把"语义等价"提升到"token 级一致"。这是 EAGLE 系列、Medusa 都用过的标准 trick(Cai 2024、Yang 2024c),只是 FastMTP 把它显式应用到 MTP fine-tune 上。

5 · 方法③ · Language-aware Dynamic Vocabulary Compression

5.1 问题:MTP head 的 vocab projection 是大头

MiMo-7B-RL 的 vocab size 是 152K(用 Qwen2 tokenizer)。MTP head 每生成一个 draft token,要做 logits = h · W^T,其中 W ∈ ℝ^{152K × 4096}。对于一个 7B 模型,这个 projection 在 draft 阶段(短序列、batch=1)不可忽略——FR-Spec (Zhao 2025a) 已经指出过这是个瓶颈。

5.2 机制图

5.3 选 vocab size 的实验数据 (Table 2)

反直觉结论:中文用 更激进的 16K 压缩,英文反而需要 32K。论文解释是:中文 token 的 head distribution 更集中(很多高频字 + 复合 token),英文 BPE 的长尾更长。

5.4 反向思考 — 为什么不直接用 top-K 的全局 vocab?

原 FR-Spec 用 SlimPajama (英语为主)统计 top-K → 中文 task 上 acceptance 跳水。FastMTP 的修补就是按语言切多个 V_high,运行时 detect 切换。这其实是个朴素工程,但效果显著——提醒我们 vocab compression 的"高频"要 conditional on 输入分布。

6 · 公式 · 加权 cross-entropy 损失

6.1 损失定义

其中第 k 步的权重按指数衰减:

6.2 代入数值

6.3 物理直觉

权重指数衰减,代表"越远的 token 越不可能预测准,降低它的损失贡献"。如果不衰减(β=1),三个步骤等权——但第 3 步的 CE loss 自然就大(理论下界 = data entropy + 累积 mismatch),会主导梯度,反而让第 1 步没学好。β=0.6 把第 3 步权重压到 18.4%,让"近期对、远期凑合"成为优化目标。

6.4 β 的敏感度(论文未做完整 sweep, 反推)

7 · Worked Example · 一个数学题片段在 K=3 上的轨迹

设输入 prompt 是 "求 sin²(x) 的不定积分。",main model 已经写到 "sin²(x)dx = (1−cos(2x))/2 dx, 所以 积分"。下一步 main model 会预测 "等于"。

FastMTP 的一次 step:

7.1 期望产出推导

记 a_k = step k 的条件接受率,τ = 1 + Σ_{k=1}^{K} ∏_{j=1}^{k} a_j。代入 FastMTP 的 (0.81, 0.56, 0.36):

论文实测 K=3 mean τ=2.73,略高,因为某些任务(数学)acceptance 显著偏高(数学第 3 步 acceptance >50%)。

同样代入 vanilla MTP (0.70, 0.11, 0.02):

实测 1.83 接近。这就是 vanilla MTP 在 K=3 上 speedup 反而退回到 1.21× 的根本原因——多走的 step 几乎都不接受,白付 draft 计算。

8 · 实验关键结果

8.1 主结果 (Table 1, MiMo-7B-RL on A10 24GB, 7 任务平均)

看点:

• Vanilla MTP 在 K=2→K=3 时 speedup 下降(1.31× → 1.21×): draft 计算白花了
• FastMTP 在 K=1→K=3 单调上升,符合"训练 reward 多步预测"
• Math 上 K=3 speedup 最高(单任务 2.31×, math τ=3.07)
• QA 上最低(1.84×),非结构化文本 acceptance 自然低

8.2 Acceptance rate breakdown (Figure 4a) — 全篇最 killer 的图

这张图的 take-away:vanilla MTP 不是"draft 不行",而是"递归不行"。修补方式不是换更强的 head,而是教它"被反复递归用"——FastMTP 的 fine-tune 改变的不是表达力,而是训练-推理输入分布对齐。

8.3 Optimal draft length (Figure 3, on A100)

论文额外训了一个 K=7 的 head:

K=3 是 sweet spot。τ 单调涨但 token/s 在 K=4 就开始降——边际 acceptance 抵不过递归 forward 成本。

9 · 与同类工作对比

9.1 vs EAGLE-3

FastMTP 与 EAGLE-3 在形式上非常像:都是单个轻量 head + autoregressive 递归 + self-distill 数据。差别在 head 来源:

• EAGLE-3: 从零训一个 transformer + projection,要选 hidden state 哪几层、要重新对齐 vocab
• FastMTP: 直接复用模型预训练就有的 MTP head 的架构和初始化,只 fine-tune;serving 框架已经支持 EAGLE-style verify,所以"假装是 EAGLE-3" 跑得通

论文措辞:"ensures compatibility with EAGLE-style speculative decoding for seamless integration with existing inference frameworks such as SGLang"。

9.2 vs DeepSeek-V3 sequential MTP

V3 部署时一般只用 first MTP module 做 K=1,因为加载多个 module 需要框架特别支持。FastMTP 的视角是:既然多 module 部署难,就把 K 个权重合并成一个(用 share 约束),同时通过 self-distill fine-tune 把它训得递归 robust。从模型容量上确实失去了一部分(K 份独立 weight 退到 1 份),但部署上换来了 EAGLE-style 框架原生支持。

9.3 vs Medusa

Medusa 的 head 之间不递归:K 个 MLP head 各看 main 的同一个 hidden state,各自预测 t+1, t+2, ...。所以 Medusa 不会出现 FastMTP 担忧的 "递归输入 OOD",但反过来也吃了 t+k 没法 condition on t+1, ..., t+k-1 的亏 — 因此 Medusa 通常需要 tree-attention 来弥补。FastMTP 是另一个极端:严格 sequential,但靠 share weight + 训练对齐保证递归不崩。

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

论文承认或暗含的局限

• 只在 MiMo-7B-RL 上验证。MTP 模块本身得是预训练就送的,模型必须有这个 head 才能套 FastMTP。Llama / Qwen3-Dense 没有原生 MTP,套不了——只能上 EAGLE-3。
• K=3 是 A10 单 batch 下的甜点,生产 batch=64+ 时 draft 开销与 verify 开销比例会变,K 甜点可能后退到 2(但论文未给 batch sweep)。
• 所有实验 greedy + temp=0。RL/创造性场景需要 sampling,acceptance 会下降(rejection sampling 会拒一部分概率不匹配 token)。
• 语言只测了中英二选一。多语言混合(code + 中 + 日 + ...)如何选 V_high?context 切换频繁怎么办?

个人 take

• 这是一篇"工程把 MTP 这个已付费的赠品兑现"的论文。想象一下买了 SaaS pro 但没用一半功能的味道。
• 核心 insight 不算新——EAGLE 系列早就用 self-distill 训过 draft——但把 self-distill 应用到 share-weight MTP head fine-tune是干净的组合创新。
• "vanilla MTP K=3 反而比 K=2 慢" 这个数据应该被钉在每个 RL/inference 工程师的工位上。盲目加 K 是最常见的本能错误。
• FR-Spec 那部分相对 minor — 1.81× → 2.03× 大约 +12% 的 speedup,但代价是引入"语言切换"的工程复杂度。如果只追求工程简洁,不上 FR 也已经 1.81×。

我会想验证的问题

1. K=3 的甜点会不会在大 batch 下后退?vLLM/SGLang 多并发场景下 draft 的固定开销摊薄了,理论上 K 甜点会前移(不是后退),但需实测。
2. 把 share-weight 改成 LoRA-per-step(K 个 LoRA adapter,共享 base)能否进一步涨 acceptance 而几乎不增显存?这是一个介于"K 独立 module"和"完全 share"中间的折中。
3. MiMo-7B-RL 本身已是 RL 后训过的模型;在没有 RL 后训的 base model 上(纯 pretrain)做 FastMTP 效果会不会差很多?数学上的高 acceptance 是不是模型本身被 RL 训得"模式化"了?
4. 对 long-CoT(>4k)长文本的 acceptance 衰减如何?论文 max-length=1024,实际推理场景动辄数万 token,KV growth + 长上下文的 token 多样性可能让 acceptance 单调下降。
5. 在 sampling 模式(temp>0)下,FR-Spec 的 vocab 截断会不会让"被截断的低概率 token"在 rejection sampling 时永远拿不到——是不是技术意义上仍然 lossless?(论文说 verify 用 full vocab 所以 lossless,但 draft 提议本身受截断影响,会改变 acceptance distribution。)

记忆点 (cold recall)

精读笔记 v1 · 2026-05-07 · PDF: /data/szhang967/papers/paper-notes/models/FastMTP_2509.18362.pdf