SGLang 和 Miles 为月之暗面 2.8T 参数 Kimi K3 模型提供发布当日支持

我们很高兴地宣布，在 SGLang 和 Miles 中对 Kimi K3 的 Day-0 支持：https://platform.kimi.ai/docs/guide/kimi-k3-quickstart。K3 是第一个拥有三万亿参数级别的开源模型，其混合架构在几乎每一个服务栈假设的地方都与传统不同。通过与 Moonshot AI 和 NVIDIA 团队的合作，这两个团队在发布当天全面覆盖了 K3：SGLang 用于推理，Miles 用于强化学习训练。本文介绍了实现这一目标所需的步骤。

启动命令和每个工作负载的配置指南在 Kimi K3 食谱中：https://docs.sglang.io/cookbook/autoregressive/Moonshotai/Kimi-K3。

Kimi K3：https://platform.kimi.ai/docs/guide/kimi-k3-quickstart 是第一个拥有三万亿参数级别的开源模型：参数为 2.8T，拥有 1M token 的上下文窗口，并具备原生视觉理解能力。其前身 K2.5 在架构上与 SGLang 已经很好支持的模型非常接近，因此大部分栈直接适用。K3 并非如此 — 它在多个独立位置同时偏离了传统：

每一行都是一个服务问题。混合注意力栈使得 1M 上下文成为可承受，但这意味着服务器同时持有两种状态：每个请求旁边有一个固定大小的 KDA 状态以及 MLA 的每 token KV — 下文的内存管理部分正是关于这个。注意力残差将一组注意力输出贯穿整个栈，这破坏了 SGLang 标准层工具链的假设，并在诸如 DP 注意力等地方迫使使用 K3 专用路径。LatentMoE 在下投影潜在空间中使用 SiTU 激活函数而非 SwiGLU 来路由 16 个 896 个专家，因此没有现有的 MoE 内核可以开箱即用。视觉路径——塔、投影器、处理器以及 Kimi 的 XTML 媒体格式——都是从零开始搭建的。下面的各节都基于这一搭建过程。

注意 KV 是仅追加的。一旦一个令牌的 KV 被计算出来，它就永远不会改变，这就是调度器可以在每个共享前缀的请求之间共享一个物理副本、将其保存在基数树中，并在迭代中无须第二次考虑就进行流水线处理的原因。KDA 层的状态则恰恰相反，是一个固定大小的递归缓冲区，每个令牌都会在原地被覆盖。因此，调度器免费提供的所有 KV，从前缀缓存到重叠调度器、从推测解码到分页，都必须为一个在读取时会变化的值重建。K3 将 69 个 KDA 层与 24 个 MLA 层交错，因此这并不是一个特殊情况。这占据了模型的大部分。接下来是这种重建，它唯一令人满意的特性是每个部分都由原地覆盖的事实强制生成，没有任何是附加的。

同一路径同时服务于重叠调度器、推测解码和 page_size > 1 的情况，这些组合在以前的递归状态模型中是相互排斥的。一个请求可以击中已缓存的前缀、恢复递归状态、运行多令牌草稿、验证它并提交，同时调度器提前一步准备下一批。

一个请求的实时状态位于单个工作槽中，每次前向传播都会在原地读取并覆盖它。缓存它意味着要从 GPU 正在积极修改的内存中复制数据，这会产生两个竞争条件。恢复操作可能与正在写入该槽的前向传播冲突。而刷新缓存的快照可能与仍在读取前一个快照的 donate 操作冲突。两者都在未进行设备级同步且热路径上无锁的情况下完成。

首先，每个状态副本都是串行前向流上的一个内核。将已缓存的检查点恢复到工作槽的写时复制操作，以及在 track 边界捕获状态的快照操作，会排队在生成和使用该状态的前向传播之间，因此同一流的顺序提供了先行发生(happens-before)关系。每个 prefill 块和解码期间的每个轨道间隔都会进行一次快照。

其次，唯一不会传输字节的操作是请求操作。快照落在我们称为额外缓冲区的乒乓对中，其中一个槽位保存最新的快照，另一个槽位用于下一个快照，这个第二个槽位只在需要的边界上分配，并在之后立即释放。缓存会在每次预填充块之后、从预填充过渡到解码时，以及请求完成时，将最新快照的槽索引捐赠给树。一个新的槽会回填缓冲区，因此快照的写入目标与捐赠的读取源永远不会是同一个物理槽位。

三种状态操作。写时复制、快照和捐赠，以及它们相对于顺序前向流的位置。

因为可重用的前缀检查点存在于共享、可淘汰的基数树中，运行中的请求仅需要保留少量临时槽，最少四个。它们是工作槽、一个用于快照的额外缓冲槽，以及两个保留空间槽：一个用于必须保留的已提交前缀状态，另一个用于分支副本所需的复制。大部分缓存状态在整个树中摊销，而不是按请求收费，所以状态池不必随每个请求保持热数据的历史数量而增长。

循环状态不能在任意标记处截取，因为你无法将其倒回到早期位置，所以它仅在块边界处建立检查点，即使在边界上也很稀疏。每条路径的上限和LRU策略仅保留每条路径上的少量检查点，且检查点可以独立于其标注的键值被淘汰，使节点变为墓碑。分支点获得特别处理，该思路源自Marconi：https://arxiv.org/abs/2411.19379。一个分叉是每个未来分支都保证共享的前缀，因此当请求在边缘中间分叉时，它从上方最近的检查点重放，并在块对齐的分叉处设置新检查点。下一条分支会直接在此处恢复，无需重放。

