部署 Kimi K3：从 KDA Kernel 到 Fleet 调度的推理 Infra 全景

本文是 Kimi K3 系列的第三篇，接续 《从线性注意力到 KDA》 和 《从残差到分位数平衡》，主要基于 Moonshot AI 发表的技术报告 Kimi K3: Open Frontier Intelligence（arXiv 2607.24653）第 5 章 Infrastructure 与 §4.1.4 整理。前两篇讲的是"模型长什么样、数学怎么推"，这一篇讲"这个模型部署上线时，工程侧要专门解决哪些问题"。文中数字如无特别说明均引自官方技术报告。

Serving Kimi K3 要解决三个问题：混合 KDA–MLA 架构维护着两种根本不同的缓存，必须在百万 token 上下文规模下联合管理；KDA、AttnRes、Stable LatentMoE 这些新模块和高度稀疏的专家都需要专门定制的 kernel；生产流量里单请求成本能相差三个数量级，普通调度策略会失效。对应的解法分三层：engine 级的 KDA-aware prefix cache、device 级的专用 kernel、fleet 级的调度策略。本文按请求生命周期的顺序展开：先 prefill，再缓存，再 decode，最后是集群调度。

一、Prefill 阶段：KDA 的 kernel 设计

1.1 FlashKDA：chunk 内计算和跨 chunk 状态传播重叠

KDA 的 chunkwise 算法（详见第一篇的推导）本质是"chunk 内并行矩阵乘法 + chunk 间串行传状态"。朴素实现会让这两个阶段交替执行——算完一个 chunk 的内部矩阵乘法，SM 就要空等状态从上一个 chunk 传过来。K3 的解法是 FlashKDA：一个基于 CUTLASS 的 chunkwise kernel，专门把 chunk 内计算和跨 chunk 状态传播重叠起来执行，拆成"token 并行阶段"和"head 并行的递归阶段"两部分分别调度调优。这个 kernel 同时服务训练和推理 prefill，是 flash-linear-attention 库的自动分发后端。

1.2 单机内 SM 级上下文并行——跟跨设备 KCP 是两个不同层级

问题：张量并行（TP）只是把不同 head 分给不同设备，不会缩短递归本身的长度。如果某个 rank 只分到几个 head，处理超长序列 prefill 时，这个 rank 的大部分 SM 会在"等状态传过来"这件事上闲置。

解法：利用一个关键观察——一个片段的状态转移，可以先不依赖"传入状态是什么"独立算出来，之后再精确地和传入状态组合（这正是下一节 KCP 数学的核心思想，只是这里应用在单个 rank 内部）。K3 设计了一个自动的 SM 级上下文并行 planner，把序列切分给同一个 rank 内部的多个 SM（不是切给不同设备），并行算出每一段的转移，再精确合并还原出每段真正的初始状态。跟跨设备的 KCP 相比，这一层完全在单个设备内部完成，不产生跨设备通信。

二、跨设备长上下文 prefill：KDA Context Parallelism（KCP）完整数学

标准 softmax attention 做上下文并行，要交换的 KV block 大小随序列长度增长。朴素线性注意力（无衰减、无 delta 纠错）反而好办——它的状态更新是纯加法 St=St−1+ktvt⊤S_t=S_{t-1}+k_tv_t^\top，一段序列的贡献可以独立地从 S=0S=0 算出来，之后直接加上传入状态就行（叠加原理成立）。

但 KDA 的更新是 St=MtSt−1+βtktvt⊤S_t = \mathbf{M}_t S_{t-1} + \beta_tk_tv_t^\top，其中 Mt=(I−βtktkt⊤)Diag(αt)\mathbf{M}_t=(I-\beta_tk_tk_t^\top)\text{Diag}(\alpha_t)——转移矩阵先作用在传入状态上，再加新写入。这意味着一段序列产生的效果依赖于传入这段序列时的状态是什么，不能像朴素线性注意力那样"从 0 算一遍，最后直接加总"。

KCP 的解法：把每一段的效果分解成两个只用本地 token 就能算出来的量：

