Kimi K3 在 AMD Instinct MI355X 上优化：单 GPU 吞吐超 NVIDIA HGX B300 24.4%

深入 Kimi K3 优化工作：理解瓶颈以及让推理变得更快的关键决策

在本文中

• AgentX：工作负载与结果
• 深入 Kimi K3：专家、注意力与循环状态
• 1. 缓存：保留对话状态以供复用
• 2. MoE：消除跨 GPU 的重复工作
• 3. 内存流量：融合操作并消除不必要的拷贝
• 4. MLA：借助 AMD AITER 汇编实现更快的长上下文解码
• 5. KDA：并行化预填充与投机验证
• 6. 投机解码：选择草拟多少 token
• 下一步优化什么
• 贡献与致谢

在 AgentX 上，一个使用 vLLM 服务 Kimi K3 的八 GPU AMD Instinct MI355X 系统，在 50 tokens/s/user P90 交互性下，相比 NVIDIA HGX B300 每 GPU 每秒总 token 数高出 24.4%。在该基准的拥有成本模型下，每美元可获得的 token 数高出 87.4%。

在这篇文章中，我们将剖析 AMD Instinct MI355X 上 Kimi K3 背后的工程实现，把 vLLM 的服务决策与 GPU 内核及其底层硬件联系起来。我们将解释这些优化为何有效，并建立直觉，以便在你自己的推理栈中发现无用功、识别瓶颈并找到下一个优化点。

已公布的稠密矩阵峰值为 AMD Instinct MI355X 在 MXFP4 下 10.1 PFLOP/s，NVIDIA B300 在 FP4 下 13.5 PFLOP/s。两者均配备 288 GB HBM3E 和 8 TB/s 峰值内存带宽。达到这一服务性能，体现了 AMD 及其 vLLM 社区合作者持续不断的软件工程努力。

AgentX：工作负载与结果

公开的 AgentX 基准使用合成 token 回放真实的编码智能体会话轨迹，保留请求长度、时序和前缀共享模式。在缓存不受限的情况下，该工作负载提供超过 96% 的理论前缀复用率。投机运行使用固定的合成接受参考。

在 50 tokens/s/user P90 交互性下，AMD Instinct MI355X 每秒每 GPU 服务约 7,567 个总 token，而 NVIDIA HGX B300 为 6,083 个，吞吐量优势为 24.4%。按 SemiAnalysis 建模的 TCO 分别为每 GPU 小时 1.50 美元与 2.26 美元，这相当于每美元 TCO 可获得 1,816 万与 969 万个总 token：每美元 token 数高出 87.4%。

公开曲线比较的是完整的服务栈。每项优化旁边的测量数据来自单独的测试，并在其所述配置下解释各自的机制。

深入 Kimi K3：专家、注意力与循环状态

返回的智能体轮次可以复用其大部分前缀，但处理一段简短的续写仍会触及三类数据。

Kimi K3 使用两种注意力。它的 24 个 MLA 层保留随对话增长的压缩 KV 历史。新的查询仍然会查阅该历史。它的 69 个 KDA 层则携带固定维度的状态矩阵，并在新 token 到来时更新它们。模型配置。

在 MoE 层中，路由器选择专家网络来变换每个 token。它们的权重在各请求之间共享。

跨轮次：缓存复用与卸载（第 1 节）。
模型单步内：专家投影、内核交接以及 MLA/KDA 调度（第 2–5 节）。
选择草稿深度：已提交输出与完整周期成本（第 6 节）。

1. 缓存：让对话状态可供复用

DCP 减少 GPU 内存中的重复历史；CPU 卸载将可复用状态保留在其之外。

在 TP8 下，每块 GPU 处理 Kimi K3 的 96 个 MLA 查询头中的 12 个，但存储整个压缩 KV 历史。张量并行划分了头，却没有划分缓存。

使用 八路解码上下文并行（DCP8）时，每块 GPU 存储八分之一的缓存位置。在解码期间，各 GPU 交换查询，针对本地历史进行计算，并合并各自的部分答案。

每个部分答案都携带归一化信息，以便合并时能正确地对各分片加权。#51705 将 ROCm AITER DCP 路径扩展到多 token 验证。

对于 142,000 token 的对话，估计的 MLA 载荷为 24 层 × 142,000 个位置 × 576 FP8 字节 ≈ 每块 GPU 1.96 GB。DCP8 将其减少到约 245 MB，尚未计入分配开销。每块 GPU 针对八分之一的历史处理八倍数量的头，使得主导的注意力算术大致保持不变。

CPU 卸载：保留不再能放入 HBM 的状态

当状态不再能放入 GPU 内存时，AMD recipe 的原生 vLLM CPU 卸载将 MLA 历史和 KDA 检查点保留在主机 DRAM 中。

缓存命中需要匹配的恢复点。如果 MLA 历史到达 token 3,072，但最新的兼容 KDA 检查点在 2,048，则请求必须从 2,048 恢复并重新计算剩余的前缀。

将 attention-block 尾部与 KDA checkpoint 一起保存，可使它们的恢复位置保持一致（#54736）。恢复该状态会以向 GPU 回传一次传输为代价，换取避免前缀重计算。请求恢复后，其新 token 仍可能让多个 GPU 重复相同的计算。

2. MoE：消除跨 GPU 的重复工作

在单个 GPU 上，路由专家的输出投影会返回其消费者所期望的完整向量。在原始 ROCm 路径中，先前的归约与归一化让全部八个 GPU 获得相同的 latent 输入，且每个 GPU 都持有完整权重。因此同一个投影运行了八次。

输出坐标可以在 GPU 之间划分，但每个 rank 仍然需要完整结果。Kimi K3 已经为其共享专家部分和实现了 all-reduce。

每个 GPU 只将它自己拥有的路由输出加到其共享专家部分和中。现有的 all-reduce 会组装出完整结果，无需新增集合通信（#51253）。ROCm 实现基于相应的 CUDA 策略构建。

每个输出坐标恰好接收一次其路由贡献。完整投影权重仍保持复制。

在 PR 于八块 AMD Instinct MI355X GPU 上以 TP8 进行的测试中，在并发数 8–128、随机 8K 输入/1K 输出请求且禁用前缀缓存的条件下，输出吞吐量提升了 3.05–4.91%。

释放 GPU 内存以保留更多缓存状态

现在每个 rank 只计算分配给它的投影坐标，同时保留完整权重。

仅加载每个 rank 使用的行，在 TP8 下跨 92 个 MoE 层可为每块 GPU 释放约 3.85 GiB（#59591）。融合路径的计算保持不变；多出的空间可用于存放 MLA 历史和 KDA checkpoint。

这一改动是在用于头条公开结果的镜像之后合入的。其单独的 AgentX 实验发现，在并发数 32 和 48 时吞吐量提升不到 1%，但在并发数 70 时提升 28.9%，此时报告的前缀缓存命中率从 74.8% 升至 80.1%。

高并发下更大的提升与缓存压力一致。在相同缓存容量下比较两种实现，将有助于分离出这一效应。

更多内容见 专家路由与四位激活。

3. 内存流量：融合操作并消除不必要的拷贝

在每个 GPU 内部，下一个值得关注的地方是内核之间传递的数据。

Kimi K3 的单 token KDA 路径运行了三个计算内核和两次输出拷贝。卷积在一个已经融合了 Q/K 归一化、decay 与 beta 门以及循环更新的内核之前运行，输出归一化则在其之后运行。融合的 HIP 路径在一次启动中完成这些操作，并直接写入目标缓冲区（#50654）。

中间交接留在芯片上。卷积历史和循环状态仍会写入 GPU 内存，以供下一个 token 使用。

AttnRes 将较早层块的输出与运行中的残差混合。将这种混合与 RMSNorm 结合，可使临时混合结果保留在寄存器中（#50593）。为后续层保存的运行残差保留其 BF16 舍入；混合结果在归一化过程中保持 FP32。测试要求持久更新精确，并允许归一化输出存在数值容差。

使用 AttnRes 融合后，在 TP8 配置的八块 AMD Instinct MI355X GPU 上，输出吞吐量提升了 14.2%，测试使用 100K 输入/1K 输出请求，并禁用了前缀缓存。长预填充延迟基本没有变化。

只为使用这些输入的 kernel 拷贝输入

ROCm 矩阵乘法调度器在选择 kernel 之前会打包输入。当它选择回退 kernel 时，该 kernel 读取原始输入，而打包副本未被使用。将打包移入专用分支后消除了这些拷贝（#58566）。

性能剖析让这种浪费变得可见：每个 decode 步骤有 69 次不必要的拷贝，每个 KDA 层一次。每次仅移动 1,792 字节，却耗时约 4.8 微秒。更快地移动这些字节几乎无济于事；启动开销才是主要因素。

4. MLA：使用 AMD AITER 汇编实现更快的长上下文解码

Kimi K3 在 TP8 下每块 GPU 有十二个查询头。填充到十六个使 AMD AITER 的汇编解码路径可用。在报告的长上下文测试中，每个输出 token 的时间从 43.8 毫秒降至 23.6 毫秒。

该比较在并发为 1 的情况下使用不带 DCP 的 TP8。#50578 中的适配器拷贝四个查询头，将它们打包到十六头布局中，并丢弃其输出。DCP8 呈现的是另一种 96 头形状。

这保持了模型的运算不变，因为每个查询头都会在缓存位置上独立地进行注意力计算。增加另一个头不会改变原始头的 softmax。在此解码路径中仅对查询进行填充；存储的压缩 KV 历史大小保持不变。

长对话会让注意力有大量工作，但这些工作需要分摊到整个 GPU 上。在我们检查的历史代码中，一个保守的调度提示将 Gluon 路径限制为每个请求一个工作组。该工作组仍然处理整个对话。填充使 AMD AITER 的汇编路径可用，其不同策略将历史划分为并行片段并合并结果。

PR 报告的每个输出 token 平均时间：八块 AMD Instinct MI355X GPU，TP8，并发数为 1，MXFP4 模型权重，BF16 注意力以及普通的单 token 解码。

短输入场景几乎不变；长输入场景的每个输出 token 时间几乎减半。这些服务测量包括查询准备和输出处理。

5. KDA：并行化预填充与推测验证

MLA 可以将保留的历史拆分到多个工作组中。KDA 提出了更难的调度问题：每个 token 都依赖于前一个状态。#54038 中的融合预填充会在每个 token 组的传入状态已知之前，准备好该组的组合变换。

例如，“将状态减半并加二”，然后“将其减半并加一”，合并为“取起始状态的四分之一并加二”。这条规则可以在起始状态到达之前计算出来。

逐步了解 KDA 依赖序列

KDA 通过矩阵变换实现这一点。对组摘要进行有序扫描可确定每组的起始状态；随后各组并行重放其 token。Token 顺序保持不变，但每组不再等待前一组重放完成。

分组增加了第二遍遍历，并将融合分块路径从两次启动扩展为四次。当序列或头数很少、暴露的并行工作不足时，这种做法有帮助；否则，额外的遍历可能得不偿失。

在每个组内，循环遍历在分块之间将状态保存在矩阵累加器寄存器中，避免全局内存的写入和重新加载。状态存储示意图。

在另一项三后端对比中，融合 KDA 在 4K 和 16K 输入长度下比 Triton 的输入吞吐量高 5.3–5.6%：八块 AMD Instinct MI355X GPU，TP8，并发数 32，禁用前缀缓存且输出为一个 token。这测量的是完整后端。基准测试结果。

为推测验证使用更小的分块

较短的验证批次需要不同的拆分方式。在 TP8 下，每块 GPU 有 12 个 KDA 头和 128 个值列。将每个头拆分为四个 32 列的分块时，仅能在 256 个计算单元上启动 48 个工作组。#58045 改为四列分块，将数量提升到 384 个。

在一项针对单块 AMD Instinct MI355X GPU 的独立图重放内核测试中（TP8 维度，一个序列，四个验证位置），每层时间从 16.20 微秒降至 5.61 微秒（2.89×）。这包括更窄的分块、单波次工作组、统一长度特化以及提升的 token 不变工作。不规则的批次保留动态路径；测试会比较输出和循环状态。

6. 推测解码：选择草拟多少个 token

更快的验证会改变有多少草稿 token 值得检查。有用的比率是已提交输出除以完整的“草稿—验证—提交”循环时间。

例如，AgentX 的固定接受参考值：在三份草稿下，每循环的预期已提交 token 数为 3.00；在四份草稿下为 3.36，其中包含恢复或奖励 token。也就是每循环的已提交输出多出 12%。只有当完整循环耗时比三份草稿时增加不到 12% 时，四份草稿才能提升吞吐量。

公开的 AMD 配方使用三到七份草稿并采用 DCP1，在更大的 DCP8 并发设置下则禁用推测。

接下来优化什么

正在进行的预填充/解码分离工作又提供了一个值得关注的方面：worker 负载均衡，以及传输是否在需要其数据之前完成。

对于另一个模型或工作负载，从已部署的请求形态入手，做三项检查：

比较 rank。 是否有多块 GPU 在某个集合通信可以汇总结果之前，就各自计算出相同结果？

追踪缓冲区。 所选路径上的每个临时张量是否都有消费者？数据能否留在片上？

