| |
| |
| import gc |
| import weakref |
|
|
| import mlx.core as mx |
| import numpy as np |
|
|
| from mlx_ternary import kernel, pack_trits, tern_matmul |
|
|
|
|
| def _weights(n=64, k=256, seed=0): |
| rng = np.random.default_rng(seed) |
| trits = rng.integers(-1, 2, size=(n, k), dtype=np.int8) |
| packed, _ = pack_trits(trits) |
| scales = rng.uniform(0.02, 0.08, size=(n, k // 64)).astype(np.float32) |
| return mx.array(packed), mx.array(scales) |
|
|
|
|
| def test_cache_bounded_across_many_batches(): |
| kernel._COMPILED.clear() |
| codes, scales = _weights() |
| for b in range(1, 300): |
| mx.eval(tern_matmul(mx.ones((b, 256)), codes, scales, 256)) |
| assert len(kernel._COMPILED) <= kernel.COMPILE_MAX_BATCH |
| mx.synchronize() |
| mx.clear_cache() |
| base = mx.get_active_memory() |
| for b in range(300, 600): |
| mx.eval(tern_matmul(mx.ones((b, 256)), codes, scales, 256)) |
| mx.synchronize() |
| mx.clear_cache() |
| assert mx.get_active_memory() <= base + (1 << 20) |
| assert len(kernel._COMPILED) <= kernel.COMPILE_CACHE_SIZE |
|
|
|
|
| def test_cache_does_not_keep_weights_alive(): |
| kernel._COMPILED.clear() |
| for seed in range(3): |
| codes, scales = _weights(seed=seed) |
| ref = weakref.ref(codes) |
| for b in (1, 4, 64): |
| mx.eval(tern_matmul(mx.ones((b, 256)), codes, scales, 256)) |
| del codes, scales |
| gc.collect() |
| assert ref() is None |
| assert len(kernel._COMPILED) <= 2 |
|
|
