Use a 32-row block in qmm_t_nax when one block covers all of M

[dwijenpatel]

A small follow-up to #3925 and #4023, which tuned the block size for
gather_qmm_rhs_nax: the same question applied to the plain quantized
matmul. qmm_t_nax always runs 64-row blocks; for M <= 32 we measured
modest gains from a 32-row block.

Only wide matrices reach this kernel at small M: below the qmv batch
limit the vector kernel serves the product, and up to 256 column tiles
the split-K path does, so N > 8192 is what arrives here with M <= 32.
In practice that is vocabulary heads and wide MLP projections during
speculative-decode verification, batch serving, and short prompts.

Unlike the gather case, a shorter block here is not free: every row
block reads all the weight columns it touches, so BM=32 pays only while
one block still covers all of M (forcing it at M=128 measures 13%
slower). The dispatch shrinks the block exactly when it stays a single
block.

Measured on a base M5 (10-core GPU), fp16, 4-bit gs64, against merged
main (5ec30ac), both arms built from one tree, alternating in one
session, fresh process per cell, three passes with the last recorded
(drift on the repeated first cell: 0.7%):

M	K	N	main	this branch
14	5120	13824	1011.0 us	833.1 us	1.21x
32	5120	13824	980.7 us	804.2 us	1.22x
48	5120	13824	same path in both builds

BM=16 was also measured and dropped: 6% over BM=32 at M=14 on the one
shape where it showed, not worth its instantiations.

Two notes for review:

1. The fp_quantized_nax.metal instantiation macros did not forward
   their tile-size arguments into the template, so every fp tile landed
   on the defaults; a new tile would have compiled to a mis-labeled
   64-row kernel. They now forward bm/bk/bn/wm/wn. No behavior change
   for existing instantiations, which all pass the defaults.

2. The new instantiations grow the metallib by 4.02%; half of that is
   the batch_1 and alN_false variants, which I can drop behind a
   dispatch guard if size matters more.

Tests: the existing test_qmm_large_dims shapes (16 and 33 rows at
N=32840) already land on either side of the new dispatch, so one added
row pins the boundary itself (M=32). A small new test covers the two
variants no existing shape reaches, batched and fp-mode; the fp row is
what exercises the macro fix.

Repro:

python -c "
import time
import mlx.core as mx
M, K, N = 32, 5120, 13824
x = (mx.random.normal((M, K)) / K**0.5).astype(mx.float16)
w = (mx.random.normal((N, K)) / K**0.5).astype(mx.float16)
wq = mx.quantize(w, group_size=64, bits=4)
def f(): mx.eval(mx.quantized_matmul(x, *wq, transpose=True, group_size=64, bits=4))
for _ in range(30): f()
t0 = time.perf_counter()
for _ in range(300): f()
print(f'{(time.perf_counter()-t0)/300*1e6:.1f} us per matmul')"

[pierre427]

Validated on M5 Max (128 GB, macOS 26.5.2), since CI has no NAX hardware. Built this branch (46c9b85) and its effective baseline (main 36fde27 — the two *_nax.metal kernels and quantized.cpp are unchanged from your base 5ec30ac) with MACOSX_DEPLOYMENT_TARGET=26.2, chained dependent transposed 4-bit gs64 ops (bf16 activations), idle GPU, warmup + 6 reps.

Speedup (base us/op divided by PR us/op):

K x N	M=8	14	16	24	32	33	48
5120 x 13824	0.97	1.44	1.60	1.54	1.45	0.97	1.00
6656 x 19968	1.01	1.35	1.44	1.37	1.32	0.99	0.99
6656 x 202048	1.00	1.40	1.52	1.46	1.33	1.00	1.00

Bigger wins here than the 1.21-1.22x in the PR description (that was a base M5) — M=14..32 lands 1.32-1.60x on M5 Max across all three shapes. M>=33 is unchanged as intended (the 0.97 at 5120x13824/M=33 is run noise; absolute base/PR there is 214.8 vs 221.4 us). M=8 is flat, which matches expectation — at M=8 the 32-row block still wastes most of its rows.

python -m unittest test_quantized on this branch: 36/36 pass on M5 Max, including the new M=32 boundary case.

Related: #4198 reports that the split-K heuristic keeps M<=32 transposed qmm at N<=8192 off the NAX path entirely; if that band is routed back to qmm, this PR's small-M block would apply there too — the two compose.

[dwijenpatel]

Thank you @pierre427 for testing on an M5 Max! Your numbers cover a gap I can't test myself. I also appreciate you confirming the kernels were unchanged between 5ec30ac and 36fde27 before comparing.

On M=8: I do not think that band reaches this kernel. At K=5120 and N=13824, get_qmv_batch_limit returns 13, so M=8 dispatches qmv rather than qmm. I measured the boundary on a base M5 (this branch): M=12 takes 1189 us, M=13 takes 879 us. Below 13 the block size cannot matter, so a flat result is expected.

On #4198: I agree that the two compose. At N <= 8192 the split-K routing keeps small-M traffic off qmm. This PR does not reach that band today. If the routing changes, the 32-row block applies there as well.