M[i+1]t←1:=∏r←1tMr(本段累积转移矩阵)\mathbf{M}_{[i+1]}^{t\leftarrow1} := \prod_{r\leftarrow1}^{t}\mathbf{M}_r \qquad(\text{本段累积转移矩阵})

S~[i+1]t(假设从 S=0 开始，本段单独产生的状态)\widetilde{\mathbf{S}}_{[i+1]}^{t} \qquad(\text{假设从 } S=0 \text{ 开始，本段单独产生的状态})

递归展开后：

S[i+1]t=S~[i+1]t+M[i+1]t←1 S[i]Ti\mathbf{S}_{[i+1]}^{t} = \widetilde{\mathbf{S}}_{[i+1]}^{t} + \mathbf{M}_{[i+1]}^{t\leftarrow1}\,\mathbf{S}_{[i]}^{T_i}

关键是：右边这两项，每个 rank 完全不知道"传入状态是什么"之前，只用本地 token 就能提前算好——这正是每个 rank 要交换的两个东西。

因为这些"rank 级更新"满足结合律（可以像 scan 操作一样两两合并），每个 rank 的入口状态可以通过一次前缀扫描（prefix scan）还原：每个 rank 先本地算好 (M,S~)(\mathbf{M},\widetilde{\mathbf{S}})，用一次 all-gather 交换给所有 rank，之后每个 rank 按顺序把前面 rank 的这些小矩阵依次组合起来，还原出自己真正的入口状态。这次 all-gather 的通信量是固定大小的（只交换 dk×dkd_k\times d_k 和 dk×dvd_k\times d_v 的小矩阵，不随每个 rank 分到的 token 数增长），这正是"KCP 能做到计算量线性扩展"的原因——这个设计建立在 DeltaNet context parallelism 的基础上，针对 KDA 多出来的衰减门做了扩展。

三、KDA-aware Prefix Cache 完整设计

3.1 统一缓存布局

MLA 的 KV cache 随序列长度增长、按 token 分页；KDA 的循环状态大小固定、每个请求只有一份。K3 把 KDA 状态打包进跟 MLA KV 一样的 paged block pool，统一页的字节大小，让两种页共享同一套分配/引用计数/淘汰逻辑。页内部按 head 连续存储，每个 head 的字节流自成一段，作为跨节点传输的最小单元；在 prefill/decode 分离部署、且两端 TP 度不同的情况下，重新布局的工作放在传输路径上完成，GPU 侧不需要重新洗牌。有个有意思的细节：这种"类型混淆访问"会直接产生乱码而不是看起来合理的数据——相当于一个零开销的正确性自检。

3.2 粒度错配问题：为什么粗粒度会让缓存几乎没用

标准的 block-hash 前缀缓存以一个物理 block 为单位做哈希——只有完整的 block 才会被哈希，只有 block 对齐的前缀才能被复用。问题出在 KDA 身上：KDA 每个序列只维护一份大的循环状态，状态快照只能在稀疏的边界上做才划算——这把共享的 block size 强行拉大到 1024–6144 token。在这么粗的粒度下，缓存几乎失去意义：短于一个 block 的请求永远无法被复用，chunked prefill 在跨过一整个 block 边界之前，导出不了任何可缓存的前缀。

解法：把两种粒度解耦。前缀哈希跑在 MLA 页内部更细的 hash block（比如 512 token）上，物理 block 仍是粗粒度的分配单元；KDA 反过来对齐——只在 MLA hash 端点的一个稀疏子集上保存状态 checkpoint（也只有这些位置是查找时可能被引用到的边界）。

物理 block（分配单元） Hash block（哈希/缓存粒度） KDA checkpoint
粒度 1024–6144 token 512 token 只在部分 hash 端点，常与对话轮次边界重合
作用 内存分配的基本单位 前缀复用判断的最小粒度 决定 KDA 状态能否在这个边界被还原

prefill 过程中，一个部分填充的 MLA 页以"最后一个完整 hash block 的链式哈希"注册进前缀缓存索引——每个哈希覆盖它之前的所有 hash block，命中一个端点就等于验证了到这个端点为止的整段前缀；每次前向传播后，KDA kernel 在处理到的最后一个 hash 对齐位置持久化状态。checkpoint 体积很大，请求推进过程中被淘汰的中间 checkpoint 会被回收，只有落在对话轮次边界上的才保留用于跨请求复用。checkpoint 是只读快照：命中时把状态拷贝进请求私有的运行状态里，绝不在原地修改一个对其它请求可见的 checkpoint。

