Add a test

mlx/backend/metal/device_info.cpp

@@ -13,7 +13,12 @@ bool is_available() {
 }
 
 int device_count() {
-  return 1;
+  try {
+    metal::device(Device::gpu);
+    return 1;
+  } catch (...) {
+    return 0;
+  }
 }
 
 const std::unordered_map<std::string, std::variant<std::string, size_t>>&

mlx/backend/metal/kernels/fp_quantized.h

@@ -533,6 +533,7 @@ METAL_FUNC void fp_qvm_impl(
     device T* y,
     const int in_vec_size,
     const int out_vec_size,
+    const int in_vec_stride,
     uint3 tid [[threadgroup_position_in_grid]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
@@ -563,7 +564,7 @@ METAL_FUNC void fp_qvm_impl(
   int out_col = pack_factor * tn * (tid.y * num_simdgroups + simd_gid);
   ws += out_col * bytes_per_pack / pack_factor + simd_lid * out_vec_size_w;
   scales += out_col / group_size + simd_lid * out_vec_size_g;
-  x += tid.x * in_vec_size + simd_lid;
+  x += tid.x * in_vec_stride + simd_lid;
   y += tid.x * out_vec_size + out_col;
 
   if (out_col >= out_vec_size) {
@@ -1122,7 +1123,16 @@ template <typename T, const int group_size, int bits, bool batched>
         tid);
   }
   fp_qvm_impl<T, group_size, bits>(
-      w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid);
+      w,
+      scales,
+      x,
+      y,
+      in_vec_size,
+      out_vec_size,
+      in_vec_size,
+      tid,
+      simd_gid,
+      simd_lid);
 }
 
 template <typename T, const int group_size, int bits, int split_k = 32>
@@ -1164,8 +1174,20 @@ template <typename T, const int group_size, int bits, int split_k = 32>
   int in_vec_size_adj =
       tid.z % split_k == split_k - 1 ? final_block_size : in_vec_size;
 
+  // The in_vec_stride is the full K dimension, not the partition size
+  int in_vec_stride = (split_k - 1) * in_vec_size + final_block_size;
+
   fp_qvm_impl<T, group_size, bits>(
-      w, scales, x, y, in_vec_size_adj, out_vec_size, tid, simd_gid, simd_lid);
+      w,
+      scales,
+      x,
+      y,
+      in_vec_size_adj,
+      out_vec_size,
+      in_vec_stride,
+      tid,
+      simd_gid,
+      simd_lid);
 }
 
 template <
@@ -1423,7 +1445,16 @@ template <typename T, int group_size, int bits>
       s_strides,
       tid);
   fp_qvm_impl<T, group_size, bits>(
-      w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid);
+      w,
+      scales,
+      x,
+      y,
+      in_vec_size,
+      out_vec_size,
+      in_vec_size,
+      tid,
+      simd_gid,
+      simd_lid);
 }
 
 template <

mlx/backend/metal/kernels/indexing/scatter.h

