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

过去几个月里，我们见证了开源模型能力的爆发式增长。随着 DeepSeek V4-Pro 和 GLM5.2 达到接近 Opus 的智能水平，开源模型已经成为我们长期依赖的闭源模型之外，一种真正高性价比的替代选择。

但我们此前还没有见过像 Kimi K3 这样的模型。Kimi K3 承诺提供 Fable/Sol 级别的智能，标志着开源进入了一个新时代。

但更聪明的模型意味着更大的模型——这些模型的规模扩张速度与能力提升速度一样快。GLM5.2 拥有 753B 参数，DeepSeek V4-Pro 拥有 1.6T 参数，而 Kimi K3 的参数量高达 2.8T（！！）。在为 100 万 token 上下文分配 KV Cache 之前，其权重就已需要超过 1.5TB VRAM。即使是一台 B200 节点（8 张 GPU）也无法容纳 Kimi K3。因此，你的选择很有限：要么在 B300 节点上部署——每张 GPU 配备 288GB VRAM；要么使用两台 B200 节点（TP16）来服务 Kimi。

不过，猜猜还有哪一种非 NVIDIA GPU 同样拥有 288GB VRAM？AMD 的 MI355X。现在能看出来我们有多喜欢这些芯片了吗？MI355X 单卡平均价格约比 B300 低 ~2.4 倍，也比 B200 低约 1.7 倍；在硬件规格相当的情况下，它是 Blackwell 的高性价比替代方案。AMD 唯一的问题在于软件支持：较慢的内核以及推理框架缺乏 day-0 支持，使得在 AMD 上部署前沿模型成为一项真正的工程工作。Wafer 的观点是，Agent 正在改进内核和模型优化，并持续缩小这一差距。但随着 AMD 为 Kimi K3 提供 day-0 支持，大部分工作实际上已经替我们完成。

结果非常出色：在 1,024-token 输入 / 400-token 输出基准中，MI355X 达到 952 tok/s/node 和 118 tok/s 单流速度——其单节点总吞吐量超过我们的 TP16 B200 部署的 3.8 倍，单流解码速度超过 1.3 倍。后者的 498 tok/s 是由 16 张 GPU、2 个节点合计实现的，即约 249 tok/s/node。B300 节点的总吞吐量仍比 MI355X 高约 1.65 倍，但价格高出 2.4 倍，因此 MI355X 在每美元性能上大幅领先 B300。

每美元性能按 MI355X 为 $2.50/GPU-hr、B300 为 $6.00/GPU-hr、B200 为 $4.25/GPU-hr 计算。

为 B200 辩护的是，它的结果在一定程度上被压低了：在解码关键路径上，它需要承担跨节点 all-reduce 的成本（RoCE v2，约 195 Gb/s）。它是这里唯一跨越两个节点的配置，因为 Kimi K3 的权重加上 100 万 token 的 KV Cache 池，无法放入单台配备 8×192GB GPU 的节点中。但这恰恰说明了问题：以 Kimi K3 的体量来看，它是我们所见首批让 MI355X 依靠 HBM 容量优势、相较 B200 获得实际且可量化优势的模型之一。

我们是如何做到的

虽然 Kimi K3 可以开箱即用地部署，但要达到当前的吞吐量数字，仍有一些工作需要完成。

最主要的优化杠杆是推测解码。K3 没有提供任何 draft tensor——没有 MTP，也没有 EAGLE——因此唯一的推测路径是使用 RadixArk 的外部块扩散 draft 模型：Kimi-K3-DSpark。在 CUDA 上，它可以直接运行；但在 ROCm 上，我们的第一个真实请求就因以下错误导致调度器崩溃：

NameError: name 'top_k_renorm_prob' is not defined. Did you mean: 'top_p_renorm_prob'?

sglang 的 accept-sampling 验证器有两种构建目标分布的方式：一条是调用 top_k_renorm_prob 的稠密路径，另一条是直接通过 torch.topk 的稀疏快速路径。CUDA 构建会从 sgl_kernel 导入 top_k_renorm_prob；而 ROCm 构建只为 Triton top-p 内核创建别名，却没有定义 top_k_renorm_prob——因为不存在可供别名到 gfx950 的 top-k renorm 内核。因此，请求一旦进入稠密路径，验证器就会触发这一 NameError，并连带使调度器崩溃。

