【轻量化复刻大模型 / KIMI K3】KDA 对 AI Infra 的挑战与各家应对

【轻量化复刻大模型 / KIMI K3】KDA 对 AI Infra 的挑战与各家应对

目录

0x00 概要

KDA 的固定状态设计让标准 serving 栈的若干组件失效，因此，各大推理引擎都对其进行适配：

基础设施层 标准假设 KDA 挑战 应对（参考实现）
前缀缓存 Prefix Cache 缓存 KV 前缀 需缓存"状态 S"（无逐 token 条目） SGLang radix cache 纳入 temporal checkpoint（copy-on-write）；kimi3-master prefix_cache.py 两级 hash
KV 池 Memory Pool 统一 KV 页 混合模型（KDA+MLA）异构缓存 SGLang MambaPool（conv + temporal）vs MLA KV 池；kimi3-master paged_cache.py::PageKind
上下文并行 CP KV 广播 状态跨段不可加 KCP 前缀扫描
varlen 批量 cu_seqlens 每序列独立状态推进 cu_seqlens + 段偏移
投机解码 KV 截断回滚 状态需回滚/分支/重放 SGLang ReplaySSM（轨迹重放）；kimi3-master ReplaySSM
CUDA Graph 静态 KV 布局 递推状态无法标准捕获 ReplaySSM 折叠机制

本文前半篇用概念性语言解释 KDA 为什么改变了推理缓存的正确性条件，以及推理框架如何在不破坏正确性的前提下节省显存和重算；后半篇给出源码导航，对应到 SGLang 与 vLLM 的具体文件与关键函数。

注：

  • 项目https://github.com/RossiLH/Model-Viz 和系列博客是我的一个尝试，也是我学习的过程。不确定是否满足读者的需求，我自己会反思改进，也希望大家提出宝贵意见。
  • 涉及 KIMI K3，DeepSeek V4，PI0.5，PI0.7 等模型。后续可能会跟踪新模型。
  • 全 CPU 复刻，不涉及 GPU，也不涉及大规模训练（这是另外一个庞杂深邃的领域）和硬件。
  • 模型的组件尽量是原始大模型的缩小版，如果组件没有开源版本，则用近似组件替代。
  • 会解释每个组件在解决一个什么问题、这个问题难在哪、为什么只能这么解、解法的原理是什么，以及组件之间如何咬合成一个整体。重点不是记名词，而是理解一条主线。
  • 每个主要组件都有可直接运行的 Gradio demo——无需 GPU。另有若干可视化脚本（plot_*.py），可生成静态图。
  • 所有文件行号基于本地源码快照；源码后续变化时，行号可能移动，但函数名和调用关系仍可作为定位入口。
  • 为了保证每个文章可以尽量独立成文，所以有些内容会在不同文章中都存在，可能在系列层面上存在冗余，还请大家理解。

0x01. 先建立正确的心智模型

1.1 混合注意力的完整执行链

Kimi K3 的混合注意力可以粗略画成：

输入 token
    ├─ MLA ──> token-addressable KV/latent history
    └─ KDA ──> convolution state
              └── recurrent state
                    └── KDA 输出
    └──────────> 混合输出

K3 特色

MLA 仍保留可按 token 或 block 访问的历史；

KDA 不保存每个历史 token 的完整地址，而是逐步更新一个固定大小的状态：

S_t = M_t S_{t-1} + b_t
  • S_{t-1}：前一时刻长期记忆；
  • M_t：当前 token 决定旧内容保留、衰减或变换多少；
  • b_t：当前 token 写入的内容；
  • S_t：处理完当前 token 后的长期状态。

KDA 还常使用短卷积，因此要保存最近若干 token 的 convolution state。

区别

因此，传统 Transformer 的 KV Cache 可以看成 "按 token 编号的历史档案"：

token 0 -> K,V
token 1 -> K,V
token 2 -> K,V
...

新 Query 到来时，系统可以读取任意历史 token。前缀缓存通常可以理解为：token 前缀相同，就能复用对应的 KV block。

而 KDA 则把历史持续写入一个固定大小的 recurrent state。

所以，KDA 不是简单地把 KV Cache 换成另一种张量，而是把 "模型状态" 从一种扩展成了多种语义不同的对象。

1.2 三种状态对比

状态 记住什么 是否随新 token 更新 典型恢复方式 是否可按 token 任意截断
MLA KV 历史 token 的可寻址表示 追加 读取 KV block 基本可以
KDA recurrent state 整个前缀压缩后的长期摘要 是 恢复 checkpoint 不可以随意截断
KDA convolution state 最近局部窗口 是 恢复局部窗口快照 需在合法边界保存

Recurrent state 决定长期记忆，convolution state 决定下一步局部卷积输出。两者都影响后续 token，因此不能只恢复 MLA KV。

1.3 KDA 的顺序性：可分段但必须保序

对于KDA，我们把 token 段 C 对状态的整体影响记为：

F_C(S) = M_C S + b_C

两段连续执行的合并是：

F_{C_2}(F_{C_1}(S)) = M_{C_2} M_{C_1} S + M_{C_2} b_{C_1} + b_{C_2}

中间的 M_{C_2} b_{C_1} 很重要：第一段写入的信息，还要接受第二段的衰减或擦除。

这解释了两件事：

  1. 可以把一段 token 浓缩为一个 (M, b)，所以能 chunkwise 并行；
  2. 不能把任意两个 token 段交换顺序，因为后一段会作用于前一段的写入。

它满足结合律，所以可以先算小段、再合并；通常不满足交换律，所以 token 顺序不能随意调换。这个 "可分段但必须保序" 的性质，正是 KDA 能并行又能编码时序的原因。

Checkpoint 保存的就是某个前缀边界的结果，例如：

token boundary 4096
    ├─ MLA KV prefix
    ├─ KDA recurrent state S_4096
    └─ KDA convolution state C_4096

0x02. KDA 的优势与结构性限制

2.1 优势

KDA 把长历史压进固定大小 state，主要收益是：

  • decode 不必反复扫描完整历史；
  • 历史读取带宽不再随上下文长度线性增长；
  • 状态可按 chunk 计算并通过仿射复合合并；
  • 衰减可以表达慢遗忘；
  • delta 更新可以沿当前 key 定向覆写；
  • 适合作为长程 mixing 分支。

2.2 代价：固定状态不是完整档案

状态是压缩摘要，不是完整档案。一个固定 state 必须让任意长历史竞争有限容量：

很多 token
    ↓ 压缩
一个固定状态矩阵
    ↓
相似 key 可能碰撞
旧信息可能被衰减
未来 query 不能重新寻找原 token

两个相似 key 可能写入相近方向并互相干扰；一旦旧信息被衰减或覆盖，未来 query 无法像 Full Attention 那样重新找到原 token。

因此 KDA 的 state 更像 "持续更新的寄存器文件"，而 Full/MLA KV 更像 "可按地址读取的缓存"。Kimi K3 同时使用两者，正是为了让 KDA 负责高效 mixing、MLA 负责精确检索。

KDA 的优势不是 "在所有任务上替代 Full Attention"，而是用固定状态承担大多数长程 mixing，再由 MLA 等 token‑addressable 分支补足精确检索。

0x03. 为什么 Prefix Cache 与 KV Cache 不再等价

3.1 普通 Prefix Cache 为什么简单

