mx.compile — Kernel Fusion You Can Actually Measure

What mx.compile actually does

mx.compile wraps a function and returns a new function that, on first call, traces the operations into a graph and compiles them into a single (or smaller-set-of) Metal kernel. On every subsequent call with the same input shapes and dtypes, it skips the tracing and dispatches the compiled kernel directly. The result, when the function is hot, is fewer kernel launches and tighter fusion of intermediates that never have to be written to memory.

You don't need to understand Metal to use it. You do need to understand when to use it.

The benchmark, on this machine

The code block below times a synthetic hot function — twenty serial tanh ops on a (512, 512) array — both with and without compilation. On my office Mac (M3 Ultra Studio, mlx 0.31.2, 2026-05-03):

Plain: ~0.97 ms / call
Compiled: ~0.31 ms / call
Speedup: ~3.18x

Your numbers will vary by machine, by what else is competing for the GPU, by the specific shape of the function. The shape of the result — meaningful but not order-of-magnitude — is what you should expect.

When to use it (and when not)

Use it on hot, repeatedly-called functions where input shapes are stable. Training-step functions, inference-token-step functions, anything called millions of times with the same shape signature.

Don't use it on functions you only call a handful of times — the tracing-and-compilation cost is paid on first call and amortizes across many calls. For one-off computations, the overhead exceeds the saving.

Be careful with functions whose input shapes vary — every new shape signature triggers a recompile, which can dominate runtime if your shapes are constantly shifting. A function that's compiled once and dispatched a thousand times is great; a function that's recompiled a thousand times is the opposite of great.

What it doesn't do

mx.compile doesn't change the result of your function — output is bit-identical (or floating-point-equivalent) to the uncompiled version. It also doesn't magically speed up everything by 3x — for memory-bound ops or ops that are already a single kernel call, the speedup is small or zero. The 3x in the benchmark above comes from fusing twenty tanh ops into one kernel; on a function that's already one matmul, you'll see almost nothing.

Code

Plain vs compiled — measure for yourself·python

import mlx.core as mx
import time


def heavy(x):
    y = x
    for _ in range(20):
        y = mx.tanh(y * 1.001 + 0.001)
    return y.sum()


heavy_compiled = mx.compile(heavy)

x = mx.random.normal((512, 512))

# Warm up both versions to amortize first-call costs (compile traces here).
mx.eval(heavy(x))
mx.eval(heavy_compiled(x))


def bench(fn, n=20):
    mx.eval(fn(x))   # warm
    t0 = time.perf_counter()
    for _ in range(n):
        r = fn(x)
        mx.eval(r)
    return (time.perf_counter() - t0) / n * 1000   # ms


t_plain = bench(heavy)
t_comp  = bench(heavy_compiled)

print(f'plain    : {t_plain:.3f} ms / call')
print(f'compiled : {t_comp:.3f} ms / call')
print(f'speedup  : {t_plain/t_comp:.2f}x')

# Verified on M3 Ultra Studio, mlx 0.31.2 (2026-05-03):
#   plain    : 0.969 ms / call
#   compiled : 0.305 ms / call
#   speedup  : 3.18x

Exercise

Run the benchmark on your machine. Then change one variable at a time and re-measure: (a) change the array shape from (512, 512) to (4096, 4096) — does the speedup hold, grow, shrink? (b) change the loop count inside heavy from 20 to 1 (i.e., a single op) — what happens to the speedup? The intuition you should walk away with: compile helps most when there are many small ops to fuse and shapes stay stable.