统计独立工作。 实际的 head、token 和 tile 是否产生足够多的工作组来分摊负载？

以你所需的响应速度测量完整请求，并同时检查输出和保存的状态。更快的 kernel 之所以重要，是因为它缩短了阻碍下一个 token 的工作。

贡献与致谢

Kimi K3 的工作涵盖来自 AMD、Embedded LLM 及更广泛社区的50 多个已合并的 AMD/ROCm 相关 vLLM PR，以及 AITER 的开发、集成、评审和基准测试。贡献涵盖注意力、递归、MoE、量化、图执行、缓存支持和平台集成。

特别感谢：Pin Siang Tan、Jun Kang Chow、Ye Hur Cheong、Hongxia Yang、Seungrok Jung、Bill He、Yichao Zhu、Hyukjoon Lee、Andy Luo、Mrinal Karvir、Richard Li、Sun Peng

Andy Friedrich, Ethan Yang, Tang Rui, Andreas Karatzas, Ankur Singh, Yihua Cheng, Badr Ba Sowid, Benjamin Chislett, Taoyu Zhu, Bill He, Bin Ding, Bowen Bao, Brandon Zhang, Chaemin Lim, Charles Wu, Chuan (Richard) Li, Felix Li, Cameron Quilici, Weishu Deng, Dewei Wang, Divakar Verma, Douglas Lehr, Edwin Lim, Yudong Pan, Elvir Crnčević, Fang Che, Fangzhou Ai, Xingran Fan, Fulu Li, Felix Marty, Canlin Guo, Summer Yang, Haichen Zhang, Yufeng He, He Liu, Hongxia Yang, Hyukjoon Lee, Yifan Qiao, James E T Smith, Jiahui Cao, Yanyuan Qin, Jonathan Chen, Juntian Liu, Shao-Chun Lee, Kedar Potdar, Kevin H. Luu, Kuan-Fu Liu, Kuntai Du, Kyle Sayers, @lalala-sh, @linnea-lin-00638949, Yinfeng Liu, Lucas Wilkinson, Gabriel Wu, Markus Hartikainen, Maithili Joshi, Jeff (Junze) Ma, Matthew Wong, Misha Goin, Micah Williamson, @minmengdie, Matvei Pashkovskii, Mathews Neha, Nicolò Lucchesi, Nick Hill, Nicholas Susanto, Oxana Korzh, Pei Zhang, @ppalanga, Peiran Qin, Qiuyang Yue, Rita Brugarolas, Rebecca Lee, Ryan Rock, Rohan Potdar, @Saddss, Samuel Shen, Sean Filimon, Seungrok Jung, Shantipriya Parida, Shanshan Shen, Shengnan Xu, Simon Danielsson, Siraj Shaikh, Stacy Roberts, Stefan Koncarevic, Asan Stefanski, Minghui Jiang, Tarık Saraç, Tun Jian Tan, Lingpeng Jin, Vansh Bhatia, Maksim Vedenev, Ye Hur Cheong, @wenwzhan, Wei Zhao, Xavier Aguilar, Xiaobing Zhang, Xiaohu Guo, Xin Li, Xudong Lyu, Yanxue Ren, Wentao Ye, Yongqin Wang, Riggins Chen, Chao Hong, Yuchen Shen, Matt Qin, Yichao Zhu, Roger Wang, Zheng Gong, Jiangyun Zhu, Zvi Grinberg, Yongye Zhu.

基准测试说明

对比目标处的服务配置

两套系统均使用八块 GPU、TP8、FP4 权重和 FP8 KV。数值在匹配交互性下、于各自发布的 offload 或仅 GPU 系列内插值；并发数、DCP 和 draft 深度在相邻两组运行之间可能不同。固定的 AMD recipe（YAML） 和 NVIDIA launcher 给出了镜像 tag 和完整设置。

Replay 使用 cc-traces-weka-062126 数据集和固定的 AIPerf harness，排除 warmup，并在固定的合成接受曲线下进行一小时的 profiling。recipe 给出了 seed、session-start 和 warmup 设置。所下载数据集的 revision 未记录。闭环 replay 可能产生不同的已完成请求组合。

软件与 offload。提交使用 vLLM 36768d1bf（AMD）和 5894fdf98（NVIDIA）。配置的 host-cache 预算分别为每 GPU 224.875 和 281 GB；连接器与镜像设置在链接的 recipe 中。

