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("frac", [1.0, 0.5, 0.25, 0.125])54def test_logits_match_mlx_lm(model_dir, ref, frac):55 refm, tok = ref56 from titan.registry import load_spec57 capacity = max(8, int(load_spec(model_dir).num_experts * frac))58 tm = TitanModel(model_dir, capacity=capacity)59 trace = _route_trace(tm)60 rec = {"capacity": capacity, "prompts": [], "tol": TOL, "ts": time.strftime("%Y-%m-%dT%H:%M:%S")}61 worst = 0.062 for p in PROMPTS:63 ids = mx.array(tok.encode(p))[None]64 ref_logits = refm(ids); mx.eval(ref_logits)65 trace.clear(); tm.reset_stats()66 tl = tm(ids); mx.eval(tl)67 d = float(mx.max(mx.abs(tl.astype(mx.float32) - ref_logits.astype(mx.float32))))68 same_argmax = bool(mx.array_equal(mx.argmax(tl, -1), mx.argmax(ref_logits, -1)))69 # reference expert choices: recompute the router exactly as mlx-lm does, from the reference model's own hidden states70 worst = max(worst, d)71 st = tm.cache_stats()72 rec["prompts"].append({"prompt": p[:40], "tokens": int(ids.shape[1]), "max_abs_dlogit": d, "argmax_equal": same_argmax,73 "hit_rate": st["hit_rate"], "misses": st["misses"], "stall_s": st["stall_time_s"],74 "layers_routed": len(trace)})75 assert same_argmax, (capacity, p)76 assert d <= TOL, (capacity, p, d)77 assert len({l for l, _ in trace}) == tm.spec.num_layers78 rec["max_abs_dlogit_worst"] = worst79 os.makedirs("benchmarks/raw/tests", exist_ok=True)80 rec["model"] = os.path.realpath(model_dir); rec["model_type"] = tm.spec.model_type81 with open(f"benchmarks/raw/tests/numerical_{tm.spec.model_type}_cap{capacity}_{rec['ts'].replace(':','')}.json", "w") as f:82 json.dump(rec, f, indent=2)8384def test_greedy_decode_identical(model_dir, ref):85 """Greedy continuation (32 tokens) must be token-identical to mlx-lm with a bounded (16-slot) table."""86 from mlx_lm import generate87 refm, tok = ref88 from titan.registry import load_spec89 tm = TitanModel(model_dir, capacity=max(8, load_spec(model_dir).num_experts // 4))90 outs = []91 for p in PROMPTS[:2]:92 a = generate(refm, tok, prompt=p, max_tokens=32, verbose=False)93 b = generate(tm.model, tok, prompt=p, max_tokens=32, verbose=False)94 outs.append({"prompt": p[:40], "ref": a, "titan": b, "equal": a == b})95 assert a == b, (a, b)96 with open(f"benchmarks/raw/tests/greedy_{tm.spec.model_type}_cap{tm.capacity}_{time.strftime('%Y%m%dT%H%M%S')}.json", "w") as f:97 json.dump(outs, f, indent=2)98