普通 Transformer 的 KV Cache 近似是一个追加式数组：

prefix tokens: [A, B, C, D]
KV cache:  [K_A,V_A] [K_B,V_B] [K_C,V_C] [K_D,V_D]

后续生成 E 时，只会追加 [K_E,V_E]，不会修改 A、B、C、D 的 KV。因此：

token 前缀相同 ≈ 对应 KV 相同 ≈ 可以直接复用

这就是 "Prefix Cache 与 KV Cache 几乎等价" 的来源。在普通 Transformer 中，前缀的 KV 是 "写完即稳定" 的：后续 token 不会修改前面 token 的 K/V。只要 token 前缀一致，缓存内容就一致。

3.2 KDA 为什么打破这个等价关系

KDA 的 state 不是每个 token 的独立记录，而是整个前缀递推后的结果：

S_n = F_n(F_{n-1}(… F_2(F_1(S_0))))

其中每个 F_t 都代表一个 token 的状态更新。

假设两个请求共享：

请求 1: A B C D E
请求 2: A B C D X

它们在 A B C D 处的 KDA state 可以共享，但在 E 或 X 之后必须分开。共享的是一个稳定的前缀快照，不是可以被两个请求同时原地修改的活动状态。

更重要的是，若系统只知道 token 前缀命中到 D，却没有保存 S_D，就不能仅凭 D 的 token 恢复 KDA。系统必须从更早的 checkpoint 重新执行：

checkpoint at B
    ↓ replay C, D
    ↓ obtain S_D
    ↓ continue E or X

因此，对于 KIMI K3 来说，恢复位置 n 必须同时知道：

  1. MLA 在位置 n 以前的完整 KV；
  2. KDA 在位置 n 的 recurrent state；
  3. 卷积分支在位置 n 的 convolution state；
  4. 所有状态确实来自同一条 token 前缀和同一套模型执行路径。

单独命中 token 前缀，只能说明 "文本相同"，不能证明 "递推状态已保存并且可恢复"。

常见疑问：为何 token 前缀命中并不意味着 KDA 缓存可以恢复？

Radix Tree 或 prefix hash 可能只证明 token 0 … token N 相同，但恢复还需要：S_N 已保存、C_N 已保存或可重建、MLA KV_N 可用。若 state 只在 0、4096、8192 保存，而 token 匹配到 6000，系统不能凭文本直接得到 S_6000，只能退到 4096 并 replay 4096‑6000。

0x04. KDA 给推理系统带来的新挑战

4.1 缓存对象从一种变成多种

系统不再只管理 KV block，而要管理：

状态 语义 是否可按 token 任意截断 典型管理方式
MLA KV 可寻址历史 基本可以 token/block cache
KDA recurrent state 前缀递推摘要 不可以随意截断 checkpoint/state block
KDA convolution state 局部卷积上下文 需在合法边界保存 small rolling state

挑战不只是 "多了一种 cache"，而是这些对象的语义、更新方式、可共享性都不同。

一个可扩展的缓存条目至少应包含：

CacheEntry {
    model_version
    prefix_range
    tokenization_fingerprint
    state_kind: mla_kv | kda_recurrent | kda_conv
    physical_location
    dtype / layout
    writable
    parent_checkpoint
}

其中最关键的是 prefix_range 和 state_kind。没有边界，就无法判断恢复是否连续；没有类型，就无法区分 "可寻址历史" 和 "顺序摘要"。

因此，核心是把缓存键从 "token prefix" 升级为：

(model/version, tokenizer, prefix boundary,
 MLA KV, KDA recurrent state, convolution state)

实际系统还要处理：

  • 活动 state 与共享 checkpoint 分离；
  • state 压缩和恢复；
  • checkpoint 淘汰；
  • 本地与远端缓存协调；
  • 不同并行 rank 的对象对齐。

4.2 状态恢复必须满足边界一致性

一个合法恢复点不是 "最长 token 前缀" 这么简单，而是一个所有相关状态都可用的共同边界。若 MLA KV 恢复到 4096 token，而 KDA state 只恢复到 3968 token，后续计算就不再等价于连续执行。

混合模型的输出是多个分支共同决定的。假设：

MLA 恢复到 N=6000
KDA 恢复到 N=4096

那么两个分支看到的历史长度不同，后续混合结果自然不是连续执行在 N=6000 的结果。因此恢复对象必须像一张 "同一时间点的系统快照"，而不是几个只按长度看起来相近的缓存片段。

常见疑问：为何 KDA 状态不能像 KV 一样任意截断？

截断 KV 意味着以后不再读取某些 token；剩余 KV 仍然是有效记录。截断 KDA state 则不同。state 的每个元素都可能是很多历史 token 混合后的结果。你不能从一个大 state 中随便删掉前 10% 的 "历史"，因为它没有独立 token 地址。正确做法只有两种：使用已保存的更早 checkpoint，再 replay；或者使用已经在目标边界保存的完整 state。

4.3 活动态不能直接共享

MLA KV 通常写入后不变，可以被多个请求引用。KDA state 会随着 decode 继续更新，如果多个请求直接指向同一份可写 state，一个请求的更新就会污染其他请求。

因此必须把 "共享的稳定 checkpoint" 和 "请求自己的活动副本" 分开。共享 checkpoint 只读，活动 state 私有可写。

4.4 PD/EPD 分离不再只搬 KV

传统 Prefill‑Decode 分离主要传输 KV blocks。KDA 需要传输或重建 recurrent state、convolution state，结合多模态后还要考虑 Encoder 产生的视觉 embedding。于是接口抽象从：

KV Cache transfer

扩展为：

model state transfer

通信层、调度层和缓存层都必须知道状态类型、版本、边界和生命周期。

0x05. SGLang 的应对思路

SGLang 在 Radix Tree 的前缀匹配能力上增加 KDA‑aware checkpoint。Radix Tree 仍负责回答 "哪些 token 前缀匹配"，但只有带有对应 KDA 状态的节点，才是可直接恢复的执行点。这相当于把缓存节点从 prefix -> KV blocks 扩展成：

prefix boundary -> MLA KV + KDA recurrent state (+ convolution state)

5.1 KDA‑aware Radix Tree

SGLang 原有 Radix Tree 擅长匹配请求之间的 token 前缀。KDA‑aware 之后，每个候选节点还要知道是否有可恢复 state：

Radix Tree
root
└─ system prompt
    └─ tools
        ├─ user branch A
        └─ user branch B

每个节点可能带：token prefix、MLA KV blocks、KDA checkpoint、checkpoint boundary。

匹配流程变成：

  1. 找最长 token prefix
  2. 沿匹配路径向前找最近合法 KDA checkpoint
  3. 取同一边界的 MLA KV + KDA state
  4. 把命中后的尾部重新计算

因此最长 token 命中点和最终恢复点可能不同。它把缓存正确性从字符串匹配提升为模型状态匹配。

常见疑问：SGLang 的 KDA‑aware Checkpoint 是什么？

普通 Radix Tree 节点回答 "前缀匹配到哪里"；KDA‑aware 节点还回答 "这里有没有对应可恢复 state"。只有 token 匹配和状态匹配同时成立，节点才是有效恢复点。命中流程是：匹配 token 路径 → 回退到最近合法 checkpoint → 联合恢复 MLA/KDA → replay 剩余 token。

