From 500cda8e8baa572e4eaa5bc02aeebae64ed59b4a Mon Sep 17 00:00:00 2001 From: Philip John Basile Date: Wed, 12 Aug 2026 01:32:19 -0400 Subject: [PATCH 1/3] Fix non-transposed NAX quantized matmul at group_size < 64 The group_size=32 QuantizedBlockLoader specialization advanced the scale (and bias) pointer by n_groups * group_stride when stepping between reduction tiles. For the non-transposed kernel the scale plane is [K, N / group_size], so a 64-row reduction tile is group_stride = 64 * N / group_size scale entries, not n_groups times that. The doubled step read out of bounds from the second K-tile on, producing garbage (and NaN under some inputs) for group_size=32 quantized_matmul(transpose=False) on Metal whenever K > 64. Also fixes the same pattern in the fp (fp8/MX) loader. Adds a regression over the gs=32 non-transposed path covering a single and multiple 64-row K-tiles. Fixes #4201. --- mlx/backend/metal/kernels/fp_quantized_nax.h | 5 ++- mlx/backend/metal/kernels/quantized_nax.h | 10 ++++-- python/tests/test_quantized.py | 33 ++++++++++++++++++++ 3 files changed, 45 insertions(+), 3 deletions(-) diff --git a/mlx/backend/metal/kernels/fp_quantized_nax.h b/mlx/backend/metal/kernels/fp_quantized_nax.h index 7c452e6a64..eb747d79c2 100644 --- a/mlx/backend/metal/kernels/fp_quantized_nax.h +++ b/mlx/backend/metal/kernels/fp_quantized_nax.h @@ -178,7 +178,10 @@ struct QuantizedBlockLoader { if (reduction_dim == 1) { scales += n_groups; } else { - scales += n_groups * group_stride; + // Advance from one BROWS-row reduction tile to the next in a + // [K, N / group_size] scale plane; see the identical fix in + // quantized_nax.h (QuantizedBlockLoader<..., 32, bits>). + scales += group_stride; } } }; diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index 81f6a7a252..a519ccdf5a 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/mlx/backend/metal/kernels/quantized_nax.h @@ -826,8 +826,14 @@ struct QuantizedBlockLoader< biases += n_groups; // } } else { - scales += n_groups * group_stride; - biases += n_groups * group_stride; + // Advance from one BROWS-row reduction tile to the next. The scale + // plane is [K, N / group_size], so the step is BROWS rows of N/gs + // entries = group_stride. Advancing by n_groups * group_stride skips + // past n_groups weight tiles of scales, reading out of bounds from the + // second reduction tile on (visible at group_size=32 where one 64-wide + // tile contains two 32-value groups). + scales += group_stride; + biases += group_stride; } } }; diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 56ba75db32..96e20196c3 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -400,6 +400,39 @@ def check(M, K, N, group_size, bits, batch=()): with self.subTest(shape=(33000, 128, 64)): check(33000, 128, 64, 64, 4) + def test_qmm_non_transposed_group_size_lt_64(self): + # Regression for the non-transposed NAX loader with group_size < 64. + # A 64-wide weight tile holds multiple 32-value groups, and the gs=32 + # loader advanced the scale pointer by n_groups * group_stride per + # reduction tile, reading out of bounds from the second K-tile on + # (garbage, and NaN under some inputs, at K > 64). + 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): + x = mx.random.normal(shape=(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) + + # K=64 is the single reduction-tile control; K > 64 spans two or more + # tiles, which exposed the over-advanced scale pointer. + for bits in [2, 4, 8]: + for M in [8, 33, 65]: + for K in [64, 128, 256]: + with self.subTest(M=M, K=K, bits=bits): + check(M, K, 128, 32, bits) + def test_qmm_vjp(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) From ec02465422d29dec8420e7e6269d838df867605e Mon Sep 17 00:00:00 2001 From: Cheng Date: Wed, 12 Aug 2026 00:01:32 -0700 Subject: [PATCH 2/3] Re-disable qmm_n_nax --- mlx/backend/metal/kernels/fp_quantized_nax.h | 3 - mlx/backend/metal/kernels/quantized_nax.h | 6 - mlx/backend/metal/quantized.cpp | 7 +- python/tests/test_quantized.py | 121 ++++++++++--------- 4 files changed, 67 insertions(+), 70 deletions(-) diff --git a/mlx/backend/metal/kernels/fp_quantized_nax.h b/mlx/backend/metal/kernels/fp_quantized_nax.h index eb747d79c2..cf64ff7f46 100644 --- a/mlx/backend/metal/kernels/fp_quantized_nax.h +++ b/mlx/backend/metal/kernels/fp_quantized_nax.h @@ -178,9 +178,6 @@ struct QuantizedBlockLoader { if (reduction_dim == 1) { scales += n_groups; } else { - // Advance from one BROWS-row reduction tile to the next in a - // [K, N / group_size] scale plane; see the identical fix in - // quantized_nax.h (QuantizedBlockLoader<..., 32, bits>). scales += group_stride; } } diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index a519ccdf5a..31e51a5b7e 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/mlx/backend/metal/kernels/quantized_nax.h @@ -826,12 +826,6 @@ struct QuantizedBlockLoader< biases += n_groups; // } } else { - // Advance from one BROWS-row reduction tile to the next. The scale - // plane is [K, N / group_size], so the step is BROWS rows of N/gs - // entries = group_stride. Advancing by n_groups * group_stride skips - // past n_groups weight tiles of scales, reading out of bounds from the - // second reduction tile on (visible at group_size=32 where one 64-wide - // tile contains two 32-value groups). 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 96e20196c3..8727df786e 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -360,78 +360,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) - # 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_non_transposed_group_size_lt_64(self): - # Regression for the non-transposed NAX loader with group_size < 64. - # A 64-wide weight tile holds multiple 32-value groups, and the gs=32 - # loader advanced the scale pointer by n_groups * group_stride per - # reduction tile, reading out of bounds from the second K-tile on - # (garbage, and NaN under some inputs, at K > 64). - 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): - x = mx.random.normal(shape=(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 - ) + 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) - # K=64 is the single reduction-tile control; K > 64 spans two or more - # tiles, which exposed the over-advanced scale pointer. - for bits in [2, 4, 8]: + 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]: - with self.subTest(M=M, K=K, bits=bits): - check(M, K, 128, 32, bits) + 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) From ca42d4c7f15a5a4a3b34dce741b537fe890a769f Mon Sep 17 00:00:00 2001 From: Cheng Date: Wed, 12 Aug 2026 00:25:37 -0700 Subject: [PATCH 3/3] Disable test_qmm_non_transposed in CI --- python/tests/test_quantized.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 8727df786e..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