Profiling MLX — Finding the Slow Op

Find the slow op, fix it, re-measure

When MLX inference is slower than you want, the answer is rarely "the framework is slow." It's almost always "this specific operation in your specific workload is the bottleneck, and you didn't know which one until you measured." This lesson is the practical profiling toolkit.

The two-tier approach

Tier 1: MLX-level profiling. mlx-lm and mlx itself expose enough metadata to tell you where time goes within MLX's view of the world — generation_tps, peak_memory, per-eval timing. Use this to identify which call is slow.

Tier 2: Metal-level profiling. When MLX-level data isn't enough (e.g. you've narrowed down to one matmul but want to know which Metal kernel is slow), Xcode's Metal Debugger gives you GPU-level traces. Heavier setup, but the only path when the bottleneck is below MLX's API surface.

Tier 1 — Built-in mlx-lm metadata

The GenerationResponse objects from stream_generate (covered in lm.lesson2) carry generation_tps and peak_memory. Log them in production and you have throughput + memory traces for free without a separate observability stack.

For non-LLM workloads, wrap the suspect MLX call in time.perf_counter() with mx.eval() at the boundary. Don't measure with print() in the hot loop — terminal I/O contaminates the measurement (covered in core.lesson7).

Tier 2 — Xcode Metal Debugger

Open Xcode → Debug → Capture GPU Frame while your MLX workload runs. The capture shows you each Metal kernel dispatch with timing. The output is dense — you'll see kernels named after the MLX operations that produced them — but identifying the longest-running kernel almost always points at the right hot spot.

This is the right tool when you're trying to understand why a particular MLX operation is slow at the kernel level. For most users this is overkill; learning Tier 1 is enough.

The fix-and-re-measure loop

Profiling without action is debugging tourism. Once you've identified the slow op, the fixes generally fall into:

Use mx.compile on the hot function (core.lesson7). Often the cheapest 2-5× speedup available.
Quantize if you haven't (model size affects memory bandwidth, which is often the real bottleneck on Apple Silicon).
Reshape your computation to match MLX's strengths — fewer, larger ops fuse better than many small ones.
Replace a CPU fallback if you discover one (rare in MLX, common in PyTorch MPS).
Upgrade MLX — the framework moves fast and kernel improvements ship every few weeks.

Then re-profile to confirm the fix. Don't skip the re-profile step; intuition about what's faster is unreliable.

Code

Tier 1 — log generation_tps and peak_memory in production·python

from mlx_lm import load, stream_generate
import time

model, tok = load("mlx-community/Llama-3.2-1B-Instruct-4bit")

last = None
t0 = time.perf_counter()
for chunk in stream_generate(model, tok, prompt="Tell me about MLX:", max_tokens=100):
    last = chunk

elapsed = time.perf_counter() - t0
print(f"Generated {last.generation_tokens} tokens in {elapsed:.2f}s")
print(f"  generation_tps : {last.generation_tps:.1f}")
print(f"  peak_memory MB : {last.peak_memory / 1024 / 1024:.1f}")
print(f"  finish_reason  : {last.finish_reason}")

# Sample on M3 Ultra Studio with Llama-3.2-1B Q4 (verified 2026-05-03):
#   generation_tps : ~600-800
#   peak_memory MB : ~685

Tier 1 — time arbitrary MLX calls correctly·python

import mlx.core as mx
import time

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

def benchmark(fn, *args, n=10):
    # Warm up
    mx.eval(fn(*args))
    t0 = time.perf_counter()
    for _ in range(n):
        r = fn(*args)
        mx.eval(r)
    return (time.perf_counter() - t0) / n * 1000   # ms per call

t = benchmark(lambda x: x @ x.T, x)
print(f"matmul 4096x4096 : {t:.2f} ms / call")

# Don't put print() in the timed loop — terminal I/O contaminates timing.
# Always wrap with mx.eval() to materialize before reading the clock.

Exercise

Take an MLX workload you actually run. Add tier-1 profiling — log generation_tps and peak_memory for every generation. Run it under a realistic load and capture the numbers. Identify whether the bottleneck is throughput, memory, or something else (e.g. preprocessing latency outside MLX). Two sentences on what you'd optimize first based on the numbers.