(self)
| 433 | self.assertLess((y_q - y_hat).abs().max(), 1e-3) |
| 434 | |
| 435 | def test_qvm_splitk(self): |
| 436 | key = mx.random.key(0) |
| 437 | k1, k2 = mx.random.split(key) |
| 438 | tests = product( |
| 439 | [128, 64, 32], # group_size |
| 440 | [2, 4, 8], # bits |
| 441 | [128], # M |
| 442 | [16384], # N |
| 443 | [1, 3], # B |
| 444 | ) |
| 445 | for group_size, bits, M, N, B in tests: |
| 446 | with self.subTest(shape=(B, M, N), group_size=group_size, bits=bits): |
| 447 | x_shape = (1, N) if B == 0 else (B, 1, N) |
| 448 | w_shape = (N, M) if B == 0 else (B, N, M) |
| 449 | x = 1e-1 * mx.random.normal(shape=x_shape, key=k1) |
| 450 | w = 1e-1 * mx.random.normal(shape=w_shape, key=k2) |
| 451 | w_q, scales, biases = mx.quantize(w, group_size, bits) |
| 452 | w_hat = mx.dequantize(w_q, scales, biases, group_size, bits) |
| 453 | y_q = mx.quantized_matmul( |
| 454 | x, w_q, scales, biases, False, group_size, bits |
| 455 | ) |
| 456 | y_hat = x @ w_hat |
| 457 | self.assertEqual(y_q.shape, y_hat.shape) |
| 458 | self.assertLess((y_q - y_hat).abs().max(), 2e-3) |
| 459 | |
| 460 | # Test with 1D vector |
| 461 | group_size = 32 |
| 462 | bits = 8 |
| 463 | N = 2048 |
| 464 | x = 1e-1 * mx.random.normal(shape=(N,), key=k1) |
| 465 | w = 1e-1 * mx.random.normal(shape=(N, N), key=k2) |
| 466 | w_q, scales, biases = mx.quantize(w, group_size, bits) |
| 467 | w_hat = mx.dequantize(w_q, scales, biases, group_size, bits) |
| 468 | y_q = mx.quantized_matmul(x, w_q, scales, biases, False, group_size, bits) |
| 469 | y_hat = x @ w_hat |
| 470 | self.assertEqual(y_q.shape, y_hat.shape) |
| 471 | self.assertLess((y_q - y_hat).abs().max(), 2e-3) |
| 472 | |
| 473 | def test_fp_qvm(self): |
| 474 | key = mx.random.key(0) |
nothing calls this directly
no test coverage detected