vLLM (4) - LLMEngine上篇

系列文章目录

vLLM (1) - Qwen2推理&部署
vLLM (2) - 架构总览
vLLM (3) - Sequence & SequenceGroup
vLLM (4) - LLMEngine上篇
vLLM (5) - LLMEngine下篇
vLLM (6) - Scheduler & BlockSpaceManager

---

---

前言

经过前面两篇的铺垫，终于来到了解析LLMEngine的篇章。如下图所示，LLMEngine主要有两部分构成，右边部分包括Worker、CacheEngine和ModelRunner等重要的类，它们在LLMEngine的初始化阶段就会用到，工作内容包括模型加载，KV Cache初始化等等，这是本文中重点；左边部分包括Scheduler和BlockSpaceManger，用于调度用户请求，并在过程中管理显存和内存，这部分发生在LLMEngine的（generate）生成阶段，将放到后续文章中。

---

一、类图

本篇重点讲述LLMEngine的初始化部分。由于代码调用相对复杂，下面我使用类图的方式来表示不同的类之间的关系。同时，在类图中只写上本篇所涉及的相关属性和方法，避免其他属性和方法对本篇阅读造成干扰。建议该类图当结合后续代码一起使用。

# 类图

+-------------------------+
|           LLM           |
+-------------------------+
| + llm_engine: LLMEngine |
+-------------------------+
            |
            |
            v
+-------------------------+
|        LLMEngine        |
+-------------------------+
| + model_executor: GPUExecutor |  # 执行器，名字有点歧义，项目有个子目录也叫model_exectuor
| - _initialize_kv_caches()     |  # 初始化kv_caches
| + scheduler: Scheduler  |        # 调度器
| + output_processor      |        # 输出处理器
+-------------------------+
            |
            |
            v
+-------------------------+
|        GPUExecutor      |        
+-------------------------+
| - _init_executor()      |        # 初始化执行器
| + driver_worker: Worker |        # worker
|                         |
| + determine_num_available_blocks: Tuple[int, int] |   # 确认可用的gpu blocks和cpu blocks
| + initalize_cache()     |        # 初始化缓存，先用全0张量为kv_cache占住内存
+-------------------------+
            |
            |
            v
+-------------------------+
|          Worker         |
+-------------------------+
| + model_runner: ModelRunner |    # 加载和执行模型的部分
| + cache_engine: CacheEngine |    # 初始化和更新kv_cache的部分
| + init_device()         |        # 初始化设备，gpu
| + load_model()          |        # 加载模型
+-------------------------+
            |                              	|
            |                              	|
            v                              	v
+-------------------------+    +-------------------------+
|       ModelRunner       |    |       CacheEngine       |
+-------------------------+    +-------------------------+
| + loader_model()        |    | + gpu_cache             | 
| + profile_run()         |    | - _allocate_kv_cache(): List[torch.Tensor] |
| + capture_model()       |    | + get_cache_block_size(...): int           |
+-------------------------+    +-------------------------+

二、LLM

LLM是一个在给定prompt和sample paramters时，使用指定的大语言模型生成文本的类；其核心组件为self.llm_engine（LLMEngine的实例化对象），LLM的绝大多数工作由它来完成。
使用LLM的示例代码如下所示。1）构建LLM实例化对象，其初始化部分将完成llm_engine: LLMEngine的创建（本文将重点）；2）处理请求，使用self.generate()方法，完成了资源调度，高效的应对用户请求，输出文本（后续文章讲述）。

# 完整示例见系列文章的Qwen2推理篇
from vllm import LLM

llm = LLM(model=DEFAULT_CKPT_PATH)      # DEFAULT_CKPT_PATH为模型名称或下载到本地的目录
outputs = llm.generate(text, sampling_params)  # text为输入文本，sampling_params是采样参数

三、LLMEngine

LLMEngine主要包含两个部分：1）model_executor；2）scheduler。model_executor主要负责模型相关的部分，比如设备的选择，模型的加载等等；而scheduler用于资源的调度，这部分在会模型推理阶段频繁使用。
结合代码来看一下LLMEngine在初始化环节都在干什么：

