moonshotai/Kimi-K3.1-Pro

Overview

Kimi-K3.1 keeps the Kimi-K3 layout (Kimi Delta Attention, 896 routed experts with 16 active per token, a latent MoE projection and a vision encoder) and adds sparse attention (DSA, top-2048 indexer) to 24 of its MLA layers. The checkpoint ships MXFP4 routed-expert weights; everything else is BF16.

The baseline is the simplest working setup on NVIDIA: one tensor-parallel group of 8 GPUs, either two GB300 trays or one B300 node. Throughput was not tuned.

VR200 (preview) takes the same TP8 layout on two trays, in the vllm/vllm-openai:kimi-k31-cu134 image (CUDA 13.4). It uses the Blackwell flags and hasn't been tuned or verified on Rubin.

On MI355X, one node takes the same TP8 layout with the AMD profile's AITER environment, in the day-0 vllm/vllm-openai-rocm:kimi-k31 image. So far only the AgentX TP8 configuration has run on MI355X: the tuned single-node AgentX configuration for concurrency 48, with DCP8, the K3.1 DSpark draft and host KV offload.

On Zhenwu PPUs, the default MXFP4 checkpoint runs without speculative decoding on two 16-PPU 810E nodes with TP32 / EP32, or four 8-PPU M890 nodes with TP8 / DP4 / EP32. On 810E, ppu_deep_gemm_w4a16 loads the checkpoint with a W4A16 compute path. Use the PPU image specified for each hardware configuration.

On TPU, the model is served by vLLM torch-TPU, text-only: the vision tower is not loaded. The validated layout is TP32 with expert parallelism on a 2x2x4 TPU v7x slice (16 chips, 32 devices, on 4 hosts). See TPU serving.

Tensor + expert parallelism

On one B300 node, the Tensor + Expert Parallel strategy keeps attention at TP8 and runs the routed experts with expert parallelism on the MegaMoE kernel. It has the lowest long-prompt time to first token of the single-node layouts: a single 340K-token prompt prefills in about 10 s, against about 15 s at plain TP8 (8x B300, one request at a time, pre-release checkpoint).

Long context: decode context parallelism

The Decode Context Parallelism row shards the MLA latent and indexer-key caches across the 8 GPUs of the TP group instead of replicating them, so DCP8 holds about 7x the KV cache tokens of TP8 (about 18.5M vs 2.6M tokens with FP8 KV on 8x B300). Prefill is as fast as at TP8. Decode steps are 7–25% slower than TP8 at the same batch, so use DCP when KV capacity is the limit: very long contexts or many long-lived agentic sessions.

Low-latency decode

The Low-latency decode feature switches decode steps to a fused CuTe DSL path tuned for small batches; prefill and mixed steps still run the default path. On 8x B300 at TP8 it cuts the decode step by 1.6–1.7x at batch 1 and 1.2–1.3x at batch 32–64. Under DCP8 it brings decode to within 6% of TP8. It requires the V2 model runner and FULL decode CUDA graphs (the defaults), TP8 or TP8 + DCP8, and no expert parallelism in this recipe.

NVFP4 KV cache

The NVFP4 KV cache feature stores the MLA latent in NVFP4 (352 bytes per token per layer instead of 576) and holds about 1.5x the KV cache tokens of FP8. The indexer keys stay FP8, so the top-2048 token selection is unchanged. GSM8K matched FP8 on 8x GB300 (96.1% vs 96.2%); long-context accuracy has not been measured yet.

TPU serving

TPU serving runs through vLLM torch-TPU and is text-only: the vision tower is not loaded. K3.1 support is under review as vllm-torchtpu#2775; until it merges, the Install tab on TPU builds from the kimi-k31pro staging branch (Python 3.12, uv). No published image or wheel carries it yet.

The generated command uses vllm-torchtpu's default multiprocessing multi-host backend: run the head command on the slice's first host and the worker commands on the remaining hosts (the builder renders both). The Multi-Node TEP strategy at 4 nodes matches the validated flags of vllm-torchtpu#2775 with two differences: #2775 ran one driver over a Ray cluster (TPU_MULTIHOST_BACKEND=ray), and its GSM8K validation used the completions API, so the default-on chat and tool-call parsers were not exercised on TPU. Validated on a 2x2x4 TPU v7x slice with an FP8 latent KV cache, the unified block pool and prefix caching: the complete GSM8K split scored 97.50% (TP32) and 97.42% (TP32 + DCP8), prefix-cache hits returned identical tokens and logprobs, and the BEAM-1M long-context benchmark matched the GPU reference. The run-specific sizing of that validation (10K context, 16 sequences, compile sizes) is not the TPU default — it lives in the two Validated presets: TEP32 and TEP32 + DCP8 for KV-capacity-bound workloads.

Spec decoding is gated off on TPU for now: DSpark with --async-scheduling (always on in this configuration) needs vllm-torchtpu#2770, still open. #2775 validated DSpark with 7 draft tokens once that fix is in.

Prefill/decode disaggregation

The Prefill/Decode Disaggregation strategy runs one TP8 prefill pool (with expert parallelism) and one DEP16 decode pool (TP1 x DP16 with expert parallelism), and decode pulls each request's KV from prefill over NIXL. Every PD worker sets VLLM_SSM_CONV_STATE_LAYOUT=DS: NIXL transfers the Kimi Delta Attention conv state in three reads, which needs that layout, and the NIXL connector refuses to start without it.

KV cache offloading

The Simple option on the KV Offload row adds SimpleCPUOffloadConnector, a host-DRAM prefix tier for each GPU. Under Prefill/Decode Disaggregation it is added to the prefill pool only, next to NIXL in a MultiConnector; decode keeps plain NIXL. The tier pins cpu_bytes_to_use_per_rank of host memory for every GPU, so size it from the host DRAM that is free after the engines are up. The tested GB300 prefill workers used 128 GiB per GPU (512 GiB per tray).

Client usage

curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "moonshotai/Kimi-K3.1-Pro",
    "messages": [{"role": "user", "content": "Explain sparse attention in two sentences."}],
    "temperature": 0.7,
    "top_p": 0.95
  }'

References

• Model card
• Kimi-K3 recipe