Apple Silicon and MLX

Apple Silicon and MLX

MLPL's MLX backend runs on Apple Silicon and exposes the same device-scoped language shape used by CUDA. Metal is enabled, and a TensorHandle seam carries lazy MLX arrays through the autograd tape so forward values, gradients, weights, and optimizer moments stay on the GPU across train steps.

Basic source shape

device("mlx") {
  X = randn(7, [1024, 64])
  mdl = chain(linear(64, 256, 0), relu_layer(), linear(256, 64, 1))
  y = apply(mdl, X)
}

Run an MLX demo from a build that enables the mlx feature. The CLI workspace owns the user-facing binaries, while components/native-rt contains the MLX runtime crates.

Two MLX deployment shapes

flowchart TD
    Client[REPL or script] --> Choice{MLX placement}
    Choice --> InProc[In-process MLX feature]
    Choice --> Server[mlpl-serve orchestrator]
    Server --> Peer[mlpl-mlx-serve peer]
    InProc --> Metal[Apple GPU through MLX]
    Peer --> Handle[Opaque device tensor handles]
    Handle --> Metal
    Metal --> Materialize[Explicit CPU materialization when needed]

Loading

• In-process mode links the Apple-native runtime into the evaluator.
• Peer mode routes a whole device("mlx") block to mlpl-mlx-serve; opaque handles keep tensors at the peer until materialized.

Capabilities

• MLX-backed arithmetic, matrix operations, activations, reductions, normalization, and relevant model-forward primitives.
• Device scopes and to_device movement with cross-device error handling.
• CPU/MLX parity tests within fp32 tolerance.
• General on-device autograd and optimizer execution through the resident tensor tape.
• MLX-specific demos including LoRA, tic-tac-toe, neural-thicket, tiny-LM-oriented work, and remote/service flows, with availability depending on the current registry and build.
• A separate peer service integrated with the session server.

Execution model

device("mlx") { train N { adam(...) } } runs the loop with weights, tape intermediates, gradients, and optimizer moments resident on Metal. One boundary and one performance reality remain:

• A few structural gradient kernels (notably fused cross-entropy backward) still run on the exact-f64 CPU path; the mixed-residency accumulator re-uploads their results, and seam counters (uploads / downloads / submits / cpu_fallbacks) expose exactly what crossed the boundary.
• Per-operation dispatch is the floor: a tiny-LM training step submits roughly 200 lazy MLX ops, so at very small model sizes the CPU interpreter still wins. The benchmark suite shows CPU winning at d=32, MLX crossing over around d=128, and tiny_lm_train_step_d256 running about 3x faster on MLX.

See the benchmarks document for timing tables and seam profiles.

When MLX helps

• Models around d=128 and up, where kernel work amortizes the per-op dispatch floor (see the crossover table in the benchmarks doc).
• General adam/train loops -- no longer only the recognized LoRA/MLP fast-path shapes.
• Engram-in-chain training runs fully resident (selection-matmul gather with exact scatter-ADD backward; one CPU fallback per step -- fused cross-entropy backward); the "MLX Tiny LM + Engram" demo is the worked example, and the engram chain crosses over to an MLX win near d=128 like the base tiny LM.
• Remote Apple compute accessed through the browser or another machine.
• Research on unified-memory and persistent device-tensor execution.

Tiny models remain faster on CPU (per-op dispatch dominates below roughly d=128). Measure release builds and inspect the seam counters.

Source references: MLX guide, MLX service guide, benchmarks, and components/native-rt.