EAGLE-3 在 AMD Instinct GPU 上的推测解码

EAGLE-3 在 AMD Instinct GPU 上的推测解码

vLLM 和 AMD Quark 团队在 AMD Instinct GPU 上推出了 EAGLE-3 推测解码的端到端管道，使大型混合专家（MoE）和注意力密集型模型能够实现无损推理加速。此集成在 AMD Instinct MI355X GPU 上实现了 Kimi-K2.5 的吞吐量提升范围为 1.69x 到 2.00x，MiniMax-M2.5 的吞吐量提升范围为 1.38x 到 1.79x。

推测解码及 EAGLE-3 方法

推测解码是一种无损加速技术，可减少昂贵目标模型解码迭代的次数。与逐个生成 token 不同，轻量级草稿模型会提出多个未来 token，目标模型随后在一次前向传递中验证这些 token。与目标模型分布匹配的 token 会被接受，生成将从第一个被拒绝的 token 继续。

EAGLE-3 通过利用目标模型的多层特征，改进了之前的版本。EAGLE-3 不使用无关的小型语言模型，而是训练一个与目标模型紧密对齐的草稿模块，利用低级、中级和高级语义特征来提高草稿的接受率和整体推理速度。

AMD Quark MXFP4 量化

为了解决大型 MoE 模型在内存带宽和容量方面的限制，AMD Quark 团队使用 MXFP4（OCP Microscaling 4-bit 浮点）格式为主流 LLM 提供“day-0”量化。

MXFP4 实现的关键技术细节包括：

  • 硬件加速：AMD Instinct MI350 系列 GPU（MI350X/MI355X）提供原生 FP4 矩阵加速。
  • 内存效率：MXFP4 将 4-bit 元素分组到具有共享缩放因子的小块中，在减少内存占用的同时保持优于 INT4 的数值行为。
  • vLLM 集成：这些检查点通过 FP4 ASM GEMM 路径（VLLM_ROCM_USE_AITER_FP4_ASM_GEMM=1）和 AITER MoE 内核在 ROCm 上被 vLLM 消费。

以 vLLM 为中心的 EAGLE-3 训练管道

训练高接受率的草稿模型被视为一个系统问题，vLLM 作为训练循环的核心，贯穿五个阶段：

  1. 策略内数据合成：vLLM 提供 AMD Quark MXFP4 目标模型，使用精确的服务聊天模板生成响应，以确保训练与部署的一致性。
  2. 隐藏状态提取：vLLM 的隐藏状态提取钩子直接从运行中的目标引擎暴露辅助层（低级、中级和高级隐藏状态以及 fc_norm）。这支持在线、离线和流式模式，后者使得在单个节点上可以训练 420B MXFP4 MoE 目标模型。
  3. 冷启动 FSDP2 训练：单层 EAGLE-3 草稿头在 FSDP2 下从零开始训练，使用训练时测试（TTT）损失和位置衰减权重。
  4. 循环内服务-评估：当前检查点会定期导出并在 vLLM 推测解码下服务，以测量真实的接受长度，确保所选检查点针对生产引擎进行了优化。
  5. 导出与部署：最终的草稿导出为 Hugging Face 格式，并通过 vLLM-ROCm 进行部署。

草稿质量与上下文缩放

在 SPEED-Bench 上的评估表明，MiniMax-M3 EAGLE-3 草稿在 11 个领域中的平均接受长度（AL）为 2.77。最高的接受率出现在结构化技术内容中：

  • 编码：3.16 AL
  • 数学：3.12 AL
  • RAG：3.11 AL
  • 多语言：3.07 AL

值得注意的是，随着提示长度从 1K 增加到 32K 个 token（2.64 到 2.63），接受长度保持稳定，表明在长上下文中加速不会衰减。

AMD Instinct MI355X 上的性能基准

在 AMD Instinct MI355X（TP=4）上使用 1K 输入和 1K 输出 token（1K/1K）进行的基准测试表明，相较于非推测基线，吞吐量显著提升：

Kimi-K2.5 结果

  • BF16 草稿路径：吞吐量提升 1.69x 至 1.90x。
  • AMD Quark FP8 草稿路径：吞吐量提升 1.76x 至 2.00x。

MiniMax-M2.5 结果

  • BF16 草稿路径：吞吐量提升 1.38x 至 1.79x。

在两种模型中，性能提升在较低的并发级别时最为显著。

Sources