Squeezing a 26B diffusion LLM onto a Mac

The goal was to make DiffusionGemma - a 26B-parameter, ~4B-active mixture-of-experts block-diffusion text model - fast enough to be a real interactive tool on a single Apple M5 Pro (48 GB, macOS 27 beta). The stock MLX inference path ran at ~18 tok/s on the original short repro, and ~23-27 tok/s on a fair 512-token run.

This is a write-up of what was tried, what actually worked, what it cost, and where the hard floor turned out to be.

An opt-in turbo engine was built. On a like-for-like benchmark, GSM8K[:20] at temperature 0, in a fresh process per config:

• 35.3 → 106.1 tok/s median excluding warmup, a 3.0× throughput gain.
• 17/20 → 13/20 accuracy, an 85% to 65% drop.
• The four lost questions (q7, q13, q14, q17) are all long multi-step problems.

On other workloads it runs ~65-90 tok/s on sustained technical and long-form text, bursting to ~140 on short Q&A. The speedup is real and repeatable. It is also not free, and the accuracy cost lands entirely on the hard problems, which is the least convenient place for it to land.

The stretch goal was 300+ tok/s. It was not reached, and the most useful outcome was not the speedup but the reason it could not be: a measured bandwidth ceiling. Single-stream 4-bit MoE on this machine is weight-streaming bound at ~80 GB/s, and that ceiling holds even for Apple's own tuned Metal-4 neural-accelerator kernel. 300+ tok/s single-stream is not a tuning problem; it requires either beating Apple's kernel (open research) or dropping the experts below 4-bit. The numbers behind that claim are shown below.

Along the way, two real MLX bugs turned up, both reproduced on-device.

---

What DiffusionGemma is, and why naive inference is slow

Autoregressive LLMs generate one token at a time. Block-diffusion models work differently: they hold a fixed-size canvas of token positions (256 here) and denoise the whole canvas over several forward passes, freezing positions as they become confident. Instead of 256 sequential decode steps, ~8-16 full-canvas forwards are performed. In principle that is a throughput win, because each forward commits many tokens at once.

The catch is what each forward costs. DiffusionGemma is a mixture of experts: 128 experts, top-8 routing, 30 layers, 4-bit quantized, 262,144-token vocabulary. One denoising step was profiled at the native 256-token canvas on the M5 Pro:

The MoE dominates everything. The reason, in one measured fact: at a 256-token canvas, the 2,048 token-expert pairs (256 × top-8) light up essentially all 128 experts in all 30 layers, every step. That is ~12.8 GB of 4-bit expert weights read per denoising step. The whole single-stream speed question collapses to one: how fast can 4-bit expert weights be streamed?

Everything else in the table is "soft" - work that is necessary for the reference sampler but cheap to approximate or skip. The MoE is the hard floor. This split matters: the optimizations below attack the soft rows, while the ceiling analysis is entirely about the hard one.

---

The optimizations that worked

Every technique below is opt-in behind --gen-kwargs '{"diffusion_turbo": true, ...}'. With the flag off, the code path is byte-for-byte identical to stock mlx-vlm, and the reference loop is never touched. Each technique is independently toggleable.

1. Hierarchical exact top-K sampler chain - ~35 ms → ~3 ms/step

A direct mx.topk over 262,144 logits costs ~40 ms on this GPU. But the sampler only ever needs the top ~64 candidates per position. A two-stage chunk-max selection is used instead: the logits are reshaped into chunks, chunk maxima are taken, the top-K chunks are picked, then top-K within. ~1.5 ms.

It is provably exact, not an approximation. If a value is in the global top-K, fewer than K chunk-maxima can exceed it, so its chunk is among the top-K chunks, and the selected chunks therefore contain every global top-K value. And because softcap (tanh) and temperature scaling are both monotonic, top-K membership computed on the raw logits is identical to membership after softcap. So softcap, softmax, entropy, and sampling all run on a [T, 64] tensor instead of [T, 262144]. The only quantity truncated is the entropy tail mass beyond 64 candidates, negligible at the low temperatures where positions actually commit.

This is the piece of the engine most worth keeping. It is exact, it is model-agnostic, and it costs nothing in quality.

2. Top-K self-conditioning - drops a 378-GFLOP matmul and a 1.5 GB table

DiffusionGemma feeds its own predictions back as conditioning: the reference computes softmax(logits) @ embedding_table. That is a ~378-GFLOP matmul and it requires keeping a dequantized ~1.5 GB copy of the embedding table resident in memory.

