Unified Memory: Why Apple Built a New Array Library

The border that isn't there

PyTorch was born in a world where the GPU is a separate continent. NVIDIA cards have their own VRAM, physically distinct from system RAM, connected over PCIe. Every operation that touches the GPU has to declare its citizenship — tensor.to('cuda') — and pay the customs duty of a memory copy.

Apple Silicon doesn't have that border. The M-series chips integrate the CPU, the GPU, and the Neural Engine into one piece of silicon, and they all reach into the same physical pool of RAM. There is no separate VRAM. There is no PCIe bus between processors. There is just memory, and the question of which compute unit reads from it next.

MLX is the array library written for that physical reality. The "unified memory" feature you see touted in MLX's marketing isn't a software trick — it's the framework refusing to pretend the hardware is something it isn't.

What that looks like in code

The contrast is starkest when you put two snippets side by side. PyTorch on CUDA is a series of customs declarations:

(This PyTorch snippet is for contrast — don't run it in the mlx env. Just read the ritual.)

And here is the equivalent in MLX. Notice what isn't there.

That's the entire change. No .to(device). No .cpu() on the way back out. No async copy stream to manage. The array's storage is the unified memory the GPU and CPU share, and the framework just runs the right kernels on the right compute unit.

The proof: cross-device arithmetic without a copy

Here is the experiment that made me fully internalize this. We make two arrays, one declared on CPU and one declared on GPU, and we add them together. In PyTorch on CUDA, this is an error (you have to .to() one of them first). In MLX, it just works:

The two "devices" in MLX aren't separate memory pools. They're labels on which compute unit will execute the next operation. The data is in one place the whole time.

Why this is the design center, not a feature

Most of MLX's other choices are downstream of this one fact. Lazy evaluation exists because the framework gets to decide which compute unit runs each op based on the whole graph, not on individual call sites — that decision needs the graph in hand before it can be smart. Function transforms like mx.grad and mx.vmap are easier because there's no transfer cost to schedule around. The training loop doesn't need to interleave compute and copy; it just runs.

If you keep one mental image from this entire quest, make it this one: on Apple Silicon, the array lives in a single pool of bytes that every compute unit can read. MLX is the framework that takes that fact seriously.

Code

PyTorch on CUDA — the customs ritual (read for contrast, do not run in mlx env)·python

# Don't run this in the mlx env — it's a contrast, not a recipe.
import torch
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

x = torch.randn(1024, 1024).to(device)   # ← copy bytes from RAM to VRAM
y = torch.randn(1024, 1024).to(device)   # ← copy bytes from RAM to VRAM
z = (x @ y).cpu()                        # ← copy result back from VRAM to RAM

# Forget any one of these .to() / .cpu() calls and you get a runtime error
# or a silent slowdown.

MLX equivalent — notice what isn't there·python

import mlx.core as mx

x = mx.random.normal((1024, 1024))
y = mx.random.normal((1024, 1024))
z = x @ y

mx.eval(z)   # materialize the lazy graph (more on this in core.lesson4)
print("shape :", z.shape)
print("dtype :", z.dtype)
print("device:", mx.default_device())   # → Device(gpu, 0) on Apple Silicon

# No .to('cuda'). No .cpu(). The unified-memory pool is the storage,
# always, and 'device' is just a label about which compute unit ran it.

Cross-device add: the proof that 'device' is a label, not a pool·python

import mlx.core as mx

mx.set_default_device(mx.cpu)
a_cpu = mx.array([1.0, 2.0, 3.0, 4.0])
print("default device:", mx.default_device())   # Device(cpu, 0)
print("a_cpu          :", a_cpu)                 # array([1, 2, 3, 4], dtype=float32)

mx.set_default_device(mx.gpu)
a_gpu = mx.array([1.0, 2.0, 3.0, 4.0])
print("default device:", mx.default_device())   # Device(gpu, 0)
print("a_gpu          :", a_gpu)                 # array([1, 2, 3, 4], dtype=float32)

# In PyTorch on CUDA, this next line is a runtime error (cross-device op).
# In MLX, it just works. There is one pool of bytes; the labels were a hint
# about scheduling, not a partition of memory.
b = a_cpu + a_gpu
mx.eval(b)
print("a_cpu + a_gpu  :", b)                    # array([2, 4, 6, 8], dtype=float32)

Exercise

Run the cross-device add code block in your mlx env. Confirm array([2, 4, 6, 8]) comes back. Then change one line: declare a_cpu first as before, but immediately do b = a_cpu + a_cpu (same device on both sides) and inspect the device of b. Now do mx.set_default_device(mx.gpu) and add a_cpu + a_cpu again. What changed about the result, and what didn't? Two sentences. The point of the exercise is to feel where the labels matter and where they don't.