返回文章列表

LMCache 视角下的 KV Cache:对象模型、数据路径与一致性

在推理引擎内部,KV Cache 有明确的物理归属:它位于某个设备,由具体的 tensor、block table 和计算 stream 管理。这些信息只在当前进程和当前实例中成立。LMCache 希望把 KV 留到请求结束之后,供其他请求或其他实例复用。原来的物理坐标随进程边界失效,缓存必须重新获得一套稳定的身份和生命周期。

Admin
16 分钟阅读
8 次阅读
最后编辑于 2026-08-31 08:20
LMCache 视角下的 KV Cache:对象模型、数据路径与一致性

LMCache 对 KV Cache 的处理,起点不是存储容量,而是状态有效性。

在推理引擎内部,KV Cache 有明确的物理归属:它位于某个设备,由具体的 tensor、block table 和计算 stream 管理。这些信息只在当前进程和当前实例中成立。LMCache 希望把 KV 留到请求结束之后,供其他请求或其他实例复用。原来的物理坐标随进程边界失效,缓存必须重新获得一套稳定的身份和生命周期。

KV Cache 是一段计算历史

在标准自回归 Transformer 中,第 l 层、第 t 个 token 会产生 Query、Key 和 Value:

qt(l)=xt(l)WQ(l),kt(l)=xt(l)WK(l),vt(l)=xt(l)WV(l)q_t^{(l)}=x_t^{(l)}W_Q^{(l)},\qquad k_t^{(l)}=x_t^{(l)}W_K^{(l)},\qquad v_t^{(l)}=x_t^{(l)}W_V^{(l)}

生成下一个 token 时,当前 Query 只参与本轮计算;历史 Key 和 Value 会被后续 attention 反复读取。推理引擎保存 K/V,由此避免在 decode 阶段重新生成整段历史状态。

对于普通 MHA 或 GQA,在未考虑并行分片的情况下,每个 token 的 KV 容量近似为:

Stoken=L×2×Hkv×Dhead×BS_{\mathrm{token}} = L\times 2\times H_{kv}\times D_{head}\times B

L 是层数,Hkv 是 KV head 数,Dhead 是 head dimension,B 是单个元素的字节数。以 32 层、8 个 KV head、head dimension 128、BF16 为例,结果是 128 KiB/token;256 个 token 约 32 MiB。

这个估算只适用于上述假设。TP 可能分片或复制部分 KV,DCP 沿 token 维度切分,MLA 保存压缩 latent,recurrent 层维护的状态也不是传统 K/V。LMCache 必须接受这些差异,所以源码里同时存在多种 memory format、layout descriptor 和 object group。

跨请求复用的目标是省掉已命中前缀的 prefill。它不会消除 decode 阶段对历史 K/V 的 attention 读取。要跳过前 p 个 token,推理引擎必须先拿到这段前缀在所有必要层上的状态,把它们写入本次请求分配的正确位置,并在 forward 开始前完成同步。LMCache 保存的因此不是一批可以任意拼接的 tensor,而是特定模型、特定 token 历史和特定并行切片下的计算结果。

从 token 序列得到稳定身份

MP 模式默认以 256 个 token 为一个 chunk,只为完整 chunk 生成存储对象。尾部不足一个 chunk 的部分不会进入普通 MP 缓存路径,仍由推理引擎计算。

第 i 个 chunk 使用滚动前缀哈希:

hi=H(hi−1∥tokens[iC:(i+1)C])h_i = H\left( h_{i-1}\parallel tokens[iC:(i+1)C] \right)

默认 BLAKE3 实现先写入前一 chunk 的哈希,再把当前 token IDs 按大端 uint32 序列写入,最终得到 32 字节 digest。计算某个后段 chunk 时,前面的哈希仍需按顺序生成,因为当前键依赖完整前缀。

绑定前缀不是为了提高哈希质量。它处理的是模型状态依赖。同一段 token 出现在两个不同前缀之后,第一层投影可能只受 token 与位置影响,后续层的输入却已经包含上下文信息。LMCache 保存多层状态,当前 chunk 的内容相同不足以证明整组 KV 相同。

哈希随后被装入 ObjectKey。MP 版本的键还包含 model_name、kv_rank、object_group_id 和 cache_salt。模型名区分权重与配置,KV rank 标识并行切片,object group 区分独立存储的状态组。cache salt 负责命名空间隔离;它不提供加密。

客户端在 ZMQ 请求中发送的是 IPCCacheServerKey。其中保留完整 token IDs、start/end、worker 信息和 request ID。server 计算 chunk hashes 后再生成存储键。request ID 只用于 session 跟踪,不参与缓存身份比较。这个区分很重要:请求会结束,缓存对象不能以请求编号命名。

逻辑对象与物理布局分开管理

