Kimi K3.1 Pro on AMD Instinct GPUs

TL;DR

• Kimi K3.1 Pro, Moonshot AI’s open-weight, natively multimodal agentic model (2.8T total and 104B activated parameters, 1M-token context), serves on 8× AMD Instinct™ MI355X GPUs with both ATOM and vLLM on day 0, from the released MXFP4 checkpoint and with its DSpark speculative-decoding draft.
• K3.1 Pro keeps Kimi K3’s hybrid architecture: 69 Kimi Delta Attention (KDA) layers and 24 latent-attention layers over an 896-expert MXFP4 MoE. It adds a sparse indexer (DSA) to every one of the 24 attention layers, which selects 2,048 tokens of history per query. ATOM serves the indexed layers on its shared DSA path, with an FP8 index cache next to the FP8 latent KV cache
• ATOM: native KDA and sparse MLA kernels, the routed experts on FlyDSL’s MXFP4 grouped MoE, per-channel FP8 attention applied at start-up, and a DSpark draft that shares the target’s KV cache and runs as a captured CUDA graph.
• vLLM: memory-efficient sparse attention with decode context parallelism (DCP) for long-context serving, alongside a specialized decode path for low-latency interactive workloads
• Accuracy on Moonshot AI’s Kimi Vendor Verifier, at reasoning effort max: OCRBench 0.871 (Moonshot reference 0.872); tool-call JSON schema 0.991.

Introduction

Kimi K3.1 Pro runs on AMD Instinct MI350 Series GPUs from day 0 with ATOM, the AMD inference engine, on ROCm™ software with AITER and FlyDSL kernels. The model is post-trained from Kimi K3 with large-scale joint reinforcement learning. Moonshot AI describes it as its most capable 3T-class model, with a 1M-token context window, selectable reasoning effort (low, high, max), native tool calling and image input.

This post covers what K3.1 Pro changes relative to Kimi K3, how ATOM serves it, the serving recipe, and its accuracy on Moonshot AI’s vendor verification suite.

Kimi K3.1 Pro at a glance

For inference, five properties matter:

• Hybrid attention, now sparse: Of the 93 decoder layers, 69 use Kimi Delta Attention, a gated linear attention with a fixed-size recurrent state per request. The other 24 are gated latent attention (MLA) layers: every fourth layer, plus the last. In K3.1 Pro each of them also has an indexer (32 heads of dimension 128) that scores the request’s history and selects the top 2,048 tokens for attention. Attention then reads at most 2,048 tokens per query however long the context grows; only the lightweight indexer scores the whole history.
• Two small caches per attention layer: Each MLA layer caches one 576-wide latent per token (512 compressed KV plus 64 RoPE dimensions) and one index key per token.
• A wide, latent MoE: There are 896 routed experts, 16 active per token, plus two shared experts. The routed experts work in a 3,584-wide latent space, half the 7,168-wide hidden size, and use the SiTU-GLU activation. Only the first layer is dense.
• MXFP4, quantization-aware trained: The routed experts are stored in MXFP4 (one E8M0 scale per 32 values). Attention, the indexer, the shared experts, the dense MLP, the latent up/down projections, the router, the LM head and the vision encoder are BF16.
• DSpark speculative decoding: A separate five-layer MLA draft reads hidden states from five of the target’s last layers and proposes a block of draft tokens per step.

Item

Configuration

Parameters

2.8T total, 104B activated

Decoder layers

93: 69 KDA, 24 gated DSA (MLA + indexer); layer 1 dense, the rest MoE

KDA

96 heads, head dim 128, short convolution of width 4

MLA

96 heads; q LoRA rank 1,536; KV LoRA rank 512 + 64 RoPE dims; output gate

DSA indexer

32 heads × 128 dims; top-2,048 tokens; every MLA layer

Experts

896 routed (top-16, sigmoid routing) + 2 shared; expert intermediate 3,072 in a 3,584-wide latent

Vocabulary / context

163,840 / 1,048,576 tokens

Vision

MoonViT-V2 encoder (401M parameters)

DSpark draft

5 MLA layers (96 heads), interleaved YaRN RoPE, low-rank Markov head (rank 512)

