Apple M5 (base · 10 GPU cores · 24 GB) · kernel engineering merged into MLX · Part 3 of the measurement project · 2026-08-02 → 09-02

Taking the measurement discipline into MLX itself

Part 1 measured one laptop and built a map: which operations fall short of the ceilings the machine sustains. Part 2 pointed an evolutionary search at that map and shipped a private speedup. Both end the same way: the kernels are shape-specific and will age out. The one place kernel work outlives the hardware it was tuned on is upstream, in the library itself. So this part moves the work upstream: the findings from the private measurement rig, applied to kernels inside MLX itself, where every user of the library runs them.

What transferred was not any particular kernel. It was the instruments. Every contribution below was built and defended with the same rig Parts 1 and 2 forced into existence: ceilings measured on the actual machine, A/Bs run inside one binary so build noise cannot vote, causal controls that separate mechanism from correlation, and correctness judged against an independent reference rather than against the code under test. Each section shows the mechanism first and the measurement that pinned it down second.

The short version: five changes merged into MLX between August and the start of September. A block-size heuristic for MoE prefill, worth 1.30–1.45× at the shape a real model runs, built on a one-line law the tiles obey. A small-batch matmul tuning validated at 1.32–1.60× on M5 Max silicon. Two correctness fixes in MLX's JIT build path, which no automated test compiles. An independent verification of a third-party one-line fix for a kernel that silently returned wrong results above 32,767 rows. And a two-line dispatch change that lets head-dimension-256 prefill with an array mask run the fused NAX kernel, worth 1.44–1.56× at kernel level and 5–6% of time-to-first-token on Gemma-3-12B.

5 merged changes · 4 kernels + 1 verification · every number from project telemetry or the public threads · generated 2026-08-27 23:55 · #4416 added 2026-09-02

#4023 · MoE block size

1.30–1.45×
routed-expert prefill kernel at the shape one real model actually runs. The maintainer reproduced it before approving: 2.74 → 2.23 ms on their own benchmark.

#4171 · small-batch matmul

1.32–1.60×
quantized matmul at small row counts. Validated by a third contributor on M5 Max silicon.

#4416 · hd=256 array-mask gate

1.44–1.56×
head-dimension-256 prefill with an array mask now runs the fused NAX kernel instead of the unfused graph. Merged 2026-09-02. 1.064× and 1.050× time-to-first-token on Gemma-3-12B at 8k and 16k.

#4372 · JIT build fixes

2 bugs
one reported by a user, one found while verifying the first. Neither shows up in the precompiled build that CI tests.

#3922 · verification

32,768
the row count past which the kernel silently returned wrong results. Not this project's fix: its measurements confirmed the one-line repair on current code, including its untested interaction with a later rewrite of the same loop.

#4023 · the law the tiles obey · % of ceiling vs rows per expert

Throughput of gather_qmm_rhs_nax divided by the fp16 ceiling, measured on this machine at E=256, K=2048, N=512, 4-bit, as rows per expert varies. Bar height is % of that ceiling.
87% · long-run saturation of this op% of the fp16 ceiling (y) 0 R=4: 5.3% of the fp16 ceiling 5.3% R=4 R=8: 10.4% of the fp16 ceiling 10.4% R=8 R=16: 20.8% of the fp16 ceiling 20.8% R=16 THE SHAPE ONE REAL MODEL RUNS R=32: 42.0% of the fp16 ceiling 42.0% R=32 R=64: 77.9% of the fp16 ceiling 77.9% R=64 R=128: 84.8% of the fp16 ceiling 84.8% R=128 R=256: 87.2% of the fp16 ceiling 87.2% R=256 rows per expert (prefill chunk × top-k / experts) · min(1, R/64) predicts every bar to within 0.02
<30% of ceiling30–70%≥70%measured at E=256, K=2048, N=512, 4-bit, σ < 0.2% per point
The kernel walks sorted expert assignments in 64-row tiles and re-runs its whole inner loop once per distinct expert in the tile, computing all 64 rows each pass and keeping one expert's slice. Useful work is therefore min(1, R/64) of the arithmetic, where R is rows per expert, and R is not a property of the model: it is chosen by the serving loop's chunk size. Qwen3.6-35B at its default chunk lands on R=16, the red end of this chart. The causal control: identical work aligned to the tile boundary ran at 11.94 TF; the same work shifted by 32 rows ran at 6.85 TF. Same arithmetic, same bytes; only tile phase differs.

#4023 · from law to heuristic

