KVCacheManager:請求級的 KV 快取生命週期

源码版本v0.25.1

職責

KVCacheManager 是排程器 (scheduler) 看到的 KV 快取 (KV cache) 介面。它把"給這條請求分配幾個 block"、"這條請求完成了把 block 還回去"、"看一下這條 prompt 命中了多少前綴快取"這些動作封裝成 allocate_slots / free / get_computed_blocks 三個高層 API,內部把不同 KV cache 類型(full attention / sliding window / Mamba 狀態等)的協調全部委託給 KVCacheCoordinator,把物理 block 池子管理交給 BlockPool。排程器只需要拿 KVCacheBlocks 這個不可變結果(kv_cache_manager.py:29-50),不需要知道下面是單類型還是混合模型。

構造時 KVCacheManager.__init__ 通過 get_kv_cache_coordinator 把 coordinator 裝起來,順手把 block_pool / num_kv_cache_groups / kv_cache_config 掛到自己身上(kv_cache_manager.py:147-183)。watermark(watermark_blocks = int(watermark * num_blocks))是預留給已排程請求的餘量,只對 WAITING / PREEMPTED 狀態生效,避免新請求一進來就把跑了一半的請求擠掉(kv_cache_manager.py:164-167)(kv_cache_manager.py:367-374)。整個 manager 幾乎不持有 per-request 狀態,所有 req_to_blocks 都在 SingleTypeKVCacheManager 裡。

設計動機

為什麼 manager 看起來這麼薄,真正邏輯全在 coordinator?

• 職責分層:KVCacheManager 負責"面向請求"的語義(allocate_slots / free / get_computed_blocks),KVCacheCoordinator 負責"面向 KV cache group"的協調(混合模型時多 group 怎麼一起分配),BlockPool 負責"面向物理 block"的池化。三層各管各的,manager 做膠水。
• KVCacheBlocks 是穩定介面:blocks[i][j] 表示第 i 個 KV cache group 的第 j 個 block,get_block_ids() 轉成 tuple[list[int], ...] 給 attention backend 用(kv_cache_manager.py:73-88)。__add__ 讓兩次 allocate_slots 的結果能拼起來(<SrcLink path="vllm/v1/core/kv_cache_manager.py" lines="52-59" label="kv_cache_manager.py"/>),new_empty 用來返回"無 block"的佔位。
• 延遲快取提交:allocate_slots 的 delay_cache_blocks=True 分支在 P/D 非同步傳輸時跳過 cache_blocks(kv_cache_manager.py:448-451),等遠端 KV 真傳到再標 hash,避免提前快取導致命中錯誤。
• watermark 只對等待中的請求生效:has_scheduled_reqs and request.status in (WAITING, PREEMPTED) 才加 watermark,已經在跑的請求增量分配不需要預留(kv_cache_manager.py:367-374)。
• 全序列准入檢查:full_sequence_must_fit=True 時先估算整條 prompt 需要多少 block,超出就直接返回 None,防止 chunked prefill 把請求放進來後又卡死(kv_cache_manager.py:376-391)。
• reserved_blocks 給非同步 KV connector 留位:已 in-flight 的 prefill 序列依賴某些 block,新進來的非同步 connector load 不能搶這些,required_blocks > available_blocks 時拒絕分配(kv_cache_manager.py:420-426)。

關鍵檔案

