When RL Meets Adaptive Speculative Training: A Unified Training–Serving System

作者/机构: Junxiong Wang, Fengxiang Bie, Jisen Li, Zhongzhu Zhou, Zelei Shao, Qingyang Wu, Yinghui Liu, Yubo Wang, Avner May, Sri Yanamandra, Tri Dao, Percy Liang, Ce Zhang, Ben Athiwaratkun, Shuaiwen Leon Song, Chenfeng Xu, Xiaoxia Wu (http://together.ai)

A1 主要贡献

推测式解码(Speculative decoding)能够显著加速大语言模型(LLM)的服务,但目前的部署大多将推测器(speculator)的训练与服务解耦,将其视为一个独立的离线建模问题。本文指出这种解耦的范式引入了严重的部署和适应延迟:
1. 高昂的上线时间(time-to-serve):因为推测器在部署前必须经过长时间的离线训练。
2. 延迟的效用反馈:因为真正的端到端解码加速效果只有在训练和部署后才能知晓,而不能仅从接受率(acceptance rate)可靠地推断,这受到模型架构、多样的提示工程和系统级开销的影响。
3. 领域漂移导致的性能下降:当目标模型被用于新领域时,推测器会变得陈旧且效率降低。

为了解决这些问题,本文提出了 Aurora,一个统一的训练-服务系统,它通过直接从实时的推理追踪数据中持续学习推测器来形成闭环。Aurora 将在线推测器学习问题重构为一个异步强化学习问题:被接受的 token 提供正反馈,而被拒绝的推测器提议则提供隐式的负反馈,这些负反馈继承了在线流量的失败信号。该设计集成了基于 SGLang 的推理服务器和异步训练服务器,支持在不中断服务的情况下对推测器进行热插拔更新。

核心贡献
1. 实现 Day-0 服务和实时可观测性:Aurora 框架支持从零开始部署推测器。即使是一个冷启动的推测器也可以被安全地引入,并立即开始收集在线服务数据,在实际场景中快速适应,将推测式解码从“先训练后服务”的工作流转变为“服务驱动训练”的飞轮,实现了实时的反馈和可观测性。
2. 快速适应新流量并直接缓解分布失配:通过在线数据收集和领域控制的缓冲区,Aurora 显著减少了推测器训练方式与使用方式之间的分布差异。实验结果显示,相较于一个训练良好但静态的推测器,Aurora 实现了 1.25 倍的性能提升。
3. 降低基础设施内存占用和成本:通过在实时流量上训练推测器,Aurora 无需大规模收集和重放目标模型激活值的流程。
4. 面向多样化用户需求的推测式解码系统:Aurora 支持异构请求混合(例如,在线流量加上离线语料库),并利用受强化学习系统启发的分布式设计模式扩展到大型 GPU 部署。
5. 算法无关性:该设计不与特定的推测器算法绑定,能够与多种推测式解码变体清晰地接口。这种通用性使得统一的训练-服务循环可以作为一个即插即用的层,放大未来推测器算法的进展。


图1:Aurora。一个统一的训练-服务框架,用于在线推测式训练,具有异步、RL风格的更新。一个生产推理服务器使用固定的目标(验证器)和轻量级的草稿模型(推测器)执行推测式解码,在验证期间接受或拒绝提议的token。服务追踪数据——包括接受和拒绝的前缀——被流式传输到一个数据缓冲区和训练管道中。一个独立的训练服务器从收集到的离策略数据中持续更新推测器,并定期将异步模型更新热插拔到推理服务器中,而不中断请求。底部:(左)随时间变化的每个请求的吞吐量,在每次异步更新后呈现阶梯式变化;(右)连续训练期间的接受长度统计数据,显示接受率随时间提高(或维持)。

A3 背景知识

2.1 推测式解码

推测式解码的基本原理。推测式解码(Speculative decoding)【索引14,Fast inference from transformers via speculative decoding,2023】是一种快速推理技术,它将一个轻量级的草稿模型(draft model)与一个更强的目标模型(target model)配对。草稿模型会预先提议多个token,然后由目标模型进行验证,并接受最长的匹配前缀,这样系统就可以跳过许多昂贵的目标模型解码步骤,同时保持输出质量。在实践中,加速效果取决于草稿模型的提议与目标模型的匹配程度(即接受率)。在此基础上,MTP 和 Eagle【索引15,Eagle: Speculative sampling requires rethinking feature uncertainty,2024,International Conference on Machine Learning】【索引16,Eagle-2: Faster inference of language models with dynamic draft trees,2024】【索引17,Eagle-3: Scaling up inference acceleration of large language models via training-time test,2025】是最流行的框架,因为它们利用了目标模型的隐藏状态,这可以提高接受率。推测式解码的一个关键特性是,其预期的加速效果主要由平均接受率 α 和前瞻步数 γ(每次迭代提议的草稿token数量)决定。在独立同分布(i.i.d.)接受的假设下,每个验证步骤产生的预期token数量为:

$$\mathbb{E}[L]=\sum_{i=0}^{\gamma} \alpha^i=\frac{1-\alpha^{\gamma+1}}{1-\alpha}.$$

加速效果的估算。如果我们还考虑草稿模型的开销,通过成本比率 $c = T_q/T_p$(草稿模型步骤时间与目标模型步骤时间之比)来衡量,那么预期的实际时间改进因子大约为:

$$\text{Speedup} \approx \frac{\mathbb{E}[L]}{1+\gamma c} = \frac{1-\alpha^{\gamma+1}}{(1-\alpha)(1+\gamma c)}.$$

实践中的挑战。然而,在实践中,准确估计 c 并非易事:它取决于服务堆栈(内核、精度、批处理/调度和硬件),并且在训练与服务不匹配的情况下可能会有很大差异。更深的草稿模型可以实现更高的接受率但运行更慢,而较浅的模型速度更快但通常接受率较低。在实践中,很难权衡这些因素并选择一个模型。因此,大多数团队会训练多个草稿模型,但最终只选择一个。相比之下,我们的系统能够直接比较加速效果,因为它是在线运行的。

2.2 在线推测式解码

与在线推测式解码(OSD)的比较。在线推测式解码(Online Speculative Decoding, OSD)【索引18,Online speculative decoding,2023】提出了一种引人注目的推测式解码变体。我们的工作受其启发,并在此强调关键差异。从实用性和部署角度看,Aurora 是一个闭环的统一训练-推理系统,它将推测式解码作为一个持续改进的服务来运作。具体来说,它闭合了三个反馈路径:(1)推测器在线上根据流式验证器输出进行训练;(2)新训练的推测器被推送回服务堆栈;(3)实时的服务信号被反馈到训练中,以将学习集中在能产生可衡量效益的地方。相比之下,OSD 实际上是开环的:它通过从目标模型在线蒸馏来解决(1),但没有通过(2)和(3)来闭合循环。弥合这些缺失的环节并非易事,需要跨系统和算法的协同设计。

Aurora解决的系统与算法挑战
1. 在系统方面,闭合循环意味着在严格的延迟服务等级目标(SLO)下,使模型更新高效且稳健。这需要 (a) 一种将草稿模型部署到推理服务器而不中断推理的同步策略;(b) 低开销的权重和训练数据传输机制,以使更新和重放不会成为瓶颈;以及 (c) 一个能够高效处理流式在线策略数据、控制陈旧性并在请求分布变化时保持稳定的训练管道。
2. 在算法方面,关键问题是哪种训练信号与生产效用相匹配。OSD 使用来自目标激活值的序列级蒸馏,但一旦架构不匹配和系统开销占主导地位,该信号与端到端加速效果的关联性可能很差。相反,我们把训练表述为异步强化学习:接受/拒绝的结果提供了直接反映实时流量和验证器在线行为的奖励。这将优化与在线接受动态以及由此产生的吞吐量/延迟对齐,使得闭环在部署中既实用又有效。

2.3 异步强化学习和在线训练系统

异步强化学习范式。异步强化学习已成为扩展RL后训练的实用系统范式,其动机在于,在长且重尾的轨迹中,生成 rollout 通常是主要瓶颈,同步管道会因掉队者(stragglers)和GPU空闲时间而受影响【索引26,Beat the long tail: Distribution-aware speculative decoding for rl training,2025】【索引31,Angles don’t lie: Unlocking training-efficient rl through the model’s own signals,2025】。通过将推理/rollout生成与训练/优化解耦,异步RL消除了全局同步障碍,并提高了端到端训练效率,通常通过一个服务器化的 actor 设置,持续产生经验,而学习器在后台更新。AReaL【索引11,Areal: A large-scale asynchronous reinforcement learning system for language reasoning,2025】通过提供一个明确将训练和推理分为独立组件的开源框架,以及在异步下保持训练性能的算法设计,体现了这一范式。同时,像 slime【索引29,SLiME: A post-training framework for reinforcement learning scaling,2024】和 miles【索引25,MiLeS: A high performance rl framework,2025】这样的开源系统建立在高性能开源推理引擎SGLang之上,专注于提供一个可扩展的 rollout 引擎,以支持服务器化生成和大规模后训练工作流的高效数据收集。

Aurora与现有系统的区别。然而,这些系统主要为RL吞吐量而非生产级的推测式解码进行优化,后者的目标是在严格的SLO下实现端到端的服务效率。因此,它们没有将推测器的在线学习循环视为一流的闭环服务组件,也没有解决关键的部署挑战,如领域漂移/验证器漂移下的草稿-验证器不匹配和非破坏性同步。

A2 方法细节

我们提出了一个统一的、生产就绪的系统Aurora,它将推测式推理和在线训练紧密集成到一个闭环中。与传统的推测式解码系统在部署前需要大量离线预训练不同,我们的框架支持“第0天”(day-0)服务:推测器可以从头开始部署,即使完全未经训练,也能在实时流量上就地快速适应。这从根本上将推测式解码从一个“先训练后服务”的管道转变为一个“服务驱动训练”的飞轮。

3.1 系统架构

系统架构。如图1所示,我们的框架由两个主要的解耦组件组成:一个推理服务器(Inference Server)和一个训练服务器(Training Server)。推理服务器负责处理用户请求。它运行一个基于SGLang【索引37,Sglang: Efficient execution of structured language model programs,2024,Advances in neural information processing systems】的推测式解码引擎,该引擎使用一个目标模型和一个草稿模型。对于每个请求,草稿模型会提议一个token序列,然后由目标模型并行验证。接受和拒绝的token结果都会被流式传输到一个分布式数据缓冲区。为了支持EAGLE【索引15,Eagle: Speculative sampling requires rethinking feature uncertainty,2024,International Conference on Machine Learning】【索引16,Eagle-2: Faster inference of language models with dynamic draft trees,2024】【索引17,Eagle-3: Scaling up inference acceleration of large language models via training-time test,2025】风格的训练,隐藏状态也会被发送到数据缓冲区。

异步训练与热插拔。训练服务器异步运行。它从数据缓冲区中获取成批的训练数据,并对草稿模型的副本执行梯度更新。一旦一个新的、改进后的推测器准备就绪,其权重就会被异步推送回推理服务器。这种更新是一种热插拔(hot-swap),意味着推理服务器可以在没有停机或服务中断的情况下开始使用新的草稿模型。

3.2 系统实现

系统实现细节。Aurora旨在实现生产环境中经济高效的部署。我们实现了批处理RPC传输(使用带有TensorPipe后端的“torch.distributed.rpc”),并配备了可扩展的CUDA内存段,以防止碎片化并确保高效的GPU到GPU通信。训练循环维护一个线程安全的已传输数据缓存,在执行反向传播前累积样本直到达到微批次大小。可选的检查点功能使得更新后的草稿模型能够热插拔回推理服务器而无需停机,从而创建了一个能够实时响应生产工作负载中分布变化的闭环自适应系统。


图2:树状注意力机制的图示。它支持对整个推测树进行高效的批处理计算,包括接受(绿色)和拒绝(红色)的token。

高效的树状注意力机制。我们定制了训练掩码,因为草稿模型在每一步都可能犯错,而对每个验证步骤单独进行训练的成本过高。为了高效处理推测式解码结果的复杂分支结构,我们采用了一种专门的树状注意力机制(Tree Attention),如图2所示。通过构建一个遵循推测树因果结构的自定义注意力掩码,我们可以在单次批处理的前向和后向传播中处理所有接受和拒绝的分支。这使得从整个推测树中学习在计算上变得可行。

系统设计的核心挑战。这种解耦和异步的架构是我们的系统Aurora的关键。它允许训练在专用资源上连续运行,利用大批量和高效的数据处理,同时推理保持响应迅速并为低延迟服务进行优化。然而,要使这个设计奏效远非易事。核心挑战是构建一个可靠的闭环:我们需要一个有效的推测器更新方案,该方案基于在线策略的服务追踪数据。一旦产生了一个新的推测器,我们还需要一个谨慎的同步策略,将其部署到生产中而不会干扰延迟或吞吐量,即选择一个合适的同步频率并最小化同步延迟,以便更新既及时又无干扰。最后,系统必须能迅速克服领域漂移:随着请求组合的演变,推测器可能在数小时内就变得陈旧。为了克服这些挑战,我们将在线推测式训练重新表述为异步强化学习问题。

3.3 将在线推测式训练视为异步强化学习

将在线推测式训练视为异步强化学习。我们可以将在线推测式解码视为一个异步强化学习(RL)系统,以揭示当训练直接嵌入到实时服务中时出现的学习信号和系统约束。在这个表述中,草稿模型是策略π,而目标模型加上验证器则实现了环境。每个推测步骤都构成一个短暂的回合(episode):策略提议一个候选续写的树,验证器接受一个前缀,其结果提供了结构化的反馈。接受的token对应正奖励,而拒绝的token提供零/负奖励。因此,最大化期望回报等同于最大化接受长度,这直接决定了解码的加速效果。

系统实现与同步策略。我们通过AReaL【索引11,Areal: A large-scale asynchronous reinforcement learning system for language reasoning,2025】中的异步actor-learner设计来实现这一构想。如图1所示,SGLang副本作为actor,持续地从实时请求中生成经验。一个分布式缓冲区聚合了接受和拒绝的分支。一个独立的多GPU学习器异步地更新草稿模型,并定期将更新后的权重推送回服务副本。重要的是,我们采用了一种惰性同步策略(lazy synchronization policy):我们容忍有限的陈旧性,以避免中断推理的关键路径。

从接受和拒绝中学习。验证过程产生的监督信号比仅从接受的token进行模仿学习要丰富得多。我们用两种互补的信号来训练草稿模型:接受损失(模仿学习):对接受的token计算交叉熵,鼓励草稿模型重现验证器批准的续写。拒绝损失(反事实反馈):被拒绝的分支指明了策略不应该提议的内容。通过丢弃采样(Discard Sampling),我们应用一个基于KL散度的目标(由λdiscard加权),将概率质量从不正确的预测上移开。联合优化这些目标可以产生更密集的训练信号:接受教会草稿模型在当前流量下哪些是成功的,而拒绝则对否则会被丢弃的不匹配提议提供了即时的负反馈。总损失是两个KL散度项的加权组合:

$$\mathcal{L}=\mathbb{E}_{x \sim p_{\text {accept }}}\left[\operatorname{KL}\left(p_{\text {target }} \| p_{\text {draft }}\right)\right]+\lambda_{\text {discard }} \mathbb{E}_{x \sim p_{\text {discard }}}\left[\operatorname{KL}\left(p_{\text {target }} \| p_{\text {draft }}\right)\right]$$

损失函数详解。在这里,第一项训练草稿模型在接受的序列上模仿目标模型。第二项,即我们新颖的丢弃采样(Discard Sampling),在拒绝的序列上训练模型,明确地教它纠正错误。我们对被丢弃的token应用top-k过滤,将训练重点放在高概率的不一致上,从而减少梯度噪声。

A7 补充细节

4 新功能研究:Day Zero 支持推测器

4.1 在线流量模拟

在线流量模拟设置。我们将一个固定的提示语料库视为一个推理请求流,而不是一个监督训练数据集。请求被顺序消费以模拟实时服务流量,并且不使用任何真实的助手响应。该数据流包含40k个提示,涵盖五个领域(见表2),包括数学推理【索引9,Training verifiers to solve math word problems,2021,arXiv preprint arXiv:2110.14168】、文本到SQL【索引35,Spider: A large-scale human-labeled dataset for complex and cross-domain semantic parsing and text-to-sql task,2018,Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing】、代码生成【索引13,Codesearchnet challenge: Evaluating the state of semantic code search,2019,arXiv preprint arXiv:1909.09436】、金融【索引2,finance-alpaca dataset,https://huggingface.co/datasets/gbharti/finance-alpaca】和通用对话指令【索引1,chatbot_instruction_prompts,https://huggingface.co/datasets/alespalla/chatbot_instruction_prompts】。这种构成反映了现实部署场景,其中服务流量表现出异构且不断变化的任务分布 。

流量模式。我们评估了两种流量模式:(i)有序流,其中请求按领域分组以引发突发的分布变化;(ii)混合流,其中提示被随机打乱以近似平稳流量。这些设置对异步策略更新的不同方面进行了压力测试,使我们能够研究更新延迟和服务稳定性之间的权衡。每个请求都使用EAGLE-3【索引17,Eagle-3: Scaling up inference acceleration of large language models via training-time test,2025】推测式解码进行处理,产生接受的前缀和拒绝的分支。这些结果是进行在线更新的唯一学习信号。


图3:混合流。一个未经训练的推测器的Day-0自适应。(a) 接受长度从1开始迅速增加,并与预训练基线收敛。(b) 每个请求的吞吐量(定义为(T_input + T_output)/t_request,其中T_input和T_output是输入和输出的token数,t_request是端到端延迟)最初会下降,但随着推测器的自适应而恢复,展示了服务驱动训练飞轮的有效性。(c) 在训练好的模型上继续微调可以取得更好的结果。

模型和配置。我们采用EAGLE-3【索引16,Eagle-2: Faster inference of language models with dynamic draft trees,2024】作为我们的推测式解码框架。我们的目标模型是Qwen3-8B【索引33,Qwen3 technical report,2025】,它使用了4倍小的语言模型头(词汇量32k),并在超过20万个数据集上进行了训练。对于Day-0实验,草稿模型使用随机权重进行初始化。

对比方法。我们比较了四种配置:(i)没有任何推测的推理系统,(ii)带有静态、预训练推测器的推理系统,(iii)Aurora(从零开始):我们的统一框架从一个随机初始化的推测器开始,以及(iv)Aurora(已训练):我们的统一框架从静态推测器初始化。


图4:有序流。一个未经训练的推测器的Day-0自适应。(a) 接受长度从1开始迅速增加,收敛并有时甚至超过预训练基线。(b) 吞吐量(定义见B.1节)最初会下降,但随着推测器的自适应而恢复,展示了服务驱动训练飞轮的有效性。(c) 在训练好的模型上继续微调起初会下降,但经过一些训练后取得了更好的结果。

结果与启示。图3和图4展示了在两种现实流量模式下,我们的服务驱动训练飞轮的显著效果。在混合流量场景中(请求被随机打乱以近似平稳分布),系统实现了更强的性能:接受长度达到3.08(超过了静态基线的2.63和预训练后微调基线的2.99),吞吐量稳定在302.3 tokens/s。在领域漂移场景中(请求按领域分组以引发突发的分布变化),未经训练的推测器从零接受长度开始,在约10,000个请求内收敛到2.46,几乎与预训练基线的2.57持平。吞吐量稳定在295.6 tokens/s,与静态推测器的288.8 tokens/s相当。关键是,这种适应在两种场景中都是在实时服务期间发生的。混合流量的结果尤其引人注目,表明从零开始的在线训练可以超过精心预训练的推测器的性能。这从根本上挑战了推测式解码需要大量离线预训练的传统观念。我们的统一框架证明,一个完全未经训练的推测器可以在第零天部署,并仅通过在线自适应就达到生产就绪状态,从而消除了数月的离线训练周期,并能够在可能不存在预训练数据的新领域中立即部署。

4.2 同步策略的权衡研究

同步策略的权衡。尽管该架构类似于异步强化学习,但主要的系统权衡是不同的:适应速度与服务稳定性。在经典的异步RL中,频繁的同步是有益的,因为它减少了策略的陈旧性并提高了样本效率。然而,在在线服务中,同步具有一阶成本:它会中断推理路径、使缓存失效并引入延迟抖动。关键问题变成了:我们能多频繁地刷新草稿策略,而同步开销又不会抵消服务增益?


图5:推测器异步策略研究。更频繁的策略刷新可以改善漂移后的适应性(更高的接受长度),但由于同步开销可能会降低服务吞吐量。一个适度惰性的调度(例如,Trained w Async 48)提供了一个强大的帕累托点,既保留了吞吐量,又保留了大部分的适应性优势。

实验设置与结果。为了量化这一点,我们在一个具有突发领域漂移的相同请求流下,扫描了从激进(每48个请求)到惰性(每1600个请求)的策略更新间隔。图5显示了明显的权衡。激进的更新能更快地恢复接受长度,但由于频繁的权重同步,吞吐量最低。相比之下,一个适度惰性的调度(每80个请求)实现了几乎相同的接受长度恢复,同时提供了最佳的整体吞吐量,甚至超过了基线。

5 推测算法探索

本节研究在现实的服务动态下,哪种训练目标和在线更新策略对推测式解码最为重要。在所有设置中,我们发现了一个一致的模式:一旦推测器在在线策略上进行训练并与当前流量保持一致,简单的token级目标就能捕捉到大部分可获得的增益,而更复杂的监督(例如,树状结构损失)在领域漂移下提供的回报递减。

5.1 基线和变体

基线和变体。我们比较了一个静态部署的推测器与一系列在线自适应方法(所有方法都在相同的训练-服务循环内):(1) 冻结草稿(静态基线)。一个预训练的推测器被部署用于推测式解码,没有在线更新,代表了标准的非自适应设置。(2) Aurora (FKL)。仅使用来自验证的接受token,并采用前向KL目标进行在线微调。(3) Aurora (RKL)。仅使用来自验证的接受token,并采用反向KL目标进行在线微调。(4) Aurora (RKL + NTP)。方法(2)的增强版,在接受的token上增加了一个辅助的下一token预测损失,以加强训练信号。(5) Aurora (w discard)。将训练扩展到包括来自被拒绝的推测分支的token,让草稿模型直接从其错误中学习。

评估指标。我们使用两个指标在多个模型和数据集配置上评估这些方法:推测接受长度(每个验证步骤平均接受的草稿token数量)。


图6:Qwen-8BInstruct在推理请求上的移动平均推测解码接受长度。在线微调显著提高了接受长度,优于冻结基线。在领域漂移下,使用不同策略进行训练只产生微小的差异。

在线微调的显著效果。在所有设置中,在线微调始终比静态基线获得更高的吞吐量,并且随着用于训练的token数量增加,增益也随之增加。相比之下,结合NTP和被丢弃的token只带来了边际的额外改进。

在领域漂移下丢弃token的边际效益。在所有设置中,主要的增益来自于闭合循环:在线策略的在线更新在流量变化后显著提高了接受长度并稳定了吞吐量。相比之下,添加被丢弃的token只提供了增量改进,并且通常与最简单的在线目标没有明显区别。图6说明了这种行为:接受长度表现出与分布变化一致的明显区域变化;冻结的基线性能下降并保持最低,而所有在线变体都遵循一个更高的平台,并在每次变化后恢复。在在线变体集群中,树/丢弃变体与RKL(+NTP)基本重叠,表明在这种在线设置下,超越token级自适应的空间有限。

丢弃token起作用的场景。在之前前瞻步数=5的实验中,预训练的推测器已经达到了相对较高的接受长度(例如,约3.6),为更具表现力的监督留下有限的改进空间。这启发了一个简单的假设:当服务设置需要更长的接受前缀时,被丢弃的token会变得更有用。为了测试这一点,我们将前瞻步数增加到10。


图7:Llama-3.1-8B-Instruct【索引10,The llama 3 herd of models,2024】在lookahead=10的编码任务上的移动平均推测解码接受长度。在丢弃的token上进行训练为Llama模型带来了额外的增益,但不同的丢弃策略(top-k = 0, 10, 50)只产生微小的差异,表明top-k策略节省了内存并保持了性能。

长前瞻场景下的结果。如图7所示,在这种更激进的前瞻设置下,从被丢弃的token中学习提供了明显更强的训练信号,并比仅接受token的微调产生了更大的增益。


图8:Aurora增加了推测接受长度并提升了吞吐量,在较小的批次大小时加速效果更明显。这里,BS4和BS12分别表示每批次有4个和12个并发请求。

Aurora在不同批次大小下的性能。我们将批次大小从12减少到4来测试Aurora的有效性。如图8所示,这些追踪图显示,在线训练的推测式解码Aurora增加了接受长度,并将其转化为更高的端到端吞吐量——但在较小的批次大小时,吞吐量的好处要大得多。在上图中,Aurora在整个流量阶段都保持了比静态草稿模型更高的推测接受长度(在请求分布变化时有可见的下降,随后重新适应),而静态基线保持平坦。接受的前缀转化为更高的tokens/s,但对于批次大小为4的情况,相对提升大于批次大小为12的情况:在小批次下,基线解码更多地受每步开销和较差的硬件利用率主导,因此通过更长的接受前缀跳过目标模型步骤会带来很大的比例增益;在较大批次下,目标模型已经得到了更好的分摊/效率,而推测开销(草稿+验证+同步)在管道中占的比例更大,即使接受率仍然提高,净加速效果也会缩小。

6 在前沿开源模型上的可扩展性

我们评估了我们的系统在两个最新的开源前沿模型上的可扩展性,证明了随着模型大小、架构复杂性和上下文长度的增加,Aurora 仍然能提供稳定的性能增益。
- MiniMax M2.1【索引6,Minimax-m1: Scaling test-time compute efficiently with lightning attention,2025】【索引21,Minimax-m2.1,2024】,发布于2025年12月23日,是一个229B参数的仅解码器Transformer-MoE LLM,拥有62层和一个256专家的路由器(每个token top-8)。
- Qwen3-Coder-Next【索引24,Qwen3-coder-next technical report,2026】,发布于2026年2月2日,是一个MoE设计,总参数80B,激活参数3B,48层,原生上下文长度262K。其混合骨干结构交错了Gated DeltaNet【索引34,Gated delta networks: Improving mamba2 with delta rule,2024】和Gated Attention块,每个都与MoE层配对(512专家,top-10激活)。

MiniMax-M2.1实验。我们在4个H200 GPU上使用张量并行(TP)=4,以FP8精度服务MiniMax-M2.1。对于Aurora(从零开始),我们额外分配一个GPU用于在线推测器调优。图9比较了无推测基线和Aurora(从零开始)。Aurora将平均接受草稿长度提高到2.8,并将其转化为更高的每个请求吞吐量(约1.45倍加速),表明学习到的推测器即使在大型专家路由下也保持有效。


图9:MiniMax M2.1。顶部:随时间变化的接受草稿长度。底部:随时间变化的每个请求的吞吐量。Aurora(从零开始)将接受长度增加到2.8,并将其转化为相对于无推测基线1.45倍的吞吐量(BS4)增益。

Qwen3-Coder-Next实验。同样,Qwen3-Coder-Next在4个H200 GPU上以FP8精度服务,使用TP=4和专家并行(EP)=4。Aurora(从零开始)再次使用一个额外的GPU进行在线推测器调优。图10比较了无推测基线,并报告了接受草稿长度和每个请求的吞吐量。Aurora持续将接受长度提高到3以上,并在这个混合TP/EP部署上产生了1.23倍的吞吐量提升。


图10:Qwen3-Coder-Next。顶部:随时间变化的接受草稿长度。底部:随时间变化的每个请求的吞吐量。我们丢弃了前1000个预热步骤,因为混合部署在初始化期间表现出短暂的吞吐量不稳定性。尽管存在这种可变性,Aurora(从零开始)将平均接受草稿长度提高到3,并比无推测基线(在最后10k步上平均)提供了1.21倍的吞吐量改进。所有结果均使用批处理大小为8。

不同批次大小下的性能
- MiniMax M2.1:如表1所示,我们使用最终的Aurora检查点在一个留出数据集上评估了不同批次大小下的端到端服务吞吐量。Aurora在小到中等大的批次大小(32)下提高了吞吐量,加速比从1.57倍到1.25倍不等。
- Qwen3-Coder-Next:如表1所示,我们使用最终的Aurora检查点在一个留出代码数据集上测量了不同批次大小下的端到端服务吞吐量。Aurora在小到中等批次大小下提供了最大的增益(在批次大小为1时高达1.51倍),随着批次大小的增加,收益递减;在批次大小为32时,验证开销占主导,推测式解码略慢于基线。

表1:不同批次大小下的端到端吞吐量。TPS = 每秒token数(定义见附录B.1)。MiniMax M2.1 (FP8) 使用前瞻4;Qwen3-Coder-Next (FP8) 使用前瞻5。完整结果见表4、5。

A4 实验环境

  • 数据集:使用一个包含40k提示的混合领域语料库模拟在线请求流,涵盖数学推理(GSM8K)、文本到SQL(Spider)、代码生成(CodeSearchNet)、金融(Finance-Alpaca)和通用对话指令(chatbot_instruction_prompts)。
  • 模型架构

    • 目标模型:Qwen3-8B、Llama-3.1-8B-Instruct、MiniMax M2.1(229B MoE)、Qwen3-Coder-Next(80B MoE)。
    • 草稿模型(推测器):采用EAGLE-3风格,使用比目标模型小4倍的语言模型头。
  • 硬件配置

    • Qwen-8B和Llama-8B实验在H100 GPU上进行。
    • MiniMax M2.1和Qwen3-Coder-Next实验在H200 GPU上进行。
    • 采用分离式GPU架构,例如使用4个GPU进行张量并行推理,另有1个专用GPU进行草稿模型训练。
  • 软件配置

    • 推理引擎:SGLang。
    • 通信:使用torch.distributed.rpc与TensorPipe后端进行GPU间数据传输。
    • 优化器:AdamW,权重衰减为0.0,梯度裁剪范数为0.5。
    • 学习率:微调时为$1 \times 10^{-5}$,从零训练时为$1 \times 10^{-4}$,采用恒定学习率。
    • 精度:BF16计算,FP32主权重。

A4 实验结果

  • Day-0部署与适应性:Aurora能够从一个随机初始化的推测器开始,在实时流量上进行训练,其性能能够迅速达到甚至超过一个经过精心离线预训练的静态推测器。该系统能够有效适应由不同任务领域引起的分布漂移,快速恢复并提升接受长度和吞吐量(图3,图4)。
  • 同步策略权衡:实验表明,过于频繁的推测器权重更新(例如每48个请求)会因同步开销而损害整体吞吐量,而过于稀疏的更新则会导致适应速度变慢。一个适度“惰性”的同步策略(例如每80个请求更新一次)在适应速度和吞吐量稳定性之间取得了最佳平衡(图5)。
  • 算法探索:研究发现,实现训练-服务闭环是在线推测式解码获得性能提升的最关键因素。简单的token级目标函数(如反向KL散度)就能捕获大部分增益。更复杂的学习信号,如从被拒绝的token中学习(discard sampling),在需要更长前瞻(lookahead=10)的场景下才能显现出更明显的优势(图6,图7)。
  • 批次大小的影响:Aurora带来的吞-吐量提升在小批次(Batch Size=4)下更为显著。这是因为小批次下基线推理的硬件利用率较低,开销占比较大,通过推测式解码跳过目标模型步骤能带来更大的相对增益(图8)。
  • 在前沿模型上的可扩展性:Aurora成功扩展到最新的大型MoE模型。在MiniMax M2.1 (229B)上,相比无推测基线,在批次大小为4时实现了1.45倍的吞吐量提升(图9)。在Qwen3-Coder-Next (80B)上,在批次大小为8时实现了1.23倍的吞吐量提升(图10)。根据模型和批次大小的不同,加速比在1.25倍至1.57倍之间(表1)。

A5 结论

本文提出了Aurora,一个统一的训练-服务系统,它将推测式解码重新表述为一个联合的学习与服务问题。通过将基于SGLang的推理服务器与一个异步训练服务器通过GPU感知的RPC连接起来,Aurora实现了在实时流量下对草稿模型的持续在线策略自适应,从而解决了限制传统两阶段流程的训练-服务不匹配问题。我们的实验表明,简单的在线微调能够捕获大部分可获得的增益,惰性同步能够最好地平衡自适应速度与服务稳定性,并且从零开始的“第0天”部署是可行的:一个未经训练的推测器在数千个请求内就能达到有竞争力的接受率,从而消除了为新模型上线而进行离线预训练的瓶颈。

A6 附录

A Aurora系统技术细节

我们工作的一个关键技术贡献是设计了一个异步、解耦的训练架构,它能够在推理时进行训练,而不会中断生产推理管道。我们提出了一个基于GPU感知的远程过程调用(RPC)的系统,它连接了推理服务器和训练服务器,在保持低推理延迟的同时,实现了在节点间后台高效的GPU到GPU数据传输。

我们的架构由两个独立组件组成:(1)一个生产推理服务器,用推测式解码为真实用户请求提供服务;(2)一个被动训练服务器,根据传入的推理数据持续更新草稿模型。与以往需要离线训练或重复加载模型的方法不同,我们的系统通过精心设计的RPC通信渠道,以最小的内存开销实现了在线训练。

A.1 零拷贝目标模型设计

零拷贝目标模型设计。我们方法的一个关键洞见是,在推理时训练期间,我们可以消除在训练服务器上加载重复目标模型的需求。传统的推测式解码训练需要计算目标模型的隐藏状态和logits,这需要加载完整的目标模型(例如,8B参数在FP16下约占16GB)。我们的基于RPC的设计直接从推理服务器传输预先计算好的目标信息:

$$\mathcal{D}_{\mathrm{RPC}}=\{(h_t^{(l)}, \ell_t, x_{\mathrm{in}}, y_{\mathrm{out}}, \mathcal{R})\}$$

RPC传输内容。其中,$h_t^{(l)} \in R^{T \times 3d}$ 代表从三个策略性选择的目标模型层(早期、中期、晚期)拼接的隐藏状态,$\ell_t \in R^{T \times V}$ 表示输出logits,$x_{in}$和$y_{out}$是输入和输出的token序列,R编码了拒绝轨迹信息。

设计优势。这种设计实现了:
- 训练内存效率:训练节点只加载草稿模型(约1B),而不是目标模型(8B到70B)。
- 热插拔,最小化推理干扰:训练在独立的硬件上运行,导致服务开销极小。
- 水平扩展:多个推理服务器可以为一个训练服务器提供数据。

A.2 推理侧内存开销

推理侧内存开销。虽然我们的设计从训练服务器上移除了目标模型,但它在推理服务器上引入了适度的内存开销,用于在RPC传输前缓冲辅助数据。对于每个序列长度为$T = T_{prompt} + T_{output}$的请求,服务器必须临时存储:(1)来自目标层的辅助隐藏状态,在BF16模型中消耗$T \times 3d \times 2$字节;(2)输出token的目标logits,消耗$T_{output} \times V \times 2$字节。对于一批B个请求,总的辅助内存为:

$$M_{\text{aux}} = B \times (T \times 3d \times 2 + T_{\text{output}} \times V \times 2)$$

缓解瓶颈的策略。为了缓解logits瓶颈,我们利用了大多数草稿模型的词汇表比其目标模型小得多的事实(例如,Llama3模型为32K对128K,Qwen3模型为32K对152K),因此我们只需要传输草稿词汇表上的logits——这减少了4到5倍。我们可以进一步应用top-K logits过滤(例如,K=1024),只保留信息最丰富的目标logits,将每个token的logits存储从256 KB减少到2 KB(减少了128倍),使得隐藏状态成为主要的成本。

A.3 多服务器聚合

多服务器聚合。多个推理服务器可以同时向一个训练服务器发送数据。每个服务器被分配一个唯一的RPC rank,训练服务器通过线程安全的LRU缓存来聚合数据:

$$ \text{RPC World} = \{ \underbrace{S_1, S_2, \dots, S_N,}_{\text{Inference Servers}} \underbrace{T}_{\text{Training Server}} \} $$

异构集群配置。对于异构集群配置,显式的GPU设备映射使得部署更加灵活,其中推理服务器可以使用GPU 0-3进行张量并行,而训练服务器则使用完全独立的节点上的GPU。

A.4 与解耦式服务框架的兼容性

与解耦式服务框架的兼容性。我们的解耦架构自然地扩展了解耦式服务框架,如Mooncake【索引23,Mooncake: A kvcache-centric disaggregated architecture for llm serving,2024,ACM Transactions on Storage】和DistServe【索引38,{DistServe}: Disaggregating prefill and decoding for goodput-optimized large language model serving,2024,18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24)】,这些框架已经将prefill和decoding分离到不同的节点池,并具备跨节点GPU传输基础设施。我们的训练服务器作为第三个解耦角色,通过相同的通信结构(例如RDMA)接收隐藏状态和logits,而无需单独的数据路径。这使得我们的推理时训练管道可以部署在任何服务后端之上——无论是单体式还是解耦式——而无需更改训练逻辑。

