Kimi Linear 的 KDA 缓存:SGLang、vLLM 与 Mooncake Store 全链路

从 Kimi Linear 的 KDA 状态出发,追踪 SGLang、vLLM 与 Mooncake Store 的混合缓存流程,并解释细粒度检查点、写时复制、部分分组命中与远端复用。

Code walkthroughsMooncake and HiCache internalsLLM servingSGLangvllmkimi-linear

Kimi Linear 的前缀缓存不能只按普通 Transformer 的 KV 块来理解。模型里一部分层是 MLA,另一部分层是 KDA:前者需要一串可分页的 KV 或低维缓存页,后者需要“前缀处理到当前位置以后”的有限状态。工程难点不只是保存数据,而是让 token 前缀、MLA 缓存页、KDA 检查点、调度步、HiCache 层级和 Mooncake Store 对象始终落在同一个逻辑边界。

本文沿两条实现链路展开:

text
Kimi KDA
  ├─ SGLang:KDA 内核 → UnifiedRadixCache → HiCache 主机层
  │            → MooncakeStore 后端 → Mooncake Distributed Store
  │
  └─ vLLM: KimiGatedDeltaNetAttention → HybridKVCacheCoordinator
               → MooncakeStoreConnector → Mooncake Distributed Store

核心结论是:Mooncake Store 负责保存、放置和搬运不透明字节;KDA 检查点在哪个 token 边界产生、是否与 MLA 前缀对齐、缺失时是否补算,仍由 SGLang 或 vLLM 的控制面决定。

两套引擎采用不同的检查点创建方式:

text
SGLang
  一次较长的扩展计算可以让 KDA 内核返回按块对齐的中间状态
  → 推理引擎从中抽取目标边界
  → 基数树节点持有 KDA/Mamba 槽位
  → HiCache 将递推状态和卷积状态写入 Mooncake

vLLM
  KDA 预填充通常只提交当前调度步的最终状态
  → 混合缓存协调器发现 MLA/KDA 命中分歧
  → 调度器在真正需要的边界提前结束前向计算
  → 将最终状态注册为前缀检查点

SGLang 已有显式的递推状态和卷积状态远端读写链路;vLLM 的原始状态页与缓存分组结构也能承载完整 KDA 检查点。通用 Mooncake 传输基础设施已经存在,但 Kimi KDA + MooncakeStoreConnector 的专项调用链证据仍有限。两套实现共同尚未完全闭合的问题是部分分组命中:当 MLA/KV 命中得更长、同一边界的 KDA 状态却不存在时,推理引擎能否发现缺失边界,并只补算一次检查点。

源码版本与结论范围

正文统一固定到 2026-07-22 20:00(北京时间,UTC+8;即 12:00 UTC)。下面列的是各仓库 main 在该时刻之前的最后一个提交:

text
SGLang   4f88206393b620f5a856c8719289cf845a72b4fb  2026-07-22 19:45:29 UTC+8
vLLM     1a659a0c370942bae3c8567902e38bddbbae95f3  2026-07-22 19:52:56 UTC+8
Mooncake 12df58c7de65f88f37f5da99ab95ff3418a9ebef  2026-07-22 19:02:55 UTC+8

对应提交分别是 SGLangvLLMMooncake。KDA 数学以 Moonshot AI 的 Kimi Linear 技术报告(arXiv:2510.26692v2)为准;该版本于 2025-11-01 修订,运行时结论只以这里固定的三个代码快照为准。

最后修订:2026-07-24

下文区分三类判断:

text
源码调用链可确认      可以从入口追到状态读写和测试
接口形状允许          数据结构可以承载,但缺少完整专项调用链证据
调用链仍有缺口        控制面信息未传播,或远端部分命中未形成闭环

本文不做性能评测,也不把通用 Mamba/GDN 测试直接等价为 Kimi KDA 的完整生产调用链。文中的“当前实现”均指上述固定快照,而不是持续变化的最新 main

先统一记号:矩阵状态不等于完整检查点

记分词器产生的前 t 个 token ID 所组成的前缀为:

Pt=(τ1,τ2,,τt),τi{0,,|V|1}.

这里的 τi 是 token ID。后面 KDA 公式中的 xt(l) 才表示第 l 层当前位置的归一化隐藏状态,二者不能混为一谈。

对第 l 个 KDA layer,定义:

Sl(t)Rnh,llocal×dk,l×dv,l

为处理完 P_t 后的递推矩阵。再定义同一边界上的 ShortConv 窗口:

Γl(t)R(wl1)×dqkv,llocal.

这里 wl 是因果卷积核大小;卷积窗口的物理实现也可能使用转置布局 [dim, state_len]

本文按数学语义将 Sl(t) 记作 [local_heads, d_k, d_v]。部分计算内核和状态池的物理布局使用 [local_heads, d_v, d_k],即最后两个维度转置存放;两种布局表示同一个关联状态。

一个 KDA layer 的完整续算状态是:

Cl(t)=(Γl(t),Sl(t)).

整个模型在前缀 P_t 上的 KDA 检查点是:

C(Pt)={Cl(t)lKDA layers}.

后文中的 Sl(t) 始终只表示递推矩阵;Cl(t) 才表示某一层完整的卷积与递推状态;C(Pt) 则表示所有 KDA 层在同一 token 边界上的检查点。

