Inkling · MLX

Run Thinking Machines Lab's Inkling — a 975B-total / 41B-active sparse-MoE, natively multimodal model — on Apple Silicon with MLX.

This repo is a from-scratch MLX port of Inkling's inkling_mm_model architecture (not in stock mlx-lm/mlx-vlm) plus a streaming HF→MLX converter/quantizer that never loads the ~1.9 TB bf16 checkpoint into RAM.

> What is the capital of France?
The capital of France is **Paris**.

> Write one sentence about the ocean.
The ocean covers more than seventy percent of Earth's surface and remains
one of the least explored frontiers on our planet.

_{— generated by the 4-bit build on an M3 Ultra.}

✨ Features

🍎 Native Apple Silicon	pure MLX; no CUDA, no PyTorch for inference
🧩 Full architecture	hybrid local/global attention, per-head QK-norm, relative-position bias, 4 short-convs/layer, sigmoid MoE router (256 experts, top-6 + 2 shared "sink"), muP logits
🖼️ Multimodal towers	HMLP vision + dMel audio encoders ported & numerically validated
📦 4 / 6 / 8-bit	standard MLX affine group quant; 4-bit runs the full 975B model in ~496 GB
🌊 Streaming convert	quantizes tensor-by-tensor — never holds the 1.9 TB model in memory
✅ Validated	fp32 parity vs reference + coherent real generation (see Validation)

📐 Architecture

tokens ─▶ embed ─▶ embed-norm ─┐
                                ▼
   66 × ┌───────────────────────────────────────────────┐
        │ pre-norm ▶ attention(hybrid SWA/global, QK-norm,│
        │            rel-pos bias, +short-conv) ▶ +resid  │
        │ pre-norm ▶ MoE(256 experts, top-6 + 2 shared)   │
        │            or dense SwiGLU (layers 0–1) ▶ +resid │
        └───────────────────────────────────────────────┘
                                ▼
                 final-norm ▶ ÷muP ▶ unembed
   (image/audio features scatter into the token stream before the stack)

🗂️ Which build do I download?

Build	Size	Runs on
4-bit	~496 GB	512 GB Mac (M3 Ultra)
6-bit	~717 GB	≥ ~768 GB
8-bit	~937 GB	≥ ~1 TB

The quantized quality is effectively lossless here — even 4-bit reproduces the model's structure with ~100% confidence on the correct next token.

🧬 REAP-pruned builds (smaller, expert-pruned)

REAP (Cerebras, arXiv:2510.13999) drops the lowest-saliency routed experts per MoE layer, where saliency = mean over active tokens of router_gate × ‖expert_output‖₂. The router renormalizes over survivors; the 2 shared experts, attention, and embeddings are untouched.

