File size: 1,469 Bytes
a9bf912 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | # Copyright 2026 Jakevin
# SPDX-License-Identifier: Apache-2.0
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
|