在MI355X上运行Kimi K3：每美元性能优于B300

过去几个月，开源模型的能力快速提升。DeepSeek V4-Pro 和 GLM5.2 已接近 Opus 级智能，开源逐渐成为闭源之外具有成本效益的选择。而 Kimi K3 的出现更进一步，宣称达到 Fable/Sol 级别的智能，标志着开源模型的新阶段。

不过，更强的能力也意味着更大的规模。GLM5.2 有 7530 亿参数，DeepSeek V4-Pro 有 1.6 万亿，Kimi K3 更是达到 2.8 万亿参数，仅权重就需要超过 1.5TB 显存，还没算上 100 万 token 上下文所需的 KV 缓存。即使是 8 卡 B200 节点也放不下，只能选择 B300 节点或两个 B200 节点组成的 TP16 配置。

此时 AMD MI355X 成为另一个选择：它同样拥有 288GB 显存，价格却比 B300 平均低约 2.4 倍，比 B200 低约 1.7 倍。AMD 的短板一直是软件生态，但 Wafer 认为 AI agent 正在加速缩小这一差距，而且 AMD 为 Kimi K3 提供了 day-0 支持，大部分工作已经完成。

实际测试中，在 1024 token 输入、400 token 输出的基准下，MI355X 达到每节点 952 tok/s 的聚合吞吐和 118 tok/s 的单流解码。相比之下，TP16 B200 配置的聚合吞吐为 498 tok/s（两个节点的总和，约每节点 249），单流仅 90 tok/s。B300 的聚合吞吐为 1,568 tok/s，仍比 MI355X 高约 1.65 倍，但价格是后者的 2.4 倍，因此在性能/美元上 MI355X 明显胜出。按 MI355X 每 GPU 小时 2.50 美元、B300 6.00 美元、B200 4.25 美元计算，每美元吞吐分别为 48、33 和 7 tok/s。

B200 的成绩其实有点吃亏，因为它需要跨节点 all-reduce（RoCE v2，约 195 Gb/s），成为唯一跨两个节点的配置。这恰好说明，Kimi K3 的规模让 MI355X 的高 HBM 容量优势第一次转化为实际可测的领先。

如何做到的？Kimi K3 虽然开箱即用，但要达到这个吞吐还需要一些工程修复。主要杠杆是投机解码。K3 没有自带草稿张量，唯一的外部捷径是 RadixArk 的 Kimi-K3-DSpark 块扩散草稿。CUDA 上可以直接运行，但在 ROCm 上会报错 NameError: name 'top_k_renorm_prob' is not defined。原因是 sglang 的 accept-sampling 验证器有两条路径：密集路径调用 top_k_renorm_prob，稀疏快速路径直接用 torch.topk。CUDA 构建从 sgl_kernel 引入了该函数，而 ROCm 构建只引入了 Triton top-p 内核，没有 gfx950 可用的 top-k renorm 内核，导致变量未定义。

修复只需要一个简单的 PyTorch 函数：对概率向量做 top-k 截断、其余置零并重新归一化，也就是排序、masked_fill 和除法，直接放入 ROCm 采样分支即可。无需自定义内核——这次的问题是缺少定义，而不是缺少内核。修复后单流性能提升约 2.2 倍，中等负载下每流约 1.7 倍，峰值聚合提升 18%，而且峰值出现在更高并发（c64 vs c24）。

预填充优化同样关键。MI355X 在 TTFT 上原本很挣扎：同样的 172k token 冷预填充，MI355X 需要约 51 秒，B300 只要约 23 秒。对 100 万上下文模型来说，大量工作负载会有巨大的预填充，GPU 空转几分钟会让整个节点阵列失效。

差距几乎全部来自一个内核。K3 在 ROCm 上回退到慢速的通用 Triton attention，因为 AITER 的快速 MLA 预填充内核无法加载。原因是形状不匹配：K3 在 TP8 下每个 rank 有 12 个 attention 头，而 AITER 的 MLA 路径只支持 4、8 或 16 的倍数。修复方法很简单：把头数从 12 零填充到 16，运行快速内核，再从输出中取出真正的 12 个头。结果同一预填充负载下，AITER MLA 内核达到约 13k tok/s，而 Triton 回退只有约 4–7k，预填充速度提升约 2–3 倍。这直接影响 TTFT，不影响聚合吞吐，但用户等待首个 token 的时间明显缩短。

总的来说，在 MI355X 上获得最佳性能/美元基本是开箱即得，虽然也遇到一些框架相关 bug，但比 GLM5.2 少，而且完全不需要自定义内核。AMD 上达到 SOTA 指日可待——CUDA 的护城河是否正在消失？