基数树上的检查点。稀疏检查点叠加层和分支点。

这些都不是独立添加到其它功能上的附加特性。它们是同一事实的四个后果，也就是状态会自我覆盖，这就是为什么设计即使同时吸收了重叠调度器、推测解码、分页和前缀缓存，仍然保持紧凑的原因。

以上内容管理每种状态在其自己的池中，而这些池本身就是设计中唯一的猜测。两个分配单元相差三个数量级：每个请求一个大的 KDA 状态块（在 TP=8 时约 54 MB，涵盖所有 69 个 KDA 层），每个 token 一个小的 MLA KV 块（约 27 KB，涵盖所有 24 个 MLA 层），所以它们今天存在于两个独立的池中，并在启动时确定大小。这个大小是对流量的一种押注，当押注失误时，一个池的服务器会耗尽内存，而另一个仍有剩余。

统一内存用一个池代替两个池：KDA 状态从一端填充，MLA KV 块从另一端填充，中间未使用的字节形成一个单一的空闲区域。

统一内存。上图：今天两种状态有独立的池，在启动时确定大小，所以一个可以闲置而另一个已满。中图：使用统一内存时，两种状态都从同一个池分配，从相对两端增长，中间有一个单一的空闲区域。下图：释放中间的状态块会留下一个间隙，然后将末端的块移动到该间隙中，使空闲区域保持完整。

释放同样简单。当请求完成、被中止或在压力下收回时，其块会被释放。如果中间因此留下间隙，则从末端移动一个块填补，以保持空闲区域完整。

这种布局提供了灵活的页面大小且没有内存碎片：54 MB 的 KDA 状态块和 27 KB 的 MLA KV 块从相同的字节中抽取，且没有强制统一的页面大小，上述的移动操作以可忽略的状态移动成本保持两端紧凑，因此空闲空间始终是一片连续区域，可由任一种状态使用。所以容量随着工作负载而变化，而不是启动标志：许多短请求填满状态块的池，少数长上下文填满 KV，任一情况都无需重新配置。

统一内存通过 --enable-unified-memory 进行选择性启用。接下来的帖子将详细介绍实现过程。

K3 配备了 DSpark 阻止投机解码，这是由我们为 K3 训练的草稿模型驱动的：https://huggingface.co/RadixArk/Kimi-K3-DSpark。集成的两个部分值得单独讲述：只在有价值的地方花费验证预算，以及让 K3 的循环 KDA 状态在投机下得以存活。

DSpark 每步提出一块草稿令牌，目标会在一次前向传播中验证整块。对于批量大小 1，额外的验证位置几乎是免费的：步骤受延迟限制，多几个令牌并无大碍。随着批量增加，情况就不再如此。验证令牌现在与同一时间步的所有其他请求争夺时间，大多数令牌最终仍会失败：在聊天工作负载下，接受长度约为 2.7，所以在典型步骤中被验证的八个位置中有五个会被拒绝。验证所有令牌意味着需要为服务器最终丢弃的令牌支付全价。

解决方案的组件已经在系统中。草稿携带一个训练过的置信头，它可以预测每个位置令牌通过验证的可能性。另一方面，服务器的一次性性能分析记录了在每个负载水平下额外验证令牌的实际成本。每步规划器将两者结合：每个请求只保留那些预期价值覆盖边际成本的验证令牌，其余的在目标前向传播开始前被裁减。剩下的部分如同以前一样经过验证，因此输出保持无损。权衡是略短的接受运行以获得更低成本的步骤。

在高负载下裁剪可以获益。解码吞吐量，在聊天面板（接受 ~2.7 左）和少量示例的数学面板（接受 ~5.0 右）上，验证全部 vs 裁剪。达到批量大小 8 时持平，随后随着批量增加差距扩大：批量大小 256 时 +68% 和 +24%，接受长度从 2.7 降至 2.2，从 5.0 降至 4.3，同时吞吐量上升。

测量的成本曲线是有趣的部分。验证令牌的边际成本并不平滑。它像阶梯一样：平坦的平台上，另一个令牌几乎不费力地插入当前的内核波动，而陡峭的上升段则意味着它开始新的波动。规划器读取这个表面。当请求的下一个令牌位于廉价的平台上时，它会保留它们；当它们将开始上升段时，规划器就在此处剪断。在数据中，这表现为接受长度曲线的小幅摆动，而吞吐量保持平滑且单调：规划器是在“冲浪”真实硬件平台，而不是噪声。

在批量大小小于8时，几乎没有减负压力，因此剪裁在此时是打平或轻度负效应；在规划器中，已知的后续操作是小批量的提前退出。

推测解码一次验证 γ+1 个草稿令牌，并可能只接受前缀。对于 MLA 来说，这是免费的，因为 KV 是只追加的，拒绝的草稿仅释放其槽。正如前一节所述，KDA 层的状态每个令牌都会自我覆盖，因此基线通过暴力方法实现可逆性，并在每个草稿步骤之后对整个 K×V 状态做快照。在 K=V=128 时，每个请求、层和头部占用 64 KB，再乘以 γ+1 步。跨越 K3 的 69 个 KDA 层和完整批次，它会超过与之竞争的持久状态池，并且由于这是为每个正在执行的请求预留的，因此会限制并发。

存储输入，而不是状态。ReplaySSM：https://tridao.me/blog/2026/replayssm/ 取消了快照。验证内核读取已提交的检查点而从不写入它，在此过程中它还存储每一步的原始输入 Sᵢ = (vᵢ, kᵢ, gkᵢ, βᵢ)，大约 1 KB，相比一个快照的 64 KB 少很多。一旦采样器确定接受长度，一个覆盖所有层和头的单一折叠内核就从检查点重放被接受的前缀，并在原地推进。被拒绝的草稿永远不会重放，因此回滚成本为零。草稿窗口从 512 KB 减少到 16 KB，大约减少了 32 倍。

