Kimi K3.1 Pro, the latest model from Moonshot AI, builds on Kimi K3 with one architectural change: each of K3's 24 Gated MLA layers gains a lightning indexer that picks the 2,048 most relevant earlier tokens to attend to. The rest of K3 stays the same, and long contexts get cheaper to serve. Moonshot’s announcement calls K3.1 its most capable open-weight model yet, with a 1M-token context and clear gains in long-horizon coding, 3D creation and financial research. It is available through the Kimi API and as open weights.
Working closely with the Moonshot AI team, TokenSpeed enabled Day 0 support for Kimi K3.1 Pro on AMD MI350X, NVIDIA Blackwell and NVIDIA Vera Rubin, including EAGLE3 speculative decoding on MI350X and decode context parallelism for long contexts. Day 0 support covers text and image input and is available to try out today; see the TokenSpeed deployment guide to get started.
This post first walks through what changed from Kimi K3 to K3.1, then how TokenSpeed plugs sparse attention into the K3 stack, the AMD work behind a cheap indexer, long-context serving across GPUs, end-to-end performance on MI350X and GB300, and bringing K3.1 to NVIDIA Vera Rubin. We continue to optimize every layer of the stack—pushing tokens toward the speed of light.
Highlights
- Sparse attention on the existing stack. K3.1's sparse layers run on TokenSpeed's sparse-attention backend as the full-attention half of K3's KDA hybrid, with no vendor-specific model code.
- One more plane in the flat KV cache. Index keys sit next to the latent rows they describe, and we store them in FP8, so the cache grows by 21%.
- Indexer on AMD. FP8 scoring, multi-workgroup top-k and a prefill top-k split across tensor-parallel (TP) ranks keep the indexer's cost small on MI350X. With 16 requests, K3.1's decode attention per layer beats dense K3's from 32K tokens on and is 2.4× faster at 256K and beyond.
- Long contexts across GPUs. KV parallelism and decode context parallelism split each request's history, index keys included, across GPUs, raising the KV capacity on 8 GB300 GPUs from 2.43M to 16.35M tokens, a 6.72× increase with KVP8 + DCP8.
- Agentic serving on MI350X. On multi-turn SWE-smith coding sessions on 8 MI350X GPUs, one session gets 248 output tokens/s, and each of 16 concurrent sessions still gets 48 tokens/s.
- Day 0 on Vera Rubin. K3.1 runs end to end on 8 Vera Rubin GPUs as a single instance, with KVP8/DCP8 attention and DSpark speculative decoding.
What's new in Kimi K3.1 Pro
| Kimi K3 | Kimi K3.1 Pro | |
|---|---|---|
| Layer stack | 93 layers: 69 Kimi Delta Attention (KDA) + 24 Gated MLA | Same |
| Gated MLA attention | Dense: every earlier token | Sparse: 2,048 earlier tokens per query, chosen per layer by a lightning indexer |
| New weights | — | One indexer per MLA layer: 32 heads × 128, BF16 |
| KV cache per token, per MLA layer | One 576-value latent row | Latent row + one 128-dim index key |
| KDA, Stable LatentMoE, AttnRes | As in the K3 post | Unchanged |
Kimi K3 interleaves Kimi Delta Attention (linear attention, 69 of its 93 layers) with Gated MLA (full attention, the other 24); see our Kimi K3 Day 0 post. KDA keeps a fixed-size state, so its cost per token doesn't depend on how long the context is. The Gated MLA layers are where long contexts get expensive: in dense K3, every decode step reads every earlier token's 576-value latent row in each of them.
K3.1 adds a lightning indexer to each Gated MLA layer (Figure 1). For each query token, the indexer scores all earlier tokens with 32 lightweight heads against one 128-dim key per token, and keeps the top 2,048. All 96 attention heads of the layer then attend only to those tokens. Each layer makes its own selection.
Figure 1. K3.1 adds a lightning indexer to each of the 24 Gated MLA layers. Left: the layer stack. Right: inside one Gated MLA layer, the index key is written to the cache (①), index queries from the query latent score all cached keys (②), the top 2,048 positions are selected (③), and sparse MLA attends only to those rows (④).
The indexer still visits every earlier token, but it reads a 132-byte FP8 key instead of a 576-byte FP8 latent row, and full MLA runs on just 2,048 rows. At a 128K-token context, that cuts the attention data each MLA layer reads per decode token, on every TP rank, by about 4×: from 75.5 MB to 18.5 MB, plus about 1 MB of index scores. Speculative decoding verifies several rows per request in each step, which trims the saving somewhat. The absolute saving grows with the context, while the ratio levels off just above 4×. The price is an extra key per token in the cache and the indexer's scan over all of them. The rest of this post is about keeping that price small.
Plugging sparse attention into the K3 stack
In TokenSpeed, each Gated MLA layer of K3.1 runs in four steps, numbered as in Figure 1:
- Write. One fused kernel normalizes the new token's index key, quantizes it to FP8 with a single scale, and writes it into the cache.
- Score. 32 index-query heads, projected from the layer's query latent, are scored against every cached index key.
- Select. A radix top-k keeps the 2,048 highest-scoring positions for each query.
- Attend. Sparse MLA attends over only those 2,048 latent rows.
TokenSpeed already had a sparse-attention backend with kernels for each step on both vendors, plus a portable Triton fallback.
| Step | NVIDIA Blackwell | AMD MI350X |
|---|---|---|
| Write | Triton: one kernel for the latent row, one for the index key | Triton: one kernel for the latent row, one for the index key |
| Score | DeepGEMM, FP8 index queries and keys | Gluon, FP8 index queries and keys |
| Select | Decode and prefill (DCP): persistent CUDA radix top-k for local candidate selection, followed by global top-k merge. Prefill (w/o DCP): TensorRT-LLM top-k | Gluon radix top-k, one or many workgroups per row depending on shape |
| Attend | FlashInfer TRTLLM-GEN sparse MLA | Gluon FP8 sparse MLA |
One more plane in the flat KV cache
In TokenSpeed's flat KV layout, all of a layer's state for a run of tokens lives in one cache block under one page ID. K3.1 adds one index-key plane per MLA layer, next to that layer's latent rows, so index keys are allocated, freed and sharded together with the rows they describe, and packing and block granularity stay as in K3. Every full-attention cache block holds index keys next to its latent rows, so the overhead applies to the whole history pool, whether or not a request ever needs selection.
Figure 2. K3.1 adds one index-key plane per MLA layer to the cache block. Top: cache block size at TP8 for K3 and for K3.1 with FP8 or BF16 index keys. Bottom: one MLA layer in a 128-token cache page, and one index key: 128 FP8 values plus one FP32 scale.
A key takes 132 bytes per token per layer in FP8, instead of 256 in BF16, while the indexer’s projection weights stay in BF16. That keeps the cache-block overhead over K3 at 21% instead of 41% at TP8. The accuracy cost is small. Compared with an FP32 indexer, FP8 keys change about 6% of the selected tokens. Compared with BF16 keys, over two runs each, AIME26 is unchanged and MRCR 2-needle is about 0.003 lower, close to the run-to-run spread.
Making the indexer cheap on AMD MI350X
Sparse attention moves the cost of long contexts from attention into the indexer, which still visits every cached token. On MI350X, the challenge is the shape of that work: with EAGLE-3.1, decode hands each layer just 4 verify rows per request, while the scan covers the request's entire history. Figure 3 shows the result.
Figure 3. Decode attention per MLA layer on one MI350X at TP8, Kimi K3 (dense MLA) vs. K3.1 (indexer + sparse MLA), 16 requests with 4 verify rows each.
Sparse MLA holds steady at 35–41 µs per layer from 16K to 1M tokens, so K3.1's cost tracks the indexer. Like dense MLA, the indexer grows linearly with the context, but at about 2.4× less time per token. At 16K tokens, sparse MLA already costs most of what dense MLA does, so adding the indexer leaves K3.1 slower than K3 (0.80×). With 16 requests, K3.1 pulls ahead from 32K tokens and reaches 2.4× at 256K and beyond.
To get there, both indexer kernels are organized around the request rather than the row. The scorer keeps index queries and keys in FP8 and scores all of a request's verify rows in one program, so each key is loaded once rather than once per row. The top-k spreads each row across many workgroups that share a histogram of score prefixes; with one workgroup per row, a single request would occupy just 4 of the GPU's 256 compute units. Both kernels also size their work by each request's length, not the server's maximum context.
Prefill flips the problem. There are thousands of rows, but K3.1 replicates the indexer on every TP rank, so all eight GPUs would select the same rows. Instead, each GPU selects for its own eighth of the chunk, and an all-gather assembles the result. The speedup of this selection step grows with the context: 4.8× for a full 8,192-row chunk at a 50K-token context and 7.2× at 250K. Below 384 rows, the all-gather costs more than it saves, so every GPU selects all rows.
Sparse attention then runs in FP8 with a kernel shaped for occupancy. Its per-tile work is a short dependent chain, so it keeps the query in registers and fits three workgroups per compute unit, bringing an 8,192-row call from 2.01 to 1.37 ms.
Finally, overlap had to be checked in the real server, where the side stream for Iris (AMD's Triton-based multi-GPU programming framework), which hides prefill communication, shared a hardware queue with the main stream; moving it to a high-priority stream cut time to first token at concurrency 16 by 8%.
Long contexts with decode context parallelism
With MLA, every attention-TP rank holds a full copy of the latent cache, so TP8 adds compute but no KV capacity for a single request. K3.1 has a second reason to be sharded: attention is capped at 2,048 rows, but the indexer still scans every key, so splitting the keys across GPUs splits the one cost that keeps growing with the context.
KV parallelism (--kv-parallel-size N) splits each request's MLA history across N attention-TP ranks. Latent rows and their index keys are sharded together, and the KDA state stays replicated. Decode context parallelism (--decode-context-parallel-size N) then attends to each piece where it lives (Figure 4): each rank scores its own index keys, the ranks exchange candidates to agree on one global top-2,048, each rank attends to the selected rows it holds, and the partial outputs are merged by log-sum-exp. The price is that exchange: at KVP8, each rank sends about 119 KiB per verify row per MLA layer, covering its 2,048 candidates, its query heads and its partial outputs. On 8 GB300 GPUs, the net cost is about 8% of single-user decode speed (see below).
TokenSpeed's sparse-attention backend already had what DCP needs: owner-masked cache writes, the candidate exchange and the log-sum-exp merge. Making DCP fast for K3.1 took more work, which we describe with the GB300 results below.
The merge needs each rank's log-sum-exp. FlashInfer's TRTLLM-GEN sparse MLA kernel returns it in TokenSpeed's pinned release, so the final build runs DCP on that kernel instead of the Triton fallback (see the GB300 section below).
Figure 4. K3.1 under KV and decode context parallelism. Each rank holds 1/N of the MLA history, latent rows and index keys together. The ranks score locally, agree on one global top-2,048, attend locally, and merge by log-sum-exp.
End-to-end performance
AMD MI350X on long agentic conversations
Agentic coding workloads keep returning to a long conversation while the user waits for the next tool call. To measure this, we replayed multi-turn SWE-smith coding sessions. Each conversation opens with a ~50K-token prompt, adds ~800 tokens per turn over 10–15 turns, and generates 500 tokens per reply. K3.1 runs on eight MI350X GPUs at TP8, with EAGLE3 speculative decoding using Kimi K3's EAGLE-3.1 draft (3 steps, 4 draft tokens), FP8 KV and GPU prefix caching.
With one active session, K3.1 generates 248 output tokens/s per user (4.03 ms per output token), with a 0.53-second average time to first token. At concurrency 16, each session still gets 48 tokens/s while the server processes 8.8K tokens per second per GPU. At that point, the sessions' combined history, about 1M tokens, outgrows the GPU prefix cache, so this run adds host-memory KV cache at 8× the GPU cache.
Figure 5. Kimi K3.1 Pro agentic serving on eight MI350X GPUs. Each point is a single run at one concurrency setting; concurrency 16 adds a host-memory KV cache.
NVIDIA GB300 with KV and decode context parallelism
On NVIDIA, we focus on what decode context parallelism adds for long contexts. K3.1 runs with DeepGEMM index scoring and FlashInfer TRTLLM-GEN sparse MLA. We replay AgentX agentic coding sessions for one hour per concurrency point, with DSpark speculative decoding at 7 steps, on one TP8 instance across 8 GB300 GPUs.
| Setup | KV capacity (tokens) | Request throughput (req/s) | Per-user decode, relative to KVP1 |
|---|---|---|---|
| K3.1, KVP1 (8 GB300 GPUs, TP8) | 2.43M | 0.31 (concurrency 8) | 1.00× (concurrency 1) |
| K3.1, KVP8 + DCP8 (8 GB300 GPUs, TP8) | 16.35M | 1.37 (concurrency 64) | 0.92× (concurrency 1) |
KVP8 holds 6.72× the KV of KVP1 and serves 4.40× the requests per second (KVP1 was run up to concurrency 8). With one user, KVP8's median per-user decode speed is 8% lower than KVP1's; at concurrency 8, it is already 11% faster.
Figure 6. Total throughput per GPU (input tokens, including prefix-cache hits, plus output tokens) against per-user decode speed at the 90th-percentile inter-token latency on AgentX for TP8 on 8 GB300 GPUs with KVP8 or KVP1, and for PD disaggregation with one prefill and 1–3 decode instances on 16–32 GPUs, with the Pareto frontier across all three. Labels give the concurrency (cc) and the PD layout; prefix-cache hits make up 95–97% of input tokens in the TP8 runs.
Making DCP fast. Before this work, KVP8 served 3.5K tokens per second per GPU at concurrency 64; the final build serves 20.1K, 5.7× as much. Two changes are worth a closer look:
- Decode path: each rank now selects candidates only from the part of each row that covers the request, exchanges candidates once per layer instead of dozens of times, and attends only to the selected rows it owns, about 1/8 of the work at KVP8. In a separate 16-GPU A/B run with prefill and decode on separate instances, this raised output throughput per GPU 2.4× at concurrency 16.
- Attention kernel: switching DCP from the Triton fallback to TRTLLM-GEN made the kernel 7× faster. Sparse attention per 16K-token prefill pass then fell from 335 to 163 ms, and throughput rose 7.7%.
Bringing Kimi K3.1 Pro to NVIDIA Vera Rubin
Our Kimi K3 work brought TokenSpeed to NVIDIA Vera Rubin with TokenSpeed MLA 2, a dense MLA decode kernel tuned for the K3 family's 96 attention heads. TokenSpeed now supports K3.1 on Vera Rubin at Day 0 as well: it runs end to end on 8 Vera Rubin GPUs as a single instance, with attention at TP8 using KV parallelism and decode context parallelism across all eight ranks, MoE at TP8, and DSpark speculative decoding.
Most of the port was enabling existing Blackwell paths on Rubin rather than writing new kernels. We widened the architecture gates of the MXFP4 MoE kernels and the MoE all-reduce fusion, moved KDA prefill to FlashInfer’s CuTeDSL kernel, which compiles for Rubin and matches the reference, and made KDA and sparse-attention prefill capturable in CUDA graphs, including under KV parallelism. Performance tuning on Rubin is just starting; we will share numbers in a follow-up.
Looking forward
Kimi K3.1 Pro is another test of TokenSpeed's modular design: a new attention variant landed on the existing flat cache and sparse-attention backend with no vendor-specific model code. Next, we are bringing wide expert parallelism to K3.1, building a fast decode-context-parallel path on MI350X, and continuing to tune the K3 family on Vera Rubin. Stay tuned.
Acknowledgments
We thank Moonshot AI for the partnership and early access to Kimi K3.1 Pro, AMD for CDNA4 compiler and Gluon support, NVIDIA DevTech for Blackwell kernel support, and the FlashInfer community.