@@ -80,7 +80,9 @@ template <
     uint3 gsize [[threads_per_grid]]) {
   Op op;
 
-  IdxT idx = IdxT(gid.z) * gsize.y + gid.y * gsize.x + gid.x * NWORK;
+  IdxT idx =
+      (IdxT(gid.z) * IdxT(gsize.y) + IdxT(gid.y)) * IdxT(gsize.x) * NWORK +
+      IdxT(gid.x) * NWORK;
   IdxT out_idx;
   IdxT update_idx;
 

mlx/backend/metal/kernels/quantized.h

@@ -983,6 +983,7 @@ METAL_FUNC void qvm_impl(
     device T* y,
     const int in_vec_size,
     const int out_vec_size,
+    const int in_vec_stride,
     uint3 tid [[threadgroup_position_in_grid]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
@@ -1016,7 +1017,7 @@ METAL_FUNC void qvm_impl(
   ws += out_col * bytes_per_pack / pack_factor + simd_lid * out_vec_size_w;
   scales += out_col / group_size + simd_lid * out_vec_size_g;
   biases += out_col / group_size + simd_lid * out_vec_size_g;
-  x += tid.x * in_vec_size + simd_lid;
+  x += tid.x * in_vec_stride + simd_lid;
   y += tid.x * out_vec_size + out_col;
 
   if (out_col >= out_vec_size) {
@@ -1643,6 +1644,7 @@ template <typename T, const int group_size, const int bits, bool batched>
       y,
       in_vec_size,
       out_vec_size,
+      in_vec_size,
       tid,
       simd_gid,
       simd_lid);
@@ -1691,6 +1693,9 @@ template <typename T, const int group_size, const int bits, int split_k = 32>
   int in_vec_size_adj =
       tid.z % split_k == split_k - 1 ? final_block_size : in_vec_size;
 
+  // The in_vec_stride is the full K dimension, not the partition size
+  int in_vec_stride = (split_k - 1) * in_vec_size + final_block_size;
+
   qvm_impl<T, group_size, bits>(
       w,
       scales,
@@ -1699,6 +1704,7 @@ template <typename T, const int group_size, const int bits, int split_k = 32>
       y,
       in_vec_size_adj,
       out_vec_size,
+      in_vec_stride,
       tid,
       simd_gid,
       simd_lid);
@@ -2077,6 +2083,7 @@ template <typename T, int group_size, int bits>
       y,
       in_vec_size,
       out_vec_size,
+      in_vec_size,
       tid,
       simd_gid,
       simd_lid);

python/tests/test_array.py

@@ -1550,6 +1550,30 @@ class TestArray(mlx_tests.MLXTestCase):
         a = a.at[:, :, :].add(update)
         self.assertEqualArray(a, update)
 
+    def test_slice_update_contiguous_2d(self):
+        for shape in [(32, 32), (64, 64), (17, 33), (128, 128)]:
+            for upd_shape in [(16, 16), (8, 8), (4, 4)]:
+                if upd_shape[0] > shape[0] or upd_shape[1] > shape[1]:
+                    continue
+                y = mx.zeros(shape)
+                x = mx.random.normal(upd_shape)
+                z = y.at[: upd_shape[0], : upd_shape[1]].add(x)
+                diff = z[: upd_shape[0], : upd_shape[1]] - x
+                self.assertTrue(mx.allclose(diff, mx.zeros_like(diff)))
+
+        # Test non-zero offset slice update
+        y = mx.zeros((32, 32))
+        x = mx.random.normal((16, 16))
+        z = y.at[16:, 16:].add(x)
+        self.assertTrue(mx.allclose(z[16:, 16:] - x, mx.zeros_like(x)))
+
+        # Test with size divisible by 4, 2, and odd
+        for cols in [32, 18, 15]:
+            y = mx.zeros((32, cols))
+            x = mx.random.normal((16, cols))
+            z = y.at[:16, :].add(x)
+            self.assertTrue(mx.allclose(z[:16, :] - x, mx.zeros_like(x)))
+
     def test_slice_negative_step(self):
         a_np = np.arange(20)
         a_mx = mx.array(a_np)

python/tests/test_quantized.py

@@ -470,6 +470,30 @@ class TestQuantized(mlx_tests.MLXTestCase):
         self.assertEqual(y_q.shape, y_hat.shape)
         self.assertLess((y_q - y_hat).abs().max(), 2e-3)
 
+    def test_qvm_splitk_multi_row(self):
+        # Test qvm split_k with M > 1 to ensure the x row stride is correct
+        key = mx.random.key(0)
+        k1, k2 = mx.random.split(key)
+        tests = product(
+            [64, 32],  # group_size
+            [4, 8],  # bits
+            [128],  # out dim (N)
+            [2048, 4096],  # in dim (K) >= 1024 to trigger split_k
+            [2, 3],  # M (multiple rows)
+        )
+        for group_size, bits, N, K, M in tests:
+            with self.subTest(M=M, K=K, N=N, group_size=group_size, bits=bits):
+                x = 1e-1 * mx.random.normal(shape=(M, K), key=k1)
+                w = 1e-1 * mx.random.normal(shape=(K, N), key=k2)
+                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(), 2e-3)
+
     def test_fp_qvm(self):
         key = mx.random.key(0)
         k1, k2 = mx.random.split(key)