Speculative Decoding#

投机解码是一种优化技术,它通过引入轻量级草稿模型来预测多个后续token,再由主模型在前向推理过程中验证并选择匹配度最高的长token序列。与标准的自回归解码相比,这种方法可使系统一次性生成多个token。

备注

请注意,这是lmdeploy中的实验性功能。

示例#

请参考如下使用示例。

Eagle 3#

安装依赖#

安装 flash-atten3

git clone --depth=1 https://github.com/Dao-AILab/flash-attention.git
cd flash-attention/hopper
python setup.py install

pipeline#

from lmdeploy import PytorchEngineConfig, pipeline
from lmdeploy.messages import SpeculativeConfig


if __name__ == '__main__':

    model_path = 'meta-llama/Llama-3.1-8B-Instruct'
    spec_cfg = SpeculativeConfig(
        method='eagle3',
        num_speculative_tokens=3,
        model='yuhuili/EAGLE3-LLaMA3.1-Instruct-8B',
    )
    pipe = pipeline(model_path, backend_config=PytorchEngineConfig(max_batch_size=128), speculative_config=spec_cfg)
    response = pipe(['Hi, pls intro yourself', 'Shanghai is'])
    print(response)

serving#

lmdeploy serve api_server \
meta-llama/Llama-3.1-8B-Instruct \
--backend pytorch \
--server-port 24545 \
--speculative-draft-model yuhuili/EAGLE3-LLaMA3.1-Instruct-8B \
--speculative-algorithm eagle3 \
--speculative-num-draft-tokens 3 \
--max-batch-size 128 \
--enable-metrics

Deepseek MTP#

安装依赖#

Install FlashMLA

git clone https://github.com/deepseek-ai/FlashMLA.git flash-mla
cd flash-mla
git submodule update --init --recursive
pip install -v .

pipeline#

from lmdeploy import PytorchEngineConfig, pipeline
from lmdeploy.messages import SpeculativeConfig


if __name__ == '__main__':

    model_path = 'deepseek-ai/DeepSeek-V3'
    spec_cfg = SpeculativeConfig(
        method='deepseek_mtp',
        num_speculative_tokens=3,
    )
    pipe = pipeline(model_path,
                    backend_config=PytorchEngineConfig(tp=16, max_batch_size=128),
                    speculative_config=spec_cfg)
    response = pipe(['Hi, pls intro yourself', 'Shanghai is'])
    print(response)

serving#

lmdeploy serve api_server \
deepseek-ai/DeepSeek-V3 \
--backend pytorch \
--server-port 24545 \
--tp 16 \
--speculative-algorithm deepseek_mtp \
--speculative-num-draft-tokens 3 \
--max-batch-size 128 \
--enable-metrics

投机解码与结构化输出#

投机解码(MTP)可以与结构化输出结合使用,使草稿模型提出的 token 也遵循语法约束(如 JSON Schema、正则表达式),从而显著提高接受率。

备注

该功能支持继承自 DeepseekMTP 的投机方法,包括 deepseek_mtpqwen3_5_mtpeagle3。仅支持 PyTorch 后端。

工作原理#

语法掩码在两个阶段分别施加:

  1. 草稿模型 — 使用 fork 出的语法匹配器,逐位置串行施加掩码。每个位置的掩码依赖于前一位置接受的 token,确保草稿模型提出符合语法的 token。

  2. 主模型验证 — 对主模型的 logits 进行逐位置串行的语法掩码处理。拒绝采样后,仅将接受的 token 反馈给原始(未 fork 的)语法匹配器,使其为下一步保持正确的状态。

当草稿模型使用与主模型不同的词表时(例如 Eagle 3 使用压缩的草稿词表),xgrammar 生成的目标词表位掩码会通过高效的 scatter-add 内核转换为草稿词表位掩码,然后再应用于草稿 logits。

pipeline#

from lmdeploy import PytorchEngineConfig, pipeline
from lmdeploy.messages import GenerationConfig, SpeculativeConfig

model_path = 'deepseek-ai/DeepSeek-V3'
spec_cfg = SpeculativeConfig(method='deepseek_mtp', num_speculative_tokens=3)
pipe = pipeline(
    model_path,
    backend_config=PytorchEngineConfig(tp=16, max_batch_size=128),
    speculative_config=spec_cfg,
)

schema = {
    'type': 'object',
    'properties': {
        'name': {'type': 'string'},
        'age': {'type': 'integer'},
    },
    'required': ['name', 'age'],
}
gen_config = GenerationConfig(
    response_format=dict(type='json_schema', json_schema=dict(name='person', schema=schema)),
    max_new_tokens=256,
)

if __name__ == '__main__':

  response = pipe(['请用 JSON 格式做自我介绍。'], gen_config=gen_config)
  print(response)

api_server#

lmdeploy serve api_server \
deepseek-ai/DeepSeek-V3 \
--backend pytorch \
--server-port 24545 \
--tp 16 \
--speculative-algorithm deepseek_mtp \
--speculative-num-draft-tokens 3 \
--max-batch-size 128

客户端可以按照结构化输出文档中的方式使用 response_format

from openai import OpenAI

schema = {
    'type': 'object',
    'properties': {
        'name': {'type': 'string'},
        'age': {'type': 'integer'},
    },
    'required': ['name', 'age'],
}
response_format = dict(type='json_schema', json_schema=dict(name='person', schema=schema))

if __name__ == '__main__':

  client = OpenAI(api_key='YOUR_API_KEY', base_url='http://0.0.0.0:24545/v1')
  model_name = client.models.list().data[0].id
  response = client.chat.completions.create(
      model=model_name,
      messages=[{'role': 'user', 'content': '请用 JSON 格式做自我介绍。'}],
      response_format=response_format,
  )
  print(response)