| """The custom Metal kernels against plain MLX ops (random weights, no checkpoint needed).""" |
| import numpy as np |
| import pytest |
|
|
| mx = pytest.importorskip("mlx.core") |
|
|
| from maya_mac import mlx_kernels as K |
| from maya_mac.mlx_model import Linear, _gelu |
|
|
| F16 = mx.float16 |
|
|
|
|
| def rel(a, b): |
| a, b = np.array(a.astype(mx.float32), np.float64), np.array(b.astype(mx.float32), np.float64) |
| return float(np.linalg.norm(a - b) / np.linalg.norm(b)) |
|
|
|
|
| @pytest.mark.parametrize("bits", [None, 8]) |
| @pytest.mark.parametrize("T,Kd,N", [(1, 1024, 3072), (3, 1024, 1024), (1, 3072, 1024), (4, 1024, 2048)]) |
| def test_linear(bits, T, Kd, N): |
| mx.random.seed(0) |
| lin = Linear(mx.random.normal((N, Kd)) * 0.03, F16, bits) |
| w = (mx.dequantize(lin.wq, lin.scales, lin.biases, 64, 8) if bits else lin.w).astype(mx.float32) |
| x = (mx.random.normal((T, Kd)) * 2 + 0.5).astype(F16) |
| g, b = (1 + 0.1 * mx.random.normal((Kd,))).astype(F16), (0.1 * mx.random.normal((Kd,))).astype(F16) |
| res, c = mx.random.normal((T, N)).astype(F16), mx.random.normal((N,)).astype(F16) |
| ln = mx.fast.layer_norm(x.astype(mx.float32), g.astype(mx.float32), b.astype(mx.float32), 1e-5) |
| assert rel(K.linear(x, lin, ln=(g, b), gelu=True, res=res, add=c), _gelu(ln @ w.T) + res + c) < 2e-3 |
| assert rel(K.linear(x, lin), x.astype(mx.float32) @ w.T) < 2e-3 |
|
|
|
|
| @pytest.mark.parametrize("P,T,last", [(0, 9, False), (0, 9, True), (12, 1, False), (0, 1, False), (31, 1, False)]) |
| def test_attend(P, T, last): |
| H, HD = 16, 64 |
| mx.random.seed(1) |
| qkv = mx.random.normal((T, 3 * H * HD)).astype(F16) |
| kv = mx.random.normal((2, H, P, HD)).astype(F16) if P else None |
| out, kv_out = K.attend(qkv, kv, H, HD, last) |
|
|
| def heads(a): |
| return a.reshape(T, H, HD).transpose(1, 0, 2)[None] |
| q, k, v = heads(qkv[:, :1024]), heads(qkv[:, 1024:2048]), heads(qkv[:, 2048:]) |
| if P: |
| k, v = mx.concatenate([kv[0][None], k], 2), mx.concatenate([kv[1][None], v], 2) |
| mask = mx.arange(P + T)[None] <= (P + mx.arange(T))[:, None] |
| f32 = mx.float32 |
| ref = mx.fast.scaled_dot_product_attention(q.astype(f32), k.astype(f32), v.astype(f32), scale=HD ** -0.5, |
| mask=mask)[0].transpose(1, 0, 2).reshape(T, H * HD) |
| assert rel(out, ref[-1:] if last else ref) < 2e-3 |
| assert mx.array_equal(kv_out, mx.stack([k[0], v[0]])).item() |
|
|
|
|
| def test_sample_distribution(): |
| """Draws stay inside the top-k (ties with the k-th value kept) and follow softmax(logits / temp) over it.""" |
| mx.random.seed(2) |
| logits = (mx.random.normal((1, 2048)) * 3).astype(F16) |
| table = mx.random.normal((31, 2048, 1024)).astype(F16) |
| prefix = mx.random.normal((3, 1024)).astype(F16) |
| inv_t, k, n = 1 / 0.7, 64, 20000 |
| lf = np.array(logits.astype(mx.float32))[0] |
| p = np.where(lf >= np.sort(lf)[-k], np.exp((lf - lf.max()) * inv_t), 0.0) |
| p /= p.sum() |
| draw = mx.compile(lambda key: K.sample(logits, mx.random.gumbel((32, 2048), key=key), 4, table, prefix, |
| inv_t, k)[0]) |
| keys = mx.random.split(mx.random.key(1), n) |
| codes = np.array(mx.concatenate([draw(keys[i]) for i in range(n)])) |
| freq = np.bincount(codes, minlength=2048) / n |
| assert (freq[p == 0] == 0).all() |
| assert 0.5 * np.abs(freq - p).sum() < 0.03 |
| code, out = K.sample(logits, mx.random.gumbel((32, 2048), key=keys[0]), 4, table, prefix, inv_t, k) |
| assert mx.array_equal(out[:3], prefix).item() and mx.array_equal(out[3], table[4, code[0]]).item() |
|
|
