Save inference parity implementation and evaluation harness

tools/qwen-gdn-reference.py

@@ -0,0 +1,137 @@
+"""Generate model-free GDN boundary fixtures with the installed MTPLX/MLX.
+
+Run with MTPLX's Python environment; this never loads or downloads a model.
+Large-row math follows GatedDeltaNet.__call__ in the pinned qwen4_exp.py;
+small rows call MTPLX's actual fused kernel. All files are diagnostic outputs.
+"""
+
+import argparse
+import hashlib
+import json
+import statistics
+import time
+from pathlib import Path
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+from mtplx.kernels.gdn_conv_norm import fused_gdn_conv_norm_rows
+from mtplx.models import qwen4_exp
+
+
+def delta_reference(output):
+    from mlx_lm.models import gated_delta
+
+    rng = np.random.default_rng(12345)
+    initial = mx.array(rng.normal(0, 0.01, (1, 48, 128, 128)).astype(np.float32))
+    np.array(initial).tofile(output / "delta-initial.f32")
+    source = Path(gated_delta.__file__).resolve()
+    receipt = {"mlx_version": mx.__version__, "reference_source": str(source),
+               "reference_sha256": hashlib.sha256(source.read_bytes()).hexdigest(),
+               "cases": []}
+    for rows in (1, 2, 3, 4, 5, 6, 7, 32, 2048):
+        def normalized(scale):
+            x = rng.normal(size=(1, rows, 16, 128)).astype(np.float32)
+            return mx.array(x / np.linalg.norm(x, axis=-1, keepdims=True) * scale).astype(mx.bfloat16)
+
+        q, k = normalized(128 ** -0.5), normalized(1)
+        v = mx.array(rng.normal(0, 0.2, (1, rows, 48, 128)).astype(np.float32)).astype(mx.bfloat16)
+        g = mx.array(rng.uniform(0.8, 0.999, (1, rows, 48)).astype(np.float32))
+        beta = mx.array(rng.uniform(0.1, 0.9, (1, rows, 48)).astype(np.float32)).astype(mx.bfloat16)
+        mx.eval(q, k, v, g, beta, initial)
+
+        def forward():
+            return gated_delta.gated_delta_kernel(q, k, v, g, beta, initial)
+
+        mx.eval(*forward())
+        elapsed = []
+        for _ in range(10):
+            start = time.perf_counter()
+            result = forward()
+            mx.eval(*result)
+            elapsed.append((time.perf_counter() - start) * 1000)
+        hashes = {}
+        for name, array in zip(("q", "k", "v", "g", "beta", "out", "state"),
+                               (q, k, v, g, beta, *result)):
+            raw = np.array(array.astype(mx.float32)).astype("<f4").tobytes()
+            (output / f"delta-{rows}-{name}.f32").write_bytes(raw)
+            hashes[name] = hashlib.sha256(raw).hexdigest()
+        case = {"rows": rows, "wall_ms": elapsed,
+                "median_ms": statistics.median(elapsed), "sha256": hashes}
+        receipt["cases"].append(case)
+        print(json.dumps(case), flush=True)
+    (output / "delta-reference.json").write_text(json.dumps(receipt, indent=2) + "\n")
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("output", type=Path)
+    parser.add_argument("--delta", action="store_true", help="Check the active scalar-gated recurrence boundary instead of Conv/Norm")
+    args = parser.parse_args()
+    args.output.mkdir(parents=True, exist_ok=True)
+    if args.delta:
+        delta_reference(args.output)
+        return
+    width = 10240
+    qk = 2048
+    # Binary fractions keep fixture inputs exactly representable in BF16.
+    values = lambda size, modulus: (
+        (np.arange(size, dtype=np.int32) % modulus - modulus // 2).astype(np.float32)
+        / 32
+    )
+    state_np = values(3 * width, 17).reshape(3, width)
+    weight_np = values(width * 4, 13).reshape(width, 4)
+    state = mx.array(state_np).astype(mx.bfloat16)
+    conv = nn.Conv1d(width, width, 4, groups=width, bias=False)
+    conv.weight = mx.array(weight_np[:, :, None]).astype(mx.bfloat16)
+    state_np.tofile(args.output / "state.f32")
+    weight_np.tofile(args.output / "weight.f32")
+    receipt = {
+        "mlx_version": mx.__version__,
+        "reference_source": str(Path(qwen4_exp.__file__).resolve()),
+        "reference_sha256": hashlib.sha256(Path(qwen4_exp.__file__).read_bytes()).hexdigest(),
+        "cases": [],
+    }
+    for rows in (1, 2, 3, 4, 5, 6, 7, 32, 2048):
+        x_np = values(rows * width, 23).reshape(rows, width)
+        x = mx.array(x_np).astype(mx.bfloat16)
+        mx.eval(x, state, conv.weight)
+
+        def forward():
+            if rows <= 6:
+                return fused_gdn_conv_norm_rows(x, state, conv.weight)
+            stream = mx.concatenate([state, x], axis=0)[None]
+            activated = nn.silu(conv(stream))
+            q, k, v = mx.split(activated, [qk, 2 * qk], axis=-1)
+
+            def l2norm(a):
+                a = a.reshape(1, rows, 16, 128)
+                f = a.astype(mx.float32)
+                return (f * mx.rsqrt((f * f).sum(-1, keepdims=True) + 1e-6)).astype(a.dtype)
+
+            q = (128 ** -0.5) * l2norm(q)
+            k = l2norm(k)
+            return q, k, v, mx.contiguous(stream[0, -3:])
+
+        mx.eval(*forward())  # Compile/warmup excluded from steady timing.
+        elapsed = []
+        for _ in range(10):
+            start = time.perf_counter()
+            result = forward()
+            mx.eval(*result)
+            elapsed.append((time.perf_counter() - start) * 1000)
+        x_np.tofile(args.output / f"{rows}-input.f32")
+        hashes = {}
+        for name, array in zip(("q", "k", "v", "state"), result):
+            raw = np.array(array.astype(mx.float32)).astype("<f4").tobytes()
+            (args.output / f"{rows}-{name}.f32").write_bytes(raw)
+            hashes[name] = hashlib.sha256(raw).hexdigest()
+        case = {"rows": rows, "path": "fused" if rows <= 6 else "conv1d_silu_l2",
+                "wall_ms": elapsed, "median_ms": statistics.median(elapsed), "sha256": hashes}
+        receipt["cases"].append(case)
+        print(json.dumps(case), flush=True)
+    (args.output / "reference.json").write_text(json.dumps(receipt, indent=2) + "\n")
+
+
+if __name__ == "__main__":
+    main()