Running Kimi K3 on AMD MI355X: The $/Token Win Comes From Two Kernel Problems

A surprise bottleneck: it wasn’t the model, it was the math around it

The first time you try to serve a frontier-scale model, it feels like there are only two knobs: GPU count and model size. Push the numbers high enough and—surely—you get speed.

Then you hit the real world: a model can be too large to fit, or it can fit on paper but stall on a single slow step. That’s where serving Kimi K3 on AMD Instinct MI355X becomes a great technical story. Kimi K3 is a 2.8T-parameter open-weight model with a 1-million-token context window, and it’s built to run as a large mixture-of-experts system (not a small model wearing a big hat). (huggingface.co)

Meanwhile, each MI355X GPU has 288GB of HBM3E memory, which is exactly the kind of memory capacity you end up needing for “1M context” serving scenarios. (amd.com)

So on paper, MI355X should be a fit. In practice, getting better performance per dollar means answering a sharper question: where, in the token-generation pipeline, do you lose time—and how do you avoid those losses?

Why “$/GPU-hour” isn’t enough

People love to compare GPUs by headline price. It’s tempting. But token generation is dominated by specific computation phases, and each phase has different hardware pressure.

The simplest mental model looks like this:

• Prefill: the model reads your prompt (say 100k–1M tokens of context) and builds the internal state it needs to start generating.
• Decode: the model produces output tokens one at a time (or in small batches), repeatedly.
• TTFT (time-to-first-token): how long until the first output token shows up. Users feel this more than raw decode throughput.

When a system is slow, it’s often because one of these phases falls back to a less optimized kernel path. In that case, improving decode speed alone won’t fix the user experience. The story here ends up being two separate kernel issues: one hurts decode via speculative decoding, and another hurts prefill via attention kernel selection.

The core deployment problem: fitting Kimi K3

A 2.8T model implies a lot of weights. But “weights fit” is not the only constraint. For long contexts, you also need memory for the KV cache.

Quick definition: a KV cache stores intermediate attention data (keys and values) so the model doesn’t redo work for every new generated token across the long history.

The Wafer write-up describes why Kimi K3 is hard to fit on smaller multi-GPU groupings with limited HBM—leading to configurations where MI355X’s HBM capacity becomes the practical edge. The punchline is that MI355X can run Kimi K3 in a way that gives measurable throughput and decode interactivity, not just a “model loads successfully” demo.

Speculative decoding: faster decode by doing extra work carefully

Now to the part that feels almost paradoxical: speculative decoding.

Definition: speculative decoding attempts to generate multiple future tokens speculatively using a draft process, then verifies them against the full model. If guesses are good, the system “skips” expensive full-model steps. If they’re wrong, it falls back—but you still learn something from the trial.

Kimi K3 ships with no internal draft tensors (so it can’t generate its own drafts the way some architectures do). That means the speculative path has to come from an external draft approach—described in the source content as RadixArk’s Kimi-K3-DSpark.

What breaks on ROCm is a classic “it runs on CUDA” trap: a missing function symbol in the verifier’s sampling logic. In the sglang verifier, the dense path needs a top‑k renormalization probability helper. On ROCm, that helper wasn’t defined the same way, so the first request that hit that dense speculative path crashed scheduling.

The fix wasn’t a heroic new kernel. It was smaller: drop in the correct computation in the ROCm branch using existing PyTorch ops (sort, masked fill, renormalize), matching the mathematical intent of the CUDA side.

That matters for beginners because it changes the debugging mindset.

Instead of asking, “Do we need a custom ROCm kernel?”, you learn to ask, “Is this a correctness/missing-definition issue that prevents the optimized path from running at all?” When speculative decoding is the lever, losing it is like throwing away the gear that makes the whole car faster.

With speculative decoding working and hardened, the reported improvements were sizable: about 2.2× single-stream performance and an aggregate peak shift (and, importantly, higher concurrency before the system saturates). Those results follow naturally: speculative decode turns decode from a strictly serialized process into a verified shortcut mechanism.

Prefill optimizations: TTFT is where fleets die

Even with decode fixed, serving long-context workloads reveals a second truth: decode throughput can be misleading.

Users care about TTFT. If a prompt takes minutes of GPU time before the first token appears, an otherwise fast decoder is irrelevant—your queue backs up, concurrency collapses, and the fleet becomes unusable.

In the Wafer results, the cold prefill for an identical 172k-token prompt took about 51 seconds on MI355X versus 23 seconds on B300. That kind of gap is the definition of a silent killer: it doesn’t show up as a decode tok/s difference.

The culprit is described as attention kernel selection. On ROCm, K3 fell back to a slower generic Triton attention because the faster AITER MLA (multi-head attention acceleration) kernel wouldn’t load.

Here’s the key detail: it wasn’t because the kernel didn’t exist. It was because of a shape mismatch.

Definition: attention “shapes” here mean how many attention heads get partitioned per rank in tensor parallel execution. AITER’s MLA fast path expects head counts in particular multiples (built around specific hardware-friendly tiling patterns). For TP8, K3 gives 12 attention heads per rank, but the MLA path was built for 4, 8, or multiples of 16.

The workaround is almost embarrassingly practical: zero-pad from 12 to 16, run the fast kernel, then ignore the padded heads in the output and extract the real 12.

In other words, instead of fighting the kernel, you make your tensor look like it. The reported effect is dramatic on steady-state prefill: AITER MLA runs at roughly ~13k tok/s versus ~4–7k tok/s for the Triton fallback, translating to ~2–3× faster prefill. Decode speed stays basically unchanged—because the prefill fix targets TTFT and fleet utilization, not decode throughput.

If you’ve ever watched a UI spinner for the first response, that’s exactly why this matters.

Performance per dollar: the bottleneck shift is the whole strategy

Once speculative decode and prefill TTFT are fixed, you can compare “tokens per second” more meaningfully.

The source content reports that on a 1,024-token input / 400-token output benchmark, MI355X reaches 952 tok/s per node and 118 tok/s single stream. It also claims B300 aggregate throughput wins on raw aggregate, but MI355X wins on cost-adjusted throughput.

To ground the pricing side, GPUs.io tracks cloud pricing normalized per GPU-hour; at the time referenced in the source write-up, example figures were around $2.50/GPU-hour for MI355X, $6.00 for B300, and $4.25 for B200. (gpus.io)

That combination creates the practical metric you actually care about in production: how many verified output tokens you can generate per unit cost while still meeting interactivity constraints.

And that’s the deeper reason this story lands. The win isn’t “AMD is faster” or “NVIDIA is worse.” It’s that MI355X’s HBM capacity makes the long-context deployment feasible, and then the performance work focuses on two choke points that determine user experience and fleet throughput.

Kernel work beats brand work

There’s a takeaway tucked into every engineering detail here: if you’re serious about frontier model serving, you eventually end up doing kernel-level or kernel-path-level engineering.

Sometimes it’s about throughput kernels. Sometimes it’s about decoder algorithms. Sometimes it’s about avoiding a crash caused by a missing ROCm symbol in a sampling verifier.

But the meta-lesson is consistent. The GPU vendor name matters less than:

• whether your workload actually stays on the fast kernel paths,
• whether your speculative decoding verifier is mathematically correct on ROCm,
• and whether prefill gets accelerated instead of silently falling back.

Or, put in one line: the best performance-per-dollar outcome comes from moving the bottleneck away from the parts your users feel most.

Sources used for verification (not part of the article)

• AMD MI355X memory and platform specifications. ()
• Kimi K3 parameter count and context window. ()
• GPUs.io methodology and example pricing pages used for normalization context. ()