在这个系列中,我们将深入解读 vLLM 的源码实现,包括 vLLM 的代码结构、调度策略、PageAttention 的实现等内容。在开始详细介绍之前,本篇将对 vLLM 的代码结构进行一个简要概述,后续我们会逐步深入。

在阅读本系列源码解读文章之前,你需要了解连续批处理和 PageAttention 的基础概念。本系列会对这些概念进行简单介绍,但不会深入讲解。

本系列的代码基于 vLLM 的 0.6.3 版本介绍

1. vLLM 整体代码架构

picture 0

picture 0

vLLM 的代码架构以 LLMEngine 作为核心组件展开,是整个推理过程的起点。在架构图中,LLMEngine 是处理请求的主要入口,它和 Centralized Controller(中央控制器)在同一个进程中运行,且都位于 CPU 上。LLMEngine 的任务是接收用户请求,并与中央控制器配合,确保请求能够被分配到合适的计算资源中处理。

1.1 Centralized Controller

Centralized Controller 实际上是实现了调度功能的模块,叫做 Scheduler(调度器)。它决定每个推理阶段的任务调度,主要负责选择哪些数据需要送到模型去推理,并且管理 KV Cache 的分配。然而,调度器并不直接处理这些物理缓存块的实际存储与管理,只是分配了缓存块的 ID,而数据的实际存储则由分布式 Worker 处理,这些 Worker 分布在 GPU 上。

在调度器之下,架构图中展示了 BlockSpaceManager 和 BlockAllocator,它们的任务是管理计算所需的内存块。BlockSpaceManager 是一种抽象管理器,而 BlockAllocator 则是真正参与内存分配的类,它有 CPU 和 GPU 两种类型,用于管理不同设备上的内存块。BlockAllocator 的作用不仅限于 GPU 上的内存管理,还包括当 GPU 显存不足时,将某些缓存块卸载到 CPU 并进行管理,以便稍后再重新加载到 GPU 中。这种内存管理机制确保系统可以处理更大规模的模型而不受显存限制。

1.2 Distributed Workers

右侧的 Distributed Workers 是整个推理任务的执行者,图中绿色部分展示了这些 Worker 的结构。你可以将它们理解为多个 GPU 进程,每个 GPU 都有一个对应的 Worker,它们负责模型的实际加载和推理。架构中称它们为 Worker,但更准确的说法可能是 Executor,因为它不仅执行推理,还对这些 Worker 实例进行整体管理和协调。

每个 Worker 中都有两个关键模块:CacheEngine 和 Worker.model(model_runner)。CacheEngine 负责管理实际的缓存数据,确保所有的计算任务都有合适的资源可用。而 Worker.model 则负责模型的加载和执行,利用 PagedAttention 这样的组件来高效处理注意力机制,从而保证推理任务的性能。

2. vLLM 处理请求的流程

在了解了 vLLM 的核心模块之后,我们来看一看当一个请求过来的时候,各个模块是如何协作的。

2.1 初始化并加载模型权重

picture 1

picture 1

在具体执行之前,vLLM 需要初始化并加载模型权重。vLLM 支持从 HF Hub 加载模型,也支持从本地加载模型。在加载模型的过程中,vLLM 会将模型权重加载到 GPU 中,以便后续的推理任务可以直接在 GPU 上执行。

2.1.1 估计 KV Cache 的物理块数量