3.3 两阶段 lookup

查找分两个阶段：MLA 阶段先按链式哈希匹配完整的物理 block，遇到第一个不完整的 block 时退回去检查它内部的 hash 端点（部分填充的页依然可以被命中）；KDA 阶段要求候选边界在每一个 KDA cache group 里都存在对应 checkpoint（每个 group 各自维护一份独立的循环状态）。最终命中的是"同时满足两个阶段"的最长边界——一定是 hash block（512）的整数倍，但不要求是物理 block（6144）的整数倍。举例：一个前 2800 token 命中缓存前缀的请求，会命中在 B=2560=5×512B=2560=5\times512 这个位置——深深嵌在一个 6144-token 物理 block 内部——然后直接从 token BB 继续 prefill，不用重算 [0,B)[0,B)。

3.4 并发调度下的一致性保证

在"命中的 block 同时是共享缓存条目、又是私有请求的增长点"这种场景下，有三个具体的失效模式和对应修复：

  1. 所有 cache group 从同一个共享空闲列表取 block——如果给一个 group 分配私有拷贝时正好淘汰了另一个 group 刚命中的 block，就会出问题。修复：在真正分配任何东西之前，先把命中的 block 在所有 group 里都锁住（pin）。
  2. 拷贝进私有 block 的操作在前向传播之前才在 GPU 上执行——如果一个 block 是在当前调度步内刚被分配/注册的，读它的人可能会拿到上一个所有者遗留的字节（竞态）。修复：这类 block 在拷贝真正落地之前，排除在匹配范围之外。
  3. 一个 checkpoint 只有在每个 KDA group 里都存在时才能恢复一个请求——修复：淘汰一个 group 的 checkpoint 时，原子性地让所有兄弟 group 的对应 checkpoint 一起失效，做到"要么在所有 group 都能命中，要么都不能"。

有了这三条，混合 KDA–MLA 架构的前缀缓存，达到了跟纯 full-attention 模型同等的通用性：任何共享前缀都能在任意 512-token 边界上被复用，跟请求长度、chunking 方式、调度交错顺序都无关。

四、Decode 阶段：投机解码与 KDA 状态的冲突

KDA decode 面临一个跟 prefill 完全不同的挑战：瓶颈不再是"怎么并行"，而是"怎么高效管理一个每步都原地更新的循环状态"。

问题的根源：K3 用 MTP 层做投机解码，会先"抢跑"生成好几个候选 token，KDA 的状态跟着这些候选顺序原地更新，往前推进好几步。等目标模型验证完，如果只接受了其中一部分、拒绝了剩下的，状态已经推进过头了，没法简单地退回去。

朴素方案的问题：给每一个候选位置都存一份状态快照，拒绝时直接回滚——但状态本身是一个 dk×dvd_k\times d_v 的大矩阵，存 KK 份候选就要搬 KK 倍的状态流量，在线上 serving 常见的大 batch 场景下会成为主导开销。

K3 的方案：关键观察是——任意一段被接受的草稿前缀之后的状态，完全由这些草稿 token 的"投影后输入"决定（过完短卷积、归一化之后的 q,k,v,β,αq,k,v,\beta,\alpha），而这些投影值远比状态本身小。所以只缓存这些小的投影输入，在片上重放被接受 token 的递推重建出正确状态，最后只把验证通过的 token + 1 个 bonus token的状态写回。

这个设计跟同期一篇独立工作 ReplaySSM 的思路一致（算是一次不谋而合的验证）。效果上，验证延迟随验证的 token 数次线性增长，且始终低于"存状态快照"的基线方案。还有一个附带好处：这些投影缓存从不离开 decode 阶段，前缀缓存和 prefill-decode 分离部署处理的还是跟非投机解码场景完全一样的数据载体——这个优化跟第三节的 prefix cache 系统是正交的，不需要为了支持投机解码去改动缓存系统本身。

五、Block AttnRes 的 prefill / decode kernel 优化

