ArneshBanerjee/Kimi-K3-for-All

Kimi K3 for All

A small-scale Kimi-K3 trained from scratch on English Wikipedia, using Moonshot's original, unmodified K3 modeling code. The full 2.8T K3 is unrunnable locally; this is a faithful ~1.27B-parameter version (0.36B active) that runs on a laptop.

What this is: a real Kimi K3 — every structural switch that defines the architecture (lower-bounded KDA decay, Attention Residuals, Gated MLA, Stable LatentMoE, SiTU-GLU, Per-Head Muon) is preserved exactly, only the scale is reduced.

What this is not: a competitive general-purpose model. It is trained on one epoch of Wikipedia only (~18% of a Chinchilla-optimal token budget for this size), has no instruction tuning, and its facts are unreliable. Judge it on architectural fidelity and runnability, not benchmark scores.

Code: https://github.com/ArneshBanerjee/Kimi-K3-for-All

---

Architecture

	Kimi-K3 (official)	this model
total params	2.8T	1.270B
activated params	104B	0.364B (ex-embed)
layers	93	24
attention	69 KDA + 24 Gated MLA	18 KDA + 6 Gated MLA
hidden size	7168	1024
attention heads	96	14
experts	896	40
experts / token	16	4
shared experts	2	2
latent MoE dim	3584	512
expert FFN dim	3072	448
vocab	163,840	163,840

Ratios preserved under scaling

ratio	K3	this model
latent / hidden	0.5000	0.5000
expert FFN / hidden	0.4286	0.4375
attention width / hidden	1.7143	1.7500
KDA : MLA layers	2.875	3.000
final layer is MLA	yes	yes

K3's per-head geometry is kept exactly — MLA qk_nope=128 / qk_rope=64 / v=128 and KDA head_dim=128 are unchanged from the released config. Only the head count scales (96 → 14).

Matched verbatim from the released config.json: gate_lower_bound=-5.0, attn_res_block_size=12, mla_use_output_gate=true, mla_use_nope=true, latent_moe_use_norm=true, hidden_act="situ" (betas 4.0 / 25.0), sigmoid noaux_tc routing, moe_renormalize=true, first_k_dense_replace=1, rms_norm_eps=1e-5, tie_word_embeddings=false.

---

Training

Completed 2026-08-05: 15,600 steps, 4.601B tokens = 1.00 epoch of English Wikipedia (wikimedia/wikipedia, 20231101.en, CC BY-SA 3.0 / GFDL) — every token seen exactly once.

final validation perplexity	14.0 (CE 2.6417), best 13.3
measured val perplexity (benchmark)	13.89 (CE 2.6313, 102,400 held-out tokens)
HellaSwag zero-shot (1000 ex.)	33.7% (random = 25.0%)

33.7% against a 25% floor shows the model has learned real commonsense structure, not just memorized surface patterns — but it is far short of what a compute-optimal model this size would reach. That gap is tokens, not architecture: 4.6B trained vs. a Chinchilla-optimal 25B for this size (18%).

Known deviations from K3

1. Quantile Balancing → DeepSeek-style sign update. K3 balances experts with Quantile Balancing, a training-side procedure absent from the released inference code and described only qualitatively. We use the aux-loss-free sign update on e_score_correction_bias instead — it random-walks rather than converging.
2. Expert sparsity: 10% vs. K3's 1.8%. K3 activates 16/896 experts; with only 40 experts the closest workable setting is 4/40. 1.8% would require top-1 of 56+, too coarse to route at this scale. This is the single largest deviation, and it is forced by expert count, not chosen.
3. Text-only. K3's MoonViT-V2 vision tower (401M) is dropped.
4. bf16, not MXFP4/MXFP8. Trained in bf16 rather than the released quantized format.
5. MTP enabled (num_nextn_predict_layers=1); the released config ships 0. The mtp.safetensors head is included but optional at inference.
6. SDPA, not FlashAttention-2. Numerically equivalent, different kernel.

---

Usage

pip install torch "transformers==4.57.1" safetensors tiktoken blobfile einops numpy
git clone https://github.com/ArneshBanerjee/Kimi-K3-for-All
cd Kimi-K3-for-All
huggingface-cli download ArneshBanerjee/Kimi-K3-for-All --local-dir checkpoints/hf \
    --include "model.safetensors" "mtp.safetensors" "config.json"

python inference/generate.py -p "The history of computing began"

inference/generate.py picks its own backend: real fla Triton kernels on CUDA, and a pure-PyTorch KDA fallback (fla_shim.py) on Mac/CPU where Triton is unavailable. Both were verified to produce identical text from the same checkpoint.

Measured performance:

peak VRAM, inference (H100, bf16)	3.50 GiB (2048-token prefill + KV cache)
decode, batch 1 (H100, bf16)	14.7 tok/s
decode, batch 16 (H100, bf16)	135.5 tok/s
decode, batch 1 (Apple Silicon, MPS, fp32, pure-PyTorch KDA fallback)	1.7 tok/s

The Mac number uses the fla_shim pure-PyTorch KDA path (no Triton), float32. It has a known issue: bf16 currently fails on MPS due to a dtype mismatch in fla_shim's depthwise conv (F.conv1d gets mismatched input/weight dtypes) — use float32 there for now.

Example greedy output:

"World War II began in" → September 1939, and the first aircraft arrived in the

"The capital of France is" → in the 15th arrondissone of Paris. The

Fluent, syntactically solid, and — as expected from an undertrained base model — not always factually reliable.

---

Snapshots

snapshots/step_* contains 12 intermediate checkpoints from steps 1,656 → 15,440 across the training run — the trajectory of how the model's capabilities emerged, not just the final state. Each is model.safetensors + config.json in the same format as the top-level checkpoint.

---

Limitations

• Undertrained. ~18% of a Chinchilla-optimal token budget. Facts are unreliable.
• Wikipedia only. No code, conversation, or diverse web text.
• Expert sparsity 10% vs. K3's 1.8%, forced by the smaller expert count. Analysis shows 0 dead experts of 40 and content-dependent routing that sharpens with depth, but the fine-grained-MoE advantage K3 relies on doesn't fully survive the scale-down (see the GitHub README's MoE analysis).
• No instruction tuning. This is a base model.
• Mac/MPS bf16 is currently broken (see above); float32 works at ~1.7 tok/s.

---

Credits and licensing

The architecture and kimi_original/ code are Moonshot AI's, redistributed unmodified under the Kimi K3 License (see LICENSE; the MLA / MoE gating portions are adapted from DeepSeek-V3 under Apache 2.0). All credit for the architecture belongs to Moonshot AI — this project only shrank it and rebuilt what was needed to train it from scratch.

• Model: https://huggingface.co/moonshotai/Kimi-K3
• Paper: https://arxiv.org/abs/2607.24653
• Blog: https://www.kimi.com/blog/kimi-k3

Training data: wikimedia/wikipedia (20231101.en), CC BY-SA 3.0 / GFDL.