Add test

mlx/backend/metal/kernels/fp_quantized.h

@@ -533,7 +533,6 @@ 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]]) {
@@ -564,7 +563,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_stride + simd_lid;
+  x += tid.x * in_vec_size + simd_lid;
   y += tid.x * out_vec_size + out_col;
 
   if (out_col >= out_vec_size) {
@@ -1123,16 +1122,7 @@ 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,
-      in_vec_size,
-      tid,
-      simd_gid,
-      simd_lid);
+      w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid);
 }
 
 template <typename T, const int group_size, int bits, int split_k = 32>
@@ -1174,20 +1164,8 @@ 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,
-      in_vec_stride,
-      tid,
-      simd_gid,
-      simd_lid);
+      w, scales, x, y, in_vec_size_adj, out_vec_size, tid, simd_gid, simd_lid);
 }
 
 template <
@@ -1445,16 +1423,7 @@ 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,
-      in_vec_size,
-      tid,
-      simd_gid,
-      simd_lid);
+      w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid);
 }
 
 template <

mlx/backend/metal/kernels/quantized.h

@@ -983,7 +983,6 @@ 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]]) {
@@ -1017,7 +1016,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_stride + simd_lid;
+  x += tid.x * in_vec_size + simd_lid;
   y += tid.x * out_vec_size + out_col;
 
   if (out_col >= out_vec_size) {
@@ -1644,7 +1643,6 @@ 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);
@@ -1693,9 +1691,6 @@ 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,
@@ -1704,7 +1699,6 @@ 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);
@@ -2083,7 +2077,6 @@ template <typename T, int group_size, int bits>
       y,
       in_vec_size,
       out_vec_size,
-      in_vec_size,
       tid,
       simd_gid,
       simd_lid);

mlx/backend/metal/rope.cpp

@@ -133,8 +133,9 @@ void RoPE::eval_gpu(
   if (single) {
     compute_encoder.set_bytes(out_strides, 1, 4);
     uint32_t dim0 = dims_ / 2;
-    group_dims = get_block_dims(dim0, N, 1);
-    grid_dims = MTL::Size(dim0, N, 1);
+    uint32_t dim1 = B * N;
+    group_dims = get_block_dims(dim0, dim1, 1);
+    grid_dims = MTL::Size(dim0, dim1, 1);
   } else {
     compute_encoder.set_bytes(strides, 3, 4);
     compute_encoder.set_bytes(out_strides, 3, 5);

python/tests/test_fast.py

@@ -365,6 +365,28 @@ class TestFast(mlx_tests.MLXTestCase):
         rx = rope_orig(x, dims, traditional, base, scale, offset)
         self.assertLess(mx.abs(rx - rx_fast).max(), 1e-5)
 
+    def test_rope_single_batch(self):
+        base = 10000.0
+        scale = 1.0
+        offset = 5
+
+        for traditional in [True, False]:
+            for B in [2, 4, 8]:
+                for n_head in [1, 4, 7]:
+                    for dims in [64, 128]:
+                        x = mx.random.uniform(shape=(B, n_head, 1, dims))
+                        mx.eval(x)
+                        rx_fast = mx.fast.rope(
+                            x,
+                            dims,
+                            traditional=traditional,
+                            base=base,
+                            scale=scale,
+                            offset=offset,
+                        )
+                        rx = rope_orig(x, dims, traditional, base, scale, offset)
+                        self.assertLess(mx.abs(rx - rx_fast).max(), 1e-5)
+
     def test_rope_with_large_offset(self):
         x = mx.random.normal(shape=(1, 1, 1024, 32))
         rx_fp32 = mx.fast.rope(

python/tests/test_quantized.py

@@ -470,30 +470,6 @@ 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)