Normalize biases before encoding in gather_qmm_rhs

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<array> 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<array> 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++);

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: