avlp12 / Kimi-K3-Alis-MLX-Dynamic

Kimi-K3-Alis-MLX-Dynamic

Kimi K3 (2.8T MoE, multimodal) converted to MLX for Apple silicon, with a sensitivity-graded mixed-precision recipe. Vision is preserved — this is a multimodal build, not text-only: the MoonViT tower is verified against the reference implementation (cos = 1.000000000), and image-feature injection is validated through the full 93-layer two-box expert-parallel forward (finite outputs, rank-deterministic).

Build recipe

Sensitivity is computed per expert from calibration activations propagated through the actual quantized stack, layer by layer.

Why promotion instead of higher affine precision

The source checkpoint is MXFP4 QAT — the 4-bit values are the trained model. Re-quantizing them onto MLX's affine grid costs accuracy even when the target has more bits than the source, because the grids are geometrically mismatched (MXFP4: e2m1 floating, power-of-two block scales; affine: uniform, scale+bias). Measured on real weights and activations:

(single-layer signal on the calibration token set — the selection is post-hoc on that set; end-to-end quality is what the KL/PPL section measures)

With 10% sensitivity-based promotion the routed-path error drops to 0.211 (−44% vs uniform 2.5 bpw), and propagated hidden-state error over the first three quantized MoE layers stays at 0.13–0.18 with no compounding blow-up.

So the experts that matter are kept in their original format, bit-for-bit — with one runtime caveat: the serving config caps the MXFP4 path at moe_k_hi = 4 promoted experts per token (of 16 routed). When a token routes to a 5th+ promoted expert (~1.7%/token/layer at Binom(16, 0.1)), that expert's contribution is dropped rather than downgraded. All quality numbers below are measured under exactly this deployment condition. (The sensitivity selection itself was computed uncapped.)

Quality

Measured against the true MXFP4 teacher (layer-streamed reference pass over the original checkpoint — same 40 eval windows, bit-identical tokens):

The published weights include block-aligned layerwise DWQ (8 blocks — 12 layers each, the last being 9 layers + the final head), trained from the pristine quantized build against boundary activations of the full uncapped teacher (all 16 routed experts — the same reference the quality table uses). Trainables = 2-bit expert scales/biases + 4-bit dense scales/biases; MXFP4 promotions kept frozen bit-exact. Per-block validation nMSE improved −13.6…−31.1% (mean −18.7%), with the largest gains in the deep blocks:

The reference is the full uncapped teacher : a layer-streamed pass over the original MXFP4 checkpoint with all 16 routed experts evaluated per token (dense bf16, experts via bit-exact MXFP4 repack). PPL rows use 24 windows × 2047 tokens (49,128 tok) per dataset for teacher and build alike — bit-identical token windows, deterministic to 4 decimals across re-runs.