修复方法只需要一个 PyTorch 函数。Top-k renorm 是一个小操作：获取模型的概率向量，保留其中概率最高的 k 项，其余置零，然后将剩余概率重新缩放为总和 1。只需一个 sort、一个 masked_fill 和一次除法，就可以直接加入 sglang 的 ROCm sampling 分支，执行与 CUDA 构建通过 sgl_kernel 获得的相同计算。无需自定义内核：在 ROCm 上，人们的条件反射是认为必须编写一个内核，但这里缺少的是定义，而不是内核。

在修复并加固推测解码后，我们将单流性能提升约 2.2 倍，中等负载下的单流性能提升约 1.7 倍，峰值总吞吐量提升 18%。更重要的是，我们的峰值总吞吐量出现在高得多的并发度下（c64，而非未使用推测解码时的 c24）。

Prefill 优化

围绕模型性能的讨论往往突出每秒解码 token 数。但在许多情况下，解码 tok/s 是一种虚假的黄金指标——解码被过度神化，而用户最能直接感受到的首 token 时间（time-to-first-token，TTFT）却被忽略了。

MI355X 在这里表现欠佳：相同的 172k-token 冷 prefill，在 MI355X 上需要约 51 秒，而在 B300 上仅需约 23 秒。对于支持 100 万上下文的模型而言，许多工作负载都有非常大的 prefill，有时还是冷 prefill；让 GPU 为 prefill 忙碌数分钟，可能会使整个节点集群失去可用性。

差距几乎完全来自一个内核。ROCm 上的 K3 回退到了较慢的通用 Triton attention，因为快速的 AITER MLA prefill 内核无法加载。问题并不在于缺少内核，而是形状不匹配：TP8 下的 K3 每个 rank 有 12 个 attention head，而 AITER 的 MLA 路径是为 4、8 或 16 的倍数设计的。修复方法极其简单：将 head 数量从 12→16 以零填充，运行快速内核，然后从输出中提取真实的 12 个 head。

结果是：在同样的 172k 冷 prefill 上，AITER MLA prefill ASM 的稳定运行速度约为 13k tok/s，而 Triton 回退路径仅为约 4—7k tok/s，令 prefill 加速约 2—3 倍。这是改善 TTFT 的杠杆，而非提升总吞吐量的手段——解码性能没有变化，因此不会改变上文的数字；它改变的是用户在第一个 token 出现前需要等待的时间。

结论

在 MI355X 上实现最佳每美元性能，整体上相对可以开箱即用。期间确实遇到了一些预期内的框架相关 bug——但比 GLM5.2 更少，而且这次肯定不需要自定义内核。

AMD 上的 SOTA 已近在眼前。CUDA 的护城河是否已经消失？

Over the past several months, we’ve seen an explosion in the capabilities of open source models. With DeepSeek V4-Pro and GLM5.2 reaching near-Opus levels of intelligence, open source has emerged as a real, cost-efficient alternative to the closed source models we’ve been married to.

But we have yet to see one like Kimi K3. Promising Fable/Sol levels of intelligence, Kimi K3 marks the start of a new era for open source.

But a smarter model means a bigger model — and these models are expanding in size just as fast as they are in capabilities. GLM5.2 has 753B parameters, DeepSeek V4-Pro 1.6T, and Kimi K3 weighs in at 2.8T (!!) parameters. That’s over 1.5TB of VRAM before allocating a KV cache for 1M tokens of context. Not even a B200 node (8 GPUs) can fit Kimi K3. That leaves you with limited options: serve on a node of B300s, which have 288GB of VRAM per GPU, or commit two B200 nodes (TP16) to serving Kimi.

But guess which other non-NVIDIA GPU has 288GB of VRAM? AMD’s MI355X. Can you tell we like these chips yet? At around ~2.4× cheaper per GPU on average versus a B300 and ~1.7× cheaper than a B200, the MI355X is a cost-efficient alternative to Blackwells with comparable hardware specs. The only problem with AMD is software support — slower kernels and less day-0 support on inference frameworks make serving frontier models on AMD a real engineering effort. Our claim at Wafer is that agents are improving at kernel and model optimization, closing this gap as we speak. But with AMD shipping day-0 support for Kimi K3, most of the work was already done for us.

The results are great: on a 1,024-token input / 400-token output benchmark, the MI355X reaches 952 tok/s/node and 118 tok/s single stream — over 3.8× the aggregate throughput per node and over 1.3× the single-stream decode of our TP16 B200 deployment (whose 498 tok/s is a 16-GPU, 2-node total — ~249/node). B300 nodes still win ~1.65× on aggregate throughput over the MI355X, but at 2.4× the price, the MI355X crushes the B300 on performance per dollar.

