arXiv'26 | Kimi K3 的训练与 Serving Infra：KDA 换掉 attention 后，系统要重做什么

September 02, 2026

2026   ·   LLM   MoE   分布式训练   Serving   KV Cache   论文解读  

arXiv’26 | Kimi K3 的训练与 Serving Infra：KDA 换掉 attention 后，系统要重做什么

原文：Kimi K3: Open Frontier Intelligence

这是 Kimi K3 系列的第三篇。前两篇分别讲了 KDA/Hybrid Attention，以及 AttnRes/Stable LatentMoE。本篇看系统：当模型不再只有一种 attention、MoE 专家多到 896 个、上下文拉到 1M token 时，训练和 serving 系统原来默认的假设哪些会失效。

系列导航：导读 ｜ 第一篇：KDA 与 Hybrid Attention ｜ 上一篇：AttnRes 与 Stable LatentMoE

1. 先把并行术语放到同一张图里

训练 2.8T 参数模型时，单卡肯定放不下，因此要把模型、数据和层拆到多张 GPU。常见的几种拆法是：

  • TP（Tensor Parallelism）：把一个矩阵切到多张卡，每张卡算一部分列或行；
  • EP（Expert Parallelism）：不同 GPU 保存不同的 MoE experts，token 通过通信送到对应专家；
  • PP（Pipeline Parallelism）：把连续的模型层切成多个 stage，数据像流水线一样经过这些 stage；
  • DP（Data Parallelism）：每张卡/每组卡保存模型副本，处理不同 batch，再同步梯度；
  • CP（Context Parallelism）：把一条很长的序列切给不同 GPU，各卡负责不同的上下文片段。

K3 的训练栈把这些方式组合起来，再加上 ZeRO（把 optimizer state、梯度等分片保存）和 offload。组合的原因不是“并行方式越多越先进”，而是每种切分正好解决不同的容量或通信问题：TP 解决单个矩阵太大，EP 解决专家太多，PP 解决层数太深，CP 解决序列太长。

2. MoonEP：为什么专家并行最怕动态负载

2.1 一个 batch 就能制造负载倾斜

假设有 4 个 experts，8 个 token，每个 token 只选 1 个 expert。一次 router 可能给出：

expert 1: 4 tokens
expert 2: 3 tokens
expert 3: 1 token
expert 4: 0 token

如果 4 个 expert 分别放在 4 张 GPU 上，GPU 1 要处理 4 个 token，GPU 4 没有工作；即使 GPU 2、3 算完了，整个 batch 也要等 GPU 1。专家并行的吞吐由最忙的 rank 决定。

通常 token 会先通过 all-to-all 发送到目标专家所在的 GPU，再做 grouped GEMM（把不同专家的 token 合并成一批矩阵乘法）。如果每步 token 数不稳定，通信 buffer 和 GEMM shape 也跟着变化，host-device synchronization 和动态 kernel launch 会进一步拖慢训练。

2.2 MoonEP 的思路：先规划，再迁移专家

MoonEP 是 K3 报告中的专家并行方案。它允许把热门 expert 的副本迁移到其他 rank，让 token 可以被多个副本分担。目标不是“平均差不多”，而是让每个 EP rank 在一层恰好接收 $S\times K$ 个 token，其中 $S$ 是序列长度，$K$ 是每个 token 选择的专家数。

这个固定数量带来两个系统收益：

  1. 每个 rank 的 GEMM shape 在编译或计划阶段就知道，不需要等 GPU 统计完 token 数再决定下一步；
  2. 通信 buffer 可以按 $S\times K$ 预分配，而不用按最坏情况扩成 $S\times K\times R$，其中 $R$ 是 EP rank 数。

MoonEP 的 planning kernel 会先计算每个 token 的目的地和专家分组位置，token 直接写入远端卡上已经排好的 buffer。后续通信 buffer 可以用 view 交给计算 kernel，避免再做一次 copy。