† pre-DWQ Korean KL/flip were captured on 16 of the 24 windows (the published build's numbers cover all 24); wikitext rows are exact same-window comparisons. All published-build numbers were measured with the absorbed MLA cache enabled (the shipped default), so they double as its end-to-end validation.

Gap framing, honestly told: DWQ closed 24.3% of the wikitext PPL gap to the teacher and 60.3% of the Korean gap — the quantization damage concentrates in distributions where the teacher is extremely sharp (wikitext PPL 1.58), which is also where 2-bit experts blunt the most. For scale context, a 4-bit build of this model is 1.5 TB and does not fit two 512 GB machines at all.

Multimodal sanity, end to end: given a synthetic RGB gradient image through the full two-box EP stack, greedy generation produced HTML/CSS that recreates a full-viewport gradient — the model recognizably sees the image through the quantized pipeline (weights, injection, and generation all as shipped).

Performance

Measured on 2× M3 Ultra 512GB (Thunderbolt 5 ring, expert-parallel: dense replicated, expert stacks split 50/50 — 491.4 GB per box):

Numbers above are with the absorbed MLA cache (default); at short context it is performance-neutral vs the expanded cache (5.75 step/s both) — its win is the 28× smaller KV footprint at long context.

The bottleneck is per-box compute, not the interconnect. Single-stream numbers; batching is the obvious next lever.

Operational notes that matter at this scale (M3 Ultra, ~491 GB resident): warm up all graph shapes ring-free before touching collectives (JIT windows on the ring cause GPU watchdog kills), keep a heartbeat op on the GPU during collective waits (idle re-wake costs ~7 s/layer otherwise), and launch loads only with ≥ model-size free pages (macOS compressor cascade kills otherwise). All three are encoded in the included harness.

Note on active parameters: K3 reads ~104B parameters per token, not the ~50B figure that circulated at release (that number omits the always-active dense path — attention, latent projections, shared experts, which are more than half the bytes read per token).

Usage

This repo ships the model weights plus the harness that runs them ( kimi_k3_mlx.py , kimi_k3_vision_mlx.py , k3_load.py , k3_ep.py , launch_full.sh , hostfile.example.json ) — K3 is not yet in upstream mlx-lm.

Two-box expert-parallel (Thunderbolt bridge between the Macs):

pip install mlx==0.31.2 mlx-lm==0.31.3 safetensors numpy transformers pillow
# 1) slice the expert stacks (once, one command per rank):
python3 slice_rank_half.py --src /path/to/this/repo --rank 0 --size 2 --out ~/k3_r0
python3 slice_rank_half.py --src /path/to/this/repo --rank 1 --size 2 --out ~/k3_r1
# 2) copy hostfile.example.json → hostfile.json, put your two TB-bridge IPs,
#    edit the VARs at the top of launch_full.sh, then:
bash launch_full.sh   # loads both ranks, runs a smoke forward + decode timing

Memory: raise the wired limit on both boxes first ( sudo sysctl iogpu.wired_limit_mb=500000 , re-run after every reboot).

Multimodal: kimi_k3_vision_mlx.py + inject_image_features reproduce the official image path (MoonViT 3D, bf16); see e2e_eval.py --image for the end-to-end example. The image preprocessor (tokenizer + processing configs) is not redistributed here — fetch those small files from the base repo: hf download moonshotai/Kimi-K3 tokenizer_config.json tiktoken.model preprocessor_config.json kimi_k3_processor.py kimi_k3_vision_processing.py media_utils.py tokenization_kimi.py --local-dir ./release

Notes

• Requires two 512 GB Macs; single-box is arithmetically impossible at any usable precision.
• Context: NoPE throughout, so 1M extrapolates without RoPE rescaling. KDA state is fixed-size (~0.4 GB, context-free); MLA KV grows with context. MLA absorb is implemented and shipped ( K3_MLA_ABSORB=1 , on by default in the launch scripts): only the compressed kv (512) + shared key (64) are cached — 61,440 B → 2,176 B per token per MLA layer (28.2×, parity-verified against the expanded path within the layer's bf16 noise floor). Budget math: ~52 KB/token across the 24 MLA layers → 128K context ≈ 6.7 GB per box ; 0.5M+ fits the KV budget. Long-context quality beyond the eval windows is not yet benchmarked; the practical limit is prefill time.
• Do not set a harness-level wired cap below the model footprint (e.g. K3_WIRED_CAP_GB=500 ): near-cap residency thrashes (~7 s per eval boundary measured). Use the OS sysctl above and leave the harness at mrws.
• Weights are 96 shards; the repo additionally carries the harness *.py , launch_full.sh , hostfile.example.json , LICENSE, and this card.
• License: Kimi K3 License (MIT-like). Copyright (c) 2026 Moonshot AI. Derivative distribution is permitted with notice (this repo complies); the commercial-scale UI-attribution clause (>100M MAU / >$20M monthly revenue) is summarized in the LICENSE file.
• Upstream fix contributed while porting: mlx-lm#1624 (KDA q/k normalization epsilon semantics).