File size: 4,647 Bytes
6154ddb d3c3658 | 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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | # Copyright 2026 Jakevin
# SPDX-License-Identifier: Apache-2.0
import mlx.core as mx
import numpy as np
import pytest
from mlx_ternary import dequantize, pack_trits, tern_matmul
# Shapes probed with the v3 kernel. The published maximum was 1.44e-6.
CASES = [
(8, 64, 1),
(16, 2048, 1),
(16, 2048, 128),
(40, 128, 3),
(64, 256, 64),
(64, 2048, 32),
(128, 192, 32),
(64, 6144, 32),
(41, 512, 1),
(41, 512, 7),
(640, 2560, 1),
(640, 2560, 128),
(2560, 640, 1),
(2560, 640, 128),
(2048, 6144, 1),
]
def _rel(y, ref) -> float:
err = mx.abs(y - ref)
mx.eval(err)
denom = float(mx.max(mx.abs(ref)).item())
return float(mx.max(err).item()) / max(denom, 1e-8)
def _layer(n, k, trits, scales):
packed, _row = pack_trits(trits)
return mx.array(packed), mx.array(np.asarray(scales, dtype=np.float32)), mx.array(dequantize(trits, scales))
def test_exact_ones_cover_qmv_and_gemm():
# Tail GEMV (K=64), 512-wide GEMV, 640 = 512+128, and the 32x32 GEMM.
specs = [
(2, 64, 1, 0.5, 32.0),
(1, 512, 1, 0.25, 128.0),
(4, 640, 1, 1.0, 640.0),
(32, 64, 32, 0.5, 32.0),
(32, 128, 32, 0.25, 32.0),
]
for n, k, batch, scale, expect in specs:
trits = np.ones((n, k), dtype=np.int8)
scales = np.full((n, k // 64), scale, dtype=np.float32)
codes, scale_mx, _ref = _layer(n, k, trits, scales)
x = mx.ones((batch, k), dtype=mx.float32)
y = tern_matmul(x, codes, scale_mx, k)
mx.eval(y)
assert y.shape == (batch, n)
assert float(mx.max(mx.abs(y - expect)).item()) == pytest.approx(0.0, abs=1e-4)
def test_negative_and_zero_trits():
n, k = 2, 64
trits = np.zeros((n, k), dtype=np.int8)
trits[0] = -1
scales = np.ones((n, 1), dtype=np.float32)
codes, scale_mx, _ref = _layer(n, k, trits, scales)
x = mx.ones((1, k), dtype=mx.float32)
y = tern_matmul(x, codes, scale_mx, k)
mx.eval(y)
assert float(y[0, 0].item()) == pytest.approx(-64.0, abs=1e-4)
assert float(y[0, 1].item()) == pytest.approx(0.0, abs=1e-4)
def test_vector_input_is_one_row():
trits = np.ones((3, 64), dtype=np.int8)
scales = np.ones((3, 1), dtype=np.float32)
codes, scale_mx, _ref = _layer(3, 64, trits, scales)
y = tern_matmul(mx.ones((64,), dtype=mx.float32), codes, scale_mx, 64)
mx.eval(y)
assert y.shape == (1, 3)
def test_random_shapes_under_1e_3():
rng = np.random.default_rng(7)
worst = 0.0
for n, k, batch in CASES:
trits = rng.integers(-1, 2, size=(n, k), dtype=np.int8)
scales = rng.uniform(0.02, 0.2, size=(n, k // 64)).astype(np.float32)
codes, scale_mx, ref_w = _layer(n, k, trits, scales)
x = mx.array(rng.standard_normal((batch, k)).astype(np.float32))
y = tern_matmul(x, codes, scale_mx, k)
ref = x @ ref_w.T
mx.eval(y, ref)
rel = _rel(y, ref)
worst = max(worst, rel)
assert np.isfinite(rel) and rel < 1e-3, f"{n}x{k} B={batch} rel={rel:.3e}"
del codes, scale_mx, ref_w, x, y, ref
mx.clear_cache()
assert worst < 1e-3
def test_rejects_bad_shapes():
trits = np.ones((2, 64), dtype=np.int8)
codes, scale_mx, _ref = _layer(2, 64, trits, np.ones((2, 1), np.float32))
with pytest.raises(ValueError):
tern_matmul(mx.ones((2, 3, 64)), codes, scale_mx, 64)
with pytest.raises(ValueError):
tern_matmul(mx.ones((1, 32)), codes, scale_mx, 32)
def _case(n, k, b, seed=7):
rng = np.random.default_rng(seed)
trits = rng.integers(-1, 2, size=(n, k), dtype=np.int8)
scales = rng.uniform(0.05, 0.3, size=(n, k // 64)).astype(np.float32)
packed, _ = pack_trits(trits)
x = rng.standard_normal((b, k)).astype(np.float32)
return x, packed, scales, x @ dequantize(trits, scales).T
@pytest.mark.parametrize("b", [1, 32])
def test_non_contiguous_x(b):
x, packed, scales, ref = _case(64, 512, b)
xt = mx.array(np.ascontiguousarray(x.T)).T # transposed view
y = np.array(tern_matmul(xt, mx.array(packed), mx.array(scales), 512))
assert np.abs(y - ref).max() / np.abs(ref).max() < 1e-3
def test_shape_mismatch_raises():
with pytest.raises(ValueError):
tern_matmul(mx.zeros((1, 128)), mx.zeros((4, 16), dtype=mx.uint8), mx.zeros((4, 2)), 128)
with pytest.raises(ValueError):
tern_matmul(mx.zeros((1, 128)), mx.zeros((4, 32), dtype=mx.uint8), mx.zeros((4, 3)), 128)
def test_empty_batch():
y = tern_matmul(mx.zeros((0, 128)), mx.zeros((4, 32), dtype=mx.uint8), mx.ones((4, 2)), 128)
assert tuple(y.shape) == (0, 4)
|