KVCacheCoordinator + BlockPool:混合模型的协调与块池
职责
KVCacheCoordinator 是 vLLM v1 在多 KV cache 类型(全注意力 / 滑窗 / Mamba 状态 / MLA 等)之上的一层协调器。每个 KV cache group 配一个 SingleTypeKVCacheManager(如 FullAttentionManager / SlidingWindowManager / RSWAManager),coordinator 把"分配 N 个 token 的 block"、"找最长前缀命中"、"释放请求"、"cache 已填满的 block"这些动作按 group 拆开调用对应的 single-type manager。BlockPool 是共享的物理块池子,所有 group 共用同一组 KVCacheBlock——只是 hash key 上附带 group_id 区分,所以同一个块可以同时被多个 group 缓存(例如全注意力和滑窗的同一 token)。
构造时 KVCacheCoordinator.__init__ 直接 new 一个 BlockPool,然后 get_manager_for_kv_cache_spec 给每个 group 起一个 single-type manager(kv_cache_coordinator.py:91-120)。具体策略由 get_kv_cache_coordinator 工厂选择:不开缓存走 KVCacheCoordinatorNoPrefixCache,单 group 走 UnitaryKVCacheCoordinator,多 group 走 HybridKVCacheCoordinator(kv_cache_coordinator.py:782-822)。BlockPool 自己维护一个 FreeKVCacheBlockQueue(双向链表)做 LRU 淘汰,一个 cached_block_hash_to_block 哈希表做前缀缓存索引(block_pool.py:144-197)。
设计动机
为什么要把 coordinator 分成 NoPrefix / Unitary / Hybrid 三种?
- 空跑兜底:
KVCacheCoordinatorNoPrefixCache支持任意 group 数(包括 0),不实现任何前缀缓存相关接口——find_longest_cache_hit直接返回(空 blocks, 0),get_num_common_prefix_blocks返回全 0(kv_cache_coordinator.py:413-424)。这样禁用 prefix caching 时不会走任何缓存路径。 - 单 group 快路径:
UnitaryKVCacheCoordinator假设hash_block_size == block_size且只有一个 group(kv_cache_coordinator.py:471-476),find_longest_cache_hit直接调 single manager 不用合并(kv_cache_coordinator.py:480-496),避免 hybrid 那套 spec-group 合并开销。 - Hybrid 按 spec 聚合:
HybridKVCacheCoordinator.verify_and_split_kv_cache_groups把相同 spec 的多个 group 合成一个SpecGroup,一起做 cache hit 查找(kv_cache_coordinator.py:560-588),并且把 full attention 排在最前面给后续 group 一个紧的上界(kv_cache_coordinator.py:590-595)。 - 共享 BlockPool:所有 group 共用
num_gpu_blocks个物理块,不同 group 的缓存 hash 用make_block_hash_with_group_id(block_hash, group_id)拼接,get_cached_block按 group_id 列表分别查(block_pool.py:199-224),BlockHashToBlockMap一个 hash 可对应多 group 的块。 - LRU + 尾块优先:
FreeKVCacheBlockQueue用双向链表实现 O(1) 中间删除,free 时按反向顺序塞回让 tail block 先淘汰(kv_cache_utils.py:179-197);free_blocks还把无 hash 的块先于有 hash 的块淘汰(block_pool.py:622-635)。 - null_block 占位:block_id=0 是
null_block,被remove_skipped_blocks用来替换滑窗外或被抢占的块,ref_cnt 不维护,释放时要小心跳过(block_pool.py:188-192)。
关键文件
KVCacheCoordinator.__init__:61-128— 抽象基类,起 BlockPool、起 single-type managers、读VLLM_PREFIX_CACHE_RETENTION_INTERVAL。get_num_blocks_to_allocate:130-185— 遍历 single-type managers 求和,cross-attention 走单独路径。allocate_new_blocks:233-266— 按 group 调manager.allocate_new_blocks,encoder token 单独处理。cache_blocks:268-283— 调每个 manager 的cache_blocks,带retention_interval。remove_skipped_blocks:331-352— 滑窗外的块替换成 null_block,R-SWA 用num_prompt_tokens决定 gap。KVCacheCoordinatorNoPrefixCache:377-424— 禁用前缀缓存时用,所有 hit 返回空。UnitaryKVCacheCoordinator:427-496— 单 group 快路径。SpecGroup + HybridKVCacheCoordinator:499-588—SpecGroupNamedTuple + 按 spec 聚合 group。find_longest_cache_hit_per_group:742-779— 每 group 独立查 hit,返回(blocks_per_group, hit_lengths_per_group)。get_kv_cache_coordinator:782-822— 工厂:NoPrefix / Unitary / Hybrid 三选一。BlockPool.__init__:144-197— 起FreeKVCacheBlockQueue、cached_block_hash_to_block、null_block。BlockPool.get_cached_block:199-224— 按 group_id 列表查同一 hash 的块,任一 miss 返回 None。BlockPool.cache_full_blocks:226-263— 填满的块打 hash 写入cached_block_hash_to_block,支持block_mask跳过永不会命中的块。BlockPool.get_new_blocks:542-572— 从 free queue 头部取 N 个,顺手 evict 已缓存的旧 hash。BlockPool.touch:597-612— ref_cnt +1,ref_cnt 从 0 升到 1 时从 free queue 摘出。BlockPool.free_blocks:614-635— 无 hash 块 prepend_last 优先淘汰,有 hash 块 append(LRU 尾)。BlockPool.reset_prefix_cache:656-690— 只在只有 null block 占用时才允许清空。KVCacheBlock:118-176— 块元数据:block_id/ref_cnt/_block_hash/ 链表指针 /is_null。FreeKVCacheBlockQueue:179-197— 双向链表,O(1) 中间删除,LRU 顺序。
数据流
BlockPool.get_new_blocks 是物理分配的真正入口——从 free queue 头部弹出 N 个块,逐个 evict 旧 hash、把 ref_cnt 设为 1:
def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]:
if num_blocks > self.get_num_free_blocks():
raise ValueError(f"Cannot get {num_blocks} free blocks from the pool")
ret: list[KVCacheBlock] = self.free_block_queue.popleft_n(num_blocks)
if self.enable_caching:
for block in ret:
self._maybe_evict_cached_block(block)
assert block.ref_cnt == 0
block.ref_cnt += 1
if self.metrics_collector:
self.metrics_collector.on_block_allocated(block)
else:
for block in ret:
assert block.ref_cnt == 0
block.ref_cnt += 1
if self.metrics_collector:
self.metrics_collector.on_block_allocated(block)
return ret(block_pool.py:542-572)被 SingleTypeKVCacheManager.allocate_new_blocks 调(single_type_kv_cache_manager.py:279-311),后者把新块 append 到 req_to_blocks[request_id]。释放走 BlockPool.free_blocks,先 ref_cnt -1,ref_cnt 归零的块按"有 hash / 无 hash"分流塞回 free queue 尾部(block_pool.py:614-635)。当一条请求的 prompt 填满一个 block 时,cache_full_blocks 把 block 的内容算 hash 并写入 cached_block_hash_to_block,下一条相同前缀的请求过来 get_cached_block 直接命中、touch 把 ref_cnt +1 而不必重新算(block_pool.py:199-224)(block_pool.py:597-612)。
边界与失败
- 超过空闲块数:
get_new_blocks不预检查,直接raise ValueError(block_pool.py:553-554),调用方(KVCacheManager.allocate_slots)必须先用get_num_blocks_to_allocate算好(kv_cache_coordinator.py:130-185)。 - Hybrid 不支持 DCP / PCP:
assert dcp_world_size == 1 and pcp_world_size == 1(kv_cache_coordinator.py:556-557)。 - Unitary 强制 hash_block_size == block_size:
assert not enable_caching or hash_block_size == self.block_size(kv_cache_coordinator.py:471-473)。 - Hybrid 至少两个 attention group:
assert len(self.attention_groups) > 1(kv_cache_coordinator.py:586-588)。 - block 已有 hash 还要 set:
set_block_hash用 assert 防重复(kv_cache_utils.py:153-157)。 reset_prefix_cache拒绝清理:除了 null block 还有别的块占用就logger.warning+ 返回False(block_pool.py:665-672)。evict_blocks接收外部 block_id:KV connector 报上来的 block id 越界立刻 assert(block_pool.py:647-652)。- retention_interval 校验:
_validate_prefix_cache_retention_interval要求正整数是scheduler_block_size的倍数(kv_cache_coordinator.py:31-58)。
小结
Coordinator + BlockPool 构成了 KVCacheManager 之下的物理层:coordinator 按策略分发到 single-type managers,BlockPool 管物理块池 + 前缀缓存索引。块分配的结果最终变成 block_id 列表,被 块表 包装成 attention kernel 能直接消费的张量。