diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index afbd14bf..b0369790 100644 --- a/python/tests/test_quantized.py +++ b/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)