Fix sorted gather_qmm NAX row overflow above 32K

ACKNOWLEDGMENTS.md

@@ -7,6 +7,7 @@ with a short description of your contribution(s) below. For example:
 
 MLX was developed with contributions from the following individuals:
 
+- Philip John Basile: Fixed sorted `gather_qmm` NAX row-bound overflow above 32K.
 - Nripesh Niketan: Added `softsign`, `softmax`, `hardswish`, `logsoftmax` activation functions. Added `dropout3d` ops. Added `LogicalAnd` and `LogicalOR` ops. Added `clip_grad_norm` along with `tree_reduce`. Added `cross`. Added `orthogonal` initializer.
 - Juarez Bochi: Fixed bug in cross attention.
 - Justin Deschenaux: Sine, Cosine, arange, randint, truncated normal, bernoulli, lion optimizer, Dropout2d, linear and logistic regression python example.

mlx/backend/metal/kernels/quantized_nax.h

@@ -1529,8 +1529,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)))));
 
@@ -1678,4 +1677,4 @@ template <
       });
     });
   }
-}
+}

python/tests/test_quantized.py

@@ -1440,6 +1440,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