vLLM 源码完全拆解(四):KVCacheManager 如何分配、复用与回收物理 Block

10371 字
52 分钟
vLLM 源码完全拆解(四):KVCacheManager 如何分配、复用与回收物理 Block

从 Scheduler 的内存关口继续向下#

前三篇已经依次建立了 vLLM 的整体模块地图、追踪了 一次请求如何穿过 V1 引擎,并拆开了 Scheduler.schedule() 的预算分配过程。上一站停在一条看似普通、实际决定整轮计划能否成立的调用上:

new_blocks = self.kv_cache_manager.allocate_slots(
request,
num_new_tokens,
num_lookahead_tokens=self.num_lookahead_tokens,
)

Scheduler 已经决定请求本轮应该算多少 token,但这些 token 的 K/V 是否有位置可写、命中的前缀块会不会在 GPU 执行前被别人覆盖、新分配的物理块怎样同步到 Worker,都由这条调用负责。allocate_slots() 返回 None,计划就必须收缩、抢占或延后;返回一组 block,Scheduler 才能把它写进 SchedulerOutput。

vLLM 由 UC Berkeley 团队在 2023 年的 SOSP 论文中提出,PagedAttention 把操作系统分页思想引入大模型推理的 KV Cache。V1 在此基础上把块管理继续拆成可复用的控制面组件。本文沿用上一篇锁定的 vLLM main 分支提交 3b4566c5(2026 年 9 月 26 日访问),而不是混用持续变化的最新行号。本文只讨论推理时的缓存分配、复用与回收,不涉及模型训练。

先给出本篇要追踪的闭环:

Scheduler
├─ get_computed_blocks() 查找可复用前缀
├─ allocate_slots() 先验证容量,再提交所有权变化
│ ├─ Coordinator 汇总各 KV cache group 的需求
│ ├─ SingleTypeManager 处理一种注意力的块语义
│ └─ BlockPool 分配、touch、缓存、淘汰与回收
└─ SchedulerOutput 只发送新增 block ID 差量
↓
ModelRunner
├─ 更新 CPU 侧 block table
├─ 拷贝到 GPU
└─ 计算 slot mapping,供 KV 写入与注意力读取

读完后,应该能够解释三个容易混淆的问题:为什么一个 block 可以同时是“free”和“cached”;为什么容量检查必须发生在 touch 之前;为什么 Manager 返回的是元数据 block,而真正的 K/V 张量一直由 Worker 持有。

先画清职责边界:Manager、Coordinator 与 BlockPool#

KVCacheManager 这个名字容易让人误以为它直接管理一整块 GPU 张量。实际 V1 把控制面和数据面分开:Engine Core 侧管理 block ID、哈希、引用计数和请求到 block 的映射;Worker 侧持有真正的 KV Cache 张量,并按 block ID 寻址。Engine Core 不需要搬运每层 K/V,Worker 也不自行决定一个请求该拿哪些物理块。

vLLM V1 前缀缓存的数据结构
vLLM V1 前缀缓存的数据结构

图源:vLLM Automatic Prefix Caching 设计文档。上半部分是预创建的 block pool 与侵入式双向链表;下半部分分别是“块哈希到物理块”和“请求 ID 到物理块序列”的索引。图中没有画出 K/V 张量,因为这里展示的是控制面元数据。

这套实现可分为四层。

第一层:KVCacheManager 是 Scheduler 的窄接口#

Scheduler 主要调用以下方法:

方法回答的问题是否改变所有权
get_computed_blocks()prompt 的最长可复用前缀在哪里否,只查询
allocate_slots()本轮全部位置能否落入物理块是,成功时原子式提交
free()请求结束后怎样释放引用是
get_kv_cache_events()哪些块被存入或移出缓存否,导出事件
reset_prefix_cache()能否清空哈希索引是,但要求没有活跃使用者

它返回 KVCacheBlocks,其外层维度是 KV cache group,内层才是该 group 的 block 序列。即使普通全注意力模型只有一个 group,接口仍保留这层嵌套,因为混合注意力模型可能同时存在全注意力、滑动窗口注意力和 Mamba 状态。

第二层:KVCacheCoordinator 负责跨 group 协调#

如果模型所有层采用相同缓存语义,UnitaryKVCacheCoordinator 只协调一个 group。若不同层的缓存生命周期不同,HybridKVCacheCoordinator 会把请求的需求拆给多个 SingleTypeKVCacheManager,再把结果拼回统一结构。

Coordinator 的价值不只是一个循环。一次前缀命中必须对所有必要 group 都成立。全注意力 group 命中 2048 token,而滑动窗口 group 只能恢复到 1536 token,请求不能谎称自己已计算到 2048;否则 num_computed_tokens 与某些层实际拥有的状态不一致。混合 Coordinator 因而要寻找各 group 都可恢复的共同边界。

第三层:SingleTypeKVCacheManager 实现具体缓存语义#

每个 manager 维护 req_to_blocks[request_id],也就是该请求在本 group 中的逻辑 block table。不同子类对“需要保留多少块”的答案不同:

  • 全注意力需要从序列开头保留到当前位置的全部 K/V;
  • 滑动窗口只需保留窗口覆盖的后缀,已经永远不会再被读到的前部块可以提前回收;
  • Mamba 保存的是递归状态,block 的含义与普通逐 token K/V 不完全相同;
  • encoder-decoder 的 cross-attention 缓存通常按 encoder token 一次性静态分配。

上层只问“为这些 token 还要多少物理块”,具体子类负责把 token 数翻译成符合自身语义的 block 数。

第四层:BlockPool 管理有限物理页#

BlockPool 启动时一次性创建 num_gpu_blocks 个 KVCacheBlock 元数据对象。每个对象至少包含:

