Making Mac Studio inference faster with MLX
A Mac Studio can have enough unified memory to load a model and still deliver disappointing token latency. Capacity answers whether the allocations fit. Speed depends on how often kernels read those allocations, how much useful arithmetic each read supports, and how the runtime schedules the work.
Start with a supported model and an explicit workload. The MLX introduction covers the array framework. Here the target is an existing decoder served through MLX-LM.
Establish the execution path
MLX-LM accepts a local model directory or a Hugging Face repository. Its documented examples include mlx-community/Mistral-7B-Instruct-v0.3-4bit; architecture support comes from the installed MLX-LM implementation, rather than the presence of a Safetensors file alone.1
On an Apple silicon Mac, create an isolated environment and record its packages:
uv venv
uv pip install mlx-lm
uv pip freeze > versions.txt
uv run mlx_lm.generate \
--model mlx-community/Mistral-7B-Instruct-v0.3-4bit \
--prompt 'Explain why decode and prefill stress different hardware.' \
--max-tokens 128 --temp 0
Also record the downloaded model revision, macOS version, chip, memory capacity, prompt token IDs, and power conditions. A repository name can point to different files next month. An instruction model needs its intended chat template; comparing different templates measures different inputs.
MLX on this path executes CPU and Metal GPU operations. A large Neural Engine specification does not establish that this program uses it. Apple inference backends explains how to inspect that distinction.
Compute two rooflines
For a kernel doing operations and transferring bytes, an approximate lower bound is
is sustainable arithmetic throughput for the selected kernel and dtype; is sustainable bandwidth for its access pattern. Advertised peaks are weaker inputs than measured values.
For a dense weight matrix with elements and token rows, the matrix multiplication performs about operations. If weights dominate traffic and occupy bytes per element, its weight-only arithmetic intensity is roughly . Decode at one token has little weight reuse. Prefill has many rows that can reuse each weight tile.
One weight bank, several token rows
Move across the token rails or press play. Batching lets several token rows reuse the same weights. This schematic shows the reuse mechanism; latency and KV traffic still depend on the workload. Open the figure.
A toy decoder reading 4 GiB of weights per step at an assumed effective 400 GiB/s has a weight-transfer lower bound of 10 ms, or 100 steps/s. This excludes attention, sampling, launches, and synchronization. Neither number is a Mac benchmark. Batching four sequences could reuse weights across four token rows, but their KV reads and output latency still grow.
Long contexts can shift decode toward cache traffic. Mixture-of-experts routing changes active weight traffic; small kernels can instead hit dispatch overhead. Use the equation to form an experiment, then inspect the trace.
Count the bytes behind “4-bit”
In MLX’s affine group quantization, values share a scale and bias. Assume 64 weights per group, four bits per weight, and FP16 scale and bias. The group occupies
That is 4.5 effective bits per weight before padding and unquantized tensors. Smaller groups spend more metadata to fit local ranges more closely. MLX’s other quantization modes have different scale formats; this arithmetic describes affine groups only.2
To make a controlled variant from an unquantized supported checkpoint:
uv run mlx_lm.convert \
--hf-path mistralai/Mistral-7B-Instruct-v0.3 \
--mlx-path ./mistral-affine-q4-g64 \
--quantize --q-mode affine --q-bits 4 --q-group-size 64
The destination must be a new path. Keep the source revision and conversion settings with the artifact.3 Quantized storage saves reads; it does not prove every multiplication uses native four-bit arithmetic. Kernels may unpack and combine values with higher-precision activations or accumulators.
For full-attention layers, an unquantized cache estimate is
With 32 layers, eight KV heads, head width 128, FP16 state, and one sequence, each token costs 128 KiB. An 8,192-token context costs 1 GiB. Four independent contexts cost 4 GiB. This is illustrative, excludes allocator overhead, and does not describe recurrent or mixed-cache architectures. The cache memory guide treats sharing and fragmentation in detail.
MLX-LM exposes --kv-bits, --kv-group-size, and --quantized-kv-start. Quantization applies where the cache supports conversion. --max-kv-size can select bounded rotating state on supported paths; it changes retained context and must undergo long-context quality testing. It is not lossless memory compression.4
Time completed work
Constructing an MLX expression can merely construct a lazy graph. An unevaluated matrix multiplication is a poor stopwatch target.5 This isolated example times completed repeated operations with inputs already materialized:
import time
import mlx.core as mx
x = mx.random.normal((8, 1024)).astype(mx.float16)
w = mx.random.normal((1024, 1024)).astype(mx.float16)
mx.eval(x, w)
for _ in range(3):
mx.eval(x @ w)
mx.synchronize()
start = time.perf_counter()
for _ in range(20):
y = x @ w
mx.eval(y)
elapsed = time.perf_counter() - start
print({"milliseconds_per_call": 1000 * elapsed / 20})
This deliberately includes one completion boundary per call. It measures neither model throughput nor the best pipelined schedule. For generation, separate model loading, first-run specialization, prompt processing, and steady decode. Record client TTFT independently of the runtime’s prompt throughput.
mx.async_eval submits evaluation without blocking the caller and remains an experimental API. It can overlap CPU token handling with device work. Immediately calling .item() on every intermediate array can remove that overlap. Enqueuing an unbounded chain creates another problem: growing pending work and delayed cancellation.6 Place evaluation boundaries where results or bounded resource lifetimes require them.
Reuse prefixes and control admission
For a fixed document prefix, MLX-LM provides a persisted prompt-cache path:
uv run mlx_lm.cache_prompt \
--model mlx-community/Mistral-7B-Instruct-v0.3-4bit \
--prompt - --prompt-cache-file document-cache.safetensors < document.txt
uv run mlx_lm.generate \
--prompt-cache-file document-cache.safetensors \
--prompt 'List the assumptions in that document.' --max-tokens 128
The saved prefix precedes the supplied prompt, and the CLI can recover the model identity from cache metadata.1 Verify the resulting token sequence against the uncached request. A document cache does not automatically reproduce the message boundaries of a chat conversation. In a service, cache identity also includes model revision, adapter, positions, and any multimodal inputs.
Use batch generation or the runtime’s scheduler for concurrency. Independent Python processes each loading the weights can consume the memory that should have held cache state. Increase concurrency while measuring per-request token gaps, rather than assuming aggregate throughput predicts responsiveness. The custom Mac stack explains admission and cancellation.
Reserve space for weights, admitted cache growth, prefill temporaries, and the rest of macOS. Unified CPU/GPU memory access removes some transfers but does not create separate RAM for each processor.7 Observe Activity Monitor memory pressure and swap growth alongside process allocations. Raising allocation limits does not add bandwidth or memory.
A repeatable optimization run
Keep a fixed set of short, long, and prefix-sharing requests. Measure uncached and reused-prefix TTFT separately. Sweep prompt chunk size and concurrency independently. Smaller --prefill-step-size values reduce peak prompt-processing memory but can add execution overhead.4
For every candidate, report median and tail token latency, total output tokens divided by wall time, peak allocation, swap growth, and quality results. Keep actual output length because early termination changes the denominator. Compare held-out token loss and task checks against the higher-precision model, including structured output, rare tokens, retrieval over long contexts, and stop behavior. A faster wrong answer is a failed configuration.
No Mac measurements accompany this article. The commands define the experiment; the byte calculations explain what to look for when running it on the target machine.
Footnotes
Prefill and decode expose different opportunities to reuse a weight tile. Measure them separately.
Prompt prefill
- Many token rows can reuse each loaded weight tile.
- Prompt chunk size changes temporary memory and execution overhead.
- Measure prompt processing and time to first token.
Autoregressive decode
- One sequence contributes a narrow token row per step.
- Weights and growing attention state can dominate memory traffic.
- Measure token gaps and aggregate throughput at each concurrency.
A faster prefill kernel need not reduce steady decode latency. Optimize the phase that limits the workload.
Qualitative model. Active experts, context length, batch size, and kernel overhead can change the limiting resource.
Illustrative full-attention cache with 32 layers, 8 KV heads, head width 128, FP16 scalars, and one sequence. Bytes = 2 × 32 × tokens × 8 × 128 × 2.
Show exact values
| Cached tokens | KV payload, GiB |
|---|---|
| 0 | 0 |
| 8,192 | 1 |
| 16,384 | 2 |
| 32,768 | 4 |
Under these assumptions, each additional 8,192 tokens adds 1 GiB per independent sequence.
Calculated payload only. Excludes allocation overhead, shared prefixes, rotating or quantized state, and non-attention layers.