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?