Instead, the 64 embedding rows for the top-K candidates are gathered and a probability-weighted sum is taken. This drops the dequantized table entirely, and at the low temperatures where positions commit, the conditioning signal stays close to the reference on short and easy inputs.

This is the only lossy shortcut in the engine, and it is the one that costs accuracy. On short and easy problems the approximation stays close to faithful, but the drift compounds over long reasoning chains, because each step's error feeds into the next step's conditioning. On GSM8K[:20] it is the direct cause of the drop from 17/20 to 13/20, and the four lost questions are all long multi-step problems. Cheap on throughput, not free on hard long-form accuracy.

3. Compacted active-set decoder - forward only the live positions

With monotone commits, once a position clears a confidence threshold it freezes, so there is no reason to keep forwarding it. The compacted decoder (turbo_compact) forwards only the live positions, plus the positions committed on the previous step so their K/V reflect the committed token. Frozen positions are served from per-layer canvas K/V buffers, with RoPE for the scattered live positions applied from precomputed cos/sin tables that replicate mx.fast.rope.

The correctness anchor here is a bit-exactness invariant: with the forward set equal to all positions and an empty freeze state, one runner step matches the stock decoder forward bit-for-bit. Both RoPE variants (the sliding and the proportional partial-rotary full-attention paths) were validated numerically against the model's own modules. The shrinking working set also helps the MoE: below ~48 live tokens MLX drops onto its faster gemv-class dispatch.

Like the sampler chain, this one is exact. It changes what gets computed, not what the answer is.

4. Active-set shape bucketing - fixing the "memory leak" that wasn't

This one was a trap. Over a long-lived process, throughput drifted from ~120 down to ~20 tok/s over ~20 minutes. It looked exactly like a memory leak.

It was not. Peak memory was flat the whole time. The real cause: the compacted decoder forwards a different number of live positions every step (256 → ~16 as tokens freeze), and each distinct shape makes MLX JIT-compile and cache a fresh graph. Over time the compiled-kernel set grew unbounded and throughput decayed. It was compiled-kernel-cache growth, not a buffer leak.

The fix is to round both the forward set and the active set up to ≤10 fixed buckets (16, 32, ..., 256), padding with duplicate positions. Duplicates recompute identical values and scatter idempotently, so results are unchanged, and only the distinct-shape count is bounded, from ~240 possible down to ≤10. (An mx.clear_cache() every N steps was the first approach tried and proved actively harmful: it zeros the buffer pool and causes realloc churn. Bucketing is the correct answer.)

5. Two-stage confidence schedule + repair + EOS tail-drop

• Two-stage schedule. A picky threshold (e.g. 0.95) while the canvas structure is still forming, then a relaxed one (e.g. 0.80) for stragglers. No forced final flush. A forced flush argmax-commits every still-uncertain position on the last step; those marginal tokens get re-encoded into the next canvas as context, compounding across blocks into single-token repetition ("the the the..." collapse). The two-stage schedule avoids it by only ever committing high-confidence positions.
• Repair pass (turbo_repair). One extra full-canvas forward after the canvas drains, re-argmaxing every position. This recovers the reference sampler's "final canvas = argmax of last forward" semantics and gives early-frozen tokens a chance to be revised. It narrows the accuracy gap. It does not close it.
• EOS tail-drop. Once a confident EOS freezes, positions after it stop being denoised. Big win on real chat answers, where the stock loop denoises all 256 positions regardless of how short the answer is.
• Repeat-guard. Stops generation if a token repeats ≥16 times in a row. Natural text never does this, so it is an unambiguous degeneration signal, and it prevents emitting a wall of garbage up to max_tokens.

---

Measured against the baseline

512-token generation, fresh process, baseline = the stock engine's own 512-token throughput:

Creative long-form gains least, because it commits fewer positions per step and drains the canvas more slowly.

GSM8K[:20], accuracy and speed together. Throughput on its own is easy to overstate, so the follow-up run measured accuracy alongside it: 20 questions, temperature 0, seed 0, each config in its own fresh process so the Metal allocator's throughput buildup does not leak between runs.

The throughput gain is a consistent 3.0× and holds across both turbo configs. Dropping from 10 steps to 8 buys essentially nothing here, which is itself informative: by 10 steps the canvas has already drained, so the step cap is not what is binding.