准确，不是近似。fold 是对 verify 递归的逐字克隆，使用相同的 tile 和相同的归约顺序，并且它直接使用 verify 内核自身存储的 gate 值，而不是重新计算它们。第二部分不是可选的。早期版本在 torch 端使用略有不同的公式重新计算 gate，这导致每个输出看起来都是正确的，但状态在底层悄悄漂移。现在重建的状态与递归基线将提交的状态在比特级别上完全相同。

收益在于容量，而不是速度。每步解码时间保持不变，因为这是内存换计算的权衡，而不是更快的内核。快照缓存只是规格上的固定成本，会随着批量大小和 γ 增长，因此在 γ 足够大值得使用时，它对状态池的压力最大。归还这部分内存可以显著提升并发上限。在旧上限以内，两条路径表现相当，ReplaySSM 为 fold 和缓冲写入付出一些代价；超过旧上限时，基线排队，而 ReplaySSM 继续接收请求，这就是差距产生的地方。

KDA ReplaySSM。Verify（每层一次融合启动）读取检查点 h₀ 而不写入，并将每个草稿步骤的原始输入存储到每个 slot 缓冲中，每层每个头一个。在采样器确定接受长度后，fold 内核（一次启动覆盖所有层和头）仅重放被接受的前缀并就地覆盖 slot，使用 verify 内核自身存储的 gate 值运行逐字递归，因此重建状态在比特层面上与递归基线完全相同。

在 2.8T 混合模型上的单序列解码步骤不是计算受限的——这是一个启动次数和延迟问题：每个 token 有 93 层注意力（69 层 KDA + 24 层 MLA）加上 92 层 latent-MoE，在启动时每步触发数百个小内核。优化过程是基于分析的：融合一件事，用固定协议进行 A/B 测试，在 GSM8K 上进行验证，然后重复。

启动与拷贝消除（P1–P4，+19.9 tok/s）。缩小步骤，而不是加快内核：MoE 前端折叠为一个 GEMM，KDA 的配对瘦投影合并（P1）；基于分析的优化扫除每层的上采样、拷贝和多余启动（P2, P3）；routing 变为对 [M, 896] logits 的一遍寄存器驻留基数选择（P4）。

NVIDIA 计算内核（P5–P8，+10.3 tok/s）。在四个阶段中切换使用与 NVIDIA 共同开发或由其提供的内核，针对 bs=1 形状默认值错误的情况：融合的 KDA 解码内核、trtllm-gen W4A8 SiTU MoE cubins（在我们的路由绕过后面）、TMA 注意力-残差聚合和 CuTe-DSL TGV bf16 GEMMs（在小 M 时替代 cuBLAS）——每一个都通过与其他所有内容相同的 A/B 和准确性门控。

通信融合（P9, P12, P13, +27.6 tok/s）。这是最大的一代：为 CustomAllReduceV2 的对称内存平面上的 MNNVL 结构设计的融合全归约系列——对小消息进行一次性多播存储，对大消息进行 NVLS 交换机内归约，残差加法和 RMSNorm 内嵌在集合操作中（P9）。然后 MoE finalize 移入集合操作的分段传递中（P12），上投影从复制 GEMM 转变为列并行 GEMM，并通过多播全聚集完成（P13）。

重叠 & 前导融合（P10, P11, P14, P15, +10.4 tok/s）。其余部分修剪了剩余的关键链条：步进输入的 MXFP8 量化（P10）、残差写回融合到上游内核尾部并在侧流上有独立分支（P11）、KDA 的 GEMV 链与 qkvg GEMM 重叠（P14）、MLA 解码前导融合到一个内核中，PDL 在其后触发注意力内核（P15）。

四条进度柱背后的时间顺序：

总结的经验。全归约是一个同步点，因此在那里节省的微秒会一比一地转换为步骤时间；一个内核驻留在另一个流的重叠空闲时间，大约转换为十分之一。在写入内核之前检查轨迹中的关键路径成员身份，是本次优化活动中杠杆作用最大的习惯。

¹ 一次重置将 P5 基线从 64.2 移动到 63.8；步骤是针对重置后的基线测量的。² P13 是合并窗口后的重新校准的规范基线，而非单一 PR 归因。³ 进入 P1–P4 的片段也包含随后被取代的临时优化（全归约拼接 → P9；微/1-CTA GEMV → P8；基数路由器 v1 → P4；Marlin top-k 求和 → P6/P12；早期注意力-残差加 → P7）；它们的短期收益在曲线上体现，但没有命名点。

对于 K3 的混合架构，传统的并行选择不再适用。张量并行无法分片 MLA 的 KV 缓存（每个 KV 头只有一个，无法拆分），因此每个节点都持有完整副本；它还将每个 GEMM 切成八份，并在每层支付一次集合通信成本。纯 DP 注意力反而在每个节点上复制注意力权重，大约 61 GB 用于 KDA，加上 MLA 的 11 GB，用于存储 KV 缓存和 KDA 状态。在这些成本下，预填充和解码的表现不同，因此 K3 按阶段拆分解决方案：预填充采用分块流水线并行，解码采用上下文并行。

在 TP 预填充中，每层都以 AllReduce 结束，这是一个无法与计算重叠的屏障。流水线并行改为按层切分模型：K3 的 93 层变成 8 个阶段，并将提示切成块流经它们：

