[scyyh11]

Proposed changes

gather_qmm_nax builds the NAX kernel name with bk = 32, but the gather qmm NAX kernels are only instantiated with _bk64_ names. On NAX devices this makes the matrix-size gather_qmm path (taken for half-precision inputs, or float32 with tf32 enabled) fail during kernel lookup with an Unable to load kernel ..._bk32... error.

This changes the gather_qmm_nax kernel-name construction to use bk = 64, matching the existing instantiations. In this path bk is only part of the kernel name; the dispatch geometry is unchanged.

The PR also adds a regression test for matrix-size half-precision gather_qmm covering both affine and mxfp4, checked against a float32 reference.

Tested with:

DEVICE=gpu .venv/bin/pytest -q python/tests/test_quantized.py::TestQuantized::test_gather_qmm_matrix_path

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]

Nice fix, thanks!

[programVeins]

Confirming this on an M5 Pro (macOS 27 beta, Metal NAX active) - we hit the
exact same bug independently and this patch fixes it.

Symptom: any transposed gather_qmm that takes the matrix-size path
(M >= vector_limit) on a NAX device fails at kernel load:

[metal::Device] Unable to load kernel
affine_gather_qmm_t_nax_bfloat16_t_gs_64_b_4_bm64_bn64_bk32_...

because the dispatcher builds the name with bk=32 while only the bk64
kernels are instantiated in quantized_nax.metal.

Where it bites: quantized MoE / switch layers. Single-token AR decode
(M=1) routes through gather_qmv and is unaffected - which is probably why
it slipped through - but anything that pushes many tokens through experts at
once trips it. For us it was a 4-bit MoE (DiffusionGemma 26B-A4B, 128 experts
top-8): a 256-token diffusion canvas hits gather_qmm at matrix size on every
denoising step, so every step died until this fix.

Rebuilt MLX from source with bk=32 → bk=64 applied → the kernel loads, the
MoE path runs, and 4-bit expert generation works end to end. The change is
name-only here (tile geometry is unchanged), so it's purely a lookup fix.

+1 to merging, thanks!