MTP投机推理¶
背景¶
MTP是一种创新的推理阶段加速技术,专注于解决大语言模型生成过程中的效率瓶颈。MTP的本质是通过预训练阶段的特殊设计,为推理阶段提供高效的草稿token预测能力,从而显著提升模型的生成速度。其核心价值在于平衡推理效率与输出质量,为大语言模型的长序列生成问题提供了一种高效的解决方案,最终实现推理性能的优化。
功能介绍¶
MTP在推理加速方面具有以下核心功能:
-
高效草稿生成:使用低成本的MTP结构快速生成草稿token,这些草稿token作为主模型验证的基础,大幅减少了传统自回归生成的计算开销。
-
批量验证机制:主模型能够同时批量验证多个MTP生成的草稿token,而不必逐个生成和验证,显著提升了推理速度。
-
高采样准确率:MTP解决了Eagle、Medusa等现有推理加速方法中的关键痛点——训练后生成的draft模块token采样率低的问题。由于MTP在预训练阶段就优化了草稿生成能力,其生成的token具有更高的准确率,减少了主模型的验证负担。
-
推理延迟降低:通过预先生成多个可能的后续token,MTP有效降低了模型生成长文本时的累积延迟,使用户体验更加流畅。
-
资源消耗优化:相比其他推理加速技术,MTP在保持加速效果的同时,对计算资源的额外需求更少,适合在资源受限环境下部署。
MTP技术为大语言模型的推理阶段提供了一种全新的效率优化方案,特别适合需要快速响应的实时应用场景,代表了语言模型推理优化的重要发展方向。
模型支持
目前支持以下模型的MTP结构导出: - DeepSeek-V3 (输入 model_type: deepseek_v3, 导出 MTP model_type: deepseek_v3_mtp) - DeepSeek-V3.2 (输入 model_type: deepseek_v3, 导出 MTP model_type: deepseek_v32_mtp) - DeepSeek-R1 (输入 model_type: deepseek_v3, 导出 MTP model_type: deepseek_v3_mtp) - GLM4 MoE (如 GLM-4.5-Air, 导出 MTP model_type: glm4_moe_mtp)
注意: - DeepSeek V3 和 R1 的输入 model_type 都是 "deepseek_v3",导出的 MTP 模型 model_type 为 "deepseek_v3_mtp" - DeepSeek V3.2 的输入 model_type 是 "deepseek_v3"(但可通过 index_head_dim 等字段自动识别),导出的 MTP 模型 model_type 为 "deepseek_v32_mtp"
使用示例¶
导出模型¶
脚本会自动检测模型类型,也可以手动指定。
DeepSeek-V3¶
python3 tools/export_mtp.py \
--input-dir /path/to/DeepSeek-V3 \
--output-dir /path/to/DeepSeek-V3-mtp
DeepSeek-V3.2¶
python3 tools/export_mtp.py \
--input-dir /path/to/DeepSeek-V3.2 \
--output-dir /path/to/DeepSeek-V3.2-mtp
DeepSeek-R1¶
python3 tools/export_mtp.py \
--input-dir /path/to/DeepSeek-R1 \
--output-dir /path/to/DeepSeek-R1-mtp
GLM4 MoE¶
python3 tools/export_mtp.py \
--input-dir /path/to/GLM-4.5-Air \
--output-dir /path/to/GLM-4.5-Air-mtp
手动指定模型类型¶
如果自动检测失败,可以手动指定模型类型:
python3 tools/export_mtp.py \
--input-dir /path/to/model \
--output-dir /path/to/model-mtp \
--model-type deepseek_v3 # 可选: deepseek_v3 (用于V3/R1), deepseek_v32 (用于V3.2), glm4_moe
输入模型参考: - DeepSeek-V3 - DeepSeek-V3.2 - DeepSeek-R1 - GLM-4.5-Air
启动脚本¶
使用MTP进行推理时,需要同时指定主模型和草稿模型(MTP模型)。
DeepSeek-V3/V3.2/R1 启动示例¶
MODEL_PATH="/models/DeepSeek-V3"
DRAFT_MODEL_PATH="/models/DeepSeek-V3-mtp"
MASTER_NODE_ADDR="127.0.0.1:42123"
START_PORT=13222
START_DEVICE=0
LOG_DIR="log"
NNODES=16
for (( i=0; i<$NNODES; i++ ))
do
PORT=$((START_PORT + i))
DEVICE=$((START_DEVICE + i))
LOG_FILE="$LOG_DIR/node_$i.log"
nohup ./xllm \
--model $MODEL_PATH \
--devices="npu:$DEVICE" \
--port $PORT \
--master_node_addr=$MASTER_NODE_ADDR \
--nnodes=$NNODES \
--draft_model $DRAFT_MODEL_PATH \
--draft_devices="npu:$DEVICE" \
--num_speculative_tokens 1 \
--max_memory_utilization=0.90 \
--max_tokens_per_batch=10000 \
--max_seqs_per_batch=256 \
--block_size=128 \
--ep_size=1 \
--dp_size=1 \
--enable_prefix_cache=false \
--enable_chunked_prefill=false \
--node_rank=$i > $LOG_FILE 2>&1 &
sleep 0.5
done
GLM4 MoE 启动示例¶
性能数据¶
基于sharegpt数据集,输入长度2500,输出长度1500,请求总数80。
| method | Concurrency | Mean TPOT(ms) | Mean TTFT(ms) | Output Tokens/s | Total Tokens/s |
|---|---|---|---|---|---|
| baseline | 1 | 40.61 | 141.80 | 24.20 | 65.77 |
| mtp | 1 | 28.33 | 142.35 | 35.19 | 95.52 |
| baseline | 2 | 42.69 | 178.59 | 45.16 | 122.74 |
| mtp | 2 | 29.81 | 187.97 | 64.75 | 175.78 |
| baseline | 4 | 46.18 | 172.34 | 79.83 | 216.96 |
| mtp | 4 | 33.54 | 194.22 | 111.18 | 301.81 |
| baseline | 8 | 53.16 | 181.49 | 110.68 | 300.81 |
| mtp | 8 | 40.99 | 203.37 | 154.46 | 419.34 |
| baseline | 16 | 68.50 | 213.89 | 143.81 | 390.84 |
| mtp | 16 | 57.04 | 254.99 | 201.89 | 548.04 |
| baseline | 20 | 74.72 | 228.80 | 154.77 | 420.65 |
| mtp | 20 | 61.73 | 264.34 | 206.24 | 559.84 |
| baseline | 40 | 119.68 | 559.32 | 180.22 | 489.80 |
| mtp | 40 | 105.70 | 544.54 | 252.91 | 686.74 |
| baseline | 80 | 180.89 | 2996.21 | 192.09 | 522.06 |
| mtp | 80 | 152.19 | 2163.72 | 278.07 | 755.12 |