/ AI Infra

[Infra-8] MTP: Multi-Token Prediction Modules

A serving-oriented guide to Multi-Token Prediction: its training objective, DeepSeek-V3 sequential MTP modules, speculative verification, KV cache commit, and scheduler integration.

Jul 14, 20264 minInfra · LLM Serving · Speculative Decoding · MTP

在大模型训练与在线 serving 里，Multi-Token Prediction（MTP） 有两层含义：它首先是一个让模型同时学习多个未来 token 的训练目标；当模型保留对应模块时，它又可以成为 speculative decoding 的原生 proposer。

这两个层次很容易被混成一句“模型一次生成多个 token”。更准确的说法是：MTP 让模型提出未来 token 候选；target model 仍然要验证这些候选，只有接受的 token 才能真正进入输出序列和 KV Cache。

本文从训练、模型结构和 serving 三条线解释 MTP，重点放在这条完整路径：

MTP proposal
    -> target verification
    -> acceptance / rejection sampling
    -> accepted token commit
    -> KV Cache update
    -> scheduler dispatches the next step

1. MTP 要解决什么问题#

传统自回归语言模型采用 Next-Token Prediction（NTP）：

p(x_{t+1} \mid x_{\leq t})

每次 forward 只预测下一个 token。decode 阶段的执行形式因此是：

context                         -> forward -> token_1
context + token_1               -> forward -> token_2
context + token_1 + token_2     -> forward -> token_3
...

即使单轮 forward 很快，token 之间仍然存在严格的时间依赖。对于短请求、小 batch 或受到调度开销限制的场景，这种串行性会直接限制 TPOT 和总吞吐。

MTP 的目标是在位置 $t$ 除了预测 $x_{t+1}$ 外，也学习更远的未来 token：

x_{t+1}, \quad x_{t+2}, \quad x_{t+3}, \quad \ldots

早期 multi-token prediction 工作采用共享主干加多个未来 token 输出头的形式。它报告了更高的样本效率，尤其在代码任务中较明显；当这些预测被用于验证式推理时，也可能减少自回归生成轮数。Gloeckle et al. 给出了这一类方法的代表性研究。

但“预测多个未来 token”不等于无条件地一次提交多个 token。训练目标、候选生成和 target verification 需要分开理解。

2. MTP objective 与 MTP modules#

2.1 MTP objective：额外训练目标#

普通 NTP 只在每个位置提供一个交叉熵监督。加入 MTP 后，可将训练损失写成：

\mathcal{L} = \mathcal{L}_{\text{next-token}} + \lambda \cdot \frac{1}{D} \sum_{k=1}^{D} \mathcal{L}_{\text{MTP}}^{(k)}

其中：

符号	含义
$\mathcal{L}_{\text{next-token}}$	主模型正常的 next-token loss
$\mathcal{L}_{\text{MTP}}^{(k)}$	第 $k$ 个 future token 的预测损失
$D$	MTP depth，即额外预测的未来深度
$\lambda$	MTP loss 的缩放权重

Megatron-Core 的 MTP 文档也采用了这个视角：对多个 prediction depth 的损失求平均，再乘以 mtp_loss_scaling_factor 形成额外训练目标。文档中的默认缩放因子为 0.1。

这个目标的价值不只在推理加速：

每个训练位置提供更密集的监督信号；
hidden state 需要编码更长视野的信息，而不只是局部 next token；
在代码、数学和结构化生成中，未来 token 与当前模式往往存在更强的可预测关系；
即使推理时丢弃 MTP 模块，主干模型也可能从额外训练信号中受益。

2.2 MTP modules：实现目标的额外结构#

MTP module 是为了产生这些额外预测而附加在主模型后的轻量结构。它通常不是另一套完整的语言模型，而是由 projection、额外 Transformer block、预测 head 等构成。

因此下面两句话分别描述不同事情：

MTP objective：训练时额外监督哪些 future tokens。
MTP modules：模型用什么额外结构产生这些预测。

如果模型只把 MTP 当作训练期辅助头，部署时可以直接移除模块；如果部署时保留模块，它们就可以参与 speculative decoding。

3. Parallel-head MTP 与 sequential MTP#

3.1 最直接的 parallel-head MTP#

一种直观结构是在主干 Transformer 的同一个 hidden state 上接多个独立 head：

main transformer trunk
        |
    hidden state h_t
        |
+-------+-------+-------+
|       |       |       |
head_1  head_2  head_3
|       |       |       |
x_t+1   x_t+2   x_t+3

它共享主干计算，分别预测多个未来位置。优点是结构和训练实现相对简单。

