Fix gather_qmm NAX kernel name mismatch

[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!