~16 min · journey, pytorch, mps, dispatch-overhead, op-coverage, measured
Level 0Spec-Sheet Skimmer
0 XP0/91 lessons0/19 achievements
0/100 XP to next level100 XP to go0% complete
"The MPS path pulls the bandwidth when the op is big and pays the dispatcher when the op is small. A token is a thousand small ops."
A Device Word for a Machine Without Devices
PyTorch reaches the Apple GPU through its mps backend, which its own note describes as mapping "computational graphs and primitives on highly efficient Metal Performance Shaders Graph framework and tuned kernels" — and which you enter by doing what the note says: "simply move your Tensor and Module to the mps device." That is the CUDA vocabulary — a device, a move — laid over a pool that has no second place to move to. The GPU track measured the verb: .to('mps') on a gigabyte costs 66 ms on office, a copy inside one pool, because the framework's model of the machine has two memories and the hardware has one. Everything else about the path follows from that mismatch, in both directions: where the framework's model is merely redundant the path is fast, and where it is load-bearing the path breaks.
Where It Is Fast: Big Ops
Measured on office, torch 2.12: an 8192² matmul runs at 18.1 TFLOP/s in fp32 and 21.9 in bf16 — the same silicon MLX drove to 19.5 in fp32 at n = 4096 in the CPU track, so the kernels are not the problem. A bf16 matrix-vector over a 2.15 GB matrix takes 3.55 ms: 605 GB/s, 95% of the streaming bandwidth this quest measured on the machine. When one op is large, the MPS path reads the pool at the pool's speed. Demucs, the music-learning engine's source separator, is a case of large ops: on torch 2.12 it separates 30 seconds of audio in 1.18 s on the GPU against 12.6 s on 24 CPU threads, outputs agreeing to 1.7% relative L2 — ten times faster on the same Mac. The image engine's whole path is this case: diffusers on MPS on server, one process, checkpoints and upscalers and inpainting all through the same door, and it is the household's one image-generation host.
Where It Breaks: Small Ops, Missing Ops, Stale Beliefs
Same Mac, the same bf16 Llama-3.2-1B checkpoint — 2.47 GB per token, tied embeddings, a 258 tokens-per-second ceiling at 638 GB/s — decoded two ways. mlx-lm: 190 tokens per second, 5.3 ms a token, 74% of the ceiling. PyTorch on MPS through transformers, greedy, eager: 32–42 tokens per second, 24–31 ms a token, 12–16%. The bytes cost 4 ms at the bandwidth the matvec just demonstrated; the other twenty-plus milliseconds are the fixed term of the decode-ceiling lesson at its worst — sixteen layers of a dozen ops each, every op dispatched separately through the framework, synchronized to a Python loop that the pool cannot speed up. Same weights, same bus, a fifth of the speed, and nothing about the chip in the difference.
Then the ops that are not there. torch 2.12 refuses float64 on MPS outright ("the MPS framework doesn't support float64"), and linalg.eigh raises the error every MPS user has read: "not currently implemented for the MPS device … you can set the environment variable PYTORCH_ENABLE_MPS_FALLBACK=1 to use the CPU as a fallback." Measured with the fallback on, a 2048² eigh took 350 ms against 371 on the CPU directly — the fallback cost nothing, because the "copy to the CPU" was inside one pool. On a discrete GPU the same fallback pays two bus crossings. That is unified memory covering for the framework's model of the machine, quietly.
And the stale belief. The music engine's code pins Demucs to the CPU with the comment that MPS is broken for it, while its own doc says the engine runs on MPS. On the torch the engine actually runs, MPS is not broken for it; it is ten times faster. A path that broke once stays broken in the code long after the runtime fixed it — this quest reports the finding to the family rather than repairing it, since the engine is not the quest's to edit. Your Mac card gets the general rule: an MPS "doesn't work" belongs on the card with a torch version beside it, and expires.
same_bytes_two_paths.py — one bf16 checkpoint, decoded through torch MPS and through mlx-lm·python
#!/usr/bin/env python3
"""The same bf16 checkpoint decoded two ways on one Mac: PyTorch on MPS through
transformers, and mlx-lm. Same bytes per token, same pool, two paths."""
import sys, time
MODEL = "unsloth/Llama-3.2-1B-Instruct" # bf16 safetensors, tied embeddings
PROMPT = "Explain, in about 300 words, why unified memory matters for language models."
N = 128
def torch_mps():
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to("mps").eval()
msgs = [{"role": "user", "content": PROMPT}]
ids = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt", return_dict=True)["input_ids"].to("mps")
nbytes = sum(p.numel() * p.element_size() for p in model.parameters())
with torch.no_grad():
model.generate(ids, max_new_tokens=8, do_sample=False) # warm-up
torch.mps.synchronize(); t0 = time.perf_counter()
out = model.generate(ids, max_new_tokens=1, do_sample=False)
torch.mps.synchronize(); ttft = time.perf_counter() - t0
t0 = time.perf_counter()
out = model.generate(ids, max_new_tokens=N, do_sample=False, min_new_tokens=N)
torch.mps.synchronize(); dt = time.perf_counter() - t0
gen = out.shape[1] - ids.shape[1]
print(f"torch {torch.__version__} MPS via transformers: {nbytes/1e9:.2f} GB of bf16 weights; prompt {ids.shape[1]} tokens; "
f"TTFT {ttft:.3f}s; {gen} tokens in {dt:.2f}s = {gen/(dt-ttft):.1f} tok/s (decode, TTFT removed)")
def mlx_path():
import mlx.core as mx
from mlx_lm import load, stream_generate
from mlx_lm.sample_utils import make_sampler
model, tok = load(MODEL) # HF bf16 loaded as-is
msgs = [{"role": "user", "content": PROMPT}]
prompt = tok.apply_chat_template(msgs, add_generation_prompt=True)
sampler = make_sampler(temp=0.0)
for _ in stream_generate(model, tok, prompt, max_tokens=8, sampler=sampler): pass # warm-up
last = None
for r in stream_generate(model, tok, prompt, max_tokens=N, sampler=sampler): last = r
print(f"mlx {mx.__version__} mlx-lm: prompt {len(prompt)} tokens; prompt {last.prompt_tps:.0f} tok/s; "
f"decode {last.generation_tps:.1f} tok/s; peak {last.peak_memory:.2f} GB")
if __name__ == "__main__":
{"torch": torch_mps, "mlx": mlx_path}[sys.argv[1]]()
# office, M3 Ultra, 2026-09-15:
# torch 2.11.0 MPS via transformers 5.14.1: 2.47 GB bf16; prompt 52; TTFT 0.024s; 42.1 tok/s, then 31.8 on a second run
# mlx 0.32.2 / mlx-lm 0.31.3: prompt 2350 tok/s; decode 190.0 tok/s; peak 2.56 GB
# ceiling for 2.47 GB per token at 638 GB/s: 258 tok/s
mps_breaks.py — big ops, a missing op, and what the fallback costs inside one pool·python
#!/usr/bin/env python3
"""Run with PYTORCH_ENABLE_MPS_FALLBACK=1 so the missing op falls back instead of raising."""
import time, torch
def timeit(fn, n=3):
fn(); torch.mps.synchronize()
t = time.perf_counter()
for _ in range(n):
fn()
torch.mps.synchronize()
return (time.perf_counter() - t) / n
x = torch.randn(8192, 8192, device="mps")
for name, a in (("fp32", x), ("fp16", x.half()), ("bf16", x.bfloat16())):
dt = timeit(lambda: a @ a, 5)
print(f"matmul 8192^2 {name} on mps: {2*8192**3/dt/1e12:.2f} TFLOP/s")
W = torch.randn(32768, 32768, device="mps", dtype=torch.bfloat16); v = torch.randn(32768, device="mps", dtype=torch.bfloat16)
dt = timeit(lambda: W @ v, 10)
print(f"matvec over {W.numel()*2/1e9:.2f} GB bf16 on mps: {dt*1e3:.2f} ms = {W.numel()*2/dt/1e9:.0f} GB/s")
s = torch.randn(2048, 2048); s = s @ s.T
s_mps = s.to("mps")
print(f"eigh 2048^2 on cpu: {timeit(lambda: torch.linalg.eigh(s))*1e3:6.1f} ms")
print(f"eigh 2048^2 on mps (CPU fallback, env): {timeit(lambda: torch.linalg.eigh(s_mps))*1e3:6.1f} ms -> result on {torch.linalg.eigh(s_mps)[0].device}")
try:
torch.randn(64, 64, device="mps", dtype=torch.float64)
except TypeError as e:
print("float64:", e)
# office, torch 2.12.0, 2026-09-15:
# matmul 8192^2 fp32 18.05 / fp16 21.09 / bf16 21.88 TFLOP/s
# matvec over 2.15 GB bf16: 3.55 ms = 605 GB/s
# eigh 2048^2 on cpu 371.3 ms; on mps with fallback 350.2 ms -> result on mps:0
# float64: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.
# without the env var, eigh raises: "The operator 'aten::_linalg_eigh.eigenvalues' is not currently implemented for the MPS device ..."
Run same_bytes_two_paths.py both ways on your Mac with a bf16 checkpoint that fits. Compute milliseconds per token for each, subtract the bytes term (bytes ÷ your measured bandwidth), and write the two fixed terms on your card. Then run mps_breaks.py and record which op your torch refuses; note the torch version next to it, because that line will expire.
Hint
If the torch fixed term is ten times the MLX one, that is normal for eager decode through a Python framework, not a broken install. If the matvec bandwidth is far below your stream figure, the matrix was not big enough to hide dispatch — use a matrix of at least a gigabyte.
Progress
Progress is local-only — sign in to sync across devices.