限制也很明确：预测 $x_{t+2}$ 或 $x_{t+3}$ 时，head 并没有显式经过 $x_{t+1} \rightarrow x_{t+2}$ 的生成链。它更像“从当前位置并行猜测未来若干 token”，而不是模拟完整的自回归因果推进。

3.2 DeepSeek-V3 的 sequential MTP modules#

DeepSeek-V3 技术报告使用的是 sequential MTP modules。它强调在预测更远 token 时保留完整 causal chain，而不是让多个独立 head 从同一个 $h_t$ 并行猜测。

把主模型的隐藏状态写成 $h_i^0$ 。第 $k$ 个 MTP module 接收上一个 depth 的表示 $h_i^{k-1}$ 和对应 future token 的 embedding，经投影与一个额外 block 后产生 $h_i^k$ ，再通过共享 output head 预测下一位置：

main hidden: h_i^0
        |
        | + embedding(x_i+1)
        v
MTP module 1 -> h_i^1 -> predict x_i+2
 
        | + embedding(x_i+2)
        v
MTP module 2 -> h_i^2 -> predict x_i+3
 
        | + embedding(x_i+3)
        v
MTP module 3 -> h_i^3 -> predict x_i+4

抽象地说，第 $k$ 层会形成：

\tilde{h}_i^k = \operatorname{Proj}\left([h_i^{k-1}; e_{i+k}]\right)

h_i^k = \operatorname{MTPBlock}_k(\tilde{h}_i^k)

p(x_{i+k+1} \mid \cdot) = \operatorname{Softmax}(W_{\text{out}} h_i^k)

其中 $e_{i+k}$ 是 token embedding， $W_{\text{out}}$ 是共享 output head。DeepSeek-V3 的每个 MTP module 还共享 embedding 与 output head，从而避免为每个 depth 复制最大的词表参数。

这带来了一个关键区别：

parallel heads：同一 h_t 同时猜多个 future tokens。
sequential modules：按 depth 推进表示，并保留 future token 的因果链。

4. 训练时到底监督什么#

假设训练序列为：

x_1, x_2, x_3, x_4, x_5, ...

普通 NTP 在 position $i$ 的任务是：

main model: predict x_i+1

加入深度为 $D$ 的 MTP 后，可以额外监督：

main model:  predict x_i+1
MTP depth 1: predict x_i+2
MTP depth 2: predict x_i+3
...
MTP depth D: predict x_i+D+1

每个 depth 都会计算 cross entropy。对一个 batch，必须注意不同 depth 的有效 token 数不同：序列尾部没有足够 future token 的位置需要用 loss mask 排除。

main_loss = cross_entropy(main_logits[:, :-1], input_ids[:, 1:])
 
mtp_losses = []
for depth, mtp_logits in enumerate(mtp_logits_by_depth, start=1):
    # depth=1 predicts x_{i+2}; the tail has fewer valid labels.
    labels = input_ids[:, depth + 1:]
    logits = mtp_logits[:, : labels.size(1)]
    mtp_losses.append(cross_entropy(logits, labels))
 
loss = main_loss + mtp_loss_scaling_factor * mean(mtp_losses)

真实实现还需要处理 packed sequences、attention mask、不同序列边界与并行切分。不过核心不变：MTP 是主目标之外的训练信号，不应意外覆盖或替代主模型 NTP loss。

5. 推理时的两种使用方式#

5.1 只把 MTP 当作训练增强#

最保守的部署方式是直接丢弃 MTP modules。主模型仍按普通自回归路径工作：

1 forward -> 1 sampled token

DeepSeek-V3 技术报告明确说明，MTP 的主要目标之一是改善主模型训练；即使不使用 MTP modules，主模型也能独立完成正常生成。

这是一个有用的工程分界：训练收益不必绑定推理路径复杂度。如果 serving 后端尚未实现 native MTP，先按普通模型部署并不影响主模型可用性。

5.2 把 MTP module 当作 speculative proposer#

启用加速时，MTP module 不直接替代 target model，而是先提出 bonus token，再由 target model 验证：

1. target path obtains the current model state
2. MTP module proposes one or more bonus tokens
3. target model verifies the proposal in a multi-token forward
4. accepted tokens are committed; rejected suffixes are discarded or resampled

vLLM 的 MTP 文档将其定义为 native multi-token prediction speculation：因为 target model 自身包含 MTP 能力，不需要额外提供单独的 draft model。

一个典型配置为：

speculative_config = {
    "method": "mtp",
    "num_speculative_tokens": 1,
}

从 1 开始很重要。它既便于观察 acceptance rate 和图执行形状，也符合一些公开权重仅提供单层 MTP module 的现实。

