6.0
关注
AI SCORE
工具产品2026-07-27 00:00
SGLang 和 Miles Day-0 支持 Kimi K3，3 万亿参数开源模型推理就绪
LMSYS / SGLang Blog#SGLang#Kimi#开源模型
Editor brief · 编辑速览
推理框架完整支持混合架构的 3T 参数开源模型。便于开发者在 K3 上进行部署和优化。
安全高效的状态复用
统一内存：一个池容纳两种状态
使用DSpark的投机解码
仅验证值得验证的部分
ReplaySSM：KDA状态的原始输入重放
预填：分块流水线并行
解码：上下文并行
合并所有策略
强化学习：基于原生MXFP4基础的LoRA训练
LoRA服务和权重同步
我们很高兴宣布在SGLang和Miles中添加对Kimi K3的第0天支持。K3是第一个开源的三万亿参数级模型，其混合架构在服务栈几乎每个有假设的地方都有所偏离。与月之暗面AI和英伟达团队合作，两个系统在发布当天全面覆盖K3：SGLang用于推理，Miles用于强化学习训练。本文介绍了所需的工作。
一个新的混合架构。2.8T参数，69个KDA线性注意力层与24个MLA、LatentMoE和注意力残差交错；服务栈的大多数假设在某个地方会被打破。
两种状态的内存管理，包括原地重写自身的循环状态、前缀缓存、重叠调度和分页的重建，以及移除最后一个大小猜测的统一池设计。
一个核心梯队，在批次1时达到约113 tok/s，投机前。
DSpark投机解码，使用我们为K3训练的草稿模型：批次1解码达到约423 tok/s。ReplaySSM处理KDA状态，重放原始输入而不是按步骤快照，将草稿窗口内存减少约32倍。
按阶段分割的并行：分块流水线并行预填和张量并行解码，在PD解聚下组合到每个GPU 2,808 tok/s。
使用Miles进行LoRA强化学习，在原生MXFP4检查点上，共置训练器和回收在相同GPU上：AIME-2024从43.3%提升到76.7%，耗时12小时。
启动命令和按工作负载的配置指导见Kimi K3 cookbook。
Kimi K3是第一个开源的三万亿参数级模型：2.8T参数、100万token上下文窗口和原生视觉理解。其前代K2.5在架构上是SGLang已经很好地服务过的模型的亲戚，所以栈的大部分可以直接应用。K3不是这样——它在多个独立的地方同时偏离约定：
这些行中的每一行都是一个服务问题。混合注意力栈是使100万上下文成本合理的原因，但这意味着服务器同时持有两种状态：每个请求一个固定大小的KDA状态，旁边是MLA的每token KV——下面的内存管理部分正是关于这一点。注意力残差在整个栈中传递一组注意力输出，这打破了SGLang标准层管道所做的假设，并在DP注意力等地方强制使用K3特定路径。LatentMoE在使用SiTU激活的下投影潜在空间内路由16个896个专家，而不是SwiGLU，所以没有现成的MoE核可以直接应用。而视觉路径——塔、投影仪、处理器和Kimi的XTML媒体格式——是从头开始带来的。下面的部分都建立在这个带来上。
K3的混合KDA和MLA架构产生两个内存管理挑战。首先是使前缀缓存对可变KDA循环状态安全有效。其次是在KDA状态和MLA KV之间动态共享容量。
注意力KV是仅追加的。一旦计算，它永远不会改变，所以调度器可以安全地在基数树中跨请求共享缓存的前缀。KDA状态不同。每一层维护一个固定大小的循环缓冲区，在每个token都被原地重写，所以前缀缓存必须管理可变状态而不是不可变KV。
在SGLang中，我们使用三个显式状态移动——在基数树和请求的工作槽之间——来安全地缓存和复用可变KDA状态。写时复制在前向变异它之前将共享检查点恢复到私有槽。快照捕获高级状态，捐献将该快照转移到树。
三个状态移动。写时复制、快照和捐献，以及每个相对于串行前向流的位置。
该设计在构造上是无竞争的。恢复和快照副本在产生和消费状态的操作之间的前向流上入队，因此流排序防止与原地更新的竞争。快照在乒乓额外缓冲区的两个槽之间交替，因此下一个快照不能覆盖正在附加到树的状态。捐献仅转移一个槽索引且不复制任何状态。第二个槽在边界处延迟分配，并在之后立即释放，避免每个请求永久槽。这既不需要设备范围同步也不需要热路径上的锁。
检查点放置决定了缓存命中率。循环状态不能向后运行，必须从较早的检查点向前重放。我们仅在对齐的基数树节点、预填期间的块边界和解码期间的固定token间隔处取检查点，并在每路径上限和LRU下保持稀疏集。受Marconi启发，我们优先考虑分支点，因为它们的前缀由所有子分支共享。当请求在边缘中途发散时，它从最近的检查点上方重放，并在对齐的分叉处植入一个，允许以后的分支直接在那里恢复。
基数树上的检查点。稀疏检查点覆盖和分支点。
无竞争状态移动和分支感知检查点放置提供了高复用KDA前缀缓存，可以与重叠调度器、投机解码和分页KV清晰组合。
上面的所有内容在各自的池内管理每种状态，而池本身是设计中剩下的唯一猜测。两个分配单位相差三个数量级：每个请求一个大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训练的草稿模型驱动。集成的两个部分值得自己讲故事：仅在需要时花费验证预算，以及使K3的循环KDA状态在投机中存活。
DSpark 每一步提出一组草稿 token，目标模型通过一次前向传播验证整个块。在批大小为 1 时，额外验证位置基本没有成本：该步骤受延迟限制，多几个 token 可以顺带完成处理。但随着批次变满，情况就不再如此。验证 token 现在会与所有其他请求争夺同一段步骤时间，而其中大多数最终都押错了：在聊天工作负载中，接受长度约为 2.7，因此典型步骤中八个待验证位置有五个会被拒绝。验证全部位置意味着服务器要为随后丢弃的 token 支付完整成本。
修复这一问题所需的组件其实已经存在于系统中。草稿模型带有一个训练过的置信度头，可以逐位置预测每个 token 通过验证的可能性。另一方面，对服务器进行一次性性能分析，可以记录在各个负载水平下增加一个验证 token 的实际成本。逐步骤规划器将两者结合起来：对于每个请求，仅在验证 token 的期望价值足以覆盖其边际成本时才予以保留，其余窗口会在目标模型前向传播启动前被裁剪。剩余部分仍按原方式精确验证，因此输出依然无损。代价是接受的连续 token 略少，但每一步的成本也更低。
裁剪在高负载下效果显著。图中展示了聊天面板（接受长度约 2.7，左）和少样本数学面板（接受长度约 5.0，右）上，全部验证与裁剪两种方案的解码吞吐量对比。批大小不超过 8 时两者基本持平，之后差距随批大小增大而拉开：批大小为 256 时，吞吐量分别提升 68% 和 24%；与此同时，接受长度分别从 2.7 缓降至 2.2、从 5.0 缓降至 4.3，而吞吐量则持续上升。
实测成本曲线才是其中最有意思的部分。验证 token 的边际成本并不平滑，而是呈阶梯状：在平坦的平台区间，新增 token 几乎可以零成本地嵌入当前的 kernel 执行波次；在陡峭的上升区间，它则会触发新的波次。规划器会读取这一成本曲面。当请求的后续 token 落在低成本平台上时，规划器会保留它们；当它们会触发一个上升阶梯时，规划器便在那里截断。数据中，这表现为接受长度曲线上的细微波动，而吞吐量始终平滑且单调递增：规划器是在利用真实硬件的平台区间，而不是追逐噪声。
批大小低于 8 时，需要缓解的压力很小，因此裁剪在这一范围内只能持平或略有负收益；已知的后续改进是在规划器中增加小批次提前退出机制。
推测解码会一次验证 γ+1 个草稿 token，但可能只接受其中一个前缀，因此状态必须支持回退。KDA 层的状态会在每个 token 到来时覆盖自身，所以基线方案会在每个草稿步骤后保存完整的 K×V 状态快照：当 K=V=128 时，每个请求、每层、每个头需要 64 KB，再乘以 γ+1 个步骤。K3 共有 69 个 KDA 层，这些暂存空间会为每个正在运行的请求预留，因此会与持久状态池争夺内存，并限制并发量。
ReplaySSM 改为保存每个草稿步骤的原始输入 Sᵢ = (vᵢ, kᵢ, gkᵢ, βᵢ)，大小约为 1 KB，由验证 kernel 在处理过程中顺手写入。采样器确定接受长度后，一个覆盖所有层和头的折叠 kernel 会从已提交的检查点开始，仅重放被接受的前缀，并原地推进状态。该折叠过程逐字复刻了验证阶段的递归关系，并直接使用验证 kernel 保存的门控值，而不是重新计算，因此重建出的状态与递归基线原本会提交的状态在比特级完全一致。
内存方面：草稿窗口从 512 KB 降至 16 KB，约缩小 32 倍。将这部分内存归还状态池后，并发上限可提升数倍。
速度方面：验证 kernel 不再需要为每一步写入状态快照，从而减少了内存访问；重放则通过一次融合启动覆盖所有层和头，因此即使在小批大小下，每步耗时也能得到改善。
KDA ReplaySSM。验证过程读取检查点 h₀，但不对其写入，并将每个草稿步骤的原始输入保存到逐槽位缓冲区中，每层、每个头各有一个缓冲区。随后，折叠 kernel 仅重放被接受的前缀，并原地覆盖该槽位。
对于一个 2.8T 混合模型，单序列解码步骤并不受计算量限制，而是受 kernel 启动次数和延迟限制：每个 token 都要经过 93 个注意力层（69 个 KDA + 24 个 MLA）和 92 个 latent-MoE 层；在初始适配阶段，每一步会触发数百个微型 kernel。这轮优化由性能分析驱动：融合一项，在固定协议上进行 A/B 测试，通过 GSM8K 准确性门禁，然后重复这一过程。
启动与复制消除（P1–P4，+19.9 tok/s）。重点是让步骤本身更精简，而不是让 kernel 更快：MoE 前端合并为一次 GEMM，KDA 成对的窄投影也被合并（P1）；通过性能分析引导的全面排查，移除了逐层的向上类型转换、复制以及多余启动（P2、P3）；路由则改为在寄存器中驻留、针对 [M, 896] logits 执行的单遍基数选择（P4）。
NVIDIA 计算 kernel（P5–P8，+10.3 tok/s）。在四个阶段中，当默认实现不适合 bs=1 的形状时，换用了与 NVIDIA 联合开发或由 NVIDIA 提供的 kernel：融合 KDA 解码 kernel、位于路由绕过路径之后的 trtllm-gen W4A8 SiTU MoE cubin、TMA 注意力-残差聚合，以及在较小 M 下取代 cuBLAS 的 CuTe-DSL TGV bf16 GEMM。每一项都经过了与其他优化相同的 A/B 测试和准确性门禁。
通信融合（P9、P12、P13，+27.6 tok/s）。这是提升幅度最大的阶段：基于 CustomAllReduceV2 的对称内存平面，为 MNNVL 互连构建了一套融合 AllReduce 实现——小消息使用单次多播存储，大消息使用 NVLS 交换机内归约，同时将残差相加和 RMSNorm 融入集合通信内部（P9）。随后，MoE 收尾操作被移入集合通信的暂存阶段（P12）；上投影则从复制式 GEMM 改为列并行 GEMM，并通过多播 AllGather 完成（P13）。
重叠与序言融合（P10、P11、P14、P15，+10.4 tok/s）。其余优化进一步缩短剩余的关键链路：针对跨步输入的 MXFP8 量化（P10）；将残差写回融合到上游 kernel 的尾部，同时把相互独立的分支放在旁路 stream 上执行（P11）；让 KDA 的 GEMV 链与 qkvg GEMM 重叠（P14）；以及将 MLA 解码序言融合为一个 kernel，并通过 PDL 预先启动其后的注意力 kernel（P15）。
四个柱形背后的时间顺序阶梯如下：
其中可推广的经验是：AllReduce 是一个同步点，因此在那里节省一微秒，就能一比一地转化为步骤耗时的缩短；而位于另一个 stream 重叠空隙中的 kernel，节省的时间通常只能按约十分之一的比例转化。在编写 kernel 前先从性能轨迹中确认它是否位于关键路径，是这轮优化中杠杆效应最高的习惯。
¹ 一次 rebase 将 P5 的基线从 64.2 调整为 63.8；该步骤相对于 rebase 后的基线进行测量。² P13 是合并窗口后重新校准的规范基线，并非单个 PR 的贡献归因。³ 通往 P1–P4 的若干阶段还包含后来被取代的临时优化（拼接 AllReduce → P9；微型/单 CTA GEMV → P8；基数路由器 v1 → P4；Marlin top-k-sum → P6/P12；早期注意力残差相加 → P7）；它们的短期收益包含在曲线中，但没有对应的命名节点。
对于 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和激活，这使得预填充非常长的提示变得容易。管道必须足够深：浅的PP4×TP2无法隐藏其交接开销，仍然要支付TP2的AllReduce成本，性能不优于TEP8。
在2×4 GB300上以8K预填充进行测量，仅改变拓扑结构：
深管道并行在两个轴上都赢了。左图：超过c1–c4交叉点后，PP8×TP1攀升到约1.7倍的TEP8天花板，且TTFT更低。右图：在每rank FLOP相等的条件下，每1k预填充令牌的成本；PP4×TP2和TEP8在计算和通信之间权衡并打平，PP8在两者上成本最低。
PP8仅在单个请求时会吃亏，而这时预填充工作线程空闲本身就说明配置有问题。在K3的分解式服务中，预填充节点运行PP8。它们每个的预填充容量是TEP8节点的1.45到1.72倍，因此一个预填充节点可以为多个解码节点供数据。解码节点运行TP或DCP，下一部分将介绍。
解码是复制KV缓存约束的地方：在TP下，接受一个新请求或增加一千个令牌的上下文，在每个rank上的成本相同。解码上下文并行化（DCP）按令牌位置而非按注意力头来分片MLA的KV。当p mod N等于r时，rank r拥有位置p，因此每个rank持有每个请求上下文的交错1/N，分片对注意力核上方是不可见的：
按位置分片。相同的16个令牌位置在4个rank上：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是解释规则的例外：其状态是每个请求一个固定大小的矩阵，而非每个令牌一个，因此没有位置轴可以分片，KDA层仍按头进行TP分片。整个功能只需一个标志，--dcp-size N。
这种方法的优势在agentic流量上显现，那里数十万令牌的会话堆积在缓存中。在2×4 GB300上重放真实的编码AI智能体会话，两组都有主机内存KV层，仅DCP不同：
DCP消除了活跃集合瓶颈。主机层抢救了重新预填充流量，但在16个并发会话时，活跃工作集超出TP8的设备KV，吞吐量崩溃。DCP8将逻辑KV从150万提升到1220万令牌，并在48个会话时将相同工作负载维持在541 tok/s。
DCP与栈的其余部分组合。DSpark验证步骤是一个解码步骤，因此它遵循相同的复制Q、一次all-to-all路径；在PD分解下，预填充侧保持DCP-无感，每个解码rank在传输边界处简单地拉取其拥有的位置，这正是使PP或TP预填充能喂给DCP解码的原因。剩余的天花板是KDA：其每个请求的状态无法进行位置分片，因此一旦DCP提升了MLA瓶颈，运行中的请求数上限就成为约束；上面内存部分的统一内存设计应对了这一点。
各部分如何组合最终是一个测量问题。将PD分解纳入循环，改变预填充拓扑、解码拓扑和预填充:解码比例，可得到服务前沿：
服务前沿。在吞吐量端，一个PP8预填充工作者喂一个TP8解码节点在fp4臂上每GPU交付2,808 tok/s，采用DCP组合，两个PP8预填充工作者喂两个DCP8解码节点，紧跟其后为2,633。向右移动是预填充:解码旋钮在工作：一个预填充工作者喂两个、三个、然后四个独立的解码实例，用聚合吞吐量换取per-us