diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index a9b708fe40..f08ae1b0d5 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -1,6 +1,7 @@ # Copyright © 2023-2024 Apple Inc. import math +import random import unittest import mlx.core as mx @@ -989,6 +990,89 @@ def call_kernel(a: mx.array, source): out = call_kernel(a, source) self.assertTrue(mx.array_equal(out, mx.ones_like(out))) + def test_spec_decode_verify(self): + # Validate mx.fast.spec_decode_verify against a pure-Python oracle that + # mirrors the greedy verify loop. On a Metal build the default device + # exercises the fused kernel and stream=mx.cpu the op-composition + # fallback; both must agree with the oracle. + def oracle(draft, tgt_tokens): + B, K = len(draft), len(draft[0]) + n_acc, committed = [], [] + for b in range(B): + na = K + for j in range(K): + if draft[b][j] != tgt_tokens[b][j]: + na = j + break + n_acc.append(na) + committed.append(draft[b][:na] + [tgt_tokens[b][na]]) + return n_acc, committed + + # One-hot logits make the greedy argmax token unambiguous while still + # driving the same large-V reduction the real caller feeds in. + def onehot_logits(tgt_tokens, V, dtype): + B, L = len(tgt_tokens), len(tgt_tokens[0]) + lg = [[[-10.0] * V for _ in range(L)] for _ in range(B)] + for b in range(B): + for j in range(L): + lg[b][j][tgt_tokens[b][j]] = 10.0 + return mx.array(lg).astype(dtype) + + def check(draft, tgt_tokens, V, dtype, stream): + exp_n, exp_c = oracle(draft, tgt_tokens) + lg = onehot_logits(tgt_tokens, V, dtype) + kw = {"stream": stream} if stream is not None else {} + n, c = mx.fast.spec_decode_verify( + mx.array(draft, dtype=mx.int32), lg, **kw + ) + mx.eval(n, c) + n, c = n.tolist(), c.tolist() + got_c = [c[b][: exp_n[b] + 1] for b in range(len(draft))] + self.assertEqual(n, exp_n) + self.assertEqual(got_c, exp_c) + + cases = [ + ("full-accept", [[5, 6, 7]], [[5, 6, 7, 9]]), + ("mismatch@2", [[5, 6, 7]], [[5, 6, 1, 9]]), + ("reject-all", [[5, 6, 7]], [[2, 6, 7, 9]]), + ("K=1", [[3]], [[3, 4]]), + ( + "batch2-K4", + [[1, 2, 3, 4], [1, 2, 3, 4]], + [[1, 2, 9, 4, 7], [1, 2, 3, 4, 8]], + ), + ("batch3-K1", [[7], [7], [7]], [[7, 1], [2, 3], [7, 9]]), + ] + # stream=None is the default device (fused kernel on Metal); mx.cpu is + # the op-composition fallback. + for path, stream in (("default", None), ("cpu-fallback", mx.cpu)): + for dtype in (mx.float32, mx.float16, mx.bfloat16): + for name, draft, tgt in cases: + with self.subTest(path=path, dtype=str(dtype), case=name): + check(draft, tgt, 32, dtype, stream) + + # Randomized sweep over prefix lengths, batch shapes, and corrected-token + # positions that are awkward to enumerate by hand. + rng = random.Random(7) + for i in range(40): + b, k, v = rng.randint(1, 4), rng.randint(1, 8), rng.randint(8, 96) + target = [[rng.randrange(v) for _ in range(k + 1)] for _ in range(b)] + draft = [ + [ + row[j] if rng.random() < 0.5 else (row[j] + rng.randrange(1, v)) % v + for j in range(k) + ] + for row in target + ] + with self.subTest(sweep=i, b=b, k=k, v=v): + check(draft, target, v, mx.float16, None) + + # target_logits.shape[1] must equal draft.shape[1] + 1. + with self.assertRaises(ValueError): + mx.fast.spec_decode_verify( + mx.zeros((1, 3), dtype=mx.int32), mx.zeros((1, 3, 8)) + ) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/test_spec_verify.py b/test_spec_verify.py deleted file mode 100644 index b09757d35f..0000000000 --- a/test_spec_verify.py +++ /dev/null @@ -1,111 +0,0 @@ -"""Validate mx.fast.spec_decode_verify against a pure-Python oracle mirroring the -easymlx-swift verify loop. Dependency-free (mlx + stdlib). - -On a Metal build this exercises BOTH paths: - - default device (GPU) -> the fused Metal kernel - - stream=mx.cpu -> the op-level composition fallback -Both must match the oracle. Run from the worktree root with the venv python: - .venv/bin/python test_spec_verify.py -""" -import random -import sys -import mlx.core as mx - - -def oracle(draft, tgt_tokens): - B, K = len(draft), len(draft[0]) - n_acc, committed = [], [] - for b in range(B): - na = K - for j in range(K): - if draft[b][j] != tgt_tokens[b][j]: - na = j - break - n_acc.append(na) - committed.append(draft[b][:na] + [tgt_tokens[b][na]]) - return n_acc, committed - - -def onehot_logits(tgt_tokens, V, dtype): - B, L = len(tgt_tokens), len(tgt_tokens[0]) - lg = [[[-10.0] * V for _ in range(L)] for _ in range(B)] - for b in range(B): - for j in range(L): - lg[b][j][tgt_tokens[b][j]] = 10.0 - return mx.array(lg).astype(dtype) - - -CASES = [ - ("full-accept", [[5, 6, 7]], [[5, 6, 7, 9]]), - ("mismatch@2", [[5, 6, 7]], [[5, 6, 1, 9]]), - ("reject-all", [[5, 6, 7]], [[2, 6, 7, 9]]), - ("K=1", [[3]], [[3, 4]]), - ("batch2-K4", [[1, 2, 3, 4], [1, 2, 3, 4]], [[1, 2, 9, 4, 7], [1, 2, 3, 4, 8]]), - ("batch3-K1", [[7], [7], [7]], [[7, 1], [2, 3], [7, 9]]), -] - -PATHS = [("kernel(gpu)", None), ("fallback(cpu)", mx.cpu)] - - -def run(): - print(f"default_device={mx.default_device()}") - V = 32 - allok = True - for path_name, stream in PATHS: - for dt in (mx.float32, mx.float16, mx.bfloat16): - for name, draft, tgt in CASES: - exp_n, exp_c = oracle(draft, tgt) - lg = onehot_logits(tgt, V, dt) - kw = {"stream": stream} if stream is not None else {} - n, c = mx.fast.spec_decode_verify( - mx.array(draft, dtype=mx.int32), lg, **kw) - mx.eval(n, c) - n, c = n.tolist(), c.tolist() - got_c = [c[b][: exp_n[b] + 1] for b in range(len(draft))] - ok = (n == exp_n) and (got_c == exp_c) - allok &= ok - if not ok: - print(f" [FAIL] {path_name} {str(dt):>16} {name}: " - f"got n={n} c={got_c} exp n={exp_n} c={exp_c}") - print(f" [{'OK ' if allok else 'FAIL'}] path={path_name}: all cases x fp32/fp16/bf16") - - # Exercise prefix lengths, batch shapes, and corrected-token positions that - # are awkward to enumerate by hand. One-hot logits make the expected greedy - # target token unambiguous while still driving the same argmax reduction. - rng = random.Random(7) - for _ in range(40): - b, k, v = rng.randint(1, 4), rng.randint(1, 8), rng.randint(8, 96) - target = [[rng.randrange(v) for _ in range(k + 1)] for _ in range(b)] - draft = [] - for row in target: - draft.append([ - row[j] if rng.random() < 0.5 else (row[j] + rng.randrange(1, v)) % v - for j in range(k) - ]) - exp_n, exp_c = oracle(draft, target) - n, c = mx.fast.spec_decode_verify( - mx.array(draft, dtype=mx.int32), onehot_logits(target, v, mx.float16) - ) - mx.eval(n, c) - got_n, got_c = n.tolist(), c.tolist() - got_c = [got_c[i][: exp_n[i] + 1] for i in range(b)] - if got_n != exp_n or got_c != exp_c: - allok = False - print(f" [FAIL] randomized b={b} k={k} v={v}: got n={got_n} c={got_c} " - f"exp n={exp_n} c={exp_c}") - break - else: - print(" [OK ] randomized GPU oracle sweep (40 cases)") - - try: - mx.fast.spec_decode_verify(mx.zeros((1, 3), dtype=mx.int32), mx.zeros((1, 3, 8))) - print(" [FAIL] expected error on bad K+1 shape"); allok = False - except Exception: - print(" [OK ] rejects target_logits.shape[1] != K+1") - - print("ALL OK" if allok else "SOME FAILED") - return 0 if allok else 1 - - -if __name__ == "__main__": - sys.exit(run())