One weight read per verify pass: wider affine row tiles in MLX

Summary

Plain decode is one row. A speculative verifier runs several rows at once: \(M = 1 + d\) for a draft of \(d\) tokens. On MLX v0.32.4 those rows land in qmv_wide, which splits them into tiles and re-reads the full quantized weights for every tile. The tile cap was 5, so 6, 7 and 8 rows paid for two reads.

The affine kernel stays register-resident up to 8 rows. Raising only the affine cap to 8 puts 6 to 8 rows in one weight pass. The result is 1.19× to 1.38× on the operation, +4.5% end to end, bit-identical output, and 11 lines changed in two files. oMLX shipped the same idea first, in PR #4475.

The dispatch ladder

Plain token-by-token decode is one row. A verifier checks the draft tokens in one pass, so the row count becomes the draft depth plus one. MLX picks the kernel from that row count:

Rows	Path	What it does
M = 1	qmv_fast	Single-row matrix-vector kernel.
M = 2..12	qmv_wide	Rows tiled across a threadgroup, weights re-read per tile.
M ≥ 13	qmm	Tiled matrix-matrix kernel.

The wide band is where verify passes live. Inside it, cost does not grow smoothly with \(M\). It steps at every tile boundary. On the 8-bit language-model head projection, the step sits between 5 and 6 rows:

Rows M	Latency	Effective bandwidth	What runs
1	1.03 ms	658 GB/s	qmv_fast
5	1.84 ms	368 GB/s	qmv_wide, one tile
6	2.86 ms	236 GB/s	qmv_wide, two tiles
12	5.86 ms	115 GB/s	qmv_wide, three tiles
13	1.79 ms	378 GB/s	qmm

Twelve rows cost 5.7 times one row. Each extra tile reads the same 636 MB of packed weights and about 40 MB of scales and biases again. At 6 rows the kernel reads all 675 MB twice, so effective bandwidth falls from 368 GB/s to 236 GB/s. At 13 rows the dispatch step to qmm makes the operation fast again.

Why every tile re-reads the weights

The wide kernel gives one threadgroup a group of activation rows. Call that group vecs_per_tg. A threadgroup streams the weights once and applies them to every row in its group. The host picks the group size from \(M\) and the tile cap:

// mlx/backend/metal/quantized.cpp, before the change
int n_tiles = (M + 4) / 5;            // ceil(M / 5); tile size caps at 5
int vecs_per_tg = (M + n_tiles - 1) / n_tiles;

With a cap of 5, six rows need two tiles, eleven rows need three, and twelve rows need three. Each tile is an independent threadgroup, so each one loads the weight matrix from memory. Unified memory removes the host-device copy, but it does not remove the read. The weight matrix is far larger than the cache, so the second read comes from DRAM.

The fix is not to make the tile arbitrarily wide. A wider tile means more live values per lane, and registers are the limit.

Why 8 and not more

The affine kernel uses eight lanes to reduce the \(K\) dimension per output row, so one SIMD group covers four rows at a time. The floating-point modes use sixteen lanes and two rows. That difference gives affine twice the row budget.

Inside the inner loop, each streamed row needs its input values and one accumulator per output column. Those values stay in registers across the whole \(K\) loop. Up to 8 rows, the affine kernel stays register-resident. Past 8 it spills. At a cap of 16, twelve rows measured 13.5 ms against 1.79 ms on the qmm path. The spill costs more than the second weight read.

The change

Two files, 11 lines. In the host code, the affine mode gets its own cap. The floating-point modes keep 5.

// mlx/backend/metal/quantized.cpp
-  int n_tiles = (M + 4) / 5; // ceil(M / 5); tile size caps at 5
+  int tile_cap = mode == "affine" ? 8 : 5;
+  int n_tiles = (M + tile_cap - 1) / tile_cap;
   int vecs_per_tg = (M + n_tiles - 1) / n_tiles;

In the kernel file, the affine dispatch table gains three entries:

// mlx/backend/metal/kernels/quantized.metal
 instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 5, 8) \
+instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 6, 8) \
+instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 7, 8) \
+instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 8, 8)

Nothing inside the kernel changes. The lane layout is the same, the accumulation order is the same, and the weight decode is the same. Only the assignment of rows to threadgroups moves. That is why the output stays bit-identical.

Results

Unless noted, the numbers below come from one machine: an M5 Max with 128 GiB, MLX v0.32.4 (commit 36969c0) built from source. Speedup is baseline divided by patched. Each cell is the median of 5 groups of 40 queued calls, with the GPU synchronized per group.

Shape (N, K)	Bits	M = 6	M = 7	M = 8
Head (248320, 2560)	8	1.38×	1.22×	1.19×
Head (248320, 2560)	4	1.34×	1.38×	1.29×
MLP (17408, 5120)	8	1.22×	1.32×	1.25×
MLP (17408, 5120)	4	1.17×	1.28×	1.17×
Narrow (4096, 4096)	8	1.24×	1.26×	1.17×
Narrow (4096, 4096)	4	1.15×	1.29×	1.07×
Small (1024, 4096)	8	1.09×	1.37×	1.20×
Small (1024, 4096)	4	1.09×	1.16×	1.07×

The clearest case is the 8-bit head. Here are the raw timings and the effective bandwidth, counting the packed weights plus the scale and bias arrays:

Rows	Baseline	Patched	Speedup	Baseline GB/s	Patched GB/s
6	2.86 ms	2.07 ms	1.38×	236	326
7	3.99 ms	3.26 ms	1.22×	169	207
8	4.65 ms	3.92 ms	1.19×	145	172

