MLX is Apple’s array framework for machine learning on Apple silicon, and it occupies an interesting spot in the ecosystem. Next to PyTorch it looks deliberately small; next to JAX it shares the functional taste — function transformations like grad, vmap, and compile are first-class — but it ships with a C++ core that speaks Metal natively rather than sitting on an existing accelerator stack. If you have read our tours of llama.cpp or GGML, you have seen what it takes to hand-roll GPU execution for inference. MLX attacks the same hardware from the opposite direction: a clean NumPy-shaped user surface on top of a graph evaluator, so research code and inference code end up looking alike.

The design bet is unified memory. On an M-series machine there is one pool of RAM visible to both CPU and GPU, so MLX can move work between the two without staging copies through host memory. That single hardware fact explains most of the architecture: a tiny Device abstraction, ops that build a graph lazily, and a scheduler that interleaves CPU and Metal work on per-device streams. Compared with heavyweight serving stacks like vLLM or the minimalism of tinygrad, MLX sits in a sweet spot: a real C++ engine, but small enough to read in an afternoon.

This tour walks that engine as it exists at version 0.32.5. The repository splits cleanly into mlx/ — the C++ core with its backends — and python/ — the nanobind bindings plus the pure-Python mlx.nn and mlx.optimizers packages. We will follow an operation from Python call to Metal kernel and back, then close with what the code genuinely cannot do today.

MLX overview: Python ops append primitives to a graph of arrays; the lazy evaluator walks that graph and hands each node to the CPU or Metal backend, coordinated by the stream scheduler over a shared buffer allocator.

Reading the overview from left to right:

• Every tensor is an array — the class comment says it plainly: an array is really a node in a graph, carrying shape, strides, dtype, and edges to its inputs.
• A Device is just cpu or gpu plus an index, and a Stream is an ordered lane on that device; every primitive is bound to one.
• Calls like add() in mlx/ops.cpp do no math — they broadcast inputs, promote dtypes, and return a new node holding an Add primitive.
• Each Primitive must implement both eval_cpu and eval_gpu; that dual contract is the whole “unified backend” story.
• Nothing runs until eval_impl in mlx/transforms.cpp walks the graph depth-first and dispatches each node to gpu::eval or cpu::eval.
• The Scheduler owns one worker thread per stream and throttles the tape when too many tasks or too many bytes are in flight.
• On the GPU side, metal::Device wraps MTL::Device with caches of Metal libraries and compiled compute-kernel states, while a MetalAllocator recycles buffers under a memory limit.
• Python sees all of this through a nanobind module built in python/src/mlx.cpp, with mlx.nn layers as thin wrappers that trace mx calls.

Core Objects: array, Device, Stream

The C++ core is refreshingly flat. mlx/array.h defines the central type in about five hundred lines: an array is a shared descriptor with a shape of int32_t, strides, a dtype, a status flag, and — crucially — a pointer to the primitive that computes it plus the vector of input arrays. Those inputs are the graph edges. Allocation is delegated to a process-wide allocator() interface, which means the core never knows whether a buffer comes from malloc or from a Metal heap.

Device and Stream complete the vocabulary. A Device is two enums wide, and a Stream is an index paired with a device. Streams matter because they are the unit of concurrency: each thread can have its own default stream, mx.new_stream(device) mints another lane, and work submitted to different streams can overlap on the GPU. Cross-stream ordering is handled by an Event type — a wait/signal pair that the evaluator attaches to arrays flowing between streams. This is the same mental model as CUDA streams, but with no vendor library behind it; the CPU backend simply treats a stream as a worker thread.

Ops, Primitives, and the Two-Backend Contract

The user-facing surface lives in mlx/ops.cpp, a four-thousand-line catalogue of NumPy-style operations from arange to einsum. Reading any one of them shows the pattern. add() at line 3266 promotes the two inputs to a common dtype, broadcasts them, and constructs the output array with a shared pointer to an Add primitive — no arithmetic anywhere. Even matmul is just validation plus node construction. Laziness is total: a thousand lines of Python-like array code can execute in microseconds because it only sketches the graph.

The interesting file is mlx/primitives.h. The base Primitive declares pure-virtual eval_cpu and eval_gpu over vectors of inputs and outputs, and documents that the evaluation function is responsible for allocating its own output storage — donation and buffer reuse fall out of that rule. Most ops derive from UnaryPrimitive, an adapter for single-output nodes that forwards the vector form to a simpler (inputs, out) signature. Beyond evaluation, primitives optionally carry vjp, jvp, and vmap rules, which is how reverse-mode autodiff and vectorized mapping stay pure library code rather than compiler magic.

A separate layer, mlx/fast.cpp, adds fused operations — rms_norm, layer_norm, rope — each with a hand-written custom backward. These are the ops transformer models from nanoGPT-style codebases call in their hot loop, and MLX implements them once, with both CPU and Metal paths, instead of hoping a generic kernel fusion pass will invent them.

The Lazy Eval Loop and Its Scheduler

