跳到正文
vLLM Blog· Benjamin Chislett (NVIDIA)·· 2 天前AI 评分58

vLLM 如何用智能体辅助优化 Kimi K3 解码性能

Hyper Optimizing Kimi K3 in vLLM: Moving Fast and Responsibly

AI 导读

vLLM 团队用智能体辅助开发为 Kimi K3 构建统一的 CuTe DSL 解码路径,在 16 块 NVIDIA GB300 上以 DSpark 达到 842 tokens/s,比原生 vLLM 快 1.73 倍。

正文 · AI 翻译
Kimi K3 decode performance on NVIDIA GB300.
Kimi K3 在 NVIDIA GB300 上的解码性能。

编码智能体如今可以帮助完成模型优化中一些最困难的部分:编写复杂的 GPU 内核、执行跨算子融合、让计算与通信重叠,以及针对不同的批大小和 GPU 配置调优执行。对于 Kimi K3,我们采用智能体辅助开发,在 NVIDIA GB300 上构建了一条统一的 CuTe DSL 解码路径,覆盖注意力、MoE 以及跨 8 或 16 个 GPU 的通信。这些优化在 16 块 NVIDIA GB300 GPU 上、批大小为 1 且使用 DSpark(合成 AL=6)时,实现了超过 840 tokens/s 的吞吐;在多个测试配置下,解码吞吐提升 1.5–2×;在 AgentX 工作负载上,P90 交互性最高提升 1.8×。

智能体优化的推理代码正在超出我们审查与维护流程的承载能力。人工审查流程需要理解每个内核的每一行代码、梳理潜在问题,并测试所有受支持的配置,这是一项非常耗时的工作。另一方面,如果我们省略这项工作,就可能发布在生产环境中出故障的有缺陷软件。

如何同时实现这两个目标已有清晰的路径:我们需要改变组织性能工作的方式。除了智能体优化之外,我们还需要设计能够自动验证可靠性和性能的系统,并需要以在设计中强制隔离和最小化风险的方式来组织软件。

拟议的快速通道模型优化 RFC为这一方法提供了基础。我们的指导原则是在管理风险的同时减轻审查负担。为此,我们:

  1. 隔离变更。将优化限制在模型代码内。
  2. 限定范围,保留核心功能。定义受支持的硬件和工作负载,但保留生产环境 LLM 使用所必需的功能。
  3. 系统化验证正确性。运行一套既定的性能和正确性基准测试,以确保完整性和可靠性,并在所有受支持的配置上严格实现更优性能。
  4. 扩展验证。让贡献者在自己的 GPU 上运行标准评估,并在请求维护者关注之前分享可复现的结果。

Kimi K3 是我们遵循这一方法的首个贡献。本文介绍我们如何优化 Kimi K3、如何建立发布和维护这些优化所需的信心,以及性能结果。

目录

第 1 部分:优化 Kimi K3

本贡献为 Kimi K3 引入了一条新的解码路径,它用定制的 CuTe DSL 内核融合并替换了几乎所有操作,并针对特定子问题进行了调优。它面向配备 8 或 16 块 GPU 的 NVIDIA Blackwell GPU,解码批大小从 1 到 64,并可选支持长上下文的解码上下文并行。这些 GPU 共享一个支持多播的 NVLink 域。更多细节请参阅 PR。

Kimi K3 kernel timeline: stock vLLM uses 130 microseconds and 23 kernels per layer; specialized decode uses 69 microseconds and 13 kernels.
Kimi K3 内核时间线:原生 vLLM 每层使用 130 微秒和 23 个内核;专用解码使用 69 微秒和 13 个内核。

1. 针对工作负载调优形状

不同的批大小需要不同的内核选择。对于注意力输入投影,我们在批大小为 1–2 时使用专门的小矩阵内核,在更大的批大小时使用 split-K 内核,将矩阵乘法分散到更多 GPU worker 上。所有此类投影都经过精细调优,以启动大小合适的内核,充分利用 GPU 资源,并尽可能重叠内存读取/通信与计算工作。

KDA 还使用了优化的 CuTe DSL 解码内核,其调度针对每个 GPU 上的批大小和注意力头数量进行了调优。

对于长上下文,解码上下文并行(DCP)将 KV 缓存拆分到多个 GPU 上。每个 GPU 在其本地部分上计算注意力,一个经过调优的融合内核将部分结果与 v-proj 和门控相结合。

2. 重叠计算与通信

我们将注意力输出投影直接与其 all-reduce 以及随后的注意力残差计算融合。GPU 内核将部分结果发布到 NVLink 可访问的缓冲区中,接收内核将它们合并并继续计算。某些小批量配置将生产者与消费者融合为一个内核。数据交换采用无 fence 的方式,开销极小。

