Re-disable qmm_n_nax and fix group_size < 64

mlx/backend/metal/kernels/fp_quantized_nax.h

@@ -178,7 +178,7 @@ struct QuantizedBlockLoader {
     if (reduction_dim == 1) {
       scales += n_groups;
     } else {
-      scales += n_groups * group_stride;
+      scales += group_stride;
     }
   }
 };

mlx/backend/metal/kernels/quantized_nax.h

@@ -826,8 +826,8 @@ struct QuantizedBlockLoader<
       biases += n_groups;
       // }
     } else {
-      scales += n_groups * group_stride;
-      biases += n_groups * group_stride;
+      scales += group_stride;
+      biases += group_stride;
     }
   }
 };

mlx/backend/metal/quantized.cpp

@@ -1035,9 +1035,10 @@ void qmm(
     metal::Device& d,
     const Stream& s,
     const std::string& mode) {
-  // 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)) {
+  bool has_nax_kernel =
+      metal::is_nax_available() && (transpose || mode == "affine");
+  if (has_nax_kernel && transpose && (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

@@ -1,5 +1,6 @@
 # Copyright © 2023-2026 Apple Inc.
 
+import os
 import platform
 import subprocess
 import unittest
@@ -353,52 +354,91 @@ 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)
 
+    @unittest.skipIf("CI" in os.environ, "too slow in CI")
     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)
+        modes = ["mxfp4", "nvfp4", "mxfp8"]
+        if mx.default_device() == mx.gpu:
+            dtypes = [mx.float16, mx.bfloat16]
+        else:
+            dtypes = [mx.float32]
+
+        def check_affine(M, K, N, group_size, bits, dtype, batch=()):
+            x = mx.random.normal(shape=(*batch, M, K), key=k1, dtype=dtype) / K**0.5
+            w = mx.random.normal(shape=(K, N), key=k2, dtype=dtype) / K**0.5
             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)
+            tol = 1e-3 if dtype == mx.float32 else 1.5e-3
+            self.assertLess((y_q - y_hat).abs().max(), tol)
+
+        def check_fp(M, K, N, mode, dtype, batch=()):
+            x = mx.random.normal(shape=(*batch, M, K), key=k1, dtype=dtype) / K**0.5
+            w = mx.random.normal(shape=(K, N), key=k2, dtype=dtype) / K**0.5
+            w_q, scales = mx.quantize(w, mode=mode)
+            w_hat = mx.dequantize(w_q, scales, mode=mode)
+            y_q = mx.quantized_matmul(x, w_q, scales, None, False, mode=mode)
+            y_hat = x @ w_hat
+            self.assertEqual(y_q.shape, y_hat.shape)
+            tol = 1e-3 if dtype == mx.float32 else 1.5e-3
             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)
+        for dtype in dtypes:
+            # 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, dtype=dtype
+                    ):
+                        check_affine(M, 512, 1024, group_size, bits, dtype)
+                for mode in modes:
+                    with self.subTest(M=M, mode=mode, dtype=dtype):
+                        check_fp(M, 512, 1024, mode, dtype)
+
+            # 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), dtype=dtype):
+                        check_affine(M, K, N, 64, 4, dtype)
+                for mode in modes:
+                    with self.subTest(shape=(M, K, N), mode=mode, dtype=dtype):
+                        check_fp(M, 512, 1024, mode, dtype)
+
+            # Batched x, unaligned M.
+            for batch in [(2,), (2, 3)]:
+                for M in [33, 250]:
+                    with self.subTest(batch=batch, M=M, dtype=dtype):
+                        check_affine(M, 512, 1024, 64, 4, dtype, batch=batch)
+                for mode in modes:
+                    with self.subTest(batch=batch, mode=mode, dtype=dtype):
+                        check_fp(M, 512, 1024, mode, dtype, 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), dtype=dtype):
+                check_affine(33000, 128, 64, 64, 4, dtype)
+                check_fp(33000, 128, 64, mode, dtype)
+
+            # K=64 is the single reduction-tile control; K > 64 spans two or more
+            # tiles, which exposed the over-advanced scale pointer.
+            for M in [8, 33, 65]:
+                for K in [64, 128, 256]:
+                    for bits in [2, 4, 8]:
+                        with self.subTest(M=M, K=K, bits=bits, dtype=dtype):
+                            check_affine(M, K, 128, 32, bits, dtype)
+                    for mode in modes:
+                        with self.subTest(M=M, K=K, mode=mode, dtype=dtype):
+                            check_fp(M, K, 128, mode, dtype)
 
     def test_qmm_vjp(self):
         key = mx.random.key(0)