vLLM BlockManager - NaiveBlockAllocator

在上一篇博客中，我们深入解析了 vLLM 调度器的调度策略，了解了它如何通过任务分配、优先级排序、抢占等机制在高并发请求环境下提升性能。而这一切的背后，内存管理的效率至关重要。调度器必须在 CPU 和 GPU 之间灵活分配内存，以支持生成任务的顺利执行。这一任务的核心正是由 BlockSpaceManager 来完成。

在 vLLM 中，BlockSpaceManager 和 BlockAllocator 共同承担了生成过程中内存分配、动态调整和缓存管理的职责。它们直接影响到高效处理 waiting、running 和 swapped 三个状态队列中的请求：如何在不同阶段为任务分配内存资源，如何优化 GPU 和 CPU 间的数据交换，如何避免内存瓶颈。这些都是 BlockSpaceManager 需要解决的问题。

本篇博客将介绍 BlockSpaceManager 的设计和实现，聚焦其在不同状态队列中如何分配和管理内存资源。需要注意的是，BlockAllocator 目前支持多种类型的内存分配器。在这里，我们将重点介绍 NaiveBlockAllocator 的实现，不涉及更为复杂的 PrefixCachingBlockAllocator。

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

1. BlockSpaceManager 的架构概览

在 vLLM 系统中，BlockSpaceManager 负责为调度器的请求提供动态内存管理，其架构设计旨在实现高效的 GPU/CPU 内存分配和切换。为了支持调度器在不同请求状态下（如 waiting、running、swapped）的内存需求，BlockSpaceManager 采用了 BlockAllocator 模块对物理内存进行细化管理。NaiveBlockAllocator 是 BlockAllocator 的一种实现，提供了基础的内存块分配和管理能力，适用于不需要复杂缓存机制的场景。

1.1 BlockSpaceManager 管理策略

BlockSpaceManager 的管理策略围绕三个核心目标展开：分配、动态扩展、和交换内存。在调度器处理请求的过程中，不同状态的请求会对内存提出不同的需求。BlockSpaceManager 通过 BlockAllocator 将内存资源拆分为小块（blocks），并灵活地将这些块分配到 waiting、running 和 swapped 状态队列中的请求。

分配：当请求进入 waiting 状态时，BlockSpaceManager 创建并分配一个初始 BlockTable，用于预填充所需的内存。
动态扩展：对于 running 状态的请求，在生成新 token 时，需要追加更多内存块以支持解码阶段。BlockAllocator 动态扩展 BlockTable，确保请求可以在 GPU 内存中进行连续解码。
交换：当 GPU 内存资源不足时，部分 running 队列的请求将转移至 swapped 队列。BlockSpaceManager 通过 NaiveBlockAllocator 将内存块从 GPU 移至 CPU，腾出 GPU 空间以供高优先级任务使用。

1.2 各组件的角色

在 BlockSpaceManager 中，内存块的分配与管理主要由以下几个关键组件协作完成：

BlockAllocator：提供底层的内存块分配接口，包括块的分配、扩展、和交换。NaiveBlockAllocator 是该接口的一种简单实现，通过一系列的块管理方法（如 allocate_immutable_blocks, append_slots, swap_out 等）来支持基本的内存操作。
BlockTable：表示一个请求的内存块表，用于记录和追踪分配给特定 Sequence 或 SequenceGroup 的物理块。BlockTable 是内存管理的核心数据结构，提供了内存块的分配和回收机制。
Sequence 和 SequenceGroup：代表了请求中的一系列 token 序列。Sequence 是基本的序列单位，而 SequenceGroup 则是多个 Sequence 的集合。当一个请求进入调度器时，BlockSpaceManager 会为 SequenceGroup 创建一个或多个 BlockTable，用于存储序列所需的内存块。
Block：Block 是内存的基本单位。在 NaiveBlockAllocator 中，Block 表示一个固定大小的内存区域，用于存储 token IDs。在解码阶段，每当需要生成新 token 时，BlockSpaceManager 会在 BlockTable 中追加新的 Block。

2. 内存分配与管理机制

2.1 调度器何时使用 BlockAllocator

在 vLLM 调度器中，BlockAllocator 作为内存管理的核心模块，负责为生成任务中的不同阶段提供动态的内存块分配。具体来说，BlockAllocator 在 prefill（预填充）和 decode（解码）两个阶段起到关键作用。这两个阶段对应着生成过程中任务状态的变化，而 BlockAllocator 则根据这些变化为 waiting 队列和 running/swapped 队列中的请求提供所需的内存支持。