5.2 Sparse Checkpoint：为什么不是简单等间隔

如果每个 token 都保存完整 KDA state，空间会快速膨胀。SGLang 因而采用稀疏 checkpoint：只在部分边界保存状态，而且不是简单地全局等间隔采样。

策略包含三类位置：

  • Prefill：优先 chunk 边界；
  • Decode：按固定间隔生成 checkpoint；
  • Radix 分叉：多条请求共享 / 分叉的位置优先保留。

原因是三类位置的复用概率不同：

  • Prefill chunk 边界本来就是执行边界，保存它不需要额外制造同步点；
  • Decode 每个 token 都可能成为未来前缀，但每 token 保存太贵，所以抽样；
  • 分叉点已经被多个请求证明有复用价值。

之后还要限制单条路径的 checkpoint 数量，并用 LRU 回收长期未访问的状态。

因此，它本质上是在缓存预算和重算成本之间做优化，而不是简单追求最大命中率。不平均地保存所有位置，而是利用执行结构和访问模式，把 checkpoint 当成有限预算下的 "缓存投资"。

常见疑问：SGLang 的稀疏 Checkpoint 如何解决空间 / 命中率矛盾？

不保存每个 token 的 state，而是选择少量高价值边界保存。这样减少驻留开销，再用命中后的局部重算弥补较粗恢复粒度。Prefill chunk 边界天然适合保存；Decode 通过间隔控制数量；Radix 分叉点通常有更高复用概率；LRU 回收冷 checkpoint；路径预算防止单条长路径吞噬显存。

5.3 Compressed Checkpoint

并非所有 checkpoint 都同样有价值。SGLang 将状态分成：

运行时活动池：当前请求正在更新，保持原始精度

缓存池：Radix Tree 上等待复用，可将 temporal recurrent state 压缩为 INT8

局部窗口状态体积较小，通常保持原始精度。命中压缩 checkpoint 时：

INT8 checkpoint
    ↓ dequantize
新的活动 state slot
    ↓ 以原始精度继续 decode

这不是让 KDA 永久在 INT8 上递推，而是只压缩 "停放在缓存里的快照"，避免连续量化 / 反量化造成误差累积。

SGLang 将正在执行的活动状态与 Radix Tree 中的缓存状态分离；不活跃 checkpoint 可放入独立缓存池，并把主要占空间的 KDA temporal state 压缩为 INT8，而较小的局部窗口状态保持原始精度。

命中压缩 checkpoint 时，先解压到新的活动 state 槽位，再以原始精度继续 token 更新。这样量化误差主要发生在缓存驻留和恢复处，不会在连续 decode 中反复累积。它仍然是有损恢复，需要在质量和容量之间权衡。

常见疑问：Compressed Checkpoint 做什么？

它降低低热度 checkpoint 的存储成本，把有限显存留给高复用状态。主要的 temporal recurrent state 使用 INT8，较小的局部窗口状态保持原始精度。命中时先解压到活动槽，再恢复正常精度执行。这样节省驻留空间，但必须接受量化恢复误差，并为不同通道保存合适的量化参数。

需要注意：INT8 checkpoint 是可选优化，不是所有 SGLang 部署都默认启用。

5.4 COW、Snapshot 与 Donate

三者共同解决 "共享缓存必须稳定、执行状态必须可变" 的矛盾。

Radix Tree shared checkpoint
    ├─ request A: copy-on-write -> private active state
    └─ request B: copy-on-write -> private active state
  • Snapshot：把请求当前状态冻结，挂到 Radix Tree，成为未来可复用快照；
  • Copy‑on‑write：命中共享快照时先复制或建立私有副本，后续更新不污染共享快照；
  • Donate：快照完成后直接转移 state slot 的所有权，只改索引，不再复制整块数据。

假设两个请求都命中 S_4096：

共享 checkpoint S_4096
    ├─ A: private S_A
    └─ B: private S_B

Snapshot 负责制造稳定的共享版本；COW 负责首次写入时保护共享版本；Donate 负责把已经完成的可写 state 直接转交缓存树，避免复制。Donate 的价值在于把 "复制一大块 state" 变成 "转交一个内存槽位"。核心原则：共享版本只读，活动版本可写。

常见疑问：COW、Snapshot、Donate 如何隔离状态？

Snapshot 固定共享的历史，活动副本负责继续递推；COW 在首次写入时复制；Donate 在旧请求结束后转交可复用的活动存储。这样共享 checkpoint 不被修改，执行又无需每步全量复制。

5.5 Unified Memory

Unified Memory 把 MLA KV block 和 KDA state block 放进同一个物理显存池，但保留各自的逻辑结构和分配单位。

显存池
MLA KV blocks ----->                 <----- KDA state blocks
                   共享空闲区域/动态边界

MLA KV 和 KDA state 仍然有不同的逻辑类型、大小和生命周期；统一的是底层容量管理。这样可以避免预先把显存永久切成两个比例失衡的池。

系统可以按实时负载动态调整两者占用，避免某一类缓存独占固定容量。统一的是物理容量管理，不是把两类状态混成同一种对象。

好处是容量可以随负载倾斜；代价是底层分配器需要处理不同大小对象、空洞回填和碎片问题。

常见疑问：Unified Memory 是什么？

它不是把 MLA KV 和 KDA state 变成同一种数据，而是让它们从同一个物理显存预算中分配。KV manager 和 State manager 在逻辑上仍然是两个，但共享物理容量池。工作负载偏向 MLA 时可以多留 KV；偏向 KDA 时可以多留 state。

5.6 Mooncake 的职责边界

职责边界可以画成：

SGLang
    ├─ 判断什么是合法恢复点
    ├─ 生成/恢复 KDA checkpoint
    └─ 管理本地请求状态
    ↓
Mooncake / HiCache
    ├─ 保存已确认合法的缓存对象
    ├─ 跨实例传输
    └─ 管理更大范围的共享缓存

Mooncake 不需要重新理解 KDA 的递推数学；它只需要可靠地存储和搬运带有边界、类型和版本信息的对象。

0x06. SGLang 源码导航

本章不是概念性总结，而是源码导航。所有文件行号基于本地源码快照；源码后续变化时，行号可能移动，但函数名和调用关系仍可作为定位入口。

源码中的 KDA 支持不是一个单独的 KDACache 类，而是分散在四层：

模型层:
    产生 KDA 的 q/k/v/g/beta

Kernel/backend 层:
    prefill、decode、conv、recurrent update

状态池层:
    recurrent state、convolution state、checkpoint、活动副本

调度/缓存层:
    prefix hit、state index、checkpoint 边界、跨实例传输

最容易误解的一点是：SGLang 的 "checkpoint pool" 不是普通 KV pool。它保存的是可恢复的 KDA 执行状态。

6.1 源码地图