MLA 缓存页不属于 C(P_t)。对于 Kimi Linear 这样的混合模型,一个可以直接恢复的前缀还必须同时拥有:

text
MLA / 全局注意力缓存 [0, t)
+
KDA 检查点 C(P_t)
+
其它必要的辅助状态与所有并行 rank 数据

KDA 的递推与充分状态

对单个注意力头,KDA 维护固定大小的关联记忆矩阵:

StRdk×dv.

当前位置的键、值和查询分别是:

kt,qtRdk,vtRdv.

KDA 先按键通道衰减旧状态,再计算旧记忆对当前值的预测误差,最后只写入增量:

St=DtSt1,v^t=Stkt,et=vtv^t,St=St+βtktet,ot=Stqt.

其中:

Dt=Diag(αt),αt(0,1]dk,βt[0,1].

等价地:

St=(Iβtktkt)DtSt1+βtktvt.

与普通线性注意力的直接外积累加相比,Delta Rule 会先判断旧状态能否重建当前值,只有预测错误的部分才被写回。KDA 又把 Gated DeltaNet 中每个头一个标量遗忘率,细化成每个键通道一个遗忘率。

Q/K/V 并不是直接进入递推计算。对第 l 层的输入 xt(l),实际模型会先做线性投影、因果短卷积和激活,Q/K 还会做 L2Norm:

qt=L2Norm(SiLU(ShortConv(Wqx(l))t)),kt=L2Norm(SiLU(ShortConv(Wkx(l))t)),vt=SiLU(ShortConv(Wvx(l))t).

因此,只有递推矩阵而没有 ShortConv 窗口,不足以继续计算下一个 token。

不需要跨请求持久化的量包括:

text
完整前缀的 q/k/v
逐 token 的 alpha / beta / delta 误差
预填充块内的 WY / UT 临时空间
注意力输出与输出门控的中间量

它们要么已经被折叠进 Sl(t),要么在当前计算内核结束后失效。

这也解释了前缀复用的单向性:C(P_1k) 可以随新增后缀向前推进成 C(P_2k),但只有 C(P_4k) 时无法倒推出 C(P_1k)。递推中的衰减和误差修正是有损更新,较早边界的 ShortConv 窗口也已经被覆盖。

预填充为什么能分块并行

逐 token 递推适合解码,但预填充如果每个 token 启动一次计算内核,GPU 利用率会很差。KDA 将一个计算块内的多次状态转移压缩成:

Sr=PrS0+Hr,Pr=TrTr1T1.

其中每次状态转移都可以写成“对角矩阵减去秩一矩阵”:

Tt=Dtβtkt(ktαt).

KDA 利用 WY 表示和 UT 变换,将块内的大部分工作转成矩阵乘法与小型下三角求解。论文与 SGLang Triton 实现中常见的 chunk_size=64 是算子内部的计算单元,不是调度器的分块预填充大小,也不是前缀检查点间隔。

五种容易混淆的“粒度”

五种相互独立的粒度

同一条前缀缓存链路里,至少有五种不同边界

它们分别服务于算子、查找、物理分配、调度和远端存储。能够形成哈希,不等于已经生成检查点,也不等于可以成为 Mooncake 对象边界。

KDA 计算块

算子内部的并行单元

中间状态 h按块对齐的状态

prefix_match_unit / 哈希块

滚动哈希与本地查找粒度

可寻址边界不自动生成状态

缓存分组的物理块

GPU 分配器与块表的管理单位

MLA 缓存页KDA 原始状态页

调度器 / 最小公倍数块

多个缓存分组与连接器的共同对齐单位

查找对齐存储对齐

Mooncake 对象块

一个键对应一个或多个内存切片

不透明字节分散聚集传输

一个逻辑检查点如何展开

逻辑上,一个检查点由前缀哈希、精确 token 长度和 C(P_t) 组成;物理上,它由所有 KDA 层的状态共同构成。

vLLM 在固定快照中把每层状态放进以块为首维的原始状态页:

text
每个 KDA 层一页原始状态
  = 卷积状态字节
  + 递推状态字节
  + 可选填充

同一个逻辑块编号会在每个 KDA 层的页表中索引出该层的 C_l(t)

SGLang 则让一个逻辑 Mamba 槽位横跨所有 KDA 层:

text
递推状态缓冲区
  [num_kda_layers, slots, local_heads, d_v, d_k]

卷积状态缓冲区
  [num_kda_layers, slots, kernel_size - 1, qkv_width]

所以:

text
vLLM
  C(P_t)
  → 每个 KDA 层各有一页原始状态

SGLang
  C(P_t)
  → 一个跨层 Mamba 槽位
  → 递推状态 + 一个或多个卷积缓冲区

逻辑原子性不要求物理上只有一个张量或对象。真正的约束是:所有组成部分必须属于同一个前缀哈希和 token 边界,并在查找时共同验证。

Kimi K3 预览版与 vLLM 的细粒度检查点

截至该代码快照,vLLM 的 PR #46384 已合入;同日 vLLM 发布了 Kimi K3 推理服务预览。这批变化进一步明确了三个必须分开的概念:

text
物理状态块大小
  管理状态页的物理分配与块表

prefix_match_unit
  管理滚动哈希与本地查找的最细 token 边界