@dataclass(slots=True)
class KVCacheBlock:
block_id: int
ref_cnt: int = 0
_block_hash: BlockHashWithGroupId | None = None
prev_free_block: KVCacheBlock | None = None
next_free_block: KVCacheBlock | None = None

block_id 对应 Worker 数据面中的物理位置;ref_cnt 表示当前有多少请求在使用;block_hash 让已计算内容可以被后续请求查到;两个链表指针把空闲块直接串成 free queue。预创建对象避免在每次 decode 的热路径中频繁创建 Python 对象,也让 block_id == blocks[block_id].block_id 成为稳定不变量。

池里还永久取出一个 null_block。滑动窗口等 manager 可以用它填补已经跳过的位置,让逻辑表的下标保持稳定,同时避免为不会再访问的 token 占用真实 KV 空间。这个占位块不维护普通引用计数,也不能进入前缀缓存。

一块 KV 到底占多少显存#

KVCacheBlock 本身只是几十字节量级的 Python 元数据;真正昂贵的是它在每一层对应的 K/V 张量切片。对采用 GQA 或 MQA 的 decoder-only Transformer,单个 token 的 KV Cache 字节数可写成:

Mtoken=2LHkvDhSM_{token}=2L H_{kv} D_h S

其中:

  • 22 分别代表 Key 和 Value;
  • LL 是需要缓存 K/V 的 Transformer 层数;
  • HkvH_{kv} 是每层 KV head 数,而不是 query head 数;
  • DhD_h 是每个 head 的维度;
  • SS 是每个缓存元素的字节数,例如 BF16 为 2,FP8 为 1。

若一个 block 容纳 BB 个 token,那么一个 block 覆盖的实际缓存为:

Mblock=B⋅Mtoken=2BLHkvDhSM_{block}=B\cdot M_{token}=2BLH_{kv}D_hS

以 Llama 3.1 8B 的常见配置为例,取 L=32L=32、Hkv=8H_{kv}=8、Dh=128D_h=128、BF16 的 S=2S=2,则每 token 需要:

Mtoken=2×32×8×128×2=131072 bytes=128 KiBM_{token}=2\times32\times8\times128\times2=131072\ \text{bytes}=128\ \text{KiB}

若 block size 为 16,一个逻辑 block 对应约 2 MiB 的 K/V。此处是整个模型的总量;张量并行下每个 rank 实际持有多少,取决于 KV head 如何切分或复制,不能简单地对所有模型一律除以 TP size。

这个量级解释了为什么 Scheduler 必须在模型执行前确认槽位。若本轮同时让 256 个 decode 请求跨过 block 边界,仅新 block 就可能需要约 512 MiB;只看“每个请求本轮都只算 1 token”会严重低估显存增量。

固定大小带来一个可控的内部碎片上界。设请求当前长度为 TT,需要的 block 数为:

Nblock=⌈TB⌉N_{block}=\left\lceil\frac{T}{B}\right\rceil

浪费槽位数为 NblockB−TN_{block}B-T,严格小于 BB。传统连续预分配常按最大输出长度预留,浪费量随“最大长度减实际长度”增长;分页后每个请求最多浪费最后一个 block 的未填满部分。PagedAttention 的详细地址翻译过程可参考 PagedAttention 完全拆解。

还要区分三个经常都被叫作 block size 的数:

  1. Scheduler block size:请求进度与前缀命中的对齐粒度;混合 group 时可能取各 group 粒度的最小公倍数。
  2. KV manager block size:某个 cache group 实际分配一次覆盖多少 token。
  3. Kernel block size:注意力 Kernel 消费 block table 时采用的粒度。

标准模型中三者常相同,所以差别不明显。混合布局中,一个 32-token 的 manager block 可以在 Worker BlockTable 中展开成两个 16-token kernel block;此时控制面一次分配一个物理页,Kernel 却看到两个连续子块 ID。

allocate_slots() 的完整决策链#

这段函数最重要的设计不是“从 free queue 弹出若干 block”,而是把查询、容量验证和状态提交分成明确阶段。若先修改引用计数,再发现容量不够而返回 None,Scheduler 重试时就会看到被半更新的状态;轻则泄漏 block,重则让一个缓存块被错误覆盖。

先把普通全注意力、单 cache group 的主路径压缩成伪代码:

allocate_slots(request, new_tokens, cached_hit_blocks, lookahead):
total_computed = request.computed + cached_hit_tokens
target_tokens = total_computed + new_tokens + lookahead
释放已经不再需要的窗口外 block
required = 预测 target_tokens 需要的新增物理 block
required += 命中但 ref_cnt=0、touch 后将离开 free queue 的 block 数
if required > free_queue.size:
return None
touch(cached_hit_blocks)
把命中块接入 request 的 block table
new_blocks = 从 free queue 分配缺少的 block
为已经完整且可提交的 block 建立哈希索引
return new_blocks

真实函数还处理外部 KV Connector、encoder cache、投机 lookahead、混合 group、watermark 与滑动窗口回收,但所有分支都必须守住同一个事务边界:None 表示本轮没有拿到任何半成品所有权;非空结果表示 Scheduler 可以放心把计划交给 Worker。

第一步:把五段 token 空间分开#

源码注释把请求位置画成五段:

| 已计算 | 新命中的本地前缀 | 外部已计算 KV | 本轮新 token | lookahead |

它们对分配的影响不同:

  • 已计算部分已经在当前请求的 block table 中;
  • 本地命中部分已有物理块,但可能处于 ref_cnt=0 的可淘汰状态;
  • 外部 KV 只有内容存在于远端,接收前仍需在本地分配目标块;
  • 本轮新 token 需要新的可写槽位;
  • lookahead 为 EAGLE 等带 KV 的草稿模型提前预留位置,最终可能因拒绝而不提交。

因此“还需多少块”不能只用 ceil(num_new_tokens / block_size)。若请求最后一个 block 还有空位,新 token 可能完全不需要新块;若前缀命中块目前无人引用,touch 又会把它从 free queue 取走,相当于减少可分配容量。

第二步:预测,而不是边分边看#

对最简单的 full-attention manager,目标长度为 TtargetT_{target}、已有请求块数为 NownedN_{owned} 时,新块基数接近:

Nnew=max⁡(⌈TtargetB⌉−Nowned,0)N_{new}=\max\left(\left\lceil\frac{T_{target}}{B}\right\rceil-N_{owned},0\right)

若命中的块中有 Nevictable_hitN_{evictable\_hit} 个当前 ref_cnt=0,它们虽然已经存着正确 K/V,却仍计入 free queue。touch 会把它们从 free queue 移除,所以容量检查必须再加上这部分:

Ncheck=Nnew+Nevictable_hitN_{check}=N_{new}+N_{evictable\_hit}

这个加项很容易漏。假设 free queue 只有 3 个 block,请求命中其中 2 个,后缀还需新分配 2 个。表面看 2 <= 3,似乎可行;实际 touch 两个命中块后只剩 1 个,后缀分配必然失败。正确检查是 2+2>32+2>3,应在任何状态修改前返回 None。

混合 Coordinator 会分别调用每个 manager 的 get_num_blocks_to_allocate(),再求和。滑动窗口 manager 还会扣掉即将离开窗口、可在本轮前安全释放的块;cross-attention manager 则按 encoder token 计算一次性静态需求。

第三步:先回收永远不会再读的块#

remove_skipped_blocks() 位于正式容量检查之前。对全注意力,它通常什么也不做;对滑动窗口或局部注意力,窗口左侧已经处理且不可能再参与未来 attention 的块可以先归还池中。

这一步看似破坏“失败不修改状态”,但被释放的是按照已提交计算进度判定、无论本轮是否执行都已无用的块。源码还特意按扣除 in-flight token 后的已处理边界释放,避免异步步骤仍在读取的窗口被过早回收。它属于独立成立的垃圾回收,而不是本轮计划的乐观提交。

第四步:容量够才接管命中块#

容量验证通过后,Coordinator 执行两阶段提交:

  1. 对所有 group 的本地命中块先调用 touch,把它们变成当前请求的强引用;
  2. 再为所有 group 的外部命中后缀和本轮新增 token 分配物理块。

“先 touch 所有 group,再给任一 group 分配新块”是跨 group 正确性要求。若 group 0 先从 free queue 取新块,它可能正好淘汰 group 1 尚未 touch 的命中块;之后 group 1 拿到的 block ID 已经代表另一份内容。两阶段操作相当于先给所有将复用的页加锁,再做可能触发淘汰的分配。

第五步:只缓存可提交的位置#

分配完成后,cache_blocks() 为已经完整、未来可复用的块登记哈希。投机解码的草稿 token 可能被拒绝,不能因为预留了 slot 就当作真实前缀写进哈希索引。因此源码把可缓存上界限制在 request.num_tokens,而不是包含全部 lookahead 的容量末端。

远程 P/D 传输还有 delay_cache_blocks 分支。目标 block 已分配但数据尚未到达时,不能把它发布成可命中的本地缓存;否则其他请求会复用一块尚未填完的数据。传输完成并确认有效后,系统才提交哈希元数据。

这里还有一个时间线细节:普通本地新块的哈希可以在本轮控制面先登记,因为下一次 Scheduler 查找发生在本轮模型执行完成之后;同一 Engine Core step 内不会有另一个调度轮插进来读取半写数据。异步外部传输不满足这个保证,所以必须延迟提交。

前缀缓存:块哈希为什么必须包含父前缀#

前缀缓存不是按 block 内 token 单独做内容寻址。第 ii 个 block 的哈希链可以写成:

hi=H(hi−1, ti, ei)h_i=H\left(h_{i-1},\ t_i,\ e_i\right)

其中 hi−1h_{i-1} 是前一块所代表的完整前缀哈希,tit_i 是当前块的 token ID 序列,eie_i 是 LoRA、多模态输入、cache_salt 或 prompt embedding 等附加键。缓存表的实际键还包含 KV cache group ID:

ki=(hi,g)k_i=(h_i,g)

同一 token 片段出现在不同历史之后,K/V 通常不同。因果注意力让当前 token 读取此前所有可见 token;RoPE 等位置编码也让绝对或相对位置影响 Key。若只哈希 tit_i,下面两个第二块会错误碰撞:

[系统提示 A][请总结以下文本]
[系统提示 B][请总结以下文本]

加入父哈希后,第二块的键同时承诺“当前内容相同”和“之前的整个前缀相同”。这相当于用哈希链隐式表示前缀树,不需要为每个 token 维护显式 trie 边。

前缀哈希、请求块表与空闲队列的关系
前缀哈希、请求块表与空闲队列的关系

图源:vLLM Automatic Prefix Caching 设计文档。同一个 KVCacheBlock 同时出现在哈希索引、请求 block table 和 free queue 中的条件不同;这些结构保存的是引用,不是三份 K/V 副本。

为什么附加键不能省略#

token ID 相同不代表模型计算相同:

  • 不同 LoRA adapter 会改变投影权重,同一 token 前缀的 K/V 不再等价;
  • 多模态占位 token 相同,但对应图像或视频特征可能不同;
  • 直接传入的 prompt embedding 可能无法由 token ID 唯一确定;
  • 多租户可用 cache_salt 隔离前缀命名空间,避免跨安全域复用。

所以缓存键的语义不是“这段文字一样”,而是“所有会影响这段前缀 K/V 的输入都一样”。哈希冲突仍是理论风险;对不可信或跨进程缓存场景,官方文档支持 SHA-256 一类密码学哈希,以更高 CPU 成本换更强的碰撞抵抗。

为什么标准路径只发布完整 block#

一个未填满尾块之后还会继续写 token。若它过早进入共享缓存,另一个请求复用后,两者可能同时把不同 token 写入同一物理块,必须立即 copy-on-write;更糟的是哈希所承诺的内容范围会随着原请求追加而变化。

标准 APC 因而只把完整 block 作为稳定缓存单元。假设 block size 为 4,两个请求共享前 10 个 token,直接命中的只有前 8 个;第 9、10 个 token 所在尾块需要重新计算。当前混合 KV 管理已支持受约束的细粒度或部分命中,但它必须记录精确 token 边界并执行 CoW,不能用普通完整块规则偷换。

为什么全命中仍要重算最后一个 token#

KV Cache 保存的是历史位置的 K/V,不自动保存“读取最后位置后得到的下一 token logits”。若 prompt 的所有 token 都被当作已计算而本轮不运行模型,就没有新 logits 可供采样。get_computed_blocks() 因而把最大命中长度限制为 request.num_tokens - 1,至少重新计算最后一个位置。

由于普通命中要求 block 对齐,这可能不是只重算 1 个 token,而是退回到上一个完整块边界。它牺牲少量重复计算,换来清晰的 block 不变量和统一 Kernel 输入。

BlockPool:free queue、引用计数与双重身份#

BlockPool 同时维护三套视图:

blocks[block_id] 所有物理块的稳定对象表
free_block_queue 当前没有活跃请求引用的块
cached_block_hash_to_block[(hash, gid)] 可按内容复用的块

这三个集合不是互斥分区。一个 block 可以同时位于 free queue 和哈希缓存中,这正是 APC 能够“几乎不额外占容量”的关键。

free 表示可覆盖,不表示内容已擦除#

当请求结束时,系统对它持有的每个 block 将 ref_cnt 减一。降到 0 说明没有活跃请求会继续读取它,于是 block 进入 free queue。但若它已经拥有 block_hash,哈希映射仍保留,显存中的旧 K/V 也没有必要立刻清零。

此时它同时满足:

  • 从分配器角度看是 free:新请求缺空间时可以覆盖;
  • 从内容缓存角度看是 cached:在被覆盖之前,相同前缀仍可复用。

这是一种惰性失效。立刻清除哈希不会增加可用物理容量,只会白白丢掉潜在命中;永久保留缓存块又会挤压新请求。把 cached block 放回 free queue,同时推迟到真正再分配时才淘汰,恰好兼顾两者。

首个请求分配并缓存完整块
首个请求分配并缓存完整块

图源:vLLM Automatic Prefix Caching 设计文档。示例 block size 为 4。请求 A–O 占用 block 0–3,其中 0–2 已填满并建立 A–D、A–H、A–L 的链式哈希;block 3 只有 MNO,仍不可作为普通完整块命中。free queue 从 block 4 开始。

请求再生成 P 与 Q 后,原来的 block 3 变满并缓存,新 token Q 写入新分配的 block 4:

尾块填满后扩展一个新块
尾块填满后扩展一个新块

图源:vLLM Automatic Prefix Caching 设计文档。req_to_blocks 的逻辑顺序是 0、1、2、3、4,物理 ID 不需要连续;这里只因初始池尚未产生淘汰,恰好连续。

为什么不用 Python deque#

命中 cached block 时,它可能位于 free queue 的中间。系统要在 O(1)O(1) 时间把它摘出,否则每次前缀命中都可能线性扫描空闲列表。KVCacheBlock 自带 prev_free_block 与 next_free_block,FreeKVCacheBlockQueue.remove(block) 只改相邻指针。

这是一种侵入式链表:队列不再额外创建包装节点,block 自己就是节点。实现还放置 fake head 与 fake tail 哨兵,减少头尾分支。代价是同一 block 在任意时刻最多属于一条 free queue,链表指针必须只由该队列修改;重复 remove 或重复 append 会直接破坏结构。

分配时才真正淘汰#

BlockPool.get_new_blocks(n) 从 free queue 头部弹出 nn 个候选。对每个候选:

  1. 若还带 block_hash,从哈希表移除所有指向该 block 的键;
  2. 清空 block 的哈希元数据;
  3. 将 ref_cnt 从 0 增为 1;
  4. 把 block 交给新请求。

从这一步起,旧前缀再也不能命中该物理块。注意淘汰的是“内容身份”,不是销毁 KVCacheBlock 对象;稳定的 block_id 仍用于寻址同一块物理槽位,只是槽位即将装入新的 K/V。

回收顺序同时编码复用策略#

free_blocks() 把引用归零的块分成两类:

  • 没有哈希的块不能产生前缀命中,放到 free queue 前端,尽快 LIFO 复用,通常也有更好的局部性;
  • 有哈希的块放到尾端,尽可能延迟覆盖,形成近似 LRU 的缓存淘汰。

同一个请求释放时还会逆序遍历 block table,让更长前缀的尾块先于短前缀块进入同一批次的淘汰顺序。原因是第 1 块可能被大量请求共享,而第 20 块只对具有很长共同前缀的请求有价值;在访问时间相同的情况下,优先牺牲尾部通常保留更多潜在命中。

这不是一个独立占容量的传统 LRU cache。free queue 既是物理分配器的空闲表,又是 cached block 的淘汰队列;缓存命中率与新请求容量通过同一顺序自然竞争。

touch、evict、free 的顺序为什么不能交换#

三个词描述的是不同状态转换:

操作哈希索引ref_cntfree queue典型触发时机
touch(block)保留加一若原为 0,则移出新请求命中前缀
free_blocks(block)通常保留减一降到 0 时加入请求结束或抢占
evict(block)删除不必改变原位置不一定改变block 即将被改写或显式失效

touch 不是“更新 LRU 时间戳”这么简单。对一个 free+cached block,它把弱缓存内容重新变成活跃请求的强引用:先从 free queue 摘除,再增加 ref_cnt。若只加引用却不移出队列,后续 get_new_blocks() 仍可能把它弹出并覆盖;若只移出却不加引用,请求结束时的减一会把计数变成负数。

free 也不等于 evict。请求结束只表示这个请求不再使用 block,不代表内容没有复用价值。只有物理槽位要改写、模型权重变化导致旧 KV 全部失效,或外部组件明确要求作废时,才需要删除哈希。

正确顺序一:先做全局容量检查,再 touch#

容量不足时,allocate_slots() 必须无副作用地返回 None。若先 touch 命中块,它们会从 free queue 消失;随后新块不足,函数退出,但请求可能还没进入 running,也没有正常 free 路径来撤销这些引用。这会造成容量凭空减少。

所以 get_num_blocks_to_allocate() 不仅预测真正的新块,还把即将被 touch 的 ref_cnt=0 命中块计入容量消耗。所有 group 的预测求和通过后,才统一接管所有权。

正确顺序二:先 touch 所有命中,再从池中取新块#

混合 group 共享一个 BlockPool。假设 group 0 需要新块,group 1 的命中块仍在 free queue 中;如果先满足 group 0,池头可能正是 group 1 将复用的 block。先 touch 全部命中,可以把所有“本轮必须保住”的页从淘汰候选中摘出,再安全分配后缀。

正确顺序三:先删除旧哈希,再写新内容#

get_new_blocks() 弹到 cached block 时,必须先从 cached_block_hash_to_block 删除旧键,再把槽位交给新请求。GPU 写入可能异步发生,但从控制面分配成功那一刻起,旧内容就不再受到保护。若哈希删除滞后,下一调度轮可能把同一个 block 同时分给“旧前缀命中者”和“新内容所有者”。

正确顺序四:最后一个使用者离开后才能进入 free queue#

多个请求复用同一完整前缀时,block 的 ref_cnt 大于 1。任一请求结束只减一;只有降到 0 才加入 free queue。这是共享块的基础所有权规则。

部分尾块的情况更严格。若请求命中到一个 group block 内部,源 block 仍要为其他命中者保持只读,当前请求则需要私有 CoW block:先分配目标块,把有效前缀从源复制过去,再把当前请求的 block table 指向目标。复制完成前,源和目标都必须额外 pin 住,防止同一步中的 free 让它们被重新分配。

把这些规则合在一起,可以得到一个实用不变量:

只要某个 SchedulerOutput 或尚未完成的 GPU 操作仍可能读取 block X,
X 就不能同时成为 get_new_blocks() 可返回的候选。

引用计数、free queue、延迟释放与 CoW pin 都是在不同时间线上维护这条不变量。

从 SchedulerOutput 到 ModelRunner 的增量 Block Table#

Engine Core 管理的是 KVCacheBlock 元数据,跨进程发给 Worker 时会转换成整数 block ID。KVCacheBlocks.get_block_ids() 的返回值仍按 cache group 分组,例如:

(
[7, 1, 9], # group 0: full attention
[12, 4], # group 1: sliding window / state group
)

对首次进入 Worker 的请求,scheduled_new_reqs 携带当前完整 block ID 序列;对已经存在的请求,scheduled_cached_reqs 通常只带 new_block_ids 差量。这样每轮 decode 不需要反复序列化从 prompt 开始的整张 block table。

GPU ModelRunner 收到计划后的核心流程是:

if new_request:
req_state.block_ids = new_req_data.block_ids
input_batch.block_table.add_row(block_ids, req_index)
else:
for old_ids, new_ids in zip(req_state.block_ids, new_block_ids):
old_ids.extend(new_ids)
input_batch.block_table.append_row(new_block_ids, req_index)
input_batch.block_table.commit_block_table(num_reqs)
input_batch.block_table.compute_slot_mapping(
num_reqs,
query_start_loc,
positions,
)

CPU 镜像为什么先更新再拷到 GPU#

BlockTable 内部维护 CPU/Numpy 可写镜像和 GPU tensor。Python 控制面擅长做请求行的 append、move、swap;注意力 Kernel 则需要设备端连续的 int32 表。commit_block_table() 只把当前有效行拷到 GPU,避免每次在 GPU 上执行大量细碎控制操作。

每一行对应一个活跃请求,第 jj 列表示该请求的第 jj 个逻辑块映射到哪个物理 block。请求在 continuous batching 中完成后,batch 行可能被其他请求复用,因此 move_row() 不仅复制目标行,还会清空源行,防止 dummy run 或异步 Kernel 沿陈旧 block ID 写入已经回收的物理页。

block table 与 slot mapping 不是同一张表#

block table 完成“逻辑 block 到物理 block”的翻译;slot mapping 则为本轮每个具体 token 找到最终写入槽位。忽略并行切分后,位置 pp 的基本映射为:

j=⌊pB⌋,o=p mod Bj=\left\lfloor\frac{p}{B}\right\rfloor,\qquad o=p\bmod Bslot=block_table[request,j]×B+oslot=block\_table[request,j]\times B+o

jj 是逻辑 block 下标,oo 是 block 内偏移。Kernel 根据 slot 把新 Key/Value 写到对应物理位置;读取历史 K/V 时则沿 block table 取得非连续物理页。DCP、PCP 或 manager block 与 kernel block 粒度不一致时,真实实现还会加入 rank 交错和虚拟子块展开,但基本的两级地址翻译不变。

为什么 Manager 不直接把完整表发给 Kernel#

职责分离带来三个收益:

  1. Engine Core 可以在没有 CUDA context 的进程里完成全局调度和所有权判定;
  2. Worker 保留持久请求状态,每轮只接收差量,IPC 成本更接近新增工作量;
  3. 不同 attention backend 可以把同一逻辑 block table 转换成自己的 Kernel metadata,而不把后端细节泄漏给 Scheduler。

代价是跨进程一致性更难:Scheduler 的 req_to_blocks、Worker 的 req_state.block_ids、GPU 上已经提交的 block table 必须代表同一版本。finished_req_ids、抢占重置、stale output 检查和延迟 free 都不能缺少,否则整数 ID 仍然合法,却可能已经指向另一请求的新内容。

用一个十块内存池走完请求生命周期#

用一个具体例子把查询、分配、缓存、释放和淘汰串起来。为简化表示,暂时忽略 null block,设池中有 10 个物理块,ID 为 0–9,block size 为 4,只有一个 full-attention group。

时刻 0:池刚初始化#

free queue: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
hash map: {}
requests: {}

所有 block 的 ref_cnt=0、block_hash=None。此时 free queue 按 block ID 排列,还没有 LRU 语义上的访问历史。

时刻 1:R0 输入 15 个 token#

R0 的 token 为 A–O,需要 ceil(15/4)=4 个 block。分配器从队首取 0、1、2、3:

R0 blocks: [0, 1, 2, 3]
free queue: [4, 5, 6, 7, 8, 9]
ref_cnt: 0..3 -> 1

模型完成 prefill 后,0、1、2 分别覆盖 A–D、E–H、I–L,是完整块;3 只覆盖 M–O,是部分块。控制面为前三块登记链式哈希,块 3 保持 unhashed。

时刻 2:R0 继续生成 P、Q#

P 填满 block 3,A–P 对应的前缀哈希可登记。Q 跨过边界,需要 block 4:

R0 blocks: [0, 1, 2, 3, 4]
free queue: [5, 6, 7, 8, 9]
cached: 0, 1, 2, 3
partial: 4

注意“新生成一个 token”不总是“新分配一个 block”。P 复用已有尾部空槽,Q 才触发物理分配。这也是 Scheduler 的 token budget 与 KV block budget不能直接互换的原因。

时刻 3:R0 完成#

free(R0) 逆序释放 4、3、2、1、0。block 4 没有哈希,被放到 free queue 前端;0–3 有哈希,被放到尾端并保留缓存身份:

free queue: [4, 5, 6, 7, 8, 9, 3, 2, 1, 0]
hash map: {A-D -> 0, A-H -> 1, A-L -> 2, A-P -> 3}
ref_cnt: all 0

所有十个 block 都可供分配,所以物理容量没有被 APC 永久占住;其中 0–3 又可被相同前缀命中。这里最能体现 free 与 cached 的正交性。

时刻 4:R1 命中 A–L,但从 M 开始不同#

get_computed_blocks(R1) 依次查到 0、1、2,到第四块哈希失配时停止。容量预测要考虑三块当前都在 free queue:touch 会把 0、1、2 从队列尾部摘走并把引用计数改为 1。R1 的不同后缀需要一个可写块,分配器从队首取 block 4:

R1 blocks: [0, 1, 2, 4]
free queue: [5, 6, 7, 8, 9, 3]
ref_cnt: 0,1,2,4 -> 1

旧 block 3 仍缓存 A–P,但 R1 不使用它。若另一个请求恰好以 A–P 开头,仍可在它被覆盖前命中。

时刻 5:一个大请求分配失败#

假设 R2 没有任何前缀命中,并且本轮需要 7 个新 block;free queue 只有 6 个。allocate_slots() 在 touch 或新分配之前返回 None。Scheduler 可以抢占一个 victim,释放其 block 后重试;也可以把 R2 留在 waiting。

关键是失败后状态仍为:R1 继续拥有 0、1、2、4,free queue 仍是 [5,6,7,8,9,3],R2 没有拿到“前 6 个成功、第 7 个失败”的半张表。容量预测让分配具有全有或全无的语义。

时刻 6:压力迫使缓存淘汰#

当后续分配最终弹到 block 3 时,_maybe_evict_cached_block(3) 先删除 A–P 的哈希,再把 ref_cnt 设为 1,随后 Worker 才会向该槽位写新 K/V。A–D、A–H、A–L 的更短前缀仍然存在,因为它们各自对应独立 block 和哈希。

这个过程说明 APC 的淘汰粒度是 block,不是整条请求。长前缀尾部可以逐步消失,公共度更高的前部仍被保留。

混合 KV Cache 为什么需要 Coordinator#

普通 Llama 类模型的每层都使用全注意力,一个 cache group 足以描述全部层。混合模型可能把 full attention、sliding window attention、Mamba 或其他状态空间层交错放置。它们对“第 tt 个 token 还需保留多少历史”的回答不同,不能共用一张简单的等长 block table。

Hybrid KV Cache Manager 的分层结构
Hybrid KV Cache Manager 的分层结构

图源:vLLM Hybrid KV Cache Manager 设计文档。KVCacheManager 对 Scheduler 保持一个接口;Coordinator 选择单一、混合或无前缀缓存路径;每个 SingleTypeKVCacheManager 实现一种注意力的分配与命中语义。

为什么要先分组#

假设模型有 10 层 full attention 和 20 层 sliding window attention。若直接要求“一个 block ID 表示所有 30 层的同一 token 区间”,滑动窗口层释放旧块时会连带破坏 full-attention 层仍需读取的历史。

V1 把具有相同缓存规格的层组织成 group。每个 group 有自己的 block 序列和生命周期,Coordinator 再为同一请求返回:

KVCacheBlocks(
blocks=(
full_attention_blocks,
sliding_window_group_0_blocks,
sliding_window_group_1_blocks,
)
)

Worker 也为每个 group 维护独立 block table。Kernel 执行某层时,使用该层所属 group 的表,而不是假设所有层共享同一映射。

相同 page size 怎样共享底层缓冲区#

混合注意力 KV Cache 的物理内存布局
混合注意力 KV Cache 的物理内存布局

图源:vLLM Hybrid KV Cache Manager 设计文档。示例把 10 层 full attention 与两组各 10 层 sliding window 映射到 10 个 KVCacheTensor;黄色、浅绿和深绿区域由不同 group 的 block ID 划分。一个逻辑 block 会对应每个相关层缓冲区中的一个 page 切片,而不是一段包含全模型所有层的单一 Python 对象。

这种布局要求各 group 的 page size 可兼容。page size 是“该 group 中每个物理 block 在底层缓冲区占多少字节”,大致由组内层数、block token 数、KV hidden size 与 dtype 决定。构建 KVCacheConfig 时,系统会选择能够共享缓冲区的分组与 page 几何;运行时的 manager 只操作统一的物理 block ID 池。

前缀命中为什么取共同可恢复边界#

全注意力的缓存命中是向下闭合的:前 12 块都在,就可以恢复到第 12 块末端。滑动窗口却可能已经释放早期块,只保留某些能恢复当前窗口的尾部;Mamba 还可能需要某个精确边界状态。

Coordinator 要求所有参与计算的 group 在同一 token 位置都有有效状态。可以把最终命中长度写成概念式:

Thit=max⁡{t∣∀g, stateg(t) 可恢复}T_{hit}=\max\left\{t\mid \forall g,\ \text{state}_g(t)\ \text{可恢复}\right\}

它不一定等于各 group 独立最长命中的简单最小值,因为某些稀疏 group 的可恢复位置不是连续区间;真实实现会按候选边界迭代收缩,直到所有 group 同意。

这层协调避免一种隐蔽错误:Scheduler 把 num_computed_tokens 设为 full-attention 的最长命中位置,但某个 Mamba group 在那里没有有效递归状态。模型仍可能运行,却会从错误状态继续生成,结果不是显式 OOM,而是静默数值错误。

为什么分配要对所有 group 原子化#

Coordinator 先汇总每个 manager 的需求,再与共享 BlockPool 的空闲数比较;成功后先 touch 全部本地命中,最后逐 group 分配新块。若每个 group 独立“检查并提交”,前几个成功、最后一个失败时,请求会持有不完整状态,回滚还要精确撤销哈希、引用计数和 free queue 顺序。

一次全局预测加两阶段提交,把复杂回滚变成简单的“容量不够就不开始”。这与数据库事务的思路相似,但对象是 GPU 物理页的所有权。

失败路径、正确性不变量与调试方法#

KV Cache 管理问题常表现为随机输出错误、偶发非法访问或吞吐抖动,而不是稳定复现的 Python 异常。沿不变量调试比盯着某个大函数有效。

六条必须一直成立的不变量#

  1. ref_cnt > 0 的普通 block 不能位于 free queue;ref_cnt == 0 且非 null 的 block 应处于可回收状态。
  2. 哈希表中指向 block 的键,必须与该 block 当前显存内容、group ID 和有效 token 边界一致。
  3. 一个请求的 num_computed_tokens 不能越过任一必要 cache group 的可恢复状态。
  4. SchedulerOutput 引用的 block 在 Worker 消费并结束相关异步读写前不能被重新分配。
  5. Worker block table 的有效前缀必须与 Engine Core 的 req_to_blocks 一致;抢占或完成后的旧行不能继续参与 Kernel。
  6. lookahead、外部传输和部分命中所产生的乐观槽位,未提交前不得发布为普通可复用前缀。

常见故障怎样定位#

症状优先检查可能根因
明明有空闲块却频繁 allocate_slots=Nonerequired_blocks 与 free queue 长度命中块 touch 成本、lookahead 或 group 总需求被忽略
前缀命中后输出错误哈希附加键、group 边界、最后 token 回退复用了不等价 K/V 或状态边界不一致
请求结束后 KV 使用率不降ref_cnt、延迟 free、Worker in-flight 状态引用泄漏或异步操作仍 pin 住 block
低负载偶发非法显存访问stale row、finished ID、block 重用时间线Worker 仍使用已回收 ID
APC 命中率低但前缀看起来相同token IDs、block 对齐、LoRA/MM extra key文本相同不等于缓存键相同
混合模型在长上下文才出错每 group 命中长度与窗口回收Coordinator 恢复边界或 null block 处理错误

推荐断点顺序#

先用单卡、关闭外部 KV 传输和投机解码,构造两个共享长前缀的短请求。每轮只记录 block ID 与计数,不打印整个 K/V 张量:

step=7 req=R1 hit_tokens=128 scheduled=1
owned=[[4, 9, 2, 15, 7, 8, 11, 3, 6]]
new=[[6]] free=213 ref(hit)=[1,1,1,1,1,1,1,1]

然后按以下位置逐层确认:

  1. KVCacheManager.get_computed_blocks():核对哈希从哪一块开始 miss,为什么最多命中 num_tokens - 1。
  2. get_num_blocks_to_allocate():分别记下目标块数、已有块数、evictable hit 数和 lookahead。
  3. BlockPool.touch():确认 ref_cnt=0 的命中块从 free queue 中间移除。
  4. BlockPool.get_new_blocks():确认弹出 cached block 时旧哈希先被清理。
  5. SingleTypeKVCacheManager.free():确认请求块逆序释放,引用只降一次。
  6. GPUModelRunner 的请求状态更新:确认 new request 使用完整表,cached request 只 append 差量。
  7. BlockTable.commit_block_table() 与 compute_slot_mapping():检查设备表行和 token 位置最终映射到的 slot。