Block AttnRes 的两阶段调度是：批量的 inter-block pass（每个 block 只读一次缓存的 block 表示）+ 每一层用 online-softmax 把 intra-block 的部分和融合进去。这类 kernel 在 prefill 和 decode 里，开销的大头都是内存访问，所以两边的优化都在抠内存效率，不是抠算力。

Prefill 侧：如果直接实现，每个 TP rank 都会各自 materialize 一份完整的 block 表示，造成大量冗余内存。解法是对激活值采用序列并行（SP）：把 TP 的 all-reduce 拆成一次 reduce-scatter + 一次 all-gather，把 intra-block 的 kernel 插在这两个 collective 之间，让它操作在按序列切分好的隐藏状态上——这样每个 token 的 block 表示只会在恰好一个 rank 上被 materialize，消除了多余的内存占用，也降低了 I/O 开销。

Decode 侧：把 inter-block kernel 放到一个 side stream 上，让它跟主 stream 上的其它独立计算重叠，隐藏掉它的延迟。intra-block 那部分则通过融合来精简：把"AttnRes 输出跟它的部分和更新做合并"，连同后面紧跟的 RMSNorm，一起融合进前面那次 TP all-reduce 里，直接省掉了 intra-block 阶段专门的 kernel launch。

六、Stable LatentMoE 的 decoding kernel

896 个专家、每 token 16 个激活，专家空间和单 token 激活数同时变大，带来的调度/协调开销让传统 MoE kernel 很难维持高硬件利用率。

Latent GEMM 侧的三个优化：

  1. 把 latent 降维投影（W↓\mathbf{W}^\downarrow）跟 MoE router 融合进一个 GEMM（两者输入相同，都是简单的线性投影，合并成一次矩阵乘法可以避免重复读一遍输入、少发一次 kernel）。
  2. latent 权重矩阵跨 rank 分片，用 multimem store 指令把输出的 all-gather 融合进 GEMM 的 epilogue（让矩阵乘法的输出写入操作本身顺带完成跨设备的数据同步，不用再单独起一次通信 kernel）。
  3. 把这部分通信跟其它算子重叠，比如共享专家的计算（共享专家的计算不依赖路由结果，可以放到另一个 stream 上去填满通信等待的时间）。

Routed 专家在小 batch 场景下的问题：单 token 只激活 16/896 个专家，摊到每一个具体专家头上的 token 数往往很少——group GEMM 退化成了"内存带宽受限的流式读权重"，传统那种为算力密集场景设计的 tile-centric kernel 反而不合适。K3 借鉴 WarpDecode 的 token-centric 设计：每个 warp 负责一个输出神经元，直接从内存流式读取对应的权重；进一步把每个 warp 拆成更细的 lane team，各自处理一个不重叠的专家子集，最后做一次跨 lane 的归约。另外，权重的内存布局会做一次性离线重排，大幅降低运行时反量化（MXFP4 → 计算精度）的开销。

七、Fleet 级调度

单个 serving 实例之外，挑战从"单请求效率"变成了"可预测性"——一次 prefix-cache miss 的代价比 hit 高出几个量级，突发的百万 token 请求可能饿死短请求。

缓存亲和调度：1M 上下文场景下，一个典型的编程输入可能携带 40 万 token 的前缀，但只需要增量 prefill 4000 个新 token——一次 cache hit 能省掉整段前缀的重新 prefill，比 miss 便宜好几个量级。所以每个请求会被路由到持有它前缀缓存的那个集群（把缓存本身搬到别的集群，走的是比集群内互联慢得多的跨集群链路）。但这样一来，一个 session 就绑死在单个集群上——那个集群一旦故障，绑定在它上面的所有 session 都会被打断。解法是用一致性哈希把每个 session 同时钉在两个集群上：一个主集群正常服务，一个预先指定的副集群在主集群故障时接管（副集群不持有该 session 的任何前缀缓存，接管时需要重新 prefill）。因为一致性哈希会把不同 session 的副集群分配得很均匀，任何单个集群故障产生的"重新 prefill"工作会被分散到整个 fleet，而不是集中砸到某一个集群上——常规情况下保留缓存局部性，单点故障的影响被限制在一个可控范围内。

