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 能直接消費的張量。