若基础路径正确,再逐项打开 prefix caching、chunked prefill、投机解码、混合注意力和 KV Connector。一次同时打开全部特性,会让“容量估计错”“哈希错”和“异步生命周期错”产生相似症状。

三个指标不能混为一谈#

  • KV cache usage 反映 free queue 之外的物理容量比例;free+cached block 仍算 free,因此它不等于“有多少旧前缀内容留在显存”。
  • prefix cache hit rate 反映请求 token 有多少跳过了重复 prefill;它不会直接提高普通 decode 单 token 的计算速度。
  • preemption count 反映调度计划经常无法兑现 KV 容量;低命中可能增加 prefill 工作,但抢占还受到并发数、输出长度、block size 与 lookahead 影响。

调优时应把三者与 TTFT、ITL 分开观察。共享系统提示词很多时,APC 主要降低 TTFT 与 prefill 计算;长输出、前缀很少重复时,增加可用 KV block 或限制活跃请求通常比调整哈希更直接。

源码阅读路线与系列进度#

到这里,vLLM V1 的控制面已经从请求入口走到物理 block ID:

篇次已完成主题建立的源码坐标
一整体架构与模块全景Entrypoint、Engine、Worker、Kernel
二请求穿过 V1 引擎add_request()、EngineCore.step()、输出提交
三Scheduler 的批次决策token budget、waiting/running、抢占
四KV Cache Managerprefix lookup、block pool、分配、淘汰、block table 交接

后续 vLLM 子系列按数据流继续向下,待讲模块暂定为:

顺序待讲模块核心问题
五GPUModelRunner 与 InputBatch变长请求怎样拼成一次前向,attention metadata 怎样生成
六Attention Backend 接口FlashAttention、FlashInfer 等后端怎样消费统一元数据
七PagedAttention Kernel 源码block table 怎样落实为合并访存、分块归约与数值稳定 softmax
八模型加载与权重映射Hugging Face 权重怎样进入 vLLM 模型类与并行分片
九采样器与结构化输出logits、top-k/top-p、grammar mask 与异步采样如何衔接
十CUDA Graph 与执行形状动态 batch 怎样落入可重放图,何时退回 eager 路径
十一Executor 与分布式 WorkerTP、PP、DP 的进程拓扑、控制消息与集合通信
十二MoE 推理路径expert 映射、路由、All-to-All 与并行配置怎样进入执行层
十三在线服务热路径API Server、流式输出、背压、取消与可观测性
十四扩展能力LoRA、多模态、KV Connector 与平台插件怎样接入主干

具体分篇会随源码职责边界调整,但在这些主干模块讲清之前不会切换到 SGLang;进入 SGLang 时再从独立的框架总览开始,并把相同机制与 vLLM 横向对照。

本篇最值得保留的源码地图有五条:

  1. KVCacheManager 是 Scheduler 的接口,不持有真正的 K/V 张量。
  2. allocate_slots() 先预测全局需求,容量够才 touch 和分配,避免半提交状态。
  3. free 与 cached 可以同时成立;只有物理槽位即将重用时才真正 evict 旧哈希。
  4. Coordinator 把不同 cache group 收敛到共同可恢复边界,并保证跨 group 原子分配。
  5. Engine Core 只把 block ID 差量交给 Worker,ModelRunner 再构造设备 block table 与 slot mapping。

下一篇继续留在 vLLM,不切换框架:进入 GPUModelRunner 的输入批处理,拆解 InputBatch 如何把不同请求的 token、position、query start location、block table 和 attention metadata 拼成一次模型前向。它会从本文最后的 commit_block_table() 接着向下,直到注意力 backend 真正消费这些地址。

参考资料#

  1. vLLM V1 KVCacheManager 源码(本文固定提交)
  2. vLLM BlockPool 源码(本文固定提交)
  3. vLLM KVCacheCoordinator 源码(本文固定提交)
  4. vLLM SingleTypeKVCacheManager 源码(本文固定提交)
  5. vLLM Worker BlockTable 源码(本文固定提交)
  6. vLLM Automatic Prefix Caching 设计文档
  7. vLLM Hybrid KV Cache Manager 设计文档
  8. Inside vLLM: Anatomy of a High-Throughput LLM Inference System
  9. PagedAttention 论文:Efficient Memory Management for Large Language Model Serving with PagedAttention
  10. Automatic Prefix Caching RFC
  11. vLLM V1 KV Cache 管理模块源码解析
  12. vLLM KV Cache Internals: Blocks, the Block Pool and the Free Queue

文章分享

如果这篇文章对你有帮助,欢迎分享给更多人!

vLLM 源码完全拆解(四):KVCacheManager 如何分配、复用与回收物理 Block
https://pinghaoyang.com.cn/aigc/posts/vllm-kv-cache-manager/
作者
平昊阳
发布于
2026-10-07
许可协议
CC BY-NC-SA 4.0

评论区

Profile Image of the Author
平昊阳
乘长风,破巨浪, 展鸿图于未央!
--
总访问量
--
访客数
公告
欢迎来到我的个人博客!欢迎关注交流吖!
更多相关公告,见
「社交-留言」。
音乐
封面

音乐

暂未播放

0:000:00
暂无歌词
站点统计
文章
212
分类
26
标签
291
总字数
2,407,536
运行时长
0 天
最后活动
0 天前

文章目录