Kimi-K3.1-Pro-2bit-UVMAX

Kimi K3.1 Pro (kimi_k31) for MLX, with the per-class bit widths of kernelpool/Kimi-K3-2bit-UVMAX: routed experts at 2 bits, the other matrices at 4 to 8, the vision tower in BF16. 2.78T parameters in 816.9 GB (2.35 bits per parameter).

Source: moonshotai/Kimi-K3.1-Pro at 6d0a020: routed experts in MXFP4, the rest in BF16, 1.56 TB in 96 shards. Upstream's d9e99b0 changes only encoding_k3.py (it drops medium from the thinking efforts). This repo ships that newer encoding_k3.py; every other file comes from 6d0a020.
K3.1's DSA indexer is held at 8 bits. Kimi K3 and kernelpool's mix have no indexer.
Runs on our tensor-parallel MLX engine (not published). mlx-lm has no kimi_k31 model; see Running.

Files

File	Bytes	Contents
`model-00001-of-00095.safetensors` ... `model-00093-of-00095.safetensors`	814,727,273,860	decoder layers 0-92, one per file (1,132,924,187 for the dense layer 0, up to 8,886,248,146)
`model-00094-of-00095.safetensors`	1,321,249,813	`embed_tokens`, `lm_head`, `norm`, `output_attn_res_norm`, `output_attn_res_proj`
`model-00095-of-00095.safetensors`	894,737,694	vision tower and `mm_projector`, BF16, as published
`config.json`	363,384	upstream's, with its `quantization_config` (compressed-tensors) replaced: `quantization` / `quantization_config` hold the defaults and 1,487 per-layer entries; `k3_whole` names the source and the mix
`model.safetensors.index.json`	574,625	5,746 tensors
`tiktoken.model`, `tokenization_kimi.py`, `tokenizer_config.json`, `encoding_k3.py`, `configuration_kimi_k3.py`, `kimi_k3_processor.py`, `kimi_k3_vision_processing.py`, `media_utils.py`, `preprocessor_config.json`, `generation_config.json`, `LICENSE`	2,884,979	upstream's

Total: 816,947,084,355 bytes in 108 files.

Quantization

MLX affine; scales and biases in BF16.

Tensors	Parameters	Format
Routed experts `switch_mlp.{gate,up,down}_proj`: 92 MoE layers x 896, 3,584 to 3,072 and back	2,722.74B	2-bit, group 128
KDA attention, 69 layers: `qkv_proj`, `g_proj`, `o_proj`, `f_a_proj`, `f_b_proj`, `b_proj`	30.62B	6-bit, group 64; `qkv_conv`, `o_norm` BF16; `A_log`, `dt_bias` FP32
MLA attention, 24 layers: `q_a_proj`, `q_b_proj`, `kv_a_proj_with_mqa`, `embed_q`, `unembed_out`, `g_proj`, `o_proj`	5.57B	6-bit, group 64; norms BF16
Shared experts, 92 layers (2 fused, intermediate 6,144)	12.16B	8-bit, group 64
MoE latent projections `routed_expert_{down,up}_proj`, 7,168 to 3,584 and back	4.73B	8-bit, group 64; `routed_expert_norm` BF16
Dense MLP, layer 0 (intermediate 33,792)	0.73B	8-bit, group 64
`embed_tokens`, `lm_head` (163,840 x 7,168)	2.35B	4-bit, group 64
Router `mlp.gate` (92 x 896 x 7,168)	0.59B	8-bit, group 64; `e_score_correction_bias` FP32
DSA indexer `wq_b`, `wk`, `weights_proj`, 24 layers	0.18B	8-bit, group 64; `k_norm` weight and bias BF16
Vision tower (27 blocks) and `mm_projector`	0.45B	BF16, unchanged
Norms, attention-residual projections (`self_attention_res_proj`, `mlp_res_proj`, `output_attn_res_proj`)	0.004B	BF16

2,780,110,361,184 parameters in 816,942,482,560 bytes of tensors: 2.35 bits per parameter. The routed experts alone take 2.25 bits per parameter (765,770,858,496 bytes).

Routed experts. Each expert's MXFP4 weights (E2M1 codes, an E8M0 scale per 32) are decoded and re-quantized to affine 2-bit in groups of 128. For each group the quantizer tries the min-max range and that range shrunk about its middle to 95, 90, 85, 80 and 75%, and keeps the one with the least squared error.
Everything else quantized. mx.quantize (min-max) of the BF16 weight.
Build. Our converter (not published) streams the original one unit at a time: a decoder layer from its own shard, then the globals, then the vision tower. It made the 95 units in 32.5 min on one Mac Studio M3 Ultra (MLX 0.32.4).
Tensor names. These are mlx-lm kimi_k3 names after sanitize(), as in kernelpool's Kimi K3 checkpoint: experts stacked in switch_mlp, MLA's kv_b_proj split into embed_q / unembed_out. The 24 DSA layers add self_attn.indexer.*.

