mx.compile

mx.compile traces a function's lazy graph once, fuses runs of elementwise operations into generated Metal kernels, and replays the optimized graph on subsequent calls.

CUDA equivalent: a deliberately scoped torch.compile. No graph breaks, no guards-and-recompile drama, no Inductor-style autotuning, because the ambition is narrower: MLX is already lazy, so the graph exists anyway; compile adds shape-specialized caching and elementwise fusion, and stops there. Matmul and attention are already fused kernels; the win compile chases is the long tail between them.

Why that tail matters here more than on NVIDIA: an unfused elementwise chain (say, SiLU → multiply → add in a transformer MLP) writes each intermediate to unified memory and reads it back. Pure bandwidth burn on the platform's scarcest resource, plus per-dispatch overhead on each tiny kernel. Fusing the chain into one generated kernel deletes both. This is the platform's lesson 1 applied automatically; Luminal↗ is the same idea pursued to its extreme (e-graph search over fusions, one command buffer per forward pass, flash attention rediscovered by the search).

Mechanics in brief: first call traces with placeholder inputs; fusable subgraphs are compiled to MSL through the runtime-compilation path (compiled.cpp in the backend generates the source) and cached against input shapes/dtypes; changed shapes retrace. shapeless=True opts hot functions out of shape-specialization where their kernels permit. Constants get baked; the usual tracing caveats (Python side effects run once, at trace time) apply as in JAX.

When to reach for it: inference step functions and training steps dominated by many small ops. Typical gains are real but modest (tens of percent, not multiples) since the heavy matmuls were already fused. When not to: code you're about to profile kernel-by-kernel (fusion renames and merges dispatches out from under you, and profiling is hard enough here), or workloads that are one giant matmul anyway.

For the kernel engineer, mx.compile also defines the boundary of custom work: anything expressible as elementwise chains, the compiler will fuse adequately. Custom metal_kernel effort belongs on the patterns compile can't see: reductions with structure, softmax-shaped streaming, quantization-aware loops.