文件 关键函数 / 类 作用
sglang/srt/models/kimi_k3.py KimiK3DeltaAttention.forward_qkvbg、forward_qkvbg_fused、forward K3 模型侧生成 KDA 输入、接入 attention backend
sglang/srt/layers/radix_linear_attention.py RadixLinearAttention.forward、unified_linear_attention_with_output 将线性注意力结果接入 forward batch 与 cache location
sglang/srt/layers/attention/linear/kda_backend.py KDAKernelDispatcher.decode、extend、target_verify 按 decode/prefill/spec‑verify 选择 KDA kernel
sglang/srt/mem_cache/mamba_checkpoint_pool.py Int8CheckpointStore、MambaCheckpointPool.store_from_active、load_to_active 缓存状态、INT8 压缩、恢复到活动池
sglang/srt/mem_cache/memory_pool.py HybridReqToTokenPool.__init__、donate_mamba_ping_pong_slot、clear 活动态分配、checkpoint pool 接入、Donate 与释放
sglang/srt/managers/schedule_batch.py _mamba_radix_cache_v2_req_prepare_for_extend、mamba_lazy_prealloc_at_boundary 选择 checkpoint 边界、准备 ping‑pong state
sglang/srt/mem_cache/unified_memory_pool.py UnifiedKVPool、MambaSubPoolSpec MLA/MHA 与 Mamba state 共享物理 byte buffer

6.2 模型层：KimiK3DeltaAttention 不是只做矩阵乘

文件：sglang-main/python/sglang/srt/models/kimi_k3.py

关键位置：

  • forward_qkvbg：第 1679 行起；
  • forward_qkvbg_fused：第 1773 行起；
  • forward：第 1821 行起。

核心逻辑：

hidden_states
    ↓
投影出 mixed_qkv、gates、beta、f/a 等 KDA 输入
    ↓
KDA/MLA attention
    ↓
AttnRes、output gate、o_proj

文件顶部第 1‑7 行直接声明了 K3‑specific features，包括 Attention Residual、Latent MoE、MLA output gate 和 full‑rank KDA gate。这说明 Infra 不能只按 "一个普通 GDN 层" 处理 K3。

6.3 统一入口：RadixLinearAttention.forward

文件：sglang-main/python/sglang/srt/layers/radix_linear_attention.py:38-121

函数：RadixLinearAttention.forward(...)

简化逻辑：

如果是 extend/prefill 且启用了 piecewise CUDA graph:
    走 unified_linear_attention_with_output
否则:
    交给 get_attn_backend().forward(...)

源码第 86‑121 行展示了这个分支。第 142‑156 行的 unified_linear_attention_with_output 临时缩窄 forward_batch.out_cache_loc，调用 backend，再把输出写回。

工程含义：

  • KDA 结果必须写入请求对应的 cache location；
  • cache location 属于 batch 元数据，不是 kernel 内部的临时变量；
  • CUDA graph/piecewise 路径也必须保持同一状态写入语义。

6.4 KDAKernelDispatcher：不同阶段用不同 kernel

文件：sglang-main/python/sglang/srt/layers/attention/linear/kda_backend.py:39-202

类：KDAKernelDispatcher

关键函数：

  • __init__：第 42‑202 行，选择 decode/prefill/verify backend；
  • packed_decode：第 204‑235 行；
  • decode：第 237‑264 行；
  • target_verify：第 266‑307 行；
  • extend：第 309‑332 行。

简化调用图：

decode:
    Triton / Helion / CuTe DSL / FlashInfer
prefill:
    Triton / FlashKDA / CuTe DSL / NVIDIA KDA / PTX
speculative verify:
    Triton / FlashInfer / nv_cutedsl 路径

源码第 64‑190 行明确列出这些 backend 和 fallback 条件。第 237‑264 行的 decode 传入 ssm_states、cache_indices；第 309‑332 行的 extend 传入同样的状态索引。

这证明：KDA state 是 kernel 的显式输入，而不是 "模型自己在 Python 对象里隐式保存" 的东西。

6.5 SGLang 的 state 不是只有 recurrent state

kda_backend.py 第 6‑18 行导入：

  • causal_conv1d_fn
  • causal_conv1d_update
  • TritonKDAKernel
  • build_verify_intermediate_state_indices

所以一次 KDA 执行同时涉及：

卷积窗口状态
recurrent SSM/KDA 状态
speculative verify 中间状态索引

常见疑问：recurrent state 和 convolution state 是什么？SGLang 有对应吗？

Recurrent state 是 KDA 长期递推矩阵，携带压缩后的长历史；convolution state 是短卷积需要的最近窗口，携带局部连续性。两者可以采用不同存储策略：temporal state 可压缩，局部窗口通常保留原精度。SGLang 的 KDA‑aware checkpoint 主要围绕 recurrent state 展开，但合法恢复不能只恢复它；对应局部窗口状态也要恢复或按明确规则重建。

6.6 Checkpoint Pool 的真实实现

6.6.1 设计分界：共享 checkpoint 只读，活动 state 私有可写

文件：sglang-main/python/sglang/srt/mem_cache/mamba_checkpoint_pool.py:15-45

模块说明明确区分：

cached states:
    radix-owned
    idle
    compressed

active MambaPool:
    running requests
    full precision
    kernel-facing

这正是 Copy‑on‑write 的根本：共享 checkpoint 只读，活动 state 私有可写。

6.6.2 Int8CheckpointStore

文件：sglang-main/python/sglang/srt/mem_cache/mamba_checkpoint_pool.py:60-184

关键函数：

  • __init__：第 79‑113 行，分配 qdata 和 scale；
  • quantize：第 117‑134 行；
  • dequantize：第 136‑138 行；
  • store：第 142‑149 行；
  • load：第 151‑153 行；
  • copy_to_pool：第 155‑164 行；
  • store_from_pool：第 166‑174 行。

数据布局：

qdata:
    [layers, slots, heads, d_v, d_k] int8
scale:
    [layers, slots, heads, 1, d_k]

第 127 行沿 d_v 做 max‑abs，第 128‑133 行完成对称量化。第 164 行把 checkpoint 解量化后直接写入 active pool，这就是 cache hit 时的 COW 恢复。

重要工程结论：

checkpoint 命中:
    int8 → 一次 dequantize → active state
继续 decode:
    active state 保持原始精度更新

不是每个 token 都在 INT8 state 上循环量化。

6.6.3 MambaCheckpointPool

文件：sglang-main/python/sglang/srt/mem_cache/mamba_checkpoint_pool.py:186-280

关键函数：

  • __init__：第 187‑225 行，创建 temporal int8 store、conv buffers、slot allocator；
  • store_from_active：第 245‑251 行；
  • load_to_active：第 253‑259 行；
  • estimate_mem_usage_bytes：第 262 行起。

store_from_active 的语义：

active_mamba_pool.temporal → checkpoint temporal
active_mamba_pool.conv     → checkpoint conv

load_to_active 的语义：

checkpoint temporal → active temporal
checkpoint conv     → active conv

SGLang 确实有 convolution state 的对应存储，不是只缓存 recurrent state。

6.6.4 可选初始化

文件：sglang-main/python/sglang/srt/mem_cache/mamba_checkpoint_pool.py:300-369

关键逻辑：

如果没有 Mamba 配置：不创建
如果没有开启 enable_int8_mamba_checkpoint：不创建
否则：估算空间、检查可用显存、创建独立 checkpoint pool

这说明 INT8 checkpoint 是可选优化，不应写成所有 SGLang 部署都默认启用。

6.7 活动状态、Donate 与边界

6.7.1 HybridReqToTokenPool.__init__

文件：sglang-main/python/sglang/srt/mem_cache/memory_pool.py:1271-1289

关键逻辑：

self.mamba_allocator = MambaSlotAllocator(...)
self.mamba_map = {...}
self.mamba_ckpt_pool = maybe_init_int8_mamba_checkpoint_pool(...)

活动 state 和 checkpoint state 使用不同池：活动请求走 mamba_allocator / active pool，Radix cached checkpoint 走 mamba_ckpt_pool。

6.7.2 donate_mamba_ping_pong_slot

文件：sglang-main/python/sglang/srt/mem_cache/memory_pool.py:1488-1509

函数：donate_mamba_ping_pong_slot(req, new_slot)

简化逻辑：

找到要保留的 ping-pong slot
    ↓
把旧 slot 作为可缓存 state 返回
    ↓
把请求切换到 new_slot 继续执行

第 1491‑1495 行的 docstring 已明确写出 "donate tracked-state slot to radix cache"。这不是复制数据，而是把 ownership/slot 交给 radix cache 使用。

6.7.3 clear

文件：sglang-main/python/sglang/srt/mem_cache/memory_pool.py:1565-1577

第 1569‑1574 行特别处理 mamba_ckpt_pool.clear()。原因是清空 Radix Tree 后，checkpoint slot 也必须释放，否则会出现 "树没了、slot 还被占用" 的泄漏。

6.8 checkpoint 何时产生

6.8.1 _mamba_radix_cache_v2_req_prepare_for_extend

文件：sglang-main/python/sglang/srt/managers/schedule_batch.py:2672-2735

函数：_mamba_radix_cache_v2_req_prepare_for_extend(req)

关键逻辑：

  1. 读取 mamba_cache_chunk_size()；
  2. 根据 mamba_checkpoint_grid(self.tree_cache.page_size) 得到 checkpoint 网格；
  3. 判断 extend_range.length >= checkpoint_grid；
  4. 计算要追踪的 seqlen；
  5. 把 state 放到 ping‑pong track buffer。

第 2677‑2680 行说明 checkpoint 网格来自 tree page；第 2699‑2730 行计算对齐边界。它解释了为什么 KDA checkpoint 不是任意 token 都自然存在：调度器会主动选择合法追踪边界。

6.8.2 mamba_lazy_prealloc_at_boundary

文件：sglang-main/python/sglang/srt/managers/schedule_batch.py:3002-3031

函数：mamba_lazy_prealloc_at_boundary(mamba_track_interval)

第 3015‑3017 行检查请求是否落在 tracking interval；第 3026‑3031 行按需申请第二个 ping‑pong slot。

这就是 lazy ping‑pong：

平时：一个活动 slot
到 checkpoint 边界：临时申请第二个 slot，新状态写入另一侧，旧状态可安全 donate

6.9 Unified Memory 的源码证据

文件：sglang-main/python/sglang/srt/mem_cache/unified_memory_pool.py

关键位置：

  • MambaSubPoolSpec：第 181‑210 行；
  • UnifiedKVPool：第 217 行起；
  • 构造函数：第 222‑280 行。

MambaSubPoolSpec 第 185‑205 行明确描述：conv_state_shapes、temporal_state_shape、conv_dtype、temporal_dtype。

UnifiedKVPool 第 217‑220 行说明两个 sub‑pool 共享一个 uint8 byte buffer；第 232‑243 行限制当前实现为一个向上增长、一个向下增长的两个池。

简化布局：

同一 byte buffer

MHA/MLA pool  --->  空闲区  <---  Mamba pool

注意：这是物理容量统一，不是逻辑对象统一。

6.10 SGLang 恢复流程总览

请求 prefix 命中
↓
Radix 选择/准备 KDA tracking boundary
↓
MambaCheckpointPool.load_to_active(...)
    ├─ temporal int8 dequantize
    └─ conv native-dtype copy
↓
请求使用私有 active slot
↓
到边界时 donate/snapshot
↓
checkpoint pool + radix cache

关键证据文件：mamba_checkpoint_pool.py:245-259、memory_pool.py:1488-1509、schedule_batch.py:2672-2735。

0x07. vLLM 的应对思路

vLLM 将不同状态纳入同一调度框架，但不强行用同一种存储方式：

同一 Scheduler
    ├─ MLA/full-attention: Paged KV blocks
    └─ KDA:
        ├─ recurrent-state blocks/checkpoints
        └─ convolution-state blocks/checkpoints

因此 "统一" 主要指请求调度和生命周期协调；底层仍承认 KV 与 KDA state 的语义差异。

7.1 State block 是什么

State block 是物理内存管理单位，不等于一个 token，也不等于一个完整请求：

State block B
┌─────────────────────────────┐
│ 一段连续执行产生的 KDA state   │
│ recurrent state             │
│ convolution state           │
│ 边界/版本/布局等元数据         │
└─────────────────────────────┘

它服务于分配、回收、引用计数、换入换出和远端传输。一个大 block 中可以登记多个逻辑 Prefix Entry。

在 vLLM 语境中：

  • recurrent state：KDA 扫描到某个位置后得到的长期状态矩阵，携带长上下文摘要；
  • convolution state：短卷积所需的最近若干 token 的局部缓存，携带局部连续性。

前者决定长期记忆，后者决定下一步局部卷积输出。两者都影响后续 token，因此不能只恢复 MLA KV。

常见疑问：vLLM 的 state block 是什么？

它是 state 的物理容器，类似 KV block 的 "内存页" 角色，但里面装的是 KDA 状态快照，不是逐 token K/V。它通常附带 state tensor、覆盖/关联的前缀范围、快照边界、dtype/layout、引用计数、是否可写。一个 state block 可以有多个逻辑恢复点，因此物理块大小不必等于 Prefix Matching 粒度。

7.2 逻辑前缀与物理 state 解耦

如果 Prefix Matching 的边界必须等于 State block 的边界，匹配粒度会被物理分块限制。vLLM 因此把两件事分开：

  • 逻辑层：记录任意已完成的 token 前缀；
  • 物理层：用较大的 state block 存储状态。

二者通过偏移、边界元数据和 state 版本关联。

假设物理 State block 按 4096 token 管理，但实际请求共享到 4480 token：

物理存储：
    State block B: [0 ————————————————— 8192]
逻辑前缀：
    Prefix Entry at 4096
    Prefix Entry at 4480

如果强制 block‑aligned，只能命中 4096；如果逻辑匹配可以落在 block 内部，就能把 4480 作为更长的合法前缀。这里没有把 state 截断成 "半块"，而是在大容器中保存或索引更细的有效边界。

常见疑问：逻辑前缀匹配与物理状态位置解耦是什么意思？

把 "地图上的地址" 和 "仓库里的箱子" 分开。逻辑层可以说 "命中 token 4480"，而物理层可能把 0‑8192 token 放在同一个 state block B。只要 B 内登记了 state after 4480，逻辑层就能把 4480 当恢复点。否则，物理 block 的存在不能自动证明 4480 有可恢复 state。

7.3 Fine‑grained Partial Hit

假设一个 State block 能容纳 128 token 的状态，但请求前缀只命中到 96 token。朴素做法是：因为 128 token 的 block 尚未完整，96 token 不能恢复，只能从更早位置重算。

Fine‑grained Partial Hit 的做法是：仍然使用这个 128‑token 物理 block，但在其中记录多个更细的 Prefix Entry：

State block B:
  entry at 32  -> state after token 32
  entry at 64  -> state after token 64
  entry at 96  -> state after token 96
  entry at 128 -> state after token 128

于是，96 token 可以成为合法恢复点，即使它位于一个尚未填满的物理 State block 中。这里 "partial hit" 不是只恢复半个状态矩阵，而是恢复该 block 内预先记录的某个有效状态快照。

完整过程：

请求 P:
    已执行 0 … 4480
    产生 state S_4480

物理上：
    S_4480 位于大 State block B 内

新请求：
    命中 token 0 … 4480
      ↓
    找到 Prefix Entry(4480)
      ↓
    取出/恢复 S_4480
      ↓
    复制到新请求私有活动槽
      ↓
    继续生成

"Partial" 指逻辑命中位置可能只占物理 block 的一部分，不是允许系统从一个没有对应快照的中间位置硬推断 state。

它允许在一个大 State block 内登记多个细粒度 Prefix Entry。新请求命中 4480 时，不必退回到 block 起点 4096，也不必复制另一份完整物理 block。它直接定位对应 state，再复制成私有活动 state。

关键边界：partial hit 允许 "逻辑边界细于物理容器"，但不允许 "从不存在的快照推断 state"。

常见疑问：Fine‑grained Partial Hit 如何工作？

大 state block 内额外记录更细粒度、链式的 Prefix Entry。比如一个 block 覆盖 4096 token，而实际共享前缀到 4480 token；逻辑匹配可以落在 block 内部的 4480 边界，并关联该位置已经生成的 KDA state，而不必把 Prefix Matching 粗化为 4096 的物理边界。它不是把一个未完成的 state "截半"，而是把多个逻辑恢复点放进较大的物理容器：物理上合并存储，逻辑上细粒度命中。恢复时仍需把对应 state 复制到请求私有空间，再继续递推。

7.4 Interval‑based 与 Marconi‑style Retention

有限缓存预算下，不能保存所有位置的 checkpoint。

  • Interval‑based Retention：按固定或结构化间隔保留，例如每隔若干 token 保存；稳定、可预测、实现简单。
  • Marconi‑style Selective Retention：根据运行时访问频率、前缀热度和复用价值选择保留位置；更灵活，但需要在线统计和淘汰策略。

前者像规划化修路指标，后者像根据交通流量设置热门出口。

Interval-based:
0 ———— K ———— 2K ———— 3K ————
固定、可预测、适合结构化前缀

Marconi-style:
第一次看到前缀：只记访问
第二次命中：才保存 state
热点继续保留，冷点淘汰

Interval‑based 还可以额外保留 Prompt 末尾状态，因为多轮对话常会完整复用上一轮 prompt。Marconi‑style 则避免一次性长 prompt 立刻占据大量 state 缓存。

常见疑问：Interval‑based Retention 与 Marconi‑style Retention 有何区别？

Interval‑based 是预先规定路标：每 K token 一个 checkpoint，Prompt 末尾额外保留。它稳定、容易预测，适合多轮对话等边界清晰的场景。Marconi‑style 是先观察再投资：第一次出现只记录访问、不保存大 state；再次命中证明它是热点，才创建 checkpoint。它能避免一次性 prompt 浪费缓存，但需要在线统计和更复杂的淘汰逻辑。两者可以并用：结构化边界用 interval，动态热点用 selective retention。

0x08. vLLM 源码导航

8.1 源码地图

文件 关键函数 / 类 作用
vllm/models/kimi_k3/nvidia/kda.py KimiK3DeltaAttention.forward、_forward、_store_cache_checkpoints_kernel KDA conv/recurrent 执行、checkpoint 写入
vllm/models/kimi_k3/nvidia/kda_metadata.py _get_aligned_state_indices_kernel、KimiK3KDAMetadataBuilder.build 把 block table 转成 state index，构造 checkpoint metadata
vllm/v1/core/single_type_kv_cache_manager.py MambaManager.find_longest_cache_hit Mamba/KDA 的 prefix lookup 与细粒度命中
vllm/v1/core/kv_cache_coordinator.py KVCacheCoordinator.__init__、_cache_hit_alignment_tokens 多 attention group 对齐、启用 partial hit
vllm/v1/core/kv_cache_utils.py resolve_block_hashes 调用点 保留细粒度 hash 以支持 partial hit

8.2 KDA 执行和 state cache

8.2.1 KimiK3DeltaAttention.get_kv_cache_spec

文件：vllm-main/vllm/models/kimi_k3/nvidia/kda.py:593-599

函数：get_kv_cache_spec(vllm_config)

关键逻辑：

spec = super().get_kv_cache_spec(...)
return replace(
    spec,
    num_prefill_checkpoint_blocks=
        int(self.kda_prefill_backend == "flashkda"),
)

KDA 的 cache spec 不只是普通 attention KV spec，还包含 prefill checkpoint block 数量。

8.2.2 forward 与 _forward

文件：

  • vllm-main/vllm/models/kimi_k3/nvidia/kda.py:601-638
  • vllm-main/vllm/models/kimi_k3/nvidia/kda.py:640-710

KimiK3DeltaAttention.forward 第 601‑638 行：

投影 hidden_states
拆出 mixed_qkv、gates、beta
调用 _forward
o_proj

KimiK3DeltaAttention._forward 第 640‑710 行：

  1. 从 forward context 取得 KDA metadata；
  2. 读取 state_indices、checkpoint；
  3. 从 self.kv_cache 取得 conv_state 和 recurrent_state；
  4. 纯 decode 条件满足时调用 fused KDA decode。

第 681‑685 行直接解包：

conv_state, recurrent_state, *recoverssm_records = self.kv_cache

第 694‑710 行把 conv_state、recurrent_state、state_indices 传入 fused kernel。

vLLM 源码直接证明：KDA cache 至少有 convolution state 和 recurrent state 两个独立 tensor 语义。

8.2.3 prefill 中的 convolution 与 recurrent state

文件：vllm-main/vllm/models/kimi_k3/nvidia/kda.py:820-860

关键逻辑：

  • 第 831‑846 行定义 _prefill_conv；
  • 第 848‑850 行分别对 q/k/v 做 causal convolution；
  • 第 858‑860 行用 gather_initial_states 取得初始 recurrent state。

简化：

prefill input
    ├─ q -> causal conv -> q
    ├─ k -> causal conv -> k
    ├─ v -> causal conv -> v
    └─ gather initial recurrent state

8.2.4 _store_cache_checkpoints_kernel

文件：vllm-main/vllm/models/kimi_k3/nvidia/kda.py:236-297

函数：_store_cache_checkpoints_kernel(...)

关键逻辑：

  • 第 262‑266 行读取 state index 和 checkpoint offset；
  • 第 267‑285 行把对应卷积历史写入 conv_state；
  • 第 287‑295 行把 recurrent checkpoint 写入 recurrent_state。

vLLM 的 checkpoint 写入不是抽象描述，而是 kernel 级别同时处理：

conv checkpoint -> convolution state
recurrent checkpoint -> recurrent state

8.3 metadata 如何找到 state

8.3.1 _get_aligned_state_indices_kernel

文件：vllm-main/vllm/models/kimi_k3/nvidia/kda_metadata.py:55-101

函数：_get_aligned_state_indices_kernel(...)

第 85 行根据 seq_lens 和 CACHE_BLOCK_SIZE 计算第一个 state slot；第 89‑100 行从 block table 读取 state indices。

它做的是：