实际检查点集合
  只包含真正物化并注册哈希的 C(P_t)

下面的 4096 / 64 / 6976 只是解释机制的示意值,不是 K3 的默认配置:

text
物理 KDA 块 = 4096
prefix_match_unit  = 64

vLLM 的本地前缀缓存可以在物理块 [4096, 8192) 内注册:

text
Hash(P6976) → C(P6976)

一个物理状态槽位只承载一个精确边界,并不会同时保存 S4160、S4224、...、S6976。一个 6976-token 请求可能保留物理边界 C(P4096) 和不满整块的尾部 C(P6976),但不会因为匹配粒度是 64 就自动生成 109 份状态。

配置的匹配粒度必须整除每个缓存分组的物理块大小;若未显式配置,vLLM 可以从各分组块大小的公约数推导匹配粒度。

当另一条请求命中 P6976 并继续 64 token 时:

text
已缓存 C(P6976) --写时复制--> 请求私有工作状态
                                      |
                                      +-- 向前计算 64 个 token --> C(P7040)

共享的 C(P6976) 不会被原地更新成 C(P7040)。如果产生该检查点的原请求继续执行,vLLM 会把 Hash(P6976) → C(P6976) 迁移到缓存持有的副本,让原工作块留在请求侧继续写。这是原持有者继续执行时的反向写时复制。

许多不同长度的请求理论上可以沿同一条前缀链逐渐贡献检查点:

text
Hash(P64)   → C(P64)
Hash(P128)  → C(P128)
...
Hash(P8192) → C(P8192)

但这是多个请求形成的唯一 (前缀哈希, token 长度) 检查点集合,不是单个请求自动保存所有中间状态。相同前缀会去重,实际驻留数量还受创建策略、保留策略、状态池容量和 LRU 淘汰限制。

检查点的创建、匹配和保留是三件事

定义实际驻留的 KDA 检查点集合:

CKDA={(HL,L)C(PL) 已物化、已注册且仍驻留},

其中 HL=Hash(PL)

三类机制分别负责:

text
匹配粒度
  哪些 L 可以形成滚动哈希并参与查找

创建策略
  哪些 L 真正生成了 C(P_L)
  例如调度步结束、算子快照或重算

保留策略
  已生成的 C(P_L) 中,哪些继续留在状态池

prefix_match_unit=64 只表示每 64 个 token 可以形成哈希边界,不表示一次较长的前向计算会自动生成所有 64-token 边界的 KDA 状态。保留策略也只能筛选已经生成的检查点,不能从 C(P_4k) 反向创造 C(P_1k)

更小的匹配粒度也不是完全没有成本:它不会自动分配更多递推状态页,但会增加块哈希、基数树与查找元数据,以及候选尾部边界数量,从而扩大可以被实际物化的检查点集合。

对于固定快照中的 MooncakeStoreCoordinator,还存在物理对齐约束:每个缓存分组的块大小必须能被哈希块大小整除,调度块大小又必须是各缓存分组块大小的整数倍。

检查点状态池的内存成本

使用当前 TP/PP 分片的局部维度,单个 KDA 层的递推状态字节数可写成:

Brec,l=nh,llocaldk,ldv,lbrec,l.

ShortConv window 的字节数近似为:

Bconv,l=(wl1)dqkv,llocalbconv,l.

一份完整模型检查点的逻辑大小是:

Blogical=lKDA(Brec,l+Bconv,l).

若状态池中驻留 N 个唯一前缀检查点:

Mlogical poolNBlogical.

对 vLLM 这类“每层一个原始状态页”的布局,完整物理检查点的大小应按所有相关层的物理页求和:

Bphysical checkpoint=lKDABphysical page,lBlogical.

每层物理页还可能包含分配器填充、页对齐和统一页大小带来的尾部空间。写时复制的峰值显存与总复制字节数应按完整物理检查点估算,而不是只按数学状态估算。固定 vLLM 快照中,KDA 的卷积状态使用配置的缓存数据类型,而递推状态使用 FP32。

SGLang:从一次较长的 extend 调用中提取检查点

SGLang 需要追踪前缀状态时,会让 KDA 内核返回中间递归状态:

python
core_attn_out = kernel_dispatcher.extend(
    ...,
    return_intermediate_states=track_ssm,
)

if track_ssm:
    core_attn_out, h = core_attn_out
    _track_mamba_state_extend(..., h, ssm_states, ...)

固定源码入口包括:

检查点同时包含两部分:

text
卷积检查点
  从 mixed_qkv 中选取边界前最近 kernel_size - 1 行

递推检查点
  若边界等于扩展阶段末尾,复制最终状态
  若边界位于扩展阶段中间,从计算内核返回的 h 中抽取

但这并不意味着可以在一次较长的 extend 调用中提取任意 token 位置的状态。固定实现先把目标长度向下对齐到 mamba_cache_chunk_size

Ltrack=Lprefix+LtargetLprefixCkernelCkernel.
  • 目标边界与 KDA 分块对齐时,直接复制最终递推状态。
  • 目标边界位于扩展阶段中间时,从 h 中取最近的分块边界状态。
  • 源码仍有 mamba_cache_chunk_size % page_size != 0 的待办项,因此内核分块大小与页大小不能任意组合。

