diff --git a/mlx/backend/metal/kernels/fp_quantized_nax.h b/mlx/backend/metal/kernels/fp_quantized_nax.h index 7c452e6a64..cf64ff7f46 100644 --- a/mlx/backend/metal/kernels/fp_quantized_nax.h +++ b/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; } } }; diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index 81f6a7a252..31e51a5b7e 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/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; } } }; diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index ab3e1deacc..f659e16c93 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/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, diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 56ba75db32..15bc892bd8 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1,5 +1,6 @@ # Copyright © 2023-2026 Apple Inc. +import os import platform import subprocess import unittest @@ -353,6 +354,7 @@ 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 @@ -360,45 +362,83 @@ def test_qmm_non_transposed(self): # 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)