菜单
主题
vLLM Logo

投机采样如何将 vLLM 性能提升高达 2.8 倍

10 分钟阅读
vLLM 团队
#推测解码#性能

vLLM 中的推测解码是一种强大的技术，通过结合使用小型和大型模型来加速 token 生成。在本篇博客中，我们将解析 vLLM 中的推测解码、其工作原理以及它带来的性能提升。

本内容基于我们双周一次的 vLLM 办公时间（Office Hours）会议，在该会议上我们讨论优化 vLLM 性能的技术和更新。您可以在此处查看会议幻灯片。如果您更喜欢观看视频，可以在 YouTube 上观看完整录像。我们非常欢迎您参加未来的会议 - 请注册！

推测解码简介

推测解码（Leviathan 等人，2023）是降低大语言模型（LLM）token 生成延迟的关键技术。这种方法利用较小的模型来处理简单的 token 预测，同时利用较大的模型来验证或调整这些预测。通过这样做，推测解码在不牺牲准确性的前提下加速了生成过程，成为优化 LLM 性能的一种无损且高效的方法。

为什么推测解码能降低延迟？ 传统上，LLM 以自回归方式一次生成一个 token。例如，给定一个提示词，模型生成三个 token T1、T2、T3，每个都需要单独的前向传递。推测解码改变了这一过程，允许在一次前向传递中提出并验证多个 token。

其工作流程如下：

  1. 草稿模型 (Draft Model)：一个更小、更高效的模型逐个提出 token。
  2. 目标模型验证 (Target Model Verification)：更大的模型在单次前向传递中验证这些 token。它确认正确的 token 并纠正任何错误的 token。
  3. 单次传递处理多个 token：该方法不是每次传递生成一个 token，而是同时处理多个 token，从而降低了延迟。

如上图所示，草稿模型提出了五个 token：["I", "like", "cooking", "and", "traveling"]。随后将这些 token 转发给目标模型进行并行验证。在这个例子中，第三个 token "cooking"（应该是 "playing"）被错误地提出了。因此，在此步骤中仅生成了前三个 token：["I", "like", "playing"]。

通过使用这种方法，推测解码加速了 token 生成，使其成为适用于小型和大型语言模型部署的有效方法。

推测解码如何在 vLLM 中工作

在 vLLM 中，推测解码与系统的连续批处理 (continuous batching) 架构集成在一起，不同的请求在同一个批次中被处理，从而实现了更高的吞吐量。vLLM 使用两个关键组件来实现这一点：

  • 草稿运行器 (Draft Runner)：负责执行较小的模型以提出候选 token。
  • 目标运行器 (Target Runner)：通过运行较大的模型来验证这些 token。

vLLM 的系统经过优化，可以高效处理此过程，使推测解码能够与连续批处理无缝协作，从而提高整体系统性能。


图示：展示了草稿运行器和目标运行器如何在 vLLM 批处理系统中交互。

为了在 vLLM 中实现推测解码，必须修改两个关键组件：

  1. 调度器 (Scheduler)：调整了调度器以处理单次前向传递内的多个 token 插槽，从而实现多个 token 的同时生成和验证。
  2. 内存管理器 (Memory Manager)：内存管理器现在同时处理草稿模型和目标模型的 KV 缓存，确保推测解码过程中的流畅处理。

vLLM 中推测解码的系统架构。

vLLM 支持的推测解码类型

vLLM 支持三种类型的推测解码，每种类型都针对不同的工作负载和性能需求量身定制：

基于草稿模型的推测解码

这是最常用的推测解码形式，即较小的模型预测下一个 token，较大的模型进行验证。一个常见的例子是使用 Llama 68M 模型为 Llama 2 70B 模型预测 token。这种方法需要仔细选择草稿模型，以平衡准确性和开销。

选择正确的草稿模型对于最大限度地提高推测解码的效率至关重要。草稿模型需要足够小以避免产生显著开销，但又要足够准确，以提供有意义的性能提升。

然而，选择合适的草稿模型可能具有挑战性。例如，在 Llama 3 等模型中，由于词汇量大小的差异，很难找到合适的草稿模型。推测解码要求草稿模型和目标模型共享相同的词汇表，在某些情况下，这会限制推测解码的使用。因此，在接下来的章节中，我们将介绍几种无需草稿模型的方法。

提示查找解码 (Prompt Lookup Decoding)


提示查找解码示例。给定提示词，我们构建所有的 2-gram 作为查找键。对应的值是紧跟在查找键后面的三个 token。在生成过程中，我们会检查当前的 2-gram 是否匹配任何键。如果匹配，我们将使用该值提出后续的 token。

这种方法也称为 n-gram 匹配，对于总结和问答等用例非常有效，因为这些任务中提示词和答案之间存在大量重叠。系统不是使用小模型来提出 token，而是基于提示词中已有的信息进行推测。当大型模型在答案中重复提示词的部分内容时，这种方法效果尤为显著。

Medusa/Eagle/MLPSpeculator


