1"""H4 cheap fusion: one gather_qmm for up+gate must be bit-identical to two (row-independent dot products)."""2import json, os, time, mlx.core as mx, pytest3from pathlib import Path4from mlx_lm.utils import load_model, load_tokenizer5from titan.runtime import TitanModel67def test_fused_upgate_bit_identical(model_dir):8 refm, _ = load_model(Path(model_dir)); tok = load_tokenizer(Path(model_dir))9 from titan.registry import load_spec10 E = load_spec(model_dir).num_experts11 rec = {"model": model_dir, "ts": time.strftime("%Y-%m-%dT%H:%M:%S"), "cases": []}12 for cap in (E, max(8, E // 4)):13 tm = TitanModel(model_dir, capacity=cap, fused_upgate=True)14 if cap == E: tm.warm_all()15 for p in ["The capital of Canada is", "def fibonacci(n):\n \"\"\"Return the n-th Fibonacci number.\"\"\"\n"]:16 ids = mx.array(tok.encode(p))[None]17 ref = refm(ids); out = tm(ids); mx.eval(ref, out)18 d = float(mx.max(mx.abs(out.astype(mx.float32) - ref.astype(mx.float32))))19 rec["cases"].append({"capacity": cap, "prompt": p[:30], "max_abs_dlogit": d})20 assert d <= 1e-3, (cap, p, d)21 os.makedirs("benchmarks/raw/tests", exist_ok=True)22 json.dump(rec, open(f"benchmarks/raw/tests/fused_upgate_{tm.spec.model_type}_{rec['ts'].replace(':','')}.json", "w"), indent=2)23