oMLX Over MLX — Complete Visual Guide

What each layer owns · which one oMLX actually replaces · and three features that only look finished · measured at HEAD, 22 Aug 2026

The Stack

Three Layers, One Uncomfortable Question

“Why does oMLX exist on top of MLX” is a trivial question — MLX is an array framework, it cannot serve HTTP. “Why does oMLX exist on top of mlx-lm” is the real one, because mlx-lm already ships a working OpenAI-compatible server, and oMLX runs that server, unmodified, on every node of a multi-Mac cluster. Answering it properly means laying out who owns what first.

The Three Layers — What Each One Owns

How Tightly These Are Bolted Together

The MLX checkout contains four Python subpackages — nn, optimizers, utils, _distributed_utils. That is PyTorch’s role in this stack, not vLLM’s.

oMLX’s dependency on mlx-lm is not a loose one. pyproject.toml pins a commit, not a release, and the build pins a nanobind ABI:

# dependencies
"mlx==0.32.0",
"mlx-lm @ git+https://github.com/ml-explore/mlx-lm@ab1806e8f5d6aa035973af194a1b9198ab4754dc",

# build-system — comment reproduced from pyproject.toml
# A mismatched nanobind isolates the `mlx` NB_DOMAIN, so custom kernel
# extensions (omlx.custom_kernels.*) reject every mlx.core.array at the
# type caster with `incompatible function arguments`.
"nanobind==2.13.0",

A commit hash plus an ABI-matched build pin plus a paragraph of justification is what coupling to private internals looks like from the outside.

One Request On One Mac

Follow a single chat completion and the division of labour becomes concrete: oMLX owns the front door, the queue and the cache; mlx-lm owns the batched decode; MLX owns the arithmetic.

Single-Mac Request Path — Who Handles What

The Cluster Path Is The Real Surprise

You would expect the multi-Mac path to be the most rewritten part of all. It is the least. Every rank in a distributed deployment is an mlx_lm.server process; oMLX plans, verifies, launches and front-ends them.

Multi-Mac Deployment — Each Rank Is An mlx_lm.server

Read The Imports, Not The README

# omlx/cluster/inference_worker.py:944
import mlx_lm.server as mlx_server
from mlx_lm.server import ModelProvider, ResponseGenerator, run

# omlx/cluster/launch.py:1414 — emitted verbatim into the remote launch script
"import mlx_lm.server\n"

# omlx/cluster/prompt_snapshot_cache.py:4
"A distributed rank runs the pinned mlx_lm.server with a rank-local ..."

# omlx/cluster/autoconfigure.py:1132
"for_vlm defaults to False because a rank is an mlx_lm.server ..."

This is the right call, and it is the strongest evidence for the argument that follows: if the serving layer had been the thing worth replacing, oMLX would not be launching mlx-lm’s server on its most demanding code path.

Layer Ownership, Summarised

Concern	MLX	mlx-lm	oMLX
Tensor Math & Metal Kernels	Owns	—	Adds 3,145 Lines Of Custom Kernels
Model Architectures	—	Owns	Patches 49 Modules
Batched Decode Loop	—	Owns	Consumes It
Queue & Admission Policy	—	Minimal	Owns
KV Cache Across Requests	—	—	Owns
Cache Persisted To Disk	—	—	Owns
Multi-Model In One Process	—	—	Owns
OpenAI HTTP Surface	—	Basic	Owns
Anthropic / Embeddings / Rerank / Audio / MCP	—	—	Owns
Distributed Collectives	Owns	Uses Them	Orchestrates
Serving On A Cluster Rank	—	Owns	Launches It

The Cache

The One Subsystem That Carries The Fork

mlx-lm ships a server, and oMLX runs it on every cluster rank. mlx-lm ships a batched decode loop, and oMLX calls it. So what is genuinely replaced? The KV cache. mlx-lm’s KVCache and RotatingKVCache live for the duration of one generation, in RAM. Nothing is reused across requests and nothing survives a restart.

oMLX States The Layering Itself