Serving with ATOM

ATOM serves Kimi K3.1 Pro with tensor parallelism over the eight GPUs of one MI355X node, for text and image input, from the released MXFP4 checkpoint without conversion. With the DSpark draft loaded, each GPU peaks at about 199 GB of its 288 GB during start-up, and a single 1M-token request fits on one node.

What we optimized

Kimi Delta Attention: Prefill runs AITER’s chunked Kimi Delta Attention kernel. Decode advances each request’s recurrent state with a fused gated delta-rule update, and the states live in a per-request pool.

Sparse MLA on the shared DSA path: ATOM’s DSA implementation, shared with the DeepSeek-V3.2 and GLM-5 class models, serves K3.1’s indexer. The latent KV and the index keys live in FP8 paged caches with 128-token blocks. The indexer’s key and head-weight projections are packed into the layer’s fused q/kv down-projection, so one GEMM produces the query and KV latents and the index keys. The top-2,048 selection then drives a sparse MLA attention.

The MXFP4 latent MoE on FlyDSL: The routed experts run on FlyDSL’s native SiTUv2 grouped-MoE kernel with MXFP4 weights and an FP8 second stage. The two shared experts run on a second stream, concurrently with the routed experts.

PTPC FP8 attention: At start-up, ATOM re-quantizes the BF16 attention and dense linears to FP8 with one scale per output channel, with activations scaled per token at run time (--online_quant_config with ptpc_fp8): 559 linears per rank, done in about 2 seconds. The routed experts, the KDA convolutions, the vision encoder and the indexer-packed down-projection keep their stored formats.

DSpark speculative decoding: The draft’s 576-wide latent has the same shape as the target’s MLA layers, so ATOM binds the five draft layers into the target’s KV pool as extra rows, and they inherit the FP8 cache dtype. The draft block, which proposes all draft tokens in one parallel pass, is captured in a CUDA graph.

Decode context parallelism. For many concurrent long-context requests, ATOM shards the latent KV cache along the sequence across all eight GPUs (--decode-context-parallel-size 8). Each GPU attends over its own shard, and the partial results are merged with an all-to-all. The MLA query projection is replicated across the group, so no per-step query all-gather is needed. The sparse attention metadata is rebuilt after each indexer layer.

Mono decode. For latency-bound serving, ATOM replaces each KDA layer of a decode or verify step with two fused kernels: K1 runs the attention path through the KDA recurrence, and K2 runs the output projection plus the latent MoE. The MLA layers keep the paged sparse attention and also run K2. Mono serves steps of up to 64 token rows, such as one request verifying seven draft token.

Performance

Single-request latency with mono decode. Mono decode lowers the mean time per output token by 33.6% compared with the same server on the standard decode path.

Metric

Standard decode

Mono decode

Improvement

Mean TPOT

5.78 ms

3.84 ms

33.6%

Throughput with decode context parallelism.

Input / output tokens

Total throughput (tokens/s)

Output throughput (tokens/s)

Median TTFT (s)

Median TPOT (ms)

8K / 1K

7,484

831.6

0.821

16.36

128K / 100

14,005

10.68

27.9

1,146

How to run

Launch the container:

Start the server (TP8 with DSpark):

Notes:

• --block-size 128 is required: it is the MLA layout, and the draft’s rows live in the target’s pages.
• Without --max-model-len, the server uses the model’s 1,048,576 tokens.
• Without speculative decoding, drop the last line.
• Reasoning effort is set per request with the top-level reasoning_effort field (low, high, max). Multi-turn and tool-call conversations must pass back the full assistant message, including reasoning_content and tool_calls.

Quick check:

The reasoning is returned in reasoning_content, and the call get_weather({"city": "Paris"}) in tool_calls.

Accuracy

Evaluated against ATOM on 8x MI355X GPUs with thinking at reasoning effort max unless noted. MMMU-Pro, OCRBench, BEAM and the tool-call schema check come from Moonshot AI’s Kimi Vendor Verifier (KVV).

Benchmark

Parameters

ATOM

MMMU-Pro (standard, 10 options), 1,730 questions

effort max, temperature 1.0, top_p 0.95, 98,304 max tokens