论文还给出一个冗余槽位的存在性结论：只要每个 rank 预留约 $E/R$ 个冗余 expert 槽位（$E$ 为 expert 数，$R$ 为 rank 数），就能找到满足均衡的专家迁移方案。这个 bound 说明“总能找到方案”的条件，但真实收益仍取决于迁移成本、expert locality 和通信拓扑。

Kimi K3 的整体架构

2.3 为什么固定 shape 对大模型很重要

在小模型里，多一次同步可能只是一点微秒；在 93 层、数万步的训练里，微秒会乘上层数、pipeline stage 和训练步数。MoonEP 把 token 数固定后，很多 kernel 可以预先编译，CUDA Graph 也更容易复用，减少了“GPU 做完一件事，CPU 才知道下一件事要多大”的等待。

所以 MoonEP 的价值不只在负载均衡公式，而在于把动态路由问题变成固定 shape 的执行问题。专家是否热门仍然会变化，但每个 rank 最终接收多少 token 被系统约束住了。

3. 2.8T 模型的显存：不是只有权重放不下

训练时显存通常被四类东西占满：模型参数、激活、梯度和 optimizer state。K3 的做法不是给所有 tensor 规定同一种存储方式，而是给 tensor 单独声明策略：重计算、FP8 量化、offload 到 CPU，或者 offload 到其他 GPU，都可以组合。

3.1 激活：用重计算换空间

如果保存每层所有中间激活，反向传播时可以直接读取，但显存随层数和序列长度增长。重计算（recomputation）则只保存检查点，反向时重新跑一小段前向，用额外计算换掉激活存储。

K3 对大块激活使用 block-wise FP8 量化再 offload，对逐元素算子更多采用重计算。量化减少传输字节数，重计算减少需要传输的 tensor 数量，两者针对的是不同的成本。

3.2 MoE 反向：不要保存完整 dispatch 中间量

MoE 前向要把 token 按 expert 重新排列，得到 dispatch buffer；反向如果把所有排列后的中间激活完整保存，显存和通信都很可观。K3 利用梯度表达式，只保存必要的 dispatch 输入，反向时重算排列和 GEMM，并与通信流水线重叠。

3.3 AttnRes 并不一定让显存线性增加

Full AttnRes 确实需要保留多层表示，但 K3 使用 Block AttnRes。一个 block 内的输出先汇总，后续层只访问少量 block 表示；把这段计算纳入 checkpointing 后，每层保存的激活量可以接近标准残差架构，而不是保存 93 份独立层输出。

4. KDA 的状态怎么做 Context Parallelism

KDA 的优势是状态 $\mathbf S$ 不随序列长度增长，难点则是状态更新具有顺序依赖。假设把一条长序列切成两段：

GPU 0: token 0 ... token 49999
GPU 1: token 50000 ... token 99999

GPU 1 不能从空状态开始算，因为 token 50000 的结果依赖前 50000 个 token 的状态。如果直接把整段状态传过去，通信会变重；如果让 GPU 1 等 GPU 0 逐 token 传状态，计算又被串行化。

4.1 普通线性 attention 为什么可以简单相加

最简单的线性 attention 更新是：

\[\mathbf S_t=\mathbf S_{t-1}+\bm k_t\bm v_t^\top\]

每段从零开始算出本地贡献，最后把两段状态相加即可，因为更新是线性的。

KDA 不满足这个条件。它的更新包含依赖当前 token 的转移算子：

\[\mathbf S_t=\mathbf M_t\mathbf S_{t-1}+\mathbf b_t\]

其中 $\mathbf M_t=(\mathbf I-\beta_t\bm k_t\bm k_t^\top)\operatorname{Diag}(\bm\alpha_t)$。后半段的 $\mathbf M_t$ 不是固定矩阵，而是由 token 决定，因此不能把两段本地状态简单相加。

4.2 KCP：传固定大小的转移信息

Kimi 的 KDA Context Parallelism（KCP）让每段 GPU 本地计算两个固定大小的 fragment：

  1. 这一段的累积转移矩阵 $\mathbf M$；
  2. 假设输入状态为零时，这一段自身产生的状态 $\tilde{\mathbf S}$。

