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。