• KVCacheBlocks dataclass:29-50 — 排程器側的不可變結果,blocks[i][j] = 第 i group 第 j block。
• KVCacheBlocks.__add__:52-59 — 兩次分配結果拼接,常用於增量分配。
• KVCacheBlocks.get_block_ids:73-88 — 轉成 tuple[list[int], ...] 給 attention backend。
• KVCacheManager.__init__:114-183 — 裝 coordinator、算 watermark_blocks、構 empty_kv_cache_blocks 快取避免 GC 開銷。
• get_computed_blocks:206-246 — 調 coordinator.find_longest_cache_hit 找前綴命中,返回 (KVCacheBlocks, num_new_computed_tokens)。
• allocate_slots:248-343 — 三段式:釋放跳過的 block → 處理 prefix token → 分配新 token 的 block,失敗返回 None。
• full_sequence_must_fit:376-391 — 准入檢查:整條 prompt 的 block 數 + watermark 必須 ≤ 空閒 block 數。
• required_blocks gate:410-426 — available = free - reserved,required > available 拒絕分配。
• allocate_new + cache_blocks:428-464 — 委託 coordinator 分配新塊 + 調 cache_blocks 提交 hash。
• free:466-474 — 調 coordinator.free(request_id),逆序釋放讓 tail block 先被淘汰。
• pop_blocks_for_free:495-506 — 拿走 per-request 記賬但不歸還 BlockPool,P/D 轉移用。
• reset_prefix_cache:516-530 — RLHF 改權重後失效整池快取。

資料流

排程器每一步對每條請求先調 get_computed_blocks 看前綴命中,再調 allocate_slots 把新 token 的 block 分出來。下面是 allocate_slots 在做完 watermark + 准入檢查後真正幹活的幾行:

if (
    new_computed_block_list is not self.empty_kv_cache_blocks.blocks
    or num_external_computed_tokens > 0
):
    # Append the new computed blocks to the request blocks until now to
    # avoid the case where the new blocks cannot be allocated.
    self.coordinator.allocate_new_computed_blocks(
        request_id=request.request_id,
        new_computed_blocks=new_computed_block_list,
        num_local_computed_tokens=num_local_computed_tokens,
        num_external_computed_tokens=num_external_computed_tokens,
    )

new_blocks = self.coordinator.allocate_new_blocks(
    request.request_id,
    num_tokens_need_slot,
    num_tokens_main_model,
    num_encoder_tokens,
)

# P/D: delay caching blocks if we have to recv from
# remote. Update state for locally cached blocks.
if not self.enable_caching or delay_cache_blocks:
    return self.create_kv_cache_blocks(new_blocks)

(kv_cache_manager.py:428-451)先吸收前綴命中的 block(增加 ref_cnt),再讓 coordinator 給新 token 分配物理 block,最後除非要延遲快取否則調 cache_blocks 把剛填滿的 block 標 hash 進 prefix cache(詳見 Coordinator 與 BlockPool)。返回的 KVCacheBlocks 經 get_block_ids 轉成 block_table,塞進 塊表 給 attention kernel 用。

邊界與失敗

• num_new_tokens == 0 且無外部 token:raise ValueError,呼叫方必須保證要麼有新 token 要算,要麼有外部已計算 token 要存(kv_cache_manager.py:346-350)。
• get_computed_blocks 全命中時少一 token:max_cache_hit_length = request.num_tokens - 1,因為全命中時還要重算最後一個 token 拿 logits,所以快取命中長度卡在 prompt_length - 1(kv_cache_manager.py:225-231)。
• 跳過 prefix cache 的請求:request.skip_reading_prefix_cache=True 或 enable_caching=False 時 get_computed_blocks 直接返回空(kv_cache_manager.py:222-223),比如需要 prompt logprobs 的請求。
• reset_prefix_cache 拒絕:num_used_blocks != 1(只算 null block)時說明還有 block 沒釋放,reset 失敗返回 False(kv_cache_manager.py:516-530)。
• pop_blocks_for_free 不歸還池子:呼叫方必須自己逆序調 block_pool.free_blocks,否則塊會洩漏(kv_cache_manager.py:495-506)。
• take_events 註解 group 元信息:BlockStored 事件只有 group_idx,manager 用 kv_cache_event_metadata 補 kind 和 sliding_window(kv_cache_manager.py:566-589)。

小結

KVCacheManager 是排程器看到的 KV 快取門面,真正幹活的 KVCacheCoordinator + BlockPool 在 Coordinator 與 BlockPool 拆解。分配出去的 block id 最終落到 塊表 上,被 attention kernel 用來定位物理 KV cache;模型權重本身在 DefaultModelLoader 那條線載入,跟本層無耦合。

對照官方資料:vLLM 文件 · README