图解大模型计算加速系列：vLLM源码解析3，Prefix Caching
原创 猛猿 大猿搬砖简记 2024年07月05日 14:24 北京
大家好，今天我们继续vllm源码的解析，一起来看下它最近总是被频繁提起、也是较不好理解的一个创新点：Prefix Caching（本文同时也是BlockManager的下篇，虽然标题没有提及）
说一些题外话，之前写vllm源码解读的文章，阅读量不是很高，再加上写这类型文章真得耗时耗力耗头发（自己看懂代码容易，但是给别人讲懂很难，把代码转变成一篇有逻辑的、兼顾全局和细节的文章就更难了。特别是mlsys的代码，懂的都懂），因此我一度丧失了对它的写作热情😢。但是这段日子打开尘封已久的私信，竟然看到有很多朋友在催更，所以动力又回来了些。不过这个系列后续的更新节奏依然还是比较慢（我的大部分文章都是在午休时间见缝插针写的），还请大家见谅哈。
【全文目录如下】
一、两种不同的BlockAllocator
二、物理块和逻辑块的结构
三、prefill阶段的物理块分配方法
3.1 allocate函数入口
3.2 计算物理块hash值的方法
3.3 使用LRUEvictor管理物理块分配细节
3.4 再探LRUEvictor，理解“prefix”
四、decode阶段的物理块分配方法
4.1 整体流程
4.2 append_slots入口函数
4.3 如何添加一个新物理块
4.4 物理块的slots满时要如何处理
【阅读本文前，建议先阅读以下文章】：
1.vllm原理篇
2.vllm源码解读1：整体架构
3.vllm源码解读2：调度器策略
4.vllm源码解读3：块管理器上篇，UncachedBlockAllocator
【如有帮助，欢迎点赞收藏在看～】
一、两种不同类型的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的思想来分配和管理物理块，是本篇讲解的重点。在原理篇中，我们提过：-
在prefill阶段，prompts中可能含有类似system message（例如，“假设你是一个能提供帮助的行车导航”）等prefix信息，带有这些相同prefix信息的prompt完全可以共享物理块，实现节省显存、减少重复计算的目的。
-
在decode阶段，我们依然可以用这种prefix的思想，及时发现可以重复利用的物理块。
-
prefill和decode阶段做prefix caching的方法有些不同，我们会在后文仔细讲解。
-
-
UncachedBlockAllocator：正常分配和管理物理块，没有额外实现prefix caching的功能。这是我们源码解读3讲解的重点，本文不再赘述。
二、物理块和逻辑块结构
首先我们来快速回顾下在vllm中一个物理块和一个逻辑块长什么样子。
2.1 物理块结构
# 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})')
2.2 逻辑块结构
一切尽在注释中：
# # 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:
"""当前逻辑块的空余槽位"""
return self.block_size - self.num_tokens
def is_full(self) -> bool:
"""判断当前逻辑块是否已经被装满"""
return self.num_tokens == self.block_size
def append_tokens(self, token_ids: List[int]) -> None:
"""将给定的一些token_ids装入当前逻辑块中"""
# 给定的token_ids的长度必须 <= 当前逻辑块剩余的槽位
assert len(token_ids) <= self.get_num_empty_slots()
# 当前逻辑块第一个空槽的序号
curr_idx = self.num_tokens
# 将这些tokens装进去
self.token_ids[curr_idx:curr_idx + len(token_ids)] = token_ids
# 更新当前逻辑块中tokens的数量
self.num_tokens += len(token_ids)
def get_token_ids(self) -> List[int]:
"""获取当前逻辑块中所有被装满的位置的token_ids"""
return self.token_ids[:self.num_tokens]
def get_last_token_id(self) -> int:
"""获取当前逻辑块所所有被装满的位置的最后一个token_id"""
assert self.num_tokens > 0
return self.token_ids[self.num_tokens - 1]
关于逻辑块，我们已在源码解读2的2.3（2）中详细介绍过，它是Sequence实例（seq）下维护的一个属性。我们也提过，在vLLM代码实现中：
-
每个seq维护自己的一份逻辑块列表，
-
BlockManager中的self.block_tables（形式如：{seq_id: List[PhysicalBlock]}）则记录者每个seq下的物理块列表
通过seq这个中介，我们维护起“逻辑块->物理块”的映射。
三、prefill阶段的物理块分配方法
在本节中，我们详细解读“如何使用CachedBlockAllocator为waiting队列中的seq_group分配做prefill需要的物理块”
3.1 allocate函数入口
如上图，当我们准备从waiting队列中调度seq_group时，我们会依次做两件事：
-
调用
self.block_manager.can_allocate(seq_group)方法，判断当前gpu上是否有充足的空间，能为当下这seq_group的prefill阶段分配充足的物理块，用于装其KV Cache（细节我们在源码解读2中已讲过，这里不再赘述） -
一旦我们认为当下空间充足，则调用
self._allocate(seq_group)方法，为waiting队列中的这个seq_group实际分配物理块，这时我们就会运用到BlockAllocator，并且BlockAllocator的类型不同（即是否做prefix caching），allocate的方法也会不同。
所以现在，我们就来看self._allocate(seq_group)函数（如何为waiting队列中的seq_group分配物理块做prefill）
self._allocate(seq_group)的入口函数如下(一切尽在注释中)：
# vllm/core/scheduler.py
def _allocate(self, seq_group: SequenceGroup) -> None:
# ==============================================================
# block_manager为当前seq_group分配物理块
# ==============================================================
self.block_manager.allocate(seq_group)
# ==============================================================
# 当前seq_group状态改为running
# ==============================================================
for seq in seq_group.get_seqs(status=SequenceStatus.WAITING):
seq.status = SequenceStatus.RUNNING
接下来我们看self.block_manager.allocate(seq_group)实现，如前文所说，本文我们解读的是BlockSpaceManagerV1，所以我们就去这个class的顶一下看allocate方法（一切尽在注释中）。
# vllm/core/block_manager_v1.py
class BlockSpaceManagerV1(BlockSpaceManager):
"""Manages the mapping between logical and physical token blocks."""
def __init__(
self,
block_size: int, # 每个block的槽位大小，默认为16
num_gpu_blocks: int, # 当前gpu上最多能分配的block数量
num_cpu_blocks: int, # 当前cpu上，用于做swap的内存中，最多能分配的block数量
watermark: float = 0.01, # 内存交换的水位线（阈值）
sliding_window: Optional[int] = None, # 滑动窗口的大小
enable_caching: bool = False, # 是否需要做prefix caching
) -> None:
self.block_size = block_size
self.num_total_gpu_blocks = num_gpu_blocks
self.num_total_cpu_blocks = num_cpu_blocks
if enable_caching and sliding_window is not None:
raise NotImplementedError(
"Sliding window is not allowed with prefix caching enabled!")
self.block_sliding_window = None
if sliding_window is not None:
assert sliding_window % block_size == 0, (sliding_window,
block_size)
self.block_sliding_window = sliding_window // block_size
self.watermark = watermark
assert watermark >= 0.0
self.enable_caching = enable_caching
# ==
最低0.47元/天 解锁文章
扫一扫
309
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