0.842 (±0.009)

OCRBench, 1,000 questions

effort max, temperature 1.0, top_p 0.95, 16,384 max tokens

0.871 (±0.011)

BEAM (1M), 700 questions

effort max, temperature 1.0, top_p 0.95; judged by GPT-4.1

0.416

GSM8K 5-shot, 1,319 questions

greedy

0.978 flexible / 0.977 strict (±0.004)

Tool-call JSON schema, 428 checks

effort max, 2,048 max tokens, one attempt per check

0.997 (427 / 428)

• MMMU-Pro and the tool-call schema check ran on the server in How to run.
• OCRBench ran on TP8/DCP8 with a BF16 KV cache, an FP8 index cache and seven draft tokens.
• GSM8K and BEAM ran on TP8 without speculative decoding; BEAM used a 1M-token context with prefix caching.
• BEAM was judged by GPT-4.1 rather than the judge model behind Moonshot AI’s reference, so the two scores are not on the same scale.

Serving with vLLM

vLLM supports Kimi K3.1 Pro on AMD Instinct MI355X GPUs with the MXFP4 checkpoint and DSpark speculative decoding. Built on AITER and FlyDSL kernels, the ROCm implementation combines memory-efficient sparse attention with decode context parallelism (DCP) for long-context serving, alongside a specialized decode path for low-latency interactive workloads.

What we optimized

Native MXFP4 indexing. The indexer scans the context to select up to 2,048 tokens for sparse attention. MXFP4 reduces each index key from 132 bytes in FP8 to 68 bytes, including scales, cutting index-cache storage by 48.5%. Scoring uses native FP4 matrix instructions, with accumulation and head-weighted reduction in FP32. During speculative verification, multiple query rows reuse the same keys to reduce repeated reads.

The smaller index cache leaves more GPU memory available for cached context, benefiting long-context serving at high concurrency. In AgentX sweeps, the MXFP4-indexer configuration delivered 5.3% higher total workload throughput at concurrency 48.

Prefill: Context-parallel sparse attention. During prefill, GPUs gather index keys and divide query rows among themselves. Each GPU scores its assigned rows against the full context and selects the global top-k directly, avoiding a merge of per-rank candidates. Main attention can separately gather KV when this reduces communication and fits the workspace budget, then process all query rows for local heads.

The combined DCP prefill optimizations improved throughput and latency in development benchmarks, with the largest gains on long inputs.

Input / output tokens

Total throughput (tokens/s)

Output throughput (tokens/s)

Median TTFT (s)

Median TPOT (ms)

8K / 1K

5,357.8 → 6,015.4 (+12%)

595.3 → 668.4

2.462 → 1.875

24.571 → 22.063

128K / 100

6,950 → 17,678 (+154%)

5.30 → 13.48

268.7 → 105.5

331.5 → 129.6

FlyDSL sparse MLA reads the selected tokens directly from the FP8 KV cache. Each workgroup processes one query row across its local attention heads, providing parallelism across the many query rows in a prefill step. Smaller batches retain the Triton path, which exposes additional parallelism within each row. In a separate controlled comparison on eight MI355X GPUs, enabling FlyDSL sparse MLA increased output throughput by 7.0% for the 128K-input workload.

Implementation

Benchmark duration (s)

Median TTFT (s)

Output throughput (tokens/s)

Triton

432.8

96.2

14.79

FlyDSL

404.4

89.8

15.83

Improvement

6.6% lower

6.6% lower

7.0% higher

Decode: Fused low-latency decode. DSA indexing and sparse attention run separately, followed by a persistent monokernel that fuses attention output projection, AttnRes and the full latent MoE, including tensor-parallel all-reduces, into one launch. KDA layers also fuse input projections, short convolution, recurrent-state updates and gated normalization. The specialized path supports verification batches of up to eight token rows, including one request with seven draft tokens.

Metrics

Normal Path

Specialized decode

Improvement

mean TPOT

6.02 ms

4.40 ms

27.0%

Communication: Efficient cross-GPU communication. During decode, each GPU exchanges candidate scores while keeping token IDs local. A shared threshold and consistent tie handling determine which local candidates are selected. For 2,048 candidates, exchanging scores alone halves the candidate payload from 16 to 8 KiB per query per GPU.

