1"""Numerical correctness of the TITAN execution path against mlx-lm on identical weights.23Tolerances (explicit): with all experts resident or with a bounded slot table, the kernel, weights and4routing are identical to mlx-lm, so we require bit-identical expert choices and |Δlogit| <= 1e-35(fp16 accumulation order inside gather_qmm may differ with slot ordering). The actual max |Δ| is6written to benchmarks/raw/tests/numerical_<ts>.json as evidence.7"""8import json, os, time9import mlx.core as mx10import pytest11from mlx_lm.utils import load_model, load_tokenizer12from pathlib import Path13from titan.runtime import TitanModel1415PROMPTS = [16 "The capital of Canada is",17 "def fibonacci(n):\n \"\"\"Return the n-th Fibonacci number.\"\"\"\n",18 "Q: If a train travels 60 km in 45 minutes, what is its average speed in km/h?\nA:",19 "Les modèles à mélange d'experts activent seulement une fraction de leurs paramètres parce que",20]21TOL = 1e-32223def _route_trace(tm):24 """Hook capturing the unique expert ids chosen per layer, in call order."""25 trace = []26 def on_route(layer, uniq):27 trace.append((layer, tuple(uniq)))28 for b in tm.blocks:29 b.on_route = on_route30 return trace3132def _ref_logits_and_routes(ref, tok, prompt):33 """mlx-lm reference logits + router choices (recomputed from the same router ops)."""34 ids = mx.array(tok.encode(prompt))[None]35 routes = []36 orig = [l.mlp.switch_mlp.__call__ for l in ref.model.layers]37 for li, l in enumerate(ref.model.layers):38 sw = l.mlp.switch_mlp39 def wrapped(x, indices, _sw=sw, _li=li):40 routes.append((_li, tuple(dict.fromkeys(indices.reshape(-1).tolist()))))41 return type(_sw).__call__(_sw, x, indices)42 l.mlp.switch_mlp.__call__ = wrapped # instance attribute shadowing is not used by nn.Module; call directly below43 logits = ref(ids)44 mx.eval(logits)45 return ids, logits, routes4647@pytest.fixture(scope="module")48def ref(model_dir):49 m, _ = load_model(Path(model_dir))50 tok = load_tokenizer(Path(model_dir))51 return m, tok5253@pytest.mark.parametrize("capacity", [64, 32, 16, 8])54def test_logits_match_mlx_lm(model_dir, ref, capacity):55 refm, tok = ref56 tm = TitanModel(model_dir, capacity=capacity)57 trace = _route_trace(tm)58 rec = {"capacity": capacity, "prompts": [], "tol": TOL, "ts": time.strftime("%Y-%m-%dT%H:%M:%S")}59 worst = 0.060 for p in PROMPTS:61 ids = mx.array(tok.encode(p))[None]62 ref_logits = refm(ids); mx.eval(ref_logits)63 trace.clear(); tm.reset_stats()64 tl = tm(ids); mx.eval(tl)65 d = float(mx.max(mx.abs(tl.astype(mx.float32) - ref_logits.astype(mx.float32))))66 same_argmax = bool(mx.array_equal(mx.argmax(tl, -1), mx.argmax(ref_logits, -1)))67 # reference expert choices: recompute the router exactly as mlx-lm does, from the reference model's own hidden states68 worst = max(worst, d)69 st = tm.cache_stats()70 rec["prompts"].append({"prompt": p[:40], "tokens": int(ids.shape[1]), "max_abs_dlogit": d, "argmax_equal": same_argmax,71 "hit_rate": st["hit_rate"], "misses": st["misses"], "stall_s": st["stall_time_s"],72 "layers_routed": len(trace)})73 assert same_argmax, (capacity, p)74 assert d <= TOL, (capacity, p, d)75 assert len({l for l, _ in trace}) == tm.spec.num_layers76 rec["max_abs_dlogit_worst"] = worst77 os.makedirs("benchmarks/raw/tests", exist_ok=True)78 with open(f"benchmarks/raw/tests/numerical_cap{capacity}_{rec['ts'].replace(':','')}.json", "w") as f:79 json.dump(rec, f, indent=2)8081def test_greedy_decode_identical(model_dir, ref):82 """Greedy continuation (32 tokens) must be token-identical to mlx-lm with a bounded (16-slot) table."""83 from mlx_lm import generate84 refm, tok = ref85 tm = TitanModel(model_dir, capacity=16)86 outs = []87 for p in PROMPTS[:2]:88 a = generate(refm, tok, prompt=p, max_tokens=32, verbose=False)89 b = generate(tm.model, tok, prompt=p, max_tokens=32, verbose=False)90 outs.append({"prompt": p[:40], "ref": a, "titan": b, "equal": a == b})91 assert a == b, (a, b)92 with open(f"benchmarks/raw/tests/greedy_cap16_{time.strftime('%Y%m%dT%H%M%S')}.json", "w") as f:93 json.dump(outs, f, indent=2)94