基于预算的准入控制：生产流量里，2K token 以下的短请求跟 100 万 token 的超长请求混在一起，单请求成本能相差三个数量级——基于"平均请求"的容量规划、排队模型、限流配额在这种方差下全部失效。典型的失效场景是：一波突发的长上下文请求把算力占满，之后到达的短请求排不上号，导致所有流量的首 token 延迟（TTFT）一起变差（不只是长请求自己变慢）。解法是给不同请求类别分配各自独立的资源预算，让突发的长上下文流量最多只能吃掉自己那一份份额，不会拖累其它类别请求的 SLO。

八、附带一提：EAGLE-3 草稿模型是怎么训练出来的

草稿模型融合了第 1、4、最后一个 AttnRes block 的特征（见第二篇），这里补一个训练目标上的细节，因为它直接决定了部署后投机解码能有多快。

投机解码的加速比由每 token 接受率 ∑x∈Vmin⁡(p(x),q(x))\sum_{x\in\mathcal{V}}\min(p(x),q(x)) 决定（p,qp,q 分别是目标模型和草稿模型的下一 token 分布，这是无损投机采样的理论结果）。问题是：常规的 KL 散度代理损失，并不保证能最大化这个接受率——尤其是当草稿模型容量有限（比目标模型小得多）时，KL 散度的最优解和"最大化 min⁡(p,q)\min(p,q) 求和"这个目标并不对齐。

K3 的解法是直接优化接受率本身的负对数，即 LK loss：

LLK=−log⁡∑x∈Vmin⁡(p(x),q(x))\mathcal{L}_{\text{LK}} = -\log\sum_{x\in\mathcal{V}}\min(p(x),q(x))

p,qp,q 都在温度 1 下评估，不掺任何有监督的交叉熵项——纯粹对齐目标模型自己的分布，而不是去拟合"标准答案"。草稿模型的微调延续跟主模型一致的 QAT 配置（MoE 专家权重 MXFP4，激活 MXFP8，非专家模块保持更高精度）——这跟前面反复出现的"消除训练-推理不一致"是同一个原则：部署后草稿模型的实际行为，要跟训练时验证过的行为完全一致。

九、小结

Serving Kimi K3 的复杂度，根本上来自"KDA 的固定大小状态"和"MLA 的传统 KV cache"这两种性质完全不同的缓存必须被联合调度、联合优化。这条主线贯穿了本文几乎所有设计：

  • Prefill：FlashKDA 把 chunk 内计算和跨 chunk 状态传播重叠；单机内 SM 级 CP 和跨设备 KCP 分别在两个层级解决"递归天生串行"的问题，KCP 靠"局部转移矩阵 + 局部零起点状态"这对可结合的量，把通信量压成固定大小的一次 all-gather。
  • 缓存：KDA-aware prefix cache 把两种缓存打包进统一的分页布局，用"哈希粒度细、checkpoint 粒度稀疏"解耦两种缓存对齐要求不同的矛盾，再用三条并发控制规则保证正确性。
  • Decode：投机解码跟 KDA 原地状态更新天然冲突，解法是只缓存小的投影输入、按需重放重建状态，避开存储整份状态快照的开销。
  • Kernel：Block AttnRes 和 Stable LatentMoE 的 decode kernel 优化都在跟内存带宽和调度开销较劲，分别用序列并行/流重叠和 token-centric 流式设计应对。
  • Fleet：缓存亲和调度和预算隔离，把"单请求效率"的优化成果转化成"整个集群可预测"的服务质量。

每一层解法都不是孤立的技巧堆砌，而是紧扣着 KDA 混合架构和极端稀疏 MoE 这两个核心设计带来的具体约束——理解了模型架构本身（前两篇），再看这些 infra 设计，会发现它们几乎都是"架构决策在部署侧必然要还的债"。

参考资料

  • Kimi Team. Kimi K3: Open Frontier Intelligence. arXiv:2607.24653
  • MoonshotAI/Kimi-K3 — 模型卡与权重
  • Yang et al. Parallelizing Linear Transformers with the Delta Rule over Sequence Length（DeltaNet context parallelism 的基础）
  • Li et al. EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test