The fix is small: pick the tile height from rows-per-expert instead of hardcoding 64. The evidence was not small. A same-binary A/B with an environment hook so both arms ran one build; the tile-phase control above, which turns a throughput curve into a mechanism; and an end-to-end number, 62.8 → 58.9 s of time-to-first-token on the real model, against a predicted 5.1 s saving.

Coverage was checked empirically, not by argument: dispatch logging across the model zoo confirmed the heuristic never selects the wrong kernel for shapes other models run. The change was then replicated independently: the maintainer measured 2.74 → 2.23 ms on their own benchmark before merging, a fourth data point from a machine this project never touched.

Configuration · what shipped in #4023

kernelgather_qmm_rhs_nax
changetile height keyed on rows-per-expert
kernel-level19.2% → 25.6–27.8% of ceiling at R=16
end-to-end1.066× TTFT, Qwen3.6-35B @ 32k
reviewer's own check2.74 → 2.23 ms
affected modelsevery MoE model on this chip generation

#4023 · one change, measured at three levels

The same diff before and after, at the kernel (% of ceiling), end-to-end (seconds of time-to-first-token on this machine), and on the maintainer's own benchmark (milliseconds). Each pair uses its own scale; the labels carry the units.
kernel · % of ceilingbefore: 19.2% of the fp16 ceiling at R=1619.2%after: 25.6-27.8% across runs26.7% (midpoint of 25.6–27.8)end-to-end TTFT · sbefore: 62.8 s time-to-first-token, Qwen3.6-35B at 32k62.8after: 58.9 s, same binary, same session58.9 s (lower is better)maintainer's bench · msbefore: 2.74 ms on gather_qmm_bench2.74after: 2.23 ms, reproduced independently before approval2.23 ms (maintainer's machine)
beforeaftereach pair uses its own scale; the labels carry the units
The same diff measured at the kernel, on the full model, and by the MLX maintainer on hardware this project never touched. The three levels agreeing is the point: a claim that only holds at one level is one of the eight hazards Part 1 cataloged.

#4171 · small-batch matmul, replicated on other hardware

The same audit surfaced a sibling: the quantized matmul's tiles are wrong for small row counts, the regime chunked serving visits. The validation worth noting came from elsewhere: a third contributor ran the change on an M5 Max and posted 1.32–1.60× on their own model shapes. Cross-machine replication is the one measurement a single-laptop rig cannot produce for itself.

Configuration · what shipped in #4171

kernelqmm_t_nax
changesmall-M tile configuration
replication1.32–1.60× on M5 Max
replication hardwareM5 Max · independent contributor

#4372 · two fixes in the untested JIT build path

MLX's GPU kernels reach users through two build paths: a precompiled library, and a source path that compiles each kernel the first time it runs. The precompiled path looks kernels up by name and never reads the source templates again; only the source path re-instantiates them. So a whole class of typo survives every test run on the normal build and detonates only for source-path users. A user reported one such breakage; verifying their diagnosis against a source build reproduced it exactly, and the verification pass found a second, unreported bug of the same class sitting next to it: a one-character template-name typo that made a kernel reference a function that does not exist.

The fix is eight lines added, five removed, across both bugs, verified by running the full quantized suite on both build paths: the source path exercised for the first time, the precompiled path with a structural argument that the diff cannot change it, since it touches no instantiation the binary library contains. The bug class remains open upstream: until an automated build compiles the source path, template-level typos keep surviving every test run.

#3922 · verifying a fix in the state a merge would produce

The fourth merged change contains none of this project's code: it is a measurement campaign on someone else's one-line fix for a wrong-results bug in the sorted MoE path. The kernel stores a remaining-row count in a 16-bit integer; above 32,767 rows it wraps negative and whole blocks of output are skipped silently.

The verification method matters more than the result. The fix was applied to current main, the state a merge would actually produce, rather than to the PR's months-old base. On that frame the outputs are simply wrong (off by 31 where the test allows 0.05), the fix makes them correct, and the full quantized suite passes on both build paths. The last check was new information for everyone: a later change had rewritten the same loop, and the fix's interaction with that rewrite had never been tested anywhere.

#4416 · the mask that forfeited the fused kernel