There is a --no-cache mode. Its startup banner is the cleanest description of the fork anywhere in the repository:

# omlx/cli.py:320
Mode: Multi-model serving (no oMLX cache, mlx-lm BatchGenerator only)

Remove the cache and what remains is described, by the project, as mlx-lm.

What It Buys You

A chat client resends a growing conversation every turn. Without cross-request reuse, the entire prefix is reprocessed on every message — the cost grows with the conversation, not with what you typed.

Prefix Reuse Across Conversation Turns — Blocks Computed Versus Restored

How The Block Pool Works

omlx/cache/paged_cache.py:2 names its source outright — “Adapted from vllm-mlx … following vLLM’s architecture (vllm/v1/core/block_pool.py), adapted for MLX on Apple Silicon.” Blocks are fixed-size, refcounted, and addressed by a hash that chains to the parent block.

Block Pool — Reference Counting, Chain Hashing, Copy On Write

Two Tiers, And One That Survives A Reboot

When the hot tier fills, blocks are not dropped — they are serialised to SSD as safetensors and can be restored on a later request with a matching prefix, including after the server has been restarted.

Prefix Lookup — Three Outcomes, Two Of Them Cheap

The Five Mechanisms

🧩

Fixed-Size Refcounted Blocks

64 tokens per block by default (paged_cache.py:505). Two conversations sharing a system prompt share the same physical blocks, and the refcount is what makes that safe.

🔗

Chain Hashing

Each block’s hash depends on its parent’s, so identifying the longest shared prefix is a hash-map walk rather than a token-by-token comparison.

✂️

Copy On Write

When two branches diverge mid-block, only that block is duplicated. Everything before it stays shared.

💾

SSD Cold Tier

paged_ssd_cache.py and boundary_snapshot_store.py serialise evicted blocks to safetensors on the inference thread — Metal-safe — buffer them in _pending_writes for instant read-back, and flush from a background writer thread.

♻️

O(1) LRU Eviction

A doubly linked free-block queue, lifted straight from vLLM’s block pool design.

Everything Else

Everything Else oMLX Adds

The cache is the argument. This tab is the rest of the surface area — what is genuinely load-bearing, what is product polish, and the third of the codebase that exists purely because the layer beneath is unfinished.

Multi-Model Serving In One Process

mlx-lm’s ModelProvider serves one model. oMLX’s engine_pool.py — 2,928 lines — runs several under a single process-memory ceiling, arbitrating between pinned models, active models, and models whose idle TTL has expired.

Engine Pool — Memory Arbitration Under A Hard Ceiling

The Enforcer Is Tier-Aware, Not A Single Number

The README says “system RAM minus 8 GB”. The code is more careful than that — the reserve is tiered, and the fraction of macOS “active” pages counted as reclaimable changes with the tier, because the compressor typically recovers 2–3×:

# omlx/process_memory_enforcer.py
_STATIC_RESERVE_LARGE = {"safe": 8*GiB, "balanced": 6*GiB,
                         "aggressive": 4*GiB, "custom": 2*GiB}
_ACTIVE_RECLAIM_RATIO = {"safe": 0.2, "balanced": 0.5, "aggressive": 0.8}
_SOFT_THRESHOLD_BY_TIER = {"safe": 0.85, "balanced": 0.90, "aggressive": 0.925}

# Systems under 24 GiB get a smaller reserve so they can still load
# something useful after the tier-scaled cut.
_SMALL_SYSTEM_RESERVE   = 4*GiB
_SMALL_SYSTEM_THRESHOLD = 24*GiB

The comment on _ACTIVE_RECLAIM_RATIO is worth reading in full: “0.8 pushes into swap territory.” That is somebody who has watched this fail.

The Cluster Approval Contract

The best-engineered thing in the cluster code is not the sharding. It is a refusal. Activation recomputes the plan signature and compares it against the one the operator approved — because for months preview and activation posted different node objects, and the stage that launched was not the stage anybody had agreed to.

Plan Approval — Activation Refuses A Plan That Changed Under It

Why That Gate Exists

The comment above it in omlx/cluster/routes.py is unusually honest for a source file:

# Preview and activation posted different node objects for months, and
# the second one dropped the split cap and the role, so the stage that
# launched was not the stage anybody approved. Refusing here costs
# nothing; discovering it after every rank has staged its weights
# costs the machine.

Activation then goes on to check peer liveness before any rank starts, because “launching into a dead peer produces a collective that blocks forever rather than an error.” This is the part of oMLX that reads like it was written by someone who had been burned. Hold that thought — tabs 4 and 5 are about the parts that were not.

Where The Python Actually Lives

Lines Of Python Per Subpackage

Anatomy Of patches/

Forty-nine modules, 51,395 lines — larger than the cluster, the dashboard or the cache. None of it is architecture. All of it is closing a gap in the layer below, against private internals.

Patch Modules Grouped By The Upstream Gap They Close

oMLX’s Own Comments Quantify The Gap

# omlx/cluster/planner.py:752
Only about seven of mlx-lm's ~120 architectures can be pipelined

# omlx/cluster/planner.py:1089
stock mlx-lm has no file for it

# omlx/patches/qwen3_sliding_window/__init__.py:5
a model_type that mlx-lm doesn't ship at all yet

# omlx/cluster/catalogue.py:474
This model cannot combine Macs because mlx-lm supports neither ...

# omlx/patches/minimax_m3_mlx_lm/__init__.py:5
every rank is an mlx_lm.server, and pinned mlx-lm has no ...

mlx-lm’s model coverage is broad — roughly 120 architectures. Its serving-grade coverage — shardable, batchable, still correct once a cache has been restored from disk — is narrow. Around seven can be pipelined at all.

The Rest Of The Surface

API Compatibility — 12,339 Lines

omlx/api/ carries OpenAI chat and completions, Anthropic Messages, embeddings, rerankers, audio, MCP routes, tool-call parsers per model family, and grammar-constrained decoding. Real work, but nothing here needed a fork of the inference stack — it needed an HTTP layer.

Admin Dashboard — 17,679 Lines

Monitoring, model management, chat, benchmarks, per-model settings, eight languages, and every CDN dependency vendored so it works offline. The single most visible part of the product and the least technically deep.

Custom Metal Kernels — 3,145 Lines

Compiled extensions for GLM-5.2, MiniMax M3 and Qwen 3.5 families, ABI-coupled to mlx==0.32.0. This is the one place oMLX reaches past mlx-lm and straight into MLX.

Cluster — 29,382 Lines

Peer discovery, SSH and runtime verification, unequal shard planning, link classification, the approval contract above, and a downloadable diagnostics bundle. Excellent in places — and the subject of tab 5’s second finding.

The SLO Postmortem

An SRE Subsystem That Never Measured Anything

Five modules, 1,028 lines of production code, 1,050 lines of tests behind 67 test cases, three HTTP endpoints and three dashboard sections. SLO definitions, a rolling-window tracker, error budgets, a deployment gate, a capacity planner. Every test passes. Every endpoint responds. It has never produced a single real reading.

The SLO Pipeline — Where It Breaks

One Missing Line Explains All Of It

SLOTracker.record() at omlx/cluster/slo_tracker.py:90 is called from nowhere inside omlx/. Only from tests/. Every rolling window is permanently empty, so every evaluation takes the zero-sample early return:

# omlx/cluster/slo_tracker.py:152
if sample_count == 0:
    return SLOEvaluation(
        slo=slo, status=SLOStatus.HEALTHY, compliance_pct=100.0,
        burn_rate=0.0, window_seconds=slo.window_seconds, sample_count=0,
    )  # <-- taken on every call, on every SLO, forever

The tracker’s docstring names the gap: “When SampleStore from #2773 is available, callers can bridge data into these windows.” SampleStore does exist, at omlx/cluster/sample_store.py:20. The bridge was never written, and nothing in the test suite could notice — every test injects samples by hand.

Four Defects Underneath The Missing Wire

Seconds Subtracted From One Hundred