| 8× MI355X (TP8) | 2×8 B200 (TP16) | B300 (TP8+DCP8)
---|---|---|---
Decode tok/s per stream | 118 tok/s | 90 tok/s | 172 tok/s
Peak aggregate | 952 tok/s | 498 tok/s | 1,568 tok/s
Peak aggregate per GPU | 119 tok/s | 31 tok/s | 196 tok/s
Peak aggregate per $/GPU-hr | 48 tok/s/$ | 7 tok/s/$ | 33 tok/s/$

Perf/dollar at$2.50/GPU-hr for the MI355X, $6.00 for the B300, and $4.25 for the B200.

To the B200’s defence, its numbers are somewhat deflated by the fact that it pays a cross-node all-reduce on the decode critical path (RoCE v2 at ~195 Gb/s) — it’s the only config here that spans two nodes, because Kimi K3 won’t fit weights plus a 1M-token KV pool on a single 8×192GB node. But that’s exactly the point: Kimi K3 at its size is one of the first models we’ve seen where the MI355X’s focus on HBM capacity gives it a practical, measurable edge over the B200.

How we did it

While Kimi K3 served out of the box, there was still work to be done to get it to its current throughput number.

The main lever was speculative decode. K3 ships zero draft tensors — no MTP, no EAGLE — so the only speculative path is an external block-diffusion draft: RadixArk’s Kimi-K3-DSpark. On CUDA it just runs. On ROCm our first real request breaks the scheduler with this error:

NameError: name 'top_k_renorm_prob' is not defined. Did you mean: 'top_p_renorm_prob'?

sglang’s accept-sampling verifier has two ways to build the target distribution: a dense path that calls top_k_renorm_prob, and a sparse fast path that routes through torch.topk directly. The CUDA build imports top_k_renorm_prob from sgl_kernel; the ROCm build aliases only a Triton top-p kernel and leaves top_k_renorm_prob undefined — there’s no top-k renorm kernel for gfx950 to alias. So the moment a request lands on the dense path, the verifier hits that NameError and takes the scheduler down with it.

The fix is a single PyTorch function. Top-k renorm is a small operation: take the model’s probability vector, keep the k highest entries, zero the rest, and rescale what’s left to sum to 1. A sort, a masked_fill, a divide — dropped straight into sglang’s ROCm sampling branch, the same computation the CUDA build gets from sgl_kernel. No custom kernel: the reflex on ROCm is to assume you need one, but here it was a missing definition, not a missing kernel.

With spec dec fixed and hardened, we gained ~2.2× performance single-stream, ~1.7× per-stream at moderate load, and +18% peak aggregate. More importantly, our peak aggregate throughput landed on much higher concurrency (c64 vs c24 no-spec).

Prefill optimizations

Discussion around model performance tends to highlight decode tokens per second. But in many cases decode tok/s is fool’s gold — decode is over-glorified, while time-to-first-token, the number users feel the most, gets overlooked.

The MI355X struggles here: an identical 172k-token cold prefill took ~51s on MI355X versus ~23s on a B300. On a 1M-context model, a lot of workloads have huge prefills (sometimes cold), and having GPUs spin on prefill for minutes can render entire fleets of nodes useless.

The gap was almost entirely one kernel. K3 on ROCm was falling back to slow generic Triton attention because the fast AITER MLA prefill kernel wouldn’t load. The problem was a shape mismatch, not a missing kernel — K3 at TP8 gives 12 attention heads per rank, and AITER’s MLA path is built for 4, 8, or multiples of 16. The fix was trivially simple: zero-pad the head count 12→16, run the fast kernel, and extract the real 12 heads from the output.

The result: on the same 172k cold prefill, the AITER MLA prefill ASM runs at ~13k tok/s steady-state vs the Triton fallback’s ~4–7k, speeding up prefill by ~2–3×. It’s a TTFT lever, not an aggregate-throughput one — decode is unchanged, so it doesn’t move the numbers above; it moves the number a user waits on before the first token appears.

Takeaways

Achieving the best performance-per-dollar ratio on the MI355X was relatively out of the box. There were some expected framework-related bugs — but fewer than GLM5.2, and this time it certainly did not require custom kernels.

SOTA on AMD is imminent. Is the CUDA moat dead?