如果配置了 FlashKDA,但推理引擎请求 return_intermediate_states=True,固定 SGLang 包装层会回退到 Triton chunk_kda。FlashKDA 负责高性能的最终状态预填充;中间检查点追踪依赖能够返回 h 的 Triton 路径。

所以 SGLang 不要求为了生成 1k 检查点,就把 4k 预填充拆成四次模型前向计算;但实际抽取的是合法的分块和分页对齐边界,而不是任意 token。

检查点怎样进入 UnifiedRadixCache

完成一次前向计算后,基数树节点不直接保存整套张量,而是保存状态池槽位:

text
UnifiedTreeNode
  token 键 / 页面哈希
  Full component
    MLA/KV device indices
  Mamba component
    value       GPU Mamba slot
    host_value  HiCache 主机内存 Mamba 槽位

MambaComponent.prepare_for_caching_req 使用 req.mamba_last_track_seqlen 作为可提交长度,并把追踪槽位转交给基数树缓存。请求命中时:

text
node.value 存在
  → 检查点仍在 GPU,可写时复制到请求槽位

node.host_value 存在
  → GPU 状态已逐出,但主机副本仍可回载

本地没有 HiCache 控制器时,如果全局注意力或 MLA 前缀比 KDA 状态更长,MambaComponent.finalize_match_result 会产生 mamba_branching_seqlen。下一次扩展计算可以继续按较大预算执行,但后端会从中间 h 抽取这个分叉检查点。

固定实现见 mamba_component.py

HiCache 与 Mooncake 的数据面

对于混合模型,HiCache 会同时迁移全局注意力或 MLA KV,以及 KDA 辅助状态:

text
全局注意力或 MLA 部分
  GPU token/KV 索引 → 主机 KV 索引

KDA 部分
  GPU Mamba 槽位 → 主机 Mamba 槽位

写入 Mooncake 时,SGLang 不上传一串历史 Sl(t),只上传基数树节点末尾的一个 C(P_t)。Mamba 传输使用 TRAILING_PAGES,因为恢复目标前缀只需要末端 KDA 检查点,而全局注意力或 MLA 仍要求前缀缓存页连续存在。

Mooncake 适配层会将一个逻辑 KDA 页展开成:

text
{base_key}_{rank}_temporal
{base_key}_{rank}_conv_0
{base_key}_{rank}_conv_1
...

batch_exists_v2 只有在递推状态和全部卷积状态都存在时,才将该逻辑检查点视为命中;任一部分缺失都必须判定未命中。固定实现见 mooncake_store.py

恢复路径反向执行:

text
Mooncake exists
  → 协调 Full/MLA 连续缓存页与末端 KDA 检查点
  → batch_get_into 主机缓冲区
  → radix node.host_value
  → 通过 load_back / CoW 恢复到请求的 GPU 槽位
  → 从 C(P_t) 继续计算后缀

SGLang + HiCache 仍缺少分支状态回填

数据面已经能完整读写 KDA 检查点,但固定快照中的 HiCache 模式会暂时跳过分支状态回填:

text
Mooncake 已有 MLA(P_1k) + C(P_1k)
  → 可以完整恢复并复用 P_1k

Mooncake 只有 MLA(P_1k),没有 C(P_1k)
  → 混合命中长度必须缩短
  → 当前 HiCache 路径不会仅凭 MLA 命中自动要求
    KDA 前向计算在 P_1k 创建检查点

问题不在 Mooncake 能否搬运递推与卷积状态,而在于谁把缺失边界传给调度器和后端。

vLLM:步末状态与调度拆分

vLLM 的 KDA 预填充路径会调用:

text
chunk_kda_with_fused_gate(..., output_final_state=True)
  → last_recurrent_state
  → 写回当前状态槽位

固定源码入口见 kimi_gdn_linear_attn.py

一次较长的预填充自然只产生该调度步末尾的检查点。vLLM 的 Kimi 路径不像 SGLang 那样暴露前向计算中多个分块的 h,因此缺失边界主要依靠调度器让前向计算恰好停在那里。

Kimi Linear 使用 mamba_cache_mode=align。逻辑块表仍覆盖完整序列,但绝大多数位置可以是空占位符;运行中的请求只需维护少量滚动状态槽位。

text
第 n 个调度步
  前一状态 --复制--> 当前工作状态
  当前工作状态 --KDA 前向计算--> C(P_end_n)

第 n+1 个调度步
  C(P_end_n) --复制--> 下一个工作状态
  更早的工作槽位可以释放

当步末状态被赋予前缀哈希后,它同时成为其它请求可以命中的检查点。KDA 查找会从目标前缀右侧向左寻找最近的状态块,不需要保留从 0 到该边界的所有历史 KDA 状态块。

第一次分叉请求为什么仍需重算

设:

text
req1 = P_4k
req2 = P_1k + A_1k
req3 = P_1k + B_1k

req1 的大段预填充可能留下:

text
MLA 缓存页:P_1k、P_2k、P_3k、P_4k
KDA 状态:只有 C(P_4k)

req2 只共享前 1k,不能使用已经混入后续 token 的 C(P_4k)。在不挂连接器的本地路径中,协调器会得到:

text
MLA 命中长度 = 1k
KDA 命中长度 = 0
最终混合命中长度 = 0
缺失的共享前缀边界 = 1k