6. 为什么 MTP 可能更快#

普通 decode 的理想化吞吐关系是：

1 target forward -> 1 committed token

若 MTP 提出的一个 bonus token 被 target model 接受，则同一次验证路径可以推进更多 token：

1 target verification -> token_1 + accepted bonus token

可以用每轮接受的平均 token 数来描述收益：

E[\text{committed tokens per step}] = 1 + E[\text{accepted bonus tokens}]

但端到端 speedup 不是上式本身，还要减去 proposer、验证、采样和缓存管理成本：

\text{speedup} \approx \frac{\text{baseline decode cost}} {\text{proposal cost} + \text{verification cost per committed token}}

DeepSeek-V3 报告中给出的第二个 token acceptance rate 在不同生成主题上约为 85%--90%，并报告约 1.8 倍的 TPS 提升。这里应把它理解为模型、实现和 workload 共同决定的结果，而不是所有 serving 环境都能复制的固定数字。

实际测量至少应同时记录：

TTFT
ITL / TPOT
output TPS
acceptance rate
accepted speculative tokens per step
decode batch size distribution
GPU utilization
KV Cache block usage

只有 TPS 上升但 TTFT、尾延迟或缓存压力变差时，MTP 并不一定改善真实线上体验。

7. Serving 路径：proposal、verification 与 commit#

从 serving 角度，MTP 最关键的部分不是多算一次 logits，而是保证 speculative token 在请求状态机中的生命周期正确。

7.1 Proposal#

当前请求已有已提交上下文和 target KV Cache：

committed context + target KV Cache
              |
              v
         MTP proposer
              |
              v
      draft token(s) and draft state

proposer 的输出只能被视为候选。它不能直接更新请求的 committed token list，也不能把自身临时状态误当作 target KV Cache。

7.2 Verification#

target model 将已提交上下文和候选 token 放进验证路径。对于 greedy decoding，可以比较 target token 与 draft token 是否一致；对于随机采样，需要使用保持目标分布正确性的 acceptance / rejection 规则。

candidate:      [token_1, draft_token_2]
target result:  verifies token_1 and scores draft_token_2
                                      |
                                      v
                         accept or reject the bonus token

关键不变量是：最终输出的分布必须由 target model 定义，而不能因为 proposer 直接提交 token 而被改变。

7.3 Accepted token commit#

真正进入请求状态的只有 target verification 后被接受的 token，以及它们对应的 target-model KV：

accepted token prefix
      |
      +-> append to output token ids
      +-> update sequence length
      +-> commit target KV blocks
      +-> release temporary proposal state

若 bonus token 被拒绝，系统必须只提交已验证的前缀，并丢弃被拒绝后缀对应的临时状态。常见错误包括：

接受 token 后没有提交正确的 target KV；
拒绝后仍保留了候选 token 的 KV 或 sequence length；
batch 内请求的 speculative depth 不同，却使用了错误的 padding 或 attention mask；
假设每轮只产生一个 token，导致 scheduler token budget 计算错误；
图执行的捕获形状与验证 token 数不匹配。

7.4 Scheduler 要看到“本轮推进了多少 token”#

continuous batching 的调度器通常为每个请求维护 token budget、已处理长度和 KV block 占用。MTP 改变的是：一次 decode step 不再必然只提交一个 token。

ordinary decode:      request A commits 1 token
MTP accepted bonus:   request A commits 2 tokens
MTP rejected bonus:   request A commits 1 token

因此 scheduler、block manager 和请求状态都必须从 verification 结果读取实际 committed token 数，而不能只根据“本轮执行过一次 decode”进行固定加一。

8. KV Cache 是 serving 实现的边界#

MTP 常被误解为“proposer 算出了未来 token，所以直接把 cache 往前推进”。这不正确。

从 target model 的语义看，只有验证并接受的 token 才对应能够长期保留的主模型 KV。系统至少应区分：

proposal state
target-model candidate KV
committed target KV
rejected suffix / rollback scope

这一区分会影响分页 KV Cache 的 block 分配策略。例如，当一个请求申请额外 speculative slots 时，block manager 可以先保留临时容量；验证后：

accept -> promote or retain the required blocks
reject -> free the uncommitted suffix blocks

在 prefix caching、PD 分离或跨实例路由存在时，还要确保缓存 key 和请求长度反映的是已提交前缀，而不是尚未验证的 draft suffix。

9. MTP、draft model、Medusa 与 EAGLE#

这些概念经常同时出现在 speculative decoding 讨论中，但工作位置不同：