1. 创建model_executor：根据model_config等一系列配置创建模型执行器；对于一个不太富裕的从业者来说，我们可能在一块单卡上跑vllm，这时候model_executor是GPUExectuor，如果你使用的硬件是Neuron或者TPU，对应的model_executor就是NeuronExecutor或TPUExecutor；另外，model_config等配置是将输入和默认参数按照功能拆分出的多个配置项，这里不赘述；
2. 初始化kv_caches：借由self.model_exectutor（下一小节展开），确定可用于kv_caches的内存空间，并创建tensor占用这部分内存；在Qwen2推理&部署中的真实显存占用这一小节中，我们已经观察到了这个动作，并做了详细分析，不清楚的可以去看一下；
3. 构建scheduler：资源调度一般都出现在模型推理阶段；
4. 其他：比如创建output_processor等，这部分不是重点。

# vllm/engine/llm_engine.py
class LLMEngine:
    def __init__(self, ...):
        # ...
        
        self.model_executor = executor_class(
            model_config=model_config,
            cache_config=cache_config,
            parallel_config=parallel_config,
            scheduler_config=scheduler_config,
            device_config=device_config,
            lora_config=lora_config,
            vision_language_config=vision_language_config,
            speculative_config=speculative_config,
            load_config=load_config,
        )    # 1) 根据输入配置构建model_executor

        if not self.model_config.embedding_mode:
            self._initialize_kv_caches()   # 2) 初始化kv caches

        # 3) 构建scheduler
        self.scheduler = Scheduler(scheduler_config, cache_config, lora_config)

        # 4) 创建输出处理器，这在最后输出的时候会用到
        # Create sequence output processor, e.g. for beam search or speculative decoding.
        self.output_processor = (
            SequenceGroupOutputProcessor.create_output_processor(
                self.scheduler_config,
                self.detokenizer,
                self.scheduler,
                self.seq_counter,
                self.get_tokenizer_for_seq,
                stop_checker=StopChecker(
                    self.scheduler_config.max_model_len,
                    self.get_tokenizer_for_seq,
                ),
            ))

    def _initialize_kv_caches(self) -> None:
        """Initialize the KV cache in the worker(s).

        The workers will determine the number of blocks in both the GPU cache
        and the swap CPU cache.
        """
        num_gpu_blocks, num_cpu_blocks = (
            self.model_executor.determine_num_available_blocks())

        if self.cache_config.num_gpu_blocks_override is not None:
            num_gpu_blocks_override = self.cache_config.num_gpu_blocks_override
            logger.info(
                "Overriding num_gpu_blocks=%d with "
                "num_gpu_blocks_override=%d", num_gpu_blocks,
                num_gpu_blocks_override)
            num_gpu_blocks = num_gpu_blocks_override

        self.cache_config.num_gpu_blocks = num_gpu_blocks
        self.cache_config.num_cpu_blocks = num_cpu_blocks

        self.model_executor.initialize_cache(num_gpu_blocks, num_cpu_blocks)

四、GPUExectuor

model_executor(比如GPUExecutor)在初始化阶段在干什么呢？GPUExecutor继承自基类ExecutorBase，在self.__init__()中调用了self._init_executor()方法，具体包括如下：

1. 使用self._create_worker()创建worker：实际上是通过WorkerWrapperBase来创建的worker，不同的配置对应不同类型的worker，默认情况下是Worker，当你使用投机采样speculative decoding的时候，则是SpecDecodeWorker（合理使用投机采样能够提升解码效率）；
2. worker初始化设备：self.driver_worker.init_device()；
3. worker加载模型：self.driver_worker.load_model()；
   前面提到，GPUExecutor在被创建之后，还用来完成kv_caches的初始化，如上一节LLMEngine._initialize_kv_caches()方法所示，这其中主要涉及GPUExecutor的两个方法：
4. self.determine_num_available_blocks()：该方法返回了当前可用的gpu_blocks和cpu_blocks的数量；block的意思是将gpu和cpu按照指定的大小block_size进行分块，每一块对应一定大小的显存/内存；
5. i