调度器将 req2 的首次预填充截到 1k。前向计算结束后生成并注册 C(P_1k),下一步再继续处理 req2 的剩余部分。req3 到来时,MLA 与 KDA 就能同时命中 1k 前缀。

这里重算的是完整模型的前向过程,而不只是 KDA 层。第一次分叉请求需要付出一次补算成本,之后共享同一前缀的请求才能获得收益。

连接器路径丢失了检查点边界

在本文固定的 vLLM 快照中,同样的本地命中分歧在是否挂接连接器时会走不同路径:

text
本地 MLA 命中长度 = 1k
本地 KDA 命中长度 = 0

未启用连接器
  → 按常规方式求混合命中不动点
  → 最终本地命中长度 = 0
  → shared_prefix_boundary = 1k
  → Scheduler 可以在 1k 处暂停并补算 C(P_1k)

启用连接器并使用 Mamba/KDA
  → shared_prefix_boundary = 0
  → num_new_local_computed_tokens = max(per_group_hits) = 1k
  → 本地缺失检查点的分叉边界不再向后传播

固定逻辑见 scheduler.py

这里的代码真实语义是采用 max(per_group_hits),并不是显式选择 Full Attention 分组;只是在本文 MLA hit = 1k, KDA hit = 0 的例子中,最大值恰好等于 MLA/FA 命中长度。

这条特殊路径原本面向 P/D 连接器:全局注意力已有的本地块可以避免重复传输,而 Mamba 状态由对应工作进程另行传递。把同一假设直接用于共享存储的前缀查找时,不能自动表达“MLA 已命中到 1k、KDA 检查点尚缺失”的补点需求。

MooncakeStoreConnector:可以存什么,缺的又是什么

vLLM 的共享 Store 路径按缓存分组建立 ChunkedTokenDatabase

text
PoolKey
  = cache_prefix
  + model basename
  + tp / pcp / dcp / pp rank
  + 缓存分组 ID
  + chained prefix hash

value
  = 该缓存分组物理块对应的一组内存段

固定实现见:

需要区分 GPU 内存布局与 Mooncake 对象布局:

text
vLLM GPU 布局
  每个 KDA 层各有原始状态页

Mooncake Store 中的值
  一个分组或分块键
  → 一个或多个底层存储切片
  → 可由分散聚集地址和长度组成

也就是说,不应简单理解成“每个 KDA 层必然对应一个独立的 Mooncake 对象”。工作进程会对底层存储去重,计算每个物理块在各存储中的地址与字节跨度,再把多个切片组成同一个 Store 值。

MLA 缓存组和 KDA 缓存组可以针对同一个前缀哈希分别保存:

text
前缀哈希 @ vLLM 缓存分组=MLA
  → MLA 缓存页

相同前缀哈希 @ vLLM 缓存分组=KDA
  → KDA 状态

如果 Store 中同时存在 MLA(P_1k) 和完整 C(P_1k),外部协调器的数据模型可以表达完整的 1k 命中,并将两个缓存分组装入各自的 GPU 块;但固定快照缺少 Kimi KDA + MooncakeStoreConnector 的专项完整调用链证据,因此这里只能判断为“接口形状允许”,不能视为已经闭环。

本地细粒度检查点不等于远端细粒度对象

本地前缀缓存可以在 prefix_match_unit 边界注册未满整块的检查点,但固定 MooncakeStoreConnector 的外部查询和写入仍按同一个调度块对齐。下文将查询与写入共用的对齐单位记作 Astore

调度器查询前会做:

Llookup=LrequestAstoreAstore.

Store worker 写入前同样会做:

Lstore=LcomputedAstoreAstore.

在固定快照中,调度器的 _block_size、工作进程使用的块大小和协调器的 lcm_block_size 来自同一个调度块大小;分开写两个公式只是为了区分查询与写入阶段。

固定实现还要求:

text
每个缓存分组的块大小 % 哈希块大小 == 0
A_store % hash block size == 0
A_store % 每个缓存分组的块大小 == 0

本地与远端的对齐差异

6976 可以是本地检查点,却不一定是远端对象边界

示意参数:本地 prefix_match_unit = 64,Mooncake A_store = 4096。虚线标出已经物化的本地 C(P6976)。

4,096 tokens8,192 tokensC(P6976): 6,976 tokens · 本地检查点

Local hash boundaries

rolling hash 可以每 64 token 寻址

every 64 tokens4096Hash(P6976)addressable8192

Actual KDA checkpoints

只有真正物化并保留的 C(P_L)

C(P4096)C(P6976)materialized locally

Mooncake Store boundaries

固定路径仍按 A_store 对齐

A_store = 4096远端 4096远端 8192

因此,一个本地 Hash(P6976) → C(P6976) 条目,只有在 6976 同时满足外部 Store 对齐约束时,才可能成为远端对象边界。prefix_match_unit=64 提供本地可寻址能力,但不会自动把 Mooncake 的传输和持久化粒度也降到 64。

远端部分分组缺少命中边界

固定快照中的 MooncakeStoreCoordinator 对外主要返回:

text
协调后的最终命中长度
各分组的加载掩码

它不会同时返回:

text
远端 MLA 的最长命中长度
远端最佳 KDA 检查点
缺失的 KDA 边界

因此:

text
Store 已有 MLA(P_1k) + C(P_1k)
  → 数据模型可以表达完整的 1k 命中

Store 只有 MLA(P_1k)
  → 最终混合命中长度被协调到更短的边界
  → Scheduler 不知道远端 MLA 曾命中到 1k
  → 不会仅凭这次查询自动在 1k 处暂停并生成 C(P_1k)

更严谨地说:截至本文固定的提交,外部查询接口没有显式返回“远端缺失的 KDA 边界”;本文也没有找到利用远端部分分组命中自动触发检查点补算的完整调用链。

要补齐这条链,需要本地与外部查询保留每个分组的命中结果,或直接返回 missing_branching_boundary,随后让调度器、尾部检查点注册和 Store 保存掩码共享同一个边界。

远端命名空间必须编码模型身份与缓存 ABI

固定 vLLM 键中的 model_name 只取模型路径的末级名称,cache_prefix 默认又是空字符串。由此可以推断:如果不同部署共享同一个 Mooncake 集群,同时拥有相同模型名和 token 前缀,但权重版本、量化格式、适配器或缓存布局不同,仅靠现有默认命名空间不足以保证安全隔离。

生产环境中的缓存命名空间至少应编码:

text
模型仓库 + 精确版本 / 权重摘要
分词器版本与对话模板标识
量化格式
KDA 状态数据类型与布局版本
适配器 / LoRA 标识
TP / PP / PCP / DCP 拓扑
推理引擎与缓存 ABI 版本

最简单的做法是为每一套不可互换的部署设置唯一 cache_prefix。前缀哈希只能证明 token 序列相同,不能证明由这些 token 计算出的模型状态可以互换。

SGLang HiCache:本地与 Mooncake 命中长度怎样合并

这一节描述的是 SGLang 从 GPU 到主机内存再到 Mooncake 的分层语义,不应机械套用到 vLLM MooncakeStoreConnector。

首先由 UnifiedRadixCache 找出 GPU 与主机内存中所有必要组成部分共同支持的本地锚点。若 L_host 表示主机内存在 GPU 命中之后额外覆盖的 token 数,则:

Llocal=Ldevice+Lhost.

Mooncake 只需要查询 L_local 之后的后缀。由于 KDA 检查点是稀疏边界集合,而不是向下闭合的逐 token 缓存,不能把每个分组的“最长长度”简单取最小值。

定义远端可恢复边界集合:

Bstore={L[Llocal,Lcommon] | MLA/KV[Llocal:L) 连续完整,C(PL) 完整,其它辅助状态与所有必要 ranks 完整}.

候选总命中是:

Lcandidate=max({Llocal}Bstore).

例如:

text
MLA/KV 缓存页连续存在到 7500
KDA 检查点只存在于 4096 和 8192
本地锚点为 0

不能取 min(7500, 8192)=7500,因为 C(P7500) 不存在。真正可恢复的边界只能是 4096。

实际传输还可能因为超时、尽力而为策略、主机内存池不足或 rank 间不一致而变短。数据插入主机基数树后,权威结果应再次执行所有组成部分的完整匹配:

Lfinal=Rematch(Device+Host+LoadedStoreData).

KDA 状态槽位数不是 token 数,不能把 mamba_host_hit_length=1 当作“多命中一个 token”相加;它只是该 token 边界有效性的必要条件。

生成请求通常还要重算提示词的最后一个 token

即使完整提示词的缓存都存在,标准自回归生成仍需要最后一个提示词位置的输出 logits。若没有另外缓存最后一层隐藏状态或 logits,语义上至少需要:

LhitLprompt1.

固定 vLLM 的本地 KVCacheManager 会把最大缓存命中设为 request.num_tokens - 1。MooncakeStoreScheduler 在全命中时还会将结果向下对齐到调度块边界。换句话说,语义上至少重算最后一个 token;受分配器和连接器的对齐约束影响,实际实现可能重算最后一个块,而不只是一个 token。

注意力缓存页与 KDA 检查点应怎样组织

注意力 KV 缓存页和 KDA 检查点应视为两种不同资源,而不是让每个 KV 缓存页都附带一份完整 KDA 状态:

text
注意力 KV 缓存页池
  第 0 页
  第 1 页
  第 2 页
  ...

KDA 检查点池
  C(P_L1)
  C(P_L2)
  ...

需要分开理解三个独立概念:

text
P_KV
  注意力/MLA 的物理页大小

G_match
  前缀滚动哈希与查找粒度

𝒞_KDA
  已实际物化、注册且仍驻留的检查点集合

最终可以跳过前向计算的长度,是最近一个完整边界:

Lhit=max{LLcommon | KV[0:L) 完整C(PL) 完整}.

如果注意力 KV 已存在到 6976,但只有 C(P4096),那么仍需从 4096 重算到请求末尾。常规推理引擎通常会在查找阶段就把 KV 命中截到 KDA 边界,避免无效加载更长的 KV。

不存在一个跨 SGLang、vLLM 和 Mooncake 通用的“每 N 个 token 自动创建 KDA 检查点”旋钮。检查点密度应该根据 B_logical、完整物理检查点大小、状态池容量、公共前缀分布、平均重算成本和实际命中收益来决定。系统提示词、对话轮次、工具结果、RAG 公共文档和请求尾部等高价值边界可以作为额外候选,但仍需由推理引擎真正物化状态。

