From 688f4de67e943dc5933149264640a59d00ba7994 Mon Sep 17 00:00:00 2001 From: Duhyeon Date: Sat, 8 Aug 2026 00:47:26 +0900 Subject: [PATCH 1/5] Use the full row-contiguity check for GatherQMM's quantized inputs --- mlx/backend/metal/quantized.cpp | 6 +++--- python/tests/test_quantized.py | 26 ++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 94c563070a..22033414da 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1563,11 +1563,11 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { out.set_data(allocator::malloc(out.nbytes())); array x = ensure_row_contiguous_matrix(inputs[0], d, s); - array w = ensure_row_contiguous_matrix(inputs[1], d, s); - array scales = ensure_row_contiguous_matrix(inputs[2], d, s); + array w = ensure_row_contiguous(inputs[1], d, s); + array scales = ensure_row_contiguous(inputs[2], d, s); std::optional biases = std::nullopt; if (inputs.size() == 6) { - biases = ensure_row_contiguous_matrix(inputs[3], d, s); + biases = ensure_row_contiguous(inputs[3], d, s); } const array& lhs_indices = inputs[inputs.size() - 2]; const array& rhs_indices = inputs[inputs.size() - 1]; diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 7fbbe6938e..74f61a0186 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1294,6 +1294,32 @@ def scatter_unsort(x, inv_order, shape=None): self.assertTrue(mx.allclose(y1, y3, atol=tol)) self.assertTrue(mx.allclose(y1, y4, atol=tol)) + def test_gather_qmm_sorted_sliced_weight(self): + E, R, D, N = 8, 64, 256, 64 + mx.random.seed(0) + w = mx.random.normal((E, 2 * R, D)) * 0.05 + qw, s, b = mx.quantize(w, group_size=64, bits=4) + x = mx.random.normal((N, 1, D)) * 0.5 + indices = mx.sort(mx.random.randint(0, E, (N,)).astype(mx.uint32)) + + for sl in (slice(0, R), slice(R, 2 * R)): + view = (qw[:, sl], s[:, sl], b[:, sl]) + copy = tuple(mx.contiguous(a) for a in view) + kwargs = dict( + rhs_indices=indices, + transpose=True, + group_size=64, + bits=4, + sorted_indices=True, + ) + self.assertTrue( + mx.allclose( + mx.gather_qmm(x, *view, **kwargs), + mx.gather_qmm(x, *copy, **kwargs), + atol=1e-4, + ) + ) + def test_gather_qmm_grad(self): def gather_qmm_ref(x, w, s, b, lhs, rhs, trans, sort): if lhs is not None: From 53fb4c985f259698231ea925570f9da9d5447b9a Mon Sep 17 00:00:00 2001 From: Duhyeon Date: Sat, 8 Aug 2026 14:12:30 +0900 Subject: [PATCH 2/5] Use fp16 on GPU in the sliced-weight test for CUDA support --- python/tests/test_quantized.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index a96cb3ba8f..09cb940cd7 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1405,10 +1405,11 @@ def scatter_unsort(x, inv_order, shape=None): def test_gather_qmm_sorted_sliced_weight(self): E, R, D, N = 8, 64, 256, 64 + dtype = mx.float16 if (mx.default_device() == mx.gpu) else mx.float32 mx.random.seed(0) - w = mx.random.normal((E, 2 * R, D)) * 0.05 + w = (mx.random.normal((E, 2 * R, D)) * 0.05).astype(dtype) qw, s, b = mx.quantize(w, group_size=64, bits=4) - x = mx.random.normal((N, 1, D)) * 0.5 + x = (mx.random.normal((N, 1, D)) * 0.5).astype(dtype) indices = mx.sort(mx.random.randint(0, E, (N,)).astype(mx.uint32)) for sl in (slice(0, R), slice(R, 2 * R)): From e07bb0e978b52b07865b9982dcde335aa32e3f84 Mon Sep 17 00:00:00 2001 From: Duhyeon Date: Sat, 8 Aug 2026 17:10:44 +0900 Subject: [PATCH 3/5] Move the fix into gather_qmm_rhs: normalize biases before encoding --- mlx/backend/metal/quantized.cpp | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 16c6d7efc3..295af1d0a7 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1460,6 +1460,10 @@ void gather_qmm_rhs_nax( array x = broadcast_with_indices(x_); array w = ensure_row_contiguous(w_, d, s); array scales = ensure_row_contiguous(scales_, d, s); + std::optional biases = std::nullopt; + if (biases_) { + biases = ensure_row_contiguous(*biases_, d, s); + } // TODO: Tune the block sizes int bm = 64, bn = 64, bk = 64; @@ -1538,9 +1542,8 @@ void gather_qmm_rhs_nax( compute_encoder.set_input_array(x, c++); compute_encoder.set_input_array(w, c++); compute_encoder.set_input_array(scales, c++); - if (biases_) { - array biases = ensure_row_contiguous(*biases_, d, s); - compute_encoder.set_input_array(biases, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); } compute_encoder.set_input_array(indices, c++); compute_encoder.set_output_array(out, c++); @@ -1611,6 +1614,10 @@ void gather_qmm_rhs( array x = broadcast_with_indices(x_); array w = ensure_row_contiguous(w_, d, s); array scales = ensure_row_contiguous(scales_, d, s); + std::optional biases = std::nullopt; + if (biases_) { + biases = ensure_row_contiguous(*biases_, d, s); + } // TODO: Tune the block sizes int bm = 16, bn = 32, bk = 32; @@ -1688,9 +1695,8 @@ void gather_qmm_rhs( compute_encoder.set_input_array(x, c++); compute_encoder.set_input_array(w, c++); compute_encoder.set_input_array(scales, c++); - if (biases_) { - array biases = ensure_row_contiguous(*biases_, d, s); - compute_encoder.set_input_array(biases, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); } compute_encoder.set_input_array(indices, c++); compute_encoder.set_output_array(out, c++); @@ -1846,11 +1852,11 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { out.set_data(allocator::malloc(out.nbytes())); array x = ensure_row_contiguous_matrix(inputs[0], d, s); - array w = ensure_row_contiguous(inputs[1], d, s); - array scales = ensure_row_contiguous(inputs[2], d, s); + array w = ensure_row_contiguous_matrix(inputs[1], d, s); + array scales = ensure_row_contiguous_matrix(inputs[2], d, s); std::optional biases = std::nullopt; if (inputs.size() == 6) { - biases = ensure_row_contiguous(inputs[3], d, s); + biases = ensure_row_contiguous_matrix(inputs[3], d, s); } const array& lhs_indices = inputs[inputs.size() - 2]; const array& rhs_indices = inputs[inputs.size() - 1]; From c53bfd2c77104a04e384c3790645284f215a9c42 Mon Sep 17 00:00:00 2001 From: Duhyeon Date: Sat, 8 Aug 2026 17:54:20 +0900 Subject: [PATCH 4/5] Skip the sliced-weight test on CUDA --- 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 09cb940cd7..2e3c430bad 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1404,6 +1404,8 @@ def scatter_unsort(x, inv_order, shape=None): self.assertTrue(mx.allclose(y1, y4, atol=tol)) def test_gather_qmm_sorted_sliced_weight(self): + if mx.default_device() == mx.gpu and not mx.metal.is_available(): + self.skipTest("Not implemented for CUDA") E, R, D, N = 8, 64, 256, 64 dtype = mx.float16 if (mx.default_device() == mx.gpu) else mx.float32 mx.random.seed(0) From 55a1eab3294bd744f9362313fd8fb87ce818cd9b Mon Sep 17 00:00:00 2001 From: Cheng Date: Sat, 8 Aug 2026 19:04:34 +0900 Subject: [PATCH 5/5] nit --- mlx/backend/metal/quantized.cpp | 4 ++-- python/tests/test_quantized.py | 3 +-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 509df75385..7c461bc1b5 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1476,7 +1476,7 @@ void gather_qmm_rhs_nax( array x = broadcast_with_indices(x_); array w = ensure_row_contiguous(w_, d, s); array scales = ensure_row_contiguous(scales_, d, s); - std::optional biases = std::nullopt; + std::optional biases; if (biases_) { biases = ensure_row_contiguous(*biases_, d, s); } @@ -1630,7 +1630,7 @@ void gather_qmm_rhs( array x = broadcast_with_indices(x_); array w = ensure_row_contiguous(w_, d, s); array scales = ensure_row_contiguous(scales_, d, s); - std::optional biases = std::nullopt; + std::optional biases; if (biases_) { biases = ensure_row_contiguous(*biases_, d, s); } diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 2e3c430bad..0b756b39d5 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1403,9 +1403,8 @@ 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(mx.cuda.is_available(), "Not implemented for CUDA") def test_gather_qmm_sorted_sliced_weight(self): - if mx.default_device() == mx.gpu and not mx.metal.is_available(): - self.skipTest("Not implemented for CUDA") E, R, D, N = 8, 64, 256, 64 dtype = mx.float16 if (mx.default_device() == mx.gpu) else mx.float32 mx.random.seed(0)