Accelerating SGLang with Multiple Token Prediction

发表时间: 2025-07 · Blog post by Eigen AI Team (lmsys.org)

原文: https://lmsys.org/blog/2025-07-17-mtp

作者/机构:Eigen AI Team
发布日期:2025年7月17日

速读

一句话结论 SGLang 推理框架成功集成了 Multiple Token Prediction (MTP) 投机解码技术,并与大规模专家并行及 Prefill-Decode 分离架构无缝兼容,在零画质损失的前提下,为 DeepSeek V3 模型带来了最高 60% 的输出吞吐量提升。

要解决什么问题 传统的大语言模型自回归解码过程在机制上是严格串行的,即每次只能依赖所有历史上下文生成一个 token。这种逐字生成的模式严重限制了计算并行度,导致推理速度受限。在面临高并发请求的生产环境中,最大化 GPU 利用率是实现高性能和低成本部署的核心诉求,而这种串行依赖使得 GPU 内部海量的并行计算单元在解码阶段无法被充分调用,成为了制约系统输出吞吐量的最关键卡点。

怎么做的 核心思路是引入 Multiple Token Prediction (MTP) 投机解码技术,将原本串行的单步生成转化为“轻量级草稿预测 + 全模型并行验证”的两阶段流水线。这种方法之所以能绕开串行卡点,是因为它将 $n$ 次低效的串行解码操作,替换成了 GPU 上极高并行度的一次验证操作。关键设计由两个部件构成:第一是起草阶段,系统调用一个轻量级的草稿模型,在一次快速的前向传播中,连续预测出未来 $n$ 个候选 token(例如当目标模型生成了“day”后,草稿模型会迅速向后推测出 $n=4$ 的后续序列“and it's so hot”);第二是验证阶段,完整的全参数目标模型接收这 $n$ 个草稿 token,在单次并行前向传播中对它们进行校验,目标模型会接受与自身输出分布相匹配的最长连续前缀,并在遇到第一个不匹配的 token 时进行重采样以纠正后续内容。整个过程的单步推进长度可由公式 $L_{step} = L_{match} + 1$ 定义,其中 $L_{match}$ 是目标模型连续接受的草稿 token 数量,加一代表在断点处重采样得到的新 token。这个机制的效率高度依赖于“接受长度”(即每次验证平均能通过的 token 数),由于每一个被接受的 token 都经过了全模型的严格校验,该方法没有任何近似或微调妥协,能保证与标准解码完全一致的确定性输出。在工程实现上,MTP 作为一个即插即用的模块,深度整合了 SGLang 的数据并行注意力、专家并行负载均衡(EPLB)、DeepEP MoE 以及 CUDA Graph 等底层特性。在超参设置上,草稿长度 draft_token_num 是核心变量:默认设为 2 风险最低;若 GPU 算力有冗余可调大至 4 以榨取更高吞吐;若 GPU 已满载,则应保持在 2 到 3 以避免引入额外计算负担。

效果如何 实验以 DeepSeek V3 为测试模型,在两种规模的集群上进行了验证。第一个是小规模部署场景:硬件为 16 张 H200 GPU(2 个解码节点),每卡并发 2 个请求,输入 65536 tokens,输出 4096 tokens;对比基线有两个,一是无 MTP 且无 Overlap 调度的纯基线(吞吐量 51.0 tokens/sec),二是仅开启 Overlap 调度的基线(吞吐量 60.4 tokens/sec);在开启 MTP(topk=1)后,3 token 窗口配置达到了 81.5 tokens/sec(平均接受长度 2.18),4 token 窗口配置达到了 82.0 tokens/sec(平均接受长度 2.44),相比纯基线实现了 60.8% 的吞吐量提升。第二个是大规模部署场景:硬件扩展至 128 张 H200 GPU(16 个节点,其中 4 个 Prefill 节点,12 个解码节点),每卡并发 128 个请求,输入 2000 tokens,输出 100 tokens;在此高负载环境下,配置 draft_token_num=2 和 topk=1,相比无 MTP 且无 Overlap 的基线,吞吐量依然提升了 14.2%,证明了该方法在大规模专家并行(EP)和 Prefill-Decode 分离架构下的有效性。该方法目前也存在局限性:当前的 MTP 实现尚未兼容 Overlap 调度,两者结合的开发仍在进行中;此外,盲目增大草稿 token 数量在系统满载时会导致接受率下降,反而可能拖累整体性能。

主要贡献

随着大型语言模型能力的不断增长,其逐个Token解码的串行本质成为了限制推理吞吐量的关键瓶颈。在最大化GPU利用率至关重要的高需求生产应用中,这一局限性尤为明显。为了解决这一问题,本文为SGLang推理框架引入了多Token预测(Multiple Token Prediction, MTP)技术。

本文的主要贡献如下:

  1. 多特性无缝融合:SGLang目前已支持将多Token预测(MTP)、大规模专家并行(Large-Scale EP)以及预填充-解码分离(Prefill-Decode disaggregation)等高级特性进行平滑结合。
  2. 显著的性能提升:通过引入这种全新的解码范式、更优的并行计算以及更高效的资源利用率,SGLang在不牺牲任何生成质量的前提下,为DeepSeek V3模型带来了高达60%的输出吞吐量提升。
  3. 开箱即用的生产级部署:MTP作为即插即用的功能被完全集成入SGLang,使得开源推理生态系统能够获得以往仅限于闭源专有系统的高级解码能力,帮助AI从业者在NVIDIA GPU上轻松实现大规模、高性能的模型推理服务。

背景知识

传统自回归解码的局限性:传统的自回归解码(Autoregressive decoding)每次只能生成一个Token,且该Token的生成依赖于所有之前的Token。这种串行处理过程极大地限制了并行性,并导致了生成速度的瓶颈。

多Token预测(MTP)的核心概念:MTP是一种投机解码(Speculative decoding)技术。它通过使用一个轻量级的草稿模型(Draft model)快速提出多个未来的候选Token,随后由全尺寸的目标模型(Target model)在单次前向传播中对这些Token进行并行验证,从而加速生成过程。

方法细节

MTP的两阶段生成机制:MTP将文本生成过程划分为起草(Drafting)和验证(Verification)两个阶段。

起草阶段(Drafting):轻量级的草稿模型在一次快速的前向传播中,预测出一个或多个包含$n$个Token的短序列候选。以生成一个序列候选为例,其具体步骤如下:

  1. 目标模型首先生成当前的上下文前缀,例如“Today is a sunny”。
  2. 目标模型的扩展/预填充(extend/prefill)阶段生成第一个Token“day”。
  3. 草稿模型的扩展/预填充阶段生成第一个草稿Token“and”。
  4. 草稿模型通过自回归解码迭代生成额外的三个草稿Token“it’s so hot”。在这个示例中,$n=4$,完整的草稿序列为“and it’s so hot”。

验证阶段(Verification):全尺寸的目标模型随后并行验证所有的$n$个草稿Token。目标模型会接受与其自身输出相匹配的最长前缀,并在必要时对剩余的部分重新采样(Resampling)。以$n=4$为例,具体的验证过程如下:

  1. 目标模型在扩展/预填充阶段后生成初始Token:“day”。
  2. 草稿模型在扩展/预填充后推测下一个Token,并通过自回归解码再推测3个Token,得出序列:“and it’s so hot”。
  3. 目标模型对整个序列进行验证:它同意前缀“and it’s”;但拒绝了后续的“so hot”,并将其重新采样为“very”。
MTP流程图
MTP流程图

MTP的加速原理与接受长度:MTP实现加速的关键在于并行化。至关重要的是,验证步骤在GPU上是完全并行化的,它用单次的并行验证通道取代了$n$个串行的解码步骤。MTP的效率取决于每次验证步骤中被接受的草稿Token数量,这一核心指标被称为接受长度(Acceptance length)。例如,平均接受长度为2.4意味着每次解码平均跳过了2.4个步骤,这在长序列生成中会带来可观的累积加速比。此外,MTP不会在生成质量或确定性上做出妥协。由于每个投机Token都由同一个全尺寸模型验证和批准,其输出与标准解码完全一致,无需任何近似处理或微调。

与SGLang高级特性的无缝集成:MTP功能已经与SGLang的多种高级特性完全集成,具体支持的特性列表如下:

  • 数据并行注意力机制(DP Attention)
  • 专家并行负载均衡器(EPLB)
  • DeepEP MoE
  • 双批次重叠(Two Batch Overlap)
  • 预填充-解码分离(Prefill-Decode Disaggregation)
  • CUDA图(CUDA Graph)
  • 多种注意力后端(Various Attention Backends)

补充细节

MTP参数配置最佳实践:要在SGLang中启用MTP,首先需要在配置中开启该功能。

草稿Token数量(draft_token_num)的设置策略:

  1. 初始推荐配置:建议将draft_token_num设置为2。这是一个平衡且低风险的选择,能够在大多数工作负载下提供可靠的性能提升。
  2. 算力充足时的配置:对于有额外GPU算力余量的系统,可以将draft_token_num增加到4甚至更大,以进一步提升吞吐量。不过,收益会随着系统维持Token接受率的能力而逐渐递减。
  3. 算力饱和时的配置:如果GPU已经在处理大批次(Large batches)或接近满载运行,将草稿大小保持在2或3通常更为高效,这可以避免引入额外的计算负载。

基于接受率的动态调整:建议通过日志监控接受率(Acceptance rates)来随着时间的推移微调此参数。如果观察到平均接受长度持续高于2,则说明有空间尝试更长的草稿序列;但如果接受率开始下降,则应考虑调低该参数,以使系统保持在舒适的运行区间内。

实验环境

  • 测试模型:DeepSeek V3。
  • 硬件配置:

    • 小规模部署:共16张 NVIDIA H200 GPU(部署为2个解码节点)。
    • 大规模部署:16节点集群,共128张 NVIDIA H200 GPU(其中配置为4个预填充节点,12个解码节点)。
  • 软件与系统配置:SGLang推理框架,启用了预填充-解码分离(Prefill-Decode disaggregation)和大规模专家并行(Large-Scale EP)。包含SGLang v0.4版本引入的重叠调度(Overlap scheduling)功能。

  • 工作负载参数:
    • 小规模部署:每个Rank并发请求数为2;输入序列长度为65,536 Tokens;输出序列长度为4,096 Tokens。
    • 大规模部署:每个Rank并发请求数为128;输入序列长度为2,000 Tokens;输出序列长度为100 Tokens;解码配置为topk = 1,step size = 1,draft_token_num = 2。

实验结果

本次评估设计了两个案例研究,分别突出了在小规模和大规模部署环境下的性能提升。

实验动机:小规模部署的配置基于某知名生成式AI公司的生产需求。该公司运营着面向开发者的延迟敏感型产品,要求每个Rank至少达到 60.4 tokens/sec 的持续输出以满足应用级的服务水平协议(SLA)。大规模集群设置则用于评估系统在更重负载下的可扩展性。

实验一:小规模部署测试

  • 实验内容:在16张H200 GPU环境下,对比基线(无MTP、无重叠调度)、仅启用重叠调度、以及启用不同窗口大小MTP的吞吐量表现。
  • 实验结果与分析:

    • 在无MTP且无重叠调度的基线配置下,系统吞吐量为 51.0 tokens/sec/rank。
    • 仅使用重叠调度时,吞吐量达到 60.4 tokens/sec/rank,刚好满足了生产阈值要求。
    • 开启MTP(3-token窗口,topk=1)时,吞吐量飙升至 81.5 tokens/sec/rank,平均接受长度为 2.18 tokens。
    • 开启MTP(4-token窗口,topk=1)时,吞吐量达到 82.0 tokens/sec/rank,平均接受长度为 2.44 tokens。
  • 结论:相比完全不使用重叠调度和MTP的基线,MTP使输出吞吐量提升了超过60%。这证明了即使在并发度较低的小型集群中,MTP也能在有限的GPU资源预算下带来显著的性能提升。

小规模部署吞吐量图表
小规模部署吞吐量图表

详细数据对比见下表:

重叠调度 MTP 吞吐量 (tokens/sec) 每 Rank
❌ ❌ 51.0 (基线)
✅ ❌ 60.4 (+20.4% ↑)
❌ ✅ 3-token 81.5 (+59.8% ↑)
❌ ✅ 4-token 82.0 (+60.8% ↑)

实验二:大规模部署测试

  • 实验内容:在128张H200 GPU集群上,结合大规模专家并行(EP)和预填充-解码分离架构,对比启用MTP与基线配置在极高并发(128并发/Rank)下的吞吐量。
  • 实验结果与分析:相比基线(无重叠调度且无MTP),启用MTP的解码吞吐量提升了 14.2%。
  • 结论:证实了在类似生产环境的高吞吐量、大规模集群负载下,MTP依然能够提供可衡量的性能增益。
大规模部署吞吐量图表
大规模部署吞吐量图表

结论

通过将多Token预测(MTP)技术深度集成到SGLang推理框架中,本文成功打破了传统自回归解码的串行瓶颈。在不牺牲生成质量的前提下,MTP为DeepSeek V3等前沿大语言模型带来了最高60%的输出吞吐量提升,并在小规模延迟敏感场景与大规模高吞吐场景中均验证了其卓越的性能表现。

未来工作展望:

  1. 大规模优化:研发团队将继续优化大规模MTP部署的性能,重点关注跨多节点系统的调度效率和内存带宽利用率。
  2. 重叠调度兼容性:当前的MTP实现尚未支持重叠调度(Overlap Scheduling)。目前相关集成开发正在进行中,预计两者结合后将带来进一步的性能飞跃。