A.5 Aurora系统的通用性

Aurora系统的通用性。我们注意到,Aurora不仅可以用于在线服务,也可以用于传统的推测器训练,只需使用预先存在的语料库模拟在线处理管道即可。

离线训练场景。对于推测器训练,用户通常会遇到两种情况:1)当目标模型生成的语料库不可用时,我们可以通过使用预先收集的提示来收集在线数据来训练推测器。2)当目标模型生成的语料库可用时,我们可以简单地以仅预填充(prefill-only)模式运行推理组件,以高效地生成推测器训练所需的激活值。

作为传统训练框架的优势。即使被视为传统的推测器训练框架,Aurora仍然具有优势:训练和推理可以并行运行,使其比以前的系统更高效,并且用户可以立即测量加速效果,避免了从接受率估计吞吐量改进或因训练-服务不匹配而错误估计吞吐量的风险。因此,我们相信Aurora可以高效地适应不同的用户场景。

A.6 数据集描述

表2:用于推理时训练实验的多领域数据集构成。

数据构成。我们策划了一个混合领域的提示流,涵盖数学推理、结构化查询生成、以代码为中心的任务、金融领域指令和通用对话提示。具体来说,数学推理采样自GSM8K【索引9,Training verifiers to solve math word problems,2021,arXiv preprint arXiv:2110.14168】,其中包含多步小学数学应用题。文本到SQL来自Spider【索引35,Spider: A large-scale human-labeled dataset for complex and cross-domain semantic parsing and text-to-sql task,2018,Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing】 ,其中模型根据自然语言问题和数据库模式生成可执行的SQL查询。代码数据采样自CodeSearchNet【索引13,Codesearchnet challenge: Evaluating the state of semantic code search,2019,arXiv preprint arXiv:1909.09436】,涵盖多种编程语言的代码及相关的自然语言描述。金融数据采样自Finance-Alpaca【索引2,finance-alpaca dataset,https://huggingface.co/datasets/gbharti/finance-alpaca】,包含与金融相关的指令跟随提示,具有领域特定的术语和推理。对话指令来自chatbot_instruction_prompts【索引1,chatbot_instruction_prompts,https://huggingface.co/datasets/alespalla/chatbot_instruction_prompts】,其中包含用于开放式对话的通用助手式提示 。