AITER accelerates DCP all-gather with a dedicated GPU kernel that reads peer buffers through HIP IPC and writes directly into the required output layout. This combines data gathering and layout conversion, avoiding a separate reordering copy. In standalone microbenchmarks on eight MI355X GPUs, the AITER path reduced all-gather latency by 53–61% versus RCCL for the tested query and candidate tensor shapes

How to run

The following configuration uses eight AMD Instinct MI355X GPUs with TP8/DCP8, FP8 MLA KV cache, MXFP4 index keys and DSpark with five draft tokens. Set the checkpoint paths to local copies of Kimi K3.1 Pro and its matching DSpark draft model.

This recipe uses standard rejection, matching the accuracy evaluations below. CPU KV offload reserves approximately 1.64 TiB of host memory across eight ranks, in addition to other runtime requirements. The 1M-token context limit and 48-sequence limit do not guarantee capacity for 48 simultaneous 1M-token requests.

Launch the container:

Start the server (DCP8 with DSpark):

vLLM accuracy

The evaluations below compare FP8 and MXFP4 indexer with DSpark using standard rejection.

Benchmark

vLLM FP8 indexer

vLLM MXFP4 indexer

GSM8K 5-shot, T = 0.7

1295/1319 (98.18%)

1291/1319 (97.88%)

Passkey, 8K–500K context, 25 cases

25/25

25/25

LongBench-v2 short, greedy

125/180 (69.44%)

124/180 (68.89%)

OCRBench, effort max, temperature 1.0, top_p 0.95, 16,384 max tokens

0.871 ± 0.0106

0.870 ± 0.0106

Tool-call JSON schema, effort max, 2,048 max tokens

407/408 (99.75%)

407/408 (99.75%)

Acknowledgements

Kimi K3.1 Pro day-0 support is a joint effort across the model vendor, inference engine, and hardware communities.

We thank the Moonshot AI team for creating Kimi K3.1 Pro, sharing architecture details ahead of the weight release, contributing the initial model integration, and collaborating closely on correctness and production validation.

We thank the Inferact team for integrating the model into vLLM, implementing serving semantics and multimodal support, building deployment recipes, and driving end-to-end performance optimization.

We thank Jin Lingpeng, Wang Yiting, Zhu Yuhua, Yan Zhiwei,Zhang Lirong and Wu Hattie for the AMD ATOM day-0 support, and Bill He, Yichao Zhu, Seungrok Jung, Hyukjoon Lee, Hongxia Yang, Richard Li for the AMD vLLM day-0 support and Carlus Huang, Gilbert Lei , Peng Sun and Andy Luo for cross-team collaboration.

Most importantly, we thank the broader open-source and AMD ROCm developer communities for the anticipation, testing, and feedback already surrounding Kimi K3.1 Pro on AMD Instinct GPUs. We look forward to putting the weights and the ATOM and vLLM inference engine support for AMD Instinct MI355X GPUs in your hands.

Test configuration and disclaimer

ATOM results were measured on 8× AMD Instinct MI355X GPUs in a ROCm ATOM development container with the Kimi K3.1 changes, in November 2026. Checkpoint: the Kimi K3.1 Pro release (MXFP4 routed experts) with the Kimi K3.1 DSpark draft. Accuracy from the Kimi Vendor Verifier. Results on other hardware, software versions, settings or workloads will differ. These figures are for illustration and are not a performance guarantee.

Footnotes

Appendix: environment variables

Variable

Set to

Description

AITER_SITUV2_A4W4

FlyDSL native SiTUv2 MXFP4 grouped MoE for the routed experts

AITER_FLYDSL_STAGE2_FP8

FP8 second stage in the FlyDSL MoE

ATOM_USE_FLYDSL_FP8_PREFILL_ATTN

FlyDSL FP8 prefill attention on the MLA layers

ATOM_GDN_SSM_DTYPE

fp16

dtype of the per-request KDA state pool

ATOM_STATE_CHECKPOINT_DEMAND

prompt-end state checkpoints only

AITER_LOG_LEVEL

WARNING

quieter AITER start-up log