vLLM Day-0 Support for Kimi K3.1 Pro: Sparse Attention, DCP, and Zero-Mismatch RL
TL;DR
vLLM shipped day-0 support for Kimi K3.1 Pro, Moonshot AI's 2.8T-parameter (104B active) open-weight multimodal MoE with a 1M-token context, on NVIDIA GPUs, AMD GPUs, Google TPUs, and Alibaba T-HEAD Zhenwu PPUs. The release matters because K3.1 Pro adds Sparse Attention (DSA) to all 24 full-attention layers of an otherwise hybrid KDA + MLA model, and vLLM had to redesign the KV cache, prefill sharding, decode parallelism, and low-latency decode path to serve that hybrid sparse architecture efficiently.
What Kimi K3.1 Pro changes
Kimi K3.1 Pro keeps the architecture of Kimi K3 but adds a sparse-attention indexer to every full-attention layer. The model remains a 93-layer hybrid with 69 KDA layers and 24 full-attention layers, 896 MoE experts with 16 active, MXFP4 MoE weights, and a 1M-token context window.
| 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 |
The central technical change is attention. Each sparse layer keeps a small indexer next to its MLA attention. The indexer stores one FP8 key per token, 132 bytes including its scale, and scores every earlier token for each query using 32 lightweight heads. It then keeps the top 2,048 tokens, and MLA attention runs exactly over only those 2,048 tokens. vLLM notes that DSA was introduced in DeepSeek-V3.2 and points to its earlier DeepSeek-V3.2 post for a detailed indexer explanation.
Sparse attention cuts long-context decode bandwidth
The payoff from DSA 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 it reads only the 132-byte index keys during that scan, 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.
This is the core serving argument for K3.1 Pro: the model keeps a 1M-token context but reduces the per-layer memory traffic of a decode step at long context to roughly a quarter of dense MLA's traffic.
A packed KV cache for three kinds of state
vLLM packs each cache block's MLA latent, indexer keys, and KDA state together rather than giving every layer pages of the same size. Each DSA layer stores two per-token caches of different widths: the MLA latent at 576 bytes in FP8 or 352 bytes in NVFP4, and the index keys at 132 bytes in FP8 or 68 bytes in MXFP4. Neither width is a multiple of the other, and the 69 KDA layers each add a recurrent state per request that is much larger than one layer's page of tokens. A conventional same-size-per-layer cache would have to pad the small pages up to the largest one.
vLLM instead makes one block of the cache a single contiguous row holding 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 running 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 enables three things:
- Any sparse MLA kernel and KV format. On Blackwell the default is FlashInfer's sparse MLA with an FP8 latent. FlashInfer's PrimsTS kernels also serve an NVFP4 latent.
- Disaggregated serving. Each block is one contiguous range of memory, so a prefill/decode 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 mixed-page-size problem.
KV cache footprint per token
Under tensor parallelism, both the MLA latent and the index keys 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 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.
Token-sharded prefill improves time to first token
For large prefill steps, vLLM shards attention by token rather than by head. Under TP8, each GPU owns 12 of K3.1 Pro's 96 attention heads, which is a poor fit for sparse MLA prefill: 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, so every GPU scores every token and its quadratic cost is repeated eight times.
vLLM instead uses an all-to-all to move 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 used in prefill context parallelism. Both all-to-alls write directly into peers' symmetric memory.
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 on 8× GB300 with FlashInfer sparse MLA. The cost is the exchange buffers, about 0.8 GiB of KV cache per GPU.
Decode context parallelism with a sparse indexer
Tensor parallelism 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 about 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.
The exact one-round merge is 1.27–1.51× faster than the generic merge. From batch 32 up, and at 128K context, K3.1 Pro's sparse attention under DCP8 is 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 across 320 BF16 shapes on 8× GB300, with bit-identical output. vLLM notes that 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 described above 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, versus 1.10, 4.39, and 14.50 s on 8× GB300 with 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 versus 3.04M for TP8, or 32.1M versus 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, and 9.45 versus 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, or 2.9 concurrent 1M-token requests.
Fused low-latency decode for DSA layers
For Kimi K3, vLLM shipped 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. 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, including the up-projection and absorbed key projection and 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, at 11.9 µs versus 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.
The fused decode path is 1.63× faster at batch 1 and still 1.18× at batch 64 on 8× GB300 with TP8 and a 1K-token input. With DSpark speculative decoding and 7 draft tokens, the fused path lowers time per token at concurrency 1 from 2.86 ms to 2.14 ms, with GSM8K accuracy 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
vLLM trained a DSpark speculator for K3.1 Pro and released a preview checkpoint, Inferact/Kimi-K3.1-Pro-DSpark-preview. With realistic sampling at temperature 1.0 and top-p 0.95, the preview 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.
vLLM describes this checkpoint as a preview only and says it is training a more accurate and more efficient speculator for release soon.
Agentic serving on NVIDIA GB300 and Rubin
vLLM benchmarked its 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. The GB300 sweep used 4 speculative tokens and a synthetic acceptance length of 3.36 tokens.
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. 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 an acceptance length of 3.0.
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 by making the trainer's forward pass reproduce the rollout bit for bit.
- 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=1, a token's log-prob does not depend on its batch, and more of vLLM's high-performance kernels are batch-invariant by design, such as the DeepSelect top-k.
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. vLLM says performance optimization of the batch-invariant path is ongoing, and the source code is in Inferact/zero-mismatch-rl.
AMD Instinct MI355X support
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.
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. In AgentX sweeps, the MXFP4-indexer configuration delivered 5.3% higher total workload throughput at concurrency 48.
Context-parallel sparse attention prefill
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.
| 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.
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 loading 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. Mean TPOT improves from 6.02 ms to 4.40 ms, a 27.0% reduction, on TP8/DCP1 with an FP8 indexer, 8K/1K, concurrency 1, 7 draft tokens, and synthetic acceptance length 3.84.
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, combining data gathering and layout conversion. 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.
Google TPU support
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 Zhenwu PPU support
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:
- 810E: two nodes, TP32 / EP32, using the W4A16 MoE backend.
- M890: four nodes, TP8 / DP4 / EP32, using DeepEP.
vLLM describes this as the first time the Zhenwu PPU offers day-0 support for a mainstream open-source large model, with continuous support coming soon.
Quick start
On NVIDIA GPUs, vLLM provides a vllm serve command that uses tensor parallelism of 8, fastsafetensors loading, the Kimi K3 tool-call and reasoning parsers, and a DSpark speculative configuration with 5 speculative tokens:
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, the command enables AITER and fast GPU safetensors loading, sets tensor parallelism of 8, limits sequences to 16, sets GPU memory utilization to 0.85, and uses the same DSpark speculative configuration:
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 the release shows about serving hybrid sparse models
Kimi K3.1 Pro is a hybrid model with three different kinds of per-token or per-request state: MLA latent, DSA index keys, and KDA recurrent state. vLLM's day-0 support shows that serving such a model efficiently requires coordinated changes across the KV cache layout, prefill sharding, decode parallelism, and low-latency kernel fusion rather than a single attention kernel.
The packed block cache is the enabling abstraction: it lets the attention kernels keep their fixed-size pages while the serving stack stores mixed-width state contiguously for RDMA transfer and prefix caching. Token-sharded prefill addresses the head-replication problem of MLA under tensor parallelism, and DCP with an exact top-k merge addresses the KV-capacity limit of TP for long contexts. The fused decode path and DSpark speculative decoding target the small-batch latency that agentic workloads depend on.
The zero-mismatch RL result is separate from inference performance but significant for training infrastructure: vLLM reports bit-identical log-probs between vLLM rollout and Megatron training for both the base model and a rank-32 LoRA adapter covering all attention projections and all 896 experts, with KL exactly 0.0 across 100 steps. That removes a source of RL instability rather than bounding it with importance sampling or routing replay.
Platform coverage
Day-0 support spans four hardware families: NVIDIA GPUs, AMD GPUs, Google TPUs, and Alibaba T-HEAD Zhenwu PPUs. The NVIDIA path includes GB300 and Rubin benchmarks on SemiAnalysis AgentX. The AMD path includes MI355X optimizations for MXFP4 indexing, context-parallel sparse prefill, fused low-latency decode, and AITER all-gather. The TPU path runs through vllm-torchtpu on a 16-chip TPU v7x slice with TP32/EP32 and DCP. The T-HEAD path runs through the vllm-sail plugin on 810E and M890 configurations.
vLLM credits Moonshot for open-sourcing Kimi K3.1 Pro, and thanks Inferact, AMD, and Alibaba T-HEAD for day-0 support, and NVIDIA for its help.