Save inference parity implementation and evaluation harness

2026-09-10 22:45:59 +02:00
parent b99ce2aa10
commit 02db0968ae
198 changed files with 111205 additions and 586 deletions
@@ -0,0 +1,138 @@
+"""Probe projection components selected by MTPLX's Qwen implementation.
+
+This is an operator diagnostic, not a full-model or serving benchmark. Only the
+selected projections are materialized; no weights are downloaded.
+"""
+
+import argparse
+import hashlib
+import json
+import statistics
+import time
+from pathlib import Path
+
+import mlx.core as mx
+import numpy as np
+from mtplx.models import qwen4_exp
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("model", type=Path)
+    parser.add_argument("--capture", type=Path, help="Capture only the 32-row router GPU operation")
+    parser.add_argument("--check-split-k", action="store_true",
+                        help="Diagnose router rounding with explicitly batched K partitions")
+    args = parser.parse_args()
+    assert mx.__version__ == "0.32.2", "Use the pinned MTPLX MLX environment"
+    index = json.loads((args.model / "model.safetensors.index.json").read_text())["weight_map"]
+    config = json.loads((args.model / "config.json").read_text())
+    quantization = config["quantization"]
+    model_args = qwen4_exp.TextArgs.from_dict(config["text_config"])
+    source = Path(qwen4_exp.__file__).resolve()
+    print(json.dumps({"reference": "MTPLX Qwen projection components",
+                      "source": str(source),
+                      "source_sha256": hashlib.sha256(source.read_bytes()).hexdigest(),
+                      "mlx_version": mx.__version__}), flush=True)
+    values = []
+    state = 0x12345678
+    for _ in range(2560):
+        state = (state * 1664525 + 1013904223) & 0xFFFFFFFF
+        values.append(((state >> 24) - 128) / 32)
+    for suffix in ("linear_attn.in_proj_qkv", "mlp.gate",
+                   "linear_attn.in_proj_b", "mlp.shared_expert_gate"):
+        name = "language_model.model.layers.0." + suffix
+        weights = {}
+        for part in ("weight", "scales", "biases"):
+            key = name + "." + part
+            weights[part] = mx.load(str(args.model / index[key]))[key]
+        mx.eval(*weights.values())
+        layout = quantization[name]
+        module = (qwen4_exp.GatedDeltaNet(model_args) if suffix.startswith("linear_attn.")
+                  else qwen4_exp.SparseMoeBlock(model_args))
+        attribute = suffix.split(".")[-1]
+        projection = getattr(module, attribute).to_quantized(**layout)
+        projection.load_weights(list(weights.items()), strict=True)
+        setattr(module, attribute, projection)
+        for rows in (1, 2, 3, 4, 31, 32, 63, 64, 2047, 2048):
+            x = mx.array([[values] * rows], dtype=mx.bfloat16)
+            mx.eval(x)
+            capture = args.capture and suffix == "mlp.gate" and rows == 32
+            if capture:
+                mx.metal.start_capture(str(args.capture))
+            try:
+                result = getattr(module, attribute)(x)
+                mx.eval(result)
+            finally:
+                if capture:
+                    mx.metal.stop_capture()
+            raw = np.asarray(result.astype(mx.float32)).astype("<f4").tobytes()
+            digest = hashlib.sha256(raw).hexdigest()
+            if suffix == "linear_attn.in_proj_qkv" and rows == 1:
+                assert digest == "8952fadfdb1ef4450fd23b0060d62ea9b118367733014985a94c2d815a2be203"
+            print(json.dumps({"mlx_version": mx.__version__, "projection": name,
+                              "quantization": layout, "rows": rows,
+                              "first_row_sha256": hashlib.sha256(raw[:result.shape[-1] * 4]).hexdigest(),
+                              "sha256": digest}), flush=True)
+            if rows in (2, 3, 4, 32, 64, 2048):
+                # Operator call + eval wall time, not a serving benchmark or
+                # an extra GPU counter probe. Discard the first warmup sample.
+                timings = []
+                for trial in range(33):
+                    started = time.perf_counter_ns()
+                    measured = getattr(module, attribute)(x)
+                    mx.eval(measured)
+                    elapsed = (time.perf_counter_ns() - started) / 1000
+                    if trial:
+                        timings.append(elapsed)
+                print(json.dumps({"diagnostic": "projection_call_eval_wall",
+                                  "projection": name, "rows": rows, "samples": len(timings),
+                                  "median_us": statistics.median(timings)}), flush=True)
+            if args.check_split_k and suffix == "mlp.gate" and rows in (32, 64):
+                # Counterfactual component operations, not a replacement oracle:
+                # a real weight batch bypasses the non-batched split-K dispatch.
+                for splits in (1, 2, 4, 5, 8, 10, 20, 40):
+                    width = x.shape[-1] // splits
+                    batches = max(2, splits)
+                    partitioned = {}
+                    for key, value in weights.items():
+                        block = value.reshape(value.shape[0], splits, -1).transpose(1, 0, 2)
+                        partitioned[key] = mx.repeat(block, batches, axis=0) if splits == 1 else block
+                    inputs = x.reshape(rows, splits, width).transpose(1, 0, 2)
+                    if splits == 1:
+                        inputs = mx.repeat(inputs, batches, axis=0)
+                    elif rows == 32:
+                        # Small K raises MTPLX's QMV threshold to 33 rows. Pad
+                        # to retain BF16 QMM dequantization in this control.
+                        inputs = mx.concatenate([inputs, inputs], axis=1)
+                    parts = mx.quantized_matmul(
+                        inputs, partitioned["weight"], partitioned["scales"],
+                        partitioned["biases"], transpose=True, **layout,
+                    )
+                    combined = parts[0] if splits == 1 else parts.sum(axis=0)
+                    combined = combined[:rows]
+                    mx.eval(combined)
+                    diagnostic_hash = hashlib.sha256(
+                        np.asarray(combined.astype(mx.float32)).astype("<f4").tobytes()
+                    ).hexdigest()
+                    first_actual = np.asarray(combined.astype(mx.float32))[0]
+                    first_expected = np.asarray(result.astype(mx.float32))[0, 0]
+                    differences = np.flatnonzero(first_actual != first_expected)
+                    print(json.dumps({"diagnostic": "explicit_batched_partitions",
+                                      "rows": rows, "splits": splits,
+                                      "partition_rows": inputs.shape[1],
+                                      "sha256": diagnostic_hash,
+                                      "different_columns": int(len(differences)),
+                                      "first_differences": [[int(i), float(first_actual[i]), float(first_expected[i])] for i in differences[:4]],
+                                      "matches_projection": diagnostic_hash == digest}), flush=True)
+                    if splits == (20 if rows == 32 else 10):
+                        assert diagnostic_hash == digest, "Split-K QMM decomposition differs from MTPLX projection"
+                    if splits == 1:
+                        expected_unsplit = {
+                            32: "3da6d25816a695eda420aea0a76ff22ee4c59689857d0ad02ed16112d66625b6",
+                            64: "8e0702490c8b94ea1c167fdbbca27107d7ecf6ec0c0b4f446ea3f8927d34c0f4",
+                        }
+                        assert diagnostic_hash == expected_unsplit[rows], "Unsplit control changed"
+
+
+if __name__ == "__main__":
+    main()