数据集代表性。我们的数据集涵盖了多样化的token级分布,这些分布直接决定了草稿-目标匹配模式,从而决定了推测式解码在分布变化下的有效性。具体而言,GSM8K的特点是数值符号和具有相对一致答案结构的多步推理轨迹;Spider需要严格结构化的SQL生成,具有严格的语法约束和频繁的关键字/标点符号;CodeSearchNet包含以代码为中心的输出,具有高频的编程token(如缩进和运算符)、长尾标识符和强烈的局部重复性;Finance-Alpaca引入了领域特定的术语(如股票代码和宏观经济概念),并伴有解释性的响应;而chatbot_instruction_prompts代表了开放式的助手对话,语言更加多样且约束更少。这种多样性使我们能够可靠地引发导致接受率明显下降的分布变化,并评估在线推测器训练是否能随时间恢复性能。

输出模式。此外,该数据集捕捉了现实世界助手流量中两种常见的输出模式:(i)长篇自然语言响应(金融和对话指令),以及(ii)格式受限的结构化生成(文本到SQL和代码)。与仅在标准问答提示上进行评估相比,这些模式更好地反映了推测式解码从可预测的结构片段(例如,SELECT ... FROM ... WHERE ...)和在结构化输出中频繁出现的重复模板中受益的场景。总的来说,我们的领域构成为一个直观且现实的转变设置提供了基础,例如,流量从以推理为主的查询过渡到代码/SQL/金融/对话工作负载。

