nielspeter/mlx-ts

A TypeScript MLX SDK over Apple's mlx-c via FFI. Runs on Bun, Deno and Node; output matches Apple's mlx-lm token for token. Apple Silicon only.

0

stars

111

commits

TypeScript

primary language

Aug 29, 2026

updated

www.npmjs.com/package/@nielspeter/mlx-ts
apple-silicon
bun
deno
ffi
inference
llm
machine-learning
metal
mlx
mlx-lm
typescript

README

mlx-ts — TypeScript → mlx-c → Metal

npm version CI License: MIT Node.js >= 24 Apple Silicon

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 MX is a pointer to GPU memory.
  • to build the model, not just call one. src/models/ are ordinary TypeScript files composing nn Modules. A new architecture is a forward pass, not a new binary.
  • training or fine-tuning from TypeScript. Real value_and_grad over 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-lm token-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 kHz

Cloning 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 store

embeddingMX 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 Macnot this — Apple Silicon + Metal only
a chat endpoint with the least possible workOllama or LM Studio; a server is less code than an SDK
the same thing in Pythonmlx-lm — same engine, far more models, maintained by Apple
GGUF, AWQ/GPTQ, or the widest model coveragethe llama.cpp ecosystem (e.g. node-llama-cpp)
to run in a browsertransformers.js — ONNX/WebGPU, a different engine entirely
high-throughput multi-tenant servingnot 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:

Pythonhere
mlx (mlx.core) — arrays, ops, memorysrc/ffi + src/core
mlx.nn / mlx.optimizerssrc/nn — Modules, Adam, valueAndGrad
mlx-lm — architectures, generation, tokenizerssrc/models + src/text
mlx-whisper — speech-to-textsrc/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 ids

src/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 op

