Save inference parity implementation and evaluation harness

tools/mtplx-binding-reference.py

@@ -0,0 +1,81 @@
+"""Real MTPLX module construction/strict lazy loading, no model inference.
+
+Only MTP's two small pre-FC norms are evaluated. Run under test-supervisor.
+"""
+import sys
+print("binding reference: importing pinned runtime", file=sys.stderr, flush=True)
+import hashlib
+import json
+import os
+from contextlib import redirect_stdout
+from pathlib import Path
+import mlx.core as mx
+import numpy as np
+from mlx.utils import tree_flatten
+from mlx_lm.utils import load_model
+from mtplx.models import qwen4_exp as qwen
+
+
+def main(root, mode=None):
+    assert mx.__version__ == "0.32.2"
+    name_in_bank = lambda name: "mtp." + name[len("language_model.mtp."):] if name.startswith("language_model.mtp.") else name
+    for mask, mtp in ((0, False), (7, False), (0, True), (7, True)):
+        print(f"binding reference: mask={mask} mtp={mtp}", file=sys.stderr, flush=True)
+        os.environ.update(MTPLX_FUSED_GATE_UP=str(mask & 1), MTPLX_FUSED_GDN_INPROJ=str((mask >> 1) & 1), MTPLX_FUSED_QSA_QKV=str((mask >> 2) & 1))
+        with redirect_stdout(sys.stderr):
+            model, _ = load_model(Path(root), lazy=True, strict=True, get_model_classes=lambda **_: (qwen.Model, qwen.ModelArgs))
+            if mtp:
+                assert model.attach_mtp(Path(root))
+        if mode == "--parameter-order":
+            names = [name_in_bank(name) for name, _ in tree_flatten(model.parameters())]
+            print(json.dumps(dict(mask=mask, mtp=mtp, names=names),separators=(",", ":")),flush=True)
+            del model, names
+            continue
+        if mode == "--post-weight-load":
+            from mtplx.memory_plan import ngram_table_resident_policy
+            assert not ngram_table_resident_policy(), "this test must not make the full table resident"
+            with redirect_stdout(sys.stderr):
+                model.post_weight_load(Path(root))
+            tables = []
+            for i, layer in enumerate(model.layers):
+                if getattr(layer, "ple", None) is None:
+                    continue
+                table = layer.ple.ple_embedding.ngram_embedding
+                sidecar = table._sidecar
+                out = table(mx.array([0, 1, 255, 1], dtype=mx.int64))
+                mx.eval(out)
+                tables.append(dict(layer=i, resident=getattr(table, "_lazy_parts", None) is not None,
+                    hot_mb=sidecar._hot_cap_rows * sidecar._hot_row_bytes // 2**20,
+                    prefetch=sidecar._pool is not None, shape=list(out.shape),
+                    sha256=hashlib.sha256(np.asarray(out.astype(mx.float32)).astype('<f4').tobytes()).hexdigest()))
+            print(json.dumps(dict(mask=mask, mtp=mtp, tables=tables, ar_ready=model.set_ar_pipeline_mode(True)),separators=(",", ":")),flush=True)
+            del model
+            continue
+        parameters = tree_flatten(model.parameters())
+        rows = sorted([[name_in_bank(name), list(a.shape), str(a.dtype).removeprefix("mlx.core.")] for name, a in parameters])
+        quantized = {}
+        for name, module in model.named_modules():
+            name = name_in_bank(name)
+            if hasattr(module, "bits") and hasattr(module, "weight") and getattr(module, "scales", None) is not None:
+                quantized[name] = [module.bits, module.group_size, list(module.weight.shape)]
+            elif hasattr(module, "bits") and hasattr(module, "gu_weight"):
+                quantized[name + ".gu"] = [module.bits, module.group_size, list(module.gu_weight.shape)]
+        norms = {}
+        if mtp:
+            for name in ("pre_fc_norm_embedding", "pre_fc_norm_hidden"):
+                a = getattr(model.mtp, name).weight
+                mx.eval(a)
+                norms["mtp." + name + ".weight"] = hashlib.sha256(np.asarray(a.astype(mx.float32)).astype('<f4').tobytes()).hexdigest()
+        print(json.dumps(dict(mask=mask, mtp=mtp, parameter_count=len(rows),
+            metadata_sha256=hashlib.sha256(json.dumps(rows,separators=(",", ":")).encode()).hexdigest(),
+            quantized=quantized, norms=norms,
+            ple=[i for i, layer in enumerate(model.layers) if getattr(layer,"ple",None) is not None],
+            linear=[i for i, layer in enumerate(model.layers) if layer.is_linear]),separators=(",", ":")),flush=True)
+        del model, parameters, quantized, rows
+
+
+if __name__ == "__main__":
+    assert len(sys.argv) in (2, 3)
+    if len(sys.argv) == 3:
+        assert sys.argv[2] in ("--post-weight-load", "--parameter-order")
+    main(sys.argv[1], sys.argv[2] if len(sys.argv) == 3 else None)