vLLM 在 AMD GPU 上的推测解码：性能与方法

摘要

推测解码在 vLLM 中通过一个快速的草稿组件在目标模型的一次前向传播中验证多个未来标记，从而在 AMD Instinct MI300X/MI355X GPU 上实现某些模型-方法组合高达约 2.9 倍的吞吐量提升，同时保持原始模型的输出行为不变。

vLLM 中推测解码的工作原理

推测解码在原始（目标）大语言模型前增加了一个 草稿-验证 阶段。草稿组件生成一系列候选标记；目标模型随后在一次前向传播中评估整个候选块。被目标模型接受的标记被提交，第一个被拒绝的标记则由目标模型自身输出替代。该过程重复进行，直到生成完成。

关键特性

目标模型仍是事实真相来源；在验证前不会输出任何标记。
一次目标模型前向传播可提交多个草稿标记，减少昂贵的目标模型前向传播次数。
接受性从左到右评估；一旦拒绝，候选块的其余部分将被中止。

评估的草稿方法

vLLM 支持五种具体的草稿方法，它们在草稿网络如何获取目标模型信息以及是否以串行或并行方式生成标记方面有所不同。

方法	草稿组件	使用的目标模型信息	标记生成方式
原生 MTP	模型原生的辅助预测路径	来自目标模型（或前一个 MTP 步骤）的隐藏表示	串行（重复调用 MTP）
Gemma 4 MTP	与目标模型配对的独立 MTP 检查点	目标激活 + 共享 KV 缓存	串行
EAGLE‑3	专用自回归推测器	来自目标模型早期、中期和晚期层的融合隐藏状态	串行，每个草稿标记依赖于下一个
DFlash	专用并行推测器	目标隐藏状态投影为每个草稿层使用的 KV 对	所有位置在一次前向传播中同时预测
DSpark	DFlash 主干 + 轻量级马尔可夫头	与 DFlash 相同的目标上下文	并行主干 + 串行修正标记选择

原生 MTP

在目标模型内部实现，MTP 使用辅助头预测固定数量的未来标记。草稿标记一个接一个生成，每一步都消耗前一个 MTP 输出的隐藏状态。当 num_speculative_tokens 超过原生深度时，vLLM 通过额外的前向传播重用 MTP 路径。

Gemma 4 MTP

一个独立的检查点（即助手模型）与目标模型共享 KV 缓存和激活，允许其重用已计算的上下文。草稿标记仍为串行生成，但由于草稿组件重用目标张量，内存开销较小。

EAGLE‑3

EAGLE‑3 训练一个推测器，接收三个目标层隐藏状态的融合表示。推测器以自回归方式运行：第一个草稿标记结合融合的目标特征与采样标记嵌入；后续标记依赖于前一个草稿输出。这导致与提议长度成比例的串行草稿工作量。

DFlash

DFlash 并行预测整个未来位置块。一个已由目标模型验证的锚点标记启动该块，其余被掩码的位置同时填充。目标隐藏状态被投影为额外的 KV 对，每个草稿层均可访问，从而在整个并行传播过程中保持目标上下文。

DSpark

DSpark 基于 DFlash，通过添加一个轻量级马尔可夫头，在并行主干之后引入标记间的依赖关系。主干为所有位置生成基础 logits；马尔可夫头使用先前选择的草稿标记调整每个位置，实现并行速度与串行一致性的混合。

在 vLLM 中启用推测解码

推测解码通过 --speculative-config 标志配置。JSON 负载指定方法、可选的草稿检查点和 num_speculative_tokens（提议长度）。

# 原生 MTP（无需额外检查点）
vllm serve <target-model> \
  --speculative-config '{"method": "mtp", "num_speculative_tokens": 4}'

# Gemma 4 MTP（需要助手检查点）
vllm serve <target-model> \
  --speculative-config '{"method": "mtp", "model": "google/gemma-4-26B-A4B-it-assistant", "num_speculative_tokens": 4}'