error_budget.py:52 computes error_budget_allowed_pct = 100.0 - slo.target. For latency_p95_small the target is 2.0 seconds, producing a “98% error budget”. For throughput_min_small it is 30.0 tokens/sec, producing 70%. Units and percentages, subtracted from each other.

A Branch With Two Identical Arms

slo_tracker.py:211-219 branches on metric_kind == LATENCY and then runs byte-identical code in both arms. error_budget_pct, computed on line 212 with a conditional expression, is never read again.

Ten Objectives, Five Windows

latency_p50_small and latency_p50_large both declare metric_key="latency_p50" and therefore share one rolling window. Same for p95, p99, and both throughput objectives. The small-versus-large split is fiction: fed real data, one is guaranteed to breach whenever the other is healthy.

Two Definitions Of Compliance

For latency, status is decided by a percentile of the window, while compliance_pct counts individual samples under target — and the burn rate is derived from the second. A card can read healthy beside a compliance figure that disagrees with it.

And One Feature Behind A 404

The capacity planner is the cleanest example of the pattern: 435 lines of production code, 17 tests, a full dashboard section — and two missing connections, either of which alone would have made it invisible.

Capacity Planner — Two Broken Connections

Was Any Of It Required?

An error budget exists to arbitrate one specific conflict: development wants to ship, operations wants stability, so both sides agree on a number that decides who wins. On a two-Mac local inference server with one operator, you are both parties. There is nobody to say no to, and a budget whose only function is to say no to somebody has no function.

The numbers make the same point. 99.9% availability over a seven-day window is ten minutes of downtime a week, on machines the operator reboots by hand. And the targets — 0.5s p50, 30 tok/s, 70% cache hit rate — are guessed constants. They were not measured from a baseline and not negotiated with a consumer. An SLO with no consequence attached is a number on a card.

The instinct was right and the layer was wrong. What this work should have produced is latency and tokens/sec histograms on the /metrics endpoint that already exists, scraped by the Prometheus already running, with alert rules shipped as a PrometheusRule. Identical visibility, no bespoke engine, and it plugs into alerting that already works. Two thousand lines here reimplement a worse Prometheus inside the server process.

What The Tests Could Not Catch

Sixty-seven tests pass against this code. Not one of them fails, and not one of them would have caught the actual defect, because every test constructs a tracker and injects samples directly:

# tests/test_slo_tracker.py — the shape of all 25 tests in this file
tracker = SLOTracker()
tracker.record("latency_p95", 1.2)     # the only caller of record()
assert tracker.evaluate(slo).status == SLOStatus.HEALTHY

The unit under test works. The wire into it was never soldered, and a unit test by construction cannot see a missing wire. A test that only ever sees data it supplied itself proves the code runs — not that the feature works.

Auth & Thunderbolt

Two More Features That Only Look Finished

The SLO subsystem is green because it measured nothing. These two fail differently: the auth check returns True without knowing who asked, and the Thunderbolt classifier gives confident, actionable, impossible advice. Both are small fixes. Both are invisible until you read the code rather than the dashboard.

1 · Authentication With No Principal

The question that started this was ordinary: can oMLX give a team of developers limited local access — Keycloak groups, AD, per-user model allowlists? The answer is that there is nowhere to hang a role, because a request never becomes a someone.

Authentication — Every Caller Collapses To A Boolean

The Loop Knows, Then Forgets

# omlx/admin/auth.py:181
def verify_any_api_key(api_key: str, main_key: str, sub_keys: list) -> bool:
    if not api_key:
        return False
    if main_key and compare_keys(api_key, main_key):
        return True
    for sk in sub_keys:
        if sk.key and compare_keys(api_key, sk.key):
            return True          # sk is in scope here — and thrown away
    return False

Two lines from the end of that loop, the code knows exactly which named sub key matched. It returns a boolean instead. Downstream there is no request principal, so there is nothing for a rate limit, a quota, an audit line or a model allowlist to attach to.

The data model confirms it. SubKeyEntry at settings.py:549 carries only key, name, created_at — no scopes, no roles, no expiry, no model restrictions, no quota. The main key doubles as the admin dashboard password, and sub keys cannot log into the dashboard at all (admin/routes.py:1499).

