运行时 Draft Model 训练：将 Speculative Decoding 适配到实际工作负载

概述

随着大语言模型（LLM）的规模和复杂性持续增长，降低推理成本已成为 AI 数据中心和服务提供商的关键优先事项。大规模部署这些模型所带来的计算开销，推动了对高效优化技术的需求，这些技术需要在不影响模型质量的前提下实现显著的成本节约。

为解决推理优化问题，多种方法应运而生，包括 disaggregation、KV cache 感知路由、test-time 路由、quantization、speculative decoding 等。在这些技术中，speculative decoding 因其独特优势而受到主要云服务提供商（CSP）的广泛关注：它在保证模型质量的同时可与其他优化方法组合使用，并在大多数实际场景中带来性能提升。

Speculative decoding 通过使用小型快速模型（称为 draft model）生成 draft token，然后由原始模型（称为 target model）并行验证来加速推理。这种方法从根本上比仅使用 target model 进行顺序 token 生成更加高效。

上图展示了使用 4 个 draft token 的 speculative decoding 与标准 decoding 生成 10 个 token 的时间线。即使平均只有 2 个 draft token 被接受，speculative decoding 也能实现约 2 倍的加速。

这种效率提升之所以可能，是因为 target model 并行验证多个 token 所需的时间与生成单个 token 所需的时间几乎相同。这是因为 LLM 推理主要是内存受限而非计算受限——瓶颈在于从内存加载模型权重，而非实际计算。无论是验证一个 token 还是同时验证多个 token，target model 都需要加载相同的权重，因此延迟相近。通过将这一内存访问成本分摊到多个 token 验证中，speculative decoding 显著减少了生成序列所需的昂贵 target model forward pass 总数。

然而，传统的 speculative decoding 方法面临一个关键限制：draft model 通常在可能与实际生产流量分布不一致的通用工作负载上进行预训练。由于工作负载分布在不同服务之间差异显著，且即使在同一服务内也会随时间演变，因此在通用数据上训练的 draft model 往往在专业化或不断变化的使用场景中表现欠佳。

用于自改进 LLM 推理的 Temporal Incremental Draft Engine (TIDE)

TIDE 通过运行时 draft model 训练来解决这一挑战。TIDE 以 SGLang 作为推理引擎，利用 SpecForge 进行 draft model 训练，并基于 EAGLE3 speculative decoding 技术。通过根据实时服务工作负载持续调整 draft model，TIDE 无需人工干预即可自动提升推理性能。

EAGLE3：从 Hidden State 训练 Draft Model

EAGLE3 采用了一种独特的 draft model 训练方法。与作为独立语言模型运行的传统 draft model 不同，EAGLE3 的 draft model 将 target model 多个中间层的 hidden state 作为输入，学习预测 target model 的输出分布。

TIDE 系统设计

TIDE 的架构旨在将运行时训练无缝集成到生产推理系统中，同时最大限度地减少开销和复杂性。

1. 带有 Hidden State 日志记录的推理引擎

推理引擎（SGLang）执行标准的 prefill 和 decoding 计算来处理用户请求。关键的是，在这些操作过程中，它会捕获并将 target model 中间层的 hidden state 转储到存储中。这些 hidden state 代表了实际生产推理过程中计算的内部表示，并作为 draft model 的训练数据。

2. 训练与模型更新

随着推理引擎持续处理请求，hidden state 在存储中不断积累。当收集到足够的数据——代表当前工作负载分布的有意义样本后——训练过程会自动触发。训练引擎随后加载累积的 hidden state，训练 EAGLE3 draft model 以根据近期工作负载更好地预测 target model 的输出分布。训练收敛后，更新的 draft model 被部署回推理引擎，替换之前的版本。这完成了一个适应周期，随着新的 hidden state 不断积累，该过程持续进行，确保对不断变化的工作负载模式进行持续适应。

该架构优雅地解决了运行时训练的关键挑战：

• 近乎零的推理开销： Hidden state 日志记录与推理计算异步运行，使得开销几乎完全被隐藏，对请求服务几乎不增加延迟。
• 异步训练： 训练独立于推理进行，不会阻塞或减慢请求处理。
• 自动适应： 整个周期自主运行，无需人工干预，持续适应工作负载变化。
• 资源效率： 在异构 GPU 系统中，训练可以卸载到不同的硬件上，而高性能 GPU 专注于推理。

该架构的简洁性和自动化使其适用于生产部署，因为在生产环境中手动调优和干预成本高昂且不切实际。

评估

为评估运行时 draft model 训练的有效性，我们将 TIDE 的性能与在通用数据上预训练的静态 draft model 进行了对比。我们使用 dbdu/ShareGPT-74k-ko 数据集进行评估，该数据集代表韩语对话查询，为工作负载特定适应提供了现实的测试平台。我们的实验使用 lmsys/gpt-oss-120b-bf16 作为 target model，lmsys/EAGLE3-gpt-oss-120b-bf16 作为基线预训练 draft model，推理在 AMD MI300X 或 NVIDIA H100 GPU 上运行，draft model 训练在 AMD MI250 GPU 上进行，展示了 TIDE 有效利用异构 GPU 资源的能力。

如 Figure 3 所示，TIDE 根据并发级别的不同，相对于预训练 draft model 实现了 1.14× 至 1.35× 的输出 token 吞吐量加速提升。

Figure 4 展示了运行时训练的一个关键优势：随着 TIDE 持续从工作负载中学习，吞吐量不断提升。预训练 draft model 由于其性能由预训练固定，在整个实验过程中保持相对恒定的吞吐量。相比之下，TIDE 在适应韩语对话数据集中的特定模式时，展现了吞吐量的渐进提升。

为了更好地理解 TIDE 性能提升背后的机制，我们分析了 acceptance length——每个验证步骤中 target model 接受的 draft token 数量——在不同并发级别下随时间的变化情况。

Figure 5 揭示了运行时训练的一个重要特征：更高的并发级别会导致 acceptance length 更快地提升。原因如下：

• 更快的数据积累： 在更高并发下，更多请求被同时处理，使 TIDE 能够以更快的速度收集训练数据。这加速了学习过程，实现了对工作负载分布的更快适应。
• 更频繁的模型更新： 单位时间内有更多的训练样本可用，draft model 可以使用统计上显著的批次进行更频繁的更新，从而实现更快的收敛和更广泛的模式覆盖。

结论

TIDE 证明了运行时 draft model 训练可以显著提升生产环境中的 speculative decoding 性能。我们在韩语对话数据集上的实验表明，相对于静态预训练 draft model，TIDE 实现了 1.14× 至 1.35× 的输出 token 吞吐量加速，且随着系统适应实时工作负载，性能持续提升。

除了性能提升，TIDE 在异构 GPU 系统中还提供了引人注目的优势。通过利用空闲或训练优化的资源——如旧一代 GPU 或利用率不足的硬件——进行 draft model 训练，同时将高性能 GPU 保留用于推理，TIDE 同时提升了资源利用率和推理效率。这在系统层面带来了更好的整体性价比。

为将运行时训练无缝集成到生产推理系统中，我们向开源项目做出了贡献，使 TIDE 的能力对更广泛的社区开放：

• SpecForge: PR#1, PR#2
• SGLang: PR#3

随着 AI 工作负载持续多样化和演变，像 TIDE 这样能够自动适应变化模式的系统，将在大规模保持高效且经济的推理方面变得越来越不可或缺。