smol-kimi-k3 is a training-oriented, text-only miniature of the language backbone described in Moonshot AI's Kimi K3 technical report. It compresses the defining ideas of a 2.8-trillion-parameter frontier model into 49,019,770 trainable parameters, then trains that model from scratch on TinyStories using a single RTX 4060 Laptop GPU.

This is an architecture study, not a claim that 49M parameters reproduce the capabilities of Kimi K3. The point is to make the interesting machinery small enough to read, run, modify, and train at home.

Most "tiny reproductions" replace the unusual parts of a frontier model with a standard Transformer. This project keeps the unusual parts:

• a 3 KDA : 1 Gated MLA hybrid attention rhythm, followed by a final MLA layer;
• channel-wise, lower-bounded Kimi Delta Attention with ShortConv;
• data-dependent Attention Residuals across network depth;
• sparse Stable LatentMoE channel mixing;
• bounded SiTU-GLU experts and exact small-batch Quantile Balancing;
• NoPE global attention;
• Muon for matrices and Per-Head Muon for attention projections;
• cached KDA/MLA autoregressive decoding;
• a tiny live dashboard backed by Matplotlib artifacts.

Kimi K3 repeats three KDA layers followed by one global MLA layer, then ends with an additional global layer. smol-kimi-k3 preserves that composition: 9 KDA layers + 4 Gated MLA layers.

Every attention and MoE sublayer receives a learned, data-dependent mixture of the embedding, completed block summaries, and the current partial block. That is the miniature implementation of Block Attention Residuals: information can be selected from earlier depths instead of being forced through one uniform residual stream.

KDA is a finite-state, linear-attention recurrence. Queries and keys pass through ShortConv, SiLU, and L2 normalization. A learned channel-wise decay forgets stale state, while the delta rule writes only the error between the current value and the state's prediction.

On CUDA, the same recurrence is evaluated with the report-linked fla-core chunkwise Triton kernel. A readable FP32 recurrent implementation remains available as the CPU fallback. Their measured mean absolute output difference is 0.00152 at BF16, against a mean reference magnitude of 0.22465.

Each layer has two always-on shared experts and 32 compact routed experts. Only the top four routed experts run for each token.

SiTU-GLU smoothly caps both multiplicative branches:

gate = 4  · tanh(Wg x / 4)  · sigmoid(Wg x)
up   = 25 · tanh(Wu x / 25)
out  = Wdown(gate ⊙ up)

The routed aggregate is normalized before projection back to model width. At this local scale, Quantile Balancing uses the exact batch quantile instead of the distributed histogram approximation required by the 2.8T model.

The vision encoder is intentionally omitted: at 401M parameters, the upstream tower alone is eight times larger than this entire model, and TinyStories contains no visual inputs. Gated MLA preserves low-rank KV compression and NoPE global attention, but does not attempt to reproduce Moonshot's production inference kernels.

The checked-in configuration lives in configs/k3_50m.json.

git clone https://github.com/cneuralnetwork/smol-kimi-k3.git
cd smol-kimi-k3
python3 -m pip install -r requirements.txt

Full training split plus 10,000 validation stories:

python3 scripts/prepare_tinystories.py \
  --output-dir data/tinystories

Small pipeline check:

python3 scripts/prepare_tinystories.py \
  --output-dir data/tinystories-smoke \
  --tokenizer-stories 10000 \
  --train-stories 10000 \
  --validation-stories 1000

In terminal A:

python3 dashboard.py \
  --run-dir checkpoints/k3-50m \
  --port 8080

Open http://127.0.0.1:8080.

In terminal B:

python3 train.py \
  --data-dir data/tinystories \
  --output-dir checkpoints/k3-50m \
  --steps 10000

The safe defaults are:

micro-batch             16
sequence length         256
gradient accumulation   2
effective token batch   8,192
warmup                  1%
schedule                cosine
weight decay            0.1
sample interval         250 steps