分块流水线并行预填充。各阶段同时处理不同的块。阶段之间的交接在该阶段计算下一个块时进行，因此在 K3 上隐藏了 91% 的通信。在 TP（底部条）下，每层都以 AllReduce 结束，所有节点都必须等待。

这一方法在三方面获胜。唯一剩下的通信，即传递给下一阶段的过程，被下一个块的计算所隐藏。每个节点运行完整的层，因此 GEMM 扩展八倍，更高效。此外，每个阶段只保存其约 12 层的 KV 和激活状态，使得处理非常长的提示进行预填充变得舒适。流水线确实需要足够深：浅层的 PP4×TP2 无法覆盖其交接成本，仍需支付 TP2 的 AllReduce，且其基准性能不优于 TEP8。

在 2×4 GB300 上测量 8K 预填充，拓扑结构为唯一变量：

深度 PP 在两个维度上都占优势。左图：超过 c1–c4 交叉点后，PP8×TP1 上升至约 TEP8 上限的 1.7 倍，TTFT 更低。右图：每节点 FLOPs 相等时每 1k 预填充 token 的成本；PP4×TP2 和 TEP8 在计算与通信之间权衡并持平，PP8 在两方面成本最低。

PP8 仅在单次请求上表现较差，而空闲的预填充工作节点本身配置不当。在 K3 的解耦服务中，预填充节点运行 PP8。它们的预填充能力是 TEP8 节点的 1.45 至 1.72 倍，因此一个预填充节点可以支撑多个解码节点。解码节点运行 TP 或 DCP，下一部分将介绍。

解码是复制的 KV 缓存绑定的地方：在 TP 下，接受一个额外的请求或多一千个上下文 token 在每个 rank 上消耗的字节数是相同的。解码上下文并行（DCP）按 token 位置而不是按 head 对 MLA 的 KV 进行分片。当 p mod N 等于 r 时，rank r 拥有位置 p，因此每个 rank 保存每个请求上下文交错的 1/N，且这种分片在注意力内核之上是不可见的：

按位置分片。在 4 个 ranks 上相同的 16 个 token 位置：TP 存储 64 个物理副本，DCP 只存储 16 个。释放的字节变成逻辑 KV 容量，在 K3 上 DCP8 可达约 7.9 倍。

位置分片会破坏 softmax，因为每个 rank 只能看到 1/N 的 key，部分 softmax 不能相加。解决方法是 FlashAttention 自己的方法：每个 rank 返回其部分注意力输出以及每个 head 的 log-sum-exp，每层进行一次全互换（all-to-all）交换，使得每个 rank 最终拥有全上下文的 1/N head。然后通过 log-sum-exp 本地合并是精确的，且结果已经是 TP 下输出投影期望的 head 布局。每层一次 collective 是整个通信开销：

解码步骤，每层。每个 rank 本地投影完整 head 的 query，仅对其拥有的位置进行 attention，然后在一次打包的 all-to-all 中发送部分输出及其 log-sum-exp。本地合并得到标准的 TP head 布局；注意力之后的部分都不变。

其他一切保持不变。DCP 组在 TP 组内构建，所以 TP8 带 DCP8 仍然是 8 个 GPU，MoE 按原本并行方式运行。KDA 是例外，它的状态是每个请求一个固定大小矩阵而非每个 token，所以没有位置维度可以分片，KDA 层保持 TP 按 head 分片。整个特性只需一个标志，--dcp-size N。

这种方法的效果在代理流量中体现出来，上百千 token 的会话堆积在缓存中。在 2×4 GB300 上重放真实编程代理会话，两个节点都用主机内存 KV 层，DCP 是唯一的区别：

DCP移除了活跃集墙。主机层拯救了重新预填充的流量，但在16个并发会话下，活跃工作集超过了TP8的设备KV，吞吐量崩溃。DCP8将逻辑KV从1.5M提升到12.2M tokens，并在48个会话下以541 tok/s处理相同的工作负载。

DCP与其余堆栈协作。DSpark验证步骤是一个解码步骤，因此它使用相同的复制Q、全对全路径；在PD解耦下，预填充端保持对DCP的不可感知性，每个解码rank只拉取它在传输边界拥有的位置，这使得PP或TP预填充可以供DCP解码使用。剩下的上限是KDA：它的每个请求状态无法按位置分片，所以一旦DCP解除MLA墙，运行请求上限成为限制因素；上文内存部分的统一内存设计处理了这个问题。

各组件如何组合最终是一个测量。将PD解耦放入循环，结合分块PP8预填充，以及解码拓扑和预填充:解码比例的变化，得到服务前沿：

服务前沿。在吞吐量方面，DCP组合：两个PP8预填充工作者喂给两个DCP8解码节点，每个GPU提供2633 tok/s，与最佳的TP8臂相当。向右移动是预填充:解码旋钮的作用：一个PP8预填充工作者喂给两个、三个、然后四个独立的TP8解码实例，交换整体吞吐量以提升每用户速度，超过每用户86 tok/s。

K3的Day-0 RL与Miles共享地点进行LoRA训练：在Miles的Megatron后端上运行BF16训练器，以及在相同64 GB 300s上共享的原生打包MXFP4 SGLang部署引擎。后端涵盖KDA、NoPE-MLA、注意力-残差库和潜在MoE，以及TP/SP/PP/CP/EP。

