Mooncake Store KV Cache
一句话结论:PagedAttention 是 GPU 内部的页式执行地址管理,Mooncake Store 是 GPU 外部的分布式分层缓存;Mooncake 命中后仍需先把 KV 数据装入 GPU Block,再更新 PagedAttention Block Table。
先区分页表与数据页
讨论 KV Cache 的“页”时,最容易把三个不同对象混在一起:
- 逻辑 KV Block:请求按 Token 顺序切分出的第 0、1、2 个块。
- GPU 物理 Block:K/V Tensor 在本次运行中实际占用的显存块。
- Block Table:保存“逻辑块序号 → GPU 物理 Block ID”的映射。
设请求的逻辑块序列为 $L_0,L_1,L_2,\ldots$,PagedAttention 为它们分配 GPU 物理块 $P_0,P_1,P_2,\ldots$,页表保存的是:
$$
L_i \rightarrow P_i
$$
严格来说,页表不是 KV Cache 数据本身。vLLM Scheduler 在 CPU 侧保存请求、引用计数、Block Hash 与所有权等调度元数据;Attention Kernel 使用的 Block Table 则位于设备侧,内容是 GPU Block ID。SSD 不保存这张临时页表,页表也不会直接指向 SSD 文件。
同一段前缀下一次被其他请求命中时,内容身份仍可由相同的 Block Hash 表示,但本次分配到的 GPU Block ID 可能完全不同。新的页表只需重新建立逻辑块到新物理块的映射。
两级管理全景
下图把两套映射放在一起:上层是 PagedAttention 的 GPU 执行地址,下层是 Mooncake Store 的内容身份与副本位置。绿色路径表示真正的 KV 数据传输,蓝色与紫色路径表示地址和元数据查询。
图中最重要的是两组映射彼此独立:
1 | PagedAttention:逻辑块序号 → 本次运行的 GPU Block ID |
二者通过 MooncakeStoreConnector 衔接。Connector 把 GPU 中已经计算好的 KV Block 保存到外部 Store,也把外部命中的 KV Block 装载到 vLLM 新分配的 GPU Block。
| 对比项 | PagedAttention | Mooncake Store |
|---|---|---|
| 管理范围 | 单个推理实例的 GPU 显存 | 跨请求、实例、节点的 DRAM/SSD 缓存池 |
| 逻辑标识 | 请求的第几个逻辑 Token Block | Block Hash / Mooncake Key |
| 物理位置 | GPU KV Tensor 的 Block ID | MEMORY、LOCAL_DISK 等副本位置 |
| 映射结构 | GPU Block Table | Master 中的对象与副本元数据 |
| 直接服务对象 | Attention Kernel | KV Connector 与 Transfer Engine |
| 主要目标 | 减少显存碎片并复用 GPU Block | 扩大可复用 KV Cache 的容量与范围 |
PagedAttention 管理显存
放入依据
一段 KV Cache 只要要参与当前或下一次 Attention 计算,就必须位于 GPU 显存中。vLLM Scheduler 为请求分配 GPU Block,把物理 Block ID 写入 Block Table;Attention Kernel 再根据表中的 ID 读取对应 K/V。
GPU 中还可能暂时保留已经完成计算、具有 Prefix Cache Hash、但当前没有活动请求引用的 KV Block。这些块能够继续提供本地前缀命中,但也已经是可回收候选。
何时回收
GPU Block 是否可回收首先取决于引用计数:
- **
ref_cnt > 0**:仍被活动请求使用,不能复用。 - **
ref_cnt == 0**:不再被请求引用,可以进入 Free Block Queue。
新请求需要更多 GPU Block 时,vLLM 从 Free Block Queue 取出候选块并复用物理空间。带 Prefix Cache Hash 的缓存块采用 FIFO 复用顺序形成近似 LRU 效果;没有缓存身份的块倾向于 LIFO 复用,以改善局部性。
这里的“GPU 淘汰”不等于一定执行 GPU → DRAM 写回:
- 如果 Mooncake 已经保存了外部副本,GPU 可以直接复用该物理块,未来仍能从 Store 恢复数据。
- 如果没有外部副本,旧数据会消失;后续相同前缀只能重新计算。
Mooncake 管理外部副本
Mooncake Store 不管理 vLLM 当前请求的 Block Table。它看到的是一个由 Block Hash 或 Key 标识的对象,以及这个对象对应的一个或多个副本:
1 | Block Hash / Key |
Master 管理键、位置、副本状态和生命周期,正常 KV 载荷由 Client 与 Transfer Engine 搬运,不经过 Master 中转。
DRAM 放置与淘汰
KV Block 由 Connector 保存到 Mooncake 时,Store 从可用的 MEMORY Segment 中分配空间。这个初始放置主要依据可用 Segment、分配策略和副本配置,不是先预测冷热再决定是否进入 DRAM。
对象被读取或查询时,Mooncake 更新其 lease_timeout。后台内存淘汰通常在两种情况下触发:
- MEMORY 使用率超过高水位,默认是 90%。
- 新对象分配失败,设置内存淘汰标志。
候选对象还必须满足:
- Lease 已过期;
- 没有 Hard Pin;
- 存在完整、可读的 MEMORY 副本;
- 副本引用计数为 0;
- 第一轮优先跳过 Soft Pin,必要时第二轮再考虑。
Mooncake 比较对象的 lease_timeout,截止时间越早,表示越久没有访问,再用 nth_element 找到淘汰边界。因此这是对象级 Lease 近似 LRU,不是经典双向链表 LRU。
SSD 下沉与淘汰
只有启用 SSD Offload 后,KV Block 才会形成 LOCAL_DISK 副本。存在两种下沉时机:
- 立即下沉:
PutEnd成功后,将完整 MEMORY 副本加入 SSD 下沉队列。因此 SSD 中不一定只有冷数据。 - 淘汰时下沉:DRAM 准备淘汰对象时,先保留并固定一个 MEMORY 副本,写入 SSD;成功生成
LOCAL_DISK副本后,再释放 MEMORY。
在淘汰时下沉模式中,如果 Offload Queue 暂时不可用,默认行为是跳过本轮并保留数据;只有显式启用强制淘汰,才会在没有 SSD 副本时直接释放 MEMORY。
SSD 容量达到限制或高水位时,FileStorage 还会执行自己的淘汰:
- **默认
fifo**:优先删除最早创建的 Bucket。 - **可选
lru**:优先删除最久没有读取的 Bucket。 - 淘汰粒度是 Bucket:一个 Bucket 中的多个 KV 对象会一起删除,而不是精确淘汰单个 GPU 风格的 KV Page。
删除 SSD 数据前,FileStorage 先通知 Master 移除 LOCAL_DISK 副本元数据;通知失败则恢复本地索引。Master 确认后,FileStorage 等待正在进行的读取结束,再删除元数据文件和 Bucket 数据文件。
一次外部命中如何完成
当新请求到达时,Mooncake Store 与 PagedAttention 的协作顺序可以概括为:
- 计算内容身份。 按 Token Block 计算 Block Hash。
- 检查本地 GPU Cache。 如果已有可用 GPU Block,直接建立或复用 Block Table 映射。
- 查询 Mooncake。 本地未命中时,根据 Block Hash 查询 Master 中的副本元数据。
- 分配 GPU Block。 外部 DRAM/SSD 命中后,vLLM 仍要先申请新的 GPU 物理块。
- 传输 KV 数据。 Connector 与 Transfer Engine 将命中副本写入这些 GPU Block。
- 更新 Block Table。 将逻辑块映射到本次分配的 GPU Block ID。
- 执行 Attention。 Kernel 此时才能从显存读取 K/V。
1 | 新请求 |
为什么两者如此相似
Mooncake Store 与 PagedAttention 都采用四个共同思想:
- 固定粒度切块。 不要求为整个请求分配一块连续空间。
- 逻辑与物理解耦。 逻辑顺序不要求物理位置连续。
- 间接寻址。 先查询映射,再找到真实数据。
- 按块复用与淘汰。 用有限容量保留更可能再次使用的数据。
但二者位于不同层级:
- PagedAttention 是执行期页式管理。 它解决 GPU 内 KV Block 的分配、映射、碎片与复用。
- Mooncake Store 是跨请求的外部缓存。 它解决 KV Block 如何按内容标识、保存副本、跨节点传输并在 DRAM/SSD 间管理生命周期。
从概念上可以把 Mooncake 看成 PagedAttention 外面的 L2/L3 缓存系统,但这不是说 Mooncake 扩展了同一张页表。它们拥有各自的元数据和淘汰策略,中间通过 Connector 显式搬运数据。
三套策略不要混用
| 层级 | 放置内容 | 放入依据 | 回收时机 | 策略 |
|---|---|---|---|---|
| GPU Block Table | 逻辑块到 GPU Block ID 的映射 | 当前调度批次与 GPU 分配结果 | 请求结束或槽位复用 | 重新构造,不下沉 SSD |
| GPU KV 数据页 | 活动请求与本地 Prefix Cache | Scheduler 分配、外部缓存装载 | ref_cnt=0 且需要新块 |
vLLM Free Queue,缓存块近似 LRU |
| Mooncake DRAM | 可跨请求、跨实例复用的 KV 副本 | Connector 保存,Store 分配 MEMORY Segment | 超过高水位或分配失败 | Lease 近似 LRU |
| Mooncake SSD | MEMORY 的下层副本 | 立即下沉或 DRAM 淘汰时下沉 | SSD 超容量或高水位 | 默认 FIFO,可选 Bucket LRU |
最容易出现的误解包括:
- 把页表当成 KV 数据。 页表很小,真正占据容量的是各层 K/V Tensor。
- 认为页表可以直接指向 SSD。 Attention 必须先看到 GPU Block,外部数据要显式装载。
- 认为存在统一的三级 LRU。 GPU、Mooncake DRAM 与 SSD 分别由不同组件管理。
- 认为 GPU 回收必然写回。 只有外部副本已经保存,数据才能在未来恢复;否则只能重算。
总结
PagedAttention 与 Mooncake Store 看起来相似,是因为它们都把连续的 KV Cache 切成块,并用间接映射代替固定物理位置。区别在于优化范围:PagedAttention 管理当前计算所需的 GPU 物理块,Mooncake Store 管理暂不在 GPU 中但值得跨请求复用的 DRAM/SSD 副本。
完整的数据生命周期不是一张页表在三级介质间移动,而是:GPU Block Table 持续描述当前执行地址;Connector 根据 Block Hash 查询外部副本;命中后将数据装入新 GPU Block,再重建执行映射。理解这一层次,才能正确区分显存复用、DRAM 淘汰、SSD 下沉和 SSD 自身淘汰。
参考资料
- Mooncake 开源仓库,本文涉及的 Mooncake 源码语义固定于
bfca1ce2af8419c50dc8d464820a95d97d43c930。 - Mooncake Store 设计。
- Mooncake 内存 near-LRU 定义。
- Mooncake 内存候选筛选。
- Mooncake SSD Bucket LRU。
- vLLM GPU Block Table。
- vLLM GPU BlockPool。
- vLLM MooncakeStoreConnector。