成本。SemiAnalysis 的大型超大规模厂商持有模型和已发布的 GPU 成本常数使用 AMD Instinct MI355X 每 GPU-小时 $1.50、NVIDIA HGX B300 每 GPU-小时 $2.26。总 tokens/$ = 总 tokens/s/GPU × 3,600 ÷ 建模 $/GPU-小时。部分资本和利用率假设仍未公开。在 50-token 交互性目标下，该对比得出每 token 建模成本低 46.6%。

快照。SemiAnalysis InferenceX 数据检索于 2026 年 10 月 9 日；提交日期为 2026 年 9 月 30 日（AMD）和 2026 年 8 月 16 日（NVIDIA）。PR 基准使用其声明的工作负载；#59591 是 headline 镜像之后的独立实验。

来源：AMD Instinct MI355X、NVIDIA HGX B300、NVIDIA 内存规格。基准详情：AgentX 方法学。

技术说明：专家路由与精度

Kimi K3 每个 token 从 896 个专家中选择 16 个，因此即使服务器 batch 很大，每个专家仍可能只留下很少的行。由此产生的专家 tile 可能包含 padding。

A4W4 减少的是另一部分流量：#53940 中的集成将激活值从八位改为四位，而权重保持四位。总流量还包括 BF16 中间输出和专家 tile padding。

技术说明：KDA 精度与片上内存

每个 chunk 仍需解决其 token 之间的依赖关系。该 kernel 将它们合并为一个变换，供后续矩阵乘法应用。这带来两个实际问题：哪种精度能保持答案，以及多少临时数据能放入片上？

一个小例子可以说明这一思路。将先前的修正代入后续修正中，每个答案都可以直接用已知输入表示。所得系数表已经考虑了依赖链：

真正的 prologue 会计算 64×64 块内依赖矩阵的逆。后续的矩阵乘积使用这一变换来考虑这些依赖关系。

逆矩阵的中间值可能会增长并相互抵消。过早对其舍入会丢失最终结果所需的信息。该实现使用 AMD 原生的 FP32 矩阵指令，为敏感的求逆过程保留 FP32 输入和累加。只有完成后的逆矩阵才会舍入为 BF16 以供后续乘积使用。内核实现。

为什么 FP32 累加无法恢复已被舍入的输入

考虑从 256.5 中减去 256。两个输入在 FP32 中都是精确的。在这个数量级下，将 256.5 转换为 BF16 会将其舍入为 256：

FP32 减法无法恢复在运算开始前就已丢失的信息。这些示例数值说明了为什么求逆使用原生 16×16×4 FP32 矩阵指令，同时保留输入和累加精度。

FP4 模型同样依赖更高精度的运算。其性能取决于能否使这些数值路径高效运行。

prologue 还有一个本地内存预算。AMD 将其片上共享暂存区称为 LDS。Q 和 K 块复用存储，因为它们在不同时间被需要；三角输入只存储其十个被占用的 16×16 块。工作组请求 79.5 KiB，向上取整为 80 KiB，是 AMD CDNA 4 计算单元 160 KiB LDS 预算的一半。

通过保留足够多的未使用 LDS 来越过该阈值，同时保持内核二进制不变，即可测试这一阈值。如果运行时几乎没有变化，那么对于该形状而言，额外驻留的工作组贡献甚微。

块间状态

返回 KDA 服务结果。

免责声明

本文档中提供的信息仅供参考，可能包含技术性错误、遗漏和排版错误。本文所含信息可能会发生变化，并可能因多种原因而变得不准确，包括但不限于产品和路线图变更、组件和主板版本变更、新模型和/或产品发布、不同制造商之间的产品差异、软件变更、BIOS 刷新、固件升级等。任何计算机系统都存在无法完全预防或缓解的安全漏洞风险。AMD 无义务更新或以其他方式更正或修订此信息。但是，AMD 保留修订此信息以及不时更改其内容的权利，且 AMD 无义务将此类修订或更改通知任何人。此信息按“原样”提供。AMD 对其内容不作任何陈述或保证，并且对此信息中可能出现的任何不准确、错误或遗漏不承担任何责任。AMD 明确否认任何关于不侵权、适销性或特定用途适用性的默示保证。在任何情况下，AMD 均不对任何人因依赖本文所含任何信息而导致的任何直接、间接、特殊或其他后果性损害承担责任，即使 AMD 已明确被告知此类损害的可能性。AMD、AMD Arrow 标识及其组合是 Advanced Micro Devices, Inc. 的商标。本出版物中使用的其他产品名称仅用于识别目的，可能是其各自公司的商标。© 2026 Advanced Micro Devices, Inc. 保留所有权利。