diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 78d8cf0df5..7c461bc1b5 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1476,6 +1476,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; + if (biases_) { + biases = ensure_row_contiguous(*biases_, d, s); + } // TODO: Tune the block sizes int bm = 64, bn = 64, bk = 64; @@ -1554,9 +1558,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++); @@ -1627,6 +1630,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; + if (biases_) { + biases = ensure_row_contiguous(*biases_, d, s); + } // TODO: Tune the block sizes int bm = 16, bn = 32, bk = 32; @@ -1704,9 +1711,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++); diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 63254ee9c7..0b756b39d5 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1403,6 +1403,34 @@ 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): + 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).astype(dtype) + qw, s, b = mx.quantize(w, group_size=64, bits=4) + 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)): + 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: