推测解码指南

推测解码指南#

本指南介绍如何在 vLLM Ascend 中使用推测解码。推测解码是一种技术，用于改善受内存带宽限制的 LLM 推理中的 token 间延迟。

通过匹配提示词中的 n-gram 进行推测#

以下代码配置 vLLM Ascend 使用推测解码，其中候选 token 通过匹配提示词中的 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 使用推测解码，其中候选 token 由基于 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 的草稿模型时需要考虑的几个重要事项：

EAGLE 模型的 HF 仓库中可用的 EAGLE 草稿模型应由 vLLM 直接加载和使用。此功能已在 PR #4893 中添加。如果您使用的 vLLM 版本早于该拉取请求合并的时间，请更新到较新的版本。
基于 EAGLE 的草稿模型需要在不使用张量并行的情况下运行（即 speculative_config 中的 draft_tensor_parallel_size 设置为 1），尽管主模型可以使用张量并行运行（参见上面的示例）。
当使用基于 EAGLE-3 的草稿模型时，选项 "method" 必须设置为 "eagle3"。也就是说，在 speculative_config 中指定 "method": "eagle3"。
启用 EAGLE 后，主模型需要在一个解码过程中验证由主模型和草稿模型生成的 (1 + K) 个 token。而全图模式将在验证阶段固定 token 的数量，因此 cudagraph_capture_sizes 必须是一个捕获大小的列表，其中每个大小计算为 n * (K + 1)，n 是您希望支持的每个批次大小。例如，要支持批次大小从 1 到 4 且 num_speculative_tokens = 4，cudagraph_capture_sizes 应设置为 [5, 10, 15, 20]。

使用 MTP 推测器进行推测#

以下代码配置 vLLM Ascend 使用推测解码，其中候选 token 由 MTP（多 Token 预测）生成，通过并行化多个 token 的预测来提升推理性能。有关 MTP 的更多信息，请参阅 Multi_Token_Prediction

在线推理

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":"deepseek_mtp", "disable_padded_drafter_batch": "False"}'

使用后缀解码进行推测#

以下代码配置 vLLM 使用推测解码，其中候选 token 使用后缀解码生成 (SuffixDecoding: Extreme Speculative Decoding for Emerging AI Applications)。

与 n-gram 类似，后缀解码可以通过使用最后 n 个生成的 token 进行模式匹配来生成草稿 token。与 n-gram 不同，后缀解码 (1) 可以针对提示词和之前的生成内容进行模式匹配，(2) 使用频率计数来提议最可能的后续内容，(3) 为每个请求在每次迭代中推测自适应数量的 token，以获得更好的接受率。

后缀解码可以在高重复性任务中实现更好的性能，例如代码编辑、智能体循环（例如自我反思、自洽性）和强化学习（RL）的 rollout。

[!注意] 后缀解码需要 Arctic Inference。您可以通过 pip install arctic-inference 安装它。

离线推理

  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 风格草稿模型的训练数据。

[!注意] 此方法每个请求仅生成 1 个输出 token。主要输出是保存到磁盘的隐藏状态，而不是生成的文本。

离线推理

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 文件的目录（每个请求一个文件）。

推测解码指南

目录

推测解码指南#

通过匹配提示词中的 n-gram 进行推测#

使用基于 EAGLE 的草稿模型进行推测#

使用 MTP 推测器进行推测#

使用后缀解码进行推测#

提取隐藏状态#