推测解码指南¶
本指南介绍如何在 vLLM Ascend 中使用推测解码。推测解码是一种在内存受限的 LLM 推理中改善令牌间延迟的技术。
概述¶
vLLM Ascend 通过**提议者-验证者**架构实现推测解码:
- 提议者 (
vllm_ascend/spec_decode/):使用多种方法生成草稿(推测)令牌——从简单的 n-gram 匹配到基于神经网络的草稿模型。 - 拒绝采样器 (
vllm_ascend/sample/):根据目标模型的输出验证草稿令牌,接受匹配项并拒绝不匹配项,可选的优化包括块验证和熵验证。
支持以下推测解码方法:
| 方法 | 描述 |
|---|---|
ngram |
从提示中匹配 n-gram |
suffix |
基于后缀的模式匹配(需要 Arctic Inference) |
medusa |
嵌入在目标模型中的 Medusa 头 |
eagle |
基于 EAGLE 的草稿模型 |
eagle3 |
基于 EAGLE-3 的草稿模型 |
mtp |
使用共享嵌入头的多令牌预测 |
dflash |
基于块扩散的并行草稿模型 |
draft_model |
通用外部草稿 LLM |
extract_hidden_states |
提取隐藏状态用于 EAGLE 训练 |
通用配置¶
所有推测解码方法在初始化模型或启动服务器时通过 speculative_config 参数进行配置:
method(str, 必需):推测解码方法。必须是上表中列出的支持方法名称之一。num_speculative_tokens(int, 必需):每次前向传播生成的推测令牌数量。当可用时,从草稿模型的n_predict配置(例如 MTP)或suffix_decoding_max_tree_depth(后缀方法)自动填充。model(str, 可选):草稿模型的路径或 HF 仓库 ID。eagle、eagle3、dflash、medusa和draft_model必需。对于mtp(重用目标模型)、ngram、suffix和extract_hidden_states自动解析。draft_tensor_parallel_size(int, 可选):草稿模型的张量并行大小。只能为1或与目标模型的张量并行大小相同。disable_padded_drafter_batch(bool, 默认值:False):禁用推测解码的输入填充。如果设置为True,推测输入批次可以包含不同长度的序列,这可能仅受某些注意力后端支持。注意: 仅对eagle、eagle3、mtp、dflash、draft_model和extract_hidden_states方法有效。
离线推理 — 将 speculative_config 作为 Python 字典传递给 LLM():
from vllm import LLM
llm = LLM(
model="path/to/target/model",
speculative_config={
"method": "eagle3",
"model": "path/to/draft/model",
"num_speculative_tokens": 3,
},
)
在线服务 — 将 --speculative-config(或 -sc)作为 JSON 字符串传递:
vllm serve path/to/target/model \
--speculative-config '{"method": "eagle3", "model": "path/to/draft/model", "num_speculative_tokens": 3}'
[!NOTE] 在 Ascend NPU 上,
npu_fused_infer_attention_score算子每轮解码最多支持 16 个令牌。因此,(num_speculative_tokens + 1)必须 ≤ 15。
通过匹配提示中的 n-gram 进行推测¶
以下代码配置 vLLM Ascend 使用推测解码,其中通过匹配提示中的 n-gram 生成提议。
-
离线推理
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={ "method": "ngram", "num_speculative_tokens": 5, "prompt_lookup_max": 4, }, ) 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}")
使用基于 EAGLE 的草稿模型进行推测¶
以下代码配置 vLLM Ascend 使用推测解码,其中提议由基于 EAGLE(Extrapolation Algorithm for Greater Language-model Efficiency) 的草稿模型生成。
在 vLLM Ascend 的 v0.12.0rc1 版本中,异步调度器更加稳定并已准备好启用。我们已对其进行适配以支持 EAGLE,您可以通过如下设置 async_scheduling=True 来使用它。如果遇到任何问题,请随时在 GitHub 上提交 issue。作为临时解决方案,您可以在初始化模型时取消设置 async_scheduling=True 来禁用此功能。
-
离线推理
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=4, distributed_executor_backend="mp", enforce_eager=True, async_scheduling=True, speculative_config={ "method": "eagle", "model": "yuhuili/EAGLE-LLaMA3.1-Instruct-8B", "draft_tensor_parallel_size": 1, "num_speculative_tokens": 2, }, ) 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}")
使用基于 EAGLE 的草稿模型时需要考虑的几个重要事项:
-
HF 仓库中的 EAGLE 模型 提供的 EAGLE 草稿模型应由 vLLM 直接加载和使用。此功能在 PR #4893 中添加。 如果您使用的 vLLM 版本在此拉取请求合并之前发布,请更新到更新版本。
-
基于 EAGLE 的草稿模型需要在没有张量并行的情况下运行 (即
speculative_config中的 draft_tensor_parallel_size 设置为 1),尽管 主模型可以使用张量并行运行(参见上面的示例)。 -
使用基于 EAGLE-3 的草稿模型时,选项 "method" 必须设置为 "eagle3"。 也就是说,在
"method": "eagle3"中指定speculative_config。 -
启用 EAGLE 后,主模型需要在一个解码过程中验证主模型和草稿模型生成的
(1 + K)个令牌。 并且全图模式将固定验证阶段的令牌数量, 因此cudagraph_capture_sizes必须是一个捕获大小列表,其中每个大小计算为n * (K + 1),n是您想要支持的每个批次大小。 例如,要支持批次大小从 1 到 4 且num_speculative_tokens = 4,cudagraph_capture_sizes应设置为[5, 10, 15, 20]。
使用 MTP 进行推测¶
MTP(多令牌预测)通过并行化预测多个令牌,将生成方式从单令牌转变为多令牌,从而提升推理性能。这种方法显著提高了生成吞吐量,并实现了推理速度的倍增加速——同时不牺牲输出质量。
-
在线推理
vllm serve /deepseek-ai/DeepSeek-V3.2-Exp-W8A8 \ --port 20004 \ --data-parallel-size 1 \ --tensor-parallel-size 16 \ --enable-expert-parallel \ --seed 1024 \ --served-model-name dsv3 \ --max-model-len 36768 \ --max-num-batched-tokens 5000 \ --max-num-seqs 10 \ --quantization ascend \ --trust-remote-code \ --gpu-memory-utilization 0.9 \ --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY"}' \ --speculative-config '{"num_speculative_tokens": 2, "method":"mtp", "disable_padded_drafter_batch": false}'
[!NOTE] 由于 DeepSeek 的 MTP 仅暴露单层权重,在
num_speculative_tokens > 1(尤其是 ≥ 3)的场景下,精度和性能无法得到有效保证。在
num_speculative_tokens > 1的全图模式下,每个 ACLGraph 的捕获大小必须是(num_speculative_tokens + 1)的整数倍。
使用后缀解码进行推测¶
以下代码配置 vLLM 使用推测解码,其中候选令牌通过后缀解码生成 (SuffixDecoding: Extreme Speculative Decoding for Emerging AI Applications)。
与 n-gram 类似,后缀解码可以通过使用最后 n 个生成的令牌进行模式匹配来生成草稿令牌。与 n-gram 不同,后缀解码 (1) 可以针对提示和先前生成的内容进行模式匹配,(2) 使用频率计数来提出最可能的续写内容,以及 (3) 在每次迭代中为每个请求推测自适应数量的令牌,以获得更好的接受率。
后缀解码在具有高重复性的任务上可以实现更好的性能,例如代码编辑、智能体循环(例如自我反思、自我一致性)和强化学习展开。
[!NOTE] 后缀解码需要 Arctic Inference。您可以使用
pip install arctic-inference进行安装。
- 离线推理
```python 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,
enforce_eager=True,
speculative_config={
"method": "suffix",
"num_speculative_tokens": 15,
},
)
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}")
```
提取隐藏状态¶
extract_hidden_states 方法是一种特殊的推测解码模式,它不执行实际的推测。相反,它从目标模型的指定层提取隐藏状态并将其保存到磁盘。这主要用于为 EAGLE 风格的草稿模型收集训练数据。
[!NOTE] 此方法每个请求仅生成 1 个输出令牌。主要输出是保存到磁盘的隐藏状态,而不是生成的文本。
-
离线推理
import tempfile from safetensors import safe_open from vllm import LLM, SamplingParams def main(): with tempfile.TemporaryDirectory() as tmpdirname: llm = LLM( model="Qwen/Qwen3-8B", tensor_parallel_size=1, speculative_config={ "method": "extract_hidden_states", "num_speculative_tokens": 1, "draft_model_config": { "hf_config": { # Layer indices to extract hidden states from "eagle_aux_hidden_state_layer_ids": [2, 18, 34], } }, }, kv_transfer_config={ "kv_connector": "ExampleHiddenStatesConnector", "kv_role": "kv_producer", "kv_connector_extra_config": { "shared_storage_path": tmpdirname, }, }, ) prompts = ["Hello, how are you?", "What is machine learning?"] sampling_params = SamplingParams(max_tokens=1) outputs = llm.generate(prompts, sampling_params) for output in outputs: print("Prompt:", output.prompt) print("Prompt token ids:", output.prompt_token_ids) hidden_states_path = output.kv_transfer_params.get("hidden_states_path") print("Hidden states saved to:", hidden_states_path) with safe_open(hidden_states_path, "pt") as f: token_ids = f.get_tensor("token_ids") hidden_states = f.get_tensor("hidden_states") print("Shape:", hidden_states.shape) # Shape: (num_tokens, num_layers, hidden_size) if __name__ == "__main__": main()
关键配置参数:
-
num_speculative_tokens:必须设置为1。此方法不执行实际推测,因此该值是固定的。 -
eagle_aux_hidden_state_layer_ids:要从中提取隐藏状态的层索引列表。例如,[2, 18, 34]从第 2、18 和 34 层提取。 -
kv_connector:必须设置为"ExampleHiddenStatesConnector"以启用将隐藏状态保存到磁盘。 -
kv_role:对于提取模式,必须设置为"kv_producer"。 -
shared_storage_path:隐藏状态将保存为.safetensors文件(每个请求一个文件)的目录。
块验证和熵验证¶
vLLM Ascend 为推测解码中的拒绝采样器提供了两个可选优化:块验证**和**熵验证。这些功能以少量输出精度为代价,换取推理吞吐量的提升。
[!WARNING] 块验证和熵验证都会修改令牌接受标准,并可能导致轻微的精度下降(例如,与标准拒绝采样器相比,输出令牌略有不同)。在生产环境中启用它们之前,请评估对特定工作负载的质量影响。
块验证¶
块验证使用累积概率乘积将所有草稿令牌作为一个整体进行评估,而不是独立检查每个令牌。这可以提高接受率并减少拒绝采样的开销,尤其是在 num_speculative_tokens >= 3 时。
熵验证¶
熵验证根据目标分布的熵调整接受阈值:
- 高熵(不确定分布)→ 较低的有效阈值 → 接受更多令牌
- 低熵(确定分布)→ 较高的有效阈值 → 更严格的拒绝
这个熵感知阈值由两个参数控制:
posterior_threshold(默认值:0.95,范围:(0, 1]):修改后阈值的上限。即使熵非常低,有效阈值也不会超过此值。posterior_alpha(默认值:0.4,范围:>= 0):控制熵对阈值的影响程度。alpha 值越高,阈值对熵变化越敏感,导致推测令牌的接受率更高,但精度损失也更大。您需要根据特定的模型和数据集调整此值。当 alpha 为0时,熵不起作用,阈值等于posterior_threshold。
使用方法¶
-
在线推理
-
离线推理
这两个功能可以独立或同时启用。同时使用时,块验证的累积接受率将与熵验证的熵调整阈值相结合。