Mooncake Store 的对象、Upsert 与分组语义

固定快照中的 Mooncake Store 是面向对象的缓存。它看到的是:

text
对象键 → 一个或多个内存切片

PutStart / PutEnd 保证读取方不会看到只写了一部分的对象。对于按内容寻址的前缀缓存,已经发布的前缀哈希键应当被视为不可变对象:新的前缀或新状态使用新键,而不是覆盖旧检查点。

Mooncake 同时提供显式的 Upsert 生命周期,可以更新已有键,甚至在布局不变时复用分配空间。因此,“不可变”是前缀缓存协议应该坚持的不变量,并不意味着 Store 完全没有更新接口。对已经发布的前缀哈希键原地执行 Upsert,会破坏哈希到检查点的稳定映射,不应作为普通缓存更新方式。

固定设计文档见 mooncake-store.md

这里还要区分两个不同的 “group”:

text
vLLM 缓存分组编号
  缓存布局 / PoolKey 命名空间
  区分 MLA、KDA、SWA 等缓存分组

Mooncake ReplicateConfig.group_ids
  Mooncake 原生对象生命周期分组
  用于元数据路由、租约续期和尽力而为的淘汰

Mooncake 原生分组可以包含不同大小、覆盖不同 token 范围的对象,但它不是事务:

text
不要求预先声明成员数量
不保证所有成员原子创建或同时可见
不保证所有成员全有全无地淘汰
已有对象的分组归属不可变

因此,一个完整的混合前缀是否可恢复,仍要由上层协调器逐项验证。

固定快照中的两种映射

text
SGLang
  检测 ReplicateConfig.group_ids 能力
  → 可为一个逻辑 HiCache 页的组成对象
    使用 sglang-hicache:<logical-key> 作为 Mooncake 原生分组

vLLM
  PoolKey 中的 @group:N 是缓存分组命名空间
  → 固定快照创建的 ReplicateConfig 只设置 preferred_segment
  → 没有把 @group:N 当成 Mooncake 原生 group_ids

SGLang 可以把一个逻辑 KDA 检查点展开为递推状态与若干卷积状态对象;vLLM 则按缓存分组和分块键组织一个或多个底层存储切片。

同一个 Mooncake 集群也不意味着 SGLang 与 vLLM 可以直接互读。两者的键命名空间、rank 编码、层顺序、数据类型、形状、字节布局和检查点策略都不同。跨引擎复用需要额外定义稳定的公共 ABI。

推荐的 Store 数据组织方式,而不是当前实现描述

若要在 Mooncake 之上实现更强的逻辑原子性,可以由上层增加清单协议:

text
注意力 KV 对象
  独立的内容寻址缓存页

KDA 检查点对象
  递推状态、卷积状态或原始状态切片
  按边界、rank 和模型命名空间组织

前缀清单
  前缀长度
  KV 缓存页引用
  KDA 组成部分引用
  数据类型、布局和拓扑元数据
  提交标记

写入顺序:

text
1. 写入所有组成对象
2. 验证写入成功
3. 最后写入清单或提交对象

失效时反过来先让清单不可见,再异步回收各组成对象。这个清单方案是推荐设计,不代表固定快照中的 SGLang 或 vLLM 已经采用。

一次完整请求的控制流

请求生命周期

一个混合前缀从本地查询到再次写回 Store

每一步都必须维持相同的 token 边界、前缀标识、位置信息和缓存 ABI。

  1. 1

    本地查询

    GPU + 主机内存

    分别查询连续的 MLA/KV 缓存页和边界检查点。

    MLA/KVC(P_t)
  2. 2

    混合结果收敛

    L_local

    求所有必要组成部分共同支持的本地锚点;必要时产生缺失边界。

    分组命中共享边界
  3. 3

    Mooncake 查询

    外部后缀

    按 Store 对齐查询连续缓存页、末端检查点与所有 rank。

    A_store命名空间
  4. 4

    恢复与 CoW

    请求私有状态

    加载缓存,并复制或回填到当前请求的可写状态。

    回载写时复制
  5. 5

    重算与发布

    模型前向计算

    重算未命中的后缀,在合法边界注册检查点并写回 Store。

    重算后缀注册写入

这条链上的 token 长度、前缀哈希、位置信息、MLA 缓存和 C(P_t) 必须始终一致。任何一环如果只知道“KV 命中了多少”,却不知道“哪个 KDA 检查点与之对应”,都无法安全跳过模型前向计算。

能力边界

text
能力                         SGLang 本地   SGLang + HiCache/Mooncake
KDA 前缀正确性               已支持        完整检查点数据面已支持
从大批前向计算提取中间状态   分块对齐后支持  追踪需使用 Triton 或回退路径
完整检查点远端读写           不适用        temporal/conv 的 put/get 已实现
自动补算缺失的分叉状态       已支持        HiCache 模式暂时跳过

能力                         vLLM 本地     vLLM + MooncakeStoreConnector
KDA 原始状态页               已支持        数据结构可承载
细粒度部分检查点             已支持        远端仍受 Scheduler/LCM 对齐限制
部分命中后的 CoW             已支持        连接器回载链路的证据有限
自动补全缺失的分叉检查点     已支持        连接器路径没有传递分叉边界
完整检查点远端读写           不适用        数据结构支持,但完整调用链证据有限
远端仅命中 MLA 时补算 KDA    不适用        尚未返回缺失边界

