Fix NAX tile bugs and make the gather_qmm block tile tunable

[vvsotnikov]

Proposed changes

A mixture of experts layer calls the NAX kernel for gather_qmm during
prefill, and that kernel is built for a single block tile: 64,64,64,2,2. The
block tile has a large effect on how fast the kernel runs, but you cannot
change it, because two bugs stop you as soon as you try.

This pull request lets you tune the block tile, and corrects both bugs.

Why the block tile matters

Each threadgroup handles a range of rows, and the rows arrive sorted by
expert. A threadgroup that starts in one expert and ends in the next must
repeat its whole K loop for the second expert, so a block tile with fewer rows
produces fewer threadgroups that cross a boundary.

We measured this on a mixture of experts prefill shape, where M is 20480, K is
3072, N is 2048, and there are 256 experts. The weights use affine
quantization with 2 bits and a group size of 128. The machine is an M5 Max.

Rows in the block tile	Speed
128	0.73x
64 (the default, 8.85 ms)	1.00x
32 (8.13 ms)	1.09x
16	0.83x

A block tile with 32 rows takes the full layer from 13.17 ms to 12.12 ms,
which is 8 percent, and makes end to end prefill 2 to 5 percent faster.

The best block tile differs between models, and it also depends on how many
rows each expert receives, so this pull request does not change the default.
It builds the kernel for 6 block tiles, and the environment variable
MLX_QMM_TILE_NAX selects one of them in the format "BM,BN,BK,WM,WN".
quantized.cpp reads the variable with env::get_var at dispatch time, in
the same way that scaled_dot_product_attention.cpp reads MLX_SDPA_BLOCKS.
If the value is not one of the 6 block tiles, MLX raises an error that lists
the ones you can use.

The kernel for hardware without NAX behaves the same way, and it already
carries a TODO: Tune the block sizes comment. A rule that picks the block
tile from the shape would help both kernels, and this pull request gives you
the tool to measure such a rule.

The first bug: the block loader accepts one shape

The block loader dequantizes the weights into threadgroup memory, but it
accepted only a block tile no wider than one quantization group. A second copy
of the loader handled the special case of a group size of 32, and that copy
contained this line:

static_assert(..., "Other configurations are not yet supported");

Every new block tile stops there first.

The first commit merges all the cases into one loader and deletes the copy. A
thread can now read less than one quantization group, part of one group, or
several whole groups. The block tiles that MLX builds today still take the
path they took before.

The second bug: the matmul does nothing

tile_matmad_nax used the M paired path when TN == 1 && TM % 2 == 0, and
the N paired path when TN % 2 == 0, but it had no other case. A block tile
with an odd TN matched neither, so the body of the function was empty: the
matmul never ran, and the output tile kept the zeros that clear() wrote. The
kernel still compiled and still ran, and it looked much faster than it was.

You can see this without building MLX. Put this file in the root of the
repository:

// probe.metal
#include "mlx/backend/metal/kernels/utils.h"
#include "mlx/backend/metal/kernels/steel/gemm/nax.h"
using namespace mlx::steel;

kernel void probe(device float* out [[buffer(0)]]) {
  NAXTile<float, 1, 1> C;          // TM is 1, TN is 1
  NAXTile<bfloat16_t, 1, 2> A;     // TK is 2
  NAXTile<bfloat16_t, 2, 1> B;
  C.clear(); A.clear(); B.clear();
  tile_matmad_nax(C, A, metal::bool_constant<false>{},
                     B, metal::bool_constant<false>{});
  out[0] = C.val_frags[0][0];
}

Then count the matmul instructions:

xcrun metal -c probe.metal -I . -S -o out.air
grep -c "matmul\|mma" out.air

On main this prints 0, and with the second commit it prints 41.

The same commit corrects the descriptor of the M paired fragment matmul. That
overload takes 2 fragments along M and 1 fragment along N, so the problem is
32x16x16, but the code built the descriptor as
matmul2d_descriptor(16, 32, 16), which is the shape of the N paired
overload. The copy into the left input tensor then ran past the end of that
tensor.