引擎以出厂状态服务于检查点，且从不重写它。每一步仅传输 BF16 LoRA 适配器，且引擎将增量作为单独的 BF16 B(Ax) 项应用于量化的基础 GEMM 之上，因此策略更新能够在 4 位基础权重上以全精度进行推理。稠密投影通过 SGLang 的 Triton LoRA 后端处理，896 个路由专家通过 Marlin 路径上的融合 MoE-LoRA 内核处理，共享专家增量折叠进融合 MoE 前 GEMM。适配器存在于 GPU 内存池中，每次同步会原地替换，因此在展开端没有 BF16 副本，不需要全量权重重同步，也无需在循环中重新量化。

流水线并行。K3 的注意力-残差快照银行必须跨阶段边界传递，但 Megatron 的点对点传递仅携带一个隐藏状态张量，因此阶段边界将 [prefix_sum, bank] 打包到其中，并在进入时解包。Megatron 未被修改。

上下文并行。银行可以免费分片，因为每个注意力-残差操作都是按 token 进行的。MLA 也不需要任何 K3 特定代码：旋转表切片是 Megatron MLA 中唯一的 CP 感知工作，而 K3 没有旋转嵌入，因此其投影使用标准 TE 注意力核并继承 CP。KDA 是需要处理的部分，通过 fla 的 CP 上下文处理循环状态和卷积 halo。它需要一个连续的 rank-local 块，Megatron 将 zigzag 顺序存储在环形注意力中所期望的位置，因此重布局仅在 KDA 周围进行。

We are excited to announce Day-0 support for Kimi K3：https://platform.kimi.ai/docs/guide/kimi-k3-quickstart in SGLang and Miles. K3 is the first open-source model in the 3-trillion-parameter class, and its hybrid architecture departs from convention in almost every place a serving stack has assumptions. In collaboration with the Moonshot AI and NVIDIA teams, the two cover K3 in full on launch day: SGLang for inference, Miles for RL training. This post covers what it took.

Launch commands and per-workload configuration guidance live in the Kimi K3 cookbook：https://docs.sglang.io/cookbook/autoregressive/Moonshotai/Kimi-K3.

Kimi K3：https://platform.kimi.ai/docs/guide/kimi-k3-quickstart is the first open-source model in the 3-trillion-parameter class: 2.8T parameters, a 1M-token context window, and native visual understanding. Its predecessor K2.5 was architecturally a close relative of models SGLang already served well, so most of the stack simply applied. K3 is not like that — it departs from convention in several independent places at once:

Each of these rows is a serving problem. The hybrid attention stack is what makes the 1M context affordable, but it means the server holds two kinds of state at once: one fixed-size KDA state per request next to MLA's per-token KV — the memory-management sections below are about exactly this. Attention Residuals thread a bank of attention outputs through the whole stack, which breaks the assumptions SGLang's standard layer plumbing makes and forced K3-specific paths in places like DP attention. LatentMoE routes 16 of 896 experts inside a down-projected latent space with the SiTU activation instead of SwiGLU, so no existing MoE kernel applied out of the box. And the vision path — tower, projector, processor, and Kimi's XTML media format — was brought up from scratch. The sections below all build on this bring-up.

Attention KV is append-only. Once a token's KV is computed it never changes, which is what lets the scheduler share one physical copy across every request that shares a prefix, keep it in a radix tree, and pipeline it across iterations without a second thought. A KDA layer's state is the opposite, one fixed-size recurrent buffer that is overwritten in place at every token . So everything the scheduler hands KV for free, from prefix caching to the overlap scheduler to speculative decoding to paging, has to be rebuilt for a value that mutates as it is read. K3 interleaves 69 KDA layers with 24 MLA layers, so this is not a corner case. It is most of the model. What follows is that rebuild, and its one satisfying property is that every piece is forced by the overwrite-in-place fact. None of it is a bolt-on.

The same path serves the overlap scheduler, speculative decoding, and page_size > 1 at once, combinations that were previously mutually exclusive for recurrent-state models. A request can hit a cached prefix, restore the recurrent state, run a multi-token draft, verify it, and commit, all while the scheduler prepares the next batch one step ahead.

A request's live state sits in a single working slot, read and overwritten in place by every forward. Caching it means copying out of memory the GPU is actively mutating, and that creates two races. A restore can collide with the forward that is writing the slot. And the snapshot that refreshes the cache can collide with a donate that is still reading the previous one. Both close with no device-wide synchronization and no lock on the hot path.

First, every state copy is a kernel on the serial forward stream. The copy-on-write that restores a cached checkpoint into the working slot, and the snapshot that captures the state at a track boundary, are enqueued between the forwards that produce and consume that state, so same-stream ordering provides the happens-before for free. A snapshot is taken once per prefill chunk and once every track interval during decode.

Second, the only move that leaves the request transfers no bytes. Snapshots land in a ping-pong pair we call the extra buffer, where one slot holds the latest snapshot while the other takes the next one, and that second slot is allocated only at the boundary that needs it and freed right after. Caching donates the latest snapshot's slot index to the tree, after each prefill chunk, at the handover from prefill to decode, and when the request finishes. A fresh slot backfills the buffer, so the snapshot's write target and the donate's read source are never the same physical slot.

The three state moves. Copy-on-write, snapshot, and donate, and where each one sits relative to the serial forward stream.

Because the reusable prefix checkpoints live in the shared, evictable radix tree, a running request reserves only a handful of transient slots, as few as four. That is its working slot, one extra-buffer slot for snapshots, and two slots of retention headroom, one for the committed prefix state that has to stay resident and one for the copy a diverging branch needs. The bulk of cached state is amortized across the whole tree rather than charged per request, so the state pool does not have to grow with how much history each request keeps warm.