The accuracy cost is specific rather than diffuse. Both turbo configs miss the identical set, and they lose exactly the four questions the baseline gets right: q7, q13, q14, q17. All four are longer multi-step problems. Three of the four run into the 320-token cap, and the detail worth noting is that this is not simple truncation: on q14 the baseline finished naturally in 296 tokens while turbo ran all the way into the cap and still got it wrong. Turbo did not run out of room so much as wander. That is the signature of the top-K self-conditioning shortcut from section 2: near-faithful on short and easy items, compounding error on long reasoning chains.

One caveat worth stating plainly, because it is easy to get wrong. The two configs above both use a forced flush. The recommended interactive default is the two-stage no-flush schedule (turbo_threshold: [0.95, 0.80, 8] with turbo_repair), and it was not part of this 20-question run. It was only measured on an 8-question subset, where it scored 7/8 at ~110 tok/s. Eight questions is far too small a sample to claim it closes the gap, and an earlier 11/12 result on a 12-item subset had looked like a match to the faithful reference only because that subset happened to exclude most of the long multi-step problems where the gap appears. The 20-item run is the honest number, and the honest number is 13/20.

---

The honest failures

Five custom Metal kernels that didn't beat MLX

The whole game is streaming 4-bit expert weights faster. Five custom Metal kernels were written to try to do it:

1. Register-resident multi-row qmv - 27-37 GB/s. Holding M×8 activation registers spills. Dead.
2. simdgroup-matrix (simdgroup_half8x8) MoE - 50-55 GB/s and numerically buggy. Dead.
3. Capacity-padded NAX gather - 74-84 GB/s. Equal to MLX, no win.
4. One-simdgroup-per-row, weights shared via L2 - 15 GB/s. The hypothesis was that 16 simdgroups reading the same weights would let L2 dedup the DRAM reads. Disproven: the GPU does not keep simdgroups in lockstep, so each one re-reads from DRAM. This is why weight reuse fundamentally needs explicit threadgroup staging.
5. Metal-4 mpp::tensor_ops::matmul2d from a JIT kernel - this one turned up something usable: #include <MetalPerformancePrimitives/...> works from a JIT mx.fast.metal_kernel on macOS 27, and the Metal 4.0 language version is honored. But the tensor-op operands want real tensor/tensor_handle device views; the pointer-backed tensor_inline that can be built from JIT's raw device T* args is not honored, and the output comes out scrambled. A correct NAX MoE kernel has to be a C++ primitive inside MLX core, reusing the steel NAXFrag machinery, and since MLX's existing NAX kernel already hits the ceiling, a faster one is open kernel research, not a tile tweak.

Best of the five: ~84 GB/s. MLX's own kernel: ~80. The wall could be matched; it could not be beaten. These kernels are kept as research artifacts only, not wired into the engine. They document the ceiling; they do not move it.

---

The physics ceiling: why 300+ tok/s is blocked

This part was measured most carefully, and the conclusion is correspondingly firm.

The single-stream MoE forward is dominated by reading ~12.8 GB of 4-bit weights per step. Exactly one thing was therefore benchmarked: 4-bit quantized matmul throughput as a function of M, the number of activation rows that reuse each weight. Same weights, varying only M:

For calibration on the same machine: elementwise streaming hits ~263 GB/s, and MLX's own qmv gemv kernel hits ~220 GB/s. The DRAM and the dequant front-end are both clearly capable of >200 GB/s.

Reading the curve: at M ≤ 2 there is no weight reuse, each weight is consumed once, it is a straight gemv, and it streams at the full ~220 GB/s. But the moment M grows past ~4, reusing a weight row across multiple activation rows requires staging it through threadgroup memory, and every staged path collapses to ~80 GB/s. From M=16 through M=64 it is a flat plateau.

And the MoE needs exactly the M that sits in the bad regime:

So the model gives up ~2.7× of available bandwidth precisely where it spends ~80% of its time.

This is not merely a case of "MLX's kernel is slow." The ~80 GB/s plateau includes Apple's own Metal-4 NAX kernel. The gather_qmm NAX path uses mpp::tensor_ops::matmul2d (the NAXFrag machinery in steel/gemm/nax.h). At the MoE's shapes, that tensor-op kernel is the 80 GB/s number. The 220 GB/s only exists at M ≤ 2, where there is no reuse to stage. This is a property of staged 4-bit weight-reuse on this hardware generation, confirmed against the best kernel Apple ships.

