Fix and enable non-transposed NAX qmm

mlx/backend/metal/kernels/quantized_nax.h

@@ -1094,7 +1094,6 @@ METAL_FUNC void qmm_n_nax_tgp_impl(
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
   (void)lid;
-  (void)M;
 
   static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE");
   static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE");
@@ -1115,23 +1114,20 @@ METAL_FUNC void qmm_n_nax_tgp_impl(
       bits>;
 
   // Set the block
-  const int K_w = K * bytes_per_pack / pack_factor;
-  const int K_g = K / group_size;
   const int y_row = tid.y * BM;
   const int y_col = tid.x * BN;
 
   auto wl = (const device uint8_t*)w;
 
+  // Here w is [K, N]: packed and group-quantized along N, with row stride N.
   x += y_row * static_cast<int64_t>(K);
-  wl += y_col * K_w;
-  scales += y_col * K_g;
-  biases += y_col * K_g;
+  wl += y_col * bytes_per_pack / pack_factor;
+  scales += y_col / group_size;
+  biases += y_col / group_size;
   y += y_row * static_cast<int64_t>(N) + y_col;
 
-  // Make the x loader and mma operation
-  // const short num_els = min(BM, M - y_row);
-  // const short num_outs = min(BN, N - y_col);
-  loader_w_t loader_w(wl, scales, biases, K, Ws, simd_gid, simd_lid);
+  // Make the weight loader
+  loader_w_t loader_w(wl, scales, biases, N, Ws, simd_gid, simd_lid);
 
   constexpr short SM = BM / WM;
   constexpr short SN = BN / WN;
@@ -1144,6 +1140,8 @@ METAL_FUNC void qmm_n_nax_tgp_impl(
   const short tm = SM * (simd_gid / WN);
   const short tn = SN * (simd_gid % WN);
 
+  const short sgp_sm = min(int(SM), M - (y_row + tm));
+
   const short ldb_tgp = BN_padded;
 
   constexpr bool transpose_a = false;
@@ -1168,7 +1166,12 @@ METAL_FUNC void qmm_n_nax_tgp_impl(
 
       volatile int compiler_barrier;
 
-      Atile.load(x + kk1, K);
+      if (sgp_sm == SM) {
+        Atile.load(x + kk1, K);
+      } else {
+        Atile.load_safe(x + kk1, K, short2(SK, sgp_sm));
+      }
+
       Btile.template load<T, BN_padded, 1>(Ws + tn + kk1 * ldb_tgp);
 
       tile_matmad_nax(
@@ -1188,7 +1191,11 @@ METAL_FUNC void qmm_n_nax_tgp_impl(
   // Store results to device memory
   threadgroup_barrier(mem_flags::mem_threadgroup);
 
-  Dtile.store(y + tm * N + tn, N);
+  if (sgp_sm == SM) {
+    Dtile.store(y + tm * N + tn, N);
+  } else {
+    Dtile.store_safe(y + tm * N + tn, N, short2(SN, sgp_sm));
+  }
 }
 
 template <
@@ -1678,4 +1685,4 @@ template <
       });
     });
   }
-}
+}

mlx/backend/metal/quantized.cpp

@@ -1011,8 +1011,9 @@ void qmm(
     metal::Device& d,
     const Stream& s,
     const std::string& mode) {
-  if (metal::is_nax_available() && transpose && (K % 64 == 0) &&
-      (env::enable_tf32() || x.dtype() != float32)) {
+  // The non-transposed kernel requires N % 64 == 0.
+  if (metal::is_nax_available() && (transpose || (N % 64 == 0)) &&
+      (K % 64 == 0) && (env::enable_tf32() || x.dtype() != float32)) {
     return qmm_nax(
         /* const array& x = */ x,
         /* const array& w = */ w,

python/tests/test_quantized.py

@@ -353,6 +353,53 @@ def test_qmm_large_dims(self):
                 tol = 1e-3 if dtype == mx.float32 else 1.5e-3
                 self.assertLess((y_q - y_hat).abs().max(), tol)
 
+    def test_qmm_non_transposed(self):
+        # The non-transposed matmul (w is [K, N]) is reachable mainly from the
+        # vjp of a quantized linear layer, so it gets much less coverage than
+        # the transposed one. Sweep it over transformer-sized K/N and over M
+        # values that leave a partial M-tile.
+        key = mx.random.key(0)
+        k1, k2 = mx.random.split(key)
+        dtype = mx.float16 if (mx.default_device() == mx.gpu) else mx.float32
+        tol = 1e-3 if dtype == mx.float32 else 1.5e-3
+
+        def check(M, K, N, group_size, bits, batch=()):
+            x = mx.random.normal(shape=(*batch, M, K), key=k1) / K**0.5
+            w = mx.random.normal(shape=(K, N), key=k2) / K**0.5
+            x = x.astype(dtype)
+            w = w.astype(dtype)
+            w_q, scales, biases = mx.quantize(w, group_size, bits)
+            w_hat = mx.dequantize(w_q, scales, biases, group_size, bits)
+            y_q = mx.quantized_matmul(x, w_q, scales, biases, False, group_size, bits)
+            y_hat = x @ w_hat
+            self.assertEqual(y_q.shape, y_hat.shape)
+            self.assertLess((y_q - y_hat).abs().max(), tol)
+
+        # M sweep. 33..63 is the interesting range: a whole simdgroup of the
+        # threadgroup's M-tile falls past the end of the matrix.
+        for M in [1, 2, 31, 32, 33, 63, 64, 65, 96, 97, 100, 127, 128, 129]:
+            for group_size, bits in [(64, 4), (128, 4), (64, 8)]:
+                with self.subTest(M=M, group_size=group_size, bits=bits):
+                    check(M, 512, 1024, group_size, bits)
+
+        # Transformer-sized K/N, aligned and unaligned M.
+        for K, N in [(2048, 2048), (512, 2048), (2048, 512), (11008, 2048)]:
+            for M in [100, 256]:
+                with self.subTest(shape=(M, K, N)):
+                    check(M, K, N, 64, 4)
+
+        # Batched x, unaligned M.
+        for batch in [(2,), (2, 3)]:
+            for M in [33, 250]:
+                with self.subTest(batch=batch, M=M):
+                    check(M, 512, 1024, 64, 4, batch=batch)
+
+        # M > 2**15 with a partial M-tile, so the per-simdgroup row count is a
+        # distance that does not fit in an int16. Same failure mode as the one
+        # test_qmm_large_dims covers for the transposed kernel.
+        with self.subTest(shape=(33000, 128, 64)):
+            check(33000, 128, 64, 64, 4)
+
     def test_qmm_vjp(self):
         key = mx.random.key(0)
         k1, k2 = mx.random.split(key)