vLLM V1 KV Cache Manager 源码学习

本文为vLLM V1源码学习系列的第二篇文章，往篇如下，

vLLM V1 Scheduler 源码学习

引言

在了解了Schedule的基本工作原理后，我们也需要知道Scheduler是如何获取KV cache所需的GPU memory的。

具体而言vLLM结合BlockPool，KVCacheManager以及基于块内容的哈希来实现底层支持，涉及如下三个文件：

Block Pool

我们知道vLLM项目的起始论文为 Paged Attention (Efficient Memory Management for Large Language Model Serving with PagedAttention)，其核心即为把一个请求拆分为固定大小的块（例如，16个token），每个块对应的KV cache（称为KV cache block，后称Block）单独管理，进而类似操作系统中的页管理实现对GPU memory的高效利用。而BlockPool就是核心管理组件，决定分配、释放和缓存Block。

除了表示Block数量和是否使用prefix caching的self.num_gpu_blocks和self.enable_caching。其核心成员变量为self.blocks，self.free_block_queue和self.cached_block_hash_to_block。

self.blocks

在实例化BlockPool时，vLLM就会创建一个self.num_gpu_blocks大小的列表。

python 复制代码

        # All kv-cache blocks.
        self.blocks: list[KVCacheBlock] = [
            KVCacheBlock(idx) for idx in range(num_gpu_blocks)
        ]

其中每个元素为KVCacheBlock（实现在kv_cache_utils.py中），包含一个Block的编号，哈希值和引用计数（用于释放Block）。

python 复制代码

class KVCacheBlock:
    """KV-cache block metadata."""
    # Block ID, ranging from 0 to num_gpu_blocks - 1.
    block_id: int
    # Reference count.
    ref_cnt: int = 0
    # The hash of the block composed of (block hash, tuple of token IDs).
    # It is only available when the block is full.
    _block_hash: Optional[BlockHashType] = None

    # Used to construct a doubly linked list for free blocks.
    # These two attributes should only be manipulated by FreeKVCacheBlockQueue.
    prev_free_block: Optional["KVCacheBlock"] = None
    next_free_block: Optional["KVCacheBlock"] = None

self.free_block_queue

是FreeKVCacheBlockQueue（一个双向链表）的一个实例，链表的头部为最优先释放的Block。在KVCacheBlock中对应的变量，即为prev_free_block和next_free_block。

self.cached_block_hash_to_block

记录缓存了的完整的块的字典，结构为{block_hash: {block ID: block}}。只有一个块满了（达到了Block Size个token）才会被缓存。缓存的Block可能是正在执行的请求所需要的，也可能是已完成请求的（可以被释放）。注意vLLM不会检查是否存在相同的Block，因为vLLM不希望分配的Block编号发生变化（Block表是仅追加的，不会修改原本的元素），例如两个相同的请求被同时执行，这两个请求的第一个块，被分配了Block 0和Block 1，此时内存中就出现了两个相同的Block但是ID不同。

Important methods

cache_full_blocks：根据一个请求涉及的完整Block缓存Block。
get_new_blocks：从空闲块池中获取指定数量的新块，
- 如果请求的块数量超过了空闲块的数量，则引发 ValueError；
- 需要确保弹出的每个块的引用计数为 0（表示它是空闲的）；
- 如果启用了缓存，则调用 self._maybe_evict_cached_block 来尝试驱逐可能缓存的当前块；
- 增加每个分配的块的引用计数 (incr_ref)。
_maybe_evict_cached_block：如果一个块在 cached_block_hash_to_block 中被缓存，则重置其哈希元数据并将其从缓存中移除。
touch： "触摸"一个或多个块，
- 增加它们的引用计数，并可能将它们从空闲队列中移除；
- 当具有相同前缀的另一个请求命中一个块时，会使用此方法。如果块的引用计数为 0 并且不是 self.null_block，则将其从 self.free_block_queue 中移除（因为它现在被使用了，不再是驱逐候选者）；
- 更新块的引用计数。
free_blocks：尝试释放一个块列表。具体操作是，
- 遍历输入的块列表减少引用计数；
- 如果块的引用计数变为 0 并且不是 self.null_block，则将其添加回 self.free_block_queue，使其成为空闲块并可以被再次分配或驱逐。
reset_prefix_cache：重置前缀缓存。
- 检查当前已使用的块数量（不包括 null_block）。如果除了 null_block 之外还有其他块被使用，则发出警告并返回 False，因为在重置缓存之前应该释放所有相关的块。
- 清空 self.cached_block_hash_to_block，移除所有已缓存的块哈希，这样新的请求就不会命中任何缓存。
- 遍历所有块，并调用 block.reset_hash() 清除每个块的哈希相关属性。
- 记录成功重置缓存的信息并返回 True。

KV Cache Manager

self.num_preallocate_tokens：每个请求预先分配的 token 数量并计算 num_preallocate_blocks（预先分配的块数量）。

self.block_pool：创建BlockPool 的实例来管理 GPU 块。

self.specialized_manager: 使用 get_specialized_manager 基于 KV 缓存规范和块池获取一个专门的管理器。

self.req_to_blocks：一个 defaultdict(list)，用于将每个请求 ID 映射到为该请求分配的 KVCacheBlock 对象列表。这用于跟踪哪些块属于哪个请求，以便稍后释放它们。

