Re-disable qmm_n_nax and fix group_size < 64

[PhilipJohnBasile]

Fixes #4201. Root-cause fix (supersedes the earlier workaround in this PR).

Root cause

The group_size=32 QuantizedBlockLoader specialization
(quantized_nax.h) advanced the scale/bias pointer by
n_groups * group_stride when stepping between reduction tiles. For the
non-transposed NAX kernel (#4051) the scale plane is [K, N / group_size],
so a 64-row reduction tile is group_stride = 64 * N / group_size scale
entries — advancing by n_groups * that (4N for gs=32) reads out of bounds
from the second K-tile on. Result: quantized_matmul(transpose=False) with
group_size=32 produced garbage (and NaN under some inputs) on Metal for any
K > 64. K = 64 (a single reduction tile) was correct — which is why the
bug was localized only after comparing K=64 vs K=128.

group_size 64/128 and transpose=True were unaffected (they use the
generic loader / the reduction_dim == 1 branch respectively).

The same over-advance existed in the fp (fp8/MX) loader and is fixed here too.

Fix

Advance by group_stride (not n_groups * group_stride) in the
reduction_dim == 0 branch of both loaders.

Tests / verification

• New regression test_qmm_non_transposed_group_size_lt_64: gs=32, bits
  2/4/8, M 8/33/65, K 64 (single-tile control) / 128 / 256 (multi-tile).
  Fails on main at K>64 (garbage/NaN), passes with the fix.
• Full test_quantized on a local Metal build (M5 Max, macOS 27, Xcode 27,
  mlx from source): previously 48 failed subtests -> now 37 passed,
  3106 subtests passed.
• Manual sweep matches fp64 reference (~1e-2) for all bits/gs, both
  transpose flags, and odd-M partial tiles.

[zcbenz]

Thanks for the fix, I should have ran the full test suite when testing the PR 🤦

So apparently the whole qmm_n is still WIP and I'm re-disabling it until a full review and benchmarking.

[gordofreemo]

Flagging a footgun this left behind, since #4051 was mine.

After this PR the guard is:

bool has_nax_kernel =
    metal::is_nax_available() && (transpose || mode == "affine");
if (has_nax_kernel && transpose && (K % 64 == 0) && ...

has_nax_kernel && transpose makes the (transpose || mode == "affine") disjunct unreachable, so it's really just is_nax_available() && transpose. What worries me is that the N % 64 == 0 term is gone while the variable now reads as "a NAX kernel exists for this config" so the natural way to re-enable later is to drop the && transpose, and that hands you the non-transposed path with no N constraint.

It needs one. qmm_n_nax_tgp_impl loads a 64×64 weight tile via load_unsafe(), and the store is clamped to the simdgroup tile rather than the matrix Dtile.store_safe(y + tm * N + tn, N, short2(SN, sgp_sm)) passes lim_y = SN = 32. So for N % 64 != 0 the last N-tile reads w/scales past the row end and writes y out of bounds. N is a multiple of group_size for any validly quantized w, so this only breaks at group_size == 32.

Folding the precondition into the variable makes the name true and keeps it from re-arming, without changing behaviour:

bool has_nax_kernel = metal::is_nax_available() &&
    (transpose || (mode == "affine" && N % 64 == 0));
if (has_nax_kernel && transpose && (K % 64 == 0) && ...

Unrelated but adjacent: fp_qmm_n_impl still has the [N, K] weight addressing and the missing M-tile bounds that #4051 fixed on the affine side. This is unreachable today for the same reason the affine one was but worth knowing before that guard widens. I can open a separate issue.