Kimi K3 Performance Optimizations in vLLM: The Road to 2.8× Throughput
Kimi K3 Performance Optimizations in vLLM: The Road to 2.8× Throughput
发表时间: 2026-09 · Blog post by Wentao Ye (Red Hat), Canlin Guo (Huawei), Yongye Zhu (Inferact), Jiangyun Zhu (Inferact), Ziming Huang (Inferact), Wei Zhao (NVIDIA), Michael Goin (Red Hat), Jie Li (Inferact) (blog.vllm.ai)
原文: https://blog.vllm.ai/blog/2026-09-13-kimi-k3-performance-optimization
作者/机构:Wentao Ye (Red Hat), Canlin Guo (Huawei), Yongye Zhu (Inferact), Jiangyun Zhu (Inferact), Ziming Huang (Inferact), Wei Zhao (NVIDIA), Michael Goin (Red Hat), Jie Li (Inferact)
速读
一句话结论 本文针对 Kimi K3 模型在 vLLM 框架中的推理服务,通过自适应调度、KDA 状态优化、零拷贝批处理及解码上下文并行等全栈工程优化,在 8K 输入 1K 输出的投机解码场景下,实现了吞吐量最高 2.8 倍的提升,并将首字延迟降低了 85%。
要解决什么问题 Kimi K3 模型在 vLLM 框架中的高效部署面临多个底层机制卡点。首先,调度器的静态 token 预算在并发请求较少时会导致大量预算闲置,甚至将单个长请求强行切分为多次前向计算,推高了首字延迟。其次,Mamba 风格的前缀缓存(Prefix cache)在处理未完全对齐的缓存块时,会强制在最后一个可缓存块边界截断预填充(Prefill),导致为了计算极短的后缀而触发第二次全模型前向传递。在投机解码(Speculative decoding)场景下,混合批处理需要频繁调用 index_select 和 index_copy_ 来拼接和拆分张量,产生大量细碎的显存拷贝开销;同时,投机解码在每个草稿位置都需要写入 KDA 循环状态以便在拒绝 token 时回滚,带来了极大的显存写入压力。此外,MXFP4 专家内核的 Top-k 聚合存在多余的内核启动和中间张量读写;而在张量并行(TP)下,MLA(Multi-Head Latent Attention)的潜在 KV 缓存在所有计算节点上被强制复制,导致增加 GPU 数量也无法扩大整体 KV 缓存容量,严重限制了长上下文代理任务的并发上限。最后,预填充与解码分离架构下的状态卸载机制,无法直接套用纯注意力模型的假设,因为 KDA 状态和 MLA 状态的分布式存储逻辑完全不同。
怎么做的 针对上述卡点,作者在 vLLM 栈内实施了一系列针对性改造。第一,引入自适应调度预算。当请求数量较少时,动态调大单次可调度的 token 数量上限(例如在单请求时将可用预算从 1024 提升至 8185),确保单个长请求能在一次前向调用中处理完毕。第二,实现内部 KDA 前缀检查点。将检查点导出逻辑下沉到单次预填充传递的内部,只需一次 FlashKDA 调用即可处理全部 token,并在内部指定位置(如第 7680 个 token 处)导出状态,彻底消除了第二次全模型前向传递(涵盖注意力、MoE、路由和 TP 集合通信)的开销。第三,构建零拷贝的混合 KDA 批处理。通过连续的零拷贝视图直接将非投机和投机输入喂给 KDA 路径,并将结果直接写入最终输出的切片中,移除了每层原有的 6 次索引选择和 2 次索引拷贝操作。第四,延迟 MXFP4 最终化。将 MXFP4 的 Top-k 最终化操作融合到潜在尾部内核中,省去了一次内核启动和中间张量的读写。第五,集成 ReplaySSM 机制。在投机解码中不再全量存储每个草稿位置的 KDA 状态,而是仅缓存最近的 SSM 输入,在接受 token 时通过 Triton 内核重建状态,回滚时只需移动缓冲指针,从而在不损失精度的前提下提升了有效缓存容量。最后,引入解码上下文并行(DCP)。针对 MLA 融合路径,将 KV 缓存沿序列维度进行分片,利用对称内存的 All-to-All 处理输出归约,并通过 NVLS 多播收集查询,使得查询分片直接写入消费者的最终缓冲区,打破了 TP 下的容量瓶颈。
效果如何 实验在 B300 节点(CUDA 13.3)上进行,采用 8K 输入、1K 输出的负载,配置为 8 路张量并行(TP8),并使用 DSpark 投机解码(每次生成 8 个投机 token)。对比的基线方法是未包含此轮优化的 vLLM v0.27.1 版本(代表基础的 Kimi K3 适配路线)。量化结果显示,在并发度为 1、4、16 的测试中,优化后的版本表现出全面压制:平均端到端延迟降低了 56% 至 60%(并发 16 时从 55.90 秒降至 22.17 秒);吞吐量提升了 2.2 至 2.8 倍(并发 16 时从 258.3 提升至 725.0 tokens/s);平均首字延迟大幅下降 72% 至 85%(并发 1 时从 2262.9 毫秒降至 376.3 毫秒)。在针对长上下文的独立测试中(120K token 负载,包含 114K 共享前缀),引入 DCP8 后,KV 缓存容量从 1.93M 跃升至 19.75M tokens,且在 GSM8K 任务上保持了 96.97% 的准确率(对比 TP8 的 96.21%),无请求错误。作者也指出了分离架构下的局限与代价:预填充与解码分离及缓存卸载必须同时传输 MLA KV 和 KDA 状态,由于 MLA KV 跨 TP 节点复制,而 KDA 状态按头和维度分片,纯注意力模型的转移假设失效,必须依赖特定的连接器来处理状态恢复并锁定边界状态直到异步写入完成。
主要贡献
Kimi K3 在发布首日即获得了 vLLM 的支持,但为了实现高效的服务部署,需要对整个技术栈进行深入的优化。研究团队发现,KDA 循环状态(recurrent state)、LatentMoE、MXFP4 专家内核(expert kernels)、投机解码(speculative decoding)以及张量/流水线并行(TP/PP)各自暴露了不同的性能瓶颈;同时,调度器的限制和小型张量的拷贝操作对性能的影响甚至可能与大型的通用矩阵乘法(GEMM)相当。
本文的核心贡献在于实现了一系列端到端的性能优化,主要创新点包括:自适应的投机 token 调度预算、内部 KDA 前缀检查点(prefix checkpoints)、零拷贝的混合 KDA 批处理(zero-copy mixed KDA batches)、延迟的 MXFP4 最终化(deferred MXFP4 finalization)、用于重建 KDA 状态的 ReplaySSM、Prefill/Decode 分离与混合状态卸载,以及解码上下文并行(Decode context parallelism)。通过这些优化,系统在延迟、吞吐量和首字时间(TTFT)上均取得了显著提升。
背景与设计原则
高效服务面临的多维度瓶颈挑战:虽然 Kimi K3 能够运行在 vLLM 上,但其复杂的架构设计带来了诸多底层性能挑战。由于模型的特殊性,传统的优化手段不足以解决所有问题。研究团队观察到,除了计算密集型的大型 GEMM 操作外,调度器的固定限制以及大量小尺寸张量的内存拷贝同样会成为严重的性能瓶颈。因此,系统设计的原则转向了全栈式的细粒度优化,旨在减少冗余计算、消除不必要的数据搬运,并最大化显存和计算资源的利用率。
实验环境
- 工作负载规模:测试采用 8K 输入 tokens 和 1K 输出 tokens 的负载配置。
- 并行与投机解码配置:采用张量并行度 8(TP8),并使用 DSpark 投机方法进行 8 个 tokens 的投机解码。
- 并发度设置:分别在并发度为 1、4 和 16 的情况下进行测试。
- 硬件配置:测试运行在配备 CUDA 13.3 的 B300 节点上。
- 软件与版本对比:对比了 vLLM v0.27.1 版本与
82a85dc1提交(0913 主分支)的性能。由于 v0.27.1 版本存在已知的 Kimi K3 前缀缓存问题,两次运行均禁用了前缀缓存机制。 - 服务启动命令:
vllm serve moonshotai/Kimi-K3 \
--trust-remote-code \
--tensor-parallel-size 8 \
--load-format fastsafetensors \
--gpu-memory-utilization 0.9 \
--reasoning-parser kimi_k3 \
--enable-auto-tool-choice \
--tool-call-parser kimi_k3 \
--host 0.0.0.0 \
--port 30000 \
--max-model-len auto \
--no-enable-prefix-caching \
--speculative-config '{"model":"RedHatAI/Kimi-K3-speculator.dspark","method":"dspark","num_speculative_tokens":8,"draft_sample_method":"probabilistic","rejection_sample_method":"standard"}'
- 基准测试运行命令:
guidellm run \
--backend '{"kind":"openai_http","target":"http://127.0.0.1:30000","model":"moonshotai/Kimi-K3","request_format":"/v1/chat/completions","timeout":100000,"extras":{"body":{"reasoning_effort":"max","temperature":1.0,"top_p":0.95}}}' \
--tokenizer '{"kind":"huggingface_auto","model":"moonshotai/Kimi-K3","load_kwargs":{"trust_remote_code":true}}' \
--data 'kind=synthetic_text,prompt_tokens=8000,output_tokens=1000' \
--profile 'kind=concurrent,warmup=0.1,cooldown=0.1' \
--override profile.streams '1,4,16' \
--constraint 'kind=max_duration,seconds=150' \
--constraint 'kind=max_errors,count=10' \
--seed 'kind=static,value=42' \
--metrics 'kind=generative,sample_size=20' \
--output 'kind=json,path=guidellm-results/output_vllm_kimik3_8k1k.json'
实验结果
在 8K/1K 负载下,优化后的主分支相比于 v0.27.1 版本取得了全面的性能飞跃。具体的实验结果和分析结论如下:
| 并发度 | 0.27.1 平均延迟 (s) | 0913 主分支平均延迟 (s) | 0.27.1 吞吐量 (tok/s) | 0913 主分支吞吐量 (tok/s) | 0.27.1 平均 TTFT (ms) | 0913 主分支平均 TTFT (ms) |
|---|---|---|---|---|---|---|
| 1 | 12.37 | 5.30 (−57.2%) | 83.3 | 183.3 (+120.0%) | 2262.9 | 376.3 (−83.4%) |
| 4 | 23.67 | 10.50 (−55.6%) | 166.7 | 416.7 (+150.0%) | 2314.9 | 640.5 (−72.3%) |
| 16 | 55.90 | 22.17 (−60.3%) | 258.3 | 725.0 (+180.6%) | 7601.1 | 1121.0 (−85.3%) |
分析结论表明,随着并发度的增加,优化版本保持了极高的扩展性。在并发度为 16 时,吞吐量实现了 180.6% 的最大增幅,同时首字时间(TTFT)在所有并发度下均实现了超过 72% 的断崖式下降。
方法细节
自适应调度预算 (Adaptive scheduling budget)
自适应调度预算解决低请求量下的预算浪费问题:在请求数量较低时,旧的调度逻辑会导致大部分的 max_num_batched_tokens 预算处于未使用状态。为了解决这一问题,https://github.com/vllm-project/vllm/pull/51725 引入了自适应的调度 token 预算机制。同时,https://github.com/vllm-project/vllm/pull/51726 将高显存 GPU 层的默认限制从 8,192 提升至 16,384。在所报告的 8K/1K 工作负载上,这一改动使得首字时间(TTFT)下降了 55%–65%,吞吐量提升了最高 41.5%。
自适应策略的具体计算逻辑与效果对比:假设配置参数为 max_num_seqs=1024,max_num_batched_tokens=8192,并且投机数量 $K=8$。在旧的逻辑中,无论实际请求数量是多少,调度的 tokens 数量固定为 1024(即 $8192 - 1024 \times (8-1)$)。而在引入自适应策略后,当实际请求数为 1 时,调度的 tokens 数量动态调整为 8185(即 $8192 - 7$);当请求数为 32 时,调度 tokens 为 7968;请求数为 128 时为 7296;只有当请求数达到 1024 时,调度 tokens 才降至 1024。该 PR 通过使用自适应策略,使得在请求数量较少时调度的 tokens 数量大幅增加,从而确保单个请求不会被强制拆分为多个前向调用(forward calls)。
内部 KDA 前缀检查点 (Internal KDA prefix checkpoints)
内部 KDA 检查点避免短后缀引发的全模型前向传播:Mamba 风格的前缀缓存(prefix cache)通常在最后一个可缓存的块边界处对 prefill 阶段进行拆分,这种机制有时会导致为了处理一个非常短的后缀而不得不进行一次完整的模型前向传播。https://github.com/vllm-project/vllm/pull/52789 实现了在单次 prefill 过程内部导出检查点(checkpoint export),这一优化使得 TTFT 下降了 9%–25%。此外,https://github.com/vllm-project/vllm/pull/53614 进一步支持了部分前缀命中(partial prefix hits)以及投机解码,其中包含了 EAGLE 重放边界的对齐(replay-boundary alignment)。
单次 FlashKDA 调用完成全量处理与检查点导出:对于一个 8K 的输入序列,在引入内部 KDA 检查点之前,8K 的 prefill 需要进行两次模型前向传播,并且每个 KDA 层需要调用两次 FlashKDA。在优化之后,只需一次 FlashKDA 调用即可处理全部 8,000 个 tokens,并在同一次循环(recurrence)内部于第 7,680 个 token 处导出检查点状态。
消除冗余计算模块的遍历开销:该 PR 显著的架构收益在于,它帮助系统避免了第二次经过注意力机制(attention)、混合专家(MoE)、路由(routing)以及张量并行集合通信(TP collectives)的全模型前向传播。
零拷贝混合 KDA 批处理 (Zero-copy mixed KDA batches)
零拷贝连续切片消除冗余的张量索引与拷贝操作:在处理混合的投机与非投机批次(batches)时,原有的实现方式在每一层都需要使用六次 index_select 和两次 index_copy_ 操作来进行数据重组。https://github.com/vllm-project/vllm/pull/56159 引入了连续的零拷贝切片(contiguous zero-copy slices)以及直接的输出写入机制。在并发度为 4 和 16,且批次大小(batch size)为 1 的扁平化场景下,该优化使吞吐量提升了 5.2%–7.7%。
直接视图馈送与切片写入重构数据流:在采用零拷贝混合批处理路径之前,每个 KDA 层通过六次 index_select 调用来收集非投机和投机的输入,并通过两次 index_copy 调用将计算结果分散写回。优化之后,系统利用连续的视图(contiguous views)同时馈送给两条 KDA 路径,并将结果直接写入最终输出张量的切片中。
延迟的 MXFP4 最终化 (Deferred MXFP4 finalization)
内核融合减少启动开销与中间张量读写:https://github.com/vllm-project/vllm/pull/53152 将 MXFP4 的 top-k 最终化(finalization)逻辑融合到了潜在尾部内核(latent-tail kernel)内部;这一架构调整成功移除了一次内核启动(kernel launch)以及一次中间张量的写入与读取操作。该优化使得端到端延迟下降了约 5%。随后,https://github.com/vllm-project/vllm/pull/53327 修复了初始化顺序问题,使得在权重加载之前就能顺利启用这种延迟最终化的路径。
融合前后的计算流对比:通过将 MXFP4 的 top-k 最终化操作融合进潜在尾部,系统不仅消除了一次额外的内核启动,还彻底避免了对已经最终化的中间张量进行显存写入和重新读取的操作。
ReplaySSM:重建 KDA 状态以替代状态存储 (ReplaySSM: reconstruct the KDA state instead of storing it)
缓冲最近输入替代高昂的循环状态写入:投机解码机制会在每一个草稿位置(draft position)写入一个 KDA 循环状态,以便在拒绝 tokens 时可以进行状态回滚。对于 $T$ 个草稿位置,这意味着每一步需要进行 $T$ 次额外的状态写入操作。https://dao-lab.ai/blog/2026/replayssm/ 提出了一种替代方案,它改为缓冲最近的 SSM 输入,并在提交(commit)阶段重建被接受的状态。在这种机制下,回滚操作被简化为仅仅移动一个缓冲区的指针。
ReplaySSM 在 Model Runner V2 上的实现与显存收益:https://github.com/vllm-project/vllm/pull/51855 在 Model Runner V2 的架构上为 Kimi K3 实现了 ReplaySSM。系统通过一个 Triton 内核在 align 模式下提交被接受的状态以及下一个前缀缓存边界。在保持相同的 46.48 GiB 缓存预算的前提下,在 TP8 配置中,系统的有效缓存容量提升了 10.97%,并且没有对模型的准确率造成任何损害。
Prefill/Decode 分离与混合状态卸载 (Prefill/decode disaggregation and hybrid state offload)
复杂的分布式状态对分离架构提出新要求:Prefill/Decode(PD)分离与缓存卸载机制必须同时传输 MLA KV 和 KDA 状态。在此架构中,MLA KV 在各个 TP rank 之间是完全复制的;而 KDA 状态则是按注意力头(head)和维度(dimension)进行分片(sharded)的。此外,Mamba 的 align 块表也可以是稀疏且可变的。因此,那些纯粹基于注意力机制的传输假设在这里完全不适用。
边界状态锁定与受限的连接器恢复机制:https://github.com/vllm-project/vllm/pull/51358 引入了 Mooncake 机制,用于存储由调度器选定的边界状态,并将这些状态锁定(pin),直到每一个 rank 上的异步写入操作全部完成。https://github.com/vllm-project/vllm/pull/50344 进一步规定,只有那些具备恢复 KDA 状态能力的连接器(connectors),才被允许服务发散的按组前缀命中(divergent per-group prefix hits)。
解码上下文并行 (Decode context parallelism)
上下文并行突破 TP 的 KV 缓存容量瓶颈:张量并行(TP)会在每一个 rank 上复制 MLA 潜在 KV(latent KV),因此单纯增加 TP 的 rank 数量并不能增加整体的 KV 缓存容量。https://vllm.ai/blog/2026-08-07-decode-context-parallelism 通过沿着序列维度对缓存进行分片来解决这一问题,这对于具有长共享前缀的https://vllm.ai/blog/2026-09-08-vllm-agentx极为有用。
多重通信原语协同实现高效的融合 MLA 路径:https://github.com/vllm-project/vllm/pull/50484 为 Kimi K3 的融合 MLA 路径实现了 DCP。在该实现中,对称内存的 All-to-All (A2A) 负责处理输出和 LogSumExp (LSE) 的归约(reduction);NVLS 组播(multicast)负责收集查询(queries);多内存(multimem)负责收集分块上下文的 KV。查询分片会直接写入消费者的最终缓冲区中。这一设计使得在 4×GB200 环境下,查询交换的延迟大幅下降了 10.4%–29.9%。
容量与延迟的双重显著改善:在一个包含 120k tokens 的工作负载(由 114k 共享前缀,6k 后缀,以及 400 个输出 tokens 组成)上,KV 缓存的有效容量从 1.93M 激增至 19.75M tokens。在并发度为 1 时,TPOT(每个输出 token 的时间)的 p50 延迟从 13.8 ms 下降到了 10.5 ms;在并发度为 2 时,从 16.2 ms 下降到了 11.8 ms。在 GSM8K 基准测试中,使用 DCP8 配置的准确率达到了 96.97%,甚至略高于 TP8 的 96.21%;并且在整个测试过程中实现了零请求错误。
结论
本文详细阐述了 Kimi K3 在 vLLM 框架中实现高达 2.8 倍吞吐量提升的深度优化路径。通过针对 KDA 循环状态、混合专家模型、并行策略以及内存管理的细粒度重构,系统成功克服了调度器限制和张量拷贝带来的性能瓶颈。未来,随着解码上下文并行和 ReplaySSM 等高级特性的进一步成熟,长序列 agentic 工作负载和投机解码的效率有望得到更深入的发掘。
补充细节
社区贡献者拓展了全栈优化工作的边界:除了上述重点讨论的优化示例外,更广泛的优化工作还涵盖了内存布局、序列与流水线并行、KDA prefill 与循环状态、MLA、MoE 以及 GEMM。通过对大型投影进行分片、共享专家、减少集合通信和数据移动,以及收紧小批量的 GPU 执行路径,整体性能得到了全方位的提升。完整的 PR 列表被详细记录在 https://github.com/vllm-project/vllm/issues/50587 中。
特定社区 PR 的详细技术贡献:选定的社区 PR 进一步拓宽了这项工作的范围:https://github.com/robertgshaw2-redhat 和 https://github.com/GirasoleY 添加了https://github.com/vllm-project/vllm/pull/50478 以及前文所描述的 DCP 支持;同时,https://github.com/gau-nernst 开发了https://github.com/vllm-project/vllm/pull/52079。https://github.com/njhill 和 https://github.com/BabyDrangoner 则分别在 https://github.com/vllm-project/vllm/pull/51540 和 https://github.com/vllm-project/vllm/pull/52458 中进一步收紧了 KDA prefill 的执行逻辑。
致谢名单与算力赞助:特别感谢 https://github.com/benchislett、https://github.com/BolinSNLHM、https://github.com/Dao007forever、https://github.com/djmmoss、https://github.com/hnover-nv、https://github.com/huangzhilin-hzl、https://github.com/mingg26、https://github.com/njhill、https://github.com/rebklee、https://github.com/robertgshaw2-redhat、https://github.com/GirasoleY、https://github.com/gau-nernst、https://github.com/tlrmchlsmth、https://github.com/BabyDrangoner 以及 https://github.com/ivanium 等人为 Kimi K3 提交的 PR 贡献。同时感谢审查者、CI 维护者、基准测试所有者以及硬件团队的付出。最后,团队特别感谢 https://novita.ai/ 团队为本文描述的部分优化工作赞助了计算资源。
💬 评论讨论
欢迎在这里分享您的想法和见解!