请求长度
↓
对应 KDA state block 的起始位置
↓
state index tensor

8.3.2 _mamba_get_block_table_tensor

文件：vllm-main/vllm/models/kimi_k3/nvidia/kda_metadata.py:104-140

函数：_mamba_get_block_table_tensor(...)

若 mamba_cache_mode 是 "all" 或 "none"，第 110 行直接返回 block table；否则第 114‑140 行创建 state indices，并通过 Triton kernel 对齐到 KDA state slot。

这就是 "逻辑前缀" 和 "物理 state block" 之间的适配层。

8.3.3 KimiK3DAMetadataBuilder.build

文件：vllm-main/vllm/models/kimi_k3/nvidia/kda_metadata.py:560-717

关键逻辑：

  1. 第 570‑575 行：只有 prefill、启用 prefill checkpoint block、并且 cache mode 为 align 时，才准备 checkpoint metadata；
  2. 第 588‑600 行：按 sequence length、block size 计算 checkpoint split 和 checkpoint column；
  3. 第 602‑624 行：把 offsets 和 state indices 组织成 KDACheckpointMetadata；
  4. 第 716 行：把 checkpoint metadata 放进最终 KimiK3DAMetadata。

这段源码说明一个重要边界：checkpoint 不是所有运行模式都自动生成；它受 prefill backend、cache mode 和对齐条件共同控制。

8.4 Fine‑grained Partial Hit 的源码证据

8.4.1 MambaManager

文件：vllm-main/vllm/v1/core/single_type_kv_cache_manager.py:1268-1345

类：MambaManager

第 1268‑1295 行：

  • supports_fine_grained_hash_lookup = True
  • 保存 mamba_cache_mode
  • 在 align 模式下维护 checkpoint block 和 partial‑tail 相关状态。

8.4.2 find_longest_cache_hit

同文件第 1297 行起：MambaManager.find_longest_cache_hit(...)

第 1315‑1321 行调用 resolve_block_hashes，并传入：

supports_fine_grained_hash_lookup=True
alignment_tokens=alignment_tokens

第 1327‑1345 行在 alignment_tokens < block_size 时，把一个大 block 拆成多个更小 hash unit 来寻找命中。

这就是 "物理 block 大、逻辑 prefix 小" 的源码基础。

8.4.3 Coordinator 如何启用 partial hit

文件：vllm-main/vllm/v1/core/kv_cache_coordinator.py:622-662

关键逻辑：

  • 第 622‑637 行检查是否存在 MambaSpec 且 mamba_cache_mode == "align"；
  • 第 638‑651 行检查其他 cache manager 是否兼容；
  • 第 655‑662 行让 _cache_hit_alignment_tokens 在 partial hit 时使用 hash_block_size，否则使用 scheduler_block_size。

换句话说：

普通命中中：只能落在 scheduler block
KDA partial hit：可以落在更细的 hash block

8.5 Retention 的源码证据

文件：

  • vllm-main/vllm/v1/core/kv_cache_coordinator.py:32-62
  • vllm-main/vllm/v1/core/kv_cache_coordinator.py:152-158

函数：_validate_prefix_cache_retention_interval(...)

第 40‑45 行说明 retention 只对 SlidingWindow 和 Mamba/linear‑attention group 生效；第 56‑60 行要求 interval 非负且是 scheduler block 的倍数。

Coordinator 第 152‑158 行保存 retention_interval。

由此可知：

retention_interval = None:
  不做稀疏化/保持默认密度
retention_interval = 0:
  对适用 group 采用特殊的最小保留策略
正数:
  按 scheduler block 对齐的周期保留

源码能够确认 interval‑based retention；Marconi‑style selective retention 若要继续考据，应在具体分支中定位其独立策略实现，不能仅凭 coordinator 的 interval 参数称二者是同一个函数。

8.6 vLLM 恢复流程总览

prefix hash lookup
↓
MambaManager.find_longest_cache_hit(...)
↓
fine-grained alignment / state block lookup
↓
KimiK3DAMetadataBuilder.build(...)
↓
state_indices + checkpoint offsets
↓
KimiK3DeltaAttention._forward(...)
    ├─ conv_state
    └─ recurrent_state
↓
kernel decode/prefill/update

关键证据文件：single_type_kv_cache_manager.py:1297-1345、kv_cache_coordinator.py:622-662、kda_metadata.py:570-624、kda.py:640-710。

0x09. 深入讨论

9.1 为什么 SWA 的 checkpoint 方法不能直接用于 KDA

SWA (Sliding Window Attention) 只需要最近窗口内的 token。窗口向前移动时，旧 token 自然离开；如果保存每隔 128 token 的窗口边界，恢复时最多重算一个窗口之间的短区间。

窗口 [t-W, …, t]

窗口向前移动后，旧 token 被明确丢弃。保存一个窗口边界就足以说明附近的局部状态。

KDA 的 recurrent state 则是全历史压缩后的结果。位置 n 的 state 不是只由最近窗口决定，而是由从开头到 n 的所有转移和写入共同决定：

S_n = M_n … M_1 S_0 + Σ_{i=1}^{n} M_n … M_{i+1} b_i

少一个 checkpoint，就不能从任意附近的 token 直接推导出完整 state；必须从更早 checkpoint 重新顺序执行所有中间更新。

而且 KDA state 往往比 SWA 的窗口边界大得多。固定间隔越密，显存成本越高；间隔越稀，命中后重算越长。这就是 KDA 比 SWA 更难做 Prefix Cache 的根本原因。

常见疑问：为什么 KDA 不能用 SWA 的固定间隔 checkpoint?

SWA 的状态只代表一个有限窗口；KDA state 代表整个前缀的递推摘要。SWA 从最近窗口边界继续，最多补一小段窗口计算；KDA 从 checkpoint 到命中点之间的每个 token 都可能改变 state，必须按序 replay。再者，KDA state 本身体积大，固定高频保存的成本远超 SWA 的轻量边界信息。

9.2 Checkpoint 密度的两难

设 checkpoint 间隔为 K：

  • 存储成本大致随 1/K 增加；
  • 命中后的平均重算长度大致随 K 增加。

这不是精确性能公式，但能说明趋势：

checkpoint 密集 -> state 占用高、恢复精细、重算少
checkpoint 稀疏 -> state 占用低、恢复粗糙、重算多

假设只每 10k token 保存一次：

checkpoint: 0 ——————— 10k ——————— 20k
共享前缀：  0 ——————————————————— 15k

两个请求实际共享到 15k，但最近合法恢复点只有 10k。系统不能从 "命中的 15k token 文本" 直接恢复，只能：

恢复 S_10k
replay 10k … 15k
再继续新请求

本来只需计算 5k 尾部，现在要重算 5k；若命中边界更不巧，浪费会更大。

因此真正目标不是最大化 "token prefix hit rate"，而是最大化：

缓存收益 − state 驻留成本 − 恢复重算成本 − 管理复杂度

这不是单纯的命中率问题，而是 state 驻留成本 + replay 成本 + 请求等待时间 的联合优化。

常见疑问：如何理解 checkpoint 密度的两难？

密集 checkpoint 用显存换恢复速度，稀疏 checkpoint 用重算时间换显存。最佳策略取决于前缀复用概率、state 大小、请求长度和显存预算，不存在脱离负载的固定最优间隔。系统 prompt、多轮对话和 Agent 工具定义通常值得高优先级保存；一次性的随机长 prompt 则不一定值得保存。