A recurrent state cannot be sliced at an arbitrary token, since you cannot run it backwards to an earlier position, so it is checkpointed only at chunk boundaries, and sparsely even there. A per-path cap and LRU keep just a few checkpoints alive on each path, and a checkpoint can be evicted independently of the KV it annotates, leaving the node as a tombstone. Branch points get special treatment, an idea from Marconi：https://arxiv.org/abs/2411.19379. A fork is the one prefix every future branch is guaranteed to share, so when a request diverges mid-edge it replays from the nearest checkpoint above and plants a new checkpoint at the chunk-aligned fork. The next branch restores there directly, no replay.

Checkpoints on the radix tree. The sparse checkpoint overlay and the branching point.

None of these is a separate feature stitched onto the others. They are four consequences of one fact, that the state overwrites itself, which is why the design stays small even as it absorbs the overlap scheduler, speculative decoding, paging, and prefix caching all at once.

Everything above manages each kind of state inside its own pool, and the pools themselves are the one guess left in the design. The two allocation units are three orders of magnitude apart: one large KDA state block per request (about 54 MB under TP=8, covering all 69 KDA layers) and one small MLA KV block per token (about 27 KB, covering all 24 MLA layers), so today they live in two separate pools, sized at startup. That sizing is a bet on the traffic, and when the bet is off the server runs out of memory in one pool while the other still has room to spare.

Unified memory replaces the two pools with one: KDA states fill in from one end, MLA KV blocks from the other, and the unused bytes between them form a single free region.

Unified memory. Top: today the two kinds of state have separate pools, sized at startup, so one can sit idle while the other is full. Middle: with unified memory both kinds allocate from the same pool, growing in from opposite ends with a single free region between them. Bottom: freeing a state block in the middle leaves a gap, and a block from the end is moved into it so the free region stays in one piece.

Freeing is just as simple. When a request finishes, is aborted, or is retracted under pressure, its blocks are released. If that leaves a gap in the middle, a block from the end is moved into it so the free region stays in one piece.

This layout gives us flexible page sizes and no memory fragmentation: a 54 MB KDA state block and a 27 KB MLA KV block draw from the same bytes with no common page size forced on them, and the moves above keep each end packed at a negligible cost for state movement, so the free space is always one contiguous region usable by either kind. Capacity therefore follows the workload rather than a startup flag: many short requests fill the pool with state blocks, a few long contexts fill it with KV, and neither case needs anything reconfigured.

Unified memory ships opt-in behind --enable-unified-memory . A follow-up post will go through the implementation in detail.

K3 ships with DSpark block speculative decoding, driven by a draft model we trained for K3：https://huggingface.co/RadixArk/Kimi-K3-DSpark. Two parts of the integration deserve their own story: spending the verification budget only where it pays, and making K3's recurrent KDA state survive speculation at all.

DSpark proposes a block of draft tokens per step, and the target verifies the whole block in one forward. At batch size 1 the extra verify positions are essentially free: the step is latency-bound, and a few more tokens ride along. As the batch fills, that stops being true. Verify tokens now compete with every other request for the same step time, and most of them lose their bet anyway: on a chat workload the accept length is around 2.7, so five of the eight verified positions on a typical step get rejected. Verifying everything pays full price for tokens the server then throws away.

The pieces to fix this were already in the system. The draft carries a trained confidence head that predicts, per position, how likely each token is to survive verification. On the other side, a one-time profile of the server records what an extra verify token actually costs at each load level. A per-step planner joins the two: each request keeps verify tokens only while their expected value covers the marginal cost, and the rest of the window is trimmed before the target forward launches. What remains is verified exactly as before, so outputs stay lossless. The trade is a slightly shorter accepted run for a cheaper step.

Trimming wins under load. Decode throughput, verify-all vs trim, on a chat panel (accept ~2.7, left) and a few-shot math panel (accept ~5.0, right). Break-even through bs 8, then the gap opens with batch size: +68% and +24% at bs 256, with accept length easing from 2.7 to 2.2 and from 5.0 to 4.3 while throughput climbs.

The measured cost curve turned out to be the interesting part. The marginal cost of a verify token is not smooth. It is a staircase: flat shelves where another token slots into the current kernel waves for almost nothing, and steep risers where it starts a new one. The planner reads this surface. When a request's next tokens sit on a cheap shelf it keeps them; when they would start a riser it cuts there. In the data this shows up as small wiggles in the accept-length curve while throughput stays smooth and monotone: the planner is surfing real hardware shelves, not noise.

Below batch size 8 there is little pressure to relieve, so trimming is break-even to mildly negative there; a small-batch early-exit in the planner is the known follow-up.

Speculative decoding verifies γ+1 draft tokens at once and may accept only a prefix. For MLA that is free, since KV is append-only and a rejected draft just releases its slots. A KDA layer's state overwrites itself every token, as the previous section covered, so the baseline buys reversibility by brute force and snapshots the whole K×V state after every draft step. At K=V=128 that is 64 KB per request, layer and head, times γ+1 steps. Across K3's 69 KDA layers and a full batch it outgrows the persistent state pool it is competing with, and because it is reserved per running request, it caps concurrency.

Store the inputs, not the state. ReplaySSM：https://tridao.me/blog/2026/replayssm/ drops the snapshots. The verify kernel reads the committed checkpoint and never writes it, and on the way through it also stores each step's raw inputs Sᵢ = (vᵢ, kᵢ, gkᵢ, βᵢ) , about 1 KB against the 64 KB a snapshot costs. Once the sampler fixes the accepted length, a single fold kernel covering every layer and head replays just the accepted prefix from the checkpoint and advances it in place. Rejected drafts are never replayed, so rollback costs nothing. The draft window goes from 512 KB to 16 KB, roughly 32×.

