Fix non transposed affine qmm dispatch logic

[RohanGautam]

The path for using the non-transposed NAX qmm kernel was unreachable in the current implementation (because of the condition if (has_nax_kernel && transpose && (K % 64 == 0) ... - lets only the transposed case through).

We needed to add some more conditions:

• having N be a multiple of 64 (the non-transposed nax kernel does not have an unaligned-N variant)
• Adding a transposed check for the block size since a non-transposed kernel with bm=32 is not instantiated. This check would not be needed if we instantiate it (adding to the metallib), but I didn't want to assume that that was an option. We can add it and simplify this condition if that's an option.

benchmarks

The benchmarks are run with mlx_lm linked to the different mlx versions. The non-transposed path is exercised for large prefills (when chunked), for the quantized kv cache computations. We use mlx_lm.generate (not benchmark, to use the kv cache quantisation options) . The command run is mlx_lm.generate --model <model> -p - --ignore-chat-template --max-tokens 8 --temp 0 --kv-bits 4 --quantized-kv-start 0 < <file withprompt of desired size>.txt

model	machine	prompt	prefill before	prefill after	Δ prefill
Qwen3-8B-4bit	M5 Pro	4096	1,424.7	1,425.6	+0.06%
		6144	1,231.3	1,290.4	+4.80%
		8192	1,092.4	1,182.4	+8.24%
		12288	907.2	1,028.1	+13.33%
	M5 Max	4096	2,791.3	2,791.1	−0.01%
		6144	2,233.7	2,277.5	+1.96%
		8192	1,969.8	2,111.2	+7.18%
		12288	1,647.2	1,819.6	+10.47%
Qwen3-4B-Instruct-2507-4bit	M5 Pro	4096	2,127.9	2,127.4	−0.02%
		6144	1,716.5	1,834.7	+6.89%
		8192	1,466.5	1,630.9	+11.21%
		12288	1,146.6	1,344.2	+17.24%
Qwen3-30B-A3B-4bit	M5 Max	4096	3,384.4	3,399.7	+0.45%
		6144	2,721.2	2,906.6	+6.81%
		8192	2,289.7	2,539.8	+10.92%
		12288	1,690.6	1,967.7	+16.39%

The 4096 one has no changes because of mlx-lm using N-1 tokens instead of N tokens. There is only one quantised kv cache to attend over, and this is the wrong size for the non-transposed qmm to kick in (it does not have a ragged-K variant). The last chunk for all others also go through the normal path, just the middle ones don't and are sped up with this PR. This boundary effect can be seen below, and is perhaps something to address in mlx-lm down the line.

model	machine	prompt	prefill before	prefill after	Δ prefill
Qwen3-8B-4bit	M5 Pro	4096	1,428.4	1,428.6	+0.01%
		4097	1,452.0	1,580.3	+8.83%
		6144	1,224.2	1,280.4	+4.60%
		6145	1,246.0	1,405.9	+12.83%
	M5 Max	4096	2,789.5	2,791.3	+0.07%
		4097	2,749.4	2,894.9	+5.29%
		6144	2,101.4	2,212.7	+5.30%
		6145	2,237.0	2,478.7	+10.80%

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: AI was used to flag candidate optimisations, which I then understood, ran benchmarks for, and simplified.

[nastya236]

Looks good to me! We probably have the same issue in gather_qmm..

Regarding mlx-lm, ml-explore/mlx-lm#1781 proposed to add a prefill size as an argument, so it is more controllable now!