Capability	Present?	Evidence
Shared secret auth	Yes	Constant-time compare, SHA-256 fingerprint on rejection
Multiple keys	Yes	sub_keys list, API-only
Per-request principal	No	verify_api_key returns bool
Roles or scopes	No	SubKeyEntry has three fields
Per-key model allowlist	No	Nowhere to store it
Quotas or budgets	No	No principal to meter
OIDC / Keycloak / AD / LDAP / SAML	No	Zero repo-wide hits outside vendored JavaScript
Auth on by default	No	server.py:312 — no key configured means every request passes

2 · Thunderbolt 4 Treated As Broken Thunderbolt 5

The symptom is a warning nobody can resolve: “RDMA is enabled, but ibv_devices reported no RDMA interfaces.” It is not a misconfiguration. RDMA over Thunderbolt requires Thunderbolt 5 — MLX’s own distributed.rst states it outright — while rdma_ctl enable sets a global flag on any Mac. A TB4 machine therefore reports enabled with an empty device list, permanently, by design.

Backend Selection — Both Branches Reach The Same Answer

Four Defects Behind One Warning

1

A Backend The Hardware Cannot Open

autoconfigure.py:374 returns "jaccl-ring" with the reason “Thunderbolt link detected, without RDMA”. But inference_worker.py:961 maps any jaccl* backend onto mx.distributed.init(backend="jaccl"), and MLX’s configure_jaccl_ring builds rdma_* device names. Upstream maps Thunderbolt-without-RDMA to ring. oMLX does not.

2

Recovery Advice That Cannot Help

classify_link() at transport.py:506 emits its rdma_not_enabled branch, telling the operator to boot into macOS Recovery and run rdma_ctl enable — on a machine where rdma_ctl status already reports enabled. The branch ignores tb_version and never receives the rdma_ctl state at all.

3

A Terminal State Reported As A Fault

probe.py:527 raises the warning as though something is misconfigured, and the transport state reads peer_linked_config_pending when nothing is pending and nothing ever will be.

4

A Successful Read Thrown Away

_extract_tb_link_speed at transport.py:293 makes its own SSH call from inside the topology try block. When that one call fails, a connectivity matrix that was read successfully is discarded along with it — distinct from the probe-synthesises-Ethernet bug, because here the read worked.

One Honest Open Question

MLX’s documentation says JACCL is necessary for tensor parallelism. oMLX’s own measurement, recorded in the comment directly above _COLLECTIVE_CAPABLE_KINDS, disagrees about how bad the alternative is — and quantifies it:

# omlx/cluster/link_bandwidth.py:55
# on this fabric the *same* Thunderbolt cable gave 28.6 tok/s under RDMA and
# 6.6 tok/s under the TCP ring — 4.3x from the transport alone, at identical
# bandwidth.

So the fix is unambiguous: stop TB4 selecting a backend that cannot initialise. What happens next is not. Whether a TB4 pair should then be offered tensor parallelism over ring at 4.3× the cost, or steered to pipeline parallelism instead, is a product decision with a real trade-off on both sides. It deserves to be argued in the open rather than smuggled into a bug fix.

What The Three Findings Have In Common

None of these is a typo. Each is a feature that looks finished from the outside — green dashboard cards, passing tests, responding endpoints, plausible log messages, actionable-sounding advice. The SLO subsystem is green because it measured nothing. The auth check succeeds without knowing who asked. The Thunderbolt classifier is confident and wrong.

All three failures live in the wiring, not in the units. SLOTracker is correct code with no producer. verify_any_api_key is correct code that discards its most useful intermediate value. classify_link is correct code that was never handed the one input that would let it decide properly. Unit tests cannot see any of that by construction, because a unit test supplies its own inputs.

The counter-example is in the same repository, on tab 3: the activation gate that recomputes a plan signature and refuses to launch something different from what was approved. That one exists because somebody watched the failure happen and wrote the check that would have caught it. That is the difference between code that was tested and code that was operated.