The arithmetic is simple. Against 12.8 GB/step at ~80 GB/s, the MoE alone is ~160 ms/step; at the 8-14 steps coherent text needs, single-stream throughput is capped well below 300 long before any other factor matters. 300+ tok/s single-stream is blocked by a measured bandwidth ceiling, not by software that can easily be improved.

Two levers could break it, and both were measured:

• Batch multiple canvases (denoise B at once; the weight read is flat in token count). Measured 1.0 / 1.26 / 1.47× at B = 1/2/4, sublinear, because attention, the dense MLP, the LM head, and the MoE compute all scale with tokens, and only the weight read is flat. This is multi-stream throughput (~190 tok/s aggregate at best), not single-stream latency.
• 2-3 bit experts - directly cut the bytes the wall is made of (~1.4-2×, plausibly enough to reach the 300 target). Untried: it needs the bf16 checkpoint (~52 GB, not local) to requantize the experts cleanly without compounding 4-bit error, and it carries real quality risk. This is the most promising next throughput experiment.

Neither is a kernel tweak. The wall is real.

The measured ceiling was written up as an mlx discussion, framed as a question to the kernel authors: is small-M tile tuning of gather_qmm_rhs_nax (which carries a literal // TODO: Tune the block sizes) tractable, or is ~80 GB/s the understood staged ceiling on this generation? The answer either way redirects future effort.

---

Two real MLX bugs found along the way

1. gather_qmm NAX kernel-name mismatch (a real, one-line bug). On a NAX-capable GPU, gather_qmm_nax() builds its Metal kernel name with the tile parameter bk = 32, but the NAX gather kernels are only ever instantiated with bk = 64. So on M5 every batched expert matmul that takes the NAX fast path fails at dispatch:

The fix is bk = 32 → bk = 64. bk only feeds the kernel-name string in that dispatch, not the geometry, so it is a pure "ask for the kernel that exists" change. The sibling qmm_nax() already uses 64, so it reads as a copy-paste slip. The same fix is ml-explore/mlx#3632, since merged; what this work adds is an independent on-device reproduction on M5 Pro / macOS 27 with the exact failing kernel name. A small MLX_DISABLE_QMM_NAX env hatch was also built for A/B-ing the NAX path against the steel kernels (on M5 the steel path is ~2× worse, so NAX is correctly the default).

2. int64 scatter crashes the Metal JIT. On macOS 27 beta, put_along_axis / scatter_add_axis with a 64-bit element dtype do not fail cleanly, they detonate the Metal library build. The fallback-atomic union computes packing_size<long> = sizeof(uint)/sizeof(long) = 4/8 = 0, which declares a zero-length array (T val[0]) and divides by zero. The plain Scatter primitive already guards 64-bit dtypes with a clean ValueError; the ScatterAxis path is missing that guard. The fix is to mirror the existing guard, or to properly support 64-bit element types. Reported as ml-explore/mlx#3690 with a minimal repro and dtype matrix, and a fix has been proposed. The turbo engine sidesteps it by using int32 throughout, which is also why the engine is int32-disciplined by design.

---

The takeaway

Measured against the faithful baseline on the same benchmark, the engine is a 3.0× throughput gain (35.3 → 106.1 tok/s on GSM8K[:20]) that costs 20 points of accuracy (85% → 65%), with every lost question a long multi-step problem. That trade is good for interactive chat, where answers are short and the approximation stays close to faithful, and bad for anything that has to reason for 300 tokens without drifting.

The optimizations that got there were the boring-but-correct kind: exact top-K instead of full-vocab, a compacted decoder with a bit-exactness invariant, shape bucketing to bound the kernel cache, and a confidence schedule that does not collapse. Three of those four are exact and cost nothing. The accuracy loss traces to exactly one shortcut, gathered self-conditioning, which makes it a tractable thing to fix rather than a diffuse tax on the whole engine.

The result worth keeping is the one that says no. Five custom kernels and a careful bandwidth sweep proved that 300+ tok/s single-stream is not an unclosed software gap, it is a hardware and precision ceiling: 4-bit weight-reuse GEMM tops out at ~80 GB/s on this machine, including Apple's own tuned kernel, exactly at the tile shape the MoE needs. Knowing precisely why further speedup is impossible, and being able to prove it, is worth as much as the speedup itself. The levers that remain are clear, lower-bit experts and multi-canvas batching for throughput, restoring self-conditioning fidelity on long generations for accuracy. What the measurements changed is that each of those is now a well-posed question with a known cost, rather than a guess.