Add a test

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)