System Design · Issue 11
Inside Wafer's Stack for Running Kimi K3 on AMD MI355X
Kimi K3 is so big that memory, not FLOPs, picks the GPU. That quietly reopens the AMD-vs-Nvidia serving fight.
The constraint
A model too big for the node
Kimi K3 is Moonshot's open-weight frontier model: roughly 2.78T total parameters with about 103B active per token, 896 routed experts picking a top-16, 93 layers, and a 1M-token max context. It is a mixture-of-experts monster where most of the weight sits idle on any given token but still has to live in VRAM.
The checkpoint lands at about 1.56TB in mixed MXFP4/BF16 storage. That is over 1.5TB of weights on the GPUs before you reserve a single byte of KV cache for that million-token context. For this model, HBM capacity rather than raw compute decides which hardware can serve it at all.
AMD's own enablement post validates exactly this: at TP8 the model needs about 191 GiB of weights per GPU plus KV state, filling roughly 71% of a 288 GiB card, and it only claims the model loads and passes a GSM8K correctness check. Turning that into fast, cheap serving was left as follow-up work, which is where Wafer comes in.
Memory is the moat
Three ways to fit it, one penalty
An MI355X and a B300 both carry 288GB of HBM3E, so a single 8-GPU node holds the weights plus a KV pool with room to spare. A B200 only has 192GB per GPU, so eight of them top out near 1.5TB, which does not leave space for K3's weights and a long-context KV pool together.
MI355X 8×288GB ─┐
B300 8×288GB ─┤ fits in ONE node (TP8)
│
B200 8×192GB ─┘ ≈1.5TB ✗ too small
└─► forced onto 2 nodes (TP16)
node A ══ RoCE ~195Gb/s ══ node B
all-reduce on the decode pathThat capacity gap forces the B200 into a two-node TP16 layout, and every decode step then pays a cross-node all-reduce over RoCE v2 at ~195 Gb/s on the critical path. It is a structural handicap, not a software one: the B200 loses tokens to the network purely because the model does not fit on one box.
The MI355X's edge over the B200 is not faster silicon. It is having enough memory to stay on a single node.
Performance per dollar
B300 wins the race, MI355X wins the receipt
On Wafer's own 1,024-in / 400-out benchmark, a single 8×MI355X node peaked at 952 tok/s, while a single 8×B300 node hit 1,568 tok/s and the two-node B200 setup managed just 498 tok/s across 16 GPUs. The B300 genuinely wins raw aggregate throughput, by about 1.65× over the MI355X.
config tok/s/node $/GPU-hr tok/s/$ 8×MI355X TP8 952 2.50 48 8×B300 TP8+DCP 1568 6.00 33 2×8 B200 TP16 498 4.25 7
The twist is price. At roughly $2.50/GPU-hr for MI355X vs $6.00 for B300, about 2.4× cheaper, the MI355X delivers 48 tok/s per dollar against the B300's 33, roughly 45% better perf/$ and about 7× better than the stranded B200 config. Independent rate checks (Vultr MI355X near $2.59/GPU-hr, B300 best rates around $6.00) land in the same ballpark.
Two caveats before you quote these numbers
The throughput and pricing figures are Wafer's self-reported benchmark; at least one outlet noted it could not independently re-verify the exact 952 tok/s/node, the $2.50-vs-$6.00 split, or the 1.65× B300 edge against third-party benchmark data. Treat the direction as solid and the precise multipliers with care.
Power partly offsets the sticker gap: MI355X draws around 1,400W against a B200's ~1,200W, so in high-power-cost regions or power-limited datacenters, the acquisition-cost advantage narrows.
So the pitch is never that MI355X is the faster card. It is that when memory forces you onto a 288GB GPU anyway, AMD hands you the same capacity for a third of the hourly rate, and the throughput you give up is smaller than the money you save.
Fix one, speculative decoding
A missing definition, not a missing kernel
Kimi K3 ships zero draft tensors, no native MTP and no EAGLE, so the only way to speed up decode is an external draft model: RadixArk's block-diffusion Kimi-K3-DSpark. On CUDA it runs out of the box. On ROCm, the very first real request took down the whole scheduler.
NameError: name 'top_k_renorm_prob' is not defined.
Did you mean: 'top_p_renorm_prob'?sglang's verifier builds the target distribution two ways: a dense path calling top_k_renorm_prob (imported from sgl_kernel on CUDA) and a sparse fast path via torch.topk. The ROCm build only aliases a Triton top-p kernel, so there is no top-k renorm for gfx950 to point at, and any request landing on the dense path kills the scheduler.
The reflex on ROCm is to assume you need a hand-written kernel. Here the fix was top-k renorm in plain PyTorch, sort, masked_fill, divide, dropped into sglang's ROCm sampling branch. That bought roughly 2.2× single-stream performance and +18% peak aggregate throughput, with the peak shifting to a higher concurrency point (c64 rather than c24).
Fix two, prefill
Padding 12 heads to 16 to feel faster
Decode tok/s is fool's gold; what a user actually feels is time-to-first-token. On a 172k-token cold prefill, MI355X took about 51s against the B300's ~23s, a 2.2× gap that never shows up in a decode-throughput table.
The cause was not a missing kernel but a shape mismatch. K3 at TP8 gives each rank 12 attention heads, and the fast AITER MLA prefill kernel only supports 4, 8, or multiples of 16 heads, so the whole thing quietly fell back to slow generic Triton attention.
TP8 → 12 heads per rank AITER MLA wants 4, 8, or ×16 ✗ 12 pad [ 12 real | 4 zero ] = 16 → run ASM kernel then slice the 12 real heads back out Triton fallback 4–7k tok/s AITER (ASM) ~13k tok/s
The workaround is almost silly: zero-pad the head count from 12 to 16, run the fast AITER kernel, then slice the real 12 heads out of the output. That took prefill from the Triton fallback's 4-7k tok/s to ~13k tok/s, a 2-3× TTFT speedup on big cold contexts. It moves only time-to-first-token, so the decode numbers in the main table stay put.
The bottom line
The CUDA moat is a patch away
Why it matters
Kimi K3 is one of the first open models where HBM capacity, rather than raw compute, decides which GPU can serve it. That reopens an AMD-vs-Nvidia question CUDA had mostly closed: the MI355X carries the same 288GB as a B300 for roughly a third of the hourly price, and the ROCm gaps Wafer hit were a missing function definition and a shape mismatch, the kind of thing an engineer patches in an afternoon rather than a silicon limit. The exact multipliers are self-reported, so weigh them accordingly, but the direction is hard to argue with.
The short version
- Kimi K3's ~1.56TB checkpoint needs 288GB-class GPUs; an 8×B200 node (192GB each) can't hold the weights plus a 1M-token KV pool, forcing a two-node topology with a cross-node all-reduce on the decode path.
- On Wafer's own benchmark, MI355X hit **952 tok/s/node at 48 tok/s/$** versus B300's **1,568 tok/s/node at 33 tok/s/$**: B300 wins raw throughput ~1.65×, MI355X wins perf/$ by ~45%.
- Two ROCm fixes carried the run: a PyTorch top-k renorm to stop the spec-decode scheduler crashing (~2.2× single-stream), and zero-padding 12→16 attention heads to load the fast AITER MLA prefill kernel (~2–3× faster TTFT).
- No custom kernels were needed this time; the blockers were a missing definition and a shape mismatch.
- Numbers are Wafer's self-reported figures, so trust the perf/$ gap directionally more than the exact values.