The MLX Path

"Prefill once. Then the same weights, one token at a time, with a cache that grows. Write the loop yourself and there is nothing left to believe."

The Loop, Written Out

The mlx quest teaches MLX as a framework — arrays, lazy evaluation, streams, quantization, the model zoo. This lesson uses it for one purpose: to show a token's path through a model with no library loop hiding it. The code block loads a model (mapped, lazy — the previous lesson), builds one KVCache object per layer with make_prompt_cache, runs the whole prompt through model(prompt, cache=cache) once — that is prefill — and then loops: one token in, one forward pass through the same weights, one token out, cache one position longer. That is decode, and every physics-track claim is visible in the code. The weights are the same object on every iteration. The cache is the thing that changes. The mx.eval on the sampled token is where lazy evaluation is forced and the GPU actually runs.

On office the hand loop decoded the 9B at 68–79 tokens per second across runs, against 95 for mlx-lm's stream_generate on the same model. The gap is the fixed term of the decode ceiling: the hand loop calls .item() to read each token back to Python, which synchronizes the GPU every step, and it drives every step from the interpreter. The library keeps the pipeline full. Same weights, same bus, same cache — a fifth to a quarter of the speed left in the loop's shape (79 and 68.5 against 95, two runs).

What Is Inside a Forward Pass

Every model(...) call is a graph MLX builds lazily and runs on the Metal stream: for each of the model's layers, a normalization, an attention block (full or linear, per the hybrid layout), and a feed-forward block, each of which is a small number of quantized matrix-vector kernels reading 4-bit weights and their scales — the bytes-per-token of the physics track, read once. For full-attention layers the kernel also reads the cache for every previous position, which is the context slope. The output head produces logits over the vocabulary, argmax picks one (the lab uses greedy sampling so runs repeat exactly), and that integer is the next input. Nothing crosses a boundary between steps; the only host-side event is reading the token id back.

Why the Household Runs This Path

Three reasons, each a track. It reads each weight once at full width — the decode ceiling lesson — and reached 74% of the ceiling on the 27B, the best fraction of any runtime the quest measured. It is unified-memory-native — the GPU track — so there is no move verb and no framework copy. And it is the door Apple maintains, so the M5 Neural Accelerators reached it first (MLX 0.30.0, 2025-11-19, by the release notes — a vendor date, not something this quest measured). The household's inference hub is an MLX server and this quest's entire lab runs on mlx-lm; the coding agent's local leg goes through the next lesson's door, which is the one most people use instead.

Code

decode_loop.py — prefill once, then one token per pass through the same weights·python

#!/usr/bin/env python3
"""The MLX path, written out: load (lazy, mmap), prefill the prompt through
the model once, then one token per step through the same weights with a
growing KV cache. No stream_generate; the loop is the lesson."""
import time
import mlx.core as mx
from mlx_lm import load
from mlx_lm.models.cache import make_prompt_cache

model, tok = load("mlx-community/Qwen3.5-9B-4bit")
prompt = mx.array(tok.apply_chat_template([{"role": "user", "content": "Name three uses of unified memory."}],
                                          add_generation_prompt=True, enable_thinking=False))[None]
cache = make_prompt_cache(model)                       # one cache object per layer

t0 = time.perf_counter()
logits = model(prompt, cache=cache)                     # PREFILL: the whole prompt, one pass
token = mx.argmax(logits[:, -1, :], axis=-1)
mx.eval(token)                                         # lazy until here; the GPU runs now
ttft = time.perf_counter() - t0

out = [int(token.item())]
t1 = time.perf_counter()
for _ in range(63):                                     # DECODE: one token per pass
    logits = model(token[None], cache=cache)            # same weights; cache grows by one
    token = mx.argmax(logits[:, -1, :], axis=-1)
    mx.eval(token)
    out.append(int(token.item()))                       # a GPU sync per token — the hand loop's cost
    if out[-1] == tok.eos_token_id:
        break
dt = time.perf_counter() - t1

print(tok.decode(out).strip()[:300])
print(f"\nprefill {prompt.shape[1]} tokens: TTFT {ttft:.3f}s; decode {len(out)-1} tokens at {(len(out)-1)/dt:.1f} tok/s")
# office, 2026-09-15, mlx 0.32.2 / mlx-lm 0.31.3: prefill 19 tokens, TTFT 0.30s; decode 68-79 tok/s across runs
# (stream_generate on the same model: 95 tok/s — the library keeps the pipeline full)

Exercise

Run decode_loop.py, then change one thing: remove the .item() call inside the loop (keep the token on the GPU, append it to an MLX array, decode at the end). Measure decode tok/s before and after and add both to your card. Then explain, in terms of the decode-ceiling formula, which term you changed.

Hint

You changed the fixed term, not the bytes. The sync per token is a round trip the GPU waits on; removing it lets MLX pipeline the next forward pass behind the current one. If the rate barely moves, the interpreter's per-step cost dominates instead — also the fixed term.