Faster Gemma 4 on MLX with multi-token prediction

gist

Ollama 0.31 adds multi-token prediction for Gemma 4 on Apple silicon, which Ollama says makes generation nearly 90% faster on coding work. It is on by default, does not change the model’s output, and rests on three engine changes: a draft length tuned at runtime, a single GPU pass per round, and a new MLX kernel for verification.

ideas

• The speedup is benchmarked on a coding agent workload. Ollama reports nearly 90% faster token generation on average across the Aider polyglot benchmark, measured with Gemma 4 in nvfp4 on an M5 Max.
• Draft length is chosen while the model runs. Ollama tracks how often proposals are accepted and how long verification takes, picks the length with the most tokens a second, and falls back to plain decoding when proposals stop being accepted.
• Each round stays on the GPU. Drafting, sampling, verification, and the following sample run as one pass without returning to the CPU, and a rollback point before each proposal makes rejected tokens cheap to undo.
• Verification needed its own kernel. Batches of 2 to 8 draft tokens fall between decode and prefill kernels, so Ollama contributed an MLX kernel that unpacks each block of weights once per batch, making Gemma 4’s largest matrix multiplications 2× to 2.5× faster on an M5 Max.

quotes