Kimi K2.5 / K2.6 / K2.7-Code in MXFP4 Set a New Benchmark for AI Inference Performance on AMD Instinct™ MI355X GPUs

Introduction

Serving Kimi-K2.5, K2.6, and K2.7-Code efficiently requires optimizing the entire serving stack: MLA attention kernels, MoE GEMMs, all-reduce collectives on the critical path of each decode step, and the scheduler that controls the interaction between prefill and decode. Building on our previous work optimizing MLA attention kernels, this post focuses on the other three critical areas—MoE GEMMs, all-reduce communication, and scheduling. We show how ATOM brings these optimizations together on AMD Instinct™ MI355X GPUs to deliver high-performance, end-to-end MXFP4 inference.

Serving Kimi K2.5/K2.6/K2.7-Code efficiently means to stressing every part of the serving stack at once: the MLA attention kernels, the MoE GEMMs , the all-reduce collectives on every decode step's critical path, and the scheduler that decides when prefill is allowed to interrupt decode. This post focuses on  optimizing the remaining three components on AMD Instinct™ MI355X GPUs with ATOM, delivering high-performance end-to-end MXFP4 inference.

At a Glance

• Kimi K2.5/K2.6/K2.7-Code (1T-parameter MoE) served in MXFP4 on a single MI355X node at TP=4 the 4-bit weights fit comfortably in 4 × 288 GB of HBM3E, and one 8-GPU node hosts two independent TP4 replicas.
• Peak throughput of 5,369.6 tok/s/GPU on the 8k/1k workload, staying ahead of single-node NVIDIA B200 across multiple concurrency levels.
• Built together with the community: optimizations co-developed by AMD engineers and open-source contributors, merged upstream into AITER/ATOM/InferenceX through an AMD developer collaboration program.

Key Optimizations: From Kernel to Engine to Benchmark