The block tile that MLX selects today reaches neither bug, but any other block
tile reaches both.

Tests

test_gather_qmm_rhs_nax_tiles in python/tests/test_quantized.py calls
gather_qmm with all 6 block tiles, at a group size of 32 and of 64, and with
a row count that divides the rows of the block tile as well as one that does
not. It compares every result against dequantize followed by gather_mm,
and it checks that an unknown block tile raises an error. The test skips on
hardware without matrix coprocessors.

Two of the 6 block tiles give each simdgroup a single fragment along N, so
they take the paths that the second commit corrects, and the group size of 32
makes one row of a block tile span several quantization groups, which is the
case the first commit adds.

These are the results on an M5 Max with macOS 26.5:

• test_quantized.py: 34 tests, all pass.
• test_blas.py: 28 tests, all pass.
• If you remove the second commit, the new test fails on exactly the block
  tiles that give one fragment along N, because the kernel returns zeros.
• If you remove the first commit, quantized_nax.metal does not compile.

The gather_qmm docstring documents MLX_QMM_TILE_NAX.

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)

[erwinzhang7]

Heads up that this overlaps with #4009, and the consolidation carries one bug forward. Flagging it so whoever reviews can sequence the two rather than have git discover it.

Merging the two loader copies is a good change on its own, and the tile_matmad_nax finding is a nice catch. The probe that counts matmul instructions is a neat way to show it.

The guard survives the merge

The surviving load_safe keeps this:

if (reduction_dim == 1 && bi >= src_tile_dim.x) {

bi indexes BROWS, and src_tile_dim.y is the row extent while src_tile_dim.x is the column extent. For the transpose case BROWS = BN, and the K tail passes tile_w = short2(k_remain, tgp_bn), so this compares a row index against the K remainder. With k_remain = 32 it zeroes tile rows 32 to 63, which are output columns, instead of masking the columns past the K edge.

Those rows are exactly one of the two N direction simdgroups (SN = BN / WN = 32), so half the output columns silently lose their tail contribution. Measured on an M5 Max, the wrong columns are precisely those with col % 64 >= 32.

This diff deletes that guard from the group_size = 32 specialization along with the rest of that struct, but no added line changes it in the loader that survives, so the behaviour is preserved.

Reachability differs

Worth separating, since it affects how urgently each needs to land:

• The two bugs fixed here are, as the description says, unreachable with the block tile MLX selects today. They gate tuning.
• The guard above is reachable at the current default tile. Any gather_qmm with sorted_indices=True and K % 64 != 0 hits it, which in a quantized MoE is a normal shape. On affine 4 and 8 bit and mxfp4 it corrupts roughly 95 percent of output elements, and on nvfp4 and mxfp8 with unaligned K it returns values around 1e38, consistent with uninitialized memory.

So they are complementary rather than competing. #4009 fixes the guard plus a K tail bound in affine_gather_qmm_rhs_nax, and #4010 fixes an int16 overflow in the row bound of the same kernel. sgp_sm does not appear anywhere in this diff, so #4010 should be independent of it.

Sequencing

#4009 patches load_safe in both loader copies, and one of those copies is deleted here, so whichever lands second needs a rebase. #4009's rebase onto this is small: the two hunks collapse into one against the merged loader.

Happy to rebase #4009 on top of this branch if that ordering is easier for review, or to fold the one line change into this PR if the maintainers would rather it arrive as a single unit. Either is fine by me, I just did not want the guard to quietly survive the consolidation.

[zcbenz]

This PR does too much things and it would take tremendous time to verify the correctness and check performance regression, and I don't enough confidence in the implementation.

[vvsotnikov]

Hello @zcbenz ! Thank you for the feedback.

Originally, I was only planning to make tile sizes tunable, but finding and fixing these two bugs made the PR scope much wider that it originally was.

If you're open to it, I can split this into 2-3 PRs so it is easier to test and review

[zcbenz]

Performant kernels are usually designed for a fixed tile size, extending it for more settings usually makes it much harder to verify correctness and check for regressions. Since the kernel was not designed for arbitrary tile size I would oppose making it tunable.