{"slug": "show-hn-i-implemented-the-kimi-k3-paper-from-scratch-in-pytorch", "title": "Show HN: I implemented the Kimi K3 paper from scratch in PyTorch", "summary": "A developer released a PyTorch implementation of the Kimi K3 architecture from the arXiv paper 'Kimi K3: Open Frontier Intelligence' (arXiv:2607.24653v1), reproducing the paper's Table 1 parameter counts (2.778T total, 104.108B activated per token, 400.9M vision tower) and running the 20M-parameter nano model on CPU. The implementation includes Kimi Delta Attention, Gated MLA, Stable LatentMoE, and MoonViT-V2, with training results showing held-out loss dropping from 6.2726 to 0.8434 and parser-valid generations rising from 0.1389 to 0.7021 over 3,000 steps on a real corpus.", "body_md": "An implementation of the architecture, pre-training protocol, systems co-designs,\nand post-training framework described in *Kimi K3: Open Frontier Intelligence*\n(arXiv:2607.24653v1).\n\nScale is configuration. Nothing in `kimi3/`\n\nbranches on model size, so the same\ncode path runs a 20M-parameter model on a laptop CPU and describes the paper's\n2.8T configuration. This repository executes the small end and the verification\nsuite. Training at cluster scale is out of scope, and `docs/PAPER_MAP.md`\n\nmarks\nwhich components are wired into a runnable path and which are libraries.\n\nParameter counts reproduce the paper's Table 1:\n\n| This implementation | Paper (Table 1) | |\n|---|---|---|\n| Total parameters | 2.778 T | 2.78 T |\n| Activated per token | 104.108 B | 104.2 B |\n| Vision tower | 400.9 M | 401 M |\n\nPython ≥ 3.10, PyTorch ≥ 2.4. **No GPU is required.** The test suite, the\nparameter audit of the 2.8T config, the demos, and the nano training run all\ncomplete on CPU.\n\nA GPU is used when one is available: `resolve_device(\"auto\")`\n\nselects CUDA if the\ncard has enough free memory, and falls back to CPU otherwise. Training is roughly\n7× faster per step on a small GPU. Triton is optional; the fused KDA path falls\nback to the reference implementation when it is absent, and never runs during\ntraining because it has no fused backward pass.\n\n```\npip install -e .          # runtime\npip install -e '.[dev]'   # adds pytest\n```\n\nOptional extras: `[hf]`\n\n, `[vision]`\n\n, `[triton]`\n\n.\n\n```\n# Inspect the 2.8T configuration without allocating anything\npython3 scripts/count_params.py configs/k3_2p8t.yaml --check\n\n# Train the nano model end to end (CPU is fine)\nbash scripts/train_nano.sh\n\n# Train on your own data: a text directory, a JSONL file, or packed shards\npython3 -m kimi3.train.pretrain --config configs/k3_nano.yaml --jsonl corpus.jsonl\n\n# Run the full gate suite\nbash scripts/verify.sh\n```\n\nFive runnable demos are in `demos/`\n\n, including a trained 19.8M model generating\nPython and a side-by-side comparison of models trained on real versus synthetic\ndata. See `demos/README.md`\n\n.\n\nThe nano model trains from a loss of 8.13 to 0.32 in 400 steps on synthetic\ndata, starting at `ln(vocab)`\n\nas expected for correct initialisation.\n\nOn a real corpus -- 47,902 documents of system documentation and source code, deduplicated and split so no held-out document shares a near-duplicate cluster with training data -- 3,000 steps give:\n\n| Metric | Start | End | Reference |\n|---|---|---|---|\n| Held-out loss | 6.2726 | 0.8434 | 6.238 = chance |\n| Parser-valid generations | 0.1389 | 0.7021 | 0.8182 = training corpus |\n\nThe model reaches 86% of the parser-validity rate of its own training data, and\nremains significantly below it (z = −2.23 over 480 generations). `docs/RESULTS.md`\n\nPhase 9 documents the scorers, confidence intervals, and why a dictionary\nword-rate metric is not a quality measure on its own.\n\nImplemented from paper §2:\n\n**Kimi Delta Attention** with the lower-bounded decay of Eq. 5 and a full-rank output gate.**Gated MLA**, with NoPE (no positional encoding) throughout.** Block Attention Residuals**.** Stable LatentMoE**with SiTU-GLU and Quantile Balancing.** MoonViT-V2**, trained from scratch under next-token prediction.\n\nAt the 2.8T configuration:\n\n**Hybrid attention schedule.** `hybrid_ratio: 3`\n\nwith `trailing_global: true`\n\nlays out 93 layers as 23 × [KDA, KDA, KDA, MLA] plus a trailing MLA: 69 KDA and\n24 MLA layers, so the tail carries two consecutive MLA layers. Under NoPE,\nposition information enters the model only through KDA's recurrent decay.\n\n**LatentMoE.** 896 routed experts at top-16 (sparsity 56), plus two always-on\nshared experts and one leading dense FFN layer. Routing is dropless: there is no\ncapacity factor in the config, every token emits exactly top-k assignments, and\nall of them are computed. Empty experts are skipped; tokens are not. The router\nscores the full-width d = 7168 token. The 3584-wide latent reduces expert compute\nand dispatch payload without shrinking the routing space. The Quantile-Balancing\nbias is a buffer rather than a parameter: mean-centred, applied one optimizer\nstep late, and frozen at inference.\n\n**Block Attention Residuals.** 8 blocks of 12 layers, so at most eight completed\nblock sources, one partial sum, and the embedding are resident simultaneously —\n`O(N·d)`\n\nlive memory instead of `O(L·d)`\n\n. Across pipeline stages, `AttnResCache`\n\ntransfers only the blocks the current stage created.\n\nPer-Head Muon, cosine schedule with 1% warmup, the four-stage 8K → 64K → 256K → 1M context curriculum, and document-isolated packing.\n\nKDA Context Parallelism (Eq. 17), MoonEP with the appendix E bound on redundant experts, Pipeline ZeRO-2 with CPU-resident gradient shards, P2P Muon orthogonalisation, block-wise FP8 activations, a unified activation manager, and cache-based AttnRes pipeline communication.\n\n**These are verified components, not an assembled distributed trainer.** Each is\ntested against a single-process reference, and KCP and P2P Muon run under real\n`torchrun`\n\nprocesses. However, `kimi3/train/pretrain.py`\n\nis single-process and\nimports nothing from `kimi3/parallel/`\n\n.\n\nTwo components are accounting rather than transport:\n\n- MoonEP's dispatch and combine are a local permutation with no\n`all_to_all`\n\n. - The pipeline schedule computes stage partitions and transfer byte counts without moving activations.\n\n`docs/PAPER_MAP.md`\n\nrecords the status of each component, and\n`tests/test_reachability.py`\n\nenforces it: a module that nothing calls fails the\nsuite unless it is listed with a reason.\n\nThe XTML chat template from appendix F, SFT with MXFP4/MXFP8 quantisation-aware training, partial rollout, per-problem reasoning-effort budgets producing nine domain × effort experts, an agentic generative reward model, Multi-Teacher On-Policy Distillation, and resumable sandboxes.\n\nA unified paged cache for the hybrid KDA–MLA state, KDA-aware prefix caching with\ndecoupled hash and physical granularities, ReplaySSM speculative decoding, and an\nEAGLE-3 draft model trained with the LK loss. A single `PagePool`\n\nserves both\nMLA-KV and KDA-state pages. It is sized by constructor argument, has no YAML\nbinding and no eviction policy, and raises `MemoryError(\"page pool exhausted\")`\n\nat the limit.\n\nThe reference implementations define the intended behaviour; the kernels are optimisations of them. The main gates are therefore equivalence checks:\n\n| Gate | Result |\n|---|---|\n| KDA chunkwise form vs Eq. 1 recurrence (fp64) | 6.7e-16, all chunk sizes, with and without document resets |\n| KDA Context Parallelism vs single-device states | 4.86e-16 across four processes |\n| MoonEP balance within the appendix E bound | 200/200 randomised skewed routings, including all-to-one |\n| Block AttnRes at block size 1 vs unblocked Eq. 8–9 | exact |\n| Triton kernel vs reference | 5e-4, cosine similarity 1.000000 |\n\nFor comparison, the naive additive scheme used for vanilla linear attention is wrong by over 10% on the KCP test case. That error is the reason KCP exists.\n\nRun everything with `bash scripts/verify.sh`\n\n. This includes 322 unit and\nequivalence tests, three multi-process `torchrun`\n\ngates, and the Triton kernel\ncheck.\n\n| Config | Total | Activated | Purpose |\n|---|---|---|---|\n`k3_nano` |\n19.754 M | 8.744 M | correctness proof |\n`k3_micro` |\n999.755 M | 273.092 M | multi-process parity, 64K context |\n`k3_small` |\n14.452 B | 1.995 B | first multi-node run |\n`k3_2p8t` |\n2.778 T | 104.108 B | the paper's configuration |\n\nEvery figure is `scripts/count_params.py`\n\noutput rather than an estimate.\n\n`docs/SCALING.md`\n\ncovers choosing `PP × EP × CP × DP`\n\ndegrees and the failure\nmodes to expect. `docs/HARDWARE.md`\n\nderives memory, cluster size, interconnect,\nand cost for each rung.\n\n**Training a 2.8T model, or reproducing any benchmark score.** The proprietary corpora and in-house evaluation suites are not available.**A multi-node launcher.** No Slurm, Ray, or`torchrun --nnodes`\n\npath exists. The only distributed entry points are the three single-node parity gates in`scripts/verify.sh`\n\n, which run on gloo.**An SFT or RL entry point.**`kimi3/posttrain/sft.py`\n\nand`kimi3/posttrain/rl/`\n\nexport pure functions — losses, schedulers, reward models — with no`main`\n\n, no CLI, and no optimizer step. MXFP4/MXFP8 QAT is wired into the expert forward and tested by`tests/test_qat.py`\n\n, but nothing enables it automatically.**MoE all-to-all.** MoonEP plans and balances the expert assignment; the token exchange itself is a local permutation.**A fused backward pass for the Triton KDA kernel.** The kernel therefore never runs during training: the module upcasts to fp32 and the fused path declines any call requiring gradients.**The Firecracker backend for AgentENV.** What exists is an in-process lifecycle state machine (`start/pause/resume/fork/snapshot/destroy`\n\nover a dictionary) that executes no code. The repository contains no`subprocess`\n\n,`exec`\n\n, or`os.system`\n\ncall, and the`Sandbox`\n\nprotocol has no execute method. A real backend would need to define its isolation contract — process, filesystem, network, timeout, teardown — before running model-generated code.**A vision pathway in the trained model.**`kimi3/vision/`\n\nimplements MoonViT-V2 and is tested against the paper's parameter count, but`kimi3/model.py`\n\ndoes not import it. The backbone is text-only.\n\n| File | Contents |\n|---|---|\n`docs/PAPER_MAP.md` |\nevery paper section mapped to its implementation and test |\n`docs/RESULTS.md` |\nmeasured results for every gate |\n`docs/AMBIGUITIES.md` |\nthe nine fields Table 1 leaves unspecified, and the evidence for each default |\n`docs/SCALING.md` |\nparallelism degrees and failure modes |\n`docs/HARDWARE.md` |\nmemory, cluster size, and cost per rung |\n`demos/README.md` |\nfive runnable demonstrations |\n\nCopyright © 2026 Tim Rots.\n\nLicensed under the **GNU Affero General Public License, version 3 or later**\n([ LICENSE](/TimRots/kimi3/blob/master/LICENSE)). You may use, study, modify and redistribute this code.\nIf you modify it and make it available to others — including over a network, as\na hosted service — you must release your complete source under the same licence.\n\nThat is deliberate. The AGPL does not forbid commercial use; it forbids building a closed product on this work. Anyone is free to run it, learn from it, and contribute back.\n\n**Commercial licences are available.** If the AGPL's source-disclosure\nrequirement does not suit your organisation, contact the author to arrange\nalternative terms.\n\n**The Kimi K3 architecture.** This is an independent implementation of a public paper (arXiv:2607.24653v1). The licence covers this source code, not the ideas or the paper, which belong to their authors.**Model weights.** Any checkpoint produced by training this code inherits obligations from the data it was trained on. The demo weights come from a corpus that is roughly 69 % copyleft-licensed by document count; the legal status of weights derived from such data is unsettled and is not resolved by this licence. No trained weights are distributed in this repository.", "url": "https://wpnews.pro/news/show-hn-i-implemented-the-kimi-k3-paper-from-scratch-in-pytorch", "canonical_source": "https://github.com/TimRots/kimi3", "published_at": "2026-08-02 21:09:06+00:00", "updated_at": "2026-08-02 21:52:49.292937+00:00", "lang": "en", "topics": ["artificial-intelligence", "machine-learning", "large-language-models", "ai-research", "developer-tools"], "entities": ["Kimi K3", "PyTorch", "arXiv", "MoonViT-V2", "LatentMoE", "Kimi Delta Attention"], "alternates": {"html": "https://wpnews.pro/news/show-hn-i-implemented-the-kimi-k3-paper-from-scratch-in-pytorch", "markdown": "https://wpnews.pro/news/show-hn-i-implemented-the-kimi-k3-paper-from-scratch-in-pytorch.md", "text": "https://wpnews.pro/news/show-hn-i-implemented-the-kimi-k3-paper-from-scratch-in-pytorch.txt", "jsonld": "https://wpnews.pro/news/show-hn-i-implemented-the-kimi-k3-paper-from-scratch-in-pytorch.jsonld"}}