Exact, not approximate. The fold is a verbatim clone of the verify recurrence, same tiles and same reduction order, and it consumes the gate values the verify kernel itself stored rather than recomputing them. That second half is not optional. An early version recomputed the gate on the torch side with a subtly different formula, which left every output looking correct while the state quietly drifted underneath. The rebuilt state is now bit-identical to what the recurrent baseline would have committed.

The payoff is capacity rather than speed. Per-step decode time is unchanged, since this is a memory-for-recompute trade and not a faster kernel. The snapshot scratch was a spec-only fixed cost growing with both batch size and γ, so it squeezed the state pool hardest exactly when γ was large enough to be worth using. Handing that memory back lifts the concurrency ceiling several-fold. Below the old ceiling the two paths are a wash, with ReplaySSM paying a little for the fold and the buffer writes; above it the baseline queues while ReplaySSM keeps admitting requests, and that is where the gap opens.

KDA ReplaySSM. Verify (one fused launch per layer) reads the checkpoint h₀ without writing it and stores each draft step's raw inputs into a per-slot buffer, one per layer and head. After the sampler fixes the accepted length, the fold kernel (one launch for all layers and heads) replays only the accepted prefix and overwrites the slot in place, running the verbatim recurrence with the verify kernel's own stored gate values, so the rebuilt state is bit-identical to the recurrent baseline.

A single-sequence decode step on a 2.8T hybrid model is not compute-bound — it is a launch-count and latency problem: 93 attention layers (69 KDA + 24 MLA) plus 92 latent-MoE layers per token, and at bring-up the step fired hundreds of tiny kernels. The campaign was profile-driven: fuse one thing, A/B it on the fixed protocol, gate on GSM8K, repeat.

Launch & copy elimination (P1–P4, +19.9 tok/s). Make the step smaller, not the kernels faster: the MoE front collapses into one GEMM and KDA's paired skinny projections merge (P1); a profile-guided sweep removes per-layer upcasts, copies and spare launches (P2, P3); routing becomes a one-pass register-resident radix select over the [M, 896] logits (P4).

NVIDIA compute kernels (P5–P8, +10.3 tok/s). Four stages swap in kernels co-developed with or provided by NVIDIA where the default was wrong for the bs=1 shape: the fused KDA decode kernel, the trtllm-gen W4A8 SiTU MoE cubins behind our routing bypass, the TMA attention-residual aggregation, and the CuTe-DSL TGV bf16 GEMMs that replace cuBLAS at small M — each through the same A/B and accuracy gate as everything else.

Communication fusion (P9, P12, P13, +27.6 tok/s). The largest era: a fused all-reduce family for the MNNVL fabric on CustomAllReduceV2's symmetric-memory plane — one-shot multicast stores for small messages, NVLS in-switch reduction for large ones, with the residual add and RMSNorm riding inside the collective (P9). The MoE finalize then moves into the collective's staging pass (P12), and the up-projection flips from a replicated to a column-parallel GEMM completed by a multicast all-gather (P13).

Overlap & prologue fusion (P10, P11, P14, P15, +10.4 tok/s). The rest trims the remaining critical chain: a strided-input MXFP8 quant (P10), residual write-back fused into upstream kernel tails with independent branches on a side stream (P11), KDA's GEMV chain overlapped with the qkvg GEMM (P14), and the MLA decode prologue fused into one kernel with PDL arming the attention kernel behind it (P15).

The chronological ladder behind the four bars:

The lesson that generalizes. All-reduce is a synchronization point, so a microsecond saved there converts one-for-one into step time; a kernel sitting in another stream's overlap slack converts at roughly one-tenth. Checking critical-path membership in the trace before writing a kernel was the single highest-leverage habit of this campaign.

¹ A rebase moved the P5 baseline 64.2 → 63.8; the step is measured against the post-rebase base. ² P13 is a re-calibrated canonical baseline after a merge window, not a single-PR attribution. ³ Segments leading into P1–P4 also carry interim optimizations later superseded (concat all-reduce → P9; tiny/1-CTA GEMV → P8; radix router v1 → P4; Marlin top-k-sum → P6/P12; early attn-res add → P7); their transient gains are in the curve but have no named point.

For K3's hybrid architecture, the conventional parallelism choices do not hold up. Tensor parallelism cannot shard MLA's KV cache (one KV head, nothing to split by), so every rank holds a full copy; it also slices every GEMM eight ways and pays a collective per layer. Pure DP attention instead replicates the attention weights on every rank, around 61 GB for KDA plus 11 GB for MLA, memory the KV cache and KDA states need. Prefill and decode fail differently under these costs, so K3 splits the answer by phase: chunked pipeline parallelism for prefill, context parallelism for decode.

In TP prefill, every layer ends in an AllReduce, a barrier that cannot overlap with compute. Pipeline parallelism cuts the model by layer instead: K3's 93 layers become 8 stages, and the prompt is cut into chunks that stream through them:

Chunked pipeline-parallel prefill. The stages work on different chunks at the same time. The hand-off between stages runs while the stage computes its next chunk, so 91% of it is hidden on K3. Under TP (bottom strip), every layer ends in an AllReduce that all ranks wait for.

This wins three ways. The only communication left, the hand-off to the next stage, hides behind the next chunk's compute. Each rank runs whole layers, so the GEMMs are eight times wider and more efficient. And each stage holds KV and activations for only its ~12 layers, which makes very long prompts comfortable to prefill. The pipeline does have to be deep: a shallow PP4×TP2 cannot cover its hand-offs and still pays TP2's AllReduce, and it benchmarks no better than TEP8.

