Fix and enable non-transposed NAX qmm

[gordofreemo]

What this does

qmm() only lets NAX through when transpose == true, so affine_qmm_n_nax is compiled but never actually run and is currently unreachable. It has two separate problems, both of which this PR fixes, and then enables the path to be reachable.

FIX: Weight addressing uses the transposed layout

qmm_n_nax_tgp_impl uses the same QuantizedBlockLoader instantiation as the
generic qmm_n_impl in kernels/quantized.h:

QuantizedBlockLoader<T, BK, BN, BN_padded, 0, WM * WN * SIMD_SIZE, group_size, bits>

but set the pointers up for w = [N, K] rather than [K, N]:

	qmm_n_impl (correct)	qmm_n_nax_tgp_impl (before)
weight offset	wl += y_col * bytes_per_pack / pack_factor	wl += y_col * K_w
scales offset	scales += y_col / group_size	scales += y_col * K_g
biases offset	biases += y_col / group_size	biases += y_col * K_g
leading dim	loader_w(..., N, ...)	loader_w(..., K, ...)
implied layout	w = [K, N], grouped along N	w = [N, K], grouped along K

This fix makes those four lines match qmm_n_impl. affine_gather_qmm_n_nax shares the impl, so this fixes that too.

FIX: No partial M tile handling

This same function had (void)M;, its bounds line commented out, and called Atile.load / Dtile.store unconditionally. qmm_nax() sizes the grid with (M + bm - 1) / bm, so partial M tiles are dispatched and any M % 64 != 0
read x and wrote y out of bounds.

Fixed by porting what qmm_t_nax_tgp_impl already does: sgp_sm plus a dispatch_bool picking load_safe/store_safe. That split is compile time, so the aligned path keeps the original unguarded loads and costs nothing, and only one branch is instantiated per call so the x += BK walk isn't duplicated. sgp_sm is min(int(SM), M - (y_row + tm)) rather than truncating the distance to short, matching the int16 fix already made for the transposed kernel.

FIX: The dispatch guard

if (metal::is_nax_available() && (transpose || (N % 64 == 0)) &&
    (K % 64 == 0) && ...

The non-transposed weight tile is BK x BN over w = [K, N] and the loader isn't clamped in N, so it also needs N % 64 == 0. Affine quantization with group_size >= 64 gives that structurally, since the groups run along N. This only rejects group_size == 32 with N not a multiple of 64, which falls back to the generic kernel like before.

Scope: this patches qmm() only. gather_qmm() stays gated on transpose.

Tests

test_qmm_non_transposed in python/tests/test_quantized.py:

• transformer sized K/N: (2048, 2048), (512, 2048), (2048, 512), (11008, 2048)
• M aligned and unaligned: 1, 2, 31, 32, 33, 63, 64, 65, 96, 97, 100, 127, 128, 129, 250, 256. The 33 to 63 range is where a whole simdgroup of the M tile lands past the end of the matrix
• batched x, 3D and 4D, with unaligned M
• M = 33000 with a partial M tile, for the int16 distance case
• group_size 64 and 128, bits 4 and 8

Tolerances and the 1/sqrt(K) operand scaling follow the existing test_qmm. Worth noting the existing test_qmm already covers transpose=False at M in {8, 32, 33, 64} with N, K in {128, 256}.

Validation

I don't have M5 hardware. All NAX validation was on an iPhone 17 Pro (A19 Pro), iOS 26.5.2, through the Swift bindings, against a dequantized reference: 52 shape/M/batch cases covering every unaligned M above plus batched. Before the fix, relative error ran 1.5 to 7826, with an output norm of 427,429 against an expected 24,009 at K=N=2048, M=500. After, it's bit exact against the reference on that device.

On the M3 I have, is_nax_available() is false, so the full test_quantized.py suite passes there only through the generic path. That confirms no regression on the fallback, not the NAX path. The kernel itself does compile (quantized_nax.air builds clean).

Performance

Same iPhone 17, but this time running LoRA fine tuning on SmolLM3-3B 4-bit (group_size 64, r=8, batch 1).

One full training run with 405 examples and 1215 iterations. It took 2.9 hours before the fix and 1.5 after the fix (~1.93x speedup). Final loss 0.8829 with the fix against 0.8813 without, a 0.18% difference, and the two loss traces correlate at 0.9999 across 1210 steps.

A paired per op benchmark on fixed synthetic shapes puts the backward phase at 1.65 to 2.13x and the whole iteration around 1.55x.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[zcbenz]

dispatch_bool is useful when you have to call some function that requires a template argument, in this case you should just do normal if:

if (sgp_sm == SM) {
  load()
} else {
  load_safe()
}

[gordofreemo]

Done, thanks, switched to a plain if (sgp_sm == SM) for both the load and the store.

[zcbenz]

Thanks!