vLLM MLP 投机解码（MLP Draft Models）实战指南：配置、原理与加速器模型

投机解码（Speculative Decoding）允许用一个小型"草稿模型"（draft model）快速预生成若干候选 token，再交给目标大模型一次性并行校验，从而在保证输出分布不受影响的前提下显著提升解码吞吐。在 vLLM 中，有一类特殊的草稿模型——MLP 加速器（MLP Speculator / "accelerator"），它不使用自回归小语言模型，而是用一组轻量 MLP 层同时预测未来多个 token。本文以 docs/features/speculative_decoding/mlp.md 为骨架，结合仓库内配置与模型实现源码，完整讲解如何在 vLLM 中启用 MLP 投机解码、配置项含义、预训练加速器模型清单以及底层工作原理，帮助读者直接照做即可跑通并理解其取舍。

什么是 MLP 草稿模型：同时预测多个未来 token

传统投机解码（如 n-gram、独立小型 LM）中，草稿模型逐 token 自回归地生成候选序列；而 MLP 类型的加速器采用完全不同的思路：它以目标大模型某一层的上下文向量（context vectors / hidden states） 作为输入，在同一时刻并行输出未来多个位置的 token 预测。

具体到 vLLM 的实现，这类模型的出处是 IBM Research 发表于论文 Accelerating Production LLMs with Combined Token/Embedding Speculators（arXiv:2404.19124）的技术方案，同时可参考 PyTorch 官方博客 The Hitchhiker's Guide to Speculative Decoding 一文理解投机解码的通用背景（两者均为原文档中给出的参考资料）。其核心思想可概括为：

• 条件双重来源：草稿预测同时依赖于"已采样到的 token"（sampled tokens）和"目标模型提取的上下文向量"（context vectors），因此草稿质量远高于纯 n-gram 等无条件方案；
• 多头一步前向：不同 stage 的轻量 MLP 层分别负责预测第 1、第 2、第 3…… 个未来 token，一次前向即可产出多个候选，而非逐 token 展开；
• 与目标模型可解耦：加速器权重独立于目标模型发布（典型 HF 命名空间为 ibm-ai-platform 与 ibm-granite），可按需挂载到任意兼容的 Llama/Granite 等目标模型上。

在 vLLM 中，这一方法被标记为 mlp_speculator，是 SpeculativeConfig 中根据草稿模型 model_type 自动识别的投机方法之一。

快速上手：一个完整的 MLP 投机解码示例

原文档给出了可直接运行的离线推理示例。在 vLLM 中以目标模型 meta-llama/Meta-Llama-3.1-8B-Instruct、草稿加速器 ibm-ai-platform/llama3-8b-accelerator 组合为例：

from vllm import LLM, SamplingParams

prompts = ["The future of AI is"]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(
    model="meta-llama/Meta-Llama-3.1-8B-Instruct",
    tensor_parallel_size=1,
    speculative_config={
        "model": "ibm-ai-platform/llama3-8b-accelerator",
        "draft_tensor_parallel_size": 1,
        "method": "mlp_speculator",
    },
)
outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

speculative_config 关键字段解读

上述字典最终会解析为 SpeculativeConfig 的构造参数（对应 vllm/config/speculative.py）：

字段	含义	本示例取值
model	草稿（加速器）模型的 HF 仓库标识，例如 ibm-ai-platform/llama3-8b-accelerator	必填
draft_tensor_parallel_size	草稿模型使用的张量并行规模	1
method	投机方法；MLP 加速器对应 "mlp_speculator"	"mlp_speculator"

需要补充的是，该投机方法实际通过草稿模型的 HuggingFace 配置完成自动识别与校验，因此 method 也可以省略：当草稿 checkpoint 的 model_type == "mlp_speculator" 时，vllm/config/speculative.py 会自动将其解析为 mlp_speculator 方法（源码中 MLPSpeculatorConfig 的 model_type 即为 "mlp_speculator"，见 vllm/transformers_utils/configs/mlp_speculator.py）。

MLP 加速器模型内部结构

尽管原文档只给出了使用示例，仓库内 vllm/model_executor/models/mlp_speculator.py 完整实现了该模型的推理结构，理解它有助于把握草稿预测如何"一步多 token"。

逐 stage 的 MLP 预测模块

对于每个预测位置，模型都维护一组子模块（数量由配置中的 lookahead 步数决定）：

• token 嵌入 VocabParallelEmbedding：把"前序已采样 token"映射为向量；
• 线性投影 nn.Linear：融合上下文向量与嵌入信号，首层从目标模型的 emb_dim 投影，后续 stage 从 inner_dim 投影到 inner_dim；
• 输出头 ParallelLMHead：每个 stage 配一个独立的词表头，直接给出该位置的 token logits，因此一次前向即可同时得到第 1、第 2、…… 个未来 token 的预测分布，无需自回归循环；
• L2 归一化层 MLPSpeculatorLayerNorm：实现 x / ||x||（带 eps 与可学习的 scale/shift），对中间信号做长度归一，提升数值稳定性。

tie_weights 与 scale_input 两种配置形态

从 vllm/model_executor/models/mlp_speculator.py 的分支可见两类参数化：

• 当 tie_weights=False 时，每个 stage 拥有独立的嵌入、投影、输出头与 LayerNorm，参数随 stage 数量线性增长；
• 当 tie_weights=True 时，第 2 个 stage 及之后的所有 stage 共享同一套嵌入与投影（首个 stage 的输入维度可能不同，故单独保留），以更少参数量近似相同的建模能力——该选项要求 n_predict > 1，否则会触发断言错误。