常见疑问：将 checkpoint 间隔拉大有什么问题？

间隔变大后，前缀可能在 token 层面命中很多，但在 state 层面只能恢复到更早位置。于是命中收益被 replay 吞掉。例如共享 905k token，但最近 checkpoint 在 900k，仍需重算 5k；若只在 800k 保存，就需重算 105k。token 命中率看起来很高，实际 TTFT 未必改善。

9.3 跨实例复用与 PD/EPD 分离

跨实例复用

传统系统只需寻找最长匹配 token prefix，再迁移 KV blocks。KDA 要求同时满足：

  1. MLA KV 在该边界可用；
  2. KDA recurrent state 在该边界可用；
  3. convolution state 在需要时可用或能被安全重建；
  4. 所有状态来自同一模型版本、同一 tokenization 和同一前缀。
最长 token prefix
    ↓
找到共同合法边界
    ↓
联合迁移：
    MLA KV
    KDA recurrent state
    KDA convolution state
    ↓
恢复到目标实例

不同来源的缓存不能简单拼接。一个来源的 MLA KV 加另一个来源的 KDA state，即使长度相同，也不保证计算一致。不能把实例 A 的 MLA KV 与实例 B 的 KDA state 按长度拼在一起，因为它们可能来自不同 tokenization、模型版本、并行布局或前缀边界。

所以跨实例缓存查找从最长前缀搜索变成了 "最长可联合恢复前缀" 搜索。

常见疑问：跨实例复用为何不能只寻找最长 token prefix?

因为还要找到该前缀对应的 KDA state checkpoint，并确认它和 MLA KV 在同一位置、同一版本、同一布局下有效。不同实例的缓存不能按长度简单拼接。跨实例不仅要问 "文本相同到哪里"，还要问：哪个边界有 KDA checkpoint？MLA KV 和 KDA state 是否同边界？远端对象是否来自同一模型 / 并行配置？恢复后的 state 能否安全转为本地活动副本？

PD 分离

传统 PD 分离传输：

Prefill → KV blocks → Decode

Prefill 端产生的不是一组 KV block，而是一组带边界和版本信息的模型状态。Decode 端需要知道：

  • 状态类型；
  • 对应 token 范围；
  • tensor layout 和 dtype；
  • checkpoint 是否可写；
  • 后续是否需要 COW；
  • 状态是否适用于当前并行切分。
Prefill
    ├─ MLA latent history
    ├─ KDA recurrent state
    ├─ KDA convolution state
    └─ boundary/version/layout metadata
    ↓
Decode

这要求通信协议从 "搬运缓存页" 升级为 "搬运可恢复模型状态"。若为每种模型单独写通信协议，基础设施会被模型结构大量污染；统一的模型状态协议能把 "对象类型" 和 "恢复元数据" 作为一等公民。

常见疑问：PD 分离为何要从 KV Cache 传输演进为模型状态传输？

因为 Decode 需要的输入不再只有 KV，而是 MLA KV + recurrent state + convolution state。每类状态可能有不同布局、精度、所有权和恢复规则。Prefill 到 Decode 的输入现在包含三者及元数据。把它们分别塞进 KV 协议会让通信层被模型细节绑死，因此更自然的抽象是统一传输 "可恢复模型状态"。

多模态 EPD

多模态推理还要把 Encoder 输出的视觉 embedding 纳入数据平面：

Encoder embeddings
    ↓
Prefill 生成 MLA/KDA states
    ↓
Decode 使用活动副本

视觉 embedding 可能在多个请求间复用，也可能只服务当前请求；它不是 KDA state，也不能简单当成 KV block。

数据平面要支持：视觉 embedding、MLA token history、KDA recurrent snapshot、KDA convolution snapshot，并分别维护生命周期和所有权。

统一数据平面需要同时处理不同生命周期、不同大小、不同可共享性的对象，而不是把所有内容都伪装成 KV。

常见疑问：多模态 EPD 的额外挑战是什么？

除了 Prefill 生成的注意力状态，还要流动 Encoder 的视觉 embedding。它们大小、生命周期和复用方式不同，系统需要统一数据平面，同时保留类型和所有权信息。EPD 比 PD 多了 Encoder 输出，统一数据平面必须保留类型、所有权、版本和恢复边界，不能只依赖一个 token prefix hash。

9.4 概念与源码对应表

概念 SGLang 源码 vLLM 源码
recurrent state MambaCheckpointPool.temporal self.kv_cache 中的 recurrent_state
convolution state MambaCheckpointPool.conv self.kv_cache 中的 conv_state
活动副本 MambaPool/allocator slot request 对应 state slot
checkpoint MambaCheckpointPool slot state block + metadata/checkpoint slot
COW load_to_active 恢复到请求自己的 state slot
Donate donate_mamba_ping_pong_slot state block/slot 生命周期管理
细粒度命中 Radix/KDA tracking boundary supports_fine_grained_hash_lookup
稀疏 checkpoint schedule boundary + tracking interval retention interval / state retention
物理统一内存 UnifiedMemoryPool Hybrid cache manager / common block pool
kernel 状态索引 backend cache indices state_indices metadata

9.5 复现源码时的最小阅读顺序

如果要自己实现一个教育版 KDA Infra，建议按这个顺序阅读：

  1. SGLang mamba_checkpoint_pool.py:117-174：先理解量化、加载和写回；
  2. SGLang memory_pool.py:1271-1289,1488-1509：理解活动 slot 与 Donate；
  3. vLLM kda.py:640-710：理解 conv/recurrent state 如何进入 kernel；
  4. vLLM kda_metadata.py:570-624：理解 checkpoint offsets/state indices；
  5. vLLM single_type_kv_cache_manager.py:1297-1345：理解细粒度命中；
  6. SGLang unified_memory_pool.py:181-280：理解物理容量共享。

0x10. 最终结论

从源码看，KDA‑aware Infra 的最小正确抽象不是 prefix → KV blocks，而是：

prefix boundary
→ token/hash identity
→ MLA KV blocks
→ recurrent state slot
→ convolution state slot
→ writable/readonly ownership
→ dtype/layout/version metadata

KDA 对 AI Infra 的核心影响可以概括为：

Prefix Cache 不再只是 "找到相同 token，然后搬 KV"；它变成 "找到共同前缀边界，并恢复一组语义一致、生命周期不同的模型状态"。

SGLang 更强调在 Radix Tree 上增加 KDA‑aware checkpoint、稀疏 / 压缩保存和共享状态隔离，实现重点是 "Radix + Mamba checkpoint pool + ping‑pong/COW + UnifiedKVPool"；

vLLM 更强调 state block、逻辑边界与物理存储解耦，以及细粒度 partial hit 和状态保留策略，实现重点是 "Hybrid cache manager + MambaManager fine‑grained hash alignment + KDA metadata/kernel pipeline"。

两者方向不同，但共同结论一致：面向 KDA 的推理系统必须从 KV‑centric cache 演变为 state‑centric cache。只有在同一个前缀边界上，完整的模型状态都可定位、可恢复、可安全写入，Prefix Cache 才真正成立。

TransFormer-封面

0xFF 参考

posted @ 2026-11-05 20:37  罗西的思考  阅读(0)  评论(0)    收藏  举报