vLLMの物理块管理

原创

已于 2024-07-19 16:43:37 修改 · 1.2k 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能

于 2024-05-02 22:44:04 首次发布

这里讲解了vLLM的物理块管理（block manager）的细节，包括物理块结构，逻辑块-物理块映射，物理块新增与释放，prefix caching等等。

大家好，本篇我们进入“有趣的”（😊，反复告诉自己它很有趣，那么它一定能变得有趣起来）且“重要的”（这是真得很重要）的vllm块管理器相关代码解读。

vllm块管理器又分成朴素块管理器（UncachedBlockAllocator）和prefix caching型块管理器（CachedBlockAllocator）。本篇我们先讲比较简单的前者，下篇我们来细看更有趣也是更难的后者。

一、前情提要

在之前对调度器策略（Scheduler）的讲解中，我们主要说明了以下几点：

从vLLM批处理的入口函数开始，介绍了其推理内核LLMEngine的两个重要函数add_request()和step()
在LLMEngine开始处理请求前（实例化阶段），它会先做一次模拟实验，来估计gpu上需要预留多少显存给KV Cache block。
当LLMEngine开始处理请求时(add_request)，它会把每个prompt当成一个请求，同时把它包装成一个SequenceGroup对象。
当LLMEngine开始执行1次调度时（step），调度器策略(Scheduler)会选择要送哪些seq_group去做新一轮推理。注意，在1次推理中，所有seq_group要么一起做prefill，要么一起做decode。

调度器策略流程图清晰版可参见下图

同时，我们遗留了以下问题：

问题1：vLLM的物理块管理（block manager）的细节，包括物理块结构，逻辑块-物理块映射，物理块新增与释放，prefix caching等等
问题2：step()其余步骤：调度器只是决定了要送哪些seq_group去做推理，但是“每1个推理阶段结束后，如何根据推理结果更新seq_group，并将其送入下一次调度”这块不是调度器的职责，这也是后面我们要讲解的“step()的其余步骤”.

今天我们就要对问题1进行解答。问题2我们放在源码解读第四篇进行讲解。

二、两种不同类型的BlockAllocator

在源码解读2中，我们画过Schduler的架构图，它的下面维护着今天我们要细讲的块管理器（BlockManager），这也是vLLM自定义的一个class。

截止本文写作时，vLLM提供了BlockSpaceManagerV1和BlockSpaceManagerV2两个版本的块管理器。V1是vLLM默认的版本，V2是改进版本（但还没开发完，例如不支持prefix caching等功能）。所以本文依然基于BlockSpaceManagerV1进行讲解。

BlockManager这个class下又维护着两个重要属性：

BlockAllocator：物理块分配者，负责实际为seq做物理块的分配、释放、拷贝等操作。其下又分成self.gpu_allocator和self.cpu_allocator两种类型，分别管理gpu和cpu上的物理块。
self.block_tables：负责维护每个seq下的物理块列表，本质上它是一个字典，形式如{seq_id: List[PhysicalTokenBlock]}。注意，这个字典维护着【所有】seq_group下seq的物理块，而不是单独某一个seq的。因为调度器是全局的，所以它下面的的BlockManager自然也是全局的。

其中，BlockAllocator又分成两种类型：

CachedBlockAllocator：按照prefix caching的思想来分配和管理物理块。在原理篇中，我们提过又些prompts中可能含有类似system message（例如，“假设你是一个能提供帮助的行车导航”）等prefix信息，带有这些相同prefix信息的prompt完全可以共享用于存放prefix的物理块，这样既节省显存，也不用再对prefix做推理。
UncachedBlockAllocator：正常分配和管理物理块，没有额外实现prefix caching的功能。

在块管理器的上篇中，我们介绍UncachedBlockAllocator，在下篇中我们介绍更为复杂的CachedBlockAllocator。

三、物理块和逻辑块结构

首先我们来快速回顾下在vllm中一个物理块和一个逻辑块长什么样子。

物理块结构（一切尽在注释中）：

# vllm/block.py
class PhysicalTokenBlock:
    """Represents the state of a block in the KV cache."""

    def __init__(
        self,
        device: Device,
        block_number: int,
        block_size: int,
        block_hash: int,
        num_hashed_tokens: int,
    ) -> None:
        # ==============================================================
        # 设备，gpu/cpu
        # ==============================================================
        self.device = device
        # ==============================================================
        # 该物理块在对应设备上的全局block index
        # ==============================================================
        self.block_number = block_number
        # ==============================================================
        # 该物理块的尺寸（即槽位数量，默认为16）
        # ==============================================================
        self.block_size = block_size
        # ==============================================================
        # 该物理块的hash值
        # （在prefix caching场景下使用，非此场景则附值为-1）
        # ==============================================================
        self.block_hash = block_hash 
        # ==============================================================
        # 该物理块的hash值是由多少个前置token计算而来的
        # （prefix caching场景下使用，非此场景则附值为0）
        # ==============================================================
        self.num_hashed_tokens = num_hashed_tokens 
        # ==============================================================
        # 该物理块被多少个逻辑块引用
        # ==============================================================
        self.ref_count = 0
        # ==============================================================
        # 该物理块最后一次被使用的时间
        # （prefix caching场景下使用，非此场景则附值为-1）
        # ==============================================================
        self.last_accessed = DEFAULT_LAST_ACCESSED_TIME
        # ==============================================================
        # 该物理块是否被计算过
        # （prefix caching场景下使用）
        # ==============================================================
        self.computed = False

    def __repr__(self) -> str:
        return (f'PhysicalTokenBlock(device={self.device}, '
                f'block_number={self.block_number}, '
                f'num_hashed_tokens={self.num_hashed_tokens}, '
                f'ref_count={self.ref_count}, '
                f'last_accessed={self.last_accessed}, '
                f'computed={self.computed})')

这里有一些和prefix caching相关的物理块属性，大家现在可能还看得一头雾水，不要担心，在块管理器的下篇中我们再来细讲，这里可以忽略。

逻辑块结构（一切尽在注释中）：

# # vllm/block.py
class LogicalTokenBlock:
    """A block that stores a contiguous chunk of tokens from left to right.

    Logical blocks are used to represent the states of the corresponding
    physical blocks in the KV cache.
    
    KV cache的逻辑块
    """

    def __init__(
        self,
        block_number: int, # 逻辑块的序号
        block_size: int, # 每个逻辑块中有多少个槽位（默认为16）
    ) -> None:
        self.block_number = block_number
        self.block_size = block_size

        # 逻辑块刚初始化时，将其中的每个token_id都初始化为_BLANK_TOKEN_ID（-1）
        self.token_ids = [_BLANK_TOKEN_ID] * block_size 
        # 当前逻辑块中已经装下的token的数量
        self.num_tokens = 0

    def is_empty(self) -> bool:
        """判断当前逻辑块是为空"""
        return self.num_tokens == 0

    def get_num_empty_slots(self) -> int:
        """当前逻辑块的空余槽位"