vLLM 解析 Kimi K3 在 AMD Instinct MI355X 上的推理优化，单 GPU 吞吐超 NVIDIA HGX B300 24.4%

Inside the Kimi K3 optimization effort: understanding the bottlenecks and the decisions that made inference faster

In this article

• AgentX: the workload and results
• Inside Kimi K3: experts, attention and recurrent state
• 1. Cache: keep conversation state available for reuse
• 2. MoE: remove duplicate work across GPUs
• 3. Memory traffic: fuse operations and remove unnecessary copies
• 4. MLA: faster long-context decoding with AMD AITER assembly
• 5. KDA: parallelize prefill and speculative verification
• 6. Speculative decoding: choose how many tokens to draft
• What to optimize next
• Contributions and acknowledgments

On AgentX, an eight-GPU AMD Instinct MI355X system serving Kimi K3 with vLLM delivers 24.4% more total tokens per second per GPU than NVIDIA HGX B300 at 50 tokens/s/user P90 interactivity. Under the benchmark’s ownership-cost model, it delivers 87.4% more tokens per dollar.

In this post, we’ll unpack the engineering behind Kimi K3 on AMD Instinct MI355X, connecting vLLM’s serving decisions to GPU kernels and the hardware underneath. We’ll explain why the optimizations worked and build the intuition to spot wasted work, identify bottlenecks and find the next optimization in your own inference stack.

Published dense matrix peaks are 10.1 PFLOP/s MXFP4 for AMD Instinct MI355X and 13.5 PFLOP/s FP4 for NVIDIA B300. Both have 288 GB HBM3E and 8 TB/s peak memory bandwidth. Reaching this serving performance reflects sustained software engineering by AMD and collaborators across the vLLM community.

AgentX: the workload and results

The public AgentX benchmark replays real coding-agent session traces with synthetic tokens, retaining request lengths, timing and prefix-sharing patterns. The workload offers over 96% theoretical prefix reuse with unlimited cache. Speculative runs use a fixed synthetic acceptance reference.

At 50 tokens/s/user P90 interactivity, AMD Instinct MI355X serves approximately 7,567 total tokens/s/GPU, compared with 6,083 for NVIDIA HGX B300, a 24.4% throughput advantage. Using SemiAnalysis’s modeled TCO of $1.50 versus $2.26 per GPU-hour, respectively, that translates to 18.16 million versus 9.69 million total tokens per dollar of TCO: 87.4% more tokens per dollar.

The public curves compare complete serving stacks. The measurements beside each optimization come from separate tests and explain individual mechanisms under their stated configurations.

Inside Kimi K3: experts, attention and recurrent state

A returning agent turn can reuse most of its prefix, but processing a short continuation still touches three kinds of data.

Kimi K3 uses two kinds of attention. Its 24 MLA layers retain a compressed KV history that grows with the conversation. A new query still consults that history. Its 69 KDA layers instead carry forward state matrices with fixed dimensions, updating them as new tokens arrive. Model configuration.

In MoE layers, a router selects expert networks to transform each token. Their weights are shared across requests.

Across turns: cache reuse and offload (Section 1).
Within a model step: expert projections, kernel handoffs, and MLA/KDA scheduling (Sections 2–5).
Choosing draft depth: committed output versus complete cycle cost (Section 6).

1. Cache: keep conversation state available for reuse

DCP reduces duplicate history in GPU memory; CPU offload preserves reusable state outside it.

Under TP8, each GPU handles 12 of Kimi K3’s 96 MLA query heads but stores the entire compressed KV history. Tensor parallelism divides the heads without dividing the cache.

With eight-way decode context parallelism (DCP8), each GPU stores one eighth of the cached positions. During decode, the GPUs exchange queries, compute against their local history and merge their partial answers.

Each partial answer carries normalization information so the merge can weight shards correctly. #51705 extends the ROCm AITER DCP path to multi-token verification.

For a 142,000-token conversation, the estimated MLA payload is 24 layers × 142,000 positions × 576 FP8 bytes ≈ 1.96 GB per GPU. DCP8 reduces this to about 245 MB, before allocation overhead. Each GPU processes eight times as many heads against one eighth of the history, leaving the leading attention arithmetic roughly unchanged.