Resume from the most recent atomic checkpoint:

python3 train.py \
  --data-dir data/tinystories \
  --output-dir checkpoints/k3-50m \
  --steps 10000 \
  --resume checkpoints/k3-50m/latest.pt

Generate a story:

python3 generate.py checkpoints/k3-50m/best.pt \
  --prompt "Once upon a time, a little fox"

The best 10,000-step TinyStories checkpoint is published in the v0.1.0-tinystories GitHub Release. It uses BF16 SafeTensors and does not include optimizer state.

mkdir -p pretrained/smol-kimi-k3

gh release download v0.1.0-tinystories \
  --repo cneuralnetwork/smol-kimi-k3 \
  --dir pretrained/smol-kimi-k3

python3 generate.py pretrained/smol-kimi-k3/model.safetensors \
  --prompt "Once upon a time, a little fox"

Published bundle:

The selected checkpoint is step 9,750 with validation loss 1.6277. To reproduce the export from a local training checkpoint:

python3 scripts/export_model.py checkpoints/k3-50m/best.pt \
  --dashboard-state checkpoints/k3-50m/dashboard_state.json

The monitor deliberately stays simple: the trainer writes durable files, and the local web page polls them every two seconds.

The page shows:

• current step and run status;
• train and validation cross-entropy;
• measured tokens per second;
• peak allocated VRAM;
• the Matplotlib loss curve;
• the latest cached autoregressive sample.

Change sample cadence with --sample-interval, change its length with --sample-tokens, or use --skip-samples for pure throughput benchmarking.

Synthetic full optimizer steps at BF16, sequence length 256:

Batch 16 is the default because it leaves room for the display driver, evaluation, other desktop applications, and dashboard sample generation. If other applications consume substantial VRAM, use:

python3 train.py \
  --batch-size 8 \
  --gradient-accumulation 4

An end-to-end TinyStories verification step measured:

These values verify the pipeline; they are not final model-quality results.

smol-kimi-k3/
├── configs/
│   └── k3_50m.json              # authoritative model shape
├── dashboard/
│   ├── index.html               # barebones live UI
│   ├── styles.css
│   └── app.js
├── k3/
│   ├── config.py
│   ├── model.py                 # KDA, MLA, AttnRes, LatentMoE, decoding
│   ├── monitoring.py            # JSONL state + Matplotlib rendering
│   └── optim.py                 # Muon, Per-Head Muon, AdamW split
├── scripts/
│   └── prepare_tinystories.py
├── tests/
│   └── test_model.py
├── dashboard.py                 # local HTTP server
├── generate.py
└── train.py

python3 -m pytest -q

The test suite checks:

1. the report's 3 KDA + 1 MLA layer composition;
2. forward, loss, backward, and router-bias behavior;
3. causal invariance;
4. cached token-by-token decoding against full-sequence logits;
5. the 48–52M parameter budget;
6. dashboard state, metric, sample, and Matplotlib artifacts.

Expected result:

6 passed

• KDA training: official FLA chunkwise Triton kernel on CUDA; readable recurrent FP32 fallback on CPU.
• KDA decoding: recurrent state and ShortConv caches are carried one token at a time.
• MLA decoding: global key/value caches are retained per layer.
• Routing: biases select experts, but do not alter normalized mixture weights, matching the auxiliary-loss-free design.
• Balancing: exact quantiles replace distributed histograms because a single-device batch is small.
• Optimizer split: matrix parameters use Muon; Q/K/V matrices are orthogonalized per head; embeddings, norms, biases, and scalars use AdamW.
• Checkpointing: latest.pt and best.pt are written through a temporary file and atomically renamed.

This project is derived from ideas presented by the Kimi Team and uses the Flash Linear Attention KDA implementation. TinyStories was introduced by Eldan and Li and is consumed through the Hugging Face dataset.

Please consult the upstream repositories and technical report for their respective licenses, citations, and authoritative architectural details.