此外若开启 scale_input，模型会先对目标模型传入的隐状态做一次无 scale/shift 的 L2 归一化（即 ln0），再进入各预测 stage。所有 stage 共享一个激活函数 GELU，最终经 LogitsProcessor 输出采样分布。

预测步数的取值逻辑

草稿模型配置中的 n_predict（即每个 stage 预测的 lookahead token 数，等价地写入 num_lookahead_tokens）决定一次前向的候选长度：

• 若用户在 speculative_config 中未显式设置 num_speculative_tokens，引擎会默认采用草稿配置定义的 n_predict（见 vllm/config/speculative.py）；
• 若显式设置的投机步数大于 n_predict，则要求两者满足整除关系，否则抛错——这与 MTP/MLP 这类"复用同一 stage 前向"的机制一致。

配置项定义与默认值（MLPSpeculatorConfig）

加速器的 HuggingFace 配置文件定义在 vllm/transformers_utils/configs/mlp_speculator.py，各字段含义如下（源码 docstring 整理）：

配置项	默认值	说明
vocab_size	32000	词表大小，与目标模型一致
emb_dim	4096	模型（目标大模型）嵌入维度，即上下文向量的维度
inner_dim	0	加速器内部维度；为 0 时取 emb_dim
n_predict	3	加速器预测的 lookahead 步数，即 stage 数量
top_k_tokens_per_head	[5, 4, 3]	每个 head 在构成候选树时考虑的 token 数（需满足长度为 n_predict；源码注释标注当前推理路径暂未使用）
n_candidates	5	每条序列创建的子候选数量
tie_weights	False	是否让首个 stage 之后的各 stage 共享一组权重
scale_input	False	是否先对目标模型输入的隐状态做缩放/归一

其中 top_k_tokens_per_head 与 n_predict 之间存在 assert len(top_k_tokens_per_head) == n_predict 的硬约束，自定义 checkpoint 时须保持一致。预训练权重文件以 speculator. 为前缀，模型加载逻辑 会在加载时去掉该前缀以匹配内部参数名。

张量并行：MLP 加速器目前仅支持 TP=1

原文档示例中草稿与目标模型均使用 tensor_parallel_size=1。这一点并非随意选择，而是由引擎强制约束的。在 vllm/config/speculative.py 的 _verify_and_get_draft_tp 中，当草稿模型 model_type == "mlp_speculator" 且用户未显式设置 draft_tensor_parallel_size 时，引擎会强制将其置为 1；即便目标模型本身使用了 tp > 1，也会输出告警并保持草稿 TP=1（而其他类型的草稿模型默认会跟随目标模型 TP）。显式传入的值只允许是 1 或目标模型 TP，否则校验失败。

预训练 MLP 加速器模型（可直接使用）

原文档列出了 HF Hub 上可用的已训练加速器 checkpoint，覆盖 Llama 2/3 与 IBM Granite 系列，均按"目标系列 + 加速器命名"组织，命名空间为 ibm-ai-platform 与 ibm-granite：

• Llama 系列：ibm-ai-platform/llama-13b-accelerator、ibm-ai-platform/llama3-8b-accelerator、ibm-ai-platform/llama3-70b-accelerator
• Code Llama：ibm-ai-platform/codellama-34b-accelerator
• Llama 2 大尺寸：ibm-ai-platform/llama2-70b-accelerator
• Granite 代码系列：ibm-granite/granite-3b-code-instruct-accelerator、ibm-granite/granite-8b-code-instruct-accelerator、ibm-granite/granite-20b-code-instruct-accelerator
• Granite 指令系列：ibm-granite/granite-7b-instruct-accelerator

选用时请确保其对应的目标大模型与草稿模型的 emb_dim/vocab_size 相互匹配（草稿模型已按目标系列预训练，实践中直接按"目标模型同系列"搭配即可，如示例中 Llama-3.1-8B-Instruct 配 llama3-8b-accelerator）。

已知问题：llama3-70b-accelerator

原文档特别标注了一个社区已知缺陷：加载 ibm-ai-platform/llama3-70b-accelerator 时可能报错：

AttributeError: 'MLPSpeculatorConfig' object has no attribute 'num_attention_heads'

原因是 vLLM 引擎在推导部分配置（如注意力头相关属性、日志记录 head_dim 等）时会对任意草稿模型配置做通用属性访问，而该类 checkpoint 的 MLPSpeculatorConfig 并不携带大模型的 num_attention_heads 等字段。使用 70B 加速器前请先确认当前 vLLM 版本是否已合入相应修复，或改用同系列其余 checkpoint。该问题在社区中由 issue #34106 与 PR #34163 跟踪。

何时选择 MLP 投机解码

从仓库实现可总结出该方案的适用画像与注意点：

• 算力开销小：草稿模型不复制完整 Transformer 主干，仅由若干轻量 MLP stage 组成，额外显存与计算远低于同规模自回归小模型；
• 候选质量高：由于显式利用目标模型的上下文向量 + 已采样 token 双重条件，长序列上接受率通常优于纯 n-gram 等无条件草稿；
• 配置简捷：无需逐 token 的采样循环调度，vLLM 会自动从 model_type 识别方法并将默认投机步数对齐到 n_predict；
• 约束明确：当前仅支持草稿 TP=1，且顶层投机步数需能被 n_predict 整除（vllm/config/speculative.py）。

若希望在不改变目标模型输出分布的前提下提升解码吞吐，且目标模型属于 Llama 2/3、Code Llama 或 Granite 生态，MLP 加速器是一个低成本、开箱即用的选项。对于其他目标模型家族，可参考仓库 docs/features 下投机解码子目录中的其余方案（如 EAGLE、MTP 等）对比选择。