# EAGLE‑3
vllm serve <target-model> \
  --speculative-config '{"method": "eagle3", "model": "RedHatAI/<target>-speculator.eagle3", "num_speculative_tokens": 3}'

# DFlash（并行块）
vllm serve <target-model> \
  --speculative-config '{"method": "dflash", "model": "z-lab/<target>-DFlash", "num_speculative_tokens": 15}'

# DSpark（并行 + 马尔可夫修正）
vllm serve <target-model> \
  --speculative-config '{"method": "dspark", "model": "RedHatAI/<target>-speculator.dspark", "num_speculative_tokens": 7}'

内存考虑 – 原生 MTP 与目标模型共享权重，因此额外 GPU 内存极少。其他所有方法均加载独立的草稿检查点；所需预留空间取决于检查点大小、精度（FP16/INT4）、张量并行度和运行时缓冲区。

基准测试方法

作者在 AMD Instinct MI300X 和 MI355X GPU 上使用 ROCm™ 堆栈测量了 输出标记吞吐量（每秒标记数）。基准测试使用任务导向的数据集（GSM8K、MATH500、HumanEval、MBPP）以反映现实的接受模式。对于每个目标-方法组合，他们对 num_speculative_tokens（N）进行扫描，并记录：

吞吐量比（推测 / 基线）
平均接受长度（MAL） – 每次验证轮次中平均接受的草稿标记数
接受率（AR） – 通过验证的提议标记比例

所有测量值均相对于在同一硬件和软件堆栈上运行的非推测自回归基线进行相对比较。

关键性能发现

观测到的最高加速比

目标模型	方法	数据集	最佳 N	吞吐量比
`gemma‑4‑26B‑A4B‑it`	DFlash	MATH500	7	2.87×
`gemma‑4‑26B‑A4B‑it`	Gemma 4 MTP	MATH500	5	2.74×
`gemma‑4‑26B‑A4B‑it`	EAGLE‑3	GSM8K	5	2.27×
`gemma‑4‑31B‑it`	DFlash	MATH500	7	2.41×
`Qwen3‑8B`	DSpark	GSM8K	7	1.63×
`Qwen3‑6‑35B‑A3B`	DFlash	MATH500	7	2.06×
`Kimi‑K2.5`	DFlash	GSM8K	7	2.37×

各方法的趋势

串行草稿（原生 MTP、Gemma 4 MTP、EAGLE‑3） 通常在适度的提议长度（N ≈ 3–5）内提升吞吐量。超过此范围后，额外的串行草稿工作量超过节省的目标模型前向传播次数，导致吞吐量趋于平稳或下降。
并行草稿（DFlash、DSpark） 通常在较大的 N（N ≈ 7–11）时达到最高比率。后期位置的接受率下降，但由于并行成本低，整体吞吐量仍持续上升。
DSpark 添加了轻量级马尔可夫头；对于相同的 N，其吞吐量通常比纯 DFlash 低几个百分点，因为额外的串行修正步骤增加了开销。
模型家族很重要 – Qwen 3.5/3.6 模型内置的原生 MTP 在相同模型尺寸下表现优于 DFlash，而 Gemma 4 模型中并行方法占主导地位。
工作负载依赖性 – 代码密集型数据集（HumanEval、MBPP）通常偏好较短的提议，因为后期标记的可预测性降低；而数学密集型数据集（GSM8K、MATH500）可容忍更长的提议。

接受行为

在所有实验中，第一个草稿标记的接受率超过 90%。接受率随位置单调下降；在 N = 15 时，DFlash 的第 15 个标记可能仅在 < 5% 的情况下被接受。平均接受长度（MAL）与吞吐量强相关：MAL 越高 → 目标模型前向传播次数越少 → 加速比越高。

实用调优指南

