推测解码指南#

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

概述#

vLLM Ascend 通过提议器-验证器架构实现推测解码:

  1. 提议器 (vllm_ascend/spec_decode/):使用多种方法生成草稿(推测)令牌——从简单的 n-gram 匹配到基于神经网络的草稿模型。

  2. 拒绝采样器 (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(后缀方法)自动填充。

    注意:对于 PD 分离部署,num_speculative_tokens 应满足以下条件之一:

    1. 混合 Mamba 模型(例如 Qwen-Next 和 Qwen3.5 系列):P 节点和 D 节点上的 num_speculative_tokens 应相等。

    2. 其他模型:P 节点上的 num_speculative_tokens 应为 1,D 节点上的 num_speculative_tokens 应大于或等于 1。

  • model (str, 可选):草稿模型的路径或 HF 仓库 ID。eagleeagle3dflashmedusadraft_model 必需。对于 mtp(重用目标模型)、ngramsuffixextract_hidden_states 自动解析。

  • draft_tensor_parallel_size (int, 可选):草稿模型的张量并行大小。只能是 1 或与目标模型的张量并行大小相同。

  • disable_padded_drafter_batch (bool, 默认值:False):禁用推测解码的输入填充。如果设置为 True,推测输入批次可以包含不同长度的序列,这可能仅受某些注意力后端支持。**注意:**仅对 eagleeagle3mtpdflashdraft_modelextract_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] 在昇腾 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 的草稿模型时需要考虑的几个重要事项:

  1. HF 仓库中的 EAGLE 草稿模型应由 vLLM 直接加载和使用。此功能在 PR #4893 中添加。如果您使用的 vLLM 版本在此拉取请求合并之前发布,请更新到较新版本。

  2. 基于 EAGLE 的草稿模型需要在没有张量并行的情况下运行(即 speculative_config 中的 draft_tensor_parallel_size 设置为 1),尽管主模型可以使用张量并行运行(参见上面的示例)。

  3. 使用基于EAGLE-3的草稿模型时,选项"method"必须设置为"eagle3"。即在speculative_config中指定"method": "eagle3"

  4. 启用EAGLE后,主模型需要在一次解码过程中验证主模型和草稿模型生成的(1 + K)个token。由于全图模式会固定验证阶段的token数量,因此cudagraph_capture_sizes必须是一个捕获大小列表,其中每个大小按您要支持的每个批次大小n计算为n * (K + 1)。例如,要支持批次大小1到4且num_speculative_tokens = 4,应将cudagraph_capture_sizes设置为[5, 10, 15, 20]

使用MTP进行推测#

MTP(多token预测)通过并行化多个token的预测来提升推理性能,从单token生成转变为多token生成。这种方法显著提高了生成吞吐量,并实现了推理速度的倍增加速——且不影响输出质量。

  • 在线推理

    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个生成的token来生成草稿token。与n-gram不同,后缀解码(1)可以同时对提示和先前生成内容进行模式匹配,(2)使用频率计数来提议最可能的续写内容,(3)在每次迭代中为每个请求推测自适应数量的token以获得更好的接受率。

后缀解码在重复性高的任务上可以获得更好的性能,例如代码编辑、智能体循环(如自我反思、自一致性)和强化学习展开。

[!NOTE] 后缀解码需要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风格草稿模型的训练数据。

[!NOTE] 此方法每个请求仅生成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()
    

关键配置参数:

  1. num_speculative_tokens:必须设置为1。此方法不执行实际推测,因此该值固定。

  2. eagle_aux_hidden_state_layer_ids:要提取隐藏状态的层索引列表。例如,[2, 18, 34]从第2、18和34层提取。

  3. kv_connector:必须设置为"ExampleHiddenStatesConnector"以启用将隐藏状态保存到磁盘。

  4. kv_role:对于提取模式,必须设置为"kv_producer"

  5. shared_storage_path:隐藏状态将保存为.safetensors文件的目录(每个请求一个文件)。

块验证与熵验证#

vLLM Ascend为推测解码中的拒绝采样器提供了两个可选优化:块验证熵验证。这些功能以少量输出精度换取推理吞吐量的提升。

[!WARNING] 块验证和熵验证都会修改token接受标准,可能导致轻微精度下降(例如,与标准拒绝采样器相比输出token略有不同)。在生产环境中启用前,请评估对特定工作负载的质量影响。

块验证#

块验证使用累积概率乘积将所有草稿token作为一个整体进行评估,而非独立检查每个token。这可以提高接受率并减少拒绝采样的开销,尤其在num_speculative_tokens >= 3时。

熵验证#

熵验证根据目标分布的熵调整接受阈值:

  • 高熵(不确定分布)→ 较低的有效阈值 → 接受更多token

  • 低熵(确定分布)→ 较高的有效阈值 → 更严格的拒绝

这个基于熵的阈值由两个参数控制:

  • posterior_threshold(默认值:0.95,范围:(0, 1]):修改后阈值的上限。即使熵值非常低,有效阈值也不会超过此值。

  • posterior_alpha(默认值:0.4,范围:>= 0):控制熵对阈值的影响强度。alpha值越高,阈值对熵变化越敏感,导致推测令牌的接受率更高,但精度损失也更大。您需要根据具体模型和数据集调整此值。当alpha为0时,熵不起作用,阈值等于posterior_threshold

使用方法#

  • 在线推理

    vllm serve <model> --additional-config \
        '{"rejection_sampler_config": {"enable_block_verify": true, \
        "enable_entropy_verify": true, "posterior_threshold": 0.95, \
        "posterior_alpha": 0.4}}'
    
  • 离线推理

    llm = LLM(
        model,
        additional_config={
            "rejection_sampler_config": {
                "enable_block_verify": True,
                "enable_entropy_verify": True,
                "posterior_threshold": 0.95,
                "posterior_alpha": 0.4,
            }
        },
    )
    

这两个功能可以独立启用,也可以同时启用。同时使用时,Block Verify的累积接受率将与Entropy Verify的熵调整阈值相结合。