self.req_to_block_hashes：一个 defaultdict(list)，用于将每个请求 ID 映射到预先计算的块哈希值列表。这用于高效地查找前缀缓存。

self.num_cached_block：一个字典，用于跟踪每个正在运行的请求的完全缓存的块数量。

self.prefix_cache_stats：一个 PrefixCacheStats 的实例，用于收集关于前缀缓存的统计信息，是否统计取决于 log_stats 是否启用。

Important methods

get_computed_blocks ：这是前缀缓存机制的核心。它尝试查找与当前请求的提示开头匹配的先前计算并缓存的 KV 缓存块。
- 如果 enable_caching 为 False，则立即返回一个空列表和 0 个已计算的 token。
- 从 self.req_to_block_hashes 中检索请求的预计算 block_hashes。如果尚未计算，则使用 hash_request_tokens 计算它们。
- 查找最长缓存命中： 调用 self.specialized_manager.find_longest_cache_hit 并传入 block_hashes，以找到缓存中存在的最长连续块序列。
- 处理边界情况： 如果提示长度可以被块大小整除并且所有块最初都被缓存，它会临时删除最后一个块哈希，以强制重新计算最后一个块。这是因为 allocate_slots 期望 num_computed_tokens 是 block_size 的倍数。这是一个已知的限制。
- 更新 prefix_cache_stats.queries 和 prefix_cache_stats.hits。
- 计算 num_computed_tokens（始终是 block_size 的倍数，因为只有完整块才被视为可缓存的）。
- 返回 computed_blocks 列表和 num_computed_tokens。
allocate_slots：为给定的请求分配新的 KV 缓存块，以容纳新的 token。
- 接受 request、要分配的 num_tokens 数量以及可选的 new_computed_blocks 列表（来自前缀缓存命中）。
- 块布局可视化： 提供了一个有用的图示，说明了涉及的不同类型的块（已计算、新计算、新分配、预分配）。
- 错误处理： 如果 num_tokens 为 0，则引发 ValueError。
- 如果 new_computed_blocks 为 None，则初始化它。
- 从 self.req_to_blocks 中检索请求的现有块。
- 释放跳过的块： 调用 self.specialized_manager.remove_skipped_blocks 来释放不再需要的块（例如，由于滑动窗口注意力）。在分配新块之前执行此操作可以减少块的驱逐。
- 通过添加来自 new_computed_blocks 的 token 来更新 num_computed_tokens。
- 基于 token 总数（已计算 + 新分配）和 block_size 计算 num_required_blocks（所需的块数）。
- 计算所需的 num_new_blocks（新块数）。
- 检查是否有足够的空闲块： 它检查 block_pool 中是否有足够的空闲块（考虑到引用计数为 0 的已计算块可能是驱逐的候选者）。如果空闲块不足，则返回 None，表示分配失败。
- 标记已计算的块： 如果启用了缓存，则调用 self.block_pool.touch 来标记 new_computed_blocks 为最近使用，以防止它们被驱逐。
- 将 new_computed_blocks 添加到 req_blocks 中。
- 分配新块： 如果 num_new_blocks 为正数，则从 block_pool 获取新块（考虑预分配和每个请求的最大块数）。
- 将新分配的 new_blocks 添加到 req_blocks 中。
- 缓存完整块： 如果启用了缓存，则调用 self.block_pool.cache_full_blocks 将完全使用的块标记为可缓存，并将它们与请求及其哈希值关联起来。它只缓存生成的（已接受的）token 的块，而不缓存推测的 token 的块。
- 更新请求的 self.num_cached_block。
- 返回分配的 new_blocks 列表。
free ：
- 释放与已完成或中止的 request 关联的 KV 缓存块。
- 从 self.req_to_blocks 中检索请求的块列表并删除该条目。
- 如果启用了缓存，则按相反的顺序释放块。这种策略优先释放最近添加的块，因为这些块不太可能是公共前缀的一部分。
- 调用 self.block_pool.free_blocks 将块释放回块池。
- 从 self.num_cached_block 中删除请求的条目。

reset_prefix_cache 方法：
- 重置由 BlockPool 管理的前缀缓存。
- 如果 BlockPool 成功重置，则将 self.prefix_cache_stats.reset 设置为 True 并返回 True。否则，返回 False。
get_num_common_prefix_blocks 方法：
- 计算当前所有处于 RUNNING 状态的请求共享的公共前缀块的数量。
- 它遍历给定 request（必须处于 RUNNING 状态）的块。
- 如果一个块在 BlockPool 中的引用计数 (ref_cnt) 等于 num_running_requests（当前正在运行的请求总数），则该块被认为是公共前缀块。这表明所有正在运行的请求当前都引用了这个块。
- 一旦遇到不是所有正在运行的请求共享的块，它就会停止计数。
- 重要提示： 该函数承认一个边界情况，即即使所有 已调度 的请求共享一个前缀，报告的公共前缀块数量也可能为 0。这可能是因为存在其他未调度的 RUNNING 请求不共享相同的前缀。目前，这种情况不容易检测到。
free_block_hashes 方法：
- 丢弃与已完成的 request 关联的预计算块哈希值。
- 当一个请求完全处理完毕并且不再需要其缓存的前缀信息时，会调用此方法。重要的是只对已完成的请求调用此方法，而不是对被抢占的请求调用，因为被抢占的请求可能会稍后恢复，并且它们的块哈希值可能仍然有用。