推理引擎通常使用 paged KV。逻辑上连续的 token 可以落在不连续的 GPU blocks 中,block ID 也会随着请求分配而变化。它适合描述本次执行,不适合做持久化键。

LMCache 在注册阶段接收 GPU KV tensor、shape、dtype、memory format、group metadata 和 IPC event。注册只建立访问关系,不把整块 GPU KV 复制到 server。之后的 store 或 retrieve 请求才携带当前 block IDs。

存储身份和物理解释被分成两部分。ObjectKey 回答对象属于谁;MemoryLayoutDesc 记录对象由哪些 shape 和 dtype 组成;MemoryObj 管理实际地址、大小、视图、引用和 pin 状态。block table 只出现在传输计划里,由 gather/scatter kernel 在 GPU 分页布局和 LMCache 连续 chunk 之间转换。

混合注意力模型还需要区分几种容易混淆的 group。Engine KV group 服从推理引擎的 block 分配规则。Kernel group 聚合能够共用传输 kernel 的层,其分组身份考虑 KV size、head 数、head dimension、block size、dtype 和 KV format。Object group 决定存储键与分配单元。三者可能重合,也可能不同。

当前 MP 的 ObjectKey 不直接包含 dtype 或 shape;layout 由注册表提供,注册表以 (model_name, world_size) 为索引。这减少了键对具体布局的耦合,同时要求同一索引下的活跃实例使用兼容布局。它是当前实现的部署前提。

命中长度不是对象计数

读 LMCache lookup 代码时,最容易误解的地方是把它当作普通 key-value 查询。后端找到一个对象,并不等于推理引擎可以多跳过一个 chunk。

假设请求有 N 个 chunks、G 个 object groups,每个 group 需要 W 个 KV ranks。server 最多要检查 N×G×W 个 ObjectKey。对第 c 个 chunk,模型级命中可以写成:

Hit(c)=⋀g∈Required(c) ⋀r∈Ranks(g)Present(c,g,r)Hit(c) = \bigwedge_{g\in Required(c)} \ \bigwedge_{r\in Ranks(g)} Present(c,g,r)

普通全注意力模型可复用的前缀是满足下式的最大 k:

PrefixHit=max⁡{k | Hit(c)=1, ∀c<k}PrefixHit = \max\left\{ k\ \middle|\ Hit(c)=1,\ \forall c<k \right\}

第 3 个 chunk 缺失时,即使第 4 个 chunk 仍在后端,普通 lookup 也只能返回前两个。推理引擎的 prefix-skip 接口要求连续前缀;深层状态的计算关系同样要求此前状态完整。存在性位图必须先折叠成模型可用的 token 前缀。

Recurrent state 不能跨缺口续接;MLA 与 DCP 会改变每个 chunk 所需的 rank fan-out。由此得到的命中率是 token 级、模型级结果,不能用 L1 或 L2 中的对象数量代替。

MP 架构中的控制面与数据面

推理引擎和 LMCache server 分进程运行,ZMQ 负责请求协议,GPU 数据通过 IPC、共享内存或 engine-driven 传输路径移动。

scheduler 掌握 token 进度和 block 分配,因此由它询问可复用前缀。worker 持有 GPU tensor 与计算 stream,实际传输必须经过 worker 对应的 cache context。LMCache server 管理逻辑对象、L1/L2 位置和锁。Coordinator 维护 fleet 元数据,不承载 KV 字节。

请求首先经过 lookup。Scheduler 发送 token IDs、model、world size、cache salt 和 reader count。server 计算完整 chunks 的滚动哈希,展开 object groups 与 ranks,在 L1 上获取读锁,再把缺失后缀交给 L2 prefetch。L2 数据进入预留的 L1 buffer 后,写锁原子转换为读锁。scheduler 收到最终命中 token 数,随后为这些 token 分配 GPU blocks。

retrieve 由 worker 发起。请求带上当前 block IDs、IPC event,以及与 vLLM Automatic Prefix Cache 重叠的 token 数。server 等待 producer event,只读取 lookup 阶段已经锁定的对象,再按 object group 执行 H2D scatter。copy 入队后记录完成 event,读锁由 stream callback 释放。APC 已覆盖的位置会被跳过,避免 LMCache stream 写入其他请求可能正在读取的 shared blocks。

prefill 完成后才可能 store。适配器用请求长度、各 engine groups 的实际 block 覆盖和已计算 token 数共同限制可写范围,并向下取整到完整 LMCache chunks。server 为新键 reserve_write,从 GPU paged KV gather 到连续 MemoryObj。D2H 完成后 finish_write,对象才进入可读状态。若配置 L2,StoreController 在持有 L1 读锁时异步写后端。

