Kimi K3.1 Pro: Day-0 Inference and RL Support with SGLang and Miles
SGLang and Miles provide Day-0 inference and reinforcement learning (RL) support for Kimi K3.1 Pro. Kimi K3.1 Pro is Moonshot AI's open-weight, natively multimodal agentic model, post-trained from Kimi K3 with large-scale RL. SGLang serves it from Day 0 on NVIDIA, AMD, Google TPU, AWS Trainium and Zhenwu PPU.
Get Started
See the Kimi K3.1 Pro SGLang cookbook for the launch commands for each platform and serving setup.
Model Overview
2.8T total parameters · 104B active per token · 1M-token context · text and image input
Kimi K3.1 Pro keeps the Kimi K3 architecture. Its 93 layers combine 69 KDA (Kimi Delta Attention) linear-attention layers with 24 attention layers that read the full history. The MoE layers route each token to 16 of 896 experts, plus 2 shared experts.
What's new in Kimi K3.1 Pro
- Sparse attention. The 24 full-history layers become Gated Sparse-MLA (DSA) layers. In each one, a small indexer scores every earlier token, and attention reads only the top 2,048. Long contexts cost much less attention work.
Hybrid attention. KDA layers keep a fixed-size state per request. In each DSA layer, an indexer scores earlier tokens, and attention reads the top 2,048.
SGLang Day-0 Support
AgentX performance
We ran Kimi K3.1 Pro (MXFP4) with SGLang on NVIDIA GB300 NVL72 on AgentX at 1 to 300 concurrent sessions, in three layouts:
- 1 to 4 sessions: one aggregated server with MegaKernels. TP8 on 8 GPUs, or TP16 on 16 GPUs for a single session.
- 24 to 56 sessions: prefill-decode disaggregation with Dynamo, 1P2D on 24 GPUs. One prefill worker and two decode workers, each TP8 with DCP8 and EP2.
- 160 to 300 sessions: 2P1D or 3P1D on 32 or 40 GPUs. Two or three TP8 · DCP8 · EP2 prefill workers feed one wide-EP16 decode server: TP16 with attention DP2, DCP8 and MegaMoE EP16.
DFlash2 runs at an acceptance length of 3.764 at block size 7 (24 to 56 sessions), and 3.508 at block size 5 (160 to 300 sessions).
Throughput grows from 1,563 tokens/s per GPU at one session to 33,688 at 300 sessions, while per-user speed goes from 449.0 to 69.8 tokens/s.
Launch commands for the disaggregated deployments are in the Kimi K3.1 Pro cookbook.
AgentX on NVIDIA GB300 NVL72. Each point is one deployment, labeled with its concurrency and parallelism. Disaggregated serving gives the most throughput per GPU; aggregated TP8 and TP16 serving gives the fastest per-user speed, up to 449 tok/s per user at concurrency 1.
Multi-hardware support
Day-0 support is available on NVIDIA Blackwell, NVIDIA Vera Rubin NVL72, NVIDIA Hopper, AMD Instinct (MI350X and MI355X), Google TPU (TPU7x), AWS Trainium (Trainium 2 and 3) and Zhenwu PPU (M890 and 810E). All Day-0 recipes use preview images; support has not shipped in a stock SGLang release yet. Launch recipes are in the Kimi K3.1 Pro cookbook.
AMD Instinct MI350X and MI355X. From Day 0, one node of 8 MI350X or MI355X GPUs serves Kimi K3.1 Pro with TP8, DCP8 and DFlash2, and two nodes run 1P1D prefill-decode disaggregation over Mooncake. New gfx950 kernels cover the hybrid model: one HIP kernel for KDA verify; for DSA, a HIP top-k that merges the DCP shards, a HIP sparse MLA prefill kernel, and an MXFP8 sparse MLA decode kernel ported from AITER's opus decode; and A8W4 grouped GEMMs for the MoE. At low concurrency, MegaKernels from sgl-amd-mono run decode and verify as persistent kernels.
Google TPU. Through SGLang-TorchTPU, Kimi K3.1 Pro runs on a 16-chip TPU7x (Ironwood) slice as one TP32 · EP8 engine. A single request decodes at about 133 tokens/s (8K input), and up to about 590 tokens/s with the experimental megakernel and DFlash.
AWS Trainium. For the first time, SGLang brings Day-0 support for an open-weight model to AWS Trainium. Kimi K3.1 Pro runs on Trainium 2/3 through the SGLang Trainium plugin, built on TorchNeuron, AWS's native PyTorch backend for Trainium. The plugin is currently in limited beta, so Trainium recipes for Kimi K3.1 Pro will be published when it becomes broadly available. To request access, contact us through this form.
Zhenwu PPU. Kimi K3.1 Pro marks Zhenwu PPU’s first Day-0 model support as an in-tree hardware backend in SGLang. Both M890 and 810E support disaggregated prefill and decode: M890 uses 2P1D with TP16/EP16 prefill and TP32/EP32/DP4 decode, while 810E uses 1P1D with TP16/PP4 prefill and TP32/EP32 decode.
Ultra-low latency with MegaKernels
For low-latency serving, SGLang runs each decode and verify step of Kimi K3.1 Pro as MegaKernels, developed together with NVIDIA. Within a step, each kernel launches while the previous one is still finishing, the GPUs exchange partial results over NVLink and signal each other on the GPU, and the all-reduces are fused into the compute kernels around them, so the GPUs move from layer to layer without waiting on the host. MegaKernels ship as the sgl-nvidia-mega plugin for NVIDIA GB300 NVL72 and NVIDIA Vera Rubin NVL72. Launch the server with --enable-nvidia-megakernel; on NVIDIA GB300 NVL72, SGLANG_ENABLE_NV_MEGA_RUNTIME_K31_GB_OPTIMIZATION=1 also applies SGLang's tuned settings.
The MegaKernels cover the whole step:
- Sparse attention (DSA): the indexer, the top-2,048 selection and sparse MLA attention for every token of a verify block. For long histories, the indexer work is split across the GPUs and the selections are merged exactly.
- KDA: the linear-attention layers over the verify block, and the state update for the accepted tokens.
- MoE: the router, the 16 routed MXFP4 experts and the shared experts, with the reduction across GPUs fused into the expert kernels.
- Layer boundaries: the attention output projection and its all-reduce, and the AttnRes mixing, are folded into the kernels next to them.
- DFlash2: the draft model also runs as a MegaKernel, from its layers to the draft head, and the target model accepts the drafted tokens, including sampled acceptance, inside its MegaKernel.
On NVIDIA GB300 NVL72, one AgentX session decodes at 449.0 tokens/s per user with TP16 on 16 GPUs, and at 354.7 with TP8 on 8 GPUs (P90 interactivity, as in AgentX performance above).
On NVIDIA Vera Rubin NVL72, one AgentX session decodes at 741.8 tokens/s per user with TP16 on 16 GPUs (P90 interactivity).
DFlash2 speculative decoding
SGLang runs Kimi K3.1 Pro with DFlash2 speculative decoding. A small draft model reads hidden states from six layers of Kimi K3.1 Pro and proposes 7 tokens in a single forward pass, and Kimi K3.1 Pro checks all 7 in one forward pass. The draft model trained by RadixArk has five MLA layers with a 1,024-token sliding window.
Parallelism Design
Prefill: Combination of pipeline parallelism and prefill context parallelism
To better meet time-to-first-token (TTFT) requirements in real-world serving, we combine pipeline parallelism (PP) with prefill context parallelism (PCP).
PP4 × PCP2 prefill. Each pipeline stage uses two PCP ranks to process token shards of the same microbatch. Different stages process different microbatches concurrently; the final stage also hosts the draft model.
Pipeline parallelism (PP) increases prefill concurrency by splitting model layers and their DSA KV caches across stages. Each stage stores KV only for its own layers, leaving more memory for concurrent requests. We use the chunked pipeline approach introduced for Kimi K3: prompt chunks form microbatches that move through the stages. While one stage processes a microbatch, the preceding stage starts the next. This overlap reduces the time needed to process a stream of microbatches.
Prefill context parallelism (PCP) interleaves token positions across ranks in round-robin order, giving each rank a shard of the sequence. KDA gathers the full chunk, splits attention heads across ranks, and reduces its outputs back to token shards. DSA keeps queries local while gathering new keys and KV to access the full history. AttnRes works on local shards, and MegaMoE uses AllToAll to route tokens to experts and return results.
Token layouts within a PP stage. KDA gathers the full sequence while splitting heads across ranks. AttnRes works on local token shards, and MegaMoE returns expert results to those shards after each attention layer.
Both PP and PCP work with HiCache, which keeps reusable prefixes in host memory and external storage. PCP currently partitions computation while DSA KV remains replicated within the CP group. We plan to add KV-cache sharding to PCP, removing duplicated KV storage across CP ranks.
Prefill workers also support decode context parallelism (DCP), which is better suited to workloads such as AgentX, where requests often append a small amount of new text to a long cached history, and the TTFT metric is not focused. At each DSA layer, DCP-enabled prefill workers use AllGather to reconstruct the batch's full KV-cache history before calling the attention kernels. With fewer DSA than KDA layers, this communication cost is manageable on modern NVLink systems.
Decode: DSA with decode context parallelism
With DSA, each query only looks at the top 2,048 previous tokens. However, the KV cache and the index cache still grow with context length. Our DCP implementation shards both caches across GPUs the same way: token p is stored on GPU p mod 8, so each GPU holds an eighth of each cache.
Because the index cache is sharded, each GPU's indexer can only score its own tokens. Each GPU then picks its local top 2,048. No token is lost at this step: any token in the global top 2,048 is also in the top 2,048 of its own GPU. To merge cheaply, each GPU sends only its best 384 candidates to all other GPUs with one NVLink multicast write. Every GPU now has the same 8 × 384 candidates and takes the 2,048th-best score among them as the cutoff. A GPU whose 384th candidate falls below the cutoff can't be hiding any winners. A GPU whose 384 candidates all beat the cutoff might be, so every GPU reads its full local top 2,048 directly over NVLink and runs the selection again. The result matches a single-GPU top-k exactly. Every GPU ends up with the same selection without another round of communication. Each GPU then attends to the winners it owns, and the partial outputs are combined.
DCP for DSA, on a toy example. Both caches are sharded by token. Each GPU shares only its best candidates, and a GPU that may hold more winners has its full list read, so the merged top-k matches a single-GPU top-k.
Miles Day-0 Support
Miles trains Kimi K3.1 Pro with reinforcement learning (RL) from Day 0 on a small cluster of 32 H200 GPUs.
Recipe. As for Kimi K3, Miles trains a BF16 LoRA adapter on top of the frozen MXFP4 base. The adapter stays a separate BF16 term in both the trainer and SGLang, and is never merged into the base.
Training parallelism.
- The trainer combines PP, TP, CP and EP.
- CP also covers the new DSA layers: each GPU keeps its own queries and gathers the latent KV and the indexer keys from its CP group.
DP attention for LoRA rollout. With TP32 attention, every GPU takes a slice of the heads for all running requests. MLA keeps one compressed KV per token instead of one per head, so tensor parallelism cannot split it: all 32 GPUs store the same MLA KV and indexer keys, and every attention layer ends with an all-reduce across four nodes. With DP attention, the 32 GPUs form four groups of 8, one per node, and each group runs attention only for its own quarter of the requests. The MoE layers still span all 32 GPUs. This changes two things:
- Less duplicated KV. Each GPU stores the KV of 64 requests instead of 256. For the same 589,824 tokens in total, the KV cache takes 4.2 GiB per GPU instead of 16.9 GiB.
- Less cross-node traffic. The attention all-reduce stays on NVLink inside one node. Only the MoE layers still exchange tokens across nodes.
On 32 H200 with 256 running requests, a decode step takes 166.6 ms, against 202.6 ms with TP32 attention. At 64 running requests DP attention is slower (136 ms against 110.5 ms), so it pays off when rollout keeps the engine full.
To make this work with LoRA, SGLang now allows attention TP above 1 under DP attention. A LoRA layer must know whether its input rows are the requests of its own group or the requests of all groups gathered together. In Kimi K3.1 Pro, the attention projections see the rows of their group, and the dense layer, the shared expert and the routed experts see all rows. The model declares this layout, so SGLang no longer requires attention TP 1.
Consistency. Rollout routing replay makes the trainer use the experts that SGLang chose for each token. Under DP attention, the GPUs of a group share one indexer selection, so they read the same 2,048 tokens. Every step logs the train/rollout KL at the same policy version.
Checked run. The run trains on DAPO-Math-17k prompts, 32 prompts with 16 samples each per step, with rank-32 LoRA and a learning rate of 1e-5. The task asks the model to finish its reasoning and answer within a fixed budget of 3,072 tokens.
| 32× H200 | 16× NVIDIA Blackwell Ultra | |
|---|---|---|
| Trainer | PP4 × TP4 × CP2, EP8 | PP2 × TP4 × CP2, EP8 |
| Rollout | MoE EP32, attention DP4 × TP8 | MoE EP16, attention DP2 × TP8 |
| Train/rollout KL | 0.44 to 0.82 × 10⁻³ | 0.36 to 0.42 × 10⁻³ |
Training curves. The reward rises from 0.63 over the first 10 steps to 0.77 over the last 10. The train/rollout KL, the log-probability gap and the gradient norm stay flat.
We evaluated the adapters with the budget, prompt template and reward of training, at temperature 1.0: 16 samples per AIME problem and 4 per MATH-500 problem. After 84 steps, AIME 2024 accuracy rises from 44.6 to 70.0.
Accuracy before and after RL. Mean accuracy over samples for the base model, and for the adapters after 39 and 84 steps.
Blackwell. The same recipe also runs on 16 NVIDIA Blackwell Ultra GPUs, half as many as on H200, with the layouts in the table above. A 15-step run with 4,096-token responses checks that training and rollout stay consistent on Blackwell. More experiment results are ongoing.
Training curves on 16 NVIDIA Blackwell Ultra GPUs. The train/rollout KL and the log-probability gap stay flat over all 15 steps. The run is too short to show a reward trend.
Acknowledgments
Kimi K3.1 Pro Day-0 support was a collaboration between the SGLang & Miles team and the Moonshot AI team, together with AMD (Zijie Chen, Wun-guo Huang, Xinyi Song, Zhaoyi Li, Hai Xiao and SkyOpt agents from Berkeley SkyDiscover), AWS and Zhenwu PPU. We are also grateful to NVIDIA for their support.