Measured on 8K prefill across 2×4 GB300, with topology as the only variable:

Deep PP wins on both axes. Left: past the c1–c4 crossover, PP8×TP1 climbs to about 1.7× TEP8's ceiling, with lower TTFT. Right: cost per 1k prefill tokens at equal per-rank FLOPs; PP4×TP2 and TEP8 trade compute against communication and tie, PP8 is cheapest on both.

PP8 loses only at a single request, and a prefill worker that idle is misconfigured anyway. In K3's disaggregated serving, the prefill nodes run PP8. Each of them has 1.45 to 1.72 times the prefill capacity of a TEP8 node, so one prefill node can keep several decode nodes fed. The decode nodes run TP or DCP, which the next part covers.

Decode is where the replicated KV cache binds: under TP, admitting one more request or one more thousand tokens of context costs the same bytes on every rank. Decode Context Parallelism (DCP) shards MLA's KV by token position instead of by head. Rank r owns position p when p mod N equals r, so every rank holds an interleaved 1/N of every request's context, and the sharding is invisible above the attention kernel:

Sharding by position. The same 16 token positions on 4 ranks: TP stores 64 physical copies, DCP stores 16. The freed bytes become logical KV capacity, about 7.9× on K3 with DCP8.

Position-sharding breaks the softmax, since each rank sees only 1/N of the keys and partial softmaxes do not add. The fix is FlashAttention's own: each rank returns its partial attention output together with a log-sum-exp per head, and one all-to-all per layer exchanges them so that each rank ends up with 1/N of the heads over the full context. A local merge by log-sum-exp is then exact, and the result is already the head layout the output projection expects under TP. One collective per layer is the entire communication cost:

The decode step, per layer. Each rank projects the full-head query locally, attends over only the positions it owns, and ships partial outputs and their log-sum-exps in a single packed all-to-all. A local merge leaves the standard TP head layout; nothing downstream of attention changes.

Everything else is left alone. DCP groups are built inside the TP group, so TP8 with DCP8 is still 8 GPUs, and the MoE runs whatever parallelism it already had. KDA is the exception that explains the rule: its state is one fixed-size matrix per request rather than per token, so there is no position axis to shard, and the KDA layers stay TP-sharded by head. The whole feature is one flag, --dcp-size N .

What this buys shows up on agentic traffic, where hundred-thousand-token sessions pile up in the cache. Replaying real coding-agent sessions on 2×4 GB300, with a host-memory KV tier on both arms and DCP as the only difference:

DCP removes the active-set wall. The host tier rescues re-prefill traffic, but at 16 concurrent sessions the active working set outgrows TP8's device KV and throughput collapses. DCP8 lifts logical KV from 1.5M to 12.2M tokens and carries the same workload to 541 tok/s at 48 sessions.

DCP composes with the rest of the stack. A DSpark verify step is a decode step, so it rides the same replicated-Q, one-all-to-all path; under PD disaggregation the prefill side stays DCP-unaware and each decode rank simply pulls the positions it owns at the transfer boundary, which is what lets PP or TP prefill feed DCP decode. The remaining ceiling is KDA: its per-request state cannot be position-sharded, so once DCP lifts the MLA wall the running-request cap becomes the binding limit; the unified-memory design in the memory section above takes that on.

How the pieces compose is ultimately a measurement. Putting PD disaggregation in the loop, with chunked PP8 prefill throughout and the decode topology and prefill:decode ratio varying, gives the serving frontier:

The serving frontier. At the throughput end the DCP composition, two PP8 prefill workers feeding two DCP8 decode nodes, delivers 2,633 tok/s per GPU, level with the best TP8 arm. Moving right is the prefill:decode knob at work: one PP8 prefill worker feeding two, three, then four independent TP8 decode instances trades aggregate throughput for per-user speed, out past 86 tok/s per user.

Day-0 RL for K3 is colocated LoRA training with Miles: a BF16 trainer on Miles's Megatron backend and native packed MXFP4 SGLang rollout engines sharing the same 64 GB300s. The backend covers KDA, NoPE-MLA, the attention-residual bank and the latent MoE, with TP/SP/PP/CP/EP.

The engines serve the checkpoint as shipped and never rewrite it. Each step transfers only BF16 LoRA adapters, and the engine applies the delta as a separate BF16 B(Ax) term on top of the quantized base GEMM, so the policy update reaches inference at full precision over 4-bit base weights. Dense projections go through SGLang's Triton LoRA backend, the 896 routed experts through a fused MoE-LoRA kernel on the Marlin path, and the shared-expert delta is folded into the fused MoE front GEMM. Adapters live in a GPU memory pool that each sync replaces in place, so there is no resident BF16 copy on the rollout side, no full-weight resync, and no requantization step in the loop.

Pipeline parallelism. K3's attention-residual snapshot bank has to cross stage boundaries, but Megatron's point-to-point carries one hidden-states tensor, so the stage boundary packs [prefix_sum, bank] into it and unpacks on entry. Megatron is untouched.

Context parallelism. The bank shards for free, since every attention-residual operation is per-token. MLA needs no K3-specific code either: rotary-table slicing is the only CP-aware work in Megatron's MLA and K3 has no rotary embedding, so its projections sit on the stock TE attention core and inherit CP. KDA is the part that needs work, going through fla's CP context for the recurrent state and the convolution halo. It wants a contiguous rank-local chunk where Megatron stores the zigzag order ring attention expects, so the relayout runs only around KDA.