Architecture

From config.json:


Layers	93: 1 dense (MLP 33,792) + 92 MoE; hidden size 7,168; activation `situ`
MoE	896 routed experts, 16 per token, + 2 shared; the routed experts work in a 3,584-wide latent (intermediate 3,072); sigmoid router, `noaux_tc`, renormalized
Attention	69 KDA layers (Kimi Delta Attention, linear: 96 heads x 128, short conv 4, full-rank gate) + 24 MLA layers (4, 8, ..., 92 and 93, 1-based: 96 heads, q_lora_rank 1,536, kv_lora_rank 512, qk 128 + 64, v 128, NoPE, output gate)
Sparse attention (DSA)	on the 24 MLA layers: indexer of 32 heads x 128 from the query latent, one key of 128 through a LayerNorm with bias, no RoPE; each query attends to its top 2,048 positions
Attention residuals	`attn_res_block_size` 12
Positions	`max_position_embeddings` 1,048,576
Vocabulary	163,840 (tiktoken)
Vision	27 blocks, hidden 1,024, MLP 4,096, 12 heads, patch 14; `patchmergerv2` projector (2 x 2 merge) to 7,168
Cache per token	MLA 27,648 bytes (24 layers, BF16) + indexer keys 6,144 = 33,792 bytes; the KDA layers keep a fixed-size state per sequence

Running

No public loader. mlx-lm (main, 2026-11-05) has kimi_k3 but no kimi_k31 model, and its kimi_k3 has no DSA indexer.
Tested only with our engine.
- Model: mlx-lm 0.32.3.dev8's kimi_k3, with an indexer added to each DSA layer's attention before quantization, so the checkpoint loads with strict=True.
- Split over 4 ranks by mlx-lm's shard(); the indexer stays whole on every rank.
- Vision tower not loaded; MLX 0.32.4.
Memory at TP4. MLX peak on rank 0: 232.7 / 236.1 / 240.9 GiB in 1K / 4K / 16K-token prefills. Served: 239-242 GB wired per Mac.
Chat encoding. encoding_k3.py (custom code); tokenizer_config.json has no chat_template. Thinking is on by default; thinking_effort takes low, high or max. generation_config.json: max_length 1,048,576, eos_token_id 163586.

Measured

Setup:

Hardware: 4 x Mac Studio M3 Ultra 512 GB, macOS 27.0, full mesh of Thunderbolt 5 cables with RDMA; TP4.
Software: MLX 0.32.4, mlx-lm 0.32.3.dev8.
Runs (2026-11-05): greedy, 256 tokens decoded, one run each.
Drafter: Inferact/Kimi-K3.1-Pro-DSpark-preview in BF16, split over the ranks.

	Chat prompt, 52 tokens	Code edit, 404 tokens
Plain decode	18.6 tok/s	18.5 tok/s
DSpark, 3 drafts	38.1 tok/s, 3.3 tokens/cycle	41.0 tok/s, 3.6
DSpark, 3, 5 or 7 drafts by acceptance	33.1-33.5 tok/s, 3.4	45.6-48.0 tok/s, 5.6-5.8

Prefill in chunks of 1,024: 1K, 4K and 16K tokens in 4.49, 16.08 and 70.56 s (228, 255 and 232 tok/s).

Through our server, thinking on, drafter on, greedy, 512 tokens:

Prompt	First token	Decode
Python function with tests, 117 tokens	1.16 s	29.8 tok/s
MoE routing explained in Russian, 137 tokens	1.35 s	23.7 tok/s
Arithmetic word problem, 124 tokens (answer correct)	1.26 s	35.4 tok/s
First prompt again, from the prompt cache	0.17 s	29.8 tok/s

Quality was checked with teacher forcing on 9 texts (2,111 tokens). No K3.1 reference was run. The reference is Kimi K3, not K3.1, in a 4-8-bit build (MXFP4 experts as published, the rest 8-bit), so the numbers include the difference between the two models as well as the quantization.

Text	KL	Same top token
All 2,111 tokens	0.174 (p99 2.39)	89%
Math	0.10	93%
Table	0.10	91%
Logic	0.15	88%
Code	0.06	95%
Translation	0.07	100%
Russian prose (2 texts)	0.21	89-90%
English explanation	0.42	83%
Corpus	0.17	87%

Perplexity was 1.73 (reference 1.51). The same check on kernelpool/Kimi-K3-2bit-UVMAX, Kimi K3 against Kimi K3, gave KL 0.157, 90% and perplexity 1.75.

License

The Kimi K3.1-Pro License (LICENSE, Moonshot AI), inherited from the base model. It is MIT-style with two conditions:

A Model-as-a-Service business with more than US$20M revenue over 12 months needs a separate agreement with Moonshot AI.
A product with more than 100M monthly active users or more than US$20M monthly revenue must display "Kimi K3.1-Pro".

AinsleyMarkov/Kimi-K3.1-Pro-2bit-UVMAX