菜单
主题
vLLM Logo

EAGLE 3.1：通过 EAGLE 团队、vLLM 和 TorchSpec 的协作推进投机采样

4 分钟阅读
EAGLE 团队、vLLM 团队和 TorchSpec 团队
#投机采样#性能

EAGLE 系列（包括 EAGLE 1、EAGLE 2 和 EAGLE 3）已成为研究和生产系统中应用最广泛、最实用的投机采样算法家族之一。

今天，EAGLE 团队、vLLM 团队和 TorchSpec 团队很高兴共同推出 EAGLE 3.1 —— 这是在投机采样的稳健性、效率和可部署性方面迈出的重要一步。

EAGLE 3.1 的创新

虽然投机采样在受控环境中表现良好，但在不同的聊天模板、长上下文输入或分布外系统提示词下，其性能往往会下降。

EAGLE 团队将这种脆弱性追溯到一种我们称为注意力漂移（attention drift）的现象——随着采样深度的增加，草稿模型逐渐将注意力从汇聚标记（sink tokens）转移到其自身生成的标记上。

我们确定了两个根本原因。首先，随着高层隐藏状态主导了草稿模型的输入，融合后的输入表示变得愈发不平衡。其次，由于未归一化的残差路径，隐藏状态的幅度在采样步骤中不断增长。综合这些因素，使得草稿模型在更深的采样深度下变得越来越不稳定。

Figure 1: EAGLE 3 vs. EAGLE 3.1 architecture comparison. EAGLE 3.1 adds FC normalization after each target hidden state and feeds post-norm hidden states into the next decoding step.
图 1：EAGLE 3 与 EAGLE 3.1 架构对比。EAGLE 3.1 在每个目标隐藏状态后增加了全连接层（FC）归一化，并将归一化后的隐藏状态输入到下一个解码步骤。

为了解决这个问题，EAGLE 3.1 引入了两个关键的架构改进：

  • 在每个目标隐藏状态之后、全连接层（FC）之前进行归一化
  • 将归一化后的隐藏状态输入到下一个解码步骤

直观地说，这种后归一化（post-norm）设计使该方法表现得更像是跨解码步骤递归地调用草稿模型，而不是仅仅向目标模型添加额外的层。

这些变化显著提高了在各种部署场景下的稳健性。与 EAGLE 3 相比，EAGLE 3.1 展示了：

  • 更好的训练到推理的外推性
  • 更强的长上下文稳健性
  • 对聊天模板和系统提示词变化的更高弹性
  • 在不同服务环境中更稳定的接受长度

在长上下文负载中，与 EAGLE 3 相比，EAGLE 3.1 的接受长度最长可提升 2 倍。

使用 TorchSpec 进行 EAGLE 3.1 训练

TorchSpec 现在为 EAGLE 3.1 以及未来的投机采样算法提供高效的训练支持。

通过降低训练开销并简化实验工作流程，TorchSpec 有助于加速下一代投机采样研究和部署的迭代与探索。

基于 TorchSpec 和 vLLM，我们还训练并开源了一个适用于 Kimi K2.6 的 EAGLE 3.1 草稿模型：

https://hugging-face.cn/lightseekorg/kimi-k2.6-eagle3.1-mla

该模型展示了如何在实际服务模型上，通过 TorchSpec 训练和 vLLM 服务支持来部署 EAGLE 3.1。

EAGLE 3.1 与 vLLM 的集成

EAGLE 3.1 以配置驱动的扩展形式落地 vLLM，作为对现有 EAGLE 3 实现的补充。

此次集成包括：

  • FC 归一化支持
  • 后归一化隐藏状态反馈
  • 移除了关于目标隐藏状态的硬编码假设

同时，与现有 EAGLE 3 检查点的向后兼容性得到完全保留。因此，EAGLE 3.1 草稿模型可以通过相同的投机采样代码路径直接插入，例如：

vllm serve nvidia/Kimi-K2.6-NVFP4 \
  --trust-remote-code \
  --tensor-parallel-size 4 \
  --tool-call-parser kimi_k2 \
  --enable-auto-tool-choice \
  --reasoning-parser kimi_k2 \
  --attention-backend tokenspeed_mla \
  --speculative-config '{"model":"lightseekorg/kimi-k2.6-eagle3.1-mla","method":"eagle3","num_speculative_tokens":3}' \
  --language-model-only

这使得在生产环境 vLLM 服务中升级草稿模型变得顺畅且容易。

该支持已合并到 vLLM 的当前主分支，并将通过 vLLM 的每日构建版（nightly release）以及即将发布的 v0.22.0 版本提供。

作为早期的基准数据，我们在 SPEED-Bench 编码数据集上，使用 vLLM（TP=4, GB200, 非解耦）在 Kimi-K2.6-NVFP4 上对 Kimi K2.6 EAGLE 3.1 草稿模型进行了基准测试。EAGLE 3.1 在并发度为 1 时，实现了 2.03 倍的单用户输出吞吐量提升，并且随着并发度的增加，加速效果依然显著（并发度为 4 时为 1.71 倍，并发度为 16 时为 1.66 倍）。

Figure 2: Per-user output throughput (TPS) on Kimi-K2.6-NVFP4 with vLLM, TP=4, GB200 on SPEED-Bench coding. EAGLE 3.1-MLA vs. no-spec baseline.
图 2：在 SPEED-Bench 编码测试集上，使用 vLLM、TP=4、GB200 在 Kimi-K2.6-NVFP4 上的单用户输出吞吐量 (TPS)。EAGLE 3.1-MLA 与无投机基准对比。

跨生态系统的开源协作

EAGLE 团队、vLLM 团队和 TorchSpec 团队之间的这种合作为算法研究、系统优化和训练基础设施之间的开源协作树立了强有力的典范。

EAGLE 团队将继续推动投机采样算法的发展，vLLM 帮助将这些创新大规模带入生产推理系统，而 TorchSpec 为未来的投机采样算法提供了高效的训练和快速实验能力。

我们也非常感谢 NVIDIA 的 GPU 支持和持续的合作伙伴关系。这种支持对于实现将 EAGLE 3.1 从算法创新转化为实际部署所需的开发、验证和基准测试工作起到了重要作用。

我们希望通过共同努力，继续提高投机采样的整体基准，并推动整个 LLM 生态系统在 Token 生成效率方面的进一步提升。