在更大的批大小时,实现切换为两阶段归约:GPU 首先合并输出的指定部分,然后分发结果。这减少了 GPU 之间的数据移动。

3. 改进 MoE 调度与重叠

尽管专家已经使用了张量并行,基线仍在每个 GPU 上计算完整的路由分数和潜在输入投影。我们将这些计算拆分到各 TP rank 上,并将每个 rank 的份额合并为一次矩阵乘法。GPU 通过 NVLink 交换结果。这减少了专家执行开始前的重复计算、权重读取和内核启动。

共享专家在单独的流上运行。其工作与路由或路由专家计算重叠,具体取决于配置。调度控制每个操作使用多少 GPU 资源,并经过仔细调优,以确保并行度良好均衡且无争用。详情见 experts_runtime。

4. 将即将使用的权重预取到缓存

在小批量推理中,我们通常受内存限制:受限于从内存加载权重的速度。某些阶段无法充分利用可用的内存带宽,例如通信:在此工作发生时,我们可以在侧流上运行一个专用内核,以确保一旦通信完成,下一个操作的权重将具有高得多的缓存命中率。

第 2 部分:快速且负责任地前进

Kimi K3 只是首批示例之一。我们相信社区可以共同构建这一点,让更多模型、GPU 和工作负载接近光速。

不同的贡献者/代理会发现不同的机会:更好的通信调度,或有助于特定部署的融合。但如果每个贡献都带着不同的假设、不完整的测试和难以复现的基准而来,验证和维护工作很快就会成为瓶颈,或者我们可能发布有问题的代码。我们需要一种通用的方式,将这些想法转化为用户可以在其长期生产用例中依赖的改进。

这改变了我们对智能体交付内容的期望。我们应当对智能体所产出的代码在测试、验证和需求方面设定明确的预期。智能体将构建测试、运行全面的评估与基准测试、调查失败原因,并记录优化生效的条件。人类审查者仍然评判设计与证据,但他们可以从高层抽象入手,而不必逐行审查代码。

基于所提出的快速通道 RFC,我们围绕四项原则组织这项工作。

1. 隔离变更

将模型特定的优化保留在模型自身的代码内。这使得受影响的路径更容易识别,并缩小审查者需要考虑的影响范围。

隔离在运行时同样需要成立。审查者应当能够追踪优化路径何时被选中、它改变了什么状态,以及执行如何返回标准路径。

这一约束确实排除了一些优化:输入准备、采样和 KV 缓存管理可能带来隔离所排除的开销。这是设计的后果,且在很大程度上是有意为之:这些组件在大多数 vLLM 部署中共享,为优化特定工作负载而改动它们,有破坏数十个其他稳定 vLLM 配置的风险。针对这些组件的工作很重要,应当让核心维护者比快速通道所允许的更密切地参与。

2. 限定范围,但保留用户功能

一项优化不必改善每一个部署才有用。它可以针对特定 GPU、并行设置或批大小范围。清晰的边界使工作更易于调优,并为我们提供一组明确的配置来验证。

对于 K3,初始硬件验证覆盖 GB300 的 TP8 和 TP16、最大 64 的解码批、不同草稿长度的投机解码,以及单个 NVLink 域内的解码上下文并行配置。

这些边界描述的是优化适用的范围,而不是削减“功能覆盖”。生产部署需要使端点有用的基础功能。全上下文服务、温度采样和结构化输出等功能不能被“优化掉”或“选择退出”,它们是一等的服务功能,与内核形状和数值检查一样属于验证计划。

3. 验证正确性

内核测试和模型评估回答的是不同的问题。内核测试检查跨形状和输入值的数值结果。模型评估检查完整执行路径是否保持模型的答案。两者都需要。

测试必须真正执行到被改动的代码。小批量优化需要在小批量下评估。注意力内核改动需要长上下文覆盖。静默使用标准路径的测试对优化提供的证据很少,这是一个常见的陷阱。

K3 的开发说明了这一点为何重要。我们的审计发现编译器将激活读取移到了同步等待之前。修复方案将该加载类型限制为静态权重。此类错误可能通过简单的功能检查而存活,需要有针对性的调查。

此外,MRCR 和 AIME25 等长上下文和长生成评估揭示了注意力和通信内核中敏感的数值问题,这些问题在使用模拟输入和/或一致启动时序的独立内核测试中无法复现。

4. 广泛检查性能

标题级的加速只是一个起点。我们还需要了解这一改动在其所影响的各种工作负载下的表现。这意味着要遍历不同的批大小、上下文长度以及所支持的并行配置。MoE 测量应包含真实提示词,因为专家选择会影响读取哪些权重以及可复用多少工作。纯解码结果也应与混合了预填充和解码的工作负载区分开来。

