Fix sorted gather_qmm NAX row overflow above 32K

mlx/backend/metal/kernels/quantized_nax.h

@@ -1543,8 +1543,7 @@ template <
   const short tm = SM * (simd_group_id / WN);
   const short tn = SN * (simd_group_id % WN);
 
-  const short sgp_sm =
-      align_M ? SM : min(SM, short(max(0, (M - (y_row + tm)))));
+  const short sgp_sm = align_M ? SM : min(int(SM), max(0, M - (y_row + tm)));
   const short sgp_sn =
       align_N ? SN : min(SN, short(max(0, (N - (y_col + tn)))));
 

python/tests/test_quantized.py

@@ -1599,6 +1599,49 @@ def scatter_unsort(x, inv_order, shape=None):
                 self.assertTrue(mx.allclose(y1, y3, atol=tol))
                 self.assertTrue(mx.allclose(y1, y4, atol=tol))
 
+    @unittest.skipIf(not mx.metal.is_available(), "requires Metal")
+    def test_gather_qmm_sorted_nax_large_m(self):
+        E, N, K, group_size = 16, 64, 64, 32
+        dtype = mx.float16
+        mx.random.seed(0)
+        w = (mx.random.normal((E, N, K)) * 0.1).astype(dtype)
+        w_q, scales, biases = mx.quantize(w, group_size=group_size, bits=4)
+        w_hat = mx.dequantize(w_q, scales, biases, group_size=group_size, bits=4)
+
+        for M in (32767, 32768, 32769, 32832):
+            with self.subTest(M=M):
+                x = (mx.random.normal((M, 1, K)) * 0.1).astype(dtype)
+                rhs_indices = (mx.arange(M) * E // M).astype(mx.uint32)
+                y_hat = mx.gather_mm(
+                    x.astype(mx.float32),
+                    mx.swapaxes(w_hat, -1, -2).astype(mx.float32),
+                    rhs_indices=rhs_indices,
+                    sorted_indices=True,
+                )
+                mx.eval(y_hat)
+                mx.synchronize()
+
+                for value in (-31.0, 47.0):
+                    poison = mx.full(y_hat.shape, value, dtype=mx.float16)
+                    mx.eval(poison)
+                    mx.synchronize()
+                    del poison
+
+                    y_q = mx.gather_qmm(
+                        x,
+                        w_q,
+                        scales,
+                        biases,
+                        rhs_indices=rhs_indices,
+                        transpose=True,
+                        group_size=group_size,
+                        bits=4,
+                        sorted_indices=True,
+                    )
+                    max_error = (y_q.astype(mx.float32) - y_hat).abs().max()
+                    self.assertLess(float(max_error.item()), 5e-2)
+                    del y_q, max_error
+
     @unittest.skipIf(mx.cuda.is_available(), "Not implemented for CUDA")
     def test_gather_qmm_sorted_sliced_weight(self):
         E, R, D, N = 8, 64, 256, 64