@nielspeter/mlx-ts
v0.5.0
Published
A TypeScript MLX SDK over Apple's mlx-c via FFI — Bun, Deno and Node. Apple Silicon only.
Maintainers
Readme
mlx-ts — TypeScript → mlx-c → Metal
npm · native runtime · FINDINGS · CHANGELOG
A TypeScript MLX SDK over mlx-c (Apple's official C API) via FFI, with
zero custom C/C++, no required runtime dependencies and no build step —
running on Bun, Deno and Node, and
numerically identical to MLX's Python reference for ops and bindings, and to
each model's original PyTorch implementation for the ports
(scripts/validate-all.sh: 69/69, of which 5 Stable Diffusion checks are
opt-in via MLXTS_SD=1, 5 Spark-TTS via MLXTS_TTS=1, and 2 Parakeet via
MLXTS_ASR=1).
Test coverage over src/ is measured and gated at 67% of functions / 74% of
lines.
Read docs/FINDINGS.md for what was proven and how. Apple Silicon + Metal only.
When you'd want this
Reach for mlx-ts when you are writing TypeScript on Apple Silicon and need the model in your process — not behind an HTTP hop, a subprocess, or a Python sidecar. Concretely, when you need:
- more than text out. Embeddings, raw logits, the KV cache, a custom sampler,
constrained decoding — anything a generate-endpoint cannot hand you. Nothing
crosses a serialization boundary here; an
MXis a pointer to GPU memory. - to build the model, not just call one.
src/models/are ordinary TypeScript files composingnnModules. A new architecture is a forward pass, not a new binary. - training or fine-tuning from TypeScript. Real
value_and_gradover a pytree of parameters, Adam, cross-entropy, LoRA — proven from a linear fit up to a full fine-tune of GPT-2-124M. - MLX specifically — Apple's own kernels and unified memory, and output that
matches
mlx-lmtoken-for-token, so a Python prototype ports without drift.
What that looks like
Each of these is exercised by a check in scripts/validate-all.sh and matched
against the Python reference — the code below is taken from the runnable files
in examples/, not written for the README.
Chat and streaming, in your process. No HTTP hop, no subprocess.
import { load, streamText } from "@nielspeter/mlx-ts";
const { model, tokenizer } = await load("mlx-community/Qwen3-0.6B-4bit");
const ids = tokenizer.encode("The capital of France is");
for await (const chunk of streamText(model, tokenizer, ids, { max: 64, temp: 0.7 }))
process.stdout.write(chunk);ChatTemplate.render([{ role: "user", content: "..." }]) turns messages into the
prompt string first when you want multi-turn — that is examples/chat.ts.
Speech to text. Token-exact against mlx_whisper, with language
auto-detection and a sliding window for unbounded audio.
import { loadWhisper, WhisperTokenizer, loadMelFilters, decodeAudio } from "@nielspeter/mlx-ts";
const model = await loadWhisper("models/config-turbo.json", "models/whisper-turbo.safetensors");
const tok = await WhisperTokenizer.fromFile();
const filters = await loadMelFilters("models/whisper-mel-filters-128.f32", 128);
const ids = model.transcribe(await decodeAudio("interview.flac"), filters);
console.log(tok.decode(ids).trim());Text to music. T5 conditioning, a codebook LM, EnCodec back to a waveform — all of it TypeScript.
import { MusicGen, saveAudio, seed } from "@nielspeter/mlx-ts";
const model = await MusicGen.fromPretrained();
seed(1234); // same seed -> same take
const audio = model.generate("trance", { maxSteps: 500 }); // 50 frames = 1s
await saveAudio("out.wav", audio.toF32(), model.samplingRate);Text to image. Stable Diffusion, end to end — CLIP conditions it, the UNet denoises, the VAE turns latents into pixels.
import { savePng, StableDiffusion } from "@nielspeter/mlx-ts";
const sd = await StableDiffusion.fromPretrained();
const img = sd.generate("a photo of an astronaut riding a horse", {
width: 384, height: 384, steps: 20, seed: 42, // same seed -> same image
});
await savePng("out.png", img.toF32(), 384, 384);Speech to text, fast. Parakeet TDT — NVIDIA's FastConformer transducer. Its decoder predicts how many encoder frames to skip at each step, so it does far less work than an autoregressive decoder over a fixed window.
import { Parakeet } from "@nielspeter/mlx-ts";
const asr = await Parakeet.fromPretrained();
console.log(await asr.transcribeFile("audio.wav"));Live transcription. A transducer emits each token once and moves on, so a streamed transcript is never revised — no sliding window, no segment seam.
const stream = new ParakeetStream(W, cfg, tok); // ~2 s behind the speaker
for await (const pcm of mic) process.stdout.write(stream.push(pcm));
process.stdout.write(stream.flush());bun examples/parakeet-live.ts audio.wav plays a file through the speakers while
feeding the model at microphone pace, so you can hear the delay rather than read
about it.
Text to speech. Spark-TTS: a Qwen2 LM predicts audio tokens, BiCodec renders them. The voice is described, not cloned — no reference clip needed.
import { SPARK_SAMPLE_RATE, SparkTTS, saveAudio } from "@nielspeter/mlx-ts";
const tts = await SparkTTS.fromPretrained();
const wav = await tts.generate("MLX runs on the GPU of your Mac.", {
gender: "female", pitch: "moderate", speed: "moderate", seed: 42,
});
await saveAudio("speech.wav", wav.toF32(), SPARK_SAMPLE_RATE); // 16 kHzCloning a voice. Six seconds of reference audio becomes 32 speaker tokens, written into the prompt so the model only has to generate the words.
const tts = await SparkTTS.fromPretrained();
const wav = await tts.clone("This was never spoken by that person.", "reference.wav");
await saveAudio("cloned.wav", wav.toF32(), SPARK_SAMPLE_RATE);Images and text in one space. CLIP's two towers project into the same 768 dimensions, so cosine similarity classifies without any training.
import { ClipVisionEncoder, fromF32, loadImage } from "@nielspeter/mlx-ts";
const px = await loadImage("photo.jpg", { size: 224 }); // sips decodes it
const vec = vision.embed(fromF32(px, [1, 224, 224, 3]), W.mx("visual_projection.weight"));
// ...then compare against text embeddings; examples/clip-zeroshot.ts is the whole thing.Embeddings for local RAG. Vectors, not a chat completion.
import { Qwen3, Tokenizer, loadSafetensors, fromI32, tidy } from "@nielspeter/mlx-ts";
const model = new Qwen3(cfg, loadSafetensors("models/model-q4.safetensors"));
const tokenizer = await Tokenizer.fromFile("models/tokenizer.json");
const ids = tokenizer.encode("a passage to index");
const vec = tidy(() => model.embeddingMX(fromI32(Int32Array.from(ids), [1, ids.length]), 1, ids.length));
const embedding = Array.from(vec.toF32()); // L2-normalized; pair with any JS vector storeembeddingMX is on the concrete model rather than the Decoder interface, so
this one takes Qwen3 directly instead of the load() helper — that is what
examples/server.ts does behind /v1/embeddings.
Training, from TypeScript. Real value_and_grad over a pytree of
parameters — the part an HTTP endpoint cannot give you at all.
import { Adam, crossEntropy, tidy, valueAndGrad, type MX, type Tree } from "@nielspeter/mlx-ts";
const forward = (p: Tree, x: MX) => { const { w, b } = p as { w: MX; b: MX }; return x.matmul(w).add(b); };
const lossFn = (p: Tree, x: MX, y: MX) => crossEntropy(forward(p, x), y);
const step = valueAndGrad(params, lossFn);
const opt = new Adam(0.1);
for (let i = 0; i <= STEPS; i++) {
const { loss, next } = tidy(() => {
const { loss, grads } = step(params, X, Y);
return { loss, next: opt.update(params, grads) };
});
params = next;
}A custom Metal kernel, written inline. EnCodec's LSTM is one of these, not a
demo — examples/metal-kernel.ts is Apple's own LSTM kernel, verbatim.
import { metalKernel, scalarI32, tidy } from "@nielspeter/mlx-ts";
const lstm = metalKernel({
name: "lstm",
inputNames: ["x", "h_in", "cell", "hidden_size", "time_step", "num_time_steps"],
outputNames: ["hidden_state", "cell_state"],
source: `/* Metal, compiled at first call */`,
});
const [hidden, cellOut] = tidy(() => lstm.apply(
[x, hIn, cell, scalarI32(H), scalarI32(0), scalarI32(T)],
[{ shape: [B, H] }, { shape: [B, H] }], // output shapes
[B, B * H, 1], // grid
[256, 1, 1], // threadgroup
));Runnable versions live in examples/: examples/chat.ts, examples/stream.ts,
examples/musicgen.ts, examples/train.ts, examples/metal-kernel.ts,
examples/hub.ts, examples/stable-diffusion.ts, examples/clip-zeroshot.ts,
examples/spark-tts.ts, examples/spark-clone.ts, examples/parakeet.ts,
examples/parakeet-live.ts, and examples/server.ts — an OpenAI-compatible endpoint with
a chat page and a live mic. CI runs them on Bun, Deno and Node.
When you'd want something else
Being honest about it is cheaper than you finding out later:
| if you need… | use |
|---|---|
| Linux, CUDA, Windows, or an Intel Mac | not this — Apple Silicon + Metal only |
| a chat endpoint with the least possible work | Ollama or LM Studio; a server is less code than an SDK |
| the same thing in Python | mlx-lm — same engine, far more models, maintained by Apple |
| GGUF, AWQ/GPTQ, or the widest model coverage | the llama.cpp ecosystem (e.g. node-llama-cpp) |
| to run in a browser | transformers.js — ONNX/WebGPU, a different engine entirely |
| high-throughput multi-tenant serving | not this yet: no continuous batching, generation is serialized |
The narrow version: if a generation API is enough, something else will get you there faster. This is for when it isn't.
Coming from Python?
Apple's stack is four packages. This repo is the analogue of all four, which is
why it is named after mlx rather than after mlx-lm:
| Python | here |
|---|---|
| mlx (mlx.core) — arrays, ops, memory | src/ffi + src/core |
| mlx.nn / mlx.optimizers | src/nn — Modules, Adam, valueAndGrad |
| mlx-lm — architectures, generation, tokenizers | src/models + src/text |
| mlx-whisper — speech-to-text | src/models/whisper.ts + src/audio |
If you only want the mlx-lm layer, that is generate / streamText plus a
model from src/models; the rest is underneath it, not in your way.
End to end: real Qwen3-0.6B generating text
mkdir -p models
curl -sL https://huggingface.co/Qwen/Qwen3-0.6B/resolve/main/config.json -o models/config.json
curl -sL https://huggingface.co/Qwen/Qwen3-0.6B/resolve/main/tokenizer.json -o models/tokenizer.json
curl -sL https://huggingface.co/Qwen/Qwen3-0.6B/resolve/main/model.safetensors -o models/model-qwen.safetensors
bun src/models/qwen.ts "The capital of France is"
# completion: " Paris. The capital of France is also the capital of the French Republic. ..."
# (24 tokens, ~190 tok/s)
python3 reference/reference-qwen.py "The capital of France is" # identical token idssrc/models/qwen.ts reads all dims from models/config.json (incl. Qwen3's explicit head_dim
and tied embeddings), loads the real weights, tokenizes with the validated
src/text/tokenizer.ts, and decodes with the KV cache — produced ids match MLX Python
token-for-token. Everything below is the validated machinery underneath it.
Install
npm i @nielspeter/mlx-ts # or: bun add / deno add npm:That is the whole install — no other tools. Audio and image decoding go
through macOS's own afconvert and sips, and playback through afplay; MLX
is Apple-Silicon-only, so there is nothing to gain by requiring ffmpeg. (The
parity suite still uses ffmpeg for the Whisper checks, because their oracle
decodes that way — a dev dependency, not a user one.)
The native runtime arrives as
@nielspeter/mlx-ts-darwin-arm64,
an optionalDependency carrying Apple's own libmlx + mlx.metallib next to
our libmlxc (199 MB unpacked). The parity suite passes forced onto it,
matching Homebrew's build exactly.
import { load, streamText } from "@nielspeter/mlx-ts";
const { model, tokenizer } = await load("mlx-community/Qwen3-0.6B-4bit");
for await (const piece of streamText(model, tokenizer, tokenizer.encode("Hello"), { max: 48 })) {
process.stdout.write(piece);
}load() fetches config, tokenizer and weights from the hub and caches them in
~/.cache/mlx-ts (MLXTS_CACHE overrides), so only the first run downloads.
Supported today: 4-bit qwen3 and olmoe checkpoints.
macOS on Apple Silicon only. Bun and Deno work as-is; Node needs 24+ (the
package ships compiled JS, because Node refuses to type-strip inside
node_modules) and pulls in koffi for FFI. The library package carries no
weights and no binaries — 162 kB packed; the native runtime is the separate
platform package above, pulled in automatically.
On Homebrew. brew install mlx-c is optional, and takes precedence when
present — useful if you already track mlx-c yourself, or want to test against a
different build. The resolver prefers, in order: MLXTS_LIB, a Homebrew
install, the platform package, a local prebuilds/. LIB_CANDIDATES shows what
it considered and what it found.
The repo itself is not the package: clone it for the examples, the parity suite
against MLX-Python, and docs/FINDINGS.md.
Try it without downloading anything
Four of the examples need no model files at all — after bun install they run
immediately:
bun examples/basics.ts # arrays, ops, and why tidy() is not optional
bun examples/module.ts # compose nn Modules into a model
bun examples/train.ts # valueAndGrad + Adam + cross-entropy, loss going down
bun examples/metal-kernel.ts # write your own Metal kernel when MLX has no fused opexamples/basics.ts measures the memory finding from docs/FINDINGS.md live — 200
[512,512] matmuls grow active memory by ~210 MB without tidy() and ~3 MB with.
Runtimes
The same code runs on all three JS runtimes, producing bit-identical output.
Only src/ffi/ differs; it selects a backend at import time. Everything in
src/ and examples/ — including the OpenAI-compatible server, SSE streaming
and audio transcription — runs on all three. training/ is still Bun-only
(Bun.mmap for the token shards, and Bun.file(...).writer() streams).
bun src/models/qwen.ts "The capital of France is"
deno run --allow-all src/models/qwen.ts "The capital of France is"
node src/models/qwen.ts "The capital of France is" # needs `npm i koffi`| runtime | binding | pointer dispatch | zero-copy readback |
|---|---|---|---|
| Bun | bun:ffi (built in) | ~12 ns/call | toArrayBuffer |
| Deno | Deno.dlopen (built in) | ~3 ns/call | UnsafePointerView |
| Node | koffi (prebuilt addon) | ~21 ns/call | koffi.view |
Measured by spikes/spike-ffi-*.ts (500k calls, warmed, best of 3). All are
cheap next to an MLX op — end-to-end generation throughput is the same on all
three within noise, which is what "compute-bound, not FFI-bound" predicts. One
sharp edge: a 64-bit return costs Deno ~52 ns because it falls off V8's
fast-call path, so hot accessors declare a 32-bit return.
Every pointer crosses the FFI boundary as a JS number — macOS user-space
addresses fit in 2^48, inside a double — so type Arr = number holds and no
model code is runtime-aware. Node runs .ts by stripping types only, so the
source stays free of enums and parameter properties.
Careful with a hand-rolled prebuilds/: a libmlxc built against a
different MLX than Homebrew's does not agree numerically — an earlier local
bundle diverged from MLX-Python on real Qwen3 and on LoRA training. That is not
true of the published platform package, which is built from Apple's own
mlx-metal binaries and passes the suite. scripts/validate-all.sh
prints which library it resolved; set MLXTS_LIB to choose.
Repo layout
src/ the SDK — ffi/ core/ nn/ text/ audio/ io/ models/, public API in index.ts
tools/ codegen.ts (headers -> src/ffi/generated.ts), inspect-real.ts
examples/ basics/module/train need no weights; server, chat UI, streaming CLI
training/ pretrain, SFT, LoRA, RL, data prep [Bun-only]
validation/ TS side of the parity suite — every file here is re-run and
diffed against reference/ on each validate-all.sh
reference/ MLX-Python / HF oracles every claim is checked against
spikes/ feasibility probes nothing depends on, kept as evidence
benchmarks/ op-level TS vs MLX-Python timings
scripts/ validate-all.sh (the full suite), run.sh (the pipeline)
docs/ FINDINGS.md first — the full write-up
models/ downloaded weights, configs, tokenizers ]
data/ corpora and tokenized shards ] all gitignored,
checkpoints/ training outputs ] created on demandNothing but source lives at the repo root. The three asset directories are
gitignored and created by the setup steps below, so a fresh clone is small and
git status stays quiet no matter how many gigabytes you download.
What runs
validation/mlx.ts— a minimal hand-written Bun-FFI binding overlibmlxc.dylib: handle management plusmatmul,rms_norm,rope,sdpa,silu, etc.tools/codegen.ts— parses the mlx-c headers and emitssrc/ffi/generated.ts: a full FFI symbol table (472 entries) + 242 typed op wrappers. The hand-writtenvalidation/mlx.tsexists only to bootstrap;src/ffi/generated.tssupersedes it.validation/block.ts— a full Qwen3 decoder block forward pass (mirrors mlx-lm'sqwen3.py): pre-norm, GQA attention with per-head q/k RMSNorm, RoPE, causal SDPA, residual, SwiGLU MLP, residual — using the hand binding.validation/block-gen.ts— the same block built entirely from the generated wrappers.validation/model-gen.ts— a small multi-layer Qwen3 model + KV-cache greedy decode loop (prefill + autoregressive steps), built from the generated wrappers.src/io/loader.ts— safetensors loading overmlx_load_safetensors: open a file into astring -> arraymap, pull tensors by name, enumerate via iterator.reference/save-model.py/validation/model-load.ts— Python writes the model to a real.safetensors; TS loads it and runs the decode loop from the loaded weights.tools/inspect-real.ts— loads a real mlx-community model shard and lists tensors.reference/reference-quant.py/validation/model-quant.ts— 4-bit quantized path: Python quantizes the Linear projections (mx.quantize) and savesweight/scales/biases; TS loads them and runs the decode withquantizedMatmul.src/text/tokenizer.ts— pure-TS byte-level BPE tokenizer (the real Qwen3models/tokenizer.json);reference/tok-reference.py/tests/tok-test.tsvalidate it against HFtokenizers(encode + decode, 11/11 cases).src/models/qwen.ts/reference/reference-qwen.py— config-driven real Qwen3-0.6B (bf16): readsmodels/config.json, loadsmodels/model-qwen.safetensors(HF key names), generates text; ids match MLX Python token-for-token.
Production runtime (mx + nn)
src/core/mx.ts—MXarray class: each wraps one handle, auto-freed by aFinalizationRegistry, plus atidy()arena for deterministic freeing, ops, and temp/top-p sampling.src/nn/nn.ts—Module,Linear,QuantizedLinear,RMSNorm,Embedding,QuantizedEmbedding.src/models/qwen-nn.ts/reference/reference-qwen-q4.py— real 4-bit Qwen3-0.6B (mlx-community format) overnn.Module; greedy ids match MLX Python token-for-token. Supports temp/top-p sampling, batching, sliding window.src/text/lm.ts— public generation surface: a model-agnosticDecoderinterface and async-generatorstreamTokens/streamText/generate. The KV cache is freed automatically (completion / earlybreak/ throw), so callers never calltidy()or free a handle;MXisDisposable.examples/stream.tsis the live demo.tests/validate-prod.ts— checks sampling reproducibility, batching, and bounded memory.tests/stream-test.ts— stream output is identical togenerate().
bun src/models/qwen-nn.ts "The capital of France is" # greedy
bun src/models/qwen-nn.ts --temp 0.8 --topp 0.95 --seed 42 "Once ..." # sampling
bun examples/stream.ts "Write a haiku about the sea" # streaming API
bun examples/server.ts # OpenAI-compatible HTTP server (:8080)
bun tests/validate-prod.ts # all threeexamples/server.ts is a working example of the local-server use case below — an
OpenAI-compatible /v1/chat/completions endpoint (streaming SSE or JSON) over
Bun.serve, generation serialized behind an async mutex. It also serves a tiny
self-contained chat web UI (examples/chat.html) at /:
bun examples/server.ts # open http://localhost:8080 for the chat UI, or:
curl localhost:8080/v1/chat/completions -H 'content-type: application/json' \
-d '{"messages":[{"role":"user","content":"Hi"}],"stream":true,"temperature":0.7,"top_k":40}'
curl localhost:8080/v1/embeddings -H 'content-type: application/json' \
-d '{"input":["a sentence to embed","another one"]}' # L2-normalized vectors for RAG
curl localhost:8080/v1/audio/transcriptions -F [email protected] # -> {"text": "..."}/v1/audio/transcriptions (multipart file) is enabled when the Whisper assets
are present (see below); otherwise it reports 501 and the rest of the server runs
normally.
The embeddings come from mean-pooling Qwen3's last-layer hidden states (same model/tokenizer, no extra weights) — RAG-useful similarity ranking, though a dedicated embedding model would rank better.
Memory — why tidy() and not just FinalizationRegistry: FR only fires
after a GC, which never happens inside a tight synchronous decode loop, so
handles pile up. Measured over a 200-token generation: FR alone grew active
memory +3034 MB; tidy() (free everything in scope except the token + KV
cache) grew it +23 MB — the KV cache only. FR remains a backstop for arrays
created outside any tidy().
reference/reference.py/reference/reference-decode.py— the same block / the same decode loop in MLX Python, same deterministic weights.
Run
bun tools/codegen.ts # parse headers -> generated.ts (+ coverage report)
bun validation/block.ts # hand binding: TS -> mlx-c -> Metal
bun validation/block-gen.ts # generated wrappers: TS -> mlx-c -> Metal
python3 reference/reference.py # MLX Python referenceAll three blocks print the same fingerprint:
sum = 0.005793
sum_sq = 0.162600The decode loop is checked the same way, but on discrete output — the greedy token ids must match exactly (any drift in cache concat, RoPE offset, masking, or sampling flips a token):
bun validation/model-gen.ts # TS + KV cache
python3 reference/reference-decode.py
# both: generated: [24, 3, 19, 2, 28, 1, 4, 14, 4, 14, 4, 14]Loading weights from safetensors
python3 reference/save-model.py # writes a real models/model.safetensors (25 tensors)
bun validation/model-load.ts # loads it via mlx_load_safetensors, decodes
# -> same ids: [24, 3, 19, 2, 28, 1, 4, 14, 4, 14, 4, 14]
# and on a genuine model file:
bun tools/inspect-real.ts ~/.cache/huggingface/hub/.../model-00001-of-00004.safetensors
# -> loaded ... — 881 tensors (real names + shapes)Gotcha: the safetensors Load primitive only implements eval_gpu == no, so
src/io/loader.ts loads on a CPU stream; the resulting concrete arrays then feed
the GPU compute graph normally.
4-bit quantized weights
Real mlx-community models are quantized. reference/reference-quant.py quantizes the
Linear projections with mx.quantize (group_size 64, 4 bits) and stores three
tensors each — weight (packed uint32), scales, biases; norms/embedding
stay fp32. validation/model-quant.ts loads them and uses quantizedMatmul(x, w, scales,
biases, transpose=true, 64, 4, "affine").
python3 reference/reference-quant.py && bun validation/model-quant.ts
# both: generated: [27, 26, 16, 11, 12, 30, 26, 16, 11, 12, 30, 26]The ids differ from the fp32 run — that is the real 4-bit quantization error, and both the TS and Python paths exhibit it identically.
Tokenizer (pure TS, no deps)
The one piece genuinely outside MLX. src/text/tokenizer.ts implements GPT-2-style
byte-level BPE over the real Qwen3 models/tokenizer.json: NFC normalization, the
special-token split, the GPT-2 pretokenization regex, the byte<->unicode map,
and rank-based merges.
mkdir -p models
curl -sL https://huggingface.co/Qwen/Qwen3-0.6B/resolve/main/tokenizer.json -o models/tokenizer.json
python3 reference/tok-reference.py && bun tests/tok-test.ts
# -> encode/decode parity vs Python tokenizers: 11/11 cases passValidated against HF tokenizers on contractions, em-dash, per-digit numbers,
tabs/newlines, Chinese, Japanese, emoji, source code, and the chat template
(special tokens like <|im_start|>).
Note: this tokenizer has Qwen3's full 151k vocab, so end-to-end text output needs a vocab-matched model (a real downloaded Qwen3) — i.e. the config-driven loading step. The toy models here use a 32-token vocab for fast parity checks.
Eval-boundary discipline
MLX is lazy: ops build a graph, nothing runs until forced. The decode loop
(validation/model-gen.ts) calls evalArray(...caches) plus reads the token each step, so
the per-layer KV caches and the next token become concrete arrays. Skip this and
the graph grows every step — unbounded memory and recompute. Per-step RoPE uses
offset = position; prefill uses a "causal" mask, single-token decode uses
none ("").
Codegen coverage
bun tools/codegen.ts reports exactly what it does, with no silent drops:
parsed 491 decls across 12 headers
symbols 472 FFI entries
wrappers 242 typed op wrappers (from ops.h, fast.h)
skipped 47: (all reported by name)Skipped functions are exotic only — nested-vector / char** / device handles,
each a one-line FFI-map addition. Every standard tensor op is wrapped. The metal
kernel-builder symbols are in the table too; src/ffi/kernel.ts puts an
ergonomic metalKernel() on top of them, which is what EnCodec's LSTM runs on.
The generator maps each C type to an FFI type, auto-supplies the trailing
mlx_stream, collapses const int* x, size_t x_num pairs into a single
number[] param, exposes nullable arrays as Arr | null, and packs the
by-value mlx_optional_float/int structs into a u64.
What it proves
Every inference-critical MLX fast op works over FFI and matches the reference:
mlx_fast_rms_norm, mlx_fast_rope (incl. the by-value mlx_optional_float
base, packed into a u64), mlx_fast_scaled_dot_product_attention with
mask_mode="causal" and GQA (4 query / 2 kv heads) handled inside the kernel.
Key ABI facts the binding relies on
- Every mlx-c handle is
struct { void* ctx; }. On Apple-silicon ARM64 a single-pointer struct is passed/returned in a register exactly like a bare pointer, so each handle is modeled asptr(a JS number). - An empty handle has
ctx == NULL, which Bun returns asnull→ normalize to0. - Out-params are
int fn(mlx_array* res, ...): pre-init the result slot with an empty handle, pass&slot, read the new handle back. mlx_optional_float { float value; bool has_value; }is an 8-byte non-HFA struct → passed in one integer register → modeled as au64with the float bits in the low 32 bits andhas_valuein byte 4.
What you can build with it
mlx-ts today is a local runtime for text LLMs, Whisper and Parakeet
speech-to-text, Spark-TTS text-to-speech, MusicGen text-to-music, Stable
Diffusion text-to-image and CLIP image embeddings, plus training (LoRA, full fine-tuning, a GRPO loss path) and
custom Metal kernels — Apple-Silicon-only, published as
@nielspeter/mlx-ts and also runnable as scripts in this repo. The library
under src/ runs on Bun, Deno and Node, as do all of examples/; only
training/ is still Bun-only. Sampling supports greedy,
temperature, top-p, top-k, and repetition penalty
(bun examples/stream.ts --temp 0.8 --topp 0.95 --topk 40 --reppenalty 1.1 "…").
✅ Buildable now (everything needed exists)
Local chat assistant / CLI — streaming replies, multi-turn via chat templates, temp/top-p/top-k/repetition-penalty sampling (
examples/chat.ts,examples/stream.ts).OpenAI-compatible server + chat web UI —
examples/server.tsoverBun.serve:/v1/chat/completions(SSE/JSON),/v1/embeddings,/v1/audio/transcriptions, and a self-contained chat page at/with a live mic (record → transcribe → edit → send). Single-process / low-concurrency, not multi-tenant.Speech-to-text (Parakeet TDT) —
examples/parakeet.ts: a recording in, a transcript out. NVIDIA's FastConformer encoder (src/models/parakeet.ts): an 8x depthwise-separable subsampling stem, 24 blocks of Macaron feed-forwards around Transformer-XL relative-position attention and a gated convolution, then a 2-layer LSTM prediction network and a joint that emits a token and a duration. That duration head is the point: a plain transducer advances one encoder frame per blank, while this one skips, so silence costs one step instead of many.Checked against the original PyTorch
transformers.ParakeetForTDT, not another MLX port: every stage matches, and the decode is token-for-token identical on real speech. 25 European languages, against Whisper's 99 — so it is an addition, not a replacement.Word timestamps (
--timestamps, or--srtfor a subtitle file) come out of the same duration head, which is why they cost nothing: the decode loop already walks encoder frames, so the pointer is a clock at 80 ms a frame, and it only records where it was. An attention decoder has no such pointer and needs a separate alignment pass. Verified by splicing clips together with exact silences between them — of 55 words, none landed in a gap. Starts are the number to trust; ends come from the duration head, capped at four frames.Streaming (
ParakeetStream): the decoder is genuinely incremental, so a token once emitted is never revised — unlike a sliding window, which re-transcribes and can rewrite what you already read. The encoder's attention is global, so each chunk is encoded with past context and a little future audio; that lookahead is the latency. Measured against transcribing the whole clip at once, ~2.2% word error at a 1.6 s average lag.--looktrades the two.examples/parakeet-live.tsplays a file to the speakers while transcribing it at microphone pace, so the delay can be heard rather than quoted.Memory is flat, which took fixing rather than luck. MLX's own accounting is the only thing that can see this — Metal buffers do not appear in process RSS, which drifts upward whether or not anything is leaking. Streaming 32 minutes of real audio, sampled every two minutes:
at 120s active 2509 MB peak 2765 MB at 960s active 2509 MB peak 2765 MB at 1920s active 2509 MB peak 2765 MBNot a megabyte of drift across the whole run, so a stream can go as long as someone keeps talking. Before
tidy()covered the encoder it grew 17.5 MB per 30 s — linear, no plateau, about 2 GB an hour:step()dropped the encoder output and 24 blocks of intermediates on the floor for the GC to find. Batch decode had its own leak, 76 MB per utterance, and is now likewise flat across repeated calls. Nearly all of the resident 2509 MB is the weights (627M parameters at F32).Which to use is decided by accuracy, not memory. Scored against FLEURS' own Danish references — not against batch, since batch is the thing in question — there is a crossover near two minutes:
| audio | batch WER | stream WER | |--------|-----------|------------| | 64 s | 10.7% | 26.2% | | 94 s | 20.7% | 30.5% | | 138 s | 34.0% | 27.7% | | 251 s | 33.3% | 23.5% |
Batch is clearly better on short clips, and 3-5x faster besides — that is what it is for. But it degrades steadily with length, because the model was trained on short utterances and a few thousand frames of global attention is out of distribution; NeMo's own long-form inference limits the attention context for the same reason. Streaming's 7 s window is therefore not a compromise, it is closer to how the model was trained, and its accuracy barely moves with length while its memory does not move at all.
On continuous speech it is worse than that table shows, because splicing read sentences hands batch a restart at every join that real talk never gives it. Take one recording, transcribe the same opening 20 s as part of files of different lengths, and count how many words in that fixed window change — the audio is identical, so only the length varies, and each run is its own control:
file 20s batch: 0 words changed stream: 0 words changed file 45s batch: 13 words changed stream: 8 words changed file 91s batch: 18 words changed stream: 8 words changedBatch rewrites audio it had already heard, worse the more you append. Streaming shifts once and then holds — and that 8 is not drift but the end of the file: in the 20 s run those last frames are flushed with no lookahead, where a longer file gives them proper future context. So the degradation is measurable by 45 seconds of continuous speech, not two minutes.
Streaming's own numbers improve with length only because its fixed 3 s warmup is amortised over more audio.
So: short clips batch, anything approaching a minute of continuous speech stream — and stream it even when the whole file is already on disk.
Speech-to-text (Whisper), multilingual —
src/audio/mel.ts(log-Mel, ~1e-6 vs numpy FFT) +src/models/whisper.ts(Conv1d stem, bidirectional encoder, cross-attention decoder, KV cache) +src/text/whisper-tokenizer.ts. Token-for-token identical tomlx_whisper(tests/whisper-transcribe-test.ts). Runslarge-v3-turbowith auto language detection and a sliding window for unbounded dictation; Danish/Swedish/ English verified.bun src/models/whisper.ts audio.flac(setup below).Text-to-music (MusicGen) —
examples/musicgen.ts: prompt in,.wavout. SentencePiece Unigram tokenizer (src/text/unigram.ts) → T5 encoder (src/models/t5.ts) → the MusicGen LM (src/models/musicgen.ts: 4 delayed EnCodec codebooks, cross-attention, classifier-free guidance) → the EnCodec decoder (src/models/encodec.ts, whose LSTM runs on a hand-written Metal kernel). LM logits match Hugging Face's own implementation.-smallis the default;jasonvassallo/mlx-musicgen-{medium,large}are the larger sizes, since Facebook ships those only as PyTorch pickles.Text-to-speech (Spark-TTS) —
examples/spark-tts.ts: a sentence in, a.wavout, no phonemizer and noespeak-ng. A Qwen2-0.5B LM (src/models/qwen2.ts) predicts audio tokens out of a 166k vocabulary, and BiCodec (src/models/bicodec.ts) renders them: a codebook quantizer, an FSQ speaker decoder, a 12-layer Vocos prenet conditioned on the speaker through AdaLayerNorm, and a Snake-activation wave generator that upsamples 320x to 16 kHz. The voice is described — gender, pitch, speed — rather than cloned. Verified two ways: stage by stage against mlx-audio, and end to end by speaking a sentence and transcribing it back with our own Whisper (validation/spark-roundtrip.ts). ~2x realtime on an M-series Mac.Voice cloning (Spark-TTS) —
examples/spark-clone.ts: a recording in, the same voice saying something else. BiCodec's speaker encoder (src/models/speaker.ts): a Slaney mel front end, an ECAPA-TDNN with Res2Net blocks and attentive statistics pooling, a perceiver resampler that squeezes any clip length into 32 latents, and FSQ to pack those into token ids. Checked against the original PyTorch Spark-TTS, not against another port: all 32 ids match on synthetic and on real audio. That mattered — the mlx-audio port left-aligns a short STFT window wheretorch.stftcentres it, which silently moved 12 of the 32 ids until it was caught. Those reference numbers are committed (validation/spark-golden.json), so the checks run with nothing installed beyond mlx-ts itself. Cloning is also checked end to end with no Python:validation/spark-clone.tsclones a voice and scores it with ECAPA's x-vector, a different head from the one the tokens come from (~0.95 against a ~0.38 floor for an unrelated voice).Multilingual chat — the server injects a system prompt so replies come back in the user's language (Danish in → Danish out).
Local RAG —
POST /v1/embeddingsreturns L2-normalized sentence vectors (mean-pooled Qwen3 hidden states); pair with any JS vector store.Prompt-driven text tools — summarize / rewrite / classify / extract / translate; agent loops (tool use via prompting + JS parsing).
LoRA fine-tuning of 4-bit Qwen3 (Adam + cross-entropy,
training/lora-train.ts).Train a transformer from scratch —
validation/spike-microgpt.tsbuilds Karpathy's ~4k-param microGPT (embeddings → attention → MLP → tied head) and trains it end-to-end on the names corpus with the autograd being real MLX over FFI (his hand-rolledValueengine replaced byvalue_and_grad); step-0 loss is exact vs the MLX-Python mirror, both converge.Train a real small GPT —
validation/spike-nanogpt.tsscales that up to nanoGPT: a multi-layer char-level GPT on tiny-shakespeare, mini-batched[B,T], AdamW + cosine LR + warmup + global grad clipping + dropout. At nanoGPT's exactshakespeare-charconfig (6 layers, 384 dim, 10.7M params) it reaches best val loss ≈ 1.50 — matching nanoGPT's ~1.47 baseline — and writes coherent Shakespeare (real character names, dialogue). The dropout-free path is bit-exact vsreference/reference-nanogpt.py(shared init + batches).Run real GPT-2-124M —
src/models/gpt2.tsloads the actual OpenAIgpt2weights and generates with a pure-TS GPT-2 BPE encoder (src/text/tokenizer.ts+GPT2_SPLIT, 8/8 token-exact vs HF) —gelu_new, LayerNorm-with-bias, tied head, KV cache, token-exact vsreference/reference-gpt2.pyat ~210 tok/s. Seedocs/GPT2.md.SFT a chatbot —
training/sft.tsfull-fine-tunes real GPT-2-124M into an instruction follower (chat format + completion-only loss), the nanochat chat stage. Step-0 loss matchesreference/reference-sft.py; after SFT it answers in-format, including a held-out question (Italy → Rome). Seedocs/SFT.md.RL with GRPO —
training/rl.tsruns Group Relative Policy Optimization on GPT-2-124M (the nanochat RL stage): sample a group of completions, reward them, normalize advantage, policy-gradient update. Positivity-reward demo: mean reward rises ~9×; GRPO loss path validated vsreference/reference-rl.py. Seedocs/RL.md.Train a tokenizer —
reference/tok-train.pytrains a byte-level BPE in native Rust (HFtokenizers, as nanochat does — training is a data-prep boundary step, not MLX compute); our pure-TSsrc/text/tokenizer.tsthen reproduces it token-exact (tests/tok-train-test.ts). Thetok_trainstage of a nanochat-style pipeline.Pretrain + checkpoint —
training/base-train.tspretrains a GPT from scratch on BPE-tokenized text and saves a safetensors checkpoint (mx.saveSafetensors, the write side of the loader) that reloads round-trip-clean — the keystone that lets pretrain → SFT/inference chain. Thebase_trainstage.The whole pipeline (
scripts/run.sh) — the TS-over-MLX analogue of nanochat'sruncpu.sh: dataset → tokenizer → data-prep → pretrain → SFT → chat, end to end on one Mac. Trains on TinyStories (coherent at this scale) via a streaming dataloader —training/data-prep.tsstream-encodes the corpus to uint16 token shards,training/base-train.tsBun.mmaps them (scales past RAM). Then SFT + chat (CLItraining/chat-ckpt.tsor web UIexamples/chat-web.ts). Seedocs/PIPELINE.md.Research / inspection — pull logits, hidden states; the
MXop surface is open.
Whisper setup (weights/assets are git-ignored — fetched, like the LLM weights):
W=https://huggingface.co/mlx-community/whisper-large-v3-turbo/resolve/main
mkdir -p models
curl -sL $W/config.json -o models/config-turbo.json
curl -sL $W/weights.safetensors -o models/whisper-turbo.safetensors
# The validation oracle. Its own venv, because it pins dependencies the other
# reference scripts do not want — and /tmp/wvenv is where validate-all.sh looks.
python3 -m venv /tmp/wvenv && /tmp/wvenv/bin/pip install mlx-whisper
curl -sL -o /tmp/jfk.flac https://github.com/openai/whisper/raw/main/tests/jfk.flac # transcription fixture
WA=$(/tmp/wvenv/bin/python -c 'import mlx_whisper,os;print(os.path.dirname(mlx_whisper.__file__))')/assets
cp "$WA/multilingual.tiktoken" models/whisper-multilingual.tiktoken
python3 -c "import mlx.core as mx,numpy as np;np.array(mx.load('$WA/mel_filters.npz')['mel_128']).astype('float32').tofile('models/whisper-mel-filters-128.f32')"
bun src/models/whisper.ts audio.flac # auto-detects language; anything macOS can decodeOLMoE-1B-7B 4-bit setup (the MoE model — weights git-ignored, ~3.9 GB):
O=https://huggingface.co/mlx-community/OLMoE-1B-7B-0125-Instruct-4bit/resolve/main
mkdir -p models
curl -sL $O/config.json -o models/config-olmoe.json
curl -sL $O/tokenizer.json -o models/tokenizer-olmoe.json
curl -sL $O/model.safetensors -o models/model-olmoe.safetensors
python3 reference/split-olmoe.py # -> models/model-olmoe-sharded/ (for the sharded-loader test)
bun src/models/olmoe.ts "The capital of France is"Note: the original 0924 checkpoint was replaced upstream by 0125 (identical
architecture: 16 layers, 64 experts, group_size 64 / 4-bit). The validate-all
OLMoE checks compare src/models/olmoe.ts against reference/reference-olmoe.py — both load the same
models/model-olmoe.safetensors — so any matching 4-bit checkpoint restores parity.
GPT-2-124M setup (real OpenAI weights — git-ignored, ~550 MB):
G=https://huggingface.co/openai-community/gpt2/resolve/main
mkdir -p models
curl -sL $G/config.json -o models/config-gpt2.json
curl -sL $G/tokenizer.json -o models/gpt2-tokenizer.json
curl -sL $G/model.safetensors -o models/gpt2-model.safetensors
bun src/models/gpt2.ts "The capital of France is" # greedy; TEMP/TOP_K/TOP_P/REP to sample (see docs/GPT2.md)🟡 Needs modest code (clear path, no feasibility risk)
- More architectures (Llama, Mistral, Gemma, Phi…) — a forward over
nnmodules + weight-key mapping (src/models/olmoe.ts/src/models/whisper.tsare templates). - Better embeddings — a dedicated embedding model (BERT encoder + WordPiece, or Qwen3-Embedding with last-token pooling) for stronger retrieval than the current mean-pooled base-LLM vectors.
❌ Not yet (substantial new code or a real gap)
- Cloning a speaker's cadence — cloning the voice ships (see above), but Spark can also take the reference's transcript, which conditions the LM on a worked example. Measured, that is worth having: without it a clone of an 11 s JFK clip runs 36% faster than he speaks, and with it 8%. Identity is unaffected either way. It needs BiCodec's content encoder plus wav2vec2-large-xlsr-53 (~317M params, PyTorch pickle only), so it is a second model rather than a tweak.
- Cross-platform — Apple-Silicon + Metal only (no Linux/CUDA, Windows, Intel). Runtime portability is done: see "Runtimes" above.
- Multimodal LLMs (LLaVA and friends) — the vision half exists now
(
src/models/clip-vision.ts), but nothing yet feeds those patch embeddings into a language model. - High-throughput multi-tenant serving — only equal-length batching; ragged prompts need padding masks + continuous batching.
- Broad model compatibility — only affine 4-bit quant (no AWQ/GPTQ), so many HF quantized checkpoints won't load.
- Constrained/JSON decoding, beam search, speculative decoding — none yet.
- Per-sample gradients — need
vmap, the one genuine mlx-c capability gap. Minibatch training is proven well past LoRA: a full fine-tune of GPT-2-124M and nanoGPT from scratch both match the Python reference.
License
MIT — see LICENSE. This project vendors and derives from MIT-licensed Apple
MLX code and depends on other third-party work; NOTICE has the attributions.
No model weights or datasets are tracked here; the setup steps download them
from their original sources under their own licenses.
Not affiliated with or endorsed by Apple.
