Multi-row decode, algorithmic discovery, and honest matched-representation benchmarks - #9
Merged
Merged
Conversation
Serve rows 2-31 in the dense BF16 decode path. MDotAccumulate and MPairedDotAccumulate now carry a `rows` dimension so each weight fragment is loaded once and reused across every activation row, keeping weight traffic flat while the work scales. MLX falls from ~110 GB/s at one row to 47-76 GB/s for rows 2-31; the multi-row QMV holds 105-122 GB/s, so the full MLP block runs 1.5-1.8x faster there and the down projection alone up to 2.4x. This is the band a speculative-decode verification pass runs in. For rows > 1 the tuner requires every row to equal MLX's own single-row result rather than its multi-row tile kernel, so a batched step stays bit-identical to decoding those tokens one at a time. Fix two defects that let the tournaments select kernels slower than native MLX: - Every tuner timed one dispatch per sample. The blocking mx.eval round trip costs ~200 us whatever the kernel does, so it was added to each candidate and compressed their ratios toward 1.0, letting switch margins admit losers. calibrate_tournament_batch now sizes a batch per timed sample, and the measurement source joins each backend signature so stale picks invalidate. - _COMPILED_SWITCH_MARGIN was 0.005, below the run-to-run noise floor, so the mx.compile variant won on noise and then measured 0.63-0.90x in steady state. Measured at matched representation, BF16 rows=128 went 0.82x -> 1.11x and every INT4 cell recovered to parity. Skip candidate filtering on schedule-cache hits: it rebuilt the filtered config tuple on every dispatch, 70% of the Python cost of a decode call. Add benchmarks/matched_representation_matrix.py, which runs identical weights in identical formats on both sides so a result is a kernel comparison rather than a representation change. Replace the model charts with three that report one unit, a multiplier against native MLX, and graph only matched-representation runs. The BF16 capacity suite compares meTile INT8 decode against MLX BF16, so it is no longer plotted; its numbers stay in the README table and committed JSON. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> (cherry picked from commit c782369)
Fold attention_discovery into algo_discovery so one module covers every rewrite the compiler can find. Both rewrites are licensed by the same proven law: the weighted-softmax monoid is proved once, flash attention uses all three components and online softmax the (maximum, normalizer) projection with the value fixed at 1. Adding a rewrite is now a matcher, not a new proof. Reduction laws and certificate types are re-exported, so composing discovery needs one import. Route discovery through the exact min-cut selector graph_fusion already used. Overlapping rewrites are mutually exclusive, so choosing between them is maximum-weight independent set, and the project-selection reduction solves it exactly on bipartite conflict components. find_flash_attention previously took whichever region it hit first. Make the online softmax rewrite work; it is now on by default and 8/8. Three fixes: - It reused the tile-max chain's Constant as the running scalar maximum, but lowering promotes that seed to a tile-typed accumulator. Two dedicated scalar seeds instead, and both post-loop reductions dropped, since re-reducing an already threadgroup-wide value inflates the normalizer by the thread count. - _emit_vec4_op fell through for MThreadgroupReduce and emitted simd_max(float4). - A threadgroup reduction cannot sit inside the ragged tail's mask branch, because masked threads never reach the barrier. ForRange.masked_identity plus _predicate_masked_load guard only the load and seed the value with an identity the law supplies: -inf is the identity of the maximum, and exp(-inf - m) = 0 is the identity of the normalizer, so one value serves both reductions. Measured 1.28x against a 1.33x four-transfers-to-three ceiling. It is a compiler result, not a competitive one: MLX's softmax is a single fused kernel and still 0.74x to 0.99x faster than ours. Fix a crash the multi-row tuner introduced. MLX-LM passes rank-3 [batch, sequence, hidden], and slicing axis 0 for the per-row reference yields empty rows once batch is smaller than the row count, which segfaults inside eval. Rows now come from a flattened view, and the reference is only built in the QMV band where it is actually checked. Stop materializing a causal bias tensor in the attention tournament's native reference. MLX masks causally without allocating anything and the two agree bitwise, including for ragged query counts. Building the bias made the baseline look slower than it is and biased selection toward the generated kernel; causal went 0.74x to 1.01x. Add three affine tilings for wide outputs. Sweeping the full legal space found the shipped six topped out at parity for N >= 8192 while 64x128 hilbert reaches 1.10x. Include the measurement source in the framework and attention persistent keys, so changing how candidates are timed invalidates stored picks instead of leaving stale ones. Add benchmarks for shape sensitivity, graph fusion, and mx.compile comparison, and chart the two findings that explain the model results: MLX changes kernel between output widths 2048 and 2560 and the one below is poor, which is why Llama 3.2 1B gains 3.18x and 3B only 1.08x; and batching should hold weight bandwidth flat, which meTile does for BF16 and nothing yet does for INT4 or INT8. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> (cherry picked from commit b2b5fa2)
Re-measure the batch-size matrix and shape sensitivity at HEAD so the charts, the JSON, and the README all quote the same run. Two results moved: - BF16 rows=512 went 0.85x to 0.99x. The earlier dip was the tournament picking a losing schedule, not a kernel limit, so the chart no longer shows a regression that does not exist. - INT4 at width 8192 went 1.01x to 1.10x, now that the wide-output tilings added in the previous commit are in the candidate list. Cut every chart subtitle to one short line. They were carrying explanation that belongs in prose, and the batch chart's ran off the canvas. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> (cherry picked from commit 759969c)
The batch chart drew MLX INT4 and INT8 with no meTile counterpart, which reads as missing data rather than as a result. The benchmark was recording None for those points on the grounds that meTile has no multi-row quantized kernel, which was an assumption rather than a measurement. Run the quantized executor path and record it. It does track MLX, within 1.0% for INT4 and 5.5% for INT8, so the label now says meTile defers to MLX rather than leaving a gap. Time the batch sweep interleaved as well. Measuring one side after the other put INT8 at 76 GB/s against MLX's 116 at one row, which contradicted the batch-size matrix; alternating the two removes it and the formats agree. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> (cherry picked from commit 67d0ca9)
For INT4 and INT8 meTile calls MLX's kernel rather than running one of its own, so a single line describes both backends. "meTile defers to MLX" explained an implementation detail; "INT4, both" says what the line is. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> (cherry picked from commit ea1ea1e)
…width The lint job caught three things the test job could not. mlx_block_scaled.py called calibrate_tournament_batch without importing it, so block-scaled tuning raised NameError as soon as it reached the tournament. Every other backend imports it; this one was missed. No test covers that path. The rest is mechanical: I had hand-wrapped near 88 columns while the project formats at 100, and the redundant "# noqa: E402" markers in benchmarks are already covered by per-file-ignores. B023 is ignored for benchmarks rather than silenced per line. The timing helpers take the work as a closure and call it before returning, so a closure built in a sweep loop never outlives the iteration that built it, and none are stored in a container. Binding each captured tensor as a default argument would only obscure what is being measured. 582 tests pass; ruff format, ruff check, and vulture are clean. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Both fused affine paths are single-row by construction: the SwiGLU kernels raise "native affine SwiGLU schedules require one decode row" and the residual QMV raises its equivalent. So from two rows up the tournament had only the scalar kernels, which lose, and it settled on native MLX for the entire batched band. That is why the matrix reports int4 at parity there. meTile already generates a multi-row affine matmul, and on a single projection it holds ~60 GB/s flat across rows 8 to 16 while MLX falls from 68 to 35. This adds two candidates built from it: a SwiGLU composed of two multi-row matmuls, and a down projection that adds the residual afterwards. Nothing is forced. Both go through the same tournament and the same numerical compatibility gate as every other candidate, and from_mlx rejects anything but 4-bit group-64 so they cannot compete at formats they do not support. Measured on the MLP block 1536 -> 8960 -> 1536, interleaved against native MLX: rows 1 2 4 8 16 32 before 0.92 0.98 0.94 1.03 1.00 1.04 after 0.95 0.96 1.01 1.28 1.31 1.24 Rows 1 to 4 are unchanged within noise and still pick native. Those sit below 1.0 in both columns, which is the pre-existing mlx_compiled selection problem, not something this introduces. 585 tests pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
AndreSlavescu
enabled auto-merge (squash)
July 29, 2026 07:16
How long a candidate measures depends on how many others share the round-robin, so ranking finalists from one big rotation can prefer a kernel that loses when the two are timed against each other. The holdout that re-measured in a small group only ran for one-row decode; every multi-row shape committed to whatever the crowded rotation ranked first. Each finalist is now timed in its own two-way round-robin against native, and ranked on the ratio to native rather than the raw time, so drift between one pairing and the next cancels. Dense SwiGLU 1536 -> 8960, paired against native MLX: rows 2 4 8 16 before 1.43 1.40 1.40 1.42 after 1.49 1.48 1.48 1.47 Selections also settle down, picking the same shape across neighbouring row counts instead of a different one each time. Two things tried and rejected, recorded so they are not retried blind: Raising _MAX_QMV_ACCUMULATOR_PAIRS from 16 to 32 does nothing. The bound looked too tight because outputs_per_simdgroup=2 at 16 rows compiles to 115 registers against G17's 140-register budget, so it cannot be spilling. But the tuner already reaches the same speed with outputs_per_simdgroup=1 at a larger simdgroup count, and the wider search bought no speed. Left at 16. The 1.28x I expected to recover here was not real. It came from comparing the tuner's pick against the best row of a partial hand-built table that omitted the simdgroup counts the tuner actually explores. Measured against the tuner's real baseline the gain is the 3 to 6% above. 586 tests pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Metal exposes no register count, and the obvious stand-in does not work: maxTotalThreadsPerThreadgroup reads 1024 on an M5 whether a kernel holds 4 live floats or 512. The compiler does record the number, in a metadata segment of the GPU binary, and MTLBinaryArchive.serialize is the only public way to get that binary onto disk. Three unwraps to reach it. The archive is a fat file whose applegpu_* slice is the GPU code; that slice's __compute section is itself a Mach-O; inside it __GPU_METADATA is a FlatBuffer and __text is the machine code. The count is field 0 of the table referenced by field 0 of the root. Read by path, never by byte offset. The buffer embeds the kernel name and signature, so a fixed offset drifts between kernels: reading byte 188 worked for one probe kernel and silently reported a neighbouring field, 24 registers and "spilling", for every real one. --self-check pins the reader against known counts so a bad read fails loudly. What it establishes on G17: the budget is 140 registers per thread, the count tracks live floats plus 4 up to 124, and kernels that reach 140 are spilling and measured 1.3x to 6.7x slower than lower-register siblings. Applied to the dense SwiGLU bound it says the bound is already right: every config admitted by _candidate_configs peaks at 99 of 140 registers, 71% of budget, so nothing admitted can spill. That is the intended use, checking a scheduling bound rather than trusting it, and it is why the tool ships despite having produced no speedup of its own. Analysis tool only. Each reading costs a Metal compile of about a second, far too slow for the tuner to call, and it degrades to a message rather than an error when swiftc is absent. 586 tests pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The compiler core has never depended on MLX: frontend, ir, compiler, codegen and runtime are 17k lines with no mlx import, and metile/runtime/metal_device.py reaches Metal directly through ctypes, which is how the gemm, softmax, rmsnorm, attention, fft and mlp benchmarks execute. MLX is a comparison target there. What had drifted is that the machinery for choosing between kernels grew inside metile/backends/mlx*.py, where eight files each carry their own round-robin, switch margins and cache handling. None of that is MLX-specific. A second execution backend would have to duplicate all of it. metile/tuning now holds the backend-agnostic part: round_robin for cheap triage of a large field, confirm_pairwise for deciding, select_fastest for the margin and tie-break policy, token_bucket for shape bucketing. It knows nothing about kernels or how to run one; callers pass a measure function that turns a thunk into seconds. mlx_dense_swiglu is ported onto it as the first consumer and its behaviour is unchanged. The remaining seven backends still hold their own copies. Correcting the previous commit while it is still unmerged: it claimed pairwise confirmation was worth 3 to 6%, reading 1.40-1.43x going to 1.45-1.49x. That was a trend read out of variance. Interleaved against the pre-change commit, four alternating pairs, medians by row count: rows 2 4 8 16 before 1.42 1.48 1.45 1.49 after 1.48 1.45 1.45 1.53 Mixed, and inside a run-to-run spread of about 0.07. There is no measurable speedup. The reason to keep pairwise confirmation is that it makes selection principled rather than fast: every candidate is measured in an identically sized context, so the ranking no longer depends on how crowded the rotation happens to be. Ranking on the ratio to the baseline is what makes separate pairings comparable when the baseline drifts. 587 tests pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this does
Three things, plus the benchmarks and charts that justify them.
Serve rows 2 to 31 in the dense BF16 decode path
MLX drops from ~110 GB/s at one row to 47-76 GB/s for rows 2-31.
MDotAccumulateandMPairedDotAccumulatenow carry arowsdimension so each weight fragment is loaded onceand reused across every activation row, keeping weight traffic flat while the work scales.
The MLP block runs 1.58x to 1.79x faster there, and the down projection alone up to
2.4x. This is the band a speculative-decode verification pass runs in.
For rows > 1 the tuner requires every row to equal MLX's own single-row result rather
than its multi-row tile kernel, so a batched step stays bit-identical to decoding those
tokens one at a time.
Consolidate algorithmic discovery into one module
attention_discoveryfolds intoalgo_discovery. Both rewrites are now licensed by thesame proven law: the weighted-softmax monoid is proved once, flash attention uses all
three components and online softmax the
(maximum, normalizer)projection with the valuefixed at 1. Adding a rewrite is a matcher, not a new proof.
Discovery candidates now go through the exact min-cut selector
graph_fusionalreadyused, so overlapping rewrites resolve optimally instead of by iteration order.
Online softmax works and is on by default (it previously failed to compile). Measured
1.28x against a 1.33x four-transfers-to-three ceiling. That is a compiler result, not a
competitive one: MLX's softmax is a single fused kernel and still faster than ours.
The ragged-tail fix is the interesting part. A threadgroup reduction cannot sit inside a
mask branch, because masked threads never reach the barrier.
ForRange.masked_identityplus
_predicate_masked_loadguard only the load and seed the value with an identity thelaw supplies:
-infis the identity of the maximum, andexp(-inf - m) = 0is theidentity of the normalizer, so one value serves both reductions.
Stop the tournaments picking kernels slower than MLX
Two defects. Every tuner timed one dispatch per sample, and the blocking
mx.evalroundtrip costs ~200 us whatever the kernel does, so it was added to each candidate and
compressed their ratios toward 1.0. And
_COMPILED_SWITCH_MARGINwas 0.005, below therun-to-run noise floor, so the
mx.compilevariant won on noise and then measured0.63-0.90x in steady state.
BF16 rows=128 went 0.82x to 1.11x and every INT4 cell recovered to parity.
Results
All matched representation, Apple M5 / MLX 0.32.0, identical weights and format on both
sides.
BF16 results are bit-identical to running those rows through MLX one at a time.
What this does not do
Stated plainly because the README now says so too:
this machine and MLX already runs at 120-126 GB/s. Batching is what moves it.
so those chart lines describe both backends. This is the largest piece of unclaimed
performance left: at 16 rows MLX moves 34 GB/s of INT4 weights out of a possible 121.
can run, not because we have not got to them.
Benchmarks and charts
New:
shape_sensitivity.py,graph_fusion_speedup.py,compile_comparison.py, and ashared
chartstyle. The charts are all in multiplier units against native MLX, and onlymatched-representation runs are graphed. The BF16-to-INT8 capacity comparison is
deliberately not plotted, since comparing two weight representations is not a speedup.
Two charts explain the model results: MLX changes kernel between output widths 2048 and
2560 and the one below is poor, which is why Llama 3.2 1B gains 3.18x and 3B only 1.06x;
and batching should hold weight bandwidth flat, which meTile does for BF16 and nothing yet
does for INT4 or INT8.
Notes
mainand cherry-picked fromalgo-discovery, whose earlier commits arealready on
mainvia the squash in Add composable Metal compiler and guarded MLX runtime #8. Tree is identical toalgo-discovery.candidates are alive, so a kernel measuring 1.10x alone can read 0.88x inside the
round-robin. Selections are unstable at some shapes as a result.
🤖 Generated with Claude Code