The inference engine — Mojo 1.0 on Metal
In the first post I described millfolio’s split: a frontier model writes a small program over an aliased schema, and a local model reads your actual files on-device. The second post was about spending that local model’s time wisely. This post is about the thing both of them take for granted: the inference engine that runs the local model at all.
Where it fits
A combined inference server serves both chat and embeddings over an OpenAI-compatible HTTP API, bound to localhost only. The local server provides embedding for the indexing process. It also classifies years of transactions in batches, generating yes/no judgments for AI tags. Its main integration point with the 3rd party model is by answering targeted question about one record. The local inference server is the only component that ever reads your data with a model — and it never listens on anything but the loopback interface.
mill install provisions it alongside the rest of the stack, and the web app can download and switch between the models it serves — currently the Qwen2.5 family (0.5B and 3B, the 3B is the default) and Gemma 4 (the dense 12B plus the effective-2B “E2B” and effective-4B “E4B” variants, whose per-layer embeddings and KV-sharing make them run like much smaller models; all quantized to int4 at load). The embedding model is Qwen3-Embedding-0.6B. The family, dimensions, and chat template are auto-detected from the checkpoint.
Mojo all the way down
Most local-LLM stacks are a thin scripting layer over a C++/Metal core (MLX) or a C/C++ engine (llama.cpp under Ollama). millfolio’s engine is written from scratch in Mojo. Every GPU kernel — matmul, attention, RMSNorm, RoPE, SwiGLU, the int4 dequant — is custom-written Mojo, reaching Apple’s simdgroup_matrix tensor units through the AIR llvm.air.simdgroup_matrix_8x8_multiply_accumulate intrinsic. There are no C++, CUDA, or Metal-shader dependencies, and no Python on the request path: sockets → tokenizer → GPU forward pass → JSON is one language. The host-side GPU plumbing — the device context that launches kernels, the tensor views they take — comes from Modular’s toolchain (the same packages the compiler ships with). No third-party GPU kernels run in the engine. Of course writing GPU code is the reason why Modular created Mojo and they have proven the language in a production stack. One difference with the Modular production stack is that millfolio does not have a MAX/graph layer.
What are the benefits?
The first is auditability. This is a privacy product, and the component that reads your data should be the most readable one. A small codebase in one language, where every kernel that touches your data is in-tree rather than vendored from a GPU library, is a surface you (or an AI reviewer) can actually walk end to end. This is also the reason we have skipped the MAX layer (currently implemented in Python).
The second is the experiment itself: how far does a compiled, Python-ish, GPU-capable language get you as a more general backend language? For the GPU code correctness is enforced the boring way — the GPU output is diffed against CPU references (HF transformers per-kernel, whole-model oracles), and the bar is token-for-token parity with the reference under greedy decoding.
Of course the audience for Millfolio is different from the MAX audience. The requests are controlled by the application serving one user so the requirements are vastly different.
Crossing Mojo 1.0
Millfolio’s engine switched from the 2026 nightlies to Mojo 1.0.0rc0 + MAX 26.5 the week the release candidate shipped: thirteen repos ported, zero deprecation warnings. Two things about that jump mattered far more than the syntax churn.
The launch-overhead wall fell. The July analysis found our decode kernels at per-op parity with MLX and blamed the rest on launch architecture: Mojo’s DeviceContext committed every kernel individually, tens of microseconds each, hundreds of times per token, while MLX batches each token’s ops into one command buffer. We filed the request upstream with measurements. The 26.5 Metal driver shipped exactly that — transparent command-buffer batching, no API change. The wall we’d been pushing against was simply gone.
The driver improved significantly . Submitting work to the GPU was a bottleneck pre-1.0 and Claude worked around it by introduction K-split decode GEMVs. The winning post-1.0 shape is a multi-row GEMV — each simdgroup owns several output rows, several independent 256-bit weight streams in flight per lane — plus on-device token selection (argmax/top-k on the GPU, saving a 600 KB logits download per token). Decode went 23.9 → ~40 tok/s on the real 3B: +66% over the best pre-1.0 number, and the M4’s weight-streaming is now the limit.
Prefill got its own 1.0 beautification: instead of dequantizing int4 inside the GEMM’s hot loop, the engine now dequantizes each weight once per prefill into a transposed half-precision scratch and runs a half-precision simdgroup_matrix GEMM (f16 — the dtype story is below; the same tiling Modular’s open Apache-2.0 kernel library uses on M1–M4). That, plus a TTFT accounting fix described below, took the 1570-token prefill from 5.7 s to 3.64 s at bit-identical answer quality. The GEMM itself now runs at ~85% of the M4’s ALU ceiling . A final pass hoisted the tile edge guards out of the GEMM’s hot loop — worth ~8% on the M4 and, it turned out, ~34% on M2-generation GPUs, whose older cores pay more for everything.
Put together, here is what crossing 1.0 and re-tuning for its driver bought, same model, same harness:
| Qwen2.5-3B | July (pre-1.0 nightly) | August (Mojo 1.0.0rc0 \n\n+ MAX 26.5) |
|---|---|---|
| M4 decode | 23.9 tok/s | 37.0 tok/s (+55%) |
| M4 prefill, 1570-token prompt | 5.7 s | 3.49 s (1.63× faster) |
| M4 batch-classify (the product workload) | 11.6 s | 8.6 s |
| M2 Pro decode | 10.8 tok/s | 32.7 tok/s (3.0×) |
Against MLX, today
The engine is a learning/research engine, there is still a gap with MLX. The cleanest comparison is both engines on the same machine, same model, same day, same harness (median of 5, one engine resident at a time, temperature 0) — Qwen2.5-3B, millfolio group-128 int4 vs a current mlx-lm 4-bit, measured 2026-08-09:
| metric (M4, Qwen2.5-3B) | millfolio int4 | MLX 4-bit |
|---|---|---|
| prefill, 71-token prompt | 202 ms | 221 ms |
| prefill, 1570-token prompt | 3 492 ms | 2 751 ms |
| decode (short / medium / long-code) | 36.6 / 37.0 / 29.9 tok/s | 51.5 / 51.9 / 48.0 tok/s |
| warm-prefix TTFT (cold → warm) | 3.51 s → 0.50 s | 2.82 s → 0.19 s |
| batch-classify, 40 snippets ×4 chunks | 8.6 s (4.6 distinct/s) | 4.9 s (8.2 distinct/s) |
The last row is the product workload — the tags post’s batched yes/no classification, chunked under a shared instruction prefix. The honest summary today: MLX is ~1.8× faster on that workload, ~1.4× on decode, ~1.27× on long prefill — and millfolio wins short-prompt prefill outright. (The engine also serves Gemma 4’s E2B/E4B/12B, which mlx_lm doesn’t support yet.)
Here are the results on an Apple M2 Pro (~200 GB/s to the M4’s 120):
| metric (M2 Pro, Qwen2.5-3B) | millfolio int4 | MLX 4-bit |
|---|---|---|
| prefill, 71-token prompt | 255 ms | 207 ms |
| prefill, 1570-token prompt | 3 273 ms | 2 267 ms |
| decode (short / medium / long-code) | 32.7 / 32.4 / 27.7 tok/s | 77.4 / 77.1 / 70.0 tok/s |
| warm-prefix TTFT (cold → warm) | 3.30 s → 0.53 s | 2.31 s → 0.18 s |
| batch-classify, 40 snippets ×4 chunks | 9.5 s (4.2 distinct/s) | 4.3 s (9.3 distinct/s) |
What’s inside the engine today:
- Tensor-core prefill. GEMM and attention on the 8×8
simdgroup_matrixunits, with the dequant-once f16 pipeline above for long prompts. - Multi-row decode GEMV. Several output rows per simdgroup keep independent weight streams in flight — the shape the batching driver rewards.
- int4 everywhere it counts. Group-128 projections and the LM head (which was reading 622 MB of bf16 per token; quantizing it cost nothing measurable in quality — the perplexity deep dive measures what int4 costs).
- GPU token selection. Greedy, top-k, and speculative-verify argmax run on-device, bit-exact against the host reference.
- Prompt-lookup speculative decode. A zero-cost draft verified greedily — token-identical output, ~1.35× on the code-heavy generations the sandbox programs produce.
- Persistent caches. The disk-backed KV prefix cache (the 8× warm prefill in the table) and a Metal pipeline cache, so warm starts recompile nothing.
What do these numbers mean in product terms? The workload that matters is the one from the tags post: batched yes/no classification over a vault’s transaction descriptions. A July run of the aggregates-only classifier eval — the same 3B — classified 400 distinct descriptions (covering 2,228 rows of a 2,930-row vault, a 3× dedup fan-out) in 257 s: ~642 ms per distinct description, 8.7 rows/s effective; with the post-1.0 decode rate the same sweep prices out well under half that. On quality, scored against the deterministic keyword rule it augments: 100% recall (nothing the rule tags is missed) at 39% precision — the model flags roughly 2.5× more candidates than the literal keyword match does. Keyword tags are imperfect truth, so a disagreement is model-vs-rule, not a certified model error — high recall with over-flagging is the right shape for a tagger whose output you can inspect and prune.
The shape of it
What the from-scratch stack buys: every GPU kernel in-tree, one language end to end, models the incumbents don’t serve yet (Gemma 4’s E2B/E4B/12B), a GEMM at the silicon’s measured ceiling, and a codebase small enough that when the last 2× wouldn’t yield, the audit could go below the kernels, into the runtime ABI, and name the exact missing method — which the very next driver release then shipped. As this project evolved and improved with Mojo/Max releases the whole stack stayed legible enough to re-tune across a major toolchain jump in days and to locate every remaining millisecond.
Appendix: the remaining gap, itemized
For readers who want the ledger rather than the story — what’s left between these tables and MLX’s, why each piece exists, and what it would cost to close.
Why not just port MLX’s kernels verbatim? MLX is MIT-licensed, so the obstacle isn’t legal. It’s three technical facts and one choice. First, MLX’s kernels are Metal Shading Language C++ built on simdgroup_matrix — an opaque, compiler-managed type with builtins like simdgroup_load and simdgroup_multiply_accumulate. Mojo reaches the same hardware through raw AIR intrinsics, and every builtin has to be reverse-engineered individually: the multiply-accumulate has been (that’s why our GEMMs work), and — since this appendix was first written — so has the hardware matrix-load. We decoded it, called it from Mojo, verified its register layout bit-exact against ours, and raced it: from device memory it is ~14% slower than our per-lane loads, and the full MLX staging pattern built on it loses ~30% on the M4 and crashes Apple’s shader compiler outright on M2-generation targets. So the “their compiler knows something” theory is now a measured negative — our direct-load GEMM equals MLX’s whole-prefill effective rate on the M4 at the kernel level, and what looked like their load-instruction advantage was never the story. Second, the quantization formats differ (their affine 4-bit vs our symmetric group-128), so a “literal” port still rewrites the entire inner unpack. Third, Apple’s Metal compiler does invisible register-allocation and scheduling work on the opaque type that a hand translation must reproduce from evidence. And the choice: we could compile MLX’s .metal source and load the metallib — but then the kernel that reads your data is no longer in-tree, in one language, auditable end to end, and that property is the point of this engine.
Is the Metal compiler open source? No — it’s a proprietary Clang/LLVM fork, and the final AIR-to-GPU-machine-code backend lives closed inside the driver. But AIR itself is LLVM bitcode, and Apple ships disassembly tools with the Metal Toolchain — so decoding simdgroup_load was a shipped-tools exercise: compile a five-line MSL kernel that uses it, ask the compiler for textual AIR, read off the intrinsic — which is exactly how the results above were obtained, one probe and one afternoon after writing this appendix. (The Asahi Linux/Mesa project’s reverse-engineered open driver is the reference for what the hardware does below that.)
The priced list, in order:
Decode theDone, and retired: decoded, layout-verified, raced, and beaten by our own loads on every configuration that compiles (see above). Its lasting value is the method — the AIR intrinsic table is now ours to extend.simdgroup_matrixload intrinsic.- Re-test the old negatives under the new driver. Two July verdicts — “f16 KV loses on decode” and “the qmv-shaped loop loses” — were measured under the pre-batching driver, and this month overturned four findings of that vintage. f16 KV already measured positive on prefill; splitting the cache dtypes captures that side regardless.
- The M2-class GEMV ceiling: single-warp routing took that machine from 16% to ~50% of memory bandwidth; MLX’s qmv reaches ~70%. The retest in (2) is the honest next experiment.
- The launch floor. On the M4, rerouting a projection to a 2.4× faster kernel changes nothing end-to-end — the per-step floor of ~250 small launches is the binding constraint, on-Metal fusion buys only 1–4%, and the real fix (explicit command-buffer control) lives below the toolchain surface we rent. This is the one item we can name but not reach.
Everything cheap and structural has been taken — including, now, the one instruction we couldn’t previously name. What remains is two possibly-stale negatives and a floor we’ve measured to the microsecond.
- The engine page: millfolio.app/server
- Demo: demo.millfolio.app
- Code: github.com/millfolio