CPU offload: keep state that no longer fits in HBM

When state no longer fits in GPU memory, the AMD recipe’s native vLLM CPU offload preserves MLA history and KDA checkpoints in host DRAM.

A cache hit needs matching resume points. If MLA history reaches token 3,072 but the latest compatible KDA checkpoint is at 2,048, the request must resume at 2,048 and recompute the remaining prefix.

Saving the attention-block tail alongside the KDA checkpoint aligns their resume positions (#54736). Restoring that state trades a transfer back to the GPU for avoided prefix recomputation. Once a request resumes, its new tokens can still make several GPUs repeat the same calculation.

2. MoE: remove duplicate work across GPUs

On one GPU, the routed-expert output projection returns the complete vector its consumer expects. In the original ROCm path, an earlier reduction and normalization gave all eight GPUs the same latent input, and each held the full weights. The same projection therefore ran eight times.

The output coordinates can be divided among GPUs, but every rank still needs the complete result. Kimi K3 already has an all-reduce for its shared-expert partials.

Each GPU adds only the routed outputs it owns to its shared-expert partial. The existing all-reduce assembles the complete result, with no new collective (#51253). The ROCm implementation builds on the corresponding CUDA strategy.

Each output coordinate receives its routed contribution exactly once. The full projection weights remain replicated.

Output throughput rose 3.05–4.91% in the PR’s test on eight AMD Instinct MI355X GPUs at TP8, across concurrency 8–128 with random 8K-input/1K-output requests and prefix caching disabled.

Free up GPU memory to keep more cached state

Each rank now computes only its assigned projection coordinates while retaining the full weights.

Loading only the rows each rank uses frees about 3.85 GiB per GPU across 92 MoE layers at TP8 (#59591). The fused path’s calculation stays the same; the extra space can hold MLA history and KDA checkpoints.

This change landed after the image used for the headline public result. Its separate AgentX experiment found less than 1% higher throughput at concurrency 32 and 48, but 28.9% at 70, where the reported prefix-cache hit rate rose from 74.8% to 80.1%.

The larger gain at high concurrency is consistent with cache pressure. Comparing both implementations at equal cache capacity would help isolate that effect.

More on expert routing and four-bit activations.

3. Memory traffic: fuse operations and remove unnecessary copies

Within each GPU, the next place to look is the data passed between kernels.

Kimi K3’s single-token KDA path ran three compute kernels and two output copies. Convolution ran before a kernel that already fused Q/K normalization, decay and beta gates, and the recurrent update; output normalization ran afterward. The fused HIP path performs these operations in one launch and writes directly to the destination buffer (#50654).

Intermediate handoffs stay on chip. Convolution history and recurrent state still go to GPU memory for the next token.

AttnRes mixes earlier layer-block outputs with the running residual. Combining that mix with RMSNorm keeps the temporary mixture in registers (#50593). The running residual saved for later layers retains its BF16 rounding; the mixture stays in FP32 through normalization. Tests require exact persistent updates and allow numerical tolerances for normalized output.

With AttnRes fusion, output throughput rose 14.2% on eight AMD Instinct MI355X GPUs at TP8, with 100K-input/1K-output requests and prefix caching disabled. Long-prefill latency was effectively unchanged.

Copy inputs only for kernels that use them

The ROCm matrix-multiplication dispatcher packed inputs before choosing a kernel. When it selected the fallback, that kernel read the original input and left the copy unused. Moving packing into the specialized branches removed these copies (#58566).

The profile made the waste visible: 69 unnecessary copies per decode step, one per KDA layer. Each moved only 1,792 bytes but took about 4.8 microseconds. Moving those bytes faster would barely help; launch overhead dominated.

4. MLA: faster long-context decoding with AMD AITER assembly

Kimi K3 has twelve query heads per GPU at TP8. Padding to sixteen made AMD AITER’s assembly decode path available. In the reported long-context test, time per output token fell from 43.8 to 23.6 milliseconds.

The comparison used TP8 without DCP, at concurrency one. The adapter in #50578 copies four query heads, packs them into the sixteen-head layout and discards their outputs. DCP8 presents a different, 96-head shape.

This preserves the model's operation because each query head attends independently over cached positions. Adding another head does not change the original head's softmax. Only the query is padded in this decode path; the stored compressed KV history stays the same size.

A long conversation gives attention plenty of work, but that work needs to be spread across the GPU. In the historical code we inspected, a conservative scheduling hint restricted the Gluon path to one workgroup per request. That workgroup still processed the entire conversation. Padding made AMD AITER’s assembly path available, with a different policy for dividing the history into parallel pieces and combining the results.

PR-reported mean time per output token: eight AMD Instinct MI355X GPUs, TP8, concurrency one, MXFP4 model weights, BF16 attention and ordinary single-token decode.

The short-input case barely changes; the long-input case nearly halves time per output token. These serving measurements include query preparation and output handling.

5. KDA: parallelize prefill and speculative verification

MLA can split a retained history among workgroups. KDA presents a harder scheduling problem: each token depends on the preceding state. Fused prefill in #54038 prepares each token group’s combined transformation before its incoming state is known.

For example, “halve the state and add two,” then “halve it and add one,” combine into “take a quarter of the starting state and add two.” This rule can be computed before the starting state arrives.

Step through the KDA dependency sequence

KDA does this with matrix transformations. An ordered scan of the group summaries determines each group’s starting state; the groups then replay their tokens in parallel. Token order stays intact, but each group no longer waits for the preceding group’s replay.

Grouping adds a second walk and expands the fused chunk path from two launches to four. It helps when few sequences or heads expose too little parallel work; otherwise, the extra passes can cost more than they save.

Within each group, the recurrent walk keeps state in matrix-accumulator registers between chunks, avoiding a global-memory write and reload. State storage diagram.

Fused KDA delivered 5.3–5.6% more input throughput than Triton at 4K and 16K input lengths in a separate three-backend comparison: eight AMD Instinct MI355X GPUs, TP8, concurrency 32, prefix caching disabled and one-token outputs. This measures the complete backend. Benchmark results.

Use smaller tiles for speculative verification

Short verification batches need a different split. At TP8, each GPU has 12 KDA heads with 128 value columns. Splitting each head into four 32-column tiles launches only 48 workgroups on 256 compute units. #58045 switches to four-column tiles, raising the count to 384.

In a separate graph-replay kernel test on one AMD Instinct MI355X GPU (TP8 dimensions, one sequence, four verification positions), time fell from 16.20 to 5.61 microseconds per layer (2.89×). This includes narrower tiles, single-wave workgroups, uniform-length specialization and hoisted token-invariant work. Ragged batches retain a dynamic path; tests compare outputs and recurrent state.

6. Speculative decoding: choose how many tokens to draft

Faster verification can change how many draft tokens are worth checking. The useful rate is committed output divided by the full draft–verify–commit cycle time.

For example, AgentX’s fixed acceptance reference gives 3.00 expected committed tokens per cycle with three drafts, and 3.36 with four, including the recovery or bonus token. That is 12% more committed output per cycle. Four drafts improve throughput only if the complete cycle takes less than 12% longer than three drafts.

The public AMD recipe uses three to seven drafts with DCP1 and disables speculation at the larger DCP8 concurrency settings.

What to optimize next

Ongoing prefill/decode disaggregation work adds another place to look: worker balance and whether transfers finish before their data is needed.

For another model or workload, start with the deployed request shapes and make three checks:

Compare ranks. Are several GPUs computing the same result before a collective that could assemble it?

Follow buffers. Does each temporary tensor have a consumer on the selected path? Can the data stay on chip?

Count independent work. Do the actual heads, tokens and tiles produce enough workgroups to spread the load?

Measure the complete request at the response speed you need, and check both outputs and saved state. A faster kernel matters when it shortens the work holding up the next token.

Contributions and acknowledgments

The Kimi K3 effort spans more than 50 merged AMD/ROCm-related vLLM PRs from AMD, Embedded LLM and the wider community, alongside AITER development, integration, review and benchmarking. Contributions cover attention, recurrence, MoE, quantization, graph execution, cache support and platform integration.

Special thanks to: Pin Siang Tan, Jun Kang Chow, Ye Hur Cheong, Hongxia Yang, Seungrok Jung, Bill He, Yichao Zhu, Hyukjoon Lee, Andy Luo, Mrinal Karvir, Richard Li, Sun Peng

Andy Friedrich, Ethan Yang, Tang Rui, Andreas Karatzas, Ankur Singh, Yihua Cheng, Badr Ba Sowid, Benjamin Chislett, Taoyu Zhu, Bill He, Bin Ding, Bowen Bao, Brandon Zhang, Chaemin Lim, Charles Wu, Chuan (Richard) Li, Felix Li, Cameron Quilici, Weishu Deng, Dewei Wang, Divakar Verma, Douglas Lehr, Edwin Lim, Yudong Pan, Elvir Crnčević, Fang Che, Fangzhou Ai, Xingran Fan, Fulu Li, Felix Marty, Canlin Guo, Summer Yang, Haichen Zhang, Yufeng He, He Liu, Hongxia Yang, Hyukjoon Lee, Yifan Qiao, James E T Smith, Jiahui Cao, Yanyuan Qin, Jonathan Chen, Juntian Liu, Shao-Chun Lee, Kedar Potdar, Kevin H. Luu, Kuan-Fu Liu, Kuntai Du, Kyle Sayers, @lalala-sh, @linnea-lin-00638949, Yinfeng Liu, Lucas Wilkinson, Gabriel Wu, Markus Hartikainen, Maithili Joshi, Jeff (Junze) Ma, Matthew Wong, Misha Goin, Micah Williamson, @minmengdie, Matvei Pashkovskii, Mathews Neha, Nicolò Lucchesi, Nick Hill, Nicholas Susanto, Oxana Korzh, Pei Zhang, @ppalanga, Peiran Qin, Qiuyang Yue, Rita Brugarolas, Rebecca Lee, Ryan Rock, Rohan Potdar, @Saddss, Samuel Shen, Sean Filimon, Seungrok Jung, Shantipriya Parida, Shanshan Shen, Shengnan Xu, Simon Danielsson, Siraj Shaikh, Stacy Roberts, Stefan Koncarevic, Asan Stefanski, Minghui Jiang, Tarık Saraç, Tun Jian Tan, Lingpeng Jin, Vansh Bhatia, Maksim Vedenev, Ye Hur Cheong, @wenwzhan, Wei Zhao, Xavier Aguilar, Xiaobing Zhang, Xiaohu Guo, Xin Li, Xudong Lyu, Yanxue Ren, Wentao Ye, Yongqin Wang, Riggins Chen, Chao Hong, Yuchen Shen, Matt Qin, Yichao Zhu, Roger Wang, Zheng Gong, Jiangyun Zhu, Zvi Grinberg, Yongye Zhu.

Benchmark notes

Serving configurations at the comparison targets

Both systems use eight GPUs, TP8, FP4 weights and FP8 KV. Values interpolate within each published offload or GPU-only series at matching interactivity; concurrency, DCP and draft depth can vary between the bracketing runs. The pinned AMD recipe (YAML) and NVIDIA launcher give the image tags and full settings.

Replay uses the cc-traces-weka-062126 dataset and a pinned AIPerf harness, with warmup excluded and one hour of profiling under fixed synthetic acceptance curves. The recipes give the seed, session-start and warmup settings. The downloaded dataset revision is unrecorded. Closed-loop replay can produce different completed request mixes.

Software and offload. The submissions use vLLM 36768d1bf (AMD) and 5894fdf98 (NVIDIA). Configured host-cache budgets are 224.875 and 281 GB per GPU, respectively; connector and image settings are in the linked recipes.

Costs. SemiAnalysis’s large-hyperscaler ownership model and published GPU cost constants use $1.50 per GPU-hour for AMD Instinct MI355X and $2.26 for NVIDIA HGX B300. Total tokens/$ = total tokens/s/GPU × 3,600 ÷ modeled $/GPU-hour. Some capital and utilization assumptions remain unpublished. At the 50-token interactivity target, the comparison gives 46.6% lower modeled cost per token.

Snapshot. SemiAnalysis InferenceX data retrieved October 9, 2026; submissions dated September 30, 2026 (AMD) and August 16, 2026 (NVIDIA). PR benchmarks use their stated workloads; #59591 is a separate experiment after the headline image.

Sources: AMD Instinct MI355X, NVIDIA HGX B300, NVIDIA memory specifications. Benchmark details: AgentX methodology.

Technical note: expert routing and precision

Kimi K3 selects 16 of 896 experts per token, so a substantial server batch can still leave few rows per expert. The resulting expert tiles may contain padding.

A4W4 reduces a different part of the traffic: the integration in #53940 changes activations from eight bits to four, while weights remain four-bit. Total traffic also includes BF16 intermediate outputs and expert-tile padding.

Technical note: KDA precision and on-chip memory

Each chunk still needs to resolve dependencies between its tokens. The kernel combines them into a transformation that later matrix products can apply. That creates two practical questions: which precision preserves the answer, and how much temporary data fits on chip?

A small example shows the idea. Substitute the earlier corrections into the later ones, and each answer can be expressed directly in terms of the known inputs. The resulting coefficient table has already accounted for the chain of dependencies:

The real prologue computes the inverse of the 64×64 within-chunk dependency matrix. Later matrix products use this transformation to account for those dependencies.

The inverse's intermediate values can grow and cancel. Rounding them too early loses information the final result needs. The implementation uses AMD's native FP32 matrix instructions, retaining FP32 inputs and accumulation for the sensitive inverse. Only the completed inverse is rounded to BF16 for subsequent products. Kernel implementation.

Why FP32 accumulation cannot recover rounded inputs

Consider subtracting 256 from 256.5. Both inputs are exact in FP32. At this magnitude, converting 256.5 to BF16 rounds it to 256:

FP32 subtraction cannot recover information lost before it starts. These toy values illustrate why the inverse uses native 16×16×4 FP32 matrix instructions, retaining both input and accumulation precision.

An FP4 model also relies on operations at higher precision. Their performance depends on making those numerical paths efficient.

The prologue also has a local-memory budget. AMD calls its on-chip shared scratch LDS. Q and K tiles reuse storage because they are needed at different times; the triangular input stores only its ten occupied 16×16 blocks. The workgroup requests 79.5 KiB, which rounds to 80 KiB, half of an AMD CDNA 4 compute unit’s 160 KiB LDS budget.

Test this threshold by reserving enough unused LDS to cross it while keeping the kernel binary fixed. If runtime barely changes, the extra resident workgroup was contributing little for that shape.

State between chunks

Return to the KDA serving result.

Disclaimers

The information presented in this document is for informational purposes only and may contain technical inaccuracies, omissions, and typographical errors. The information contained herein is subject to change and may be rendered inaccurate for many reasons, including but not limited to product and roadmap changes, component and motherboard version changes, new model and/or product releases, product differences between differing manufacturers, software changes, BIOS flashes, firmware upgrades, or the like. Any computer system has risks of security vulnerabilities that cannot be completely prevented or mitigated. AMD assumes no obligation to update or otherwise correct or revise this information. However, AMD reserves the right to revise this information and to make changes from time to time to the content hereof without obligation of AMD to notify any person of such revisions or changes. THIS INFORMATION IS PROVIDED ‘AS IS.” AMD MAKES NO REPRESENTATIONS OR WARRANTIES WITH RESPECT TO THE CONTENTS HEREOF AND ASSUMES NO RESPONSIBILITY FOR ANY INACCURACIES, ERRORS, OR OMISSIONS THAT MAY APPEAR IN THIS INFORMATION. AMD SPECIFICALLY DISCLAIMS ANY IMPLIED WARRANTIES OF NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR ANY PARTICULAR PURPOSE. IN NO EVENT WILL AMD BE LIABLE TO ANY PERSON FOR ANY RELIANCE, DIRECT, INDIRECT, SPECIAL, OR OTHER CONSEQUENTIAL DAMAGES ARISING FROM THE USE OF ANY INFORMATION CONTAINED HEREIN, EVEN IF AMD IS EXPRESSLY ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. AMD, the AMD Arrow logo, and combinations thereof are trademarks of Advanced Micro Devices, Inc. Other product names used in this publication are for identification purposes only and may be trademarks of their respective companies. © 2026 Advanced Micro Devices, Inc. All rights reserved.