Save inference parity implementation and evaluation harness

tools/mtplx-gdn-run-reference.py

@@ -0,0 +1,103 @@
+"""Whole original compiled decoder-run receipts, using the Rust smoke geometry.
+
+No downloads or real model weights. Run only under test-supervisor --command.
+The fixture uses real GDN/HC dimensions and two small MoE experts; it is a graph
+and numerical contract, not a model-throughput benchmark.
+"""
+import io
+import os
+from pathlib import Path
+import runpy
+import sys
+
+print("gdn-run reference: importing pinned runtime", file=sys.stderr, flush=True)
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten
+from mlx_lm.models.switch_layers import QuantizedSwitchLinear
+from mtplx.models import qwen4_exp as qwen
+
+
+def main():
+    assert mx.__version__ == "0.32.2", mx.__version__
+    helpers = runpy.run_path(str(Path(__file__).with_name("mtplx-kernel-fixtures.py")))
+    emit = helpers["emit"]
+    # Explicitly mirror decoder.rs's current connected model fixture.
+    os.environ.update(MTPLX_FUSED_HC="1", MTPLX_FUSED_HC_V3="0",
+        MTPLX_FUSED_GDN_STEP="1", MTPLX_FUSED_GDN_CONVNORM="0",
+        MTPLX_FUSED_CONVNORM_VERIFY="1", MTPLX_FUSED_GDN_OUT="0",
+        MTPLX_FUSED_MOE_DECODE="0", MTPLX_FUSED_MOE_VERIFY="0")
+    args = qwen.TextArgs(num_hidden_layers=1, vocab_size=32, num_experts=2,
+        num_experts_per_tok=1, moe_intermediate_size=64,
+        shared_expert_intermediate_size=64, layer_types=["linear_attention"])
+    model = qwen.Qwen4ExpTextModel(args)
+    layer = model.layers[0]
+    zero = lambda shape: mx.zeros(shape, mx.bfloat16)
+
+    def quantized(n, k, experts=None):
+        # Tiny constructors avoid evaluating random placeholder parameters.
+        module = (nn.QuantizedLinear(64, 1, bias=False, group_size=64, bits=4)
+                  if experts is None else QuantizedSwitchLinear(64, 1, 1,
+                      bias=False, group_size=64, bits=4))
+        prefix = () if experts is None else (experts,)
+        module.weight = mx.zeros((*prefix, n, k // 8), mx.uint32)
+        module.scales = zero((*prefix, n, k // 64))
+        module.biases = zero((*prefix, n, k // 64))
+        return module
+
+    norm = mx.ones((10240,), mx.bfloat16)
+    down, up, inject = zero((320, 10240)), zero((10240, 320)), zero((4, 10240))
+    for hc in (layer.attn_hyper_connection, layer.mlp_hyper_connection):
+        hc.hc_norm.weight = norm
+        hc.input_mix_weight_down.weight = down
+        hc.input_mix_weight_up.weight = up
+        hc.block_inject_weight.weight = inject
+    gdn = layer.linear_attn
+    projected = quantized(16480, 2560)
+    gdn.in_proj_fused = qwen._FusedGDNInProj(projected.weight, projected.scales,
+        projected.biases, 64, 4, "affine", [10240, 16384, 16432])
+    for name in ("qkv", "z", "a", "b"):
+        gdn.pop("in_proj_" + name, None)
+    gdn.out_proj = quantized(2560, 6144)
+    gdn.conv1d.weight = zero((10240, 4, 1))
+    gdn.A_log = zero((48,))
+    gdn.dt_bias = zero((48,))
+    gdn.norm.weight = mx.ones((128,), mx.bfloat16)
+    experts = quantized(128, 2560, 2)
+    layer.mlp.switch_mlp = qwen._FusedGateUpSwitchGLU(quantized(2560, 64, 2),
+        experts.weight, experts.scales, experts.biases, 64, 4, "affine")
+    shared = quantized(128, 2560)
+    layer.mlp.shared_expert = qwen._FusedGateUpMLP(quantized(2560, 64),
+        shared.weight, shared.scales, shared.biases, 64, 4, "affine")
+    layer.mlp.gate = quantized(2, 2560)
+    layer.mlp.shared_expert_gate = quantized(1, 2560)
+    model.eval()
+    assert not gdn.training
+    print("gdn-run reference: materializing fixture parameters", file=sys.stderr, flush=True)
+    mx.eval(*[a for _, a in tree_flatten(layer.parameters())])
+    for step, (rows, capture) in enumerate(2 * ((4, False), (1, False), (1, True), (4, True))):
+        hyper_fused = step < 4
+        if step % 4 == 0:
+            states = [zero((1, 3, 10240)), mx.zeros((1, 48, 128, 128), mx.float32)]
+            os.environ["MTPLX_FUSED_HC"] = str(int(hyper_fused))
+            model._decode_run_fns.clear()
+        print(f"gdn-run reference: step={step} rows={rows} capture={capture}", file=sys.stderr, flush=True)
+        h = mx.full((1, rows, 10240), 0.125, mx.bfloat16)
+        mx.eval(h, *states)
+        token = qwen._VERIFY_CAPTURE.set(capture)
+        try:
+            outputs = model._get_run_fn((0,), capture)(h, *states)
+        finally:
+            qwen._VERIFY_CAPTURE.reset(token)
+        dot = io.StringIO()
+        mx.export_to_dot(dot, hidden=h, conv=states[0], delta=states[1],
+            **{f"out{i}": a for i, a in enumerate(outputs)})
+        emit(f"gdn_run_r{rows}_capture{int(capture)}_hc{int(hyper_fused)}", outputs,
+            rows=rows, capture=capture, hyper_fused=hyper_fused, dot=dot.getvalue(), shapes=[a.shape for a in outputs],
+            contract=helpers["compiled_dot_contract"](dot.getvalue(),
+                [f"out{i}" for i in range(len(outputs))], {"hidden", "conv", "delta"}))
+        states = list(outputs[1:3])
+
+
+if __name__ == "__main__":
+    main()