可靠的验证同样需要消耗 GPU 时间。相关的 Bring-Your-Own-GPU 评估 RFC 提出了标准评估镜像,贡献者可以在自己的硬件上运行。可复现的结果包以及与 GitHub PR 关联的提交将使证据更易获取。不同的评估层级和共享的社区算力将有助于扩大覆盖范围,同时管理 GPU 需求并减轻项目资助 CI 的压力。

第 3 部分:性能深入剖析

使用 DSpark 的解码性能

借助 DSpark,Kimi K3 在受控的 AL=6 基准测试中,于 16 块 NVIDIA GB300 GPU 上达到每秒 842 个 token。这比采用相同投机解码设置的原始 vLLM 快 1.73 倍。完整结果表请参阅该 PR。

Kimi K3 at batch 1 with TP16, DSpark DL7, and synthetic AL=6: 842 decode tokens/s versus 488 for stock vLLM, a 1.73x speedup.
Kimi K3 在批大小为 1、TP16、DSpark DL7、合成 AL=6 条件下:解码吞吐为每秒 842 个 token,而原始 vLLM 为 488,加速 1.73 倍。

设置:TP16,DSpark 使用七个草稿 token,合成接受长度为 6,DL=7,ISL=1,OSL=1k,批大小为 1。TPS 衡量聚合解码吞吐量,不包含预填充。

无投机解码的解码性能

这些改进覆盖了所测试的批大小范围。TP16 在批大小 16–64 时保持 1.62–1.71 倍的增益,而 TP8 在相同范围内提供 1.31–1.43 倍的提升。

批大小TP8 token/秒TP8 加速比TP16 token/秒TP16 加速比
1193.31.58×256.41.88×
2317.81.50×436.71.92×
4540.61.55×702.51.81×
8820.11.45×1,157.91.79×
161,274.81.43×1,758.31.71×
321,702.61.33×2,550.91.64×
642,375.91.31×3,606.11.62×

无投机解码。吞吐量针对整个 TP 组。

AgentX 上的改进

我们优化的 Kimi K3 实现还在 AgentX 上以并发度 1 实现了高 1.8 倍的 P90 交互性:403.2 对 223.7 token/秒/用户。这是将 16 块 GB300 GPU 上的优化运行与早前 8 块 GB300 GPU 上的 vLLM 结果进行比较。

AgentX 回放真实编码代理会话的流量模式,包括长上下文、多轮对话、工具暂停、并行子代理以及 KV 缓存复用。它在保留这些工作负载模式的同时,用合成文本替换私有内容,从而在真实的代理流量下测试服务系统。AgentX 方法论。

Kimi K3 AgentX results: P90 interactivity versus total throughput per GPU, comparing the earlier GB300 vLLM curve with specialized decode.
Kimi K3 AgentX 结果:P90 交互性与每 GPU 总吞吐量的关系,比较早前的 GB300 vLLM 曲线与专用解码。

下一步

对于 DeepSeek-V4.1-Flash,#61868 遵循 K3 的设计模式:在相同的可选启用 decode_backend: latency_specialized 开关背后,该实现还支持最多 64 个并发请求、预填充和混合批的无缝回退,以及投机解码。该执行器不是追逐单一热点,而是削减内核和主机工作:它在每个子层边界融合 all-reduce + mHC,用为低延迟构建的 CuTe DSL 内核替换大批量 MoE 和 GEMM,对索引器的工作进行去重和分片,并将原始 vLLM 在 CPU 上急切执行的每步工作改为以 CUDA 图重放。

DeepSeek AgentX results on four GB300 GPUs: P90 interactivity versus total throughput per GPU with and without optimized decode.
DeepSeek 在四块 GB300 GPU 上的 AgentX 结果:P90 交互性与每 GPU 总吞吐量的关系,包含和不包含优化解码。

在 Meta 的 Muse Spark Open 发布后两天内,我们的早期结果已经超过每用户每秒 950 个 token,使用七个草稿 token(K=7),接受长度为六(AL=6)。更多优化正在进行中。

我们还将发布一个类似的机制,用于在低并发场景下优化 GLM5.3 的解码延迟。我们将在后续文章中分享其优化细节和性能结果。

对于每一项贡献,我们将继续关注具体的性能提升、生产功能支持,以及可复现的正确性和性能检查。

致谢

我们感谢 Rohan Varma(NVIDIA)和 Alex Steiner(NVIDIA)对这项工作的贡献。

我们还要感谢 Woosuk Kwon(Inferact),以及 Michael Goin 和 Lucas Wilkinson(Red Hat),感谢他们的模型重构以及宝贵的反馈和支持。

来源:vLLM Blog · vllm.ai