Every performance improvement in this post follows the same three-stage, fully public pipeline: kernels optimized in AITER (ROCm/aiter, the AI Tensor Engine for ROCm), ATOM (ROCm/atom, AMD's PyTorch-native inference engine) integrates them into the serving stack, and the resulting evaluated in InferenceX, through continuous benchmarking. We walk through each stage.

AITER: kernel-level optimization

Deeper-fused MXFP4 (A4W4) MoE backends for gfx950

The heart of Kimi K2.7-Code inference — and of every model on the shared K2 backbone — is the grouped GEMM inside the MoE block. AITER already ran this path in full A4W4 MXFP4 — activations and weights both entering the matrix units as 4-bit, using the AMD CDNA™ 4 architecture's native micro scaling instructions with per-32-element E8M0 scales. The new backend (aiter #3832, building on #3470) rebuilds the MoE path around deeper fusion and Kimi-K2-specific tuning:

• Two grouped GEMMs, one fused bridge between them. The gate/up projection (gemm1) and the down projection (gemm2) run as two FlyDSL grouped GEMMs. Between them, a single kernel fuses the SwiGLU gate activation, the elementwise multiply, the MXFP4 quantization of the intermediate activations, and the E8M0 scale write in the sorted layout gemm2 expects — so the largest activation tensor in the model is never written out in high precision or re-read for a separate quantization pass.
• Inline activation quantization for decode. For small batch decode shapes, gemm1 converts incoming bf16 hidden states to MXFP4 in its prologue as they are gathered, keeping quantization directly on the GEMM's input path.
• Sorting fused with the GEMM metadata. An adaptive MoE-sorting backend emits the token-gather metadata (m_indices / reverse-sorted indices) the grouped GEMM consumes, in the same pass rather than as a separate pre-step.
• Per-shape tuning for K2's experts. Tile shapes and split-K are tuned per expert shape (kimik2_fp4_tuned_fmoe.csv), and the per-32-element scale amax is computed with hardware DPP instructions instead of a separate reduction.

Two GEMM configurations are kept because MoE GEMMs live in two regimes — many small expert problems at decode, near-dense GEMM at prefill — so both ends of the Pareto curve benefit.

These kernels are authored in FlyDSL (ROCm/FlyDSL), ROCm's Python-first, MLIR-native kernel DSL that lets developers express tiling, layouts, and data movement at a high level without giving up expert-level performance — a big part of why community contributors can iterate on AITER kernels this quickly (see the FlyDSL introduction on ROCm Blogs).

Single-stage fused all-reduce at production batch sizes

Every transformer layer ends in an all-reduce that sits on the critical path of each decode step. AITER's Quick Reduce, the fused all-reduce ATOM uses over XGMI links, chooses between a low-latency one-stage algorithm and a lower-traffic two-stage one. aiter #3458 and aiter #3880 introduced TP-aware gating, selecting the one-stage path through batch 32 at TP4 and batch 16 at TP8 before switching to two-stage for larger decode batches. This keeps communication latency low at common serving batch sizes without sacrificing efficiency as concurrency grows.

Faster expert routing

Routing every token to 8 of 384 experts runs once per layer per decode step. aiter #3466 replaces the original top-k selection with a fused single-kernel reduction specialized for K2's expert count and group-routing scheme, trimming the per-step non-GEMM overhead that otherwise eats into decode throughput.

ATOM: wiring the kernels into the engine

Per-shape backend dispatch

Fast kernels only pay off if the engine picks them at the right moment. ATOM's dispatcher selects the MoE backend per shape, so the decode- and prefill-optimized A4W4 kernels each serve the regime they win in; AITER_MXFP4_INTERMEDIATE=1 turns on the 4-bit intermediate-activation path.

Prefill scheduling and parallelism

The engine-level tuning matters just as much. --scheduler-delay-factor 1 changes when ATOM admits new prefills into a busy decode loop: instead of greedily injecting them — which fragments decode batches and spikes every in-flight user's latency — the scheduler briefly accumulates prefills and co-schedules them, smoothing the concurrency-64/128 operating points at negligible TTFT cost. And the parallelism choice is deliberate: TP=4 keeps expert GEMMs above the kernels' efficiency knee and collectives cheap, so the InferenceX submission runs TP4 only.

InferenceX: shipping it and measuring it in the open

The last mile is the benchmark itself: bump the public rocm/atom image, set the two flags above, and let the harness sweep the workload at concurrency 4–128. Because InferenceX re-runs continuously, this stage is also the safety net — each merge is re-measured on the same public benchmark, so improvements show up on the dashboard within days, and any regression is caught before it reaches users.

Performance Results

We benchmark with InferenceX's single-node fixed-sequence-length harness on the 8k/1k workload (ISL 8192 / OSL 1024), sweeping concurrency from 4 to 128 and plotting the resulting Pareto frontier of throughput (tok/s/GPU) versus interactivity (tok/s/user).

On the 8k/1k workload, MI355X with ATOM reaches 5369.6 tok/s/GPU at concurrency 128 while sustaining 19.2 tok/s/user, and 116.4 tok/s/user at concurrency 4 on the interactive end. As Figure 1 shows, its Pareto frontier sits above single-node NVIDIA B200 (vLLM, NVFP4) across multiple concurrency levels — MI355X GPU delivers more throughput at matched interactivity along the curve, not just at the peak.

Figure 2 shows the same trend from a latency perspective: when plotting end-to-end latency against interactivity, the MI355X GPU maintains lower latency at comparable interactivity. This means users can achieve faster responses at the same per-user token rate — demonstrating that the throughput gains shown in Figure 1 do not come at the expense of responsiveness.

This comparison is deliberately scoped for single-node serving. Rack-scale, disaggregated deployments — prefill/decode split across nodes with wide expert parallelism — represent a different deployment model that we'll cover separately (see ATOM mesh in the roadmap below).

How to Reproduce

To serve Kimi K2 models with ATOM on MI355X GPUs and reproduce the numbers above, follow the official recipe: ATOM Kimi-K2 recipe — Launching Server. It pins the image, model, and serving flags used for these results.

Built With the Developer Community

The optimizations described above were co-developed through open collaboration between AMD engineers and the broader developer community. Many ideas originated from community contributions; kernels, reviews, and benchmarks flowed in both directions; and improvements were integrated through public PRs across AITER, ATOM, and InferenceX. The A4W4 MoE backend from the RadeonFlow team is one example, now shipping in the public rocm/atom image and automatically selected by the dispatcher when it delivers the best performance for a given workload.

Much of this work grew from an earlier AMD collaboration program that invited developers to push MI355X GPU inference and rewarded the teams that met its performance bar with a $300k reward pool.  While the awards helped accelerate participation, the broader goal was to build a fully open ecosystem where contributions can move directly into production, be validated transparently, and serve as a foundation for the community to build upon.

The development loop is open to anyone. Whether you have a faster kernel, a better scheduling heuristic, or an improved configuration, you can contribute directly. Tools like FlyDSL make the kernel development more accessible, contributions can be submitted as PRs to AITER or ATOM . Open benchmarks like InferenceX give one fair, level field to compare results on, even though no single benchmark captures every workload and deployment. We are also upstreaming these optimizations into vLLM and SGLang, allowing the broader ecosystem to benefit from these improvements over time.  Our goal is to continue evolving the ROCm software stack — from kernels and serving engines to the surrounding tooling — together with the community that builds and deploys on it.

Summary

Kimi K2.5/K2.6/K2.7-Code on MI355X is now a fully 4-bit serving stack: MXFP4 weights and activations through the MoE, running at TP=4 on a single node, at 5,369.6 tok/s/GPU on 8k/1k — ahead of single-node NVIDIA B200. The individual ingredients — A4W4 fused MoE kernels, the collective-strategy gating fix, a faster router, and calmer prefill scheduling — are each modest; compounded, and traveling the AITER → ATOM → InferenceX pipeline in public, they move the frontier past B200.

And more is coming: Two-Batch Overlap and DP Attention for the Kimi K2 family, ATOM mesh — ATOM's distributed serving layer with prefill/decode disaggregation for multi-node deployments — and the vLLM/SGLang upstreaming mentioned above.

Additional Resources

• ATOM on GitHub · AITER on GitHub · FlyDSL on GitHub
• InferenceX benchmark platform
• Upstream kernel PRs: aiter #3470 (A4W4 MoE), aiter #3832 (RadeonFlow A4W4 MoE), aiter #3466 (TopK), aiter #3458 / #3880 (fused all-reduce)
• Related reading: FlyDSL: Expert GPU Kernel Development with a Python-Native DSL · DP Attention and TBO for DeepSeek-V4 on MI355X

Disclaimers

Third-party content is licensed to you directly by the third party that owns the content and is not licensed to you by AMD. ALL LINKED THIRD-PARTY CONTENT IS PROVIDED “AS IS” WITHOUT A WARRANTY OF ANY KIND. USE OF SUCH THIRD-PARTY CONTENT IS DONE AT YOUR SOLE DISCRETION AND UNDER NO CIRCUMSTANCES WILL AMD BE LIABLE TO YOU FOR ANY THIRD-PARTY CONTENT. YOU ASSUME ALL RISK AND ARE SOLELY RESPONSIBLE FOR ANY DAMAGES THAT MAY ARISE FROM YOUR USE OF THIRD-PARTY CONTENT.