Kimi K3.1 Pro is Moonshot AI’s most capable open-weight model: a 2.8T-parameter (104B active) multimodal MoE with a 1M-token context that matches the strongest closed-source models on coding, agentic and knowledge benchmarks.
The main change is attention: K3.1 Pro adds Sparse Attention (DSA) to all 24 full-attention layers for efficiency. vLLM supports it on day 0 on NVIDIA GPUs, AMD GPUs, Google TPUs and Alibaba T-Head Zhenwu PPUs.
TL;DR #
- Sparse attention on a hybrid model: K3.1 Pro keeps K3’s KDA + MLA hybrid and adds a DSA indexer to all 24 MLA layers. At long context, a decode step reads about a quarter of the bytes per layer that dense MLA reads.
- One KV cache for three kinds of state: vLLM packs each block’s MLA latent, indexer keys and KDA state together. The packed cache runs with FlashInfer and FlashMLA sparse kernels with FP8 or NVFP4 KV and optional MXFP4 index keys, and P/D disaggregation transfers it unchanged.
- Token-sharded prefill attention: up to 1.69× faster time to first token (TTFT) than head-sharded prefill at 340K tokens on 8× GB300.
- Decode context parallelism (DCP) with a sparse indexer: an exact top-k merge with**~10× less traffic** makes the attention module1.3–1.5× faster at DCP8. DCP prefill attention runs 4.1–4.3× faster per layer, and DCP8 holds 7.2× the KV tokens of TP8.
- Fused low-latency decode: the specialized CuTe DSL decode path from Kimi K3 now covers DSA layers, making decode 1.2–1.6× faster on 8× GB300, and 1.9× faster under DCP8.
- Agentic serving: on NVIDIA GB300, up to 375.9 tok/s per user at the latency-sensitive end and 37.8K total tok/s per GPU at the throughput end; on NVIDIA Rubin, up to 815.6 tok/s per user and 68.7K total tok/s per GPU on the SemiAnalysis AgentX benchmark.
- Zero-mismatch RL on day 0: for the first time, day-0 support comes with zero-mismatch RL: vLLM and Megatron produce bit-identical log-probs for both the base model and the LoRA adapter, so the PPO KL is exactly 0 at every step of a 100-step full-model run.
Quick start #
On NVIDIA GPUs:
vllm serve moonshotai/Kimi-K3.1-Pro \
--tensor-parallel-size 8 \
--trust-remote-code \
--load-format fastsafetensors \
--enable-auto-tool-choice \
--tool-call-parser kimi_k3 \
--reasoning-parser kimi_k3 \
--speculative-config '{"model":"Inferact/Kimi-K3.1-Pro-DSpark-preview","num_speculative_tokens":5,"method":"dspark","draft_sample_method":"probabilistic"}'
On AMD GPUs:
VLLM_ROCM_USE_AITER=1 SAFETENSORS_FAST_GPU=1 \
vllm serve moonshotai/Kimi-K3.1-Pro \
--served-model-name kimi-k31-pro \
--tensor-parallel-size 8 \
--trust-remote-code \
--load-format fastsafetensors \
--max-num-seqs 16 \
--gpu-memory-utilization 0.85 \
--speculative-config '{"model":"Inferact/Kimi-K3.1-Pro-DSpark-preview","num_speculative_tokens":5,"method":"dspark","draft_sample_method":"probabilistic"}'
What’s new in Kimi K3.1 Pro #
| Kimi K3 | Kimi K3.1 Pro | |
|---|---|---|
| Layers | 93: 69 KDA + 24 full attention | 93: 69 KDA + 24 full attention |
| Full-attention layers | dense MLA | MLA + DSA indexer (top-2048), NoPE |
| Indexer | none | 32 heads × 128 dims, FP8 keys |
| MoE | 896 experts, 16 active, MXFP4 | unchanged |
| Context window | 1M | 1M |
| Total parameters | 2.8T | 2.8T |
Sparse Attention in K3.1 Pro
Each sparse attention layer has a small indexer next to its MLA attention. The indexer keeps one FP8 key per token (132 bytes with its scale) and scores every earlier token for each query with 32 lightweight heads. Then it keeps the top 2,048 tokens. MLA attention then runs exactly, but only over those 2,048 tokens. DSA was introduced in DeepSeek-V3.2; our DeepSeek-V3.2 post covers the indexer in more detail.
The payoff is largest in long-context decode, which is bound by memory bandwidth. Dense MLA reads 576 bytes of FP8 latent per context token per layer. DSA still scans the whole context, but only the 132-byte index keys. It then reads the full latent for just 2,048 tokens. At long context, that is about 23% of the bytes dense MLA reads per layer.
Figure 1: Sparse attention (DSA) in Kimi K3.1 Pro. The indexer scores every past token, keeps the top 2,048, and every query head attends exactly over only those tokens’ latents. Step through it.
KV cache footprint
The index keys are a second per-token cache next to the MLA latent. Under tensor parallelism (TP) both are replicated on every GPU; KDA state is per request, as in K3.
| KV cache per token, per GPU (24 full-attention layers) | Bytes |
|---|---|
| Kimi K3, FP8 MLA latent | 13,824 |
| Kimi K3.1 Pro, FP8 latent + FP8 index keys | 16,992 |
| Kimi K3.1 Pro, NVFP4 latent + MXFP4 index keys | 10,080 |
The latent and the index keys can be quantized independently. An NVFP4 latent alone already brings K3.1 Pro below K3 with FP8. Opt-in MXFP4 index keys fit 11% more tokens on top and make the indexer 10–20% faster at long context. With both, eight GB300 GPUs hold 1.7× as many tokens as with FP8, and DCP8 multiplies that by seven: about 36 full 1M-token contexts at once, against 3 for TP8 with FP8.
Figure 2: Measured KV cache capacity on 8× GB300. Change the latent and index-key formats, the parallelism, or the context length to see how many full contexts fit.
A KV cache for three kinds of state #
Each DSA layer stores two per-token caches of different widths: the MLA latent (576 bytes in FP8, 352 in NVFP4) and the index keys (132 bytes in FP8, 68 in MXFP4). Neither width is a multiple of the other. The 69 KDA layers each add a recurrent state per request that is much larger than one layer’s page of tokens. A cache that gives every layer pages of the same size would have to pad the small pages up to the largest one.
vLLM packs the cache by block instead of by layer. One block of the cache is a single contiguous row that holds that block’s MLA pages and index-key pages for all 24 DSA layers, together with the KDA states that share the block. Pages inside a row can have different sizes; the only constraint is that the row holds the KDA states. The block size is a power of two, so blocks of prefill and decode instances that run at different TP sizes divide evenly and can be paired for KV transfer.
The attention kernels still see the small fixed-size pages they expect. vLLM translates block tables and cache slots into kernel-page addresses before each step, so FlashInfer’s and FlashMLA’s sparse MLA kernels and DeepGEMM’s indexer kernels read the packed cache without changes.
This layout makes the rest of the stack work:
- Any sparse MLA kernel and KV format. On Blackwell the default is FlashInfer’s sparse MLA with an FP8 latent. FlashInfer’s new PrimsTS kernels also serve an NVFP4 latent.
- Disaggregated serving. Each block is one contiguous range of memory, so a P/D transfer moves it with one large RDMA read instead of dozens of small per-layer reads.
- Shared infrastructure. The same packed planner serves GLM-5.3-Flash and DeepSeek-V4, whose caches have the same problem of mixed page sizes (see theDeepSeek V4 post ).
Figure 3: How the packed cache is laid out at TP8. The attention group and the three KDA groups share one pool of block rows; pick a token, a DSA layer and a KDA layer to see which block and page each one lands in.
Prefill: shard the attention by token, not by head #
Under TP8, each GPU owns 12 of K3.1 Pro’s 96 attention heads. That is a poor fit for sparse MLA prefill, for the same reasons described in our GLM-5.3 post:
- MLA’s latent KV is shared by all heads, so every GPU reads the full cache to serve only 12 heads, and the sparse kernels pad those 12 heads to their minimum tile.
- The indexer is not split by head at all: every GPU scores every token, so its quadratic cost is repeated eight times.
For large prefill steps, vLLM shards the attention by token instead. An all-to-all moves each GPU’s query from “all tokens, 12 heads” to “1/8 of the tokens, all 96 heads”. Each GPU runs the indexer and sparse MLA for its own tokens, and a second all-to-all returns the output to the head layout. The KV cache is already replicated under TP, so nothing else changes. MoE, sequence parallelism and the cache layout stay as they are. To balance causal attention, each request is cut into 16 pieces and GPU r takes pieces r and 15 − r, the same balanced split as prefill context parallelism. Both all-to-alls write directly into peers’ symmetric memory.
Figure 4: Token-sharded prefill, before and after, on four GPUs. Split by head, every GPU repeats the indexer and pads its head tile; after an all-to-all, each GPU scores only a balanced slice of the tokens with every head. Step through it.
The result on 8× GB300 with TP8, real weights and one request at a time:
| Prompt tokens | TTFT, head-sharded | TTFT, token-sharded | Speedup |
|---|---|---|---|
| 27K | 0.82 s | 0.71 s | 1.16× |
| 108K | 3.72 s | 2.89 s | 1.29× |
| 340K | 15.89 s | 9.42 s | 1.69× |
Short appends on a long cached prefix, common in agentic turns, benefit too. For a 2K-token append on a 150K-token cached prefix, the indexer and sparse MLA kernels take about 15 ms per GPU instead of 56 ms (8× GB300, FlashInfer sparse MLA). The cost is the exchange buffers, about 0.8 GiB of KV cache per GPU.
Long context: decode context parallelism with a sparse indexer #
TP replicates the MLA latent on every GPU, so adding GPUs does not add KV capacity. Decode context parallelism (DCP) shards the cache by token instead: GPU r of W stores every W-th token. For dense MLA, each GPU attends over its shard and the partial outputs are combined by their log-sum-exp (LSE).
DSA adds one step: each GPU can rank only its own 1/W of the tokens, so the global top-2048 has to be merged across GPUs in every layer before attention starts.
An exact top-k merge with 10× less traffic
The generic merge sends every GPU’s full local top-2048, as FP32 score and token id pairs, to every other GPU. At DCP8 on 8× GB300, with 64 AgentX-like requests of 5 query tokens each, it takes 258 µs of a 671 µs attention module.
vLLM sends far less. Each GPU picks only its best 1.5 × 2048 / W candidates straight from the BF16 scores and packs each with its score into 32 bits, about 10× fewer bytes. Every GPU then radix-selects the same global top-2048 from the gathered candidates. The merge stays exact: if a GPU holds more of the top-2048 than it sent, the select detects it and that GPU resends its full local top-k for that row.
Figure 5: One DSA attention module at DCP8 on 8× GB300 with 5 query tokens per request. The exact one-round merge is 1.27–1.51× faster than the generic merge. Hover a bar for its value.
From batch 32 up, and at 128K context, K3.1 Pro’s sparse attention under DCP8 is now faster than K3’s dense MLA.
The partial attention outputs are then exchanged and combined by LSE with two CuTe DSL kernels on symmetric memory: up to 3.7× faster than the existing CUDA implementation (320 BF16 shapes on 8× GB300), with bit-identical output. This helps every MLA model under DCP.
DCP prefill without a merge
A plain DCP prefill would pay the merge and the LSE combine for every prompt token, and repeat the attention work on every GPU. When the DCP group is the TP group, vLLM instead gathers the step’s context, both MLA latent and index keys, from all GPUs into a symmetric-memory workspace (about 0.7 GB per 1M tokens per GPU with FP8). It then runs the token-sharded prefill from the previous section over it. No merge or combine remains, and the output matches token-sharded prefill without DCP bit for bit.
| One DSA layer, DCP8 on 8× GB300, AgentX mixed 16K-token step | DCP prefill | Gathered-context prefill |
|---|---|---|
| FP8, FlashInfer trtllm-gen | 24.8 ms | 5.8 ms (4.3×) |
| FP8, FlashInfer PrimsTS | 24.8 ms | 5.7 ms (4.3×) |
| NVFP4, FlashInfer PrimsTS | 25.7 ms | 6.2 ms (4.1×) |
End to end, a single request prefills as fast under DCP8 as under TP8: 1.10, 4.40 and 14.61 s at 27K, 108K and 340K tokens, vs 1.10, 4.39 and 14.50 s (8× GB300, FP8 KV).
When to use DCP
DCP trades per-step latency for KV capacity: DCP8 cuts the per-GPU MLA and index-key caches by eight, and KDA state is already split by head. When the context fits, plain TP decodes faster at low concurrency; DCP pays off when KV is the limit, as with very long contexts or many long-lived agentic sessions. On 8× GB300, DCP8 holds 21.8M tokens of FP8 KV vs 3.04M for TP8 (32.1M vs 4.49M with an NVFP4 latent). At concurrency 1, TPOT is 9.35 ms with TP8 and 11.48 ms with DCP8 at 32K context (9.45 vs 11.58 ms at 128K); the fused decode path brings DCP8 to 5.96 and 6.04 ms.
Eight GPUs at TP8 also serve full 1M-token requests: on 8× GB300 with FP8 KV the cache holds 3.04M tokens, 2.9 concurrent 1M-token requests.
Low-latency decode: a fused DSA decode path #
For Kimi K3, vLLM ships an opt-in decode path that replaces almost every operator with CuTe DSL kernels tuned for small batches. It fuses communication into compute, reworks the MoE schedule and prefetches weights into L2 (see Hyper Optimizing Kimi K3 in vLLM). For K3.1 Pro, this path now covers the DSA layers. A DSA layer runs as:
- one packed input projection for the query, KV latent, output gate and the indexer’s key and head weights;
- a fused indexer front: the query norm, the indexer’s query projection, its FP8 or MXFP4 quantization with the head weights folded in, and the new token’s index key, in one kernel;
- the MLA query (the up-projection and absorbed key projection, writing the new token’s latent into the KV cache), overlapped with the indexer instead of running before it;
- a fused indexer: a key-balanced scorer and an exact top-k select in two kernels that emit the physical KV rows the attention reads;
- sparse attention: FlashInfer sparse MLA (trtllm-gen or PrimsTS, picked per batch size) or FlashMLA sparse; under DCP at batch 1–8, a CuTe DSL kernel built for the shape (11.9 µs vs 21.4 µs per layer at batch 1);
- the value projection and sigmoid gate in one kernel, followed by K3’s output projection with its fused all-reduce.
Under DCP, the query is published to every GPU, each GPU’s top-k candidates are exchanged and merged exactly, and the partial outputs are combined by LSE, all through NVLink mailboxes inside the kernels, with no NCCL collective on the path.
Figure 6: Time per output token on 8× GB300 with TP8 and a 1K-token input. The fused decode path is 1.63× faster at batch 1 and still 1.18× at batch 64. Hover a bar for its value.
With DSpark speculative decoding (7 draft tokens), the fused path lowers time per token at concurrency 1 from 2.86 ms to 2.14 ms. GSM8K accuracy is unchanged. Under DCP8, the fused path cuts TPOT at concurrency 1 from 11.48 ms to 5.96 ms at 32K context and from 11.58 ms to 6.04 ms at 128K.
Speculative decoding with DSpark #
We trained a DSpark speculator for K3.1 Pro and released a preview checkpoint: Inferact/Kimi-K3.1-Pro-DSpark-preview. With realistic sampling (temperature 1.0, top-p 0.95), the preview already reaches a mean acceptance length of 4.597 on GSM8K with 7 draft tokens. At batch size 1 with TP8 on 8× GB300, that is a 3.0× decode speedup. The speculator also keeps its accuracy across K3.1 Pro’s full context window of up to 1M tokens, measured on the long-context BEAM benchmark.
This checkpoint is only a preview. We are training a more accurate and more efficient speculator right now and will release it soon.
Agentic serving: AgentX #
vLLM delivers strong performance on agentic serving workloads. We benchmark our day-0 image with the SemiAnalysis AgentX benchmark on NVIDIA GB300 and NVIDIA Rubin.
NVIDIA GB300
At the latency-sensitive end of the frontier, Kimi K3.1 Pro holds 375.9 tokens per second per user. Further along the frontier, it keeps each user above 100 tokens per second while serving 19.2K total tokens per second per GPU. At the throughput end, it reaches 37.8K tokens per second per GPU, about 26× the most interactive point, while each user still sees about 54 tokens per second.
Figure 7: AgentX Pareto frontier for Kimi K3.1 Pro on NVIDIA GB300, with 4 speculative tokens and a synthetic acceptance length of 3.36 tokens. Hover a point for its configuration and metrics.
NVIDIA Rubin
On NVIDIA Rubin, Kimi K3.1 Pro holds 815.6 tokens per second per user at the latency-sensitive end of the frontier. Further along the frontier, it keeps each user above 100 tokens per second while serving 46.0K total tokens per second per GPU. At the throughput end, it reaches 68.7K tokens per second per GPU, about 51× the most interactive point, while each user still sees about 75 tokens per second.
Figure 8: AgentX Pareto frontier for Kimi K3.1 Pro on NVIDIA Rubin. Points up to concurrency 64 use 7 speculative tokens with a synthetic acceptance length of 3.84 tokens; higher concurrencies use 3 speculative tokens with 3.0. Hover a point for its configuration and metrics.
Zero-mismatch RL on day 0 #
For the first time, vLLM’s day-0 support ships RL with zero mismatch: LoRA RL with verl, vLLM as the rollout engine and Megatron as the trainer, where the trainer’s log-probs equal the rollout’s bit for bit. Instead of bounding the gap with importance sampling or routing replay, vLLM eliminates it at the kernel level: the trainer’s forward pass reproduces the rollout bit for bit, which takes RL training stability to a new level:
- LoRA RL on the whole model. Rank-32 LoRA covers every attention projection and all 896 experts, and both the base model and the adapter are bit-identical between vLLM and Megatron.
- Batch-invariant kernels, natively. With
VLLM_BATCH_INVARIANT=1a token’s log-prob does not depend on its batch, and more and more of vLLM’s high-performance kernels are batch-invariant by design, such as the DeepSelect top-k.
Figure 9: Same tokens, same kernels, same log-probs. Switch to the default rollout to compare.
Over 100 steps on the full model, the rollout/trainer KL is exactly 0.0 at every step while AIME mean@8 rises from 0.451 to 0.664 (Figure 10).
Figure 10: A 100-step LoRA RL run on the full model.
Performance optimization of the batch-invariant path is ongoing. The source code is in Inferact/zero-mismatch-rl.
AMD #
On AMD Instinct MI355X, vLLM combines MXFP4 indexing with sparse attention and decode context parallelism. The optimizations reduce memory requirements for long contexts and improve prefill and decode efficiency.
Quantization: Native MXFP4 indexer
The indexer scans the context before selecting 2,048 tokens for 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 FP32 accumulation and weighted reduction across heads. Keys are reused across query rows during speculative verification 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.
| 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 |
Setup: MI355X, TP8/DCP8, concurrency 16, 64 prompts, prefix caching disabled.
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.
| 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 |
Setup: MI355X, TP8/DCP8, concurrency 16, 64 prompts, ISL/OSL 128K/100
Decode: Fused Low-latency decode
For small speculative verification batches, DSA indexing and sparse attention run separately, followed by a persistent monokernel that fuses attention output projection, AttnRes and the full latent MoE computation, including tensor-parallel all-reduces, into a single launch. This reduces kernel-launch overhead and overlaps MoE weight with preceding reduction and residual processing. The specialized path supports up to eight verification tokens per step, including a single request with seven draft tokens.
| Metrics | Normal Path | Specialized decode | Improvement |
|---|---|---|---|
| mean TPOT | 6.02 ms | 4.40 ms | 27.0% |
Setup: TP8/DCP1, FP8 indexer, 8K/1K, c1, draft tokens 7, synthetic AL 3.84
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
TPU #
vLLM supports Kimi K3.1 Pro on day 0 on TPU through vllm-torchtpu, validated on a 16-chip slice of TPU v7x. The layout is TP32 with EP32 for the MoE, with DCP that shards the MLA latent and index keys by token for long contexts. DSA’s work is split between the TPU’s two kinds of compute units: TensorCore scores the FP8 index keys, then SparseCore selects the top 2,048 tokens and gathers their latents into a dense buffer for TensorCore’s MXU to run attention. As on GPU, the DSA and KDA state share one block-indexed KV pool, which keeps prefix caching working across the hybrid cache.
Alibaba T-HEAD #
The Alibaba T-HEAD team added support for running Kimi K3.1 Pro on the Alibaba T-HEAD Zhenwu PPU (both M890 and 810E) through the vllm-sail plugin. The validated configurations are as follows:
- 810E: Two nodes, TP32 / EP32, using the W4A16 MoE backend
- M890: Four nodes, TP8 / DP4 / EP32, using DeepEP
This marks the first time the Zhenwu PPU offers Day-0 support for a mainstream open-source large model, with continuous support coming soon.
Acknowledgements #
We extend our sincere gratitude to the Moonshot team for open-sourcing Kimi K3.1 Pro. We also thank Inferact, AMD, and Alibaba T-Head for delivering day-0 support for the model in vLLM, and NVIDIA for their help.