架构概览#

本文档概述了 vLLM 的架构。

入口点#

vLLM 提供了许多与系统交互的入口点。下图显示了它们之间的关系。

Entrypoints Diagram

LLM 类#

LLM 类为执行离线推理提供了主要的 Python 接口,即在不使用单独的模型推理服务器的情况下与模型交互。

以下是 LLM 类的用法示例

from vllm import LLM, SamplingParams

# Define a list of input prompts
prompts = [
    "Hello, my name is",
    "The capital of France is",
    "The largest ocean is",
]

# Define sampling parameters
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

# Initialize the LLM engine with the OPT-125M model
llm = LLM(model="facebook/opt-125m")

# Generate outputs for the input prompts
outputs = llm.generate(prompts, sampling_params)

# Print the generated outputs
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

更多 API 详情可以在 API 文档的 离线推理 部分找到。

LLM 类的代码可以在 vllm/entrypoints/llm.py 中找到。

OpenAI 兼容 API 服务器#

与 vLLM 交互的第二个主要接口是通过其 OpenAI 兼容 API 服务器。可以使用 vllm serve 命令启动此服务器。

vllm serve <model>

vllm CLI 的代码可以在 vllm/entrypoints/cli/main.py 中找到。

有时您可能会看到直接使用 API 服务器入口点,而不是通过 vllm CLI 命令。例如

python -m vllm.entrypoints.openai.api_server --model <model>

该代码可以在 vllm/entrypoints/openai/api_server.py 中找到。

有关 API 服务器的更多详细信息,请参阅 OpenAI 兼容服务器 文档。

LLM 引擎#

LLMEngineAsyncLLMEngine 类是 vLLM 系统功能的核心,处理模型推理和异步请求处理。

LLMEngine Diagram

LLMEngine#

LLMEngine 类是 vLLM 引擎的核心组件。它负责接收来自客户端的请求并从模型生成输出。LLMEngine 包括输入处理、模型执行(可能分布在多个主机和/或 GPU 上)、调度和输出处理。

  • 输入处理:使用指定的 tokenizer 处理输入文本的 tokenization。

  • 调度:选择在每个步骤中处理哪些请求。

  • 模型执行:管理语言模型的执行,包括跨多个 GPU 的分布式执行。

  • 输出处理:处理模型生成的输出,将语言模型中的 token ID 解码为人类可读的文本。

LLMEngine 的代码可以在 vllm/engine/llm_engine.py 中找到。

AsyncLLMEngine#

AsyncLLMEngine 类是 LLMEngine 类的异步包装器。它使用 asyncio 创建一个后台循环,持续处理传入的请求。AsyncLLMEngine 专为在线服务而设计,它可以处理多个并发请求并将输出流式传输到客户端。

OpenAI 兼容 API 服务器使用 AsyncLLMEngine。在 vllm/entrypoints/api_server.py 中还有一个作为更简单示例的演示 API 服务器。

AsyncLLMEngine 的代码可以在 vllm/engine/async_llm_engine.py 中找到。

Worker#

Worker 是运行模型推理的进程。vLLM 遵循使用一个进程控制一个加速器设备(如 GPU)的常用做法。例如,如果我们使用大小为 2 的张量并行和大小为 2 的流水线并行,我们将总共有 4 个 worker。Worker 通过它们的 ranklocal_rank 标识。rank 用于全局编排,而 local_rank 主要用于分配加速器设备和访问本地资源,例如文件系统和共享内存。

模型运行器#

每个 worker 都有一个模型运行器对象,负责加载和运行模型。大部分模型执行逻辑都驻留在此处,例如准备输入张量和捕获 cudagraphs。

模型#

每个模型运行器对象都有一个模型对象,它是实际的 torch.nn.Module 实例。有关各种配置如何影响我们最终得到的类,请参阅 huggingface_integration

类层次结构#

下图显示了 vLLM 的类层次结构

query

此类的层次结构背后有几个重要的设计选择

