vLLM 分布式推理

2025年2月17日5 分钟阅读

vLLM 团队

#大规模服务

动机

大模型服务往往会导致内存瓶颈，例如令人头疼的 CUDA 显存溢出 (CUDA out of memory) 错误。为解决此问题，主要有两种解决方案：

降低精度 —— 利用 FP8 及更低比特的量化方法可以减少显存使用。然而，这种方法可能会影响精度和扩展性，且随着模型规模超过数千亿参数，仅靠此方法已不足够。
分布式推理 —— 将模型计算分散到多个 GPU 或节点上，从而实现扩展性和效率。这就是张量并行和流水线并行等分布式架构发挥作用的地方。

vLLM 架构与大语言模型推理挑战

与训练相比，大语言模型（LLM）推理面临独特的挑战：

训练专注于吞吐量且具有已知的静态形状，而推理则要求低延迟并处理动态工作负载。
推理工作负载必须高效管理 KV 缓存、投机解码以及预填充到解码（prefill-to-decode）的转换。
大模型往往 超过单 GPU 容量，需要高级的 并行化策略。

为了解决这些问题，vLLM 提供了：

张量并行：将每个模型层分片（shard）到节点内的多个 GPU 上。
流水线并行：将连续的模型层分段分布到多个节点上。
优化的通信内核和控制平面架构：以最小化 CPU 开销并最大化 GPU 利用率。

vLLM 中的 GPU 并行技术

张量并行 (Tensor Parallelism)

问题：模型超过单张 GPU 容量

随着模型增长，单张 GPU 已无法容纳，因此需要多 GPU 策略。张量并行 将模型权重分片到各 GPU 上，允许并行计算以降低延迟并增强扩展性。

这种方法最初是为 Megatron-LM (Shoeybi et al., 2019) 的训练而开发的，现已在 vLLM 中针对推理工作负载进行了适配和优化。

张量并行依赖于两种主要技术：

列并行 (Column Parallelism)：沿列拆分权重矩阵，并在计算后拼接结果。
行并行 (Row Parallelism)：沿行拆分矩阵，在计算后对部分结果求和。

作为一个具体示例，让我们分解这种并行方式如何应用于 Llama 模型中的 MLP（多层感知机）层：

列并行应用于上投影（up-projection）操作。
逐元素激活函数（例如 SILU）在分片后的输出上进行运算。
行并行用于下投影（down-projection），并配合 All-Reduce 操作来聚合最终结果。

张量并行确保推理计算分布在多个 GPU 上，从而最大限度地利用可用的内存带宽和计算能力。使用该技术，我们可以通过有效倍增内存带宽来改善延迟。这是因为分片模型权重允许多个 GPU 并行访问内存，从而减少了单 GPU 可能遇到的瓶颈。

然而，这需要 GPU 之间具备 高带宽互联（如 NVLink 或 InfiniBand），以最小化因通信成本增加而带来的开销。

流水线并行 (Pipeline Parallelism)

问题：模型超过多 GPU 容量

对于超大模型（例如 DeepSeek R1, Llama 3.1 405B），单个节点可能不够。流水线并行 将模型分片到不同节点上，每个节点处理特定的连续模型层。

工作原理

每个 GPU 加载并处理一组不同的层。
发送/接收操作： 随着计算进行，中间激活值会在 GPU 之间传输。

相比张量并行，这带来了更低的通信开销，因为每个流水线阶段只需要传输一次数据。

流水线并行减少了跨 GPU 的内存限制，但并不能像张量并行那样直接降低推理延迟。为了减轻吞吐量效率低下的问题，vLLM 结合了 高级流水线调度，通过优化微批处理（micro-batch）的执行，确保所有 GPU 保持活跃。

结合张量并行与流水线并行

作为一般经验法则，可以这样考虑并行应用方式：

当互联较慢时，使用 跨节点的流水线并行 和 节点内的张量并行。
如果互联高效（如 NVLink, InfiniBand），张量并行可以跨节点扩展。
明智地结合这两种技术可以 减少不必要的通信开销 并最大化 GPU 利用率。

性能扩展与内存效应

虽然并行的基本原则暗示是线性扩展，但在实践中，由于内存效应，性能提升可能呈超线性。无论是使用张量并行还是流水线并行，由于 KV 缓存可用内存的超线性增加，可能会以非显而易见的方式提升吞吐量。

这种超线性扩展效应的出现，是因为更大的缓存允许更大的批处理量以并行处理更多请求，并实现更好的内存局部性，从而导致 GPU 利用率的提升超出了仅增加计算资源所预期的水平。在上图中，你可以看到在 TP=1 和 TP=2 之间，我们能够将 KV 缓存块的数量增加了 13.9 倍，这使我们能够观察到 3.9 倍的 Token 吞吐量 —— 这远超出了从 1 张 GPU 增加到 2 张 GPU 所预期的线性 2 倍增长。

总结

高效服务大模型需要结合 张量并行、流水线并行 以及 Chunked Prefill（分块预填充） 等性能优化手段。vLLM 通过利用这些技术实现可扩展推理，同时确保跨不同硬件加速器的适应性。随着我们不断完善 vLLM，持续关注 混合专家模型（MoE）的专家并行 和 扩展的量化支持 等新进展，对于优化 AI 工作负载至关重要。

欢迎参加双周办公时间（Office Hours），了解更多关于 LLM 推理优化和 vLLM 的知识！

致谢

感谢 Sangbin Cho (xAI) 提供部分插图的原始构思。

动机