图片来自 https://github.com/FasterDecoding/Medusa。在该示例中，三个“头”（heads）被用于为接下来的三个位置提出 token。头 1 为第一个位置提出 ["is", "\'", "the"]。头 2 为第二个位置提出 ["difficult", "is", "\'"]。头 3 为第三个位置提出 ["not", "difficult", "a"]。所有头都以最后一个 Transformer 块的输出作为输入。

在这种方法中，将额外的层（或头）添加到大型模型本身，使其能够在单次前向传递中预测多个 token。这减少了对单独草稿模型的需求，而是利用大型模型自身的能力进行并行 token 生成。尽管尚处于初步阶段，但随着更优化的内核开发，这种方法在提高效率方面显示出巨大潜力。

推测解码性能洞察：加速比与权衡

推测解码在低 QPS（每秒查询数）环境中提供了显著的性能优势。例如，在 ShareGPT 数据集的测试中，当使用基于草稿模型的推测解码时，vLLM 在 token 生成方面实现了高达 1.5 倍的加速。同样，当应用于 CNN/DailyMail 等总结数据集时，提示查找解码显示出高达 2.8 倍的加速。

   

性能比较显示：在 4xH100 上运行 Llama3-70B 和 ShareGPT 数据集时，使用草稿模型（turboderp/Qwama-0.5B-Instruct）在 QPS=1 时实现了 1.5 倍加速；在 4xH100 上运行 Llama3-70B 和 CNN Dailymail 数据集时，使用 n-gram 在 QPS=1 时实现了 2.8 倍加速。

然而，在高 QPS 环境中，推测解码可能会引入性能权衡。当系统已经处于计算密集状态时，提出和验证 token 所需的额外计算有时会减慢系统速度，正如每秒请求数增加时所看到的那样。在这种情况下，推测解码的开销可能会超过其收益，导致性能下降。


在高 QPS 下，我们观察到 Llama3-70B 在 ShareGPT 和 4xH100 上有 1.4 倍的减速，在 CNN Dailymail 和 4xH100 上有 1.8 倍的减速。

路线图：实现更高性能的动态调整

为了克服推测解码在高 QPS 设置下的局限性，vLLM 正在致力于实现动态推测解码 (dynamic speculative decoding)。欢迎查阅论文了解更多细节。这也是 vLLM 的积极研究方向之一！该功能将允许 vLLM 根据系统负载和草稿模型的准确性来调整推测 token 的数量。简而言之，动态推测解码会在系统负载较高时缩短提议的长度。然而，如下图所示，当平均 token 接受率较高时，缩减幅度较小。


未来，系统将能够自动修改每一步的推测程度，确保无论工作负载如何，推测解码始终是有益的。这将允许用户在无需担心是否会拖慢系统的情况下激活推测解码。

如何在 vLLM 中使用推测解码

在 vLLM 中设置推测解码非常简单。启动 vLLM 服务器时，只需包含必要的标志来指定推测模型、token 数量和张量并行大小即可。

以下代码配置 vLLM 以离线模式使用带有草稿模型的推测解码，每次推测 5 个 token：

from vllm import LLM
 
llm = LLM(
    model="facebook/opt-6.7b",
    speculative_model="facebook/opt-125m",
    num_speculative_tokens=5,
)
outputs = llm.generate("The future of AI is")
 
for output in outputs:
    print(f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}")

以下代码配置 vLLM 使用推测解码，其中建议通过匹配提示词中的 n-gram 生成：

from vllm import LLM
 
llm = LLM(
    model="facebook/opt-6.7b",
    speculative_model="[ngram]",
    num_speculative_tokens=5,
    ngram_prompt_lookup_max=4,
    ngram_prompt_lookup_min=1,
)
outputs = llm.generate("The future of AI is")
 
for output in outputs:
    print(f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}")

有时，您可能希望草稿模型以与目标模型不同的张量并行大小运行，以提高效率。这允许草稿模型使用更少的资源并减少通信开销，将更耗费资源的计算留给目标模型。在 vLLM 中，您可以配置草稿模型使用 1 的张量并行大小，而目标模型使用 4，如下例所示。

from vllm import LLM
 
llm = LLM(
    model="meta-llama/Meta-Llama-3.1-70B-Instruct",
    tensor_parallel_size=4,
    speculative_model="ibm-fms/llama3-70b-accelerator",
    speculative_draft_tensor_parallel_size=1,
)
outputs = llm.generate("The future of AI is")
 
for output in outputs:
    print(f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}")

未来的更新（论文, RFC）将允许 vLLM 自动选择推测 token 的数量，从而无需手动配置，进一步简化流程。

请查阅我们的文档 vLLM 中的推测解码 以开始使用。加入我们的双周办公时间会议进行提问和反馈。

结论：vLLM 中推测解码的未来

vLLM 中的推测解码提供了显著的性能提升，特别是在低 QPS 环境中。随着动态调整机制的引入，它在极高 QPS 环境中也将变得非常有效，使其成为降低延迟和提高 LLM 推理效率的多功能必备功能。