From 669080ab2a6efc3995cfffda7feace92c6ddda41 Mon Sep 17 00:00:00 2001 From: Philip John Basile Date: Sat, 8 Aug 2026 21:58:58 -0400 Subject: [PATCH 1/2] Fix sorted gather_qmm NAX row overflow above 32K Keep the remaining M extent in int until after it is clamped to the SIMD-group tile size. This prevents ragged sorted-RHS workloads above the signed-short boundary from leaving output rows unwritten. Add a focused regression over the exact 32767/32768/32769 seam and an aligned control above it, using a dense fp32 oracle and allocator-poisoned outputs. Co-authored-by: OpenAI Codex --- ACKNOWLEDGMENTS.md | 1 + mlx/backend/metal/kernels/quantized_nax.h | 3 +- python/tests/test_quantized.py | 43 +++++++++++++++++++++++ 3 files changed, 45 insertions(+), 2 deletions(-) diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md index 7005396e71..50905e1f97 100644 --- a/ACKNOWLEDGMENTS.md +++ b/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. diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index ed32eb59a7..bef2634736 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/mlx/backend/metal/kernels/quantized_nax.h @@ -1543,8 +1543,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))))); diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 461175f013..3d62018b96 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1599,6 +1599,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 From 902054649703736913d9f78c47b403f2a4183156 Mon Sep 17 00:00:00 2001 From: Cheng Date: Thu, 27 Aug 2026 08:02:21 +0900 Subject: [PATCH 2/2] nit --- ACKNOWLEDGMENTS.md | 1 - 1 file changed, 1 deletion(-) diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md index 50905e1f97..7005396e71 100644 --- a/ACKNOWLEDGMENTS.md +++ b/ACKNOWLEDGMENTS.md @@ -7,7 +7,6 @@ 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.