本文参考自 2 篇博客文章《Speculative Decoding - Deep Dive》、《Speed Up Text Generation with Speculative Sampling on AMD GPUs》，来源链接见文末。
随着 Transformer 模型规模的扩大，推理成本也会随之增加，从而影响延迟和吞吐量。压缩方法（例如量化和蒸馏）以及硬件感知优化（例如 Flash Attention 和 Triton）已被提出来在不同层面上降低计算成本。然而，这些模型要么牺牲了准确性，要么需要对模型实现进行重大更改。
推测采样是一种加速推理（在 70B Chinchilla 模型上可提高 2-2.5 倍）的技术，同时不会改变数学准确性，也不会改变模型配置或训练。此外，它可以与上述优化策略结合使用，以进一步减少文本生成的延迟。在这篇博文中，我们将简要介绍推测采样、其基于草稿（draft）和目标模型的算法方法，以及它在使用 ROCm 的 AMD GPU 上的实现。
理解推测采样
DeepMind 在《使用推测采样加速大型语言模型解码》(https://arxiv.org/pdf/2302.01318)以及Google Research 在《通过推测解码从 Transformers 进行快速推理》(https://arxiv.org/pdf/2211.17192)中提出了推测采样。尽管略有不同，但它们提出了类似的策略来降低大型语言模型中的推理成本。
从高层次来看，推测采样采用辅助生成(https://huggingface.co/blog/a...)方法，使用小型草稿模型生成下一个k token，使用大型目标模型验证生成的 token。使用此方法生成和验证 token 所需的总时间小于仅使用大型目标模型生成 token 所需的时间。
使用草稿模型的理由是，大多数“下一个标记”预测都很简单且显而易见。例如，它们可以是常用短语或预期代词。想象一下，使用数十亿参数模型的前向传递只是为了生成逗号作为下一个标记！较小版本的模型可以做到这一点吗？在大多数情况下可以。
此外，Transformer 模型中的大多数计算更多地是受内存限制，而不是计算限制。矩阵乘法器和大型 KV 缓存的内存带宽使用量会随着参数数量的增加而增加，因此有必要减小模型的大小，从而减少延迟。
此外，推测采样可以并行对 K 个标记进行评分和验证，从而比传统自回归提高速度。
为了说明推测抽样，请考虑输入提示“一天一个苹果”。设置 K=3，草稿模型自回归推测接下来的 3 个标记将是“远离医生”。目标模型将输入提示与预测标记的不同延续进行批处理，以验证这些预测。当满足验证标准且草稿模型的分布接近所有三个标记的目标分布时，所有三个标记都会被接受。除了接受第一个 K 标记的预测之外，目标模型对标记的预测(K+1)也会被接受。在这种情况下，K+1 标记是在大型模型的单次前向传递和草稿模型的前向传递中生成的 K，从而大大减少了延迟。整个过程重复进行，直到达到所需的序列长度。
当然，我们现在很想知道如果任何预测的 token 验证失败会发生什么。为了解决这个问题，作者提出了一个 Modified Rejection Sampling 策略。验证策略（如下面算法 2 中所述）采用目标模型预测 token 的概率 q(x_n+1|x_1,...,x_n)与草稿模型预测 token 的概率 p(x_n+1|x_1,...,x_n)之比。如果此比率高于指定阈值，则接受 token。如果低于该比率，则拒绝 token，并根据两个概率之间的差异 q()-p()从新分布中重新采样。
图片来源为 Chen 等人的 Speculative Sampling 论文(https://arxiv.org/pdf/2302.01318)
事实证明，使用草稿模型进行的预测与使用目标模型进行的预测在数学上具有相同的结果。但是，选择与目标模型高度一致的草稿模型非常重要。选择目标模型的较小版本作为草稿模型可实现尽可能低的延迟。
推测解码 - 深入探究
如今，LLM 在线服务已成为科技行业中越来越受欢迎的服务，每天有成千上万的请求发送到 LLM 服务器，并生成响应并发送回世界各地的客户端。在线服务的性能作为衡量其用户体验和服务质量的关键指标之一，已引起业界和学术界的关注。
在本篇博文中，我们采用最常用的开源 LLM 框架之一 vLLM 作为服务引擎。然后，我们研究了旨在加快服务性能的功能推测解码的有效性。我们对 AMD Instinct™ MI300X GPU 和 AMD 软件堆栈进行了所有基准测试和分析。
如您所见，我们的结果表明，启用推测解码后，vLLM 可实现高达 2.31 倍的加速。在文章中，我们首先展示了两个常用模型（即 Llama 3.1-70B 和 Llama 3.1-405B）在启用和未启用推测解码的情况下的性能。然后，我们提供了在 AMD MI300X GPU 上重现结果的详细步骤。此外，我们还提供了一些关于推测解码所涉及的不同因素的消融研究，包括请求率、输入序列长度、数据集等。
vLLM & 推测解码
vLLM 是一个流行的开源 LLM 服务系统，以其独特的设计而闻名，几乎完全减少了 KV 缓存内存的浪费，并允许在请求内和请求之间灵活共享此缓存，从而进一步减少内存使用量。推测解码是 LLM 服务系统中常用的技术，它通过并行计算多个 token 来加速自回归模型的推理过程。在 vLLM 中，推测解码已得到支持以提高服务效率，并提供各种推测算法和优化技术，例如 FP8 量化。
在这篇文章中，我们主要关注如何采用性能更好的推测解码，例如使用更小的草稿模型、启用量化等。
Llama 3.1-70B 和 Llama 3.1-405B 的性能
我们对两种常用的 LLM 模型 Llama 3.1-70B 和 Llama 3.1-405B 的性能进行了基准测试，以揭示推测解码的有效性。
Llama 3.1-70B
本次实验中，我们以 Llama 3.1-70B 为基础模型，结合 Llama 3.1-8B、Llama 3.2-3B 和 Llama 3.2-1B 三个较小的草稿模型 ，测试推测解码的加速比。表 1 和图 1 展示了启用和未启用推测解码的 Llama 3.1-70B 端到端延迟对比。总体而言，与较小的草稿模型搭配使用时，Llama 3.1-70B 的加速比可达到 2.0 倍以上。
表 1. 使用单个 AMD MI300X GPU 运行同一基础模型 Llama 3.1 70B ，带上不同的较小草稿模型（即 Llama 3.1 8B、Llama 3.2 3B 和 Llama 3.2 1B）的延迟性能。与 Llama 3.2 1B 草稿模型结合使用时，速度可提高 2.31 倍。
数据测量时间为 2025 年 2 月 12 日。端到端延迟表示用户发送请求和收到响应所需的总时间，我们使用每个请求的毫秒数作为其单位。所有实验均在单个 AMD MI300X GPU 上进行，系统环境中安装了 ROCm TM 6.3.1 和 vLLM v0.6.7。
图 1. 以 Llama 3.1 70B 为基础模型，以 Llama 3.1 8B、Llama 3.2 3B 和 Llama 3.2 1B 为草稿模型，推测解码的端到端延迟加速。
Llama 3.1-405B
在本实验中，我们以 Llama 3.1 405B 为基础模型，并将其与三个较小的草稿模型 Llama 3.1-8B、Llama 3.2 3B 和 Llama 3.2 1B 相结合，以测试推测解码的加速比。我们在 4 个 AMD MI300X GPU 上进行了所有实验。表 2 和图 2 显示了启用和未启用推测解码的 Llama 3.1 405B 的端到端延迟比较。与较小草稿模型配对的 Llama 3.1 405B 也实现了 >2.0 倍的加速。
表 2. 使用 4 个 AMD MI300X 加速器，相同基础模型 Llama 3.1 405B 与不同的较小草稿模型（即 Llama 3.1 8B、Llama 3.2 3B 和 Llama 3.2 1B）的延迟性能。与 Llama 3.2 1B 草稿模型结合使用时，速度可提高 2.31 倍。
数据于 2025-02-12 测量。端到端延迟表示用户发送请求到收到响应所需的总时间，我们使用每个请求的毫秒数作为单位。所有实验均在 4 个 AMD MI300X GPU 上进行，系统环境中安装了 ROCm 6.3.1 和 vLLM v0.6.7。
图 2. 以 Llama 3.1 405B 为基础模型，以 Llama 3.1 8B、Llama 3.2 3B 和 Llama 3.2 1B 为草稿模型，推测解码的端到端延迟加速。
重现步骤
硬件设置
所有实验均在 AMD MI300X GPU 上进行。建议使用预装 ROCm 的近似一代 AMD 数据中心 GPU。
Docker
我们在所有实验中使用针对 AMD 优化的统一 vLLM docker 镜像：
rocm/vllm:rocm6.3.1_mi300_ubuntu22.04_py3.12_vllm_0.6.6。首先，从 dockerhub 拉取 docker 镜像：
docker pull rocm/vllm:rocm6.3.1_mi300_ubuntu22.04_py3.12_vllm_0.6.6通过以下命令初始化 docker 容器：
docker run -d -it --ipc=host --network=host --privileged --cap-add=CAP_SYS_ADMIN --device=/dev/kfd --device=/dev/dri --device=/dev/mem --group-add render
--cap-add=SYS_PTRACE --security-opt seccomp=unconfined --name vllm_spec_dec -v
/home/models/:/models -v /home/:/work rocm/vllm:rocm6.3.1_mi300_ubuntu22.04_py3.12_vllm_0.6.6注意：请用您自己的路径替换模型路径“/home/models/:/models”和工作区“/home/:/work”。
模型下载
对于模型检查点，我们从 Hugging Face 下载公共权重。按如下方式设置 Hugging Face 令牌，并请用您自己的 HF_TOKEN 替换该令牌：
export HF_TOKEN=xIxAxMxAxPxLxAxCxExHxOxLxDxExRx在本地下载特定模型。在此示例中，我们从 Hugging Face 上的 AMD repo 下载模型“Llama 3.1 70B Instruct FP8 KV”：
huggingface-cli download --resume-download --local-dir-use-symlinks False amd/Llama-3.1-70B-Instruct-FP8-KV --local-dir amd--Llama-3.1-70B-Instruct-FP8-KV基准测试
本次实验中，我们使用 vLLM 的在线服务模式进行性能基准测试。
不使用推测解码来启动服务：
export PYTORCH_TUNABLEOP_ENABLED=0
export PYTORCH_TUNABLEOP_TUNING=0
export PYTORCH_TUNABLEOP_MAX_TUNING_DURATION_MS=100
export PYTORCH_TUNABLEOP_MAX_WARMUP_DURATION_MS=10
export PYTORCH_TUNABLEOP_ROTATING_BUFFER_SIZE=1024
export PYTORCH_TUNABLEOP_FILENAME=afo_tune_device_%d_full.csv
export HIP_FORCE_DEV_KERNARG=1
export VLLM_USE_ROCM_CUSTOM_PAGED_ATTN=1
export VLLM_INSTALL_PUNICA_KERNELS=1
export TOKENIZERS_PARALLELISM=false
export RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1
export NCCL_MIN_NCHANNELS=112
export VLLM_USE_TRITON_FLASH_ATTN=0
export VLLM_FP8_PADDING=1
export VLLM_FP8_ACT_PADDING=1
export VLLM_FP8_WEIGHT_PADDING=1
export VLLM_FP8_REDUCE_CONV=1
vllm serve /models/models--amd--Meta-Llama-3.1-405B-Instruct-FP8-KV/ --swap-space 16 --disable-log-requests
--tensor-parallel-size 8 --distributed-executor-backend mp --dtype float16 --quantization fp8 --kv-cache-dtype fp8 --enable-chunked-prefill=False --max-num-seqs 300启动服务并启用推测解码：
export PYTORCH_TUNABLEOP_ENABLED=0
export PYTORCH_TUNABLEOP_TUNING=0
export PYTORCH_TUNABLEOP_MAX_TUNING_DURATION_MS=100
export PYTORCH_TUNABLEOP_MAX_WARMUP_DURATION_MS=10
export PYTORCH_TUNABLEOP_ROTATING_BUFFER_SIZE=1024
export PYTORCH_TUNABLEOP_FILENAME=afo_tune_device_%d_full.csv
export HIP_FORCE_DEV_KERNARG=1
export VLLM_USE_ROCM_CUSTOM_PAGED_ATTN=1
export VLLM_INSTALL_PUNICA_KERNELS=1
export TOKENIZERS_PARALLELISM=false
export RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1
export NCCL_MIN_NCHANNELS=112
export VLLM_USE_TRITON_FLASH_ATTN=0
export VLLM_FP8_PADDING=1
export VLLM_FP8_ACT_PADDING=1
export VLLM_FP8_WEIGHT_PADDING=1
export VLLM_FP8_REDUCE_CONV=1
vllm serve /models/models--amd--Meta-Llama-3.1-405B-Instruct-FP8-KV/ --swap-space 16 --disable-log-requests --tensor-parallel-size 4 --distributed-executor-backend mp
--dtype float16 --quantization fp8 --kv-cache-dtype fp8 --enable-chunked-prefill=False --max-num-seqs 300 --speculative-model /models/models--amd--Meta-Llama-3.1-8B-Instruct-FP8-KV/ --num_speculative_tokens 5
--speculative-model-quantization fp8对于客户端，我们使用 SGLang 的 docker 镜像：lmsysorg/sglang:v0.4.1.post4-rocm620。使用以下命令初始化 docker 镜像：
docker run -d -it --ipc=host --network=host --privileged --cap-add=CAP_SYS_ADMIN --device=/dev/kfd --device=/dev/dri --device=/dev/mem --group-add render --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --shm-size=192g
--name sglang_client -v /home/models/:/models -v /home/:/work lmsysorg/sglang:v0.4.1.post4-rocm620在客户端容器中，使用不同的请求率对性能进行基准测试：
#!/bin/bash
set -xeu
python3 -m sglang.bench_serving --backend vllm --dataset-name random --num-prompt 500 --request-rate 0.5 --random-input 8192 --random-output 256 > vllm_log0.5
python3 -m sglang.bench_serving --backend vllm --dataset-name random --num-prompt 500 --request-rate 1.0 --random-input 8192 --random-output 256 > vllm_log1.0
python3 -m sglang.bench_serving --backend vllm --dataset-name random --num-prompt 500 --request-rate 1.5 --random-input 8192 --random-output 256 > vllm_log1.5
python3 -m sglang.bench_serving --backend vllm --dataset-name random --num-prompt 500 --request-rate 2.0 --random-input 8192 --random-output 256 > vllm_log2.0
python3 -m sglang.bench_serving --backend vllm --dataset-name random --num-prompt 500 --request-rate 2.5 --random-input 8192 --random-output 256 > vllm_log2.5
python3 -m sglang.bench_serving --backend vllm --dataset-name random --num-prompt 500 --request-rate 3.0 --random-input 8192 --random-output 256 > vllm_log3.0消融研究（Ablation Study）
请求率
在线 LLM 服务系统中，一个重要的因素是请求率。一般来说，它衡量的是每秒处理的请求数。在这项消融研究中，我们希望调查推测解码的有效性和请求率之间的联系。
我们将请求率设置为 0.2、0.4、0.6、0.8、1.0、1.2 和 1.4，其他设置保持不变。所有性能数据点都记录在下表 3 中。可以推断，请求率降低时加速比会增加，但差异并不大。
表 3. 请求率的消融研究。使用 4 个 AMD MI300X 加速器（同一基础模型 Llama 3.1 405B 和草稿模型 Llama 3.2 1B）以不同的请求率（0.2、0.4、0.6、0.8、1.0、1.2 和 1.4）进行的延迟性能。随着请求率的增加，加速比略有下降。
图 3. 不同请求率如何影响服务性能的消融研究。图表展示了 Llama 3.1 405B 基础模型和 Llama 3.2 1B 草稿模型在不同请求率下的延迟加速情况。
更长的输入
很多实际场景都需要 LLM 服务系统能够处理长输入上下文，包括自然语言处理（NLP）、科学计算、金融与交易、代码生成等。具体来说，对于 NLP 任务，有文档摘要，例如学术论文摘要、根据参考网页进行问答、文本翻译和语音识别，这些任务的输入包含丰富的信息，用户希望使用 LLM 搜索关键信息。对于科学计算，有气候建模和蛋白质结构预测等情况，LLM 需要消化长输入序列以进行简短的信息预测。
本次实验旨在探究较长的输入序列长度时推测解码对 LLM 服务性能的影响，并使用离线服务模式进行基准测试，以模拟实际应用场景。我们将输入长度设置为 16384 和 32768，并将输出长度统一设置为 128。实验结果记录在下表 4 中。可以观察到，加速比随着输入长度的增加而呈现增加的趋势，这表明当输入序列长度较长时，推测解码可以进一步提高性能。
表 4. 较长输入序列的消融研究。使用同一基础模型、Llama 3.1 70B 和草稿模型 Llama 3.1 8B 的 1 个 AMD MI300X GPU 的延迟性能，输入长度分别为 16384 和 32768。输入长度增加时，加速比也会增加。
不同的数据
对于一般基准测试，人们倾向于采用随机数据集来评估 LLM 服务性能以及推测解码功能，其中输入序列是随机生成的，没有任何文字含义或语法检查。然而，这可能会给基准测试性能带来一些偏差，尤其是当我们将 LLM 服务切换到实际应用时，输入序列通常语法正确。为了更全面地评估推测解码对 LLM 服务性能的有效性，我们使用另一个名为 ShareGPT 的数据集进行了这项消融研究，其中每个输入和输出序列都有相应的含义。
我们在下表 5 中整合了启用和未启用推测解码的 vLLM 在随机数据集（记为“random”）和 ShareGPT（记为“sharegpt”）上的性能。可以观察到，在随机数据集和 ShareGPT 数据集上的加速比没有显著差异。
表 5. 不同数据集的消融研究。在随机数据集和 ShareGPT 数据集上，使用 4 个 AMD MI300X 加速器，运行相同基础模型 Llama 3.1 405B 和 3 个不同草案型号 Llama 3.1 8B、Llama 3.2 3B 和 Llama 3.2 1B 的延迟性能。不同数据集之间的加速比保持一致。
草稿模型量化（Draft Model Quantization）
量化已经成为提升 LLM 服务性能最常用的技术之一，因为它能够有效降低内存成本和计算工作量。对于推测解码的草稿模型，vLLM 支持使用量化来加载其模型权重，用户可以使用 FP8 进一步加速。这项消融研究主要探讨草稿模型量化对提升 LLM 服务性能的有效性。
在下表 6 中，我们比较了 Llama 3.1 70B 搭配启用和未启用量化的草稿模型的性能。我们可以得出结论，在推测解码中为草稿模型启用量化也有效提升了 LLM 服务的性能。
表 6. 草稿模型量化效果的消融研究。使用同一基础模型 Llama 3.1 70B 和同一草稿模型 Llama 3.1 8B 的单个 AMD MI300X 加速器，以三种不同的请求率（1.0、1.2 和 1.4）测试延迟性能。加速比表明，尽管草稿模型 Llama 3.1 8B 比基础模型 Llama 3.1 70B 小得多，但启用量化仍有助于进一步提高性能。
图 4. 量化对草稿模型影响的消融研究。图表展示了 Llama 3.1 70B 基础模型和 Llama 3.1 8B 草稿模型在未量化和量化情况下的延迟加速情况。
总结
在这篇博文中，我们展示了使用开源 vLLM 框架和 ROCm（运行在 AMD MI300X GPU 上）在 LLM 服务中进行推测解码的性能优势。我们评估了使用和不使用推测解码的 Llama 3.1-70B 和 Llama 3.1-405B，结果显示端到端延迟最多可提高 2.31 倍。
我们分享了详细的硬件和软件设置说明和基准测试，以及多项消融研究，这些研究考察了请求率、输入序列长度、数据集类型和草稿模型量化的影响。我们的结果表明，在这些条件下，性能持续改善，尤其是在输入序列较长和量化草稿模型的情况下。推测解码显著提高了 LLM 服务系统的性能，在我们使用 Llama 模型进行的测试中，速度提高了 2.31 倍。我们的消融研究表明，请求率、输入序列长度、数据集和草稿模型量化等因素会影响其有效性。
随着对高效 LLM 服务的需求不断增长，推测解码对于满足性能需求至关重要。我们鼓励进一步探索这项技术，以充分发挥其潜力。如需更多与 ROCm 相关的开发人员内容，请访问 ROCm 开发人员中心。
参考链接 https://rocm.blogs.amd.com/so...
https://rocm.blogs.amd.com/ar...
扩展阅读
《在 AMD MI300X上增强 DeepSeek-R1推理能力：调优至 7,318 Token/s》
《解锁 DeepSeek-R1 671B FP8__推理性能：5,921 token/s @ AMD MI300X》
《企业存储技术》文章分类索引更新（微信公众号合集标签）
END
作者：Chang Liu等
原文：企业存储技术
推荐阅读
- DeepSeek 能否引领 AI PC 发展趋势？(技术篇)
- 在AMD MI300X上增强DeepSeek-R1推理能力：调优至7,318 Token/s
- DeepSeek-671B 纯 CPU 部署：配置选型、性能测试与量化对比
欢迎关注企业存储技术极术专栏，欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