Everything converges on eval_impl in mlx/transforms.cpp. Given output arrays, it first runs a depth-first pass that records each node’s fan-in degree, then flattens the traversal into a tape in breadth-first order — with a width limit so a pathological graph does not blow up the queue. Before executing, it inserts fences wherever consecutive tape nodes switch streams, and events for ordering across separate asynchronous evaluations.

The execution loop is the heart of the system, at line 272: for each node on the tape, if the primitive’s device is GPU it calls gpu::eval, otherwise cpu::eval. That single if is where “unified Metal/CPU backend” is decided, one node at a time. The GPU branch does not block; it encodes a command buffer and returns immediately, so the tape keeps unrolling while the GPU chews. Two conditions force a pause: too many active tasks, or active memory crossing the configured limit — then the loop calls scheduler::wait_for_one() until pressure drops. The same file also hosts the transformations users see: grad() is built on value_and_grad by tracing with tracer arrays, and compile() runs graph simplify and fuse passes and caches the result keyed by shapes and mode.

The CPU branch of the contract is almost anticlimactic. cpu::eval invokes the primitive’s eval_cpu synchronously — SIMD kernels and hand-tuned GEMMs under mlx/backend/cpu/ — then dispatches a no-op task to the stream’s encoder purely to release donated buffers on the worker thread. The asymmetry is the point: the CPU is the fast path with nothing to schedule, and the GPU is the async path that needs fencing.

Inside the Metal Backend

The Metal backend under mlx/backend/metal/ is the largest single directory in the core, and its entry point is the GPU hook gpu::eval implemented in mlx/backend/metal/eval.cpp. For each array it resolves the stream’s CommandEncoder — an object that batches many kernel dispatches into one MTL::CommandBuffer, inserts barriers only when dependencies require them, and commits when a size threshold is hit. The actual call is one line: arr.primitive().eval_gpu(...) at line 45, followed by an async commit whose completion callback tells the scheduler the task finished.

Per-op GPU code lives in mlx/backend/metal/primitives.cpp and friends; for example Arange::eval_gpu figures out grid dimensions and records the matching kernel. The kernels themselves are Metal Shading Language sources — the Steel GEMM family in mlx/backend/metal/kernels/steel/ is the flagship, with tiled fused matmuls specialized per architecture generation, instantiated from templates in steel_gemm_fused.metal and dispatched by steel_matmul_regular_axpby_nax. If you enjoyed the tiling tour through CUTLASS or the fused-attention work in FlashAttention, this is Apple’s in-house answer to the same problem.

Memory management closes the loop. MetalAllocator implements the core’s Allocator interface with a buffer cache: freed GPU buffers are parked and reused by size class, with configurable cache and memory limits and peak tracking exposed through mx.metal.get_cache_memory() and friends. The evaluator consults those statistics between tape nodes — that is how a model larger than the free budget degrades gracefully instead of aborting.

From Install to Your First Model

There is no build step on Apple silicon; the PyPI wheel carries the compiled core:

pip install mlx

The two-line demo that shows the whole machine:

import mlx.core as mx

a = mx.array([1.0, 2.0, 3.0])
b = mx.array([4.0, 5.0, 6.0])
c = a + b + 1.0      # graph node - nothing has executed yet
mx.eval(c)           # DFS tape -> CPU or Metal, per node
print(c)             # array([6, 8, 10], dtype=float32)

Between those lines, mx.eval — bound at line 1181 of the bindings — reaches eval_impl and does all the scheduling described above. A neural network is the same story with syntax sugar: mlx.nn layers are dict subclasses that collect parameters recursively, and a forward pass just traces mx calls. Drop a mx.compile around a training step or a nn.value_and_grad around a loss and you are using the transformation machinery without leaving Python.

import mlx.core as mx
import mlx.nn as nn

model = nn.Sequential(nn.Linear(4, 8), nn.relu, nn.Linear(8, 1))
x = mx.random.normal((2, 4))
y = model(x)         # still lazy
mx.eval(y)

Honest Limits

Read the source and the boundaries are clearly marked. First, platform: the Metal backend requires Apple silicon; on Linux the same code compiles against CPU only, and the newly landed CUDA backend under mlx/backend/cuda/ is real but younger and less battle-tested than the Metal path. There is no Windows build. Second, distribution: mlx.distributed exists with MPI and a ring backend, but there is nothing like the tensor-parallel machinery of full serving stacks — MLX targets one machine with a big unified-memory GPU, which is why it pairs so well with single-node model ports rather than vLLM-style cluster serving. Third, operator coverage: the op catalogue is broad but not NumPy-complete, and a few long-tail ops fall back to slower composed paths. Fourth, the Python object model leaks a little — nn.Module is a dict subclass, which makes attribute plumbing occasionally surprising — and compiled functions must be pure, since the compile cache assumes the graph only depends on its inputs. None of these are hidden; they are all visible in the tree we just walked, and that legibility is arguably MLX’s best feature.

Enjoyed this post? Never miss out on future posts by following us