examples/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`
runtimebindingpointer dispatchzero-copy readback
Bunbun:ffi (built in)~12 ns/calltoArrayBuffer
DenoDeno.dlopen (built in)~3 ns/callUnsafePointerView
Nodekoffi (prebuilt addon)~21 ns/callkoffi.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 demand

Nothing 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 over libmlxc.dylib: handle management plus matmul, rms_norm, rope, sdpa, silu, etc.
  • tools/codegen.tsparses the mlx-c headers and emits src/ffi/generated.ts: a full FFI symbol table (472 entries) + 242 typed op wrappers. The hand-written validation/mlx.ts exists only to bootstrap; src/ffi/generated.ts supersedes it.
  • validation/block.ts — a full Qwen3 decoder block forward pass (mirrors mlx-lm's qwen3.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 over mlx_load_safetensors: open a file into a string -> array map, 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.ts4-bit quantized path: Python quantizes the Linear projections (mx.quantize) and saves weight/scales/ biases; TS loads them and runs the decode with quantizedMatmul.
  • src/text/tokenizer.ts — pure-TS byte-level BPE tokenizer (the real Qwen3 models/tokenizer.json); reference/tok-reference.py / tests/tok-test.ts validate it against HF tokenizers (encode + decode, 11/11 cases).
  • src/models/qwen.ts / reference/reference-qwen.pyconfig-driven real Qwen3-0.6B (bf16): reads models/config.json, loads models/model-qwen.safetensors (HF key names), generates text; ids match MLX Python token-for-token.

Production runtime (mx + nn)

  • src/core/mx.tsMX array class: each wraps one handle, auto-freed by a FinalizationRegistry, plus a tidy() arena for deterministic freeing, ops, and temp/top-p sampling.
  • src/nn/nn.tsModule, Linear, QuantizedLinear, RMSNorm, Embedding, QuantizedEmbedding.
  • src/models/qwen-nn.ts / reference/reference-qwen-q4.pyreal 4-bit Qwen3-0.6B (mlx-community format) over nn.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-agnostic Decoder interface and async-generator streamTokens / streamText / generate. The KV cache is freed automatically (completion / early break / throw), so callers never call tidy() or free a handle; MX is Disposable. examples/stream.ts is the live demo.
  • tests/validate-prod.ts — checks sampling reproducibility, batching, and bounded memory. tests/stream-test.ts — stream output is identical to generate().
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 three

examples/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 file=@audio.flac   # -> {"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 reference

All three blocks print the same fingerprint:

  sum    = 0.005793
  sum_sq = 0.162600

The 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 pass

Validated 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 as ptr (a JS number).
  • An empty handle has ctx == NULL, which Bun returns as null → normalize to 0.
  • 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 a u64 with the float bits in the low 32 bits and has_value in 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 UIexamples/server.ts over Bun.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 --srt for 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. --look trades the two. examples/parakeet-live.ts plays 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 MB
    

    Not 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:

    audiobatch WERstream WER
    64 s10.7%26.2%
    94 s20.7%30.5%
    138 s34.0%27.7%
    251 s33.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 changed
    

    Batch 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), multilingualsrc/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 to mlx_whisper (tests/whisper-transcribe-test.ts). Runs large-v3-turbo with 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, .wav out. 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. -small is 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 .wav out, no phonemizer and no espeak-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 where torch.stft centres 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.ts clones 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 RAGPOST /v1/embeddings returns 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 scratchvalidation/spike-microgpt.ts builds 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-rolled Value engine replaced by value_and_grad); step-0 loss is exact vs the MLX-Python mirror, both converge.

  • Train a real small GPTvalidation/spike-nanogpt.ts scales 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 exact shakespeare-char config (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 vs reference/reference-nanogpt.py (shared init + batches).

  • Run real GPT-2-124Msrc/models/gpt2.ts loads the actual OpenAI gpt2 weights 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 vs reference/reference-gpt2.py at ~210 tok/s. See docs/GPT2.md.

  • SFT a chatbottraining/sft.ts full-fine-tunes real GPT-2-124M into an instruction follower (chat format + completion-only loss), the nanochat chat stage. Step-0 loss matches reference/reference-sft.py; after SFT it answers in-format, including a held-out question (Italy → Rome). See docs/SFT.md.

  • RL with GRPOtraining/rl.ts runs 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 vs reference/reference-rl.py. See docs/RL.md.

  • Train a tokenizerreference/tok-train.py trains a byte-level BPE in native Rust (HF tokenizers, as nanochat does — training is a data-prep boundary step, not MLX compute); our pure-TS src/text/tokenizer.ts then reproduces it token-exact (tests/tok-train-test.ts). The tok_train stage of a nanochat-style pipeline.

  • Pretrain + checkpointtraining/base-train.ts pretrains 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. The base_train stage.

  • The whole pipeline (scripts/run.sh) — the TS-over-MLX analogue of nanochat's runcpu.sh: dataset → tokenizer → data-prep → pretrain → SFT → chat, end to end on one Mac. Trains on TinyStories (coherent at this scale) via a streaming dataloadertraining/data-prep.ts stream-encodes the corpus to uint16 token shards, training/base-train.ts Bun.mmaps them (scales past RAM). Then SFT + chat (CLI training/chat-ckpt.ts or web UI examples/chat-web.ts). See docs/PIPELINE.md.

  • Research / inspection — pull logits, hidden states; the MX op 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 decode

OLMoE-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 nn modules + weight-key mapping (src/models/olmoe.ts / src/models/whisper.ts are 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.

Contributors

nielspeter

111 commits

nielspeter/mlx-ts

A TypeScript MLX SDK over Apple's mlx-c via FFI. Runs on Bun, Deno and Node; output matches Apple's mlx-lm token for token. Apple Silicon only.

0

stars

111

commits

TypeScript

primary language

Aug 29, 2026

updated

www.npmjs.com/package/@nielspeter/mlx-ts
apple-silicon
bun
deno
ffi
inference
llm
machine-learning
metal
mlx
mlx-lm
typescript

README

mlx-ts — TypeScript → mlx-c → Metal

npm version CI License: MIT Node.js >= 24 Apple Silicon

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 MX is a pointer to GPU memory.
  • to build the model, not just call one. src/models/ are ordinary TypeScript files composing nn Modules. A new architecture is a forward pass, not a new binary.
  • training or fine-tuning from TypeScript. Real value_and_grad over 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-lm token-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 kHz

Cloning 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 store

embeddingMX 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 Macnot this — Apple Silicon + Metal only
a chat endpoint with the least possible workOllama or LM Studio; a server is less code than an SDK
the same thing in Pythonmlx-lm — same engine, far more models, maintained by Apple
GGUF, AWQ/GPTQ, or the widest model coveragethe llama.cpp ecosystem (e.g. node-llama-cpp)
to run in a browsertransformers.js — ONNX/WebGPU, a different engine entirely
high-throughput multi-tenant servingnot 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:

Pythonhere
mlx (mlx.core) — arrays, ops, memorysrc/ffi + src/core
mlx.nn / mlx.optimizerssrc/nn — Modules, Adam, valueAndGrad
mlx-lm — architectures, generation, tokenizerssrc/models + src/text
mlx-whisper — speech-to-textsrc/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 ids

src/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 op

examples/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`
runtimebindingpointer dispatchzero-copy readback
Bunbun:ffi (built in)~12 ns/calltoArrayBuffer
DenoDeno.dlopen (built in)~3 ns/callUnsafePointerView
Nodekoffi (prebuilt addon)~21 ns/callkoffi.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 demand