以下不变量适合作为实现与回归验证的检查清单:

text
本地分支补全
  req1=4k,req2/req3 分别共享 1k

远端完整命中
  Store 同时存在 MLA(P_1k) 与 C(P_1k)
  重算得到的 logits 与无缓存基线一致

远端部分分组命中
  Store 只有 MLA(P_1k)
  req2 只付一次补算并上传 C(P_1k)
  req3 在另一实例完整命中

外部对齐
  本地存在 Hash(P6976) → C(P6976)
  Store 的 Scheduler/LCM 块大小 = 4096
  远端只暴露合法的对齐边界

命名空间隔离
  token 前缀和模型 basename 相同
  模型 revision、量化配置或 adapter 不同
  缓存结果必须相互隔离

FlashKDA 状态追踪
  return_intermediate_states=True
  必须使用 Triton 回退路径,并生成与分块边界对齐的检查点

对象完整性
  temporal/conv 或原始状态的任一必要部分缺失都必须视为未命中

写时复制
  共享检查点在原请求继续执行和新请求命中后都必须保持不可变

完整提示词命中
  最后一个提示词 token 或 block 的 logits 计算路径必须保持正确

相关配置项如何影响检查点追踪

SGLang

text
--mamba-radix-cache-strategy extra_buffer
  使用专用追踪槽位保存前缀检查点

--page-size
  RadixCache 与 HiCache 的页面粒度,也参与可追踪边界对齐

--mamba-track-interval
  状态追踪和保留策略的一部分,不等同于强制设置 prefill 前向计算大小

--chunked-prefill-size
  Scheduler 的 prefill token 预算,可以大于检查点粒度

--enable-hierarchical-cache
--hicache-storage-backend mooncake
  启用 GPU → 主机内存 → Mooncake 链路

vLLM

text
enable_prefix_caching
  启用混合前缀缓存

mamba_cache_mode=align
  Kimi Linear 的实际状态物化模式

enable_chunked_prefill
  允许 Scheduler 在必要边界结束当前调度步

max_num_batched_tokens
  可以保持较大,不要求机械按某个 interval 切分

prefix_match_unit
  调整本地可寻址的哈希粒度,但不会自动创建状态
  也不会自动改变 Mooncake 的 Scheduler/LCM 对齐

VLLM_PREFIX_CACHE_RETENTION_INTERVAL
  调整已创建检查点的保留密度,但不会补建历史状态

kv_connector_extra_config.cache_prefix
  应包含模型 revision、缓存布局和 adapter 标识,以隔离不可互换的部署

验证实现时,不要一开始就把预填充预算降到 1k,否则无法区分“按需补点生效”和“所有请求本来就被机械切成 1k”。对于 Mooncake,也不能只检查对象键是否出现;还应验证组成部分是否完整、混合缓存命中长度、缺失边界传播、远端对齐、命名空间隔离、写时复制、最后一个 token 或块,以及重算后的 logits。

最终判断

回到 req1/req2 示例,可以分三层判断:

req1 的 4k 预填充如果只留下 C(P_4k),req2 第一次共享 1k 时确实必须重新计算这 1k;C(P_4k) 无法回滚,MLA 缓存页也不能替代 KDA 检查点。

SGLang 本地可以让 req2 的较长扩展计算跨过 1k,再从计算内核返回的分块对齐中间状态中抽出并插入 C(P_1k);需要中间状态时,FlashKDA 会回退到 Triton。vLLM 不挂连接器的本地路径则把 req2 动态截到 1k,利用调度步结束补出并注册 C(P_1k)。两边都不要求把所有请求的预填充预算固定成 1k。

再加上外部缓存限定:

SGLang + HiCache 已把递推与卷积检查点接到 Mooncake 的显式读写链路,但 HiCache 模式仍暂时跳过分支状态回填。vLLM 在本文固定快照中已支持本地物理块内的细粒度尾部检查点与写时复制;然而连接器与 Mamba/KDA 的本地路径会清零 shared_prefix_boundary 并采用 max(per_group_hits),MooncakeStoreConnector 的远端写入和查询还受同一个调度块对齐,MooncakeStoreCoordinator 也不单独返回 MLA 单组更长命中所对应的缺失 KDA 边界。

因此,可以确认 SGLang 已经存在且完整的 C(P_t) 能通过 HiCache/Mooncake 数据面读写;vLLM 的原始状态页和缓存分组数据模型,以及通用 Mooncake 传输路径,可以表示完整检查点。但在固定快照中,这一结论仍属于“接口形状允许”,不能替代一条已经确认闭环的 Kimi KDA 专项调用链。

优化目标不是机械地每隔 1k 切一次前向计算,而是在出现可复用分叉时只补算一次,并把同一个边界持续传递到本地基数树和调度器、主机暂存区、Mooncake 键值、查询、命名空间与保留策略。还要始终区分物理状态块、前缀匹配粒度、调度对齐块与实际检查点集合:细粒度可寻址不等于细粒度状态自动生成,更不等于远端对象也具备相同粒度。检查点一旦发布就应保持不可变,后续扩展通过写时复制产生新的工作状态或检查点。