vLLM 自身的 prefix hit 与 LMCache hit 可能重叠,metadata 不会把两者简单相加。全 prompt 命中时,connector 还可能留下最后一个 token 执行 forward,以满足 logits 生成接口。这些细节属于推理引擎和外部缓存之间的契约,不是存储后端能够自行决定的行为。

一致性依赖锁和 stream

L1 对象的状态机可以简化为:

exec-f5e9c5ee-bd50-40ad-9a31-7e5a3f868da1.png

write lock 防止 lookup 读到尚未完成的 DMA。read lock 防止 H2D、L2 store 或 P2P 读取期间发生 eviction。PrefetchController 从 L2 加载对象时使用 finish_write_and_reserve_read,在同一次原子操作中结束写锁并取得读锁;若分成两个调用,中间的 ready 状态可能被 eviction 线程观察到。

L2 adapter 的 lookup_and_lock 把存在性检查与后端 pin 合并。load 把数据写入调用方提供的 L1 buffer,完成后再 unlock。若 lookup 与 pin 分离,对象可能在查询成功和实际读取之间被远端驱逐。

GPU event 解决的是另一类时间关系。server 需要等待推理引擎完成生产,retrieve 结束后也不能立即释放 L1 对象;只有消费该对象的 GPU stream 越过 callback,read lock 才能归还。TTL 负责异常兜底,正常路径仍依赖明确的 completion 与 cleanup。

发生缺键、布局不一致或传输失败时,connector 回退到 recompute。对 KV 缓存而言,这是合理的失败方式。多算一次会增加延迟;使用错误或不完整的 KV 会改变模型输出,而且通常没有显式错误信号。

分层存储何时值得使用

LMCache MP 正式管理的近端层是 L1,可使用 pinned CPU DRAM、Device-DAX 或 GDS 注册的 NVMe slab。L2 adapter 连接文件系统、对象存储和分布式后端。P2P adapter 可以把另一节点的 L1 视为只读远端层,通过 NIXL/RDMA 拉取。Serde wrapper 位于 L2 边界,可执行压缩、量化或 AES-GCM 加密。

这些后端共享对象协议,不代表访问成本相同。一次复用是否节省时间,取决于:

Tlookup+Tfetch+TH2D+Tsync<Tprefill savedT_{\mathrm{lookup}} + T_{\mathrm{fetch}} + T_{\mathrm{H2D}} + T_{\mathrm{sync}} < T_{\mathrm{prefill\ saved}}

对象已在 L1 时,fetch 成本较低;跨节点或对象存储读取时,传输可能占主导。压缩减少容量和网络字节,却增加编解码时间;有损格式还要单独验证生成质量。

chunk size 也进入同一个成本模型。较大的 chunk 减少哈希、元数据、锁和 I/O 请求数,但会放大尾部浪费与无效搬运。较小的 chunk 有更细的复用和 eviction 粒度,控制面开销随之增加。默认 256 是当前配置默认值,不能代替具体负载上的测量。

因此,LMCache 的性能评价至少需要同时观察命中前缀分布、各层读取延迟、实际传输字节、H2D 时间和被省掉的 prefill。单独报告对象命中率,无法说明请求是否加速。

CacheBlend 处理另一类复用

滚动前缀哈希使普通 prefix reuse 保持上下文一致,也使相同内容在不同前缀下得到不同键。CacheBlend 增加内容指纹,在请求的非前缀位置寻找候选 chunk,并通过 sparse prefetch 加载。

候选内容相同,深层 KV 仍可能不同。K 包含 RoPE 位置变换,chunk 移动后需要 re-RoPE;V 不直接应用 RoPE,但后续层的 K/V 已受新上下文影响。CacheBlend 会复用候选状态,并选择部分 token 重新计算以修复上下文差异。

普通 prefix path 的前提是身份、布局和依赖完整,可以恢复连续前缀。CacheBlend 接受非前缀候选,再用选择性重算换取可用质量。两条路径的命中含义不同,不宜合并成一个“KV 去重率”。

LMCache 实际管理的对象

从源码抽象看,LMCache 管理的 KV 对象可以写成:

KVObject=(Identity, Layout, Dependency, Residency, Ownership)KVObject = ( Identity,\ Layout,\ Dependency,\ Residency,\ Ownership )

Identity 对应 token 历史、模型和 rank;Layout 描述对象字节;Dependency 决定对象集合能否组成可跳算前缀;Residency 记录当前存储层;Ownership 由锁、session 和 event 维护。

推理引擎内建 prefix cache 更靠近 scheduler 与 GPU block allocator,命中后通常不需要跨设备搬运。LMCache 扩大了缓存的容量、寿命和共享范围,也增加了寻址、传输与一致性成本。两者可以重叠,适配层必须明确计算覆盖范围和写入边界。

发表评论

未登录评论需提供昵称和联系邮箱,提交后默认进入审核。

0/1000