Pick BM from rows per expert in gather_qmm_rhs_nax

[dwijenpatel]

Following up on #3925, where you invited a PR that tunes the block sizes.

affine_gather_qmm_rhs_nax walks the sorted index array and re-runs its entire K loop once per distinct expert inside a BM-row block, computing all BM rows on each pass and keeping only that expert's rows through store_slice. Useful arithmetic is therefore about min(1, rows_per_expert / BM). The run length is a caller-side quantity: for MoE prefill it is tokens * top_k / num_experts. A 256-expert top-8 model prefilling 512-token chunks presents runs of 16 rows against a block of 64, so three quarters of the arithmetic is computed and discarded.

Bandwidth is set by the number of experts a block touches rather than by BM. Each expert's BN by K slab is read once per block it appears in, and a 16-row run sits inside one block whether that block is 32 rows or 64. A shorter block removes arithmetic and adds no traffic.

The change. Dispatch BM=32 when the average run is shorter than 64 rows, and add the matching instantiations for the affine and fp modes. Runs of 64 rows and longer keep BM=64 and the kernel they use today.

Measurement. Both arms were built from this tree one after the other and measured in the same session by one script: chained up and down projection pairs, 4-bit gs64, E=256, K=2048, N=512, median of five dependent strips after a warmup pass. Percentages are of this device's measured fp16 tensor-core peak of 15.4 TFLOPS, on a base M5 with a 10-core GPU. Within-arm sigma stayed at or below 0.9%.

rows/expert	main	this PR	ratio
4	5.3%	7.7%	1.45x
8	10.5%	15.3%	1.46x
16	20.8%	29.5%	1.42x
32	41.5%	53.1%	1.28x
64	75.1%	76.6%	1.02x
128	84.3%	84.4%	1.00x

The last two rows take the unchanged path in both builds, so they act as a control: they agree within 2%, which is the process-to-process floor I see on this machine. Read the rows above against that floor rather than against the within-arm sigma. The curve also tracks min(1, run / BM) across the whole range, which is the evidence that the block height drives the gap.

The benchmark case added here, which is the same shape at 16 rows per expert, reads 9.56 ms on main and 7.45 ms with the patch. Its equivalent_matmul control reads 1.814 ms and 1.818 ms across the two builds.

End to end. On a 4-bit Qwen3.6-35B-A3B under mlx-lm at 32k context with 512-token prefill chunks, warm, time to first token went from 62.8 s to 58.9 s, and prefill from 523.6 to 557.9 tokens/s. Decode was unchanged, which is the expected null because single-token decode does not take this path. That measurement predates this patch and used an earlier build where an environment variable chose BM instead of the heuristic. The dispatch decision was the same one, but I have not re-run it against this branch.

Costs and things I left out. The extra instantiations grow the metallib from 164,587,560 to 171,001,000 bytes, or 3.9%. If that is too much, dropping the fp-mode instantiations and dispatching only the affine path would cut it. I also built and measured BM=16 with WM=1 and WN=2: it reached parity with BM=32 at runs of 8 and below and fell behind at 16, so nothing dispatches it.

All of the above comes from one machine, and the crossover between block sizes may well sit elsewhere on other devices. I am happy to re-run the sweep with different thresholds, or to narrow the condition, if you have a shape in mind.

python -m pytest python/tests/test_quantized.py passes: 33 tests, 3014 subtests.

[zcbenz]

Thanks for the PR, I can verify that this is bringing a very nice improvement.

Running gather_qmm_bench.py‎:

main branch	this branch
2.74436 msec	2.22729 msec

Can you please check whether existing tests actually cover all the kernel instantiations added in this PR?

[dwijenpatel]

Thank you for testing it out, and glad to see those perf improvements on your machine as well! Great question - I wish I asked myself that before submitting the PR.

Method: added temp logging to print every gather_qmm_rhs_nax kernel name at dispatch, ran the test suite, counted distinct names. These kernels are reached via mx.gather_qmm(..., transpose=True, sorted_indices=True). Block height comes from average rows per expert (BM=32 below 64, BM=64 at or above).

I've pushed an update (fd63970) addressing a coverage regression: on main, the fp-mode rows in test_gather_qmm_sorted (64 tokens over 4 experts) dispatched the BM=64 fp kernels, and the new block-height logic moves those shapes onto BM=32, which left the mxfp4/nvfp4/mxfp8 BM=64 kernels unexercised. The new test runs a 64-rows-per-expert regime that dispatches them again.

The same commit also expands coverage of the kernels this PR adds: existing tests dispatched 4 of the 63 reachable new BM=32 kernels (affine float16 gs64 4-bit, plus the three fp modes in bfloat16). test_gather_qmm_sorted_block_sizes runs every mode, group size and bit width at float16 and bfloat16 in three regimes: 16 rows per expert (BM=32), 64 (BM=64), and a 78-row batch (unaligned path). Suite coverage goes from 4 distinct gather_qmm_rhs_nax kernels to 84: all 42 reachable new BM=32 kernels and all 42 BM=64. Runtime ~10 s on a cold pipeline cache, ~1 s warm. On devices without nax support the same test exercises the fallback path.

To your original question, there are two gaps (pre-existing):

1. The float32 variants could never run under the harness: mlx_tests.py sets MLX_ENABLE_TF32=0, this path requires TF32 for float32, and the flag is cached on first read, so no in-process test can reach them. Pushed a fix (367c7eb): a small test that runs the float32 check in a subprocess with MLX_ENABLE_TF32=1. Verified by dispatch logging that it selects the float32 nax kernels at both block heights, matching gather_mm bit for bit.

2. The _nn instantiations (transpose=false) have no route to them: gather_qmm_rhs only takes the nax path when transpose is true, so both the existing BM=64 and new BM=32 _nn kernels are unreachable. I probed them locally by lifting the transpose gate: they compute correctly (bit-identical to gather_mm at fp16), and a quick timing on one shape family (E=4, K=D=512, base M5) shows parity with the current steel path at small batch and ~25% faster at B=1024 (0.268 ms vs 0.356). Two options, happy to do either: route transpose=false through nax (the vjp flips the transpose, so the backward pass of every quantized MoE fine-tune runs this shape today, currently on the steel path), or drop the _nn instantiations, which roughly halves this PR's 3.9% metallib growth.

[zcbenz]

Thanks!