1. 可扩展性:层次结构中的所有类都接受包含所有必要信息的配置对象。VllmConfig 类是主要配置对象,它在各处传递。类层次结构相当深入,每个类都需要读取它感兴趣的配置。通过在一个对象中封装所有配置,我们可以轻松地在各处传递配置对象并访问我们需要的配置。假设我们要添加一个只触及模型运行器的新功能(考虑到 LLM 推理领域发展如此之快,这种情况很常见)。我们必须在 VllmConfig 类中添加新的配置选项。由于我们在各处传递整个配置对象,我们只需要将配置选项添加到 VllmConfig 类,模型运行器就可以直接访问它。我们不需要更改引擎、worker 或模型类的构造函数来传递新的配置选项。

2. 统一性:模型运行器需要一个统一的接口来创建和初始化模型。vLLM 支持 50 多种流行的开源模型。每个模型都有自己的初始化逻辑。如果构造函数签名随模型而变化,则模型运行器不知道如何相应地调用构造函数,而无需复杂且容易出错的检查逻辑。通过使模型类的构造函数统一,模型运行器可以轻松创建和初始化模型,而无需知道具体的模型类型。这也适用于组合模型。视觉语言模型通常由视觉模型和语言模型组成。通过使构造函数统一,我们可以轻松创建视觉模型和语言模型,并将它们组合成视觉语言模型。

注意

为了支持此更改,所有 vLLM 模型的签名都已更新为

def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

为了避免意外传递不正确的参数,构造函数现在是仅关键字的。这确保了如果传递旧配置,构造函数将引发错误。vLLM 开发者已经为 vLLM 中的所有模型进行了此更改。对于树外注册的模型,开发者需要更新他们的模型,例如通过添加 shim 代码来使旧的构造函数签名适应新的签名

class MyOldModel(nn.Module):
    def __init__(
        self,
        config,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
        lora_config: Optional[LoRAConfig] = None,
        prefix: str = "",
    ) -> None:
        ...

from vllm.config import VllmConfig
class MyNewModel(MyOldModel):
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        config = vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
        lora_config = vllm_config.lora_config
        super().__init__(config, cache_config, quant_config, lora_config, prefix)

if __version__ >= "0.6.4":
    MyModel = MyNewModel
else:
    MyModel = MyOldModel

这样,模型就可以与旧版本和新版本的 vLLM 一起工作。

3. 初始化时的分片和量化:某些功能需要更改模型权重。例如,张量并行需要对模型权重进行分片,而量化需要量化模型权重。有两种可能的方法来实现此功能。一种方法是在模型初始化后更改模型权重。另一种方法是在模型初始化期间更改模型权重。vLLM 选择后者。第一种方法无法扩展到大型模型。假设我们要使用 16 个 H100 80GB GPU 运行一个 405B 模型(大约有 810GB 权重)。理想情况下,每个 GPU 应该只加载 50GB 权重。如果我们在模型初始化后更改模型权重,我们需要将完整的 810GB 权重加载到每个 GPU,然后再对权重进行分片,从而导致巨大的内存开销。相反,如果我们在模型初始化期间对权重进行分片,则每个层将仅创建它需要的权重的分片,从而导致更小的内存开销。相同的想法适用于量化。请注意,我们还向模型的构造函数添加了一个额外的参数 prefix,以便模型可以根据前缀以不同的方式初始化自身。这对于非均匀量化很有用,其中模型的不同部分以不同的方式量化。prefix 对于顶层模型通常是一个空字符串,对于子模型则是一个字符串,如 "vision""language"。通常,它与检查点文件中模块状态字典的名称匹配。

此设计的一个缺点是很难为 vLLM 中的各个组件编写单元测试,因为每个组件都需要通过完整的配置对象进行初始化。我们通过提供一个默认初始化函数来解决这个问题,该函数创建一个默认配置对象,其中所有字段都设置为 None。如果我们想要测试的组件只关心配置对象中的几个字段,我们可以创建一个默认配置对象并设置我们关心的字段。这样,我们可以隔离地测试组件。请注意,vLLM 中的许多测试都是端到端测试,用于测试整个系统,因此这不是一个大问题。

总而言之,完整的配置对象 VllmConfig 可以被视为引擎级别的全局状态,在所有 vLLM 类之间共享。