Neural Magic专家深度解析：vLLM性能提升2.7倍的秘籍

文章导读：大家好，我是Rob，我负责Neural Magic的工程工作，加入公司已有两年时间。正如Michael提到的，我非常积极参与vLLM社区，并且是上游项目的贡献者之一。今天，我想与大家探讨最近备受关注的VM 0.6.0版本发布。

这一版本的亮点是显著提升了H100 GPU上运行Llama 8B和Llama 70B模型的速度。为了让大家更好地了解我们是如何实现这些性能提升的，我将深入讲解这背后的技术细节。

演讲的主题

今天的分享将涵盖以下几个方面：

1. LM（语言模型）推理的基本原理：为何这是一项具有挑战性的任务；

2. vLLM的内部架构：如何支持一种关键优化技术——连续批处理；

3. 性能瓶颈诊断与优化：在VM 6.0版本之前的架构中存在哪些问题，以及我们在过去几个月中如何通过关键优化来解决这些问题；

4. 未来规划：继续提升性能的一些待完成项目；

值得注意的是，本次分享将聚焦于如何降低vLLM调度器层面的开销，而不是讨论GPU的执行性能。如果对GPU执行优化感兴趣，可以参考Tyler上周关于Cutlass的演讲。

LM推理的基本原理

LM的核心功能是预测序列中的下一个Token。具体而言，它根据输入序列的前N-1个Token，计算第N个Token的概率分布。基于这个原理，LM通过逐步生成一个个Token来完成文本生成任务。例如：输入包含3个Token，模型生成第4个Token。接着使用前4个Token生成第5个Token，以此类推。

这种逐次生成意味着生成1000个Token需要通过模型进行1000次推断，因此每次生成可能耗时数秒。用户在与LM交互时往往会感受到这一延迟。这也为推理服务器的设计带来了挑战：如何高效处理多个并发请求？

连续批处理的原理

传统的推理服务器通常采用动态批处理（Dynamic Batching），通过将多个请求合并为一个批次执行，来提升GPU利用率。但这种方法在处理LM推理时效果不佳，因为LM推理的每个请求可能需要多次模型推断。举个例子：

如果一个请求需要生成1000个Token，而其他请求只需生成1个Token，那么批次中的所有请求都需要等待最长请求完成，导致GPU利用率低下且延迟增加。

为了解决这一问题，我们采用了连续批处理的策略。它的核心思路是：

1.按迭代级别而非请求级别进行批处理；

2.在某个请求完成一部分计算后，将其从批次中移除并引入新请求；

这种方法显著提高了GPU利用率并降低了用户感知的延迟。

vLLM的内部架构

为了支持连续批处理，VM架构中设计了一个核心组件：LLM引擎。这个引擎在一个循环中运行，主要包含以下三步：

1.调度：决定哪些任务需要执行，分配内存并管理KV缓存。

2.执行：将调度好的任务传递给GPU工作线程。vLLM原生支持Tensor并行，可以将模型拆分并分布到多个GPU上。

3.处理输出：包括将结果流式传递到API服务器、检查结束标志以及移除已完成的任务。

这套架构在CPU端完成大部分内存管理和状态跟踪工作，从而实现了高效的连续批处理。

挑战与未来

虽然vLLM已经在减少调度开销方面取得了很大进展，但面对H100等新一代GPU的超高性能，我们仍需优化系统的其他部分，例如内存管理和通信开销，以充分发挥硬件潜力。

通过这些改进，vLLM可以为语言模型推理提供更强大的支持，同时继续作为AI生态系统的核心平台。

在大语言模型（LLM）的解码阶段，性能主要受限于将权重从GPU内存传输到芯片内存的速度。在H100 GPU上，其内存带宽为3.35 TB/s。以一个16GB的模型（约8亿参数）为例，从GPU内存传输到芯片内存所需的时间约为5毫秒。GPU内部执行速度极快，每次迭代仅需5到7毫秒。

然而，随着批处理规模从1增加到256，由于内存带宽的限制，执行时间的增长非常有限。例如，批大小为1时，解码100个token的延迟约为0.8秒，而批大小为64时，该延迟略低于1秒。这表明GPU执行时间几乎保持不变，而输入准备、调度和输出处理（主要在CPU上完成）会随着批处理规模线性增长。

CPU的处理包括调度操作和输出处理，运行在单一Python进程中，受制于全局解释器锁（GIL），阻止了任何实际的并发执行。这种架构限制了系统的吞吐量。

性能瓶颈分析与优化方案

在一个实际运行LLM的48秒任务中，仅有25秒用于模型执行，剩余时间大多花在输入准备、调度和输出处理上。而API服务器引入的额外30%的开销，使GPU操作的占比降至38%。因此，提升CPU部分性能成为关键。

优化一：解耦API服务器和LLM引擎

通过与Neural Magic和IBM深度合作，我们将API服务器与LLM引擎拆分为两个独立的进程。Python的GIL限制了多线程的并发能力，因此这种解耦可以避免两部分争夺CPU资源。我们引入了ZeroMQ作为消息传递协议。经过多种协议的评估，ZeroMQ表现出良好的性能和开发便利性。该优化带来了25%的吞吐量提升，未来我们计划在更多核心环节中使用ZeroMQ。

优化二：异步输出处理

CUDA的异步特性使得GPU计算和主机代码可以并行执行。我们调整了LLM引擎的工作方式，使输出处理与GPU前向传播并行，而非同步执行。这一改动通过消除数据依赖性，实现了7.5%的吞吐量提升。

优化三：多步调度

传统的逐步调度需要每次迭代都进行内存分配和调度操作。通过多步调度，我们将多次调度合并为一次，大幅减少CPU工作量。此方法有效降低了输入准备的频率和函数调用的次数，实现了70%的吞吐量提升。

成果与展望

通过上述优化，LLM引擎在CPU侧的开销大幅减少。例如，运行Llama-70B模型时，GPU执行速度限制较慢的情况下，吞吐量依然提升了近2倍，而在执行速度较快的Llama-38模型上，整体速度提升达到了2.7倍。

这一成果得益于Neural Magic、IBM、Databricks、AnyScale和伯克利大学的紧密合作。未来，我们将继续优化现有技术。例如：

1.引入ZeroMQ至核心循环：目前LLM引擎内部使用asyncio管理并发，但其任务切换存在显著开销。通过ZeroMQ替代asyncio，进一步减少切换开销，我们实现了更高的吞吐量。

2.异步流式生成：多步调度的一个问题是每生成一个新token的延迟较高。通过ZeroMQ实现异步流式生成，我们成功降低了token间延迟，同时对整体输出时间几乎没有影响。

3.兼容更多优化策略：我们正在扩展多步调度的适用范围，例如结合预填充和解码批处理的Chunk Prefill技术，以进一步平滑token生成延迟。

通过这些改进，我们希望持续推动LLM解码阶段的性能边界，为用户提供更高效、更流畅的模型推理体验。

在以下情况下，我们假设有一个长度为 8000 的提示（prompt）。我们将其分成 2000 长度的块，然后在 vLLM 内部使用解码操作重叠这些块。通过块预填充（chunk prefill），我们能够实现更平滑的每个输出 token 的生成时间，避免长时间的停顿来处理新的预填充数据并将其引入 BLM。

启用块预填充后（目前已经有一个草案 PR），可以显著减少每个输出 token 所需的时间。这项改进预计将在 6.2 版本中上线，目前仍在积极开发中。我们的目标是持续提升性能。