A single row reaches 658 GB/s. The patched numbers move back toward that ceiling. Rows 9 to 16 also improve in most shapes, because the tiling drops from three passes to two. Rows 1 to 5 and rows 13 and above are unchanged.

Second machine: generation 13

The same A/B ran on a second machine with a generation-13 Apple GPU (applegpu_g13s). MLX disables affine qmv_wide on this GPU: use_qmv_wide returns false below architecture generation 15. The patched lines are never reached, so this is a neutrality check, not a speed test.

Rows	Baseline	Patched	Ratio
1	2.16 ms	2.19 ms	0.99×
5	9.06 ms	9.16 ms	0.99×
6	13.88 ms	14.43 ms	0.96×
7	24.04 ms	25.02 ms	0.96×
8	17.31 ms	17.07 ms	1.01×
12	15.13 ms	14.59 ms	1.04×

Median over eight interleaved rounds, with the build order reversed each round. The builds agree within 4%, with no benefit and no regression. Output is bit-identical for M = 1..16.

End to end

The test below uses speculative decoding with six draft tokens, so the verify pass has seven rows.

Setting	Value
Target model	Qwen3-1.7B-4bit
Draft model	Qwen3-0.6B-4bit, 6 draft tokens → M = 7
Workload	128 tokens, temperature 0, fixed seed
Runs	3 per build

Build	Run 1	Run 2	Run 3	Median
Baseline	308.7	307.7	308.9	308.7 tok/s
Patched	321.3	322.5	322.7	322.5 tok/s

That is +4.5% decode throughput. A second run with two independent builds and six interleaved rounds per build put the same change at 313.6 to 327.5 tok/s (+4.45%), with all six patched runs above all six baseline runs. The gain is smaller than the operation-level win because only the verify-pass projections land in the improved band. Attention and single-row decode work are untouched.

Correctness

The operation output is bit-identical to the baseline for every row count from 1 to 16, on the 8-bit head shape and the 4-bit MLP shape. The maximum absolute difference is 0.0.

Check	Result
Operation output vs baseline, M = 1..16	bit-exact, max absolute difference 0.0
python/tests/test_quantized.py	same pass/fail set before and after
Four quantized models, 8 greedy tokens each	identical text before and after
Numeric error vs a dequantized reference	unchanged (about 4×10−3 relative)

Lineage

The idea did not start in MLX. It shipped first in oMLX, a Mac inference server, as PR #4475 [1]: perf(mtp): share large 8-bit projections across partial row tiles. It merged on October 14, 2026 as commit 2c53654.

That patch targets partial tiles directly. A seven-row verify pass, for example, previously divided to one row per pass and ran seven separate weight reads. The oMLX kernel pads the partial tile and computes every row with the one-row arithmetic order, so the output stays bit-exact. The author measured the seven-row head at 2.798 to 1.876 ms, bit-identical. The maintainer then tested it on an M3 Ultra with a real model and reported:

“The output matched main token for token at every depth. With a fixed depth of 4, decode went from 128 to 138 tok/s, and with a fixed depth of 6 from about 123 to 136 tok/s. The default setting (adaptive, max depth 3) doesn't change, since it never hits 5 or 7 rows.”

Fixed depth 4 means 5 verify rows, and fixed depth 6 means 7 rows. Those are the row counts that the MLX tile cap handles badly.

Inside MLX, the work also has a line. #3764 added the small-batch qmv_wide kernel with the 5-row tile [2]. #3791 raised the batch limit that keeps large matrices on it and noted that the tuning limits are empirical [3]. The change here continues that line: the tile is narrower than the hardware allows.

What this does not claim

• The speedup numbers come from one machine, an M5 Max. The generation-13 run only checks neutrality.
• Earlier GPU generations gate affine qmv_wide out before this code path, so the change is neutral there. The generation-13 result confirms this.
• Only a cap of 8 was validated. A cap of 16 spills registers and measured 7.5 times slower at M = 12.
• Single-row decode, prefill, and the qmm path are untouched. The change only moves rows between tiles inside the wide band.
• The end-to-end number is one configuration. A different draft depth lands on a different row count and a different speedup.

Reproduce

The worktree is a clean MLX v0.32.4 checkout with the patch applied. The microbenchmark covers all row counts for one shape. The spec-decode script runs the end-to-end test.

cd mlx-affine-rows
.venv-src/bin/python tests/microbench.py --shape head --bits 8
bash tests/specbench.sh 3 6
.venv-src/bin/python tests/dump_qmm.py patched 8 head

The baseline worktree is a separate clean build, so neither binary is rebuilt between comparisons. The A/B scripts alternate the two builds: tests/opab.sh 5 for the operation and tests/ab2.sh 3 for all three methods. On the generation-13 machine, run_m1.sh builds both trees and m1_ab.sh runs the interleaved check.

References

1. jundot/omlx, PR #4475, “perf(mtp): share large 8-bit projections across partial row tiles”, merged October 14, 2026, commit 2c53654.
2. ml-explore/mlx, PR #3764, “Add small-batch quantized matvec kernel (qmv_wide)”, merged June 26, 2026.
3. ml-explore/mlx, PR #3791, “Raise qmv batch limit for large matrices on M5-class GPUs”, merged August 8, 2026.
4. ml-explore/mlx, v0.32.4, commit 36969c0. Files: mlx/backend/metal/quantized.cpp, mlx/backend/metal/kernels/quantized.metal.