MLX gained a fused NAX kernel for head-dimension-256 prefill on 2026-08-19 (#3842). Its dispatch gate admitted only a causal mask passed as the string "causal". The same causal mask passed as a boolean array, which is what every sliding-window layer produces, fell to the unfused path: the score tensor materialized in full, the mask applied afterwards. Measured at the Gemma-3 shape, the array form ran at less than half the fused form's throughput, and a window mask did no better than a causal array because the unfused path ignores sparsity.

The change is two lines: the gate do_causal && !has_arr_mask becomes do_causal || has_arr_mask. The masked kernel instantiations already shipped in the metal library, so nothing new is compiled. The !has_arr_mask was conservatism, not a guard against a defect: 50 correctness cases against an independent fp32 reference (f16 and bf16; bool, float, window and random masks; key lengths that are not multiples of 32) all pass, worst error 0.0081 in bf16, no NaN or Inf. Kernel-level A/B in one binary: 1.44× at 4k, 1.51× at 32k, 1.56× for a window mask at 32k; the string path unchanged. Upstream's own sdpa_bench.py shows the same rows at 1.46–1.66× (for example 324 → 206 ms per call at the largest shape), which gave the PR a one-command reproduction.

End to end on Gemma-3-12B, 4-bit, same binary with an environment hook, alternating arms, three pairs per context: 1.064× time-to-first-token at 8k and 1.050× at 16k. The modest end-to-end figure was predicted: Gemma's 1,024-token window bounds the keys each masked layer sees per chunk, so a 1.5× kernel compresses to 5–6% of the whole. What remains, named in the PR as future work: the array form still walks every key block and reads the mask, where the string form skips past the diagonal.

Configuration · what shipped in #4416

filescaled_dot_product_attention.cpp
changedispatch gate, +2 / −2 lines
kernel-level1.44–1.56× on array-masked hd=256 prefill
reviewer-native bench1.46–1.66×, sdpa_bench.py hd=256 rows
end-to-end1.064× / 1.050× TTFT, Gemma-3-12B @ 8k / 16k
correctness50/50 vs fp32 reference · test_fast_sdpa green
affected modelssliding-window layers (Gemma 3, gpt-oss), array-mask models, padded batches
merged2026-09-02

How every claim was verified · the rig behind the numbers

Every merged change above passed through the measurement rig Parts 1 and 2 built, and the rig's rules are written down in the project's methodology file, each with the incident that created it. The ones that did the work this month: performance A/Bs run same-binary where possible, with an environment hook selecting old or new behavior inside one build, because two separate builds differ by more than most kernel wins. Where two builds are unavoidable, arms alternate within one session, because this machine's day-to-day drift was measured and absolute numbers do not survive it. Kernel probes time dependent chains, report σ per cell, and repeat a drift-control cell last.

Correctness gets the same treatment as speed: outputs compared against an independent float32 reference built from basic operations, not against the code under test's own relatives; suites run on both build paths after #4372 proved they diverge; and for #3922, the fix was verified on a copy of current main with the patch applied, the state a merge would actually produce, rather than on the PR's stale base.

Numbers quoted from the public threads (the maintainer's 2.74 → 2.23 ms, the M5 Max 1.32–1.60×) are theirs, labeled as theirs, and were not re-derived here.

The instruments · one line each

same-binary A/Benv hook flips old/new inside one build
alternating armsstock/change/stock/change, one session
dependent chainseach call feeds the next; no overlap hides latency
drift controlfirst cell repeated last, must agree
σ gatea win smaller than the noise is not a win
fp32 referencecorrectness judged against independent math
dual buildsevery suite on AOT and JIT paths
merge-framefixes verified on current main + patch
GPU lockone measurement owns the GPU at a time
ceilingsmeasured 120.5 GB/s, never the spec's 153

In closing · findings, lessons, and what comes next

Findings. Five changes merged into MLX in a month, each carrying its mechanism and its measurement. The MoE prefill collapse traces to a single quantity, rows per expert, that the serving loop chooses and the kernel silently punishes; keying the tile height on it recovered 1.30–1.45× at the shape one real model runs. The small-batch matmul tuning replicated at 1.32–1.60× on M5 Max silicon. The JIT build path harbored two bugs invisible to the precompiled build every test exercises. A one-line third-party fix was confirmed against the exact state a merge would produce, including an interaction with a later rewrite nobody had tested. And a two-line gate change gave every array-masked head-dimension-256 prefill the fused kernel it was already entitled to.

Lessons. The engineering lessons generalize past MLX. Parameters chosen by the caller can dominate a kernel's throughput while belonging to nobody's benchmark: rows-per-expert moved this kernel 16× across its range, and no fixed-shape suite would ever have seen it. A claim is strongest where three independent levels agree: kernel probe, end-to-end run, and a replication on foreign hardware. And a codebase with two build paths has two behaviors: correctness proven on one says nothing about the other, as two silent template bugs demonstrated.

Next. A correctness fix for the head-dimension-256 causal path is staged on a local branch, and the decode-attention work it extends is measured and written up in the project's private records. The kernels in this story will be superseded like the ones before them. The instruments that produced and defended them are the part built to last.

Dotted terms carry definitions. Hover, tap, or tab to them. Like this.