方法	proposer 来源	与 target model 的关系	主要用途
Draft model	额外小模型	独立模型提出候选	推理加速
Medusa	主模型后多个 head	多头或 token tree 候选	推理加速
EAGLE	draft feature / head 路径	以特征预测提出候选	推理加速
MTP modules	target model 原生模块	训练目标的额外模块，也可提出候选	训练增强 + 可选加速

Speculative decoding 是上层算法框架：先 propose，再 verify。MTP、draft model、Medusa 或 EAGLE 是不同的 proposer 实现。

DeepSeek-V3 的 sequential MTP modules 与多个独立 Medusa-style head 的差异在于，它按 depth 维护了更完整的因果链。与 EAGLE 的相似点是都重视候选质量；但 DeepSeek-V3 将 MTP 同时作为训练增强目标，而不仅是 serving 加速器。

10. 模型权重与部署边界#

原生 MTP 不是给普通模型加一个配置项就能获得的能力。模型必须在训练时包含对应目标和模块，权重中也必须保留模块参数。

以 DeepSeek-V3 为例，权重说明中使用 num_nextn_predict_layers 表示 MTP module 数量；公开说明将模型权重划分为 main model weights 与 MTP modules 两部分，并描述了共享 embedding、RMSNorm、projection、额外 Transformer block 与共享 output head 的组成。DeepSeek-V3 权重说明可作为检查权重结构的入口。

这带来两条部署结论：

普通没有原生 MTP module 的模型，不能仅通过 method: "mtp" 凭空得到 MTP；应选择支持的 draft-model 或其他 speculation 方法。
模型公开了几层 MTP module，serving 配置就应先限制在相应深度内。更大的 speculative depth 既未必有权重支持，也不一定有更高 acceptance rate。

11. Benchmark 不应只看单请求速度#

MTP 的收益与 workload 强相关。单请求、低并发、短输出可能更容易看到减少 decode round 的效果；高并发场景中，target forward 已较饱和时，额外 proposal 和 verification 管理可能会吞掉部分收益。

建议至少按以下维度切分压测：

维度	建议取值
输入长度	短 prompt、中等上下文、长上下文
输出长度	短输出、中输出、长输出
并发	单请求、低并发、饱和并发
Sampling	greedy、低温度、top-p / top-k
Speculation	off、MTP depth 1、模型支持的其他深度
指标	TTFT、TPOT、TPS、acceptance、显存与 KV block 使用率

同时保留基线是必要的：

baseline:       ordinary decode, no speculation
comparison:     native MTP speculation
same model:     same weights, same precision, same batch policy
same workload:  same prompt/output distribution and sampling parameters

否则很容易把缓存命中、请求分布变化或调度策略差异误认为 MTP 本身的收益。

12. 简化伪代码#

下面伪代码省略了 batch、EOS、分页块和随机采样的细节，重点展示 commit 边界。

普通 decode#

while not request.finished:
    logits, next_kv = main_model.forward(request.token_ids, request.kv_cache)
    next_token = sample(logits, request.sampling_params)
 
    request.token_ids.append(next_token)
    request.kv_cache = commit(next_kv)

MTP speculative decode#

while not request.finished:
    # Proposals are candidates, not committed output.
    target_state = main_model.prepare_decode(request.token_ids, request.kv_cache)
    draft_tokens, proposal_state = mtp_module.propose(target_state)
 
    # The target model determines which prefix is valid.
    verified = main_model.verify(
        request.token_ids,
        request.kv_cache,
        draft_tokens,
        request.sampling_params,
    )
 
    accepted_tokens = verified.accepted_prefix
    request.token_ids.extend(accepted_tokens)
    request.kv_cache = commit(verified.target_kv_for(accepted_tokens))
 
    release(proposal_state)
    release(verified.rejected_suffix_state)

真实系统通常还需要将 accepted_tokens 的数量回写给 scheduler，并在下轮调度前更新 token budget、sequence length、KV block 表和流式输出缓冲区。

13. 总结#

MTP module 本质上是一组模型原生的未来 token 预测模块：

训练时：提供更密集、更长视野的监督信号。
推理时：作为 speculative proposer，争取一次验证接受更多 token。
服务时：必须正确处理 verification、commit、rollback、KV Cache 和 scheduler 状态。

对于大模型 serving，真正决定 MTP 是否可靠、高效的不是“额外输出了几个 logits”，而是这条状态链路是否完整：

proposal
-> target verification
-> accepted token commit
-> target KV Cache update
-> scheduler observes actual progress

如果其中任意一环把候选 token 当成已提交 token，或者让 KV Cache、序列长度和 scheduler budget 脱节，性能收益很容易变成隐蔽的正确性问题。