对于两段，第二段真实入口状态可以写成：

\[\mathbf S_{in}^{(2)}=\mathbf M^{(1)}\mathbf S_{in}^{(1)}+\tilde{\mathbf S}^{(1)}\]

各 GPU 交换这些固定大小 fragment 后，做一次 prefix scan，就能恢复每段真实的入口状态，再计算输出。通信不再随着整条序列保存每个 token 的状态而增长，而是与状态矩阵大小相关。

这和“把序列直接切两半、各算各的”不同：KCP 传的是每段如何变换状态的信息，最后仍然遵守原始 token 顺序。

5. 1M context 的 Agentic RL：状态要能暂停和恢复

Agentic RL 是让模型在环境中多轮调用工具、执行任务，再根据结果训练。一个 1M token 的 rollout 可能包含几百次工具调用，系统不只要保存模型 KV，还要保存文件、容器和工具环境的状态。

5.1 Cache 池：只为离开 GPU 的前缀付费

正在 decode 的请求需要 GPU 上的 KV；暂时等待训练或工具返回的前缀，如果一直占 GPU，cache 很快爆掉。K3 使用 write-back 策略：前缀被 GPU 驱逐时才写回 CPU DRAM，下次重新调度前再预取。

KDA 状态与对应 MLA KV block 一起 offload，因为它们必须在相同的 token 边界恢复。模型权重和 optimizer state 还会继续 offload 到 NVMe，把 DRAM 主要留给 cache 池。

5.2 为什么沙箱要用 microVM

Agent 可能会创建文件、运行代码、启动容器，甚至触发内核级操作。普通容器隔离如果不够强，一次异常就可能影响整台训练机器。K3 使用 Firecracker microVM，把每个 AgentENV 当作可暂停、可复制的轻量虚拟机。

它支持增量 checkpoint，只保存脏页；报告中的 checkpoint 约 133ms、resume 约 49ms。Agent 等待工具结果时，沙箱可以暂停，不继续占用 CPU 和内存；需要比较不同动作的 reward 时，还可以 fork 出副本，在副本里试错而不影响原环境。

这类机制对 1M context 很关键：长任务的瓶颈不只是 token 数，还包括“一个任务的外部世界状态要活多久”。

6. Prefix Cache：KDA 让已经成熟的 cache 又变复杂

传统 KV cache 可以按 token 分页，任意前缀边界都能复用。KDA 的状态是一个聚合后的矩阵，只有保存过 checkpoint 的边界才能恢复。于是一个请求的前缀可能出现这种情况：MLA 部分已经命中，但 KDA 状态没有对应 checkpoint，不能直接从任意位置继续。

K3 用两级粒度解决这个问题：

  • 逻辑匹配在 512-token hash block 上进行，保证 prefix 可以细粒度比较；
  • 物理内存用较大的 block 分配，例如一个 6144-token physical block 里包含 12 个 hash block；
  • KDA checkpoint 只在稀疏边界保存，通常优先选择对话轮次边界；
  • 命中后恢复最近 checkpoint，中间未覆盖的 token 用 copy-on-write 继续 prefill。

KDA 感知的 prefix cache

图中蓝色部分是已缓存的 MLA hash block，下面的点表示 KDA checkpoint。请求可以复用前面若干 MLA block，同时恢复最近的 KDA 状态，再从命中点继续计算，而不必把整段前缀重新跑一遍。

这里的关键不是“缓存粒度越细越好”。Checkpoint 本身比单个 token 的 KV 大得多，存得太密会抵消 KDA 的空间优势；存得太稀又会增加恢复时需要重算的 token。K3 的设计是在匹配粒度和状态保存成本之间折中。

7. Speculative Decoding：KDA 状态不能随便回滚

Speculative decoding（推测解码）用一个较小的 draft model 一次猜多个 token，再让大模型验证。接受的 token 被写入正式状态，被拒绝的 token 则需要丢弃。