B 训练细节

B.1 基础设施和效率

GPU分配。为了在不降低推理吞吐量的情况下实现在线训练,我们采用了一种分离式GPU架构,并基于RPC进行数据传输:我们将目标模型放置在GPU 0-3上进行推理,同时在另一块独立的GPU上训练草稿模型,以避免服务和学习之间的资源竞争。所有Qwen-8B和Llama-8B实验均在H100 GPU上进行,而MiniMax M2.1和Qwen3-Coder-Next实验则在H200 GPU上进行。

数据传输。隐藏状态和验证结果通过RPC从推理服务器传输到训练进程,存储缓冲区与GPU 1共置,以避免冗余的内存拷贝。

异步更新。梯度更新相对于推理是异步进行的,每100个推理请求触发一次更新。这个更新间隔允许训练缓冲区在每次优化步骤前累积多样化的样本,平衡了训练的响应性和计算效率,同时确保服务延迟不受训练开销的影响。

吞吐量指标。在本文中,我们报告的每个请求的吞吐量以每秒token数(tokens/s)为单位,定义为处理单个请求的总token数除以端到端延迟:

$$ \text{Throughput} = \frac{T_{\text{input}} + T_{\text{output}}}{t_{\text{request}}}, $$

其中,$T_{input}$是输入(提示)token的数量,$T_{output}$是生成的输出token的数量,而$t_{request}$是从请求提交到响应完成的实际时间。该指标捕获了单个用户体验到的有效token处理速率,并直接反映了推测式解码带来的延迟改进。与衡量并发请求总token数的聚合系统吞吐量不同,每个请求的吞吐量隔离了每个查询的加速效益,使其成为评估不同推测式解码配置下用户感知性能增益的合适指标。

