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_mtp、qwen3_5_mtp 和 eagle3。仅支持 PyTorch 后端。
工作原理#
语法掩码在两个阶段分别施加:
草稿模型 — 使用 fork 出的语法匹配器,逐位置串行施加掩码。每个位置的掩码依赖于前一位置接受的 token,确保草稿模型提出符合语法的 token。
主模型验证 — 对主模型的 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)