Calibrate on text, images and audio. Saliency was profiled over a mixed corpus of text (code + 15 languages + reasoning), 200 real images, and 180 speech clips run through the full vision and audio paths. This matters: a text-only calibration prunes experts that ground visual features (a Pallas's cat → "brown bear", a golf ball → "butterfly"), and text+image alone leaves audio-grounding experts unprotected (speech transcription overlap fell 0.88 → 0.57 at 25% pruning) — all while text perplexity looked fine. Profiling over every modality keeps each expert that matters to any of them. Inkling routes very uniformly (entropy 0.922; ~0 cold experts/layer under full multimodal calibration), so it is only lightly prunable:

Build	Experts kept	Size	Text ppl	vs unpruned	Vision (image ID)	Audio (speech)
4-bit (unpruned)	256	~496 GB	3.830	—	✓	✓
REAP-12	225	~470 GB	3.806	−0.6% (free)	6/6	0.88
REAP-25	192	~402 GB	3.946	+3.0%	6/6	0.87
REAP-50	128	~272 GB	4.682	+22.2% ⚠️	5/6	0.87

(Text-only calibration scored 2/6 vision; text+image left audio at 0.57 — hence the full multimodal recalibration, which recovers both at no text cost.) REAP-25 is the sweet spot: a real size cut that clears the 512 GB memory cliff (comfortable eager/wired load) for ~3% text perplexity, with vision and audio intact. REAP-50 keeps audio but text and fine-grained vision degrade — experimental only. Reproduce with scripts/profile_experts_mm.py → scripts/prune_build.py <usage.npz> mm → scripts/eval_build.py.

🛠️ Prerequisites

• macOS on Apple Silicon with enough unified memory for the build above
• Python 3.10+

git clone https://github.com/PipeNetwork/inkling-mlx.git
cd inkling-mlx
pip install -r requirements.txt        # mlx, mlx-lm, transformers, numpy, scipy

For near-capacity builds, raise the GPU wired-memory limit once per boot:

sudo sysctl iogpu.wired_limit_mb=500000   # ~524 GB; adjust to your RAM

🚀 Quick start

from inkling_mlx.load import load
from inkling_mlx.generate import greedy_generate
from transformers import AutoTokenizer

model, config = load("/path/to/Inkling-MLX-4bit")
tok = AutoTokenizer.from_pretrained("/path/to/Inkling-MLX-4bit", trust_remote_code=True)
tok.chat_template = open("/path/to/Inkling-MLX-4bit/chat_template.jinja").read()

ids = tok.apply_chat_template(
    [{"role": "user", "content": "What is the capital of France?"}],
    add_generation_prompt=True, reasoning_effort="none", tokenize=True)["input_ids"]
print(tok.decode(greedy_generate(model, config, ids, max_new_tokens=64)))

Inkling is a reasoning model: its chat template injects a Thinking effort level system message. Use reasoning_effort ∈ {none, minimal, low, medium, high, max}.

Images & audio

InklingProcessor handles the full preprocessing (image patchify + CLIP-normalize; audio log-mel → dMel bins) and inserts the placeholder soft-tokens:

from inkling_mlx.processing import InklingProcessor
from PIL import Image

proc = InklingProcessor(tok, open("/path/to/model/chat_template.jinja").read())
inputs = proc.apply([{"role": "user", "content": [
    {"type": "image", "image": Image.open("cat.jpg")},
    {"type": "text",  "text": "What's in this image?"},
    # or {"type": "audio", "audio": <16kHz mono np.ndarray>, "sampling_rate": 16000}
]}])
out = greedy_generate(model, config, inputs["input_ids"], max_new_tokens=128,
                      pixel_values=inputs.get("pixel_values"),
                      audio_input_ids=inputs.get("audio_input_ids"))
print(tok.decode(out[len(inputs["input_ids"]):]))

Speech transcription is verbatim — two held-out LibriSpeech clips through the 4-bit REAP-25 build ("Transcribe the speech in this audio."), each returned in ~7 s:

Reference	Model output
…WHEN DEATH LIKE SOME REMORSELESS CREDITOR SEIZES ON ALL WE FONDLY THOUGHT OUR OWN…	"…when death, like some remorseless creditor, seizes on all we fondly thought our own…"
…ACCORDING TO SALVIAN AND HIS CONTEMPORARIES THE VANDAL CONQUERORS WORKED IN NORTH AFRICA…	"…according to Salvian and his contemporaries, the Vandal conquerors worked in North Africa…"

Both scored 1.00 content-word overlap — including obscure proper nouns ("Salvian") and archaic phrasing — with punctuation and casing added. (This fidelity is because the pruning was calibrated on audio; the text+image-only build paraphrased and dropped proper nouns — see below.)

⚡ Performance

• Load mode. load() eagerly materializes weights wired-resident by default, so forwards don't re-read the mmap. For a model near your RAM ceiling (e.g. the 496 GB 4-bit build on a 512 GB Mac), the eager copy may not fit — pass --lazy (CLI) / load(path, lazy=True) to mmap instead (lower peak RAM, but the first forward pays the ~weight-read and decode can thrash near capacity).

• Attention uses the fused scaled_dot_product_attention kernel.

• Multimodal prefill scales with image patches (one 40 px patch = one soft-token). Big images become long prompts; cap resolution with max_long_edge to trade a little detail for a much shorter prefill:

  proc.apply(messages, max_long_edge=512)   # ~130 patches instead of ~450 for a 960px image

🔄 Converting / quantizing yourself

# one build
python -m inkling_mlx.convert_cli --src /path/Inkling-src --dst out-4bit --bits 4
# standard sweep (bf16 + 8/6/4-bit)
scripts/convert_all.sh /path/Inkling-src /path/out

# REAP-pruned build: profile expert saliency, then convert with --prune
python scripts/profile_experts.py /path/out-4bit          # -> expert_usage.npz
python scripts/prune_experts.py 0.25                      # -> keep_indices.npz (25% prune)
python -m inkling_mlx.convert_cli --src /path/Inkling-src --dst out-reap25 --bits 4 \
       --prune /path/keep_indices.npz
# (or prune an already-quantized build in one pass, bit-identically: scripts/prune_build.py)

Quantized: attention/MLP/expert projections, embeddings, vision/audio matmuls. Kept high-precision: the MoE router, RMSNorms, the four short-convs, and the relative-position bias. Conversion is streaming, so it runs in bounded memory regardless of model size.

Why affine int4 and not MXFP4 / NVFP4?

MLX also supports floating-point 4-bit modes (mx.quantize(mode="mxfp4"|"nvfp4")), and Thinking Machines ships an Inkling-NVFP4 checkpoint — so it's a fair question. We benchmarked round-trip reconstruction error (‖W − Ŵ‖ / ‖W‖ vs bf16) on real Inkling expert weights:

Scheme	bits/weight	reconstruction error
affine int4 (group 64)	4.50	~9.1%
nvfp4 (group 16)	4.50	~10.2%
mxfp4 (group 32)	4.25	~12.3%

Affine int4 is the most faithful — it's asymmetric (per-group scale and zero-point, 16 uniform levels), so it centers on Inkling's near-Gaussian expert weights better than symmetric FP4's fixed non-uniform levels (scale only, no zero-point). FP4's real payoff is heavy-tailed activations and native Blackwell FP4 tensor cores — neither helps weight fidelity on Apple Silicon (MLX would dequantize FP4 anyway, as there's no FP4 compute). So these builds use affine int4; a Mac port of the NVFP4 checkpoint would be lower quality at best-equal size.

✅ Validation

• fp32 parity of the decoder layer vs a verbatim reference — max |Δ| ~4e-3 on real weights, both attention types + full MoE.
• Real generation — coherent, correct answers at 4-bit (above).
• Vision/audio towers match the reference to ~1e-5.

Note on the gate/up layout. Inkling's checkpoint stores each fused gate+up MLP weight (w13) row-interleaved — [g0, u0, g1, u1, …], with the gemm reading 0::2/1::2. A contiguous [:half]/[half:] split silently scrambles gate↔up in every layer and yields confident-but-wrong output. The transformers PR skeleton has this same omission; the authoritative loader is SGLang's deinterleave_w13. This port de-interleaves correctly.

📁 Project structure

inkling-mlx/
├── inkling_mlx/            # the MLX model + converter/loader/generation
│   ├── model.py text.py layers.py attention.py moe.py
│   ├── vision.py audio.py common.py cache.py config.py
│   └── convert.py load.py generate.py convert_cli.py
├── scripts/               # convert_all.sh, verify_shards.py, smoke_test.py
├── tests/                 # parity / cache / driver / single-layer reference checks
└── requirements.txt

🐛 Troubleshooting

Problem	Fix
Incoherent / multilingual output	ensure you loaded the chat_template.jinja and used apply_chat_template
OOM while loading	raise iogpu.wired_limit_mb; use a smaller build; close other apps
No module named inkling_mlx.models...	this arch isn't in stock mlx-lm — load via this repo's load()
Very slow first token	MLX pages weights from disk on first use; subsequent tokens are faster

📚 Resources

• Source model: thinkingmachines/Inkling
• MLX builds: 🤗 pipenetwork
• References: SGLang PR · transformers PR · MLX

📝 License

Code: Apache-2.0. Model weights inherit the Apache-2.0 license of the base model.

⚠️ Status

Text and multimodal generation are complete and validated: the vision/audio towers and their preprocessing (InklingProcessor — image patchify/normalize, audio log-mel→dMel, validated to ~1e-7 vs the reference) are included.