
摘要
推测解码在生产环境中会过时——草稿模型可能漂移,而离线重训练无法总是跟上节奏。
Aurora 解决了这个问题。它是一个基于强化学习的开源框架,直接从在线推理轨迹中学习,并在不中断服务的情况下持续更新推测器。
关键结果:
→ 跨变化的流量域实时自适应
→ 相比训练良好的静态推测器,额外加速 1.25 倍
核心发现:从头开始的在线训练可以超越精心预训练的静态基线。
在生产环境中运行大型语言模型始终是性能与成本之间的权衡。推测解码是标准手段:原则上,它能加速推理。实际上,它常常效果不佳——草稿模型过时,接受率漂移,而离线重训练太慢、太贵,无法跟上实时流量。如果你的系统能够从正在服务的请求中持续、即时地学习呢?
去年,我们推出了 ATLAS——迈向自适应推测器的第一步。那项工作奠定了基础,但目标始终是一个完全自主的系统,能够闭环服务和训练。
今天,我们发布 Aurora,一个基于强化学习的开源框架,从在线推理轨迹中学习并异步更新推测器——将推测解码从静态的一次性设置转变为动态的自我改进飞轮。这种统一设计解锁了标准流水线难以实现的能力,包括:(1) 直接缓解分布不匹配,相比强离线基线提升 1.25 倍;(2) 通过消除大规模激活收集流水线降低基础设施成本;(3) 与算法无关的框架,兼容未来的推测器设计;(4) 支持多样化的异构用户需求。
在实验中,Aurora 在广泛使用的模型(如 Qwen3 和 Llama3)上,相比训练良好但静态的推测器实现了额外 1.25 倍的加速。
复现论文结果的代码已开源,我们欢迎社区的贡献。

不同批量大小下的端到端吞吐量
MiniMax M2.5 (FP8, lookahead 5):
Qwen3-Coder-Next-FP8 (lookahead 5):
1. 为什么标准的先训练后服务流水线会失效
离线推测训练在组织上很方便,但在生产环境中引入了几个实际问题,限制了其有效性。传统流水线是单向的——导致模型过时,并与实际性能脱节。

传统的推测解码遵循线性、静态的流程,随时间推移性能下降。Aurora 引入了循环、持续自适应的方法。
验证器在移动,但起草器滞后。 生产目标模型会因质量、安全、成本或硬件迁移而发生变化。推测器通常更新得慢得多,因此它会变得陈旧,推测性能随时间下降。
离线蒸馏流程成本高昂。 用于起草器训练的激活收集和重放流程在大规模存储和运行时可能极其昂贵。在生产规模下,存储占用可能达到PB级别,内存、带宽和操作复杂性成本高昂。Aurora通过直接从在线服务轨迹中学习来减轻这一负担。
接受率不等于实际加速。 离线训练可以在实验室环境中优化接受率,但生产加速取决于实际服务栈:内核、数值精度(FP8/FP4)、批处理、调度和硬件行为。最佳的离线草稿模型可能不是最佳的在线模型。在实践中,大多数团队训练多个起草器,但最终只选择一个——Aurora能够直接比较加速效果,因为它在线运行。
这些差距表明,推测解码不应仅仅被视为一个建模问题(“训练更好的起草器”),而应被视为一个联合学习与服务的问题。
2. 核心思想:由强化学习驱动的服务到训练飞轮
Aurora将推测解码转变为服务到训练的飞轮。它不是将推测器视为静态工件,而是从它服务的每个请求中持续学习。

该系统围绕两个解耦的组件构建。推理服务器运行一个推测解码引擎(基于SGLang或vLLM),包含一个目标模型和一个草稿模型。对于每个请求,草稿模型提出一个token序列,然后由目标模型并行验证。接受和拒绝的token的结果——以及用于EAGLE风格训练的隐藏状态——被流式传输到分布式数据缓冲区。训练服务器异步运行:它从缓冲区获取训练数据批次,对草稿模型的副本执行梯度更新,并定期将改进的权重热替换回推理服务器,而不会中断服务。
这种设计基于两个生产现实。首先,服务效率是真正的目标——延迟、吞吐量和SLO下的每token成本。其次,同步必须是惰性且无中断的——频繁的权重推送可能导致缓存失效和延迟抖动。为了使这种设计可靠地工作,我们将在线推测训练重新表述为异步强化学习问题。
这不仅仅是理论上的便利——它直接将训练信号与真实部署效用对齐,而不仅仅是离线模仿质量。推测解码自然地映射到强化学习:
在这种框架下,最大化回报直接映射到最大化接受长度——这直接映射到解码加速。Aurora的一个微妙但强大的部分是,它不仅从接受的token中学习。接受损失(模仿)使用接受token上的交叉熵,鼓励草稿重现验证器批准的延续。拒绝损失(丢弃采样)教导草稿不要提出什么,使用拒绝分支作为反事实监督。
为了高效处理推测解码结果的复杂分支结构,我们采用了专门的树注意力机制。通过构建一个尊重推测树因果结构的自定义注意力掩码,我们可以在单次批处理的前向和后向传播中处理所有被接受和被拒绝的分支。
3. 适应分布偏移
为了测试Aurora的鲁棒性,我们使用包含40,000个提示的流模拟了在线服务流量,这些提示涵盖五个领域:数学推理、文本到SQL、代码生成、金融和通用对话。这种组合反映了现实部署场景,其中服务流量表现出异构和变化的任务分布。我们评估了两种流量模式:(i) 有序流,其中请求按领域分组以引发突然的分布偏移,以及(ii) 混合流,其中提示被随机打乱以近似平稳流量。
当请求按领域分组以引发突然的分布变化时,Aurora持续适应。系统在每次偏移后大约10,000个请求内恢复接受长度,展示了鲁棒的在线适应能力。

从一个训练良好的推测器开始,Aurora通过持续适应在静态基线之上实现了额外的1.25倍加速。这表明Aurora的收益在现有离线训练投资的基础上叠加。
混合流量的结果尤其引人注目:从头开始的在线训练可以超越精心预训练的推测器的性能。接受长度达到3.08(超过了静态基线的2.63和预训练后微调基线的2.99),吞吐量稳定在302.3 tokens/s。这从根本上挑战了推测解码需要大量离线预训练的传统观点。
4. 结论
Aurora不仅仅是另一种推测解码算法。它是一个系统层面的转变。它将推测解码从静态的离线任务转变为动态的在线学习过程。
这种转变解锁了实时效用反馈、领域漂移下的适应、与大型离线蒸馏流水线相比更低的基础设施成本,以及与未来推测器算法兼容的系统层。这就是为什么推测解码的正确抽象不再仅仅是孤立的更好草稿训练——而是一个统一的训练-服务循环。