Nothing 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 over libmlxc.dylib: handle management plus matmul, rms_norm, rope, sdpa, silu, etc.
  • tools/codegen.tsparses the mlx-c headers and emits src/ffi/generated.ts: a full FFI symbol table (472 entries) + 242 typed op wrappers. The hand-written validation/mlx.ts exists only to bootstrap; src/ffi/generated.ts supersedes it.
  • validation/block.ts — a full Qwen3 decoder block forward pass (mirrors mlx-lm's qwen3.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 over mlx_load_safetensors: open a file into a string -> array map, 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.ts4-bit quantized path: Python quantizes the Linear projections (mx.quantize) and saves weight/scales/ biases; TS loads them and runs the decode with quantizedMatmul.
  • src/text/tokenizer.ts — pure-TS byte-level BPE tokenizer (the real Qwen3 models/tokenizer.json); reference/tok-reference.py / tests/tok-test.ts validate it against HF tokenizers (encode + decode, 11/11 cases).
  • src/models/qwen.ts / reference/reference-qwen.pyconfig-driven real Qwen3-0.6B (bf16): reads models/config.json, loads models/model-qwen.safetensors (HF key names), generates text; ids match MLX Python token-for-token.

Production runtime (mx + nn)

  • src/core/mx.tsMX array class: each wraps one handle, auto-freed by a FinalizationRegistry, plus a tidy() arena for deterministic freeing, ops, and temp/top-p sampling.
  • src/nn/nn.tsModule, Linear, QuantizedLinear, RMSNorm, Embedding, QuantizedEmbedding.
  • src/models/qwen-nn.ts / reference/reference-qwen-q4.pyreal 4-bit Qwen3-0.6B (mlx-community format) over nn.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-agnostic Decoder interface and async-generator streamTokens / streamText / generate. The KV cache is freed automatically (completion / early break / throw), so callers never call tidy() or free a handle; MX is Disposable. examples/stream.ts is the live demo.
  • tests/validate-prod.ts — checks sampling reproducibility, batching, and bounded memory. tests/stream-test.ts — stream output is identical to generate().
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 three

examples/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 file=@audio.flac   # -> {"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 reference

All three blocks print the same fingerprint:

  sum    = 0.005793
  sum_sq = 0.162600

The 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 pass

Validated 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 as ptr (a JS number).
  • An empty handle has ctx == NULL, which Bun returns as null → normalize to 0.
  • 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 a u64 with the float bits in the low 32 bits and has_value in 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 UIexamples/server.ts over Bun.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 --srt for 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. --look trades the two. examples/parakeet-live.ts plays 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 MB
    

    Not 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:

    audiobatch WERstream WER
    64 s10.7%26.2%
    94 s20.7%30.5%
    138 s34.0%27.7%
    251 s33.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 changed
    

    Batch 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), multilingualsrc/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 to mlx_whisper (tests/whisper-transcribe-test.ts). Runs large-v3-turbo with 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, .wav out. 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. -small is 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 .wav out, no phonemizer and no espeak-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 where torch.stft centres 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.ts clones 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 RAGPOST /v1/embeddings returns 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 scratchvalidation/spike-microgpt.ts builds 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-rolled Value engine replaced by value_and_grad); step-0 loss is exact vs the MLX-Python mirror, both converge.

  • Train a real small GPTvalidation/spike-nanogpt.ts scales 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 exact shakespeare-char config (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 vs reference/reference-nanogpt.py (shared init + batches).

  • Run real GPT-2-124Msrc/models/gpt2.ts loads the actual OpenAI gpt2 weights 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 vs reference/reference-gpt2.py at ~210 tok/s. See docs/GPT2.md.

  • SFT a chatbottraining/sft.ts full-fine-tunes real GPT-2-124M into an instruction follower (chat format + completion-only loss), the nanochat chat stage. Step-0 loss matches reference/reference-sft.py; after SFT it answers in-format, including a held-out question (Italy → Rome). See docs/SFT.md.

  • RL with GRPOtraining/rl.ts runs 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 vs reference/reference-rl.py. See docs/RL.md.

  • Train a tokenizerreference/tok-train.py trains a byte-level BPE in native Rust (HF tokenizers, as nanochat does — training is a data-prep boundary step, not MLX compute); our pure-TS src/text/tokenizer.ts then reproduces it token-exact (tests/tok-train-test.ts). The tok_train stage of a nanochat-style pipeline.

  • Pretrain + checkpointtraining/base-train.ts pretrains 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. The base_train stage.

  • The whole pipeline (scripts/run.sh) — the TS-over-MLX analogue of nanochat's runcpu.sh: dataset → tokenizer → data-prep → pretrain → SFT → chat, end to end on one Mac. Trains on TinyStories (coherent at this scale) via a streaming dataloadertraining/data-prep.ts stream-encodes the corpus to uint16 token shards, training/base-train.ts Bun.mmaps them (scales past RAM). Then SFT + chat (CLI training/chat-ckpt.ts or web UI examples/chat-web.ts). See docs/PIPELINE.md.

  • Research / inspection — pull logits, hidden states; the MX op 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 decode

OLMoE-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 nn modules + weight-key mapping (src/models/olmoe.ts / src/models/whisper.ts are 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.

Contributors

nielspeter

111 commits

Languages

TypeScript

79.8%

Python

13.6%

Shell

6.1%