SGLang 和 Miles 为 Kimi K3 提供首日支持
SGLang 团队
我们激动地宣布,SGLang 和 Miles 已为 Kimi K3 提供首日支持。K3 是首个开源的三万亿参数级别模型,其混合架构在推理栈几乎每个存在预设假设的地方都打破了常规。通过与月之暗面(Moonshot AI)和 NVIDIA 团队的合作,两者在发布当日即实现对 K3 的全面覆盖:SGLang 负责推理,Miles 负责强化学习训练。本文将介绍实现这一目标所需的工作。
- 一种全新的混合架构。2.8T 参数,69 层 KDA 线性注意力层与 24 层 MLA 交错排列,采用 LatentMoE 和注意力残差机制;推理栈的大部分预设假设在某种程度上都被打破了。
- 针对两种状态的内存管理,包括一种会在原地覆写自身的循环状态,并在此基础上重建了前缀缓存、重叠调度和分页机制,以及一种消除了最后尺寸猜测的统一池设计。
- 在推测解码之前,单批次推理速度可达约 113 tok/s 的内核阶梯。
- DSpark 推测解码,搭配我们为 K3 训练的草稿模型:单批次解码速度可达约 423 tok/s。ReplaySSM 负责处理 KDA 状态,通过重放原始输入而非每步快照的方式,将草稿窗口内存削减约 32 倍。
- 按阶段划分的并行策略:分块流水线并行预填充和上下文并行解码,在 PD 分离架构下组合,达到每 GPU 2,633 tok/s 的速度。
- 在原生 MXFP4 检查点上使用 Miles 进行 LoRA 强化学习,训练器和推理部署共用同一批 GPU:在 12 小时的运行中,AIME-2024 得分从 43.3% 提升至 76.7%。
启动命令和按工作负载的配置指南已收录于 Kimi K3 使用手册中。
Kimi K3 的部署上线
Kimi K3 是首个开源的三万亿参数级别模型:2.8T 参数,100 万 token 的上下文窗口,以及原生的视觉理解能力。其前代 K2.5 在架构上与 SGLang 已能良好服务的模型关系密切,因此大部分技术栈可以直接沿用。K3 则完全不同——它在多个独立方面同时打破了常规:
| Kimi K2.5 | Kimi K3 | |
|---|---|---|
| 规模 | 约 1T 参数 | 2.8T 参数 |
| 注意力机制 | 全 MLA,61 层 | 混合架构:69 层 KDA 线性注意力 + 24 层 MLA,共 93 层 |
| 注意力残差 | 无 | 每个注意力输出均被缓存,按块聚合 |
| 混合专家模型(MoE) | 384 个专家,Top-8 路由 | LatentMoE:896 个专家,top-16 选择,在 3584 维的潜在空间中运行 |
| 专家激活 | SwiGLU | SiTU |
| 视觉塔 | MoonViT | MoonViT3d,一个专为 K3 设计的新堆栈 |
这里的每一行都是一个服务问题。混合注意力堆栈使得 100 万上下文窗口变得可行,但这意味着服务器同时持有两种状态:每个请求一个固定大小的 KDA 状态,以及 MLA 的每个 token KV 状态——下面关于内存管理的部分正是讨论这个。注意力残差机制将一组注意力输出贯穿整个堆栈,这打破了 SGLang 标准层管道所做的假设,并迫使在 DP 注意力等地方采用 K3 特定的路径。LatentMoE 在一个经过下投影的潜在空间内,使用 SiTU 激活函数而非 SwiGLU,从 896 个专家中路由 16 个,因此没有现成的 MoE 内核可以直接应用。而视觉路径——视觉塔、投影器、处理器以及 Kimi 的 XTML 媒体格式——是从零开始构建的。下面的所有部分都建立在这个构建过程之上。
混合 KDA 内存管理
注意力 KV 是只追加的。一旦某个 token 的 KV 被计算出来,它就永远不会改变,这使得调度器可以在共享同一前缀的所有请求之间共享一份物理副本,将其保存在基数树中,并在迭代间无缝地进行流水线处理。而 KDA 层的状态则恰恰相反:它是一个固定大小的循环缓冲区,在每个 token 处都会被原地覆盖。因此,调度器免费提供的所有 KV 处理方式——从前缀缓存到重叠调度器,再到推测解码和分页——都必须为这种在读取时就会发生变动的值重新构建。K3 将 69 个 KDA 层与 24 个 MLA 层交错排列,所以这并非一个边缘情况。它占据了模型的大部分。接下来就是这种重建工作,它唯一令人满意的特性是,每一个部分都是由“原地覆盖”这一事实所决定的。没有哪部分是事后添加的。
单请求路径
同一条路径同时服务于重叠调度器、推测解码以及 page_size > 1 的情况,这些组合对于循环状态模型来说以前是互斥的。一个请求可以命中缓存的前缀,恢复循环状态,运行一个多 token 草稿,验证它,然后提交,所有这一切都在调度器提前一步准备下一批次的同时进行。
前向流上的三种状态移动
请求的实时状态位于单个工作槽中,每次前向计算都会原地读取并覆写该槽。缓存该状态意味着要从 GPU 正在主动修改的内存中拷贝数据,这会产生两种竞态条件。恢复操作可能与正在写入该槽的前向计算发生冲突;而刷新缓存的快照操作则可能与仍在读取前一个快照的捐赠操作发生冲突。这两种情况都无法通过设备级同步或在热路径上加锁来解决。
首先,每次状态拷贝都是串行前向流上的一个内核。将缓存的检查点恢复至工作槽的写时复制操作,以及在轨迹边界捕获状态的快照操作,都会被排入产生和消费该状态的前向计算之间,因此同一流上的顺序天然提供了"先发生"关系。每个预填充块会拍摄一次快照,解码阶段每个轨迹间隔也会拍摄一次。
其次,唯一离开请求的状态移动并不传输任何字节。快照被存入我们称为额外缓冲区的乒乓对中,其中一个槽保存最新快照,另一个槽接收下一个快照,且第二个槽仅在需要它的边界处分配,并在之后立即释放。在每个预填充块之后、从预填充切换到解码时以及请求完成时,缓存会将最新快照的槽索引捐赠给树结构。新的槽会回填该缓冲区,因此快照的写入目标与捐赠的读取源永远不会是同一个物理槽。
三种状态移动:写时复制、快照和捐赠,以及每种操作相对于串行前向流的位置。
按请求分配的槽与共享检查点
由于可复用的前缀检查点位于共享的、可驱逐的基数树中,正在运行的请求仅需预留少量临时槽,最少只需四个。这包括其工作槽、一个用于快照的额外缓冲区槽,以及两个保留余量槽——一个用于必须常驻的已提交前缀状态,另一个用于分叉分支所需的副本。大部分缓存状态由整棵树分摊,而非按每个请求单独计费,因此状态池无需随每个请求保持的历史记录量增长。
针对循环状态的前缀缓存
循环状态无法在任意 token 处进行切片,因为你无法将其反向运行到更早的位置,因此它仅在分块边界处设置检查点,即便在那里也是稀疏分布的。每条路径上通过每路径上限和 LRU 策略仅保留少数几个检查点,且检查点可以独立于其所标注的 KV 缓存被驱逐,留下节点作为墓碑标记。分支点会获得特殊处理,这一思路源自 Marconi。分叉是未来所有分支都保证共享的唯一前缀,因此当请求在边中间发生分叉时,它会从最近的上方检查点重放,并在按分块对齐的分叉处设置一个新检查点。下一个分支直接从此处恢复,无需重放。
基数树上的检查点。稀疏检查点覆盖层与分支点。
这些功能中没有一个是独立于其他功能拼接而成的。它们是一个事实的四个推论,即状态会覆盖自身,这正是该设计在同时吸收重叠调度器、推测解码、分页和前缀缓存时仍能保持简洁的原因。
统一内存:为两种状态共用的单一池
上述所有内容都在各自独立的池中管理每种状态,而这些池本身是设计中仅剩的一个猜测点。两种分配单元的大小相差三个数量级:每个请求一个大的 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 块,这两种情况都无需重新配置任何参数。
统一内存功能以 opt-in 方式提供,通过 `--enable-unified-memory` 启用。后续文章将详细介绍其实现细节。
基于 DSpark 的推测性解码
K3 搭载了 DSpark 块级推测性解码,由我们为 K3 训练的草稿模型驱动。该集成方案中有两个部分值得单独讲述:一是仅在值得投入验证预算的地方进行验证,二是让 K3 的循环 KDA 状态在推测过程中得以完整保留。
只验证值得验证的内容
DSpark 提出每步生成一个草稿 token 块,目标模型通过一次前向传播验证整个块。在批大小为 1 时,额外的验证位置基本是免费的:该步骤受延迟限制,多几个 token 只是顺带处理。随着批处理量增加,情况就不再如此。验证 token 现在与每个其他请求竞争相同的步骤时间,而且其中大多数 token 的赌注最终都会输掉:在聊天负载下,接受长度约为 2.7,因此典型步骤中验证的八个位置中有五个会被拒绝。验证所有内容需要为服务器随后丢弃的 token 支付全价。
解决这个问题的组件已经存在于系统中。草稿模型带有一个经过训练的置信度预测头,用于预测每个位置的 token 有多大可能通过验证。另一方面,对服务器的一次性性能分析记录了在不同负载水平下每个额外验证 token 的实际成本。一个每步规划器将两者结合起来:每个请求仅在预期价值覆盖边际成本时才保留验证 token,其余窗口在目标模型前向传播启动前被裁剪掉。剩余部分照常进行验证,因此输出保持无损。这种权衡是用稍短的接受序列换取更便宜的步骤。
在负载下,裁剪策略胜出。在聊天面板(接受长度约 2.7,左图)和少样本数学面板(接受长度约 5.0,右图)上,解码吞吐量的对比:全部验证 vs 裁剪。在批大小达到 8 之前两者持平,随后差距随批大小扩大而拉开:在批大小为 256 时,吞吐量分别提升 +68% 和 +24%,接受长度从 2.7 降至 2.2、从 5.0 降至 4.3,同时吞吐量上升。
实测的成本曲线才是有趣的部分。验证 token 的边际成本并非平滑的。它呈阶梯状:平坦的台阶上,另一个 token 几乎不花代价就能插入当前的内核计算波;陡峭的上升段则意味着需要启动新的计算波。规划器会读取这个曲面。当一个请求的下一个 token 位于廉价台阶上时,规划器会保留它们;当它们会触发上升段时,规划器就在那里截断。在数据中,这表现为接受长度曲线上的微小波动,而吞吐量保持平滑且单调递增:规划器是在利用真实的硬件台阶,而非噪声。
在批次大小低于 8 时,几乎没有需要缓解的压力,因此剪枝的效果基本持平或略有负面;在规划器中采用小批次提前退出是已知的后续改进方向。
ReplaySSM:针对 KDA 状态的原始输入重放
推测性解码会一次性验证 γ+1 个草稿 token,并且可能只接受其中一部分前缀。对于 MLA 来说,这是没有代价的,因为 KV 只支持追加写入,被拒绝的草稿只需释放其槽位即可。如上一节所述,KDA 层的状态在每个 token 处都会自我覆盖,因此基线方法通过暴力手段实现可逆性,在每个草稿步骤后对整个 K×V 状态进行快照。当 K=V=128 时,每个请求、每层、每个注意力头需要 64 KB,再乘以 γ+1 个步骤。在 K3 的 69 个 KDA 层以及完整批次下,其大小超过了与之竞争的状态池的持久化容量,并且由于它是为每个正在运行的请求预留的,因此会限制并发能力。
存储输入,而非状态。ReplaySSM 放弃了快照。验证内核读取已提交的检查点,但从不写入它,同时在此过程中还会存储每个步骤的原始输入 Sᵢ = (vᵢ, kᵢ, gkᵢ, βᵢ),大约 1 KB,而快照的成本是 64 KB。一旦采样器确定了被接受的长度,一个覆盖所有层和注意力头的单一折叠内核就会从检查点开始,仅重放被接受的前缀,并就地推进状态。被拒绝的草稿永远不会被重放,因此回滚没有成本。草稿窗口从 512 KB 降至 16 KB,大约缩小了 32 倍。
精确,而非近似。该折叠是验证递归的逐字克隆,具有相同的分块和相同的归约顺序,并且它消耗的是验证内核自身存储的门控值,而非重新计算它们。这后半部分并非可有可无。早期版本在 torch 端使用了一个略有差异的公式重新计算门控值,导致每个输出看起来都正确,但底层状态却在悄然漂移。重建后的状态现在与递归基线本应提交的状态在比特级别上完全一致。
其回报在于容量而非速度。每步解码时间保持不变,因为这是用内存换计算的权衡,而非更快的核函数。快照暂存区是推测解码特有的固定开销,随批次大小和γ值增长而增加,因此当γ值大到值得使用时,它恰恰对状态池造成了最大挤压。将这部分内存归还后,并发上限提升了数倍。在旧上限以下,两条路径效果相当,ReplaySSM 因折叠操作和缓冲区写入而略慢一些;超过旧上限后,基线方案开始排队,而 ReplaySSM 仍能持续接收请求,差距正是在这里拉开的。
KDA ReplaySSM。验证步骤(每层一次融合启动)读取检查点 h₀ 但不写入,并将每个推测步骤的原始输入存储到每个槽位的缓冲区中,每个层和每个注意力头各有一个缓冲区。在采样器确定接受长度后,折叠核函数(所有层和注意力头一次启动)仅重放已接受的前缀,并原地覆写该槽位,使用验证核函数自身存储的门控值执行精确的循环计算,因此重建后的状态与循环基线版本在比特级别完全一致。
核函数
在 2.8T 混合模型上执行单序列解码步骤并非计算受限——而是启动次数和延迟问题:每个 token 需要处理 93 个注意力层(69 个 KDA + 24 个 MLA)加上 92 个潜变量 MoE 层,在初始阶段,每一步会启动数百个小核函数。优化策略以性能剖析为导向:融合一个模块,在固定协议上进行 A/B 测试,以 GSM8K 为门控指标,重复迭代。
消除启动与拷贝(P1–P4,+19.9 tok/s)。目标是缩小步骤规模,而非加速核函数:MoE 前端合并为一个 GEMM 操作,KDA 的成对瘦投影也合并(P1);基于性能剖析的扫描消除了每层的类型提升、拷贝和冗余启动(P2、P3);路由变为一次性的寄存器驻留基数选择,作用于 [M, 896] 的 logits 上(P4)。
NVIDIA 计算内核(P5–P8,+10.3 tok/s)。四个阶段替换了与 NVIDIA 联合开发或由其提供的内核,这些内核在默认情况下对于 bs=1 的形状是不正确的:融合的 KDA 解码内核、我们路由旁路背后的 trtllm-gen W4A8 SiTU MoE cubins、TMA 注意力-残差聚合,以及在小 M 场景下替代 cuBLAS 的 CuTe-DSL TGV bf16 GEMM——每个都通过了与其他所有内容相同的 A/B 测试和精度验证关卡。
通信融合(P9、P12、P13,+27.6 tok/s)。最大的阶段:针对 CustomAllReduceV2 对称内存平面上的 MNNVL 架构,实现了一个融合的 all-reduce 系列——小消息使用单次多播存储,大消息使用 NVLS 交换机内归约,同时将残差加法和 RMSNorm 嵌入到集合通信内部(P9)。然后,MoE 的最终化步骤移入集合通信的暂存阶段(P12),而上投影则从复制式 GEMM 转变为列并行 GEMM,并通过多播 all-gather 完成(P13)。
重叠与前序融合(P10、P11、P14、P15,+10.4 tok/s)。其余部分削减了剩余的关键链路:步进输入的 MXFP8 量化(P10)、将残差写回融合到上游内核尾部并在侧流上使用独立分支(P11)、KDA 的 GEMV 链与 qkvg GEMM 重叠(P14),以及将 MLA 解码前序融合到一个内核中,并使用 PDL 为其后面的注意力内核做好准备(P15)。
四个柱状图背后的时间顺序阶梯:
| 阶段 | 优化 | tok/s | 依据 |
|---|---|---|---|
| P0 | 启动基线(Marlin W4A16 MoE) | 44.3 | 实测 |
| P1 | MoE-Front GEMM + KDA GEMV 合并 ³ | 53.2 | 基于 itl 重新缩放 |
| P2 | 启动/复制消除 + O-Gate 融合 ³ | 61.9 | 基于 itl 重新缩放 |
| P3 | 融合 Mamba 轨道 + 批处理路由器 | 62.4 | 实测 |
| P4 | Radix-Select 路由器 | 64.2 | 实测 |
| P5 | Conv + KDA + Onorm 融合(NVIDIA 内核)¹ | 65.3 | 实测 |
| P6 | W4A8 SiTU MoE Cubins(NVIDIA)+ 路由旁路 | 71.0 | 实测 |
| P7 | TMA 注意力-残差聚合(NVIDIA 内核) | 72.0 | 实测 |
| P8 | CuTeDSL TGV 密集 GEMM(NVIDIA 内核) | 74.5 | 实测 |
| P9 | NVLS AR + RMSNorm 融合 | 84.3 | 实测 |
| P10 | 零拷贝 MXFP8 量化 | 85.2 | 实测 |
| P11 | 残差写回融合 + 多流 | 90.2 | 实测 |
| P12 | LL AR 融合(MoE 最终化在集合通信内部) | 92.0 | 实测 |
| P13 | 列并行 GEMM + 多播 AG ² | 108.0 | 实测 |
| P14 | KDA GEMV 侧流重叠 | 111.4 | 实测 |
| P15 | KV-分散 + Q-拼接融合(+PDL) | 112.5 | 实测 |
具有普适性的教训。All-reduce 是一个同步点,因此在那里节省一微秒会一对一地转化为步长时间;而位于另一条流的重叠裕度中的内核,其转化率大约只有十分之一。在编写内核之前,先检查追踪中的关键路径成员资格,是本次优化行动中最高杠杆率的习惯。
¹ 一次基准线重设将 P5 的基线从 64.2 移至 63.8;步长是针对重设后的基线测量的。² P13 是在一次合并窗口后重新校准的规范基线,并非单一 PR 的归因。³ 通向 P1–P4 的各个阶段也包含了一些后来被取代的临时优化(concat all-reduce → P9;tiny/1-CTA GEMV → P8;radix router v1 → P4;Marlin top-k-sum → P6/P12;early attn-res add → P7);它们的瞬时增益体现在曲线中,但没有对应的命名点。
并行化 K3
对于 K3 的混合架构,传统的并行化选择并不适用。张量并行无法对 MLA 的 KV 缓存进行分片(只有一个 KV 头,没有可分割的对象),因此每个 rank 都持有完整副本;它还将每个 GEMM 切分为八份,并在每一层付出一次集合通信的代价。纯 DP 注意力机制则是在每个 rank 上复制注意力权重,KDA 约需 61 GB,MLA 约需 11 GB,这些内存是 KV 缓存和 KDA 状态所需要的。预填充和解码在这些代价下失败的方式不同,因此 K3 按阶段拆分答案:预填充使用分块流水线并行,解码使用上下文并行。
预填充:分块流水线并行
在 TP 预填充中,每一层都以一个 AllReduce 结束,这是一个无法与计算重叠的屏障。流水线并行则按层切分模型:K3 的 93 层被分为 8 个阶段,提示词被切分成多个块,在这些阶段中流式传输:
分块流水线并行预填充。各个阶段同时处理不同的块。阶段之间的交接在阶段计算下一个块的同时进行,因此在 K3 上,91% 的交接时间被隐藏了。在 TP(底部条带)下,每一层都以一个所有 rank 都在等待的 AllReduce 结束。
这种方案在三方面具有优势。唯一剩余的通信环节——即向下一阶段的交接——隐藏在下一块的计算过程中。每个等级(rank)运行完整的层,因此通用矩阵乘法(GEMM)的宽度是原来的八倍,效率更高。同时,每个阶段仅为其约12层保留键值缓存(KV)和激活值,这使得超长提示词(prompt)的预填充(prefill)更加从容。不过,流水线必须足够深:浅层的PP4×TP2无法掩盖其交接开销,且仍需承担TP2的全规约(AllReduce)通信,其基准测试表现也不如TEP8。
在2×4 GB300上,以8K预填充(prefill)为测量场景,仅改变拓扑结构,结果如下:
深度流水线并行(Deep PP)在两个维度上均胜出。左图:在c1–c4交叉点之后,PP8×TP1的吞吐量上限攀升至TEP8的约1.7倍,且首令牌延迟(TTFT)更低。右图:在每等级(per-rank)计算量(FLOPs)相等的情况下,每千个预填充(prefill)token的成本;PP4×TP2与TEP8在计算与通信之间相互权衡,结果持平,而PP8在两个维度上成本均为最低。
PP8仅在处理单个请求时表现不佳,而一个空闲的预填充(prefill)工作节点本身配置就有问题。在K3的解耦式服务(disaggregated serving)架构中,预填充节点运行PP8。每个预填充节点的预填充容量是TEP8节点的1.45至1.72倍,因此一个预填充节点即可为多个解码(decode)节点提供充足的数据。解码节点则运行TP或DCP,这将在下一部分进行介绍。
解码:上下文并行(Context Parallelism)
解码阶段正是复制型键值缓存(KV cache)的瓶颈所在:在TP模式下,每多处理一个请求或多增加一千个token的上下文,每个等级(rank)上都需要占用相同字节数的存储。解码上下文并行(DCP)按token位置而非注意力头(head)来分片MLA的键值缓存(KV)。当位置p满足p模N等于r时,该位置由等级r负责,因此每个等级持有所有请求上下文中交错分布的1/N,这种分片方式在注意力核(attention kernel)层面是不可见的:
按位置分片。以4个等级(rank)上的16个token位置为例:TP存储了64个物理副本,而DCP仅存储16个。释放出的字节转化为逻辑上的键值缓存(KV)容量,在K3上采用DCP8时,容量提升约7.9倍。
位置分片会破坏 softmax,因为每个 rank 只看到 1/N 的键,而部分 softmax 结果无法直接相加。解决方案来自 FlashAttention 自身:每个 rank 返回其部分注意力输出以及每个注意力头的 log-sum-exp,然后每层通过一次 all-to-all 通信交换这些数据,使得每个 rank 最终拥有完整上下文上 1/N 的注意力头。随后通过 log-sum-exp 进行本地合并是精确的,其结果已经是 TP 下输出投影层所期望的注意力头布局。每层一次集体通信就是全部通信开销:
解码步骤,每层执行一次。每个 rank 在本地对完整注意力头的查询进行投影,仅关注其拥有的位置,并通过一次打包的 all-to-all 通信发送部分输出及其 log-sum-exp。本地合并后得到标准的 TP 注意力头布局;注意力机制之后的所有环节均无需更改。
其余部分保持不变。DCP 组构建在 TP 组内部,因此 TP8 搭配 DCP8 仍然是 8 块 GPU,MoE 也沿用其已有的并行方式。KDA 是印证这一规则的例外:其状态是每个请求一个固定大小的矩阵,而非每个 token 一个,因此不存在可进行分片的位置轴,KDA 层仍按注意力头进行 TP 分片。整个功能仅由一个标志控制:`--dcp-size N`。
这一特性带来的收益在智能体流量中得以体现,这类场景下数万 token 的会话会堆积在缓存中。在 2×4 GB300 上回放真实的编码智能体会话,两侧均配备主机内存 KV 层级,唯一区别在于是否使用 DCP:
DCP 消除了活跃集瓶颈。主机层级缓解了重新预填充流量,但在 16 个并发会话时,活跃工作集超出了 TP8 设备 KV 的容量,吞吐量随之崩溃。DCP8 将逻辑 KV 从 150 万 token 提升至 1220 万 token,并在 48 个会话时仍能承载相同的工作负载,达到 541 tok/s。
DCP 与堆栈的其他部分协同工作。DSpark 的验证步骤本质上是一个解码步骤,因此它走的是相同的复制 Q、一对全路径;在 PD 分离架构下,预填充端保持对 DCP 无感知,每个解码节点只需在传输边界拉取自己拥有的位置,这使得 PP 或 TP 预填充能够为 DCP 解码提供数据。剩下的瓶颈是 KDA:其每个请求的状态无法按位置分片,因此一旦 DCP 突破了 MLA 的限制,运行中的请求上限就成为新的约束边界;上文内存部分提到的统一内存设计正是为了解决这个问题。
将所有策略组合在一起
这些组件最终如何组合,需要通过实际测量来确定。将 PD 分离纳入循环,配合全程使用的分块 PP8 预填充,并改变解码拓扑结构和预填充与解码的比例,就得到了服务前沿:
服务前沿。在吞吐量端,DCP 组合——两个 PP8 预填充工作节点为两个 DCP8 解码节点提供数据——实现了每 GPU 2633 tok/s 的性能,与最佳 TP8 方案持平。向右移动是预填充与解码比例的调节效果:一个 PP8 预填充工作节点为两个、三个、四个独立的 TP8 解码实例提供数据,以牺牲聚合吞吐量为代价换取更高的单用户速度,最高可达每用户 86 tok/s 以上。
强化学习:在原生 MXFP4 基座上进行 LoRA 训练
K3 的 Day-0 强化学习是在 Miles 上进行的共置 LoRA 训练:一个基于 Miles 的 Megatron 后端的 BF16 训练器,与原生打包的 MXFP4 SGLang 推理引擎共享同一批 64 个 GB300。该后端支持 KDA、NoPE-MLA、注意力残差库和潜在 MoE,并采用了 TP/SP/PP/CP/EP 并行策略。
LoRA 服务与权重同步
推理引擎按原样加载检查点,绝不重写。每一步仅传输 BF16 LoRA 适配器,引擎将增量作为独立的 BF16 B(Ax) 项叠加在量化后的基础 GEMM 之上,因此策略更新能以全精度作用于 4 位基础权重之上。密集投影通过 SGLang 的 Triton LoRA 后端处理,896 个路由专家通过 Marlin 路径上的融合 MoE-LoRA 内核处理,共享专家增量则被折叠进融合 MoE 前向 GEMM 中。适配器驻留在 GPU 内存池中,每次同步都会原地替换,因此推理侧不存在常驻的 BF16 副本,没有全权重重新同步,循环中也没有重新量化步骤。
并行策略
流水线并行。K3 的注意力残差快照库需要跨越流水线阶段边界,但 Megatron 的点对点通信只携带一个隐藏状态张量,因此阶段边界将 [前缀和, 快照库] 打包进该张量,并在进入时解包。Megatron 保持不变。
上下文并行。快照库天然分片,因为每个注意力残差操作都是按 token 进行的。MLA 也不需要 K3 特定的代码:旋转表切片是 Megatron 的 MLA 中唯一需要感知上下文并行的工作,而 K3 没有旋转位置嵌入,因此其投影位于标准 TE 注意力核心上,并继承了上下文并行。KDA 是需要处理的部分,它通过 fla 的上下文并行上下文来处理循环状态和卷积边缘。它需要一个连续的、按 rank 划分的本地数据块,而 Megatron 存储的是环形注意力所期望的锯齿顺序,因此重排仅在 KDA 周围进行。
专家并行。适配器将其潜在侧因子共享给全部 896 个路由专家,并保持另一个因子按专家独立,这与引擎的融合 MoE-LoRA 契约相匹配。该共享因子在专家并行(EP)间复制,但标记为专家并行,因此 DDP 仅对专家数据并行(expert-DP)进行规约,而 EP 求和是本次发布自行添加的唯一梯度规约操作。
内存
原生 MXFP4 推理侧在初始化时峰值接近每 GPU 225 GiB,BF16 训练侧接近每 GPU 155 GiB,而显存为 277 GiB,因此两者绝不会同时驻留:一方运行时另一方休眠,而共置依赖于每次交接后不留下任何残留。
基座模型始终驻留。引擎在整个运行过程中将基座权重保留在 GPU 上,仅在训练器工作时释放 KV 缓存和 CUDA 图。由于基座从不释放,因此也无需恢复,LoRA 路径完全不传输基座字节:基座同步被完全跳过。
进程组保持活跃。训练器的 NCCL 通信器缓冲区不属于 torch 分配器内存,因此卸载无法释放它们,显而易见的解决方案是在休眠时销毁进程组,在唤醒时重建。这会将常驻的几 GiB 内存换取为每个周期的重建和专家并行预热开销,只有当引擎需要回收这些字节时才值得付出——而在此场景中并不需要,因为没有基座权重经过更新路径。进程组保持活跃。
适配器传输按块限定范围。适配器包含约 2,800 个张量。每个块作为一个扁平的 CUDA IPC 桶传输,共 278 个,一旦接收方确认且所有生产者秩已越过引擎组屏障,即被收集。这限制了每个块的瞬时 IPC 内存,而非让其在传输过程中累积,峰值时约为每 GPU 48 GiB。
主机副本仅存在于被读取之处。引擎在将每个适配器安装到 GPU 池后释放其 CPU 副本,将调度器 RSS 从 76–88 GiB 降至约 17 GiB;CPU 副本已消失的适配器无法重新安装,因此池驱逐会返回错误而非静默提供过期槽位。在训练器侧,DDP 缓冲区按生命周期拆分:适配器参数缓冲区保留在 CPU 支持区域,因为权重更新在训练器休眠时读取它们,而梯度缓冲区可重建,进入休眠时丢弃的无备份区域。
验证
训练正确性是强化学习支持栈最重要的部分。我们在配方工作之前就构建了这些检查。
训练/推理展开 KL 散度,每次展开时记录,是基于采样 token 上 KL(展开 ‖ 训练) 的 Schulman k3 估计值,由引擎返回的对数概率与训练器重新计算的对数概率计算得出。这是一个诊断指标,并非目标函数的一部分,目标函数不携带任何 KL 惩罚项。其下限约为 2e-3,由以 MXFP4 格式提供基础模型所设定;随步数增长则是发散信号。
金丝雀同步校验探针。一条轨迹从首次展开起固定,每一步都由引擎(在实时适配器下)和训练器(在同一策略版本下)分别评分。两条曲线的共同移动证明了训练后的适配器确实到达了推理端:传输校验和只能证明字节已到达,无法证明有任何东西在读取它们。
张量级转储对比。相同的 token 在两个构建版本或两种并行布局上运行,转储每一次前向激活和每一个参数梯度,通过相对 L2 范数和余弦相似度与实测噪声底限进行比较。比特级精确并非标准,因为同一套算术运算在两套不同的 GPU 上已经会在最低有效位级别产生差异。这正是流水线边界和上下文并行布局的验证方式:当 CP=2 和 CP=1 构建出完全相同的调用时,两者是比特级一致的,而前向内核替换后的余弦中位数为 0.996。
权重同步断言。训练器与引擎之间的逐张量 SHA256 清单;一个验证器,在适配器更新未改变任何导出张量时报错;一项检查,确保每个 B 因子在版本 1 时均为零;以及适配器梯度和优化器步数检查。
训练结果
本次报告的运行是 16 节点 × 4 GB300 上的 DAPO 数学任务:BF16 训练器采用 TP8 / PP8 / EP8 配置,与原生 MXFP4 推理展开引擎共置,4096 token 响应,每次展开 64 个样本,每次展开进行一次优化器步进,秩 32 / α-64 的 LoRA 学习率 1e-5,无 KL 项的 GRPO,运行至 12 小时墙钟时间限制。AIME-2024 贪婪评估从 43.3% 攀升至 76.7%,经过 60 步,30 道题中从 13 题答对增至 23 题,且在截止时仍在上升,而训练/展开 KL 散度在整个运行过程中稳定保持在约 2e-3 的 MXFP4 下限水平。
12 小时的 DAPO 运行。左图:每轮 rollout 中 64 个采样响应的平均奖励,基于响应通道评分;每轮 rollout 使用不同的提示批次,因此趋势是有效信号,而单步数值是噪声。中图:AIME-2024 贪婪通过率,每 10 轮 rollout 评估一次,采用与训练相同的 4096 token 限制,因此超出该限制的正确解法会被计为未通过。右图:在采样 token 处对 KL(rollout ‖ train) 的 Schulman k3 估计值,仅报告而不加入损失函数;在量化底限处保持平坦是目标,单调增长则是散度特征。
致谢
本工作是 RadixArk 的 SGLang & Miles 团队与月之暗面(Moonshot AI)团队的合作成果,合作方还包括 NVIDIA、AMD、Approaching AI、Baseten 和 Modal。
AMD:黄文国、宋欣怡、肖海、林索加、王杜一
Approaching AI:沈焕明、张晓浩、李楠、张明兴
感谢 DigitalOcean 为我们的测试提供 AMD 实例。
感谢 Google Cloud、DigitalOcean、Nebius、fal、RunPod、DeepInfra 和 GMI Cloud 在 SGLang 上部署 Kimi K3。