从小开始 – 对任何方法，从 num_speculative_tokens = 1（无额外成本）开始，并验证正确性。
扫描 N – 在代表性工作负载上逐步增加 N（例如 1–3–5–7–11–15），同时测量吞吐量、MAL 和 AR。
观察逐位置 AR – 如果在位置 k 之后接受率急剧下降，将 N 减少到 k，以避免浪费草稿工作。
内存预算 – 确保有足够的 GPU 内存用于草稿检查点；并行方法（DFlash/DSpark）通常需要最多内存。
硬件特定调优 – AMD Instinct GPU 可从 ROCm 优化的注意力后端（triton_attn）中受益。并行草稿利用 MI300X/MI355X 的高内存带宽。
批大小与标记限制 – 更大的 max-num-batched-tokens 和更长的上下文窗口有助于摊销草稿成本，尤其对并行方法有益。

如何获取草稿检查点

发布者	支持的方法	示例检查点
Google	Gemma 4 MTP	`google/gemma-4-26B-A4B-it-assistant`, `google/gemma-4-31B-it-assistant`
LightSeek	EAGLE‑3, EAGLE‑3.1	`lightseekorg/kimi-k2.5-eagle3-mla`
Red Hat AI	EAGLE‑3, DFlash, DSpark	`RedHatAI/gemma-4-26B-A4B-it-speculator.eagle3`, `RedHatAI/gemma-4-31B-it-speculator.dspark`
Z‑Lab	DFlash	`z-lab/gemma-4-26B-A4B-it-DFlash`, `z-lab/Qwen3.8B-DFlash-b16`
DeepSeek AI	EAGLE‑3, DFlash, DSpark	`deepseek-ai/eagle3_qwen3_8b_ttt7`, `deepseek-ai/dflash_qwen3_8b_block7`
Inferact	EAGLE‑3, DSpark	`Inferact/MiniMax-M3-EAGLE3`, `Inferact/Kimi-K3-DSpark`

训练新推测器（高层次工作流）

收集代表性提示 用于预期工作负载（聊天、代码、数学等）。
使用与推理时完全相同的分词器、聊天模板和采样设置 生成目标模型响应。
选择隐藏状态提取模式 – 在线（实时）、离线（预先保存）或混合（缓存第一轮）。
提取所需的目标层（例如，EAGLE‑3 的早期、中期、晚期层；DFlash/DSpark 的所有选定层）。
使用与目标模型相同的词汇表、隐藏大小和标记嵌入 训练推测器。包含方法特定的头部（并行块、马尔可夫修正等）。
验证接受率、MAL 和端到端吞吐量；如果接受率低，迭代调整提示混合或训练超参数。
打包检查点，并使用 --speculative-config 与目标模型一起提供服务。

未来方向

非学习型推测（例如，n-gram 后缀预测）可与学习型推测器结合，用于高度重复的代码编辑工作负载。
更广泛的并发研究 – 在多用户批处理负载、不同批大小和不同采样温度下测量推测解码。
推测器数据影响 – 系统分析提示多样性与领域特定训练数据如何影响代码、数学、聊天和多语言任务中的接受率。
更深入的剖析 – 在 AMD GPU 上隔离草稿生成、目标验证、KV 缓存重用和图执行开销，以指导内核级优化。

致谢

感谢 Hongxia Yang 和 Peng Sun（AMD）以及 Pin Siang Tan、Jun Kang Chow 和 Ye Hur Cheong（Embedded LLM）提供的硬件访问与合作。

免责声明

测量在两种 AMD Instinct 配置上进行：

硬件 1：8 × MI300X GPU（gfx942），配备双 EPYC 9654 96 核 CPU。
硬件 2：8 × MI355X GPU（gfx950），配备双 EPYC 9575F 64 核 CPU（用于 MiniMax‑M3‑MXFP8）。软件堆栈：Ubuntu 22.04.5 LTS，ROCm 7.2.53211，vLLM 0.23.1rc1.dev1120+g0f0f28b53，PyTorch 2.11.0，Transformers 5.13.1，Python 3.12.13。结果可能因不同驱动程序、内核版本或硬件配置而异。