在普通 KV cache 里，回滚通常是释放被拒绝 token 对应的 cache block；KDA 每生成一个 token 都会更新同一个状态矩阵，如果 draft 连续走了 7 步，再拒绝第 3 步，状态已经多更新了几次，不能只改一个长度指针。

K3 的思路是：draft 阶段只缓存每个 token 的 q/k/v 投影输入，这些数据比完整状态快照小；验证 kernel 确认哪些 token 被接受后，在片上重放 accepted token 的状态更新，最后只写回合法的状态。这样避免为每个 draft token 保存一份完整 $\mathbf S$。

8. 量化和 vLLM：训练时就面对部署格式

K3 的 post-training 使用 MXFP4 量化感知训练：MoE expert 权重以 MXFP4 保存/计算，激活使用 MXFP8，attention projection 和 router 保持更高精度。量化感知训练的意思是，训练过程中就模拟部署时的低精度，而不是训练完成后突然把 BF16 权重转换成 MXFP4。

这样做是为了减少 train-inference mismatch：训练时模型已经见过量化误差，rollout 和训练使用同一套数值格式，部署后不容易出现突然掉点。

开源 serving 侧，vLLM 已提供 K3 的 Day-0 支持。一个最简启动方式是：

vllm serve moonshotai/Kimi-K3 \
  --tensor-parallel-size 8 \
  --trust-remote-code \
  --load-format fastsafetensors \
  --enable-prefix-caching \
  --enable-auto-tool-choice \
  --tool-call-parser kimi_k3 \
  --reasoning-parser kimi_k3

这里的 tensor-parallel-size 8 表示用 8 张 GPU 切分张量；enable-prefix-caching 打开前面讲的混合 cache；tool-call 和 reasoning parser 则负责把模型输出解析成工具调用和思考/答案结构。

vLLM 的 scheduler 需要同时管理两种 cache：MLA 的分页 KV block 和 KDA 的固定状态 block。prefill 使用 FlashKDA，decode 使用融合 CUDA kernel。高吞吐部署还可以把 prefill 和 decode 分到不同机器，前者用 TEP（attention TP + MoE EP），后者用 DEP（专家并行），中间通过 NIXL 传输 MLA KV、KDA 状态和 block table。

K3 serving 的计算与通信重叠

这张图反映的是另一个系统事实：当计算、通信和 offload 能够重叠时，单独优化某一个 kernel 未必决定最终性能；调度顺序和 buffer 生命周期同样重要。

9. 系列小结：新架构把哪些老问题重新打开了

KDA 减少了随 context 增长的状态，但引入了状态传递、checkpoint 和回滚问题；AttnRes 提供了深度方向的检索，但要求系统管理跨层表示；LatentMoE 压低了专家计算成本，却要求更严格的负载均衡和激活稳定性。

对应到 infra，就是一串连锁反应：

  1. MoonEP 把动态专家路由约束成固定 shape，减少 EP 的尾部等待；
  2. KCP 用固定大小 fragment 做 KDA 的 context parallelism；
  3. 混合 prefix cache 同时管理分页 KV 和稀疏 checkpoint；
  4. speculative decoding 通过重放状态避免保存大量快照；
  5. 量化感知训练让训练、rollout 和 serving 使用接近的数值格式；
  6. vLLM 把这些结构暴露成 scheduler、kernel 和部署参数，而不是只停留在论文里的模块图。

Kimi K3 这份报告的系统价值，主要就在这里：它没有把“新 attention”“超大 MoE”“1M context”当成互不相干的卖点，而是展示了每个结构变化会把哪些系统假设推翻。真正部署时，决定成败的往往不是某个公式单独快多少，而是状态、通信、内存和调度能不能在同一条 pipeline 里对上。

K3 的 infra 细节和 LLM 压缩不是同一件事。《动手学 AutoML：从 NAS 到大语言模型优化实战》第 8、11 章整理了剪枝、量化和后训练压缩，适合从模型压缩角度补充“如何降低模型本身的成本”。

动手学AutoML书籍封面

Flag Counter