B.2 优化和超参数

通用优化设置。所有在线训练配置共享一个通用的优化设置,该设置旨在在推理工作负载约束下实现稳定的流式更新:
- 优化器:AdamW,权重衰减β = 0.0,梯度裁剪最大范数为0.5。
- 学习率:我们对微调使用恒定的学习率η = $1 \times 10^{-5}$,对从头开始训练使用η = $1 \times 10^{-4}$,在经过400步的短暂线性预热后(占总训练的0.05%)。这种恒定学习率对于测试时训练场景至关重要,在这种场景中,模型必须保持可塑性以持续适应而不会遗忘【索引28,Testtime training with self-supervision for generalization under distribution shifts,2020,International Conference on Machine Learning (ICML)】。
- 批次配置:草稿模型训练的全局批次大小为8,草稿模型前向传播的微批次大小为1,目标模型验证的微批次大小为4。这些不对称的批次大小平衡了训练稳定性和推理吞吐量要求。
- 上下文窗口:最大序列长度为2,048个token,测试时训练(TTT)长度为5个token,定义了每个梯度更新的局部上下文窗口。
- 混合精度:BF16计算,使用FP32主权重,并在优化前将梯度转换为FP32,遵循了训练稳定性的最佳实践【索引20,Mixed precision training,2017,arXiv preprint arXiv:1710.03740】。
- 损失函数:反向KL散度$D_{KL}(p_{target} \| p_{draft})$鼓励草稿模型覆盖目标模型的分布,从而最小化推测验证期间的错误拒绝。

B.3 附加结果

表3:所有配置共享的超参数。恒定学习率保持了在线学习的可塑性。

表4:MimiMax M2.1 (BF16):在不同批处理大小和前瞻步数下的端到端吞吐量。我们报告每秒token数(TPS,定义见附录B.1)和相对于无推测基线的加速比。

表5:Qwen-Coder-Next:在不同批处理大小和前瞻步数下的端到端吞吐量。我们报告每秒token数(TPS)统计数据和相对于无推测基线的加速比。