Save inference parity implementation and evaluation harness

tools/metal-trace-summary.py

@@ -0,0 +1,202 @@
+"""Summarize exported xctrace intervals for the attached target only.
+
+Reads XML incrementally; never prints the TOC environment or other processes.
+GPU intervals are not shader timings. Nested intervals are not added to level 0.
+Per-buffer covered time merges overlaps; group sums are not GPU utilization.
+"""
+
+import argparse
+import io
+import json
+import math
+import re
+import statistics
+import sys
+import xml.etree.ElementTree as ET
+
+
+def summarize(source, schemas, pid):
+    ids, totals, buffers = {}, {}, {}
+    scalar_tags = {"start-time", "duration", "process", "gpu-channel-name",
+                   "gpu-state", "metal-nesting-level", "metal-command-buffer-id"}
+    schema = None
+
+    def value(element):
+        if element is None:
+            return None
+        if "ref" in element.attrib:
+            return ids[element.attrib["ref"]]
+        return element.get("fmt") if element.tag == "process" else element.text
+
+    for event, element in ET.iterparse(source, events=("start", "end")):
+        if event == "start":
+            if element.tag == "node":
+                index = int(re.search(r"table\[(\d+)\]$", element.attrib["xpath"])[1])
+                schema = schemas[index - 1]
+            continue
+        if element.tag in scalar_tags and "id" in element.attrib:
+            ids[element.attrib["id"]] = value(element)
+        if element.tag == "row":
+            process = value(element.find("process"))
+            if process and process.endswith(f"({pid})"):
+                duration = value(element.find("duration"))
+                if duration is not None:
+                    ms = int(duration) / 1e6
+                    channel = value(element.find("gpu-channel-name"))
+                    state = value(element.find("gpu-state"))
+                    depth = value(element.find("metal-nesting-level"))
+                    key = "/".join(str(x) for x in (schema, channel, state, depth))
+                    total = totals.setdefault(key, {"rows": 0, "sum_ms": 0, "max_ms": 0})
+                    total["rows"] += 1
+                    total["sum_ms"] += ms
+                    total["max_ms"] = max(total["max_ms"], ms)
+                    if (schema == "metal-gpu-intervals" and channel == "Compute"
+                            and depth == "0" and state == "Active"):
+                        command = value(element.find("metal-command-buffer-id"))
+                        start = int(value(element.find("start-time"))) / 1e6
+                        record = buffers.setdefault(command, {"command_buffer": command,
+                            "start_ms": start, "end_ms": start, "gpu_interval_sum_ms": 0,
+                            "max_interval_ms": 0, "intervals": 0, "_spans": []})
+                        record["start_ms"] = min(record["start_ms"], start)
+                        record["end_ms"] = max(record["end_ms"], start + ms)
+                        record["gpu_interval_sum_ms"] += ms
+                        record["max_interval_ms"] = max(record["max_interval_ms"], ms)
+                        record["intervals"] += 1
+                        record["_spans"].append((start, start + ms))
+            element.clear()
+        elif element.tag == "node":
+            element.clear()
+    for record in buffers.values():
+        covered, end = 0, float("-inf")
+        for start, stop in sorted(record.pop("_spans")):
+            covered += max(0, stop - max(start, end))
+            end = max(end, stop)
+        record["gpu_covered_ms"] = covered
+    return {"pid": pid, "interval_groups": totals,
+            "largest_compute_buffers": sorted(buffers.values(),
+                key=lambda x: x["gpu_interval_sum_ms"], reverse=True)[:12]}
+
+
+def self_test():
+    xml = """<trace-query-result><node xpath="//table[1]">
+    <row><start-time id="1">1000000</start-time><duration id="2">2000000</duration>
+    <process id="3" fmt="ds4-server (42)"/><gpu-channel-name id="4">Compute</gpu-channel-name>
+    <gpu-state id="5">Active</gpu-state><metal-nesting-level id="6">0</metal-nesting-level>
+    <metal-command-buffer-id id="7">100</metal-command-buffer-id></row>
+    <row><start-time>4000000</start-time><duration ref="2"/><process ref="3"/>
+    <gpu-channel-name ref="4"/><gpu-state ref="5"/><metal-nesting-level ref="6"/>
+    <metal-command-buffer-id ref="7"/></row>
+    <row><start-time>2000000</start-time><duration ref="2"/><process ref="3"/>
+    <gpu-channel-name ref="4"/><gpu-state ref="5"/><metal-nesting-level ref="6"/>
+    <metal-command-buffer-id ref="7"/></row>
+    <row><start-time>1000000</start-time><duration ref="2"/><process ref="3"/>
+    <gpu-channel-name ref="4"/><gpu-state ref="5"/><metal-nesting-level>1</metal-nesting-level>
+    <metal-command-buffer-id ref="7"/></row>
+    <row><duration ref="2"/><process fmt="other (43)"/></row>
+    </node></trace-query-result>"""
+    result = summarize(io.StringIO(xml), ["metal-gpu-intervals"], 42)
+    assert list(result["interval_groups"].values()) == [
+        {"rows": 3, "sum_ms": 6, "max_ms": 2}, {"rows": 1, "sum_ms": 2, "max_ms": 2}]
+    assert result["largest_compute_buffers"] == [{"command_buffer": "100", "start_ms": 1,
+        "end_ms": 6, "gpu_interval_sum_ms": 6, "gpu_covered_ms": 5,
+        "max_interval_ms": 2, "intervals": 3}]
+    print("PASS: XML references, PID filtering, depth filtering and overlap union")
+
+
+def summarize_stages(source):
+    """Validate native counter receipts; group encoder spans, never shaders."""
+    headers, stages, groups = {}, {}, {}
+    for line in source:
+        # Supervisor JSON can precede a worker's complete native stderr record.
+        # Keep strict field/completeness checks; never discard a damaged record.
+        _, marker, record = line.partition("ds4: stage-counter ")
+        if not marker:
+            continue
+        values = dict(word.split("=", 1) for word in record.split())
+        cb = int(values["cb"])
+        if "stages" in values:
+            assert cb not in headers, f"duplicate CB {cb}"
+            assert values["resolved"] == "1" and values["dropped"] == "0", values
+            headers[cb] = values
+            stages[cb] = {}
+            continue
+        assert cb in headers and values["valid"] == "1", values
+        index = int(values["stage"])
+        assert index not in stages[cb], (cb, index)
+        start, end = int(values["start"]), int(values["end"])
+        duration = float(values["ms"])
+        scale = float(headers[cb]["ns_per_tick"])
+        assert end >= start and math.isfinite(duration) and duration >= 0
+        assert abs(duration - (end - start) * scale / 1e6) <= 0.000002
+        stages[cb][index] = (start, end)
+        key = (values["first"], values["last"], int(values["qwen_dispatches"]))
+        groups.setdefault(key, []).append(duration)
+    assert headers, "no counter receipts"
+    covered_total, cb_total = 0.0, 0.0
+    for cb, header in headers.items():
+        assert set(stages[cb]) == set(range(int(header["stages"]))), cb
+        covered, prior_end = 0, 0
+        for start, end in sorted(stages[cb].values()):
+            covered += max(0, end - max(start, prior_end))
+            prior_end = max(prior_end, end)
+        covered_ms = covered * float(header["ns_per_tick"]) / 1e6
+        cb_ms = float(header["cb_ms"])
+        assert math.isfinite(cb_ms) and cb_ms > 0
+        assert covered_ms <= cb_ms * 1.05 + 0.002, (cb, covered_ms, cb_ms)
+        covered_total += covered_ms
+        cb_total += cb_ms
+    return {"scope": "Qwen four-row command buffers, existing compute encoder boundaries",
+            "command_buffers": len(headers), "encoder_spans": sum(map(len, stages.values())),
+            "counter_covered_ms": covered_total, "command_buffer_gpu_ms": cb_total,
+            "groups_are_encoder_spans_not_individual_kernels": True,
+            "encoder_groups": sorted([
+                {"first_qwen_kernel": first, "last_qwen_kernel": last,
+                 "qwen_dispatches": dispatches, "samples": len(values),
+                 "sum_ms": sum(values), "median_ms": statistics.median(values), "max_ms": max(values)}
+                for (first, last, dispatches), values in groups.items()
+            ], key=lambda row: row["sum_ms"], reverse=True)}
+
+
+def stage_self_test():
+    log = ("ds4: stage-counter cb=1 rows=4 stages=2 dropped=0 resolved=1 ns_per_tick=1 cb_ms=0.05\n"
+           "ds4: stage-counter cb=1 stage=0 valid=1 start=100000 end=130000 ms=0.03 qwen_dispatches=2 first=a last=b\n"
+           "ds4: stage-counter cb=1 stage=1 valid=1 start=120000 end=140000 ms=0.02 qwen_dispatches=1 first=c last=c\n")
+    result = summarize_stages(io.StringIO(log))
+    assert result["encoder_spans"] == 2 and result["counter_covered_ms"] == 0.04
+    assert summarize_stages(io.StringIO(log.replace("ds4:", '{"event":ds4:'))) == result
+    for bad in (log.replace("resolved=1", "resolved=0"), log.replace("dropped=0", "dropped=1"),
+                log.replace("ms=0.03", "ms=1.25"), "\n".join(log.splitlines()[:-1])):
+        try:
+            summarize_stages(io.StringIO(bad))
+        except AssertionError:
+            continue
+        raise AssertionError("invalid stage trace accepted")
+    print("PASS: stage completeness, failure reporting, clock scale and overlap coverage")
+
+
+if __name__ == "__main__":
+    if sys.argv[1:] == ["--self-test"]:
+        self_test()
+        stage_self_test()
+    else:
+        parser = argparse.ArgumentParser(description=__doc__)
+        parser.add_argument("toc", nargs="?")
+        parser.add_argument("intervals", nargs="?")
+        parser.add_argument("--stage-log", help="Native DS4_METAL_STAGE_PROFILE stderr receipt")
+        args = parser.parse_args()
+        if args.stage_log:
+            if args.toc or args.intervals:
+                parser.error("do not combine --stage-log with XML inputs")
+            with open(args.stage_log) as source:
+                print(json.dumps(summarize_stages(source), indent=2))
+            sys.exit(0)
+        if not args.toc or not args.intervals:
+            parser.error("provide both TOC and intervals, or --stage-log")
+        toc = ET.parse(args.toc)
+        if len(toc.findall("./run")) != 1:
+            parser.error("export a trace containing exactly one run")
+        pid = int(toc.find("./run/info/target/process").attrib["pid"])
+        schemas = [table.attrib["schema"] for table in toc.findall("./run/data/table")]
+        result = summarize(args.intervals, schemas, pid)
+        result["trace_duration_seconds"] = float(toc.findtext("./run/info/summary/duration"))
+        print(json.dumps(result, indent=2))

tools/mtplx-allocator-reference.py

@@ -0,0 +1,108 @@
+"""Model-free allocation/cache receipts from the installed pinned MTPLX runtime.
+
+No replacement allocator, timings, model loading or downloads. stdout is JSONL.
+Physical addresses are replaced with run-local identities before emission.
+"""
+import gc
+import json
+import os
+
+import mlx.core as mx
+import numpy as np
+
+
+def main():
+    assert mx.__version__ == "0.32.2"
+    arrays = {}
+    identities = {}
+    page = os.sysconf("SC_PAGE_SIZE")
+    mx.set_cache_limit(4 * 1024 * 1024)
+    mx.clear_cache()
+    assert mx.get_active_memory() == 0
+
+    def emit(op, **fields):
+        print(json.dumps(dict(op=op, active=mx.get_active_memory(),
+                              cached=mx.get_cache_memory(), **fields)), flush=True)
+
+    def alloc(name, size):
+        before = mx.get_active_memory()
+        arrays[name] = mx.array(np.zeros(size, dtype=np.uint8))
+        mx.eval(arrays[name])
+        mx.synchronize()
+        pointer = np.asarray(arrays[name]).__array_interface__["data"][0] if size else None
+        identity = identities.setdefault(pointer, len(identities)) if size else None
+        emit("alloc", name=name, size=size, buffer=identity,
+             storage=mx.get_active_memory() - before)
+
+    def free(name):
+        del arrays[name]
+        gc.collect()
+        mx.synchronize()
+        emit("free", name=name)
+
+    def clear():
+        assert not arrays
+        mx.clear_cache()
+        emit("clear")
+
+    def limit(size):
+        mx.set_cache_limit(size)
+        emit("limit", size=size)
+
+    emit("init", page=page, limit=4 * 1024 * 1024, runtime=mx.__version__)
+    for i, size in enumerate([0, 1, 255, 256, 257, page - 1, page, page + 1,
+                              2 * page - 1, 2 * page, 2 * page + 1]):
+        alloc(str(i), size)
+    for name in list(arrays):
+        free(name)
+    for i, size in enumerate([256, 1, page, page + 1, 2 * page + 1]):
+        alloc(str(i), size)
+    for name in list(arrays):
+        free(name)
+    clear()
+
+    # Equal-size multimap entries reuse the oldest insertion, not the newest.
+    for name in "abc":
+        alloc(name, 257)
+    for name in "bac":
+        free(name)
+    for name in "xyz":
+        alloc(name, 256)
+    for name in "xyz":
+        free(name)
+    clear()
+
+    # Strict upper bounds: min(2 * request, request + 2 * page).
+    for stored, requested in [(512, 256), (511, 256),
+                              (4 * page, 2 * page), (3 * page, 2 * page),
+                              (6 * page, 4 * page), (5 * page, 4 * page)]:
+        alloc("old", stored)
+        free("old")
+        alloc("new", requested)
+        free("new")
+        clear()
+
+    for name, size in zip("abcd", [64, 128, 256, 512]):
+        alloc(name, size)
+    for name in "cadb":
+        free(name)
+    limit(600)
+    alloc("e", 2048)  # trim from oldest, overshooting by complete buffers
+    free("e")       # free checks the pre-insertion size, allowing overshoot
+    alloc("f", 10000)
+    free("f")
+    limit(1)
+    alloc("g", 20000)  # >=90% requested release clears the whole pool
+    free("g")
+    clear()
+    limit(4096)
+    alloc("old", 64)
+    free("old")
+    limit(0)  # changing the limit does not clear; reuse is still attempted first
+    alloc("new", 64)
+    free("new")
+    clear()
+
+
+if __name__ == "__main__":
+    main()

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)

tools/mtplx-cache-reference.py

@@ -0,0 +1,50 @@
+"""Reproduce MTPLX's AR prefix-selection failure without model/GPU work.
+
+Uses the pinned reference's real bank and lookup functions with empty cache
+payloads. This verifies selection policy only, not recurrent-state correctness.
+"""
+
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+from mtplx import generation
+from mtplx.session_bank import SessionBank
+
+
+def main():
+    runtime = SimpleNamespace(model_path=Path("/diagnostic/qwen"), mtp_enabled=True)
+    bank = SessionBank()
+    prompt = list(range(7896))
+    for length, policy in ((7460, "cycle"), (7836, "committed")):
+        bank.put(runtime=runtime, token_ids=prompt[:length], cache=[],
+                 logits=None, hidden=None, mtp_history_policy=policy)
+
+    assert bank.longest_prefix(prompt).prefix_len == 7836
+    assert bank.restore(runtime, prompt, mtp_history_policy="cycle",
+                        cache_factory=list) is None
+    assert bank.last_miss_reason == "policy_mismatch"
+    candidates = bank.near_prefix_candidates(prompt, mtp_history_policy="cycle")
+    assert [(entry.prefix_len, matched) for entry, matched in candidates] == [
+        (7836, 7836), (7460, 7460),
+    ]
+    assert generation._restore_near_prefix_prompt_state(
+        runtime, prompt, base_hidden_variant=None, mtp_hidden_variant=None,
+        mtp_history_policy="cycle", session_bank=bank, template_hash=None,
+        draft_head_identity=None, policy_fingerprint=None, cache_factory=list,
+    ) is None
+
+    # Control: the shorter prefix is usable; merely removing the incompatible
+    # longest candidate makes the unchanged exact-restore function serve it.
+    del bank._entries[tuple(prompt[:7836])]
+    restored = bank.restore(runtime, prompt, mtp_history_policy="cycle",
+                            cache_factory=list)
+    assert restored is not None and restored.entry.prefix_len == 7460
+    print(json.dumps({"ok": True, "reference": "MTPLX SessionBank + generation",
+                      "incompatible_longest_prefix": 7836,
+                      "shadowed_usable_prefix": 7460, "prompt_tokens": len(prompt),
+                      "control_restored_prefix": restored.entry.prefix_len}))
+
+
+if __name__ == "__main__":
+    main()

tools/mtplx-execution-reference.py

@@ -0,0 +1,592 @@
+"""Full installed MTPLX forward oracle; run only under test-supervisor.
+
+Uses the exact DS4Server native canary entry point from a linked bridge dylib.
+Diagnostic, not a clean generation benchmark: optional original AR to EOS.
+"""
+import ctypes
+import argparse
+import hashlib
+import json
+import os
+from pathlib import Path
+import sys
+import threading
+import time
+
+
+def ar_oracle(model, tokenizer, ids, root, set_phase):
+    """Actual _prefill and generate_ar, natural EOS, no tools or downloads."""
+    import mlx.core as mx
+    import numpy as np
+    from mtplx.runtime import MTPLXRuntime
+    from mtplx.sampling import SamplerConfig
+    from mtplx.generation import (_prefill, generate_ar, _final_logits_prefill_enabled,
+        _sustained_prefill_enabled, _prefill_chunk_size)
+    rt = MTPLXRuntime(model, tokenizer, root, False, None)
+    def fingerprint(a):
+        values = np.asarray(a.astype(mx.float32)).astype('<f4')
+        assert np.isfinite(values).all()
+        return dict(shape=list(values.shape), sha256=hashlib.sha256(values.tobytes()).hexdigest())
+    set_phase("prefill")
+    cache, logits, hidden, elapsed = _prefill(rt, ids, return_hidden=False)
+    print(json.dumps(dict(event="ar_prefill_reference", token_ids=ids, logits=fingerprint(logits),
+        hidden=hidden is not None, seconds=elapsed, chunks=rt.diagnostic_counters.get("prefill_chunks", 0),
+        chunk_size=_prefill_chunk_size() if _sustained_prefill_enabled() else None,
+        final_logits_only=_final_logits_prefill_enabled())), flush=True)
+    del cache, logits, hidden
+    sampler = SamplerConfig(temperature=.6, top_p=.95, top_k=20)
+    def progress(data):
+        set_phase("decode" if data["phase"] == "completed" else "prefill")
+        print(json.dumps(dict(event="ar_prefill_progress", **data)), file=sys.stderr, flush=True)
+    out = generate_ar(rt, ids, max_tokens=2**63-1, sampler=sampler, seed=0,
+        stop_token_ids=set(tokenizer.eos_token_ids), capture_final_state=True,
+        prefill_callback=progress,
+        token_callback=lambda ids: print(json.dumps(dict(event="ar_token", ids=ids)), file=sys.stderr, flush=True))
+    print(json.dumps(dict(event="ar_reference_end", tokens=out.tokens, text=out.text,
+        finish_reason=out.finish_reason, stops=sorted(tokenizer.eos_token_ids), seed=0,
+        temperature=.6, top_p=.95, top_k=20, verify_calls=out.stats.verify_calls,
+        elapsed_seconds=out.stats.elapsed_s, prompt_seconds=out.stats.prompt_eval_time_s,
+        final_logits=None if out.final_state is None else fingerprint(out.final_state.final_logits))), flush=True)
+
+
+def mtp_prefill_oracle(model, tokenizer, ids, root, set_phase):
+    """Original streaming history prefill and the first actual draft, no generation claim."""
+    import mlx.core as mx
+    import numpy as np
+    from mtplx.runtime import MTPLXRuntime
+    from mtplx.mtp_patch import MTPContract
+    from mtplx.generation import _prefill_committed_mtp_history_streaming, _prefill_chunk_size, _final_logits_prefill_enabled
+    rt = MTPLXRuntime(model, tokenizer, root, True, MTPContract())
+    def fingerprint(a):
+        values = np.asarray(a.astype(mx.float32)).astype('<f4')
+        assert np.isfinite(values).all()
+        return dict(shape=list(values.shape), sha256=hashlib.sha256(values.tobytes()).hexdigest())
+    for window in (None, 3):
+        set_phase("prefill")
+        cache, logits, hidden, history, target_s, history_s, base = _prefill_committed_mtp_history_streaming(
+            rt, ids, history_window_tokens=window,
+            chunk_callback=lambda data: print(json.dumps(dict(event="mtp_prefill_progress", **data)), file=sys.stderr, flush=True))
+        receipt = dict(event="mtp_prefill_reference", token_ids=ids, window=window,
+            chunk_size=_prefill_chunk_size(), final_only=_final_logits_prefill_enabled(),
+            logits=fingerprint(logits), hidden=fingerprint(hidden), offset=history[0].offset,
+            base=base, target_seconds=target_s, history_seconds=history_s)
+        set_phase("decode")
+        token = int(mx.argmax(logits[0]).item())
+        draft_logits, draft_hidden = model.mtp_forward(hidden, mx.array([[token]]), mtp_cache=history, return_hidden=True)
+        mx.eval(draft_logits, draft_hidden)
+        receipt.update(token=token, draft_logits=fingerprint(draft_logits), draft_hidden=fingerprint(draft_hidden), draft_offset=history[0].offset)
+        print(json.dumps(receipt), flush=True)
+        del cache, logits, hidden, history, draft_logits, draft_hidden
+    print(json.dumps(dict(event="mtp_prefill_reference_end")), flush=True)
+
+
+def mtp_generation_oracle(model, tokenizer, ids, root, set_phase, restored_suffixes=None, snapshot_restore=False):
+    """Actual generate_mtpk to natural EOS; preserves standard family p/q policy."""
+    import mlx.core as mx
+    import numpy as np
+    from mtplx.runtime import MTPLXRuntime
+    from mtplx.mtp_patch import MTPContract
+    from mtplx.sampling import SamplerConfig
+    from mtplx.generation import generate_mtpk, _prefill_chunk_size, _final_logits_prefill_enabled
+    rt = MTPLXRuntime(model, tokenizer, root, True, MTPContract())
+    def fingerprint(a):
+        values = np.asarray(a.astype(mx.float32)).astype('<f4')
+        assert np.isfinite(values).all()
+        return dict(shape=list(values.shape), sha256=hashlib.sha256(values.tobytes()).hexdigest())
+    sampler = SamplerConfig(temperature=.6, top_p=.95, top_k=20)
+    def progress(data):
+        set_phase("decode" if data["phase"] == "completed" else "prefill")
+        print(json.dumps(dict(event="mtp_prefill_progress", **data)), file=sys.stderr, flush=True)
+    print(json.dumps(dict(event="mtp_generation_input", token_ids=ids, temperature=.6, top_p=.95, top_k=20,
+        seed=0, depth=3, stops=sorted(tokenizer.eos_token_ids), history="committed",
+        restored_suffix_cases=3 if restored_suffixes is not None else 0,
+        chunk_size=_prefill_chunk_size(), final_only=_final_logits_prefill_enabled())), flush=True)
+    out = generate_mtpk(rt, ids, max_tokens=2**63-1, sampler=sampler, speculative_depth=3,
+        seed=0, stop_token_ids=set(tokenizer.eos_token_ids), mtp_history_policy="committed",
+        verify_strategy="batched", capture_final_state=True, prefill_callback=progress,
+        token_callback=lambda ids: print(json.dumps(dict(event="mtp_tokens", ids=ids)), file=sys.stderr, flush=True))
+    cycles = [dict(primary=e["primary"], primary_already_emitted=e["primary_already_emitted"],
+        drafts=[d["token"] for d in e["drafts"]], accepted=e["accepted_depths"],
+        correction=next((d["correction"] for d in e["drafts"] if d.get("accepted") is False), None),
+        bonus=e.get("bonus_token"), capture_committed=str(e.get("capture_repair", "")).startswith("captured_prefix"),
+        **({"context_copy": {k: v for k, v in e["context_copy"].items() if k != "time_s"}}
+           if "context_copy" in e else {}))
+        for e in out.stats.events if "primary" in e]
+    final = out.final_state
+    assert final is not None
+    print(json.dumps(dict(event="mtp_generation_end", tokens=out.tokens, text=out.text, cycles=cycles,
+        raw_cycles=out.stats.events,
+        finish_reason=out.finish_reason, safe_to_commit=final.safe_to_commit,
+        logits=fingerprint(final.final_logits), hidden=fingerprint(final.final_hidden),
+        history_offset=final.final_committed_mtp_cache[0].offset,
+        verify_calls=out.stats.verify_calls, accepted=out.stats.accepted_drafts, drafted=out.stats.drafted_tokens,
+        elapsed_seconds=out.stats.elapsed_s, prefill_seconds=out.stats.prompt_eval_time_s)), flush=True)
+    if restored_suffixes is not None:
+        from types import SimpleNamespace
+        from mtplx.generation import _prefill_restored_prompt_suffix, _small_suffix_fused_max, _runtime_counter_snapshot, _runtime_counter_delta
+        from mtplx.cache_state import snapshot_cache, snapshot_cache_lazy_hybrid, restore_cache
+        def snapshot_digest(snapshot):
+            h = hashlib.sha256()
+            for state in snapshot.states:
+                h.update(b"[")
+                for a in state:
+                    if a is None:
+                        h.update(b"N")
+                    else:
+                        tag = {mx.bfloat16:"BF16", mx.float32:"F32", mx.int32:"I32", mx.int64:"I64", mx.uint32:"U32", mx.float16:"F16"}[a.dtype]
+                        h.update(tag.encode())
+                        h.update(json.dumps(list(a.shape), separators=(",", ":")).encode())
+                        h.update(np.asarray(a.astype(mx.float32)).astype('<f4').tobytes())
+                h.update(b"]")
+            return h.hexdigest()
+        restored = SimpleNamespace(cache=final.final_trunk_cache,
+            mtp_history_cache=final.final_committed_mtp_cache, hidden=final.final_hidden)
+        cached = len(ids) + len(out.tokens)
+        # Numerical warm-extension cases on the same live final cache, not a
+        # substitute for the full session-bank/ongoing-chat acceptance run.
+        texts = ["Tell me a complete short story about a lighthouse keeper. Do not ask questions.",
+            "Give a summary of the following text:\n\n" + Path(restored_suffixes).read_text(),
+            "Write a Python function is_prime(n: int) -> bool, followed by five assert examples. No tools."]
+        for step, text in enumerate(texts):
+            suffix = tokenizer.encode("\n<|im_start|>user\n" + text + "<|im_end|>\n<|im_start|>assistant\n<think>\n")
+            chunks = []
+            set_phase("prefill")
+            saved = None
+            if snapshot_restore:
+                lazy_kv = step != 1
+                capture = snapshot_cache_lazy_hybrid if lazy_kv else snapshot_cache
+                trunk_snapshot, head_snapshot = capture(restored.cache), capture(restored.mtp_history_cache)
+                saved = dict(lazy_kv=lazy_kv, trunk=snapshot_digest(trunk_snapshot), mtp=snapshot_digest(head_snapshot))
+                restored.cache, restored.mtp_history_cache = model.make_cache(), model.make_mtp_cache()
+                restore_cache(restored.cache, trunk_snapshot, clone_states=not lazy_kv)
+                restore_cache(restored.mtp_history_cache, head_snapshot, clone_states=not lazy_kv)
+                print(json.dumps(dict(event="snapshot_restored", step=step, **saved)), file=sys.stderr, flush=True)
+            def chunk(data):
+                chunks.append({k: data[k] for k in ("tokens_done", "tokens_total", "cached_tokens", "new_prefill_tokens", "chunk_size")})
+                print(json.dumps(dict(event="restored_suffix_progress", **data)), file=sys.stderr, flush=True)
+            counts_before = _runtime_counter_snapshot(rt)
+            logits, hidden, target_s, history_s = _prefill_restored_prompt_suffix(rt, restored, suffix,
+                base_hidden_variant="post_norm", mtp_hidden_variant="post_norm", mtp_history_policy="committed",
+                chunk_callback=chunk, tokens_total=cached + len(suffix), cached_tokens=cached)
+            if saved is not None:
+                saved.update(trunk_after=snapshot_digest(trunk_snapshot), mtp_after=snapshot_digest(head_snapshot))
+                assert saved["trunk"] == saved["trunk_after"] and saved["mtp"] == saved["mtp_after"], "snapshot poisoned by continuation"
+            print(json.dumps(dict(event="restored_suffix_reference", step=step, token_ids=suffix, cached_tokens=cached,
+                fused_max=_small_suffix_fused_max(), chunk_size=_prefill_chunk_size(), final_only=_final_logits_prefill_enabled(),
+                logits=fingerprint(logits), hidden=fingerprint(hidden), history_offset=restored.mtp_history_cache[0].offset,
+                chunks=chunks, snapshot=saved, runtime_counts=_runtime_counter_delta(rt, counts_before), target_seconds=target_s, history_seconds=history_s)), flush=True)
+            restored.hidden = hidden
+            cached += len(suffix)
+        print(json.dumps(dict(event="restored_suffix_reference_end", cases=len(texts))), flush=True)
+
+
+def source_chat(records, mtp, canary):
+    """Accept only the explicit, complete UI-compatible comparison contract."""
+    starts = [r for r in records if r.get("event") == "start"]
+    results = [r for r in records if r.get("event") == "result"]
+    warmups = [r for r in records if r.get("event") == "warmup_result"]
+    if len(starts) != 1:
+        raise ValueError("source requires exactly one start receipt")
+    start = starts[0]
+    settings = start["settings"]
+    if not (start["model"] == "qwen3.8-flash-next" and start["plain_chat"]
+            and not start["system_prompt"] and start["canary"] == canary
+            and settings["reasoning"] == "low" and settings["power_percent"] == 100
+            and settings["prefill_chunk"] == 2048 and settings["min_p"] == 0
+            and not settings["quality"] and settings["seed"] is not None
+            and settings["acceleration"]["kind"] == "mtp"
+            and settings["acceleration"]["enabled"] == mtp
+            and len(start["prompts"]) == 3
+            and [r["turn"] for r in results] == [1, 2, 3]
+            and all(r["ok"] and r["finish_reason"] == "stop" for r in results)
+            and start["warmup"]["enabled"] and len(warmups) == 1 and warmups[0]["ok"]):
+        raise ValueError("source must be a complete warm Power-100 Low Qwen plain chat with matching MTP/canary and supported settings")
+    return start, results, warmups[0]
+
+
+def token_report(event, mode, **metadata):
+    """Rate-limit IO, not inference; only newly emitted tokens mean progress."""
+    total, last = 0, None
+    def report(tokens):
+        nonlocal total, last
+        if not tokens:
+            return
+        total += len(tokens)
+        now = time.monotonic()
+        if mode == "full":
+            print(json.dumps(dict(event=event, ids=tokens, **metadata)), file=sys.stderr, flush=True)
+        elif last is None or now - last >= 1:
+            print(json.dumps(dict(event="chat_decode_progress", scope=event,
+                tokens=total, **metadata)), file=sys.stderr, flush=True)
+            last = now
+    return report
+
+
+def chat_oracle(model, tokenizer, ids, root, set_phase, mtp, source=None, token_events="full"):
+    """Three complete turns through the real SessionBank and generation entrypoints."""
+    import mlx.core as mx
+    import numpy as np
+    from mtplx.runtime import MTPLXRuntime
+    from mtplx.mtp_patch import MTPContract
+    from mtplx.sampling import SamplerConfig
+    from mtplx.session_bank import SessionBank, _lazy_snapshot_enabled
+    from mtplx.cache_state import snapshot_cache
+    from mtplx.generation import (generate_ar, generate_mtpk, _prefill_chunk_size, _sustained_prefill_enabled,
+        _final_logits_prefill_enabled, _small_suffix_fused_max,
+        _store_on_prefill_env_enabled, _store_on_prefill_min_suffix,
+        _gdn_boundary_capture_enabled, _gdn_boundary_tail_interval, _gdn_boundary_max_count,
+        restore_or_prefill_prompt_state, _resolve_runtime_base_hidden_variant)
+    rt = MTPLXRuntime(model, tokenizer, root, mtp, MTPContract() if mtp else None)
+    bank = SessionBank()
+    settings = source[0]["settings"] if source else dict(temperature=.6, top_p=.95, top_k=20, seed=0)
+    sampler = SamplerConfig(temperature=settings["temperature"], top_p=settings["top_p"], top_k=settings["top_k"])
+    seed = settings["seed"]
+    print(json.dumps(dict(event="chat_config", mtp=mtp, turns=3, seed=seed, token_events=token_events,
+        temperature=sampler.temperature, top_p=sampler.top_p, top_k=sampler.top_k,
+        warmup=source[0]["warmup"] if source else None, depth=3, stops=sorted(tokenizer.eos_token_ids),
+        lazy_kv=_lazy_snapshot_enabled(), chunk_size=_prefill_chunk_size(), sustained_prefill=_sustained_prefill_enabled(),
+        final_only=_final_logits_prefill_enabled(), fused_max=_small_suffix_fused_max(),
+        store_on_prefill=_store_on_prefill_env_enabled(), store_min_suffix=_store_on_prefill_min_suffix(),
+        boundary_capture=_gdn_boundary_capture_enabled(), boundary_tail=_gdn_boundary_tail_interval(), boundary_max=_gdn_boundary_max_count())), flush=True)
+    def fingerprint(a):
+        if a is None:
+            return None
+        values = np.asarray(a.astype(mx.float32)).astype('<f4')
+        assert np.isfinite(values).all()
+        return dict(shape=list(values.shape), sha256=hashlib.sha256(values.tobytes()).hexdigest())
+    def progress(data):
+        set_phase("decode" if data["phase"] == "completed" else "prefill")
+        print(json.dumps(dict(event="chat_prefill_progress", **data)), file=sys.stderr, flush=True)
+    followups = [
+        "Tell me a complete short story about a lighthouse keeper. Do not ask questions.",
+        "Write a Python function is_prime(n: int) -> bool, followed by five assert examples. No tools.",
+    ]
+    generate = generate_mtpk if mtp else generate_ar
+    mode = dict(speculative_depth=3, mtp_history_policy="committed", verify_strategy="batched") if mtp else {}
+    if source:
+        from mtplx.server.openai import ChatMessage, _encode_messages_uncached
+        start, expected_turns, expected_warmup = source
+        followups = start["prompts"][1:]
+        # Derive the UI's system-only bootstrap from the ORIGINAL template,
+        # without copying the Low instruction or hard-coding its token count.
+        rendered = tokenizer.decode(ids, skip_special_tokens=False)
+        prefix = rendered.split("<|im_start|>user\n", 1)[0]
+        bootstrap_ids = tokenizer.encode(prefix)
+        assert bootstrap_ids and list(ids[:len(bootstrap_ids)]) == bootstrap_ids
+        state = restore_or_prefill_prompt_state(rt, bootstrap_ids,
+            base_hidden_variant="post_norm" if mtp else None,
+            mtp_history_policy="committed" if mtp else "cycle",
+            capture_hidden=mtp, session_bank=bank,
+            store_prefix_snapshot=False, prefill_callback=progress)
+        head = snapshot_cache(state.committed_mtp_cache) if state.committed_mtp_cache is not None else None
+        assert bank.put(runtime=rt, token_ids=bootstrap_ids, cache=state.trunk_cache,
+            logits=state.logits, hidden=state.hidden,
+            hidden_variant=_resolve_runtime_base_hidden_variant(rt, None),
+            session_id="oracle-bootstrap", mtp_history_policy="committed" if mtp else "cycle",
+            mtp_history_snapshot=head, snapshot_epoch=len(bootstrap_ids),
+            mtp_snapshot_epoch=len(bootstrap_ids) if head is not None else None,
+            gdn_boundaries=state.gdn_boundaries) is not None
+        del state, head
+        warm_ids = _encode_messages_uncached(tokenizer,
+            [ChatMessage(role="user", content=start["warmup"]["prompt"])],
+            enable_thinking=True, reasoning_effort="low", preserve_reasoning_history=True, tools=None)
+        warm = generate(rt, warm_ids, max_tokens=start["warmup"]["max_generated_tokens"],
+            sampler=sampler, seed=seed, stop_token_ids=set(tokenizer.eos_token_ids),
+            session_bank=bank, session_id="oracle-warmup", capture_final_state=True,
+            prefill_callback=progress,
+            token_callback=token_report("chat_warmup_tokens", token_events), **mode)
+        visible = sum(t not in tokenizer.eos_token_ids for t in warm.tokens)
+        print(json.dumps(dict(event="chat_warmup", tokens=warm.tokens, text=warm.text,
+            finish_reason=warm.finish_reason, bootstrap_ids=bootstrap_ids,
+            cached_tokens=warm.stats.cached_tokens)), flush=True)
+        assert warm.stats.cached_tokens == len(bootstrap_ids), "warmup must reuse the UI bootstrap"
+        assert warm.finish_reason == expected_warmup["finish_reason"]
+        assert visible == expected_warmup["completion_tokens"], "warmup token count differs"
+        del warm
+        bank.clear(session_id="oracle-warmup")
+    for step in range(3):
+        prompt_ids = list(ids)
+        limit = min(settings["max_generated_tokens"], settings["context_tokens"] - len(ids)) if source else 2**63-1
+        assert limit > 0, "source prompt exhausts the context"
+        kwargs = dict(max_tokens=limit, sampler=sampler, seed=seed,
+            stop_token_ids=set(tokenizer.eos_token_ids), session_bank=bank, session_id="oracle-chat",
+            capture_final_state=True, prefill_callback=progress,
+            token_callback=token_report("chat_tokens", token_events, step=step))
+        out = generate_mtpk(rt, ids, speculative_depth=3, mtp_history_policy="committed",
+            verify_strategy="batched", **kwargs) if mtp else generate_ar(rt, ids, **kwargs)
+        final = out.final_state
+        assert out.finish_reason == "stop" and final is not None and final.safe_to_commit
+        cycles = [dict(primary=e["primary"], primary_already_emitted=e["primary_already_emitted"],
+            drafts=[d["token"] for d in e["drafts"]], accepted=e["accepted_depths"],
+            correction=next((d["correction"] for d in e["drafts"] if d.get("accepted") is False), None),
+            bonus=e.get("bonus_token"), capture_committed=str(e.get("capture_repair", "")).startswith("captured_prefix"),
+            **({"context_copy": {k:v for k,v in e["context_copy"].items() if k != "time_s"}} if "context_copy" in e else {}))
+            for e in out.stats.events if "primary" in e]
+        ids = prompt_ids + list(out.tokens)
+        # Same final-state payload and head snapshot rule as openai's commit.
+        head = snapshot_cache(final.final_committed_mtp_cache) if final.final_committed_mtp_cache is not None else None
+        entry = bank.put(runtime=rt, token_ids=ids, cache=final.final_trunk_cache,
+            logits=final.final_logits, hidden=final.final_hidden, hidden_variant="post_norm",
+            session_id="oracle-chat", mtp_history_policy="committed" if mtp else "cycle",
+            mtp_history_snapshot=head, snapshot_epoch=len(ids),
+            mtp_snapshot_epoch=len(ids) if head is not None else None)
+        assert entry is not None
+        print(json.dumps(dict(event="chat_turn", step=step, prompt_ids=prompt_ids, tokens=out.tokens,
+            text=out.text, finish_reason=out.finish_reason, safe_to_commit=final.safe_to_commit,
+            cached_tokens=out.stats.cached_tokens, boundaries=[b[0] for b in entry.gdn_boundaries],
+            logits=fingerprint(final.final_logits),
+            hidden=fingerprint(final.final_hidden), cycles=cycles,
+            history_offset=final.final_committed_mtp_cache[0].offset if head is not None else None,
+            prefill_seconds=out.stats.prompt_eval_time_s, elapsed_seconds=out.stats.elapsed_s,
+            prefill_scope="engine_target_and_mtp_history_excluding_restore_and_checkpoint",
+            mtp_history_seconds=out.stats.prompt_mtp_history_time_s,
+            decode_seconds=out.stats.decode_elapsed_s, decode_tps=out.stats.decode_tok_s,
+            prompt_state_seconds=out.stats.prompt_state_total_time_s,
+            restore_seconds=out.stats.cache_restore_time_s)), flush=True)
+        if source:
+            expected = expected_turns[step]
+            checks = dict(
+                prompt_tokens=len(prompt_ids) == expected["tokens"]["prompt"],
+                cached_tokens=out.stats.cached_tokens == expected["tokens"]["cached"],
+                completion_tokens=sum(t not in tokenizer.eos_token_ids for t in out.tokens) == expected["tokens"]["completion"],
+                text=out.text == (expected["reasoning"] or "") + "</think>" + expected["content"])
+            print(json.dumps(dict(event="chat_source_match", step=step, checks=checks)), flush=True)
+            assert all(checks.values()), f"source mismatch at turn {step}: {checks}"
+        if step < len(followups):
+            ids += tokenizer.encode("\n<|im_start|>user\n" + followups[step] + "<|im_end|>\n<|im_start|>assistant\n<think>\n")
+    print(json.dumps(dict(event="chat_end", turns=3)), flush=True)
+
+
+def main(root, bridge, decode=False, server_defaults=False, mtp=False, ar=False, readme=None, mtp_prefill=False, mtp_generate=False, restored_suffixes=None, snapshot_restore=False, chat=False, source=None, canary=True, token_events="full"):
+    class Sample(ctypes.Structure):
+        _fields_ = [("scheduled_seconds", ctypes.c_double), ("completed_seconds", ctypes.c_double)]
+
+    lib = ctypes.CDLL(str(Path(bridge).resolve()))
+    probe = lib.ds4_gpu_canary_probe
+    probe.argtypes = [ctypes.POINTER(Sample)]
+    probe.restype = ctypes.c_int
+    phase = "startup"
+    pending = None
+    stop = threading.Event()
+    samples = []
+
+    def watch():
+        while not stop.wait(0.01):
+            if pending is not None and time.monotonic() - pending >= 2:
+                print(json.dumps(dict(event="test_canary_abort", phase=phase, reason="GPU probe exceeds 2s")), file=sys.stderr, flush=True)
+                os._exit(86)
+
+    def probes():
+        nonlocal pending
+        while not stop.is_set():
+            at = time.monotonic()
+            probe_phase = phase
+            pending = at
+            sample = Sample()
+            ok = probe(ctypes.byref(sample))
+            pending = None
+            samples.append(dict(phase=probe_phase, scheduled_ms=sample.scheduled_seconds * 1000, completed_ms=sample.completed_seconds * 1000))
+            if not ok or sample.completed_seconds >= 2:
+                print(json.dumps(dict(event="test_canary_abort", sample=samples[-1], ok=bool(ok))), file=sys.stderr, flush=True)
+                os._exit(86)
+            stop.wait(max(0, 0.1 - (time.monotonic() - at)))
+
+    watcher = threading.Thread(target=watch, daemon=True)
+    prober = threading.Thread(target=probes, daemon=True)
+    if canary:
+        watcher.start()
+        prober.start()
+    try:
+        os.environ["HF_HUB_OFFLINE"] = "1"
+        os.environ["TRANSFORMERS_OFFLINE"] = "1"
+        print("reference: importing pinned MTPLX", file=sys.stderr, flush=True)
+        import mlx.core as mx
+        import numpy as np
+        from mlx_lm.utils import load_model
+        from mtplx.models import qwen4_exp as qwen
+        from mtplx.runtime import _load_tokenizer_resilient
+        from mtplx.server.openai import ChatMessage, _encode_messages_uncached
+        assert mx.__version__ == "0.32.2"
+        root = Path(root)
+        if server_defaults:
+            from mtplx.server.openai import _server_runtime_env_overrides
+            overrides = _server_runtime_env_overrides(
+                argparse.Namespace(model=str(root), generation_mode="mtp" if mtp or mtp_prefill or mtp_generate else "ar"), None)
+            os.environ.update(overrides)
+            print(json.dumps(dict(event="reference_environment", overrides=overrides)), flush=True)
+        if source:
+            # Match policy::turn_options, as mtplx-policy-reference.py does:
+            # bounded streaming prefill, request-local chunks, final logits.
+            # This also governs bootstrap and restored suffixes, not just the
+            # reported chunk-size number or final-logits switch.
+            os.environ["MTPLX_SUSTAINED_PREFILL"] = "1"
+            print(json.dumps(dict(event="reference_request_policy", sustained_prefill=True, final_logits_only=True,
+                prefill_chunk=source[0]["settings"]["prefill_chunk"])), flush=True)
+        tokenizer = _load_tokenizer_resilient(root, json.loads((root / "config.json").read_text()))
+        prompt = "ping" if readme is None else "Give a summary of the following text:\n\n" + Path(readme).read_text()
+        if source:
+            assert prompt == source[0]["prompts"][0], "README prompt differs from the source run"
+        print(json.dumps(dict(event="reference_prompt", text=prompt,
+            sha256=hashlib.sha256(prompt.encode()).hexdigest())), flush=True)
+        ids = _encode_messages_uncached(tokenizer, [ChatMessage(role="user", content=prompt)],
+            enable_thinking=True, reasoning_effort="low", preserve_reasoning_history=True, tools=None)
+        phase = "loading"
+        print(json.dumps(dict(event="reference_loading", token_ids=ids)), file=sys.stderr, flush=True)
+        started = time.monotonic()
+        model, _ = load_model(root, lazy=False, strict=True, get_model_classes=lambda **_: (qwen.Model, qwen.ModelArgs))
+        model.post_weight_load(root)
+        if mtp or mtp_prefill or mtp_generate:
+            assert model.attach_mtp(root), "installed MTP head missing"
+        loaded = time.monotonic()
+        if ar or mtp_prefill or mtp_generate:
+            def set_phase(value):
+                nonlocal phase
+                phase = value
+            oracle = mtp_generation_oracle if mtp_generate else mtp_prefill_oracle if mtp_prefill else ar_oracle
+            if chat:
+                from mtplx.generation import prefill_chunk_size_override
+                with prefill_chunk_size_override(source[0]["settings"]["prefill_chunk"] if source else None):
+                    chat_oracle(model, tokenizer, ids, root, set_phase, mtp_generate, source, token_events)
+            elif mtp_generate:
+                oracle(model, tokenizer, ids, root, set_phase, restored_suffixes, snapshot_restore)
+            else:
+                oracle(model, tokenizer, ids, root, set_phase)
+            phase = "finishing"
+            del model
+            return
+        print("reference: trunk loaded, starting prefill", file=sys.stderr, flush=True)
+        phase = "prefill"
+        cache = model.make_cache()
+        prefill_ids = ids[:-1] if mtp else ids
+        out = model(mx.array([prefill_ids], dtype=mx.int32), cache=cache)
+        mx.eval(out)
+        finished = time.monotonic()
+        values = np.asarray(out.astype(mx.float32)).astype('<f4')
+        assert np.isfinite(values).all()
+        print(json.dumps(dict(event="forward_reference", token_ids=prefill_ids, shape=list(values.shape),
+            row_sha256=[hashlib.sha256(row.tobytes()).hexdigest() for row in values[0]],
+            argmax=[int(row.argmax()) for row in values[0]], load_seconds=loaded-started,
+            forward_seconds=finished-loaded, mtp=mtp, prompt=prompt, power_percent=100)),flush=True)
+        if mtp:
+            def fingerprint(array):
+                values = np.asarray(array.astype(mx.float32)).astype('<f4')
+                assert np.isfinite(values).all()
+                return dict(shape=list(values.shape), sha256=hashlib.sha256(values.tobytes()).hexdigest())
+
+            def record(op, logits=None, hidden=None, **fields):
+                mx.eval(*[a for a in (logits, hidden) if a is not None])
+                print(json.dumps(dict(event="mtp_reference", op=op,
+                    logits=None if logits is None else fingerprint(logits),
+                    hidden=None if hidden is None else fingerprint(hidden), **fields)), flush=True)
+
+            body_hidden = model.language_model.model._last_widened
+            out, tail_hidden = model(mx.array([[ids[-1]]]), cache=cache, return_hidden=True)
+            record("trunk", out, tail_hidden, token_ids=[ids[-1]])
+            phase = "decode"
+            head_cache = model.make_mtp_cache()
+            history = model.mtp_update_cache(body_hidden, mx.array([ids[1:]]), mtp_cache=head_cache)
+            record("history", hidden=history, token_ids=ids[1:], offset=head_cache[0].offset)
+            base_offset = head_cache[0].offset
+            first = int(mx.argmax(out[0, -1]).item())
+            token, draft_hidden = first, tail_hidden
+            drafts = []
+            for depth in range(3):
+                out, draft_hidden = model.mtp_forward(draft_hidden, mx.array([[token]]),
+                    mtp_cache=head_cache, return_hidden=True)
+                record("draft", out, draft_hidden, token_ids=[token], depth=depth,
+                    offset=head_cache[0].offset)
+                token = int(mx.argmax(out[0, -1]).item())
+                drafts.append(token)
+            window = [first, *drafts]
+            snapshots = [None if c.is_trimmable() else list(c.state) for c in cache]
+            with model.verify_capture_scope():
+                out, verified_hidden = model(mx.array([window]), cache=cache, return_hidden=True)
+            record("verify", out, verified_hidden, token_ids=window)
+            choices = np.asarray(mx.argmax(out[0], axis=-1)).tolist()
+            keep = 1
+            for draft, target in zip(drafts, choices):
+                if draft != target:
+                    break
+                keep += 1
+            committed = model.commit_verified_window(cache, snapshots,
+                keep_tokens=keep, verified_tokens=len(window))
+            assert committed
+            # generate_mtpk retains the first authoritative draft row and
+            # replaces only recursively predicted rows after it.
+            head_cache[0].trim(head_cache[0].offset - (base_offset + 1))
+            record("commit", keep=keep, verified=len(window), offset=head_cache[0].offset)
+            history = None
+            if keep > 1:
+                history = model.mtp_update_cache(verified_hidden[:, :keep-1], mx.array([window[1:keep]]), mtp_cache=head_cache)
+            record("history_commit", hidden=history, token_ids=window[1:keep], offset=head_cache[0].offset)
+            token = int(choices[keep-1])
+            out, hidden = model(mx.array([[token]]), cache=cache, return_hidden=True)
+            record("continuation", out, hidden, token_ids=[token])
+            next_token = int(mx.argmax(out[0, -1]).item())
+            out, hidden = model.mtp_forward(hidden, mx.array([[next_token]]), mtp_cache=head_cache, return_hidden=True)
+            record("draft_continuation", out, hidden, token_ids=[next_token], offset=head_cache[0].offset)
+            print(json.dumps(dict(event="mtp_reference_end", keep=keep, verified=len(window))), flush=True)
+        if decode:
+            # Numerical cache-continuation oracle, not the serving sampler or
+            # a throughput benchmark. Follow greedy tokens to natural EOS.
+            phase = "decode"
+            token = int(values[0, -1].argmax())
+            generated = []
+            while token not in tokenizer.eos_token_ids:
+                generated.append(token)
+                started = time.monotonic()
+                out = model(mx.array([[token]], dtype=mx.int32), cache=cache)
+                mx.eval(out)
+                elapsed = time.monotonic() - started
+                values = np.asarray(out.astype(mx.float32)).astype('<f4')
+                assert np.isfinite(values).all()
+                next_token = int(values[0, -1].argmax())
+                print(json.dumps(dict(event="decode_reference", step=len(generated), token_id=token,
+                    shape=list(values.shape), argmax=next_token,
+                    row_sha256=hashlib.sha256(values[0, -1].tobytes()).hexdigest(),
+                    forward_seconds=elapsed)), flush=True)
+                token = next_token
+            print(json.dumps(dict(event="decode_reference_end", eos_token=token,
+                eos_token_ids=sorted(tokenizer.eos_token_ids), token_ids=generated,
+                text=tokenizer.decode(generated), finish_reason="eos")), flush=True)
+        phase = "finishing"
+        del out, model, cache
+    finally:
+        stop.set()
+        if canary:
+            prober.join(timeout=2)
+            watcher.join(timeout=2)
+        print(json.dumps(dict(event="reference_canary_summary", enabled=canary, samples=samples,
+            incomplete=prober.is_alive())),file=sys.stderr,flush=True)
+        if prober.is_alive():
+            os._exit(86)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("root")
+    parser.add_argument("bridge")
+    parser.add_argument("--decode", action="store_true")
+    parser.add_argument("--server-defaults", action="store_true")
+    parser.add_argument("--mtp", action="store_true", help="Numerical draft/verify/commit round, not a generation benchmark")
+    parser.add_argument("--ar", action="store_true", help="Original cold prefill and sampled AR to natural EOS, with final-state capture")
+    parser.add_argument("--mtp-prefill", action="store_true", help="Original streaming committed-history prefill and first draft")
+    parser.add_argument("--mtp-generate", action="store_true", help="Original sampled generate_mtpk to natural EOS")
+    parser.add_argument("--readme", type=Path, help="Summarize this README instead of ping; freezes the exact prompt in the receipt")
+    parser.add_argument("--restored-suffixes", type=Path, help="After MTP EOS, compare live-cache suffix extension using short text, this README, then short Python instructions")
+    parser.add_argument("--snapshot-restore", action="store_true", help="Snapshot and restore fresh trunk/MTP containers before each suffix; verify retained snapshots after mutation")
+    parser.add_argument("--chat", action="store_true", help="Full README-summary/story/Python chat through the original SessionBank; needs --readme and --ar or --mtp-generate")
+    parser.add_argument("--source-run", type=Path, help="Match a completed warm model-eval chat: settings, warmup, bootstrap, output and cache frontiers")
+    parser.add_argument("--canary", choices=("on", "off"), default="on")
+    parser.add_argument("--token-events", choices=("full", "progress"), default="full",
+        help="Chat token trace or at most one real decode-progress event per second; full output is always retained")
+    args = parser.parse_args()
+    assert sum((args.mtp, args.decode, args.ar, args.mtp_prefill, args.mtp_generate)) <= 1, "choose exactly one execution oracle"
+    assert not args.restored_suffixes or args.mtp_generate, "restored suffix oracle requires --mtp-generate"
+    assert not args.snapshot_restore or args.restored_suffixes, "snapshot restore requires suffix cases"
+    assert not args.chat or (args.readme and (args.ar or args.mtp_generate) and not args.restored_suffixes), "chat requires README and a generation lane, without suffix-only cases"
+    assert not args.source_run or (args.chat and args.server_defaults), "source matching requires chat and server defaults"
+    assert args.token_events == "full" or args.chat, "progress-only token events require chat"
+    source = source_chat([json.loads(line) for line in args.source_run.read_text().splitlines()
+        if line.startswith("{")], args.mtp_generate, args.canary == "on") if args.source_run else None
+    main(args.root, args.bridge, args.decode, args.server_defaults, args.mtp, args.ar, args.readme, args.mtp_prefill, args.mtp_generate, args.restored_suffixes, args.snapshot_restore, args.chat, source, args.canary == "on", args.token_events)

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()

tools/mtplx-host-import-reference.py

@@ -0,0 +1,63 @@
+"""Tiny host-import/lifetime receipts from the pinned MTPLX runtime; no models."""
+import gc
+import json
+import mmap
+import os
+import weakref
+
+import mlx.core as mx
+import numpy as np
+
+
+def main():
+    assert mx.__version__ == "0.32.2"
+    page = os.sysconf("SC_PAGE_SIZE")
+    mx.set_cache_limit(4 * page)
+    for name, offset, size in (
+        ("aligned", 0, page),
+        ("unaligned", 1, page),
+        ("short", 0, 17),
+        ("empty", 0, 0),
+    ):
+        mx.clear_cache()
+        baseline = mx.get_active_memory()
+        owner = mmap.mmap(-1, 2 * page)
+        source = np.ndarray((size,), dtype=np.uint8, buffer=owner, offset=offset)
+        source[:] = 37
+        pointer = source.__array_interface__["data"][0]
+        weak = weakref.ref(source)
+        # The codec calls this constructor: array.cpp passes copy=true to
+        # create_array, even though nd_array_to_mlx has a different default.
+        value = mx.array(source)
+        mx.eval(value)
+        exported = np.asarray(value)
+        adopted = bool(size and pointer == exported.__array_interface__["data"][0])
+        assert (exported == 37).all()
+        del exported, source
+        gc.collect()
+        retained = weak() is not None
+        assert not adopted and not retained
+        active = mx.get_active_memory() - baseline
+        view = value[:]
+        mx.eval(view)
+        del value
+        gc.collect()
+        retained_by_view = weak() is not None
+        assert not retained_by_view
+        assert (np.asarray(view) == 37).all()
+        del view
+        gc.collect()
+        mx.synchronize()
+        released = weak() is None
+        assert released
+        assert mx.get_active_memory() == baseline
+        cached = mx.get_cache_memory()
+        owner.close()
+        print(json.dumps(dict(event="host_import", name=name, page=page, size=size,
+                              adopted=adopted, active=active, retained=retained,
+                              retained_by_view=retained_by_view, released=released,
+                              cached_after_release=cached)), flush=True)
+
+
+if __name__ == "__main__":
+    main()

tools/mtplx-jit-reference.py

@@ -0,0 +1,162 @@
+"""Observe MTPLX's generated Metal source in this model-free reference process.
+
+The local Objective-C hook forwards the original call unchanged and is restored
+before exit. It neither modifies the reference checkout nor writes artifacts.
+Run in the pinned MTPLX environment; stdout is a JSON source receipt.
+"""
+import ctypes as c
+import argparse
+import hashlib
+import json
+from pathlib import Path
+import re
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--check",action="store_true",help="compare with the pinned captured shader receipt")
+    parser.add_argument("--rank",type=int,choices=(2,3,5),default=2,
+                        help="compiled input rank; 3/5 cover sorted/unsorted SwitchGLU")
+    parser.add_argument("--operation",choices=("silu","swiglu","compute_g","qsa-update"),default="silu",
+                        help="actual compiled activation used by fused or separate projection modules")
+    parser.add_argument("--qsa-shape",type=int,nargs=3,default=(4,4096,1024),
+                        metavar=("ROWS","RAW_CAP","POOL_CAP"))
+    parser.add_argument("--qsa-header",action="store_true",help="also retain the unchanged scalar library header")
+    parser.add_argument("--qsa-mode",choices=("update_only","blocks","row_tokens","dense_mask","prefill_blocks"),default="update_only")
+    parser.add_argument("--qsa-score-budget",type=int,default=32*1024*1024)
+    args = parser.parse_args()
+    objc = c.CDLL("/usr/lib/libobjc.A.dylib")
+    metal = c.CDLL("/System/Library/Frameworks/Metal.framework/Metal")
+    metal.MTLCreateSystemDefaultDevice.restype = c.c_void_p
+    device = metal.MTLCreateSystemDefaultDevice()
+    assert device
+    objc.object_getClass.argtypes = [c.c_void_p]
+    objc.object_getClass.restype = c.c_void_p
+    objc.sel_registerName.argtypes = [c.c_char_p]
+    objc.sel_registerName.restype = c.c_void_p
+    objc.class_getInstanceMethod.argtypes = [c.c_void_p,c.c_void_p]
+    objc.class_getInstanceMethod.restype = c.c_void_p
+    objc.method_getImplementation.argtypes = [c.c_void_p]
+    objc.method_getImplementation.restype = c.c_void_p
+    objc.method_setImplementation.argtypes = [c.c_void_p,c.c_void_p]
+    objc.method_setImplementation.restype = c.c_void_p
+    selector = objc.sel_registerName(b"newLibraryWithSource:options:error:")
+    method = objc.class_getInstanceMethod(objc.object_getClass(device),selector)
+    assert method
+    signature = c.CFUNCTYPE(c.c_void_p,*([c.c_void_p]*5))
+    original_ptr = objc.method_getImplementation(method)
+    original = signature(original_ptr)
+    string = c.CFUNCTYPE(c.c_char_p,c.c_void_p,c.c_void_p)(("objc_msgSend",objc))
+    utf8 = objc.sel_registerName(b"UTF8String")
+    receipts = []
+    errors = []
+
+    @signature
+    def observe(receiver,sel,source,options,error):
+        try:
+            text = string(source,utf8).decode()
+            match = re.search(r'\[\[host_name\("[^"]+"\)\]\]\n\[\[kernel\]\] void ',text)
+            if match and "tmp_" in text[match.start():]:
+                kernels = text[match.start():]
+                if args.operation == "qsa-update":
+                    # All moving frontiers are contiguous int32[1] leaves.
+                    # Retain the exact invoked specialization, not unused ranks.
+                    kernels = kernels[:kernels.index("[[host_name",2)]
+                receipts.append({"source_sha256":hashlib.sha256(text.encode()).hexdigest(),
+                                 "kernels":kernels})
+                if args.operation == "qsa-update" and args.qsa_header and len(receipts)==1:
+                    receipts[-1]["header"] = text[:match.start()]
+            elif args.operation == "qsa-update" and "[[kernel]] void compute_dynamic_offset_" in text:
+                receipts.append({"source_sha256":hashlib.sha256(text.encode()).hexdigest(),
+                                 "kernels":text[text.index("[[kernel]] void compute_dynamic_offset_"):]})
+        except Exception as exc:
+            errors.append(str(exc))
+        return original(receiver,sel,source,options,error)
+
+    objc.method_setImplementation(method,c.cast(observe,c.c_void_p))
+    try:
+        import mlx.core as mx
+        import mlx.nn as nn
+        assert mx.__version__ == "0.32.2",mx.__version__
+        shape = {2:(70,1280),3:(70,1,1280),5:(1,7,10,1,1280)}[args.rank]
+        x = mx.arange(70*1280).astype(mx.bfloat16).reshape(shape)/128
+        mx.eval(x)
+        gate,up = mx.split(x,2,axis=-1)
+        if args.operation == "qsa-update":
+            from mtplx.kernels.qsa_indexer_compile import QSACompiledIndexerCore
+            norm = mx.ones((128,),dtype=mx.bfloat16)
+            freq = mx.arange(32,dtype=mx.float32)/128
+            core = QSACompiledIndexerCore(n_heads=4,kv_heads=1,head_dim=128,
+                block_topk=512,compress_ratio=4,q_norm_weight=norm,k_norm_weight=norm,
+                inv_freq=freq,rms_norm_eps=1e-6,selector_scratch_bytes=args.qsa_score_budget)
+            rows,raw_cap,pool_cap = args.qsa_shape
+            qk = mx.arange(rows*640).astype(mx.bfloat16).reshape(1,rows,640)/128
+            raw = mx.zeros((1,raw_cap,128),dtype=mx.bfloat16)
+            pooled = mx.zeros((1,pool_cap,128),dtype=mx.bfloat16)
+            pos = min(2051,raw_cap-rows)
+            total = pos+rows
+            logical = total//4
+            scalars = [mx.array([n],dtype=mx.int32) for n in (pos,total,logical,max(0,logical-(rows+3)//4))]
+            mx.eval(qk,raw,pooled,norm,freq,*scalars)
+            result = core.select_qk_rows(qk,raw,pooled,pos_start=scalars[0],
+                total_tokens=scalars[1],logical_blocks=scalars[2],pooled_len=scalars[3],mode=args.qsa_mode)
+            mx.eval(result.raw_keys,result.pooled,result.pooled_len,result.offset,result.selection)
+        elif args.operation == "compute_g":
+            from mlx_lm.models.gated_delta import compute_g
+            # Actual fused-projection A view, with per-head A_log/dt_bias broadcasts.
+            ashape = {2:(70,16480),3:(1,70,16480),5:(1,7,10,1,16480)}[args.rank]
+            a = (mx.arange(70*16480).astype(mx.bfloat16)/128).reshape(ashape)[...,-48:]
+            a_log = mx.arange(48).astype(mx.bfloat16)/128
+            dt_bias = -a_log
+            mx.eval(a,a_log,dt_bias)
+            mx.eval(compute_g(a_log,a,dt_bias))
+        elif args.operation == "swiglu":
+            from mlx_lm.models.activations import swiglu
+            gate,up = mx.contiguous(gate),mx.contiguous(up)
+            mx.eval(gate,up)
+            mx.eval(swiglu(gate,up))
+        else:
+            mx.eval(nn.silu(gate)*up)
+        assert not errors,errors
+        if args.operation == "qsa-update":
+            assert receipts,"No compiled QSA kernels captured"
+            if args.check:
+                root = Path(__file__).resolve().parents[1]
+                if args.qsa_header:
+                    assert receipts[0].pop("header")== (root/"metal/mtplx-qsa-compiled-header.metal").read_text(),"Compiled scalar header differs"
+                if tuple(args.qsa_shape)==(4,4096,1024):
+                    fixture = "mtplx-qsa-select-jit.json" if args.qsa_mode=="blocks" and args.qsa_score_budget==4096 else "mtplx-qsa-update-jit.json"
+                    expected = json.loads((root/"tests/fixtures"/fixture).read_text())
+                    assert receipts==expected,"Compiled QSA shaders differ from the pinned receipt"
+                else:
+                    bank = json.loads((root/"metal/mtplx-qsa-compiled-scalars.json").read_text())
+                    rows,raw,pool = args.qsa_shape
+                    new = (rows+3)//4
+                    limit = min(raw//4-new,pool-new)
+                    candidates = [*bank["clamp"].values(),bank["check_distinct"]]
+                    expected = next(v for v in candidates if (v["max_new"],v["max_start"])==(new,limit))
+                    assert receipts[0]=={k:expected[k] for k in ("source_sha256","kernels")},"Compiled clamp differs"
+                    assert receipts[1]==bank["multiply"],"Compiled multiply differs"
+                print("Actual compiled QSA update shaders match the pinned receipt",flush=True)
+            else:
+                print(json.dumps(receipts),flush=True)
+            return
+        assert len(receipts)==1,len(receipts)
+        if args.check:
+            expected = json.loads((Path(__file__).resolve().parents[1]/f"tests/fixtures/mtplx-{args.operation}-jit.json").read_text())
+            if args.rank==2:
+                assert receipts[0]==expected,"Generated activation differs from the pinned receipt"
+            else:
+                def normalize(source):
+                    name = re.search(r'host_name\("([^\"]+)_contiguous"',source).group(1)
+                    return source.replace(name,"COMPILED_ACTIVATION")
+                assert normalize(receipts[0]["kernels"])==normalize(expected["kernels"]),"Rank-dependent activation computation changed"
+            print(f"Actual rank-{args.rank} {args.operation} compiler source matches the pinned receipt (rank-specific symbol names normalized for 3/5)",flush=True)
+        else:
+            print(json.dumps(receipts[0]),flush=True)
+    finally:
+        objc.method_setImplementation(method,original_ptr)
+
+
+if __name__ == "__main__":
+    main()

tools/mtplx-kernel-source.py

@@ -0,0 +1,592 @@
+"""Extract pinned MTPLX Metal bodies without executing its Python model code.
+
+Print the deterministic Metal translation unit; --check compares the checked-in
+copy byte for byte. Only entry-point ABI/template declarations are supplied here.
+The bodies, including whitespace and comments, come unchanged from MTPLX.
+"""
+
+import argparse
+import ast
+import hashlib
+import json
+import re
+from pathlib import Path
+import subprocess
+
+REVISION = "e652d55e2652137a4abcf1312357abbf3eb9d692"
+GATED_DELTA_SHA256 = "79c8376a51c694b03e54d2f996ced6ea6c8c42868b8571529f97334db165a3e1"
+OUTPUT = Path(__file__).resolve().parents[1] / "metal/mtplx_qwen.metal"
+QSA_SELECT_OUTPUT = OUTPUT.with_name("mtplx-qsa-select.json")
+RUNTIME_REVISION = "1f8e74e3f12f31365464a6867c6579f0e9b29d85"
+
+# (module, source variable, entry point, inputs, outputs, builtins, extra template,
+#  (entry-point suffix, complete template arguments) specializations)
+KERNELS = [
+    ("hyper_connection", "_SOURCE", "hyper_read",
+     "T:x,gamma,wd,wu,wi", "T:mixed,inject",
+     "thread_position_in_grid,thread_position_in_threadgroup", "int HAS_INJECT",
+     [("mix_bf16", "bfloat, 0"), ("inj_bf16", "bfloat, 1")]),
+    ("hyper_connection_v3", "_SRC_R1", "hyper_v3_r1",
+     "T:x,wn;uint32_t:qw;T:qs,qb", "T:mix_out,inject_out;float:rms_out",
+     "thread_position_in_threadgroup,threadgroup_position_in_grid", "",
+     [("bf16", "bfloat")]),
+    ("hyper_connection_v3", "_SRC_R2", "hyper_v3_r2",
+     "T:x,wn,mixv;float:rms_in;uint32_t:qw;T:qs,qb", "T:y",
+     "thread_position_in_threadgroup,threadgroup_position_in_grid", "",
+     [("bf16", "bfloat")]),
+    ("gdn_conv_norm", "_SRC", "gdn_conv_norm",
+     "T:xnew,state,cw", "T:q_out,k_out,v_out,state_out",
+     "thread_position_in_threadgroup,threadgroup_position_in_grid", "",
+     [("bf16", "bfloat")]),
+    ("gdn_conv_norm", "_SRC_ROWS", "gdn_conv_norm_rows",
+     "T:xnew,state,cw", "T:q_out,k_out,v_out,state_out",
+     "thread_position_in_threadgroup,threadgroup_position_in_grid", "int S",
+     [(f"s{s}_bf16", f"bfloat, {s}") for s in range(2, 7)]),
+    ("gdn_step_fused", "_SRC", "gdn_step_fused",
+     "T:xnew,z_row,a_row,b_row,state,cw,A_log,dt_bias;StT:dstate;T:norm_w",
+     "T:y,state_out;StT:dstate_out",
+     "thread_position_in_threadgroup,threadgroup_position_in_grid", "typename StT",
+     [("f32_state_bf16", "bfloat, float")]),
+    ("gdn_out_fused", "_SRC", "gdn_out_fused",
+     "T:x;bfloat:z,wn;uint32_t:qw;bfloat:qs,qb", "T:y",
+     "thread_position_in_threadgroup,threadgroup_position_in_grid", "int GS_C",
+     [(f"gs{gs}_{tag}", f"{dtype}, {gs}")
+      for gs in (32, 64) for tag, dtype in (("bf16", "bfloat"), ("f32", "float"))]),
+    ("moe_glu_decode", "_SRC_A", "moe_glu_h",
+     "T:x;uint32_t:gw;T:gs,gb;uint32_t:experts", "T:h",
+     "thread_position_in_threadgroup,threadgroup_position_in_grid",
+     "int n_inter_c, int topk_c, int GS_GU",
+     [(f"g{gs}_bf16", f"bfloat, 640, 10, {gs}") for gs in (32,64)]),
+    ("moe_glu_decode", "_SRC_B", "moe_down_y",
+     "T:h;uint32_t:dw;T:ds,db;uint32_t:experts;float:rw", "T:y",
+     "thread_position_in_threadgroup,threadgroup_position_in_grid",
+     "int dmodel_c, int n_inter_c, int topk_c, int GS_DN",
+     [(f"g{gs}_bf16", f"bfloat, 2560, 640, 10, {gs}") for gs in (32,64)]),
+    ("moe_glu_decode", "_SRC_A_M", "moe_glu_h_m",
+     "T:x;uint32_t:gw;T:gs,gb;uint32_t:experts", "T:h",
+     "thread_position_in_threadgroup,threadgroup_position_in_grid",
+     "int n_inter_c, int topk_c, int m_rows_c, int GS_GU",
+     [(f"m{m}_g{gs}_bf16", f"bfloat, 640, 10, {m}, {gs}") for m in (2,3,4) for gs in (32,64)]),
+    ("moe_glu_decode", "_SRC_B_M", "moe_down_y_m",
+     "T:h;uint32_t:dw;T:ds,db;uint32_t:experts;float:rw", "T:y",
+     "thread_position_in_threadgroup,threadgroup_position_in_grid",
+     "int dmodel_c, int n_inter_c, int topk_c, int m_rows_c, int GS_DN",
+     [(f"m{m}_g{gs}_bf16", f"bfloat, 2560, 640, 10, {m}, {gs}") for m in (2,3,4) for gs in (32,64)]),
+]
+
+
+def parameters(groups, output=False):
+    return [(f"device {'const ' if not output else ''}{kind.strip()}*", name)
+            for group in groups.split(";") for kind, names in [group.split(":")]
+            for name in names.split(",")]
+
+
+def kernel_source(source, raw, body, function, args, builtins, templates, variants):
+    signature = [f"    {kind} {key} [[buffer({i})]]" for i, (kind, key) in enumerate(args)]
+    signature += [f"    {'uint' if key == 'thread_index_in_simdgroup' else 'uint3'} "
+                  f"{key} [[{key}]]" for key in builtins.split(",")]
+    output = [f"\n// Source: {source}\n"
+              f"// File SHA256: {hashlib.sha256(raw).hexdigest()}\n"
+              f"// Body SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
+              f"template <{templates}>\n"
+              f"kernel void {function}(\n" + ",\n".join(signature) + ") {" + body + "}\n"]
+    for suffix, argument in variants:
+        specialization = f"{function}<{argument}>"
+        alias = f"{function}_{suffix}"
+        output.append(f"typedef decltype({specialization}) {alias}_type;\n"
+                      f'template [[host_name("{alias}")]]\n'
+                      f"kernel {alias}_type {specialization};\n")
+    return "".join(output)
+
+
+def gated_delta_source(path):
+    raw = path.read_bytes()
+    if hashlib.sha256(raw).hexdigest() != GATED_DELTA_SHA256:
+        raise ValueError("Expected the pinned mlx-lm 0.31.3 gated_delta.py used by MTPLX")
+    function = next(n for n in ast.parse(raw).body if isinstance(n, ast.FunctionDef)
+                    and n.name == "_make_gated_delta_kernel")
+    output = ["\n// mlx-lm 0.31.3 dependency kernel: Copyright 2023 Apple Inc.\n"
+              "// SPDX-License-Identifier: MIT; see MLX-LM-LICENSE.txt.\n"]
+    # Qwen compute_g is scalar per head; the masked branch also covers ragged rows.
+    for masked in (False, True):
+        values = {"has_mask": masked, "vectorized": False}
+
+        def literal(node):
+            if isinstance(node, ast.Constant):
+                return node.value
+            if isinstance(node, ast.Name):
+                return values[node.id]
+            if isinstance(node, ast.IfExp):
+                return literal(node.body if literal(node.test) else node.orelse)
+            if isinstance(node, ast.JoinedStr):
+                return "".join(literal(n) for n in node.values)
+            if isinstance(node, ast.FormattedValue) and node.format_spec is None and node.conversion == -1:
+                return str(literal(node.value))
+            raise ValueError(f"Unexpected source expression: {ast.dump(node)}")
+
+        for statement in function.body:
+            if isinstance(statement, ast.If) and isinstance(statement.test, ast.Name):
+                for assignment in statement.body if values[statement.test.id] else statement.orelse:
+                    values[assignment.targets[0].id] = literal(assignment.value)
+            elif isinstance(statement, ast.Assign):
+                name = statement.targets[0].id
+                values[name] = literal(statement.value)
+                if name == "source":
+                    break
+        args = parameters("InT:q,k,v;float:g;InT:beta;StT:state_in")
+        args += [("constant const int32_t&", "T")]
+        if masked:
+            args += parameters("bool:mask")
+        args += parameters("InT:y;StT:state_out", output=True)
+        suffix = "_masked" if masked else ""
+        output.append(kernel_source(
+            "mlx_lm/models/gated_delta.py::_make_gated_delta_kernel" + suffix,
+            raw, values["source"], "kernel_qwen_mtplx_gated_delta" + suffix, args,
+            "thread_index_in_simdgroup,thread_position_in_grid,thread_position_in_threadgroup",
+            "typename InT, typename StT, int Dk, int Dv, int Hk, int Hv",
+            [("bf16", "bfloat, float, 128, 128, 16, 48")]))
+    return "".join(output)
+
+
+def gather_front_source(runtime):
+    revision = subprocess.check_output(["git","-C",str(runtime),"rev-parse","HEAD"],text=True).strip()
+    if revision != RUNTIME_REVISION:
+        raise ValueError(f"Expected MTPLX runtime {RUNTIME_REVISION}, got {revision}")
+    output = ["\n// MTPLX runtime JIT shaders: Copyright Apple Inc. SPDX-License-Identifier: MIT\n"
+              "// See MLX-LM-LICENSE.txt; only includes are resolved and instantiations supplied.\n"]
+    for name, digest in (
+        ("indexing.h","e820b8ee2b5132a97122780c12433ebb5100d8078d31e211d0429400a11415bb"),
+        ("gather_front.h","64aacebf6576dfcd389383564fa1214bc87f2a091dd33cc64c598c5367ecab96")):
+        raw = (runtime/"mlx/backend/metal/kernels/indexing"/name).read_bytes()
+        if hashlib.sha256(raw).hexdigest() != digest:
+            raise ValueError(f"Pinned runtime shader changed: {name}")
+        body = "".join(line for line in raw.decode().splitlines(keepends=True)
+                       if not line.startswith(("#include", "#pragma once")))
+        output.append(f"// Runtime unit: {name}; file SHA256: {digest}\n"
+                      f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
+                      "// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
+    for tag, kind, n in (("bf16","bfloat",1),("bf16","bfloat",2),("f16","half",1),("f16","half",2),("u32","uint32_t",1),("f32","float",1)):
+        for loc in ("int","int64_t"):
+            specialization = f"gather_front<{kind}, uint32_t, {loc}, {n}>"
+            output.append(f'template [[host_name("kernel_qwen_mtplx_gather_front_{tag}_{loc}_{n}")]]\n'
+                          f"[[kernel]] decltype({specialization}) {specialization};\n")
+    for tag,kind,n in (("bf16","bfloat",1),("bf16","bfloat",2),("f16","half",1),("f16","half",2),("u32","uint32_t",1),("f32","float",1)):
+        for loc in ("int","int64_t"):
+            specialization=f"gather_front<{kind}, int64_t, {loc}, {n}>"
+            output.append(f'template [[host_name("kernel_qwen_mtplx_gather_front_{tag}_idxi64_{loc}_{n}")]]\n'
+                          f"[[kernel]] decltype({specialization}) {specialization};\n")
+    for tag,kind,n in (("bf16","bfloat",1),("bf16","bfloat",2),("f16","half",1),("f16","half",2),("u32","uint32_t",1),("f32","float",1)):
+        for loc in ("int","int64_t"):
+            specialization=f"gather_front<{kind}, int32_t, {loc}, {n}>"
+            output.append(f'template [[host_name("kernel_qwen_mtplx_gather_front_{tag}_idxi32_{loc}_{n}")]]\n'
+                          f"[[kernel]] decltype({specialization}) {specialization};\n")
+    return "".join(output)
+
+
+def gather_rows_source(runtime):
+    output=[]
+    path=runtime/"mlx/backend/metal/kernels/indexing/gather.h"
+    raw=path.read_bytes()
+    digest="3b2f5b21cd2e71427c9641368a457840f5cb131e202095cd1352f983ce7541ce"
+    if hashlib.sha256(raw).hexdigest()!=digest: raise ValueError("Pinned gather.h changed")
+    body="".join(line for line in raw.decode().splitlines(keepends=True) if not line.startswith(("#include","#pragma once")))
+    output.append(f"// Runtime unit: gather.h; file SHA256: {digest}\n"
+                  f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
+                  "// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
+    raw=(runtime/"mlx/backend/metal/jit/indexing.h").read_bytes()
+    digest="1b38dbdf3120eca3e5266ffd6d69691591c9432c0b5acece4c5ab2de7a9e802f"
+    if hashlib.sha256(raw).hexdigest()!=digest: raise ValueError("Pinned gather JIT wrapper changed")
+    template=raw.decode().split('gather_kernels = R"(',1)[1].split(')";',1)[0]
+    wrappers=[]
+    for tag,kind in (("bf16","bfloat"),("f16","half"),("u32","uint32_t"),("f32","float")):
+        for rank in range(5):
+            for loc in ("int","int64_t"):
+                body=template.format(f"{tag}idxi64",kind,"int64_t",1,
+                    "const device int64_t *idx0 [[buffer(20)]],","idx0",rank,loc)
+                body=body.replace(f"void gather{tag}idxi64_",f"void kernel_qwen_mtplx_gather_{tag}_idxi64_",1)
+                wrappers.append(body)
+    for tag,kind in (("bf16","bfloat"),("f16","half"),("u32","uint32_t"),("f32","float")):
+        for rank in range(5):
+            for loc in ("int","int64_t"):
+                body=template.format(f"{tag}idxi32",kind,"int32_t",1,
+                    "const device int32_t *idx0 [[buffer(20)]],","idx0",rank,loc)
+                body=body.replace(f"void gather{tag}idxi32_",f"void kernel_qwen_mtplx_gather_{tag}_idxi32_",1)
+                wrappers.append(body)
+    for tag,kind in (("bf16","bfloat"),("f16","half"),("u32","uint32_t"),("f32","float")):
+        for rank in range(5):
+            for loc in ("int","int64_t"):
+                body=template.format(f"{tag}idxu32",kind,"uint32_t",1,
+                    "const device uint32_t *idx0 [[buffer(20)]],","idx0",rank,loc)
+                body=body.replace(f"void gather{tag}idxu32_",f"void kernel_qwen_mtplx_gather_{tag}_idxu32_",1)
+                wrappers.append(body)
+    body="".join(wrappers)
+    output.append(f"// Runtime unit: gather JIT wrappers; file SHA256: {digest}\n"
+                  f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
+                  "// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
+    return "".join(output)
+
+
+def silu_source(runtime):
+    output = ["\n// Actual runtime-generated SiLU, captured by tools/mtplx-jit-reference.py.\n"
+              "// Copyright Apple Inc. SPDX-License-Identifier: MIT. Only host aliases change.\n"]
+    def unit(name, body, digest):
+        output.append(f"// Runtime unit: {name}; file SHA256: {digest}\n"
+                      f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
+                      "// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
+    for name,digest,sections in (
+        ("bf16.h","abd87446a310b77ac530ef52a324feae5cb285d03ec9613e3a88ebb71410fdcb",None),
+        ("bf16_math.h","1f374f8380f756eb89acf6a847741cb8fecbe642945e159fb6208d804cc06496",None),
+        ("utils.h","5e1568e9edde9d05dbf86f68fa0d6c6240f2c32b973c7c6a76166b9c0d91543d",
+         [("template <typename IdxT = int64_t>\nMETAL_FUNC IdxT elem_to_loc_1", "///////////////////////////////////////////////////////////////////////////////"),
+          ("template <typename U, typename T>\ninline U cast_to", "template <>")]),
+        ("unary_ops.h","0a5492b65ae39ecb6d8b04e64ea5007e0a4ff60d8d9428559bfbfd03387ece2a",[("struct Sigmoid {","struct Sign {")]),
+        ("binary_ops.h","2dd13c2496f5d6f0856e4ca99db2ceb5c6d7dc0c344b9ba8e41ef3b7d7ed8a97",[("struct Multiply {","struct NotEqual {")]),
+    ):
+        raw = (runtime/"mlx/backend/metal/kernels"/name).read_bytes()
+        if hashlib.sha256(raw).hexdigest()!=digest:
+            raise ValueError(f"Pinned SiLU dependency changed: {name}")
+        text = raw.decode()
+        if sections:
+            body = "".join(text[text.index(start):text.index(end,text.index(start))] for start,end in sections)
+        else:
+            body = "".join(line for line in text.splitlines(keepends=True) if not line.startswith(("#include","#pragma once")))
+        unit(name,body,digest)
+    raw = (OUTPUT.parents[1]/"tests/fixtures/mtplx-silu-jit.json").read_bytes()
+    digest = hashlib.sha256(raw).hexdigest()
+    if digest!="ed0ccd2cbbbcdacbb3e1c1fe8eb7f6b4d104a00197c97a43c93904fc0124d0af":
+        raise ValueError("Pinned SiLU JIT receipt changed")
+    receipt = json.loads(raw)
+    prefix = "BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263"
+    body = receipt["kernels"].replace(f'host_name("{prefix}', 'host_name("kernel_qwen_mtplx_silu_bf16')
+    assert len(re.findall(r'\[\[host_name',body))==19
+    unit("captured-silu-jit",body,digest)
+    return "".join(output)
+
+
+def gather_axis_source(runtime):
+    output = ["\n// Runtime GatherAxis shader: Copyright Apple Inc. SPDX-License-Identifier: MIT\n"]
+    for name,digest,start,end in (
+        ("utils.h","5e1568e9edde9d05dbf86f68fa0d6c6240f2c32b973c7c6a76166b9c0d91543d",
+         "template <typename IdxT = int64_t>\nMETAL_FUNC IdxT elem_to_loc(\n    IdxT", "// Non templated version"),
+        ("indexing/gather_axis.h","e1a745391ff4990f3f1ad75c5687c3b102dcdc4833d8fbbac38e10f54af29af4",None,None),
+    ):
+        raw = (runtime/"mlx/backend/metal/kernels"/name).read_bytes()
+        if hashlib.sha256(raw).hexdigest()!=digest:
+            raise ValueError(f"Pinned GatherAxis source changed: {name}")
+        text = raw.decode()
+        body = text[text.index(start):text.index(end,text.index(start))] if start else "".join(
+            line for line in text.splitlines(keepends=True) if not line.startswith(("#include","#pragma once")))
+        output.append(f"// Runtime unit: {name}; file SHA256: {digest}\n"
+                      f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
+                      "// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
+    for dtype,type_tag in (("bfloat16_t","bf16"),("half","f16"),("float","f32"),("uint32_t","u32")):
+        for idx,tag in (("uint32_t",""),("int32_t","idxi32_")):
+            for loc in ("int","int64_t"):
+                for sc in (False,True):
+                    for ic in (False,True):
+                        args = f"{dtype}, {idx}, {loc}, {str(sc).lower()}, {str(ic).lower()}"
+                        name = f"kernel_qwen_mtplx_gather_axis_{type_tag}_{tag}{loc}_{int(sc)}{int(ic)}"
+                        output.append(f'template [[host_name("{name}")]]\n[[kernel]] decltype(gather_axis<{args}>) gather_axis<{args}>;\n')
+    for loc in ("int","int64_t"):
+        for sc in (False,True):
+            for ic in (False,True):
+                args = f"bool, int64_t, {loc}, {str(sc).lower()}, {str(ic).lower()}"
+                name = f"kernel_qwen_mtplx_gather_axis_bool_idxi64_{loc}_{int(sc)}{int(ic)}"
+                output.append(f'template [[host_name("{name}")]]\n[[kernel]] decltype(gather_axis<{args}>) gather_axis<{args}>;\n')
+    for loc in ("int","int64_t"):
+        for sc in (False,True):
+            for ic in (False,True):
+                args = f"int64_t, int64_t, {loc}, {str(sc).lower()}, {str(ic).lower()}"
+                name = f"kernel_qwen_mtplx_gather_axis_i64_idxi64_{loc}_{int(sc)}{int(ic)}"
+                output.append(f'template [[host_name("{name}")]]\n[[kernel]] decltype(gather_axis<{args}>) gather_axis<{args}>;\n')
+    return "".join(output)
+
+
+def scatter_axis_source(runtime):
+    output = ["\n// Runtime ScatterAxis: Copyright Apple Inc. SPDX-License-Identifier: MIT\n"
+              "#include <metal_atomic>\nnamespace mtplx_eager_scatter {\n"]
+    for name,digest,start,end in (
+        ("complex.h","16e8a815b2cbdb6070e0824e64fe33fccb6e918f1b84ea5c792bd89d33e57bf1",None,None),
+        ("atomic.h","4c35ea2798a2335502865247aee878149fc9ada0d7e84c05d771baef0c7fcc60",None,None),
+        ("reduction/ops.h","78d06730fc9564a73944e7f1fe3897d25c8789b28a939bf418e1968db311da41",
+         "#define DEFINE_SIMD_REDUCE()", "template <typename U>\nstruct Prod"),
+        ("indexing/indexing.h","e820b8ee2b5132a97122780c12433ebb5100d8078d31e211d0429400a11415bb",None,None),
+        ("indexing/scatter.h","fa799c286378c59fbb3aeb973e74bf471861e3356ce2817968fe5d6872c3d9ad",None,None),
+        ("indexing/scatter_axis.h","43eabd0216101f8e32f5cdd19ce40b7f954564be27fad98a5e0fe345e7b94ce5",None,None),
+    ):
+        raw = (runtime/"mlx/backend/metal/kernels"/name).read_bytes()
+        if hashlib.sha256(raw).hexdigest()!=digest:
+            raise ValueError(f"Pinned ScatterAxis source changed: {name}")
+        text = raw.decode()
+        body = text[text.index(start):text.index(end,text.index(start))] if start else "".join(
+            line for line in text.splitlines(keepends=True) if not line.startswith(("#include","#pragma once")))
+        output.append(f"// Runtime unit: {name}; file SHA256: {digest}\n"
+                      f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
+                      "// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
+    for loc in ("int","int64_t"):
+        for uc in (False,True):
+            for ic in (False,True):
+                args = f"bool, int64_t, {loc}, None, {str(uc).lower()}, {str(ic).lower()}"
+                name = f"kernel_qwen_mtplx_scatter_axis_bool_idxi64_{loc}_{int(uc)}{int(ic)}"
+                output.append(f'template [[host_name("{name}")]]\n[[kernel]] decltype(scatter_axis<{args}>) scatter_axis<{args}>;\n')
+    raw = (runtime/"mlx/backend/metal/jit/indexing.h").read_bytes()
+    digest = "1b38dbdf3120eca3e5266ffd6d69691591c9432c0b5acece4c5ab2de7a9e802f"
+    if hashlib.sha256(raw).hexdigest() != digest:
+        raise ValueError("Pinned Scatter JIT wrapper changed")
+    template = raw.decode().split('scatter_kernels = R"(',1)[1].split(')";',1)[0]
+    wrappers = []
+    # One index array/axis, rank-one indices: the original planner always
+    # selects nwork=1 for the sparse raw-logit penalty path.
+    for tag,kind in (("bf16","bfloat"),("f16","half"),("f32","float")):
+        for contiguous in ("false","true"):
+            for loc in ("int","int64_t"):
+                body = template.format(f"{tag}idxi64_sum",kind,"int64_t",f"Sum<{kind}>",1,
+                    "const device int64_t *idx0 [[buffer(20)]],","idx0",contiguous,1,loc)
+                body = body.replace(f"void scatter{tag}idxi64_sum_",f"void kernel_qwen_mtplx_scatter_{tag}_idxi64_sum_",1)
+                # This dependency lives in a namespace alongside ScatterAxis.
+                # Preserve the original unqualified entry ABI explicitly.
+                name = f"kernel_qwen_mtplx_scatter_{tag}_idxi64_sum_1_updc_{contiguous}_nwork1_{loc}"
+                body = body.replace("[[kernel]]",f'[[host_name("{name}")]] [[kernel]]',1)
+                wrappers.append(body)
+    body = "".join(wrappers)
+    output.append(f"// Runtime unit: scatter JIT wrappers; file SHA256: {digest}\n"
+                  f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
+                  "// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
+    output.append("} // namespace mtplx_eager_scatter\n")
+    return "".join(output)
+
+
+def swiglu_source():
+    raw = (OUTPUT.parents[1]/"tests/fixtures/mtplx-swiglu-jit.json").read_bytes()
+    digest = hashlib.sha256(raw).hexdigest()
+    if digest != "8cb2a51a7d9f0cb91a8a7669ef844c3964dc469ec9e4436a719e47f025b8e781":
+        raise ValueError("Pinned SwiGLU JIT receipt changed")
+    receipt = json.loads(raw)
+    prefix = re.search(r'host_name\("([^\"]+)_contiguous"',receipt["kernels"]).group(1)
+    body = receipt["kernels"].replace(f'host_name("{prefix}', 'host_name("kernel_qwen_mtplx_swiglu_bf16')
+    assert len(re.findall(r'\[\[host_name',body))==19
+    return ("\n// Actual runtime-generated SwiGLU: Copyright Apple Inc.; SPDX-License-Identifier: MIT\n"
+            f"// Runtime unit: captured-swiglu-jit; file SHA256: {digest}\n"
+            f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
+            "// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
+
+
+def compute_g_source(runtime):
+    output = ["\n// Original compute_g JIT dependencies; Apple MIT except cexpf.h (Apache-2.0).\n"]
+    def unit(name,body,digest):
+        output.append(f"// Runtime unit: {name}; file SHA256: {digest}\n"
+                      f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
+                      "// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
+    for name,digest,sections in (
+        ("complex.h","16e8a815b2cbdb6070e0824e64fe33fccb6e918f1b84ea5c792bd89d33e57bf1",None),
+        ("cexpf.h","88b6e15a52a5800d98d9bc6da840ca5cf70bf572fda136409580c1f17b1e0aab",None),
+        ("utils.h","5e1568e9edde9d05dbf86f68fa0d6c6240f2c32b973c7c6a76166b9c0d91543d",
+         [("template <typename U>\nstruct Limits", "///////////////////////////////////////////////////////////////////////////////"),
+          ("inline float log1p(float x)", "///////////////////////////////////////////////////////////////////////////////")]),
+        ("unary_ops.h","0a5492b65ae39ecb6d8b04e64ea5007e0a4ff60d8d9428559bfbfd03387ece2a",
+         [("struct Exp {","struct Expm1 {"),("struct Negative {","struct Real {")]),
+        ("binary_ops.h","2dd13c2496f5d6f0856e4ca99db2ceb5c6d7dc0c344b9ba8e41ef3b7d7ed8a97",
+         [("struct Add {","struct FloorDivide {"),("struct LogAddExp {","struct Maximum {")]),
+    ):
+        raw = (runtime/"mlx/backend/metal/kernels"/name).read_bytes()
+        if hashlib.sha256(raw).hexdigest()!=digest:
+            raise ValueError(f"Pinned compute_g dependency changed: {name}")
+        text = raw.decode()
+        body = "".join(text[text.index(start):text.index(end,text.index(start))] for start,end in sections) if sections else "".join(
+            line for line in text.splitlines(keepends=True) if not line.startswith(("#include","#pragma once")))
+        unit(name,body,digest)
+    raw = (OUTPUT.parents[1]/"tests/fixtures/mtplx-compute_g-jit.json").read_bytes()
+    digest = hashlib.sha256(raw).hexdigest()
+    if digest != "701e2f54b7cb8bf97616f83256657e6c5e8cc8b46f4b65c559ccc030ab111dbf":
+        raise ValueError("Pinned compute_g JIT receipt changed")
+    kernels = json.loads(raw)["kernels"]
+    prefix = re.search(r'host_name\("([^\"]+)_contiguous"',kernels).group(1)
+    body = kernels.replace(f'host_name("{prefix}', 'host_name("kernel_qwen_mtplx_compute_g_bf16')
+    assert len(re.findall(r'\[\[host_name',body))==19
+    unit("captured-compute_g-jit",body,digest)
+    return "".join(output)
+
+
+def qsa_prepare_source(reference):
+    path = reference/"mtplx/kernels/qsa_indexer_prepare.py"
+    raw = path.read_bytes()
+    if hashlib.sha256(raw).hexdigest() != "a77f6ca5ae805e729519c4629ae88b455a6dbf473a457a6e1c8219174eb59091":
+        raise ValueError("Pinned QSA preparation source changed")
+    constants = dict(heads=4,head_dim=128,rotary_dim=64,half_rotary=32,ratio=4,
+                     eps=1e-6,attention_scaling=1.0)
+    def literal(node):
+        if isinstance(node,ast.Constant): return node.value
+        if isinstance(node,ast.Name): return constants[node.id]
+        if isinstance(node,ast.Call) and isinstance(node.func,ast.Name) and node.func.id=="float":
+            return float(literal(node.args[0]))
+        if isinstance(node,ast.FormattedValue):
+            value=literal(node.value)
+            return repr(value) if node.conversion==ord("r") else str(value)
+        if isinstance(node,ast.JoinedStr): return "".join(literal(v) for v in node.values)
+        raise ValueError(f"Unsupported QSA header expression: {ast.dump(node)}")
+    output=[]
+    for function,name,inputs,outputs in (
+        ("_prepare_queries_kernel","qsa_prepare_q",
+         "T:raw_q;int64_t:raw_q_strides;T:norm_weight;int64_t:norm_weight_strides;float:inv_freq;int64_t:inv_freq_strides;int:pos_start","T:prepared_q"),
+        ("_pool_keys_kernel","qsa_pool_k",
+         "T:raw_keys;int64_t:raw_keys_strides;T:norm_weight;int64_t:norm_weight_strides;float:inv_freq;int64_t:inv_freq_strides;int:block_start","T:pooled"),
+    ):
+        fn=next(n for n in ast.parse(raw).body if isinstance(n,ast.FunctionDef) and n.name==function)
+        values={n.targets[0].id:n.value for n in fn.body if isinstance(n,ast.Assign)
+                and isinstance(n.targets[0],ast.Name)}
+        # metal_stdlib is already included at translation-unit scope; Metal
+        # module imports cannot be nested inside our collision-avoiding scope.
+        header = "".join(line for line in literal(values["header"]).splitlines(keepends=True)
+                         if not line.lstrip().startswith("#include"))
+        output.append(f"\nnamespace mtplx_{name} {{\n"+header)
+        args = [("constant const int64_t*" if key.endswith("_strides") else kind,key)
+                for kind,key in parameters(inputs)+parameters(outputs,output=True)]
+        output.append(kernel_source(f"mtplx/kernels/qsa_indexer_prepare.py::{function}",raw,
+            literal(values["source"]),f"kernel_qwen_mtplx_{name}",
+            args,
+            "threadgroup_position_in_grid,thread_index_in_simdgroup","typename T",[("bf16","bfloat")]))
+        output.append(f"}} // namespace mtplx_{name}\n")
+    return "".join(output)
+
+
+def qsa_select_source(reference):
+    path = reference / "mtplx/kernels/qsa_indexer_select.py"
+    raw = path.read_bytes()
+    digest = hashlib.sha256(raw).hexdigest()
+    if digest != "a3c74af27a7045c12f2893a8b7a91724c00d8a4148315c3165f3480c83016cf3":
+        raise ValueError("Pinned QSA selector source changed")
+    tree = ast.parse(raw)
+    header = next(ast.literal_eval(n.value) for n in tree.body
+                  if isinstance(n, ast.Assign) and isinstance(n.targets[0], ast.Name)
+                  and n.targets[0].id == "_HEADER")
+    fn = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "_selector_kernel")
+    body = next(n.value for n in fn.body if isinstance(n, ast.Assign)
+                and isinstance(n.targets[0], ast.Name) and n.targets[0].id == "source")
+    assert isinstance(body, ast.BinOp) and isinstance(body.op, ast.Add) and body.right.id == "epilogue"
+    body = ast.literal_eval(body.left)
+    epilogues = {}
+    ep = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "_epilogue")
+    for branch in ep.body:
+        if isinstance(branch, ast.If):
+            mode = ast.literal_eval(branch.test.comparators[0])
+            names, source = ast.literal_eval(branch.body[0].value)
+            epilogues[mode] = dict(outputs=names, source=source,
+                                   sha256=hashlib.sha256(source.encode()).hexdigest())
+    prefill_raw = (reference / "mtplx/kernels/qsa_indexer_prefill.py").read_bytes()
+    prefill_digest = hashlib.sha256(prefill_raw).hexdigest()
+    if prefill_digest != "4d6fd428243c001746f69f8aed45991356772c2bd4a45586eb3c6813c91998d3":
+        raise ValueError("Pinned QSA prefill source changed")
+    prefill_tree = ast.parse(prefill_raw)
+    literals = {n.targets[0].id: ast.literal_eval(n.value) for n in prefill_tree.body
+                if isinstance(n, ast.Assign) and isinstance(n.targets[0], ast.Name)
+                and isinstance(n.value, ast.Constant)}
+    topk = next(n for n in prefill_tree.body if isinstance(n, ast.FunctionDef) and n.name == "_prefill_topk_kernel")
+    assignments = {n.targets[0].id:n.value for n in topk.body if isinstance(n,ast.Assign) and isinstance(n.targets[0],ast.Name)}
+    assert assignments["source"].right.id == "epilogue"
+    topk_body = ast.literal_eval(assignments["source"].left)
+    # Preserve literal f-string segments verbatim; Rust fills only these
+    # named header substitutions, including constants and the output mode.
+    topk_header = "".join(n.value if isinstance(n,ast.Constant) else "@"+ast.unparse(n.value)+"@"
+                          for n in assignments["header"].right.values)
+    prefill = dict(file_sha256=prefill_digest, mpp_header=literals["_MPP_SCORE_HEADER"],
+                   mpp_body=literals["_MPP_SCORE_SOURCE"],topk_header=topk_header,topk_body=topk_body)
+    for key in ("mpp_header","mpp_body","topk_header","topk_body"):
+        prefill[key+"_sha256"] = hashlib.sha256(prefill[key].encode()).hexdigest()
+    return json.dumps(dict(revision=REVISION, file_sha256=digest, header=header, body=body,
+                           header_sha256=hashlib.sha256(header.encode()).hexdigest(),
+                           body_sha256=hashlib.sha256(body.encode()).hexdigest(),
+                           epilogues=epilogues,prefill=prefill), indent=2) + "\n"
+
+
+def generate(reference, gated_delta):
+    revision = subprocess.check_output(
+        ["git", "-C", str(reference), "rev-parse", "HEAD"], text=True).strip()
+    if revision != REVISION:
+        raise ValueError(f"Expected MTPLX {REVISION}, got {revision}")
+    subprocess.run(["git", "-C", str(reference), "diff", "--exit-code", "HEAD"],
+                   check=True, stdout=subprocess.DEVNULL)
+    output = ["// MTPLX bodies: Copyright 2026 MTPLX. SPDX-License-Identifier: Apache-2.0\n"
+              "// Dependency bodies carry their own license notice below.\n"
+              "// Generated by tools/mtplx-kernel-source.py; do not edit bodies.\n"
+              f"// MTPLX revision: {REVISION}\n"
+              "// Only entry-point ABI and template instantiations are adapted.\n"
+              "#include <metal_stdlib>\nusing namespace metal;\n"]
+    for module, variable, name, inputs, outputs, builtins, extra, variants in KERNELS:
+        path = reference / f"mtplx/kernels/{module}.py"
+        raw = path.read_bytes()
+        constants = {node.targets[0].id: node.value.value
+                     for node in ast.parse(raw).body
+                     if isinstance(node, ast.Assign) and len(node.targets) == 1
+                     and isinstance(node.targets[0], ast.Name)
+                     and isinstance(node.value, ast.Constant)
+                     and isinstance(node.value.value, str)}
+        body = constants[variable]
+        args = parameters(inputs) + parameters(outputs, output=True)
+        output.append(kernel_source(
+            f"mtplx/kernels/{module}.py::{variable}", raw, body,
+            f"kernel_qwen_mtplx_{name}", args, builtins,
+            f"typename T{', ' + extra if extra else ''}", variants))
+    output.append(gated_delta_source(gated_delta))
+    output.append(gather_front_source(reference.parent/"mtplx-runtime-0.32.2"))
+    output.append(silu_source(reference.parent/"mtplx-runtime-0.32.2"))
+    output.append(gather_axis_source(reference.parent/"mtplx-runtime-0.32.2"))
+    output.append(gather_rows_source(reference.parent/"mtplx-runtime-0.32.2"))
+    output.append(scatter_axis_source(reference.parent/"mtplx-runtime-0.32.2"))
+    output.append(swiglu_source())
+    output.append(compute_g_source(reference.parent/"mtplx-runtime-0.32.2"))
+    output.append(qsa_prepare_source(reference))
+    output.append(qsa_attention_source(reference))
+    return "".join(output)
+
+
+def qsa_attention_source(reference):
+    output=[]
+    for module,variable,digest,inputs,builtins,extra,variants in (
+        ("qsa_flash_skip","_SRC","27d49eb93f0de82013fdc0b271e77076b19d1972450e1bbc44cdd92984669103",
+         "T:q,k,v;int:blocks,params;float:scale",
+         "threadgroup_position_in_grid,thread_position_in_threadgroup","int GQA",
+         [(f"g{g}_{tag}",f"{kind}, {g}") for g in (1,12) for tag,kind in (("bf16","bfloat"),("f16","half"))]),
+        ("qsa_prefill_flash","_SOURCE","5ca852ea441810a47b374a1c94a49a636055dab31f7f11c4ce0118f7f70b388d",
+         "T:q;int64_t:q_strides;T:k;int64_t:k_strides;T:v;int64_t:v_strides;int:block_ids;int64_t:block_ids_strides;bool:block_valid;int64_t:block_valid_strides;int:params;float:scale",
+         "threadgroup_position_in_grid,thread_index_in_simdgroup","",
+         [("bf16","bfloat"),("f16","half")]),
+    ):
+        raw=(reference/f"mtplx/kernels/{module}.py").read_bytes()
+        if hashlib.sha256(raw).hexdigest()!=digest: raise ValueError(f"Pinned {module} source changed")
+        values={n.targets[0].id:ast.literal_eval(n.value) for n in ast.parse(raw).body
+                if isinstance(n,ast.Assign) and isinstance(n.targets[0],ast.Name) and isinstance(n.value,ast.Constant)}
+        header=values["_HEADER"]
+        includes="".join(line for line in header.splitlines(keepends=True) if line.lstrip().startswith("#include"))
+        header="".join(line for line in header.splitlines(keepends=True) if not line.lstrip().startswith("#include"))
+        output.append(includes+f"\nnamespace mtplx_{module} {{\n"+header)
+        args=[("constant const int64_t*" if name.endswith("_strides") else kind,name)
+              for kind,name in parameters(inputs)+parameters("T:out",output=True)]
+        output.append(kernel_source(f"mtplx/kernels/{module}.py::{variable}",raw,values[variable],
+            f"kernel_qwen_mtplx_{module}",args,builtins,"typename T"+(f", {extra}" if extra else ""),variants))
+        output.append(f"}} // namespace mtplx_{module}\n")
+    return "".join(output)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("reference", type=Path)
+    parser.add_argument("--gated-delta-source", type=Path, required=True,
+                        help="gated_delta.py from MTPLX's pinned mlx-lm 0.31.3 installation")
+    parser.add_argument("--check", action="store_true")
+    parser.add_argument("--qsa-select", action="store_true", help="Print original dynamic QSA selector source sections")
+    args = parser.parse_args()
+    source = generate(args.reference, args.gated_delta_source)
+    selector = qsa_select_source(args.reference)
+    if args.check:
+        if OUTPUT.read_text() != source:
+            raise SystemExit(f"MTPLX kernel source differs: {OUTPUT}")
+        if QSA_SELECT_OUTPUT.read_text() != selector:
+            raise SystemExit(f"MTPLX selector source differs: {QSA_SELECT_OUTPUT}")
+        print(f"All {len(KERNELS) + 6} Metal bodies and {sum(len(k[-1]) for k in KERNELS) + 10} "
+              f"custom entry points and {source.count('// BEGIN RUNTIME UNIT')} runtime units plus dynamic QSA decode/prefill sources match the pinned reference export")
+    else:
+        print(selector if args.qsa_select else source, end="")
+
+
+if __name__ == "__main__":
+    main()

tools/mtplx-policy-reference.py

@@ -0,0 +1,70 @@
+"""Record pinned Qwen server policy functions; no model loading or generation."""
+import argparse
+import itertools
+import json
+import os
+
+parser = argparse.ArgumentParser(description=__doc__)
+parser.add_argument("model")
+args = parser.parse_args()
+print("reference: importing policy functions without model loading", flush=True)
+from mtplx.server.openai import _server_runtime_env_overrides
+from mtplx import generation as g
+from mtplx.session_bank import _lazy_snapshot_enabled
+
+keys = [
+    "MTPLX_AR_PIPELINE", "MTPLX_FAMILY_CAPTURE_COMMIT", "MTPLX_FUSED_HC_V3",
+    "MTPLX_FUSED_GDN_INPROJ", "MTPLX_FUSED_GATE_UP", "MTPLX_FUSED_GDN_CONVNORM",
+    "MTPLX_FUSED_GDN_STEP", "MTPLX_FUSED_CONVNORM_VERIFY", "MTPLX_QSA_GATHER",
+    "MTPLX_NAX_VERIFY", "MTPLX_SKIP_VERIFY_SNAPSHOT",
+]
+cases = [{}]
+cases += [{key: value} for key in keys for value in ("0", " YES ", "", "unknown")]
+cases += [dict((k, v) for k, v in zip(
+    ("MTPLX_COMPILED_GDN", "MTPLX_QWEN4EXP_COMPILE"), pair) if v is not None)
+    for pair in itertools.product((None, "1", "0", "", "junk"), repeat=2)]
+cases += [{key: value} for key in (
+    "MTPLX_SESSION_LAZY_SNAPSHOT", "MTPLX_SESSION_STORE_ON_PREFILL",
+    "MTPLX_GDN_BOUNDARY_CAPTURE", "MTPLX_ASYNC_AR", "MTPLX_EVAL_AUDIT",
+    "MTPLX_LAZY_MTP_HISTORY_APPEND", "MTPLX_DEFER_REPAIR_EVAL",
+    "MTPLX_PREFILL_EXTERNAL_EMIT_LOGITS",
+) for value in ("0", "on", "", "junk")]
+cases += [{key: value} for key in (
+    "MTPLX_SESSION_STORE_ON_PREFILL_MIN_SUFFIX", "MTPLX_SMALL_SUFFIX_FUSED_MAX",
+    "MTPLX_GDN_BOUNDARY_MAX", "MTPLX_GDN_BOUNDARY_TAIL_INTERVAL",
+) for value in ("-3", "0", "17", "1_024", "bad")]
+cases += [{"MTPLX_PREFILL_OMLX_EXTERNAL": "1"},
+          {"MTPLX_PREFILL_STOCK_CACHE_ONLY": "1"},
+          {"MTPLX_PREFILL_STOCK_CACHE_ONLY": "1", "MTPLX_ALLOW_UNSAFE_PREFILL_STOCK_CACHE_ONLY": "1"},
+          {"MTPLX_ASYNC_AR": "1", "MTPLX_EVAL_AUDIT": "audit"}]
+cases += [{key: value} for key in ("MTPLX_FAMILY_CAPTURE_COMMIT", "MTPLX_DEFER_REPAIR_EVAL")
+          for value in ("enable", "enabled", "disable", "disabled")]
+baseline = {k: v for k, v in os.environ.items() if not k.startswith("MTPLX_")}
+for mtp, env in itertools.product((False, True), cases):
+    os.environ.clear()
+    os.environ.update(baseline)
+    os.environ.update(env)
+    overrides = _server_runtime_env_overrides(argparse.Namespace(
+        model=args.model, generation_mode="mtp" if mtp else "ar"), None)
+    os.environ.update(overrides)
+    # Native app selects bounded prefill; emulate that request-local policy.
+    os.environ["MTPLX_SUSTAINED_PREFILL"] = "1"
+    error = None
+    options = None
+    try:
+        with g.prefill_chunk_size_override(2048):
+            options = dict(chunk=g._prefill_chunk_size(), final_only=g._final_logits_prefill_enabled(),
+                external_cache_only=g._prefill_external_cache_only_enabled(),
+                external_emit_logits=g._prefill_external_emit_logits_enabled(),
+                lazy_kv=_lazy_snapshot_enabled(), store=g._store_on_prefill_env_enabled(),
+                store_min=g._store_on_prefill_min_suffix(), fused_max=g._small_suffix_fused_max(),
+                boundaries=g._gdn_boundary_capture_enabled(), boundary_max=g._gdn_boundary_max_count(),
+                boundary_tail=g._gdn_boundary_tail_interval(), pipeline=g._env_truthy("MTPLX_AR_PIPELINE"),
+                sync_eval=not g._env_truthy("MTPLX_ASYNC_AR") or bool(os.environ.get("MTPLX_EVAL_AUDIT")),
+                capture=g._family_capture_commit_enabled() if mtp else None,
+                lazy_history=g._env_truthy("MTPLX_LAZY_MTP_HISTORY_APPEND"),
+                defer_repair=g._defer_repair_eval() if mtp else None)
+    except ValueError as exc:
+        error = str(exc)
+    print(json.dumps(dict(event="policy_reference", mtp=mtp, environment=env,
+        overrides=overrides, options=options, error=error)), flush=True)

tools/mtplx-session-codec-reference.py

@@ -0,0 +1,70 @@
+"""Pinned SessionBank codec fixture; no model weights, downloads or inference."""
+import hashlib
+import json
+import sys
+from pathlib import Path
+
+import mlx.core as mx
+from mtplx.cache_bank.codec import encode_payload, decode_payload
+from mtplx.cache_state import CacheSnapshot
+
+
+def digest(value):
+    if value is None:
+        return None
+    mx.eval(value)
+    return {"shape": list(value.shape), "dtype": str(value.dtype).removeprefix("mlx.core."),
+            "sha256": hashlib.sha256(bytes(memoryview(value))).hexdigest()}
+
+
+def state_digests(snapshot):
+    return [None if state is None else [digest(v) for v in state] for state in snapshot.states]
+
+
+def main(directory):
+    directory = Path(directory)
+    directory.mkdir(parents=True, exist_ok=True)
+    base = mx.arange(2 * 513 * 4, dtype=mx.float32).reshape(1, 2, 513, 4) / 16
+    bf16 = base.astype(mx.bfloat16)
+    # Actual cache shapes can be strided, reversed, broadcast, empty or scalar.
+    small = mx.arange(12, dtype=mx.int32).reshape(1, 3, 4)
+    reversed_view = small[:, :, ::-1]
+    broadcast = mx.broadcast_to(mx.array([7], dtype=mx.uint16), (1, 2, 3))
+    trunk = CacheSnapshot(states=([bf16, reversed_view, broadcast, mx.array([], dtype=mx.float16)],
+        (base, base[:, :, ::-1, :], small.transpose(0, 2, 1), None), None),
+        meta_states=("", None, None))
+    head = CacheSnapshot(states=((bf16[:, :, :5, :], base[:, :, :5, :], None, None),), meta_states=(None,))
+    boundary = CacheSnapshot(states=([bf16[:, :, :2, :], reversed_view], None, None), meta_states=("", None, None))
+    cases = []
+    for block in (0, 256, 1024):
+        # Scalar codec quirk: decode preserves the one-element import vector.
+        encoded = encode_payload(cache_snapshot=trunk, logits=mx.array(1.25),
+            hidden=bf16[:, :1, :1, :], mtp_history_snapshot=head,
+            gdn_boundaries=[(2, boundary, None)], has_recurrent=True, block_size=block)
+        folder = directory / str(block)
+        folder.mkdir(exist_ok=True)
+        for name, raw in encoded.tensors.items():
+            (folder / (name + ".bin")).write_bytes(raw)
+        decoded = decode_payload(encoded.spec, encoded.tensors.__getitem__)
+        legacy = dict(encoded.spec)
+        del legacy["gdn_boundaries"]
+        legacy = decode_payload(legacy, encoded.tensors.__getitem__)
+        assert not legacy.gdn_boundaries
+        no_hidden = {**encoded.spec, "gdn_boundaries": [
+            {k: v for k, v in b.items() if k != "hidden_last"}
+            for b in encoded.spec["gdn_boundaries"]]}
+        assert decode_payload(no_hidden, encoded.tensors.__getitem__).gdn_boundaries[0][2] is None
+        cases.append(dict(block_size=block, spec=encoded.spec, nbytes=encoded.nbytes,
+            legacy_boundary_count=len(legacy.gdn_boundaries),
+            blobs={k: hashlib.sha256(v).hexdigest() for k, v in encoded.tensors.items()},
+            decoded=dict(trunk=state_digests(decoded.cache_snapshot),
+                mtp=state_digests(decoded.mtp_history_snapshot), logits=digest(decoded.logits),
+                hidden=digest(decoded.hidden), boundaries=[dict(tokens=n, states=state_digests(s), hidden=digest(h))
+                    for n, s, h in decoded.gdn_boundaries])))
+        print(json.dumps({"event": "codec_reference_case", "block_size": block,
+                          "blobs": len(encoded.tensors), "bytes": encoded.nbytes}), flush=True)
+    (directory / "reference.json").write_text(json.dumps({"cases": cases}))
+
+
+if __name__ == "__main__":
+    main(sys.argv[1])

tools/mtplx-weight-fusion-reference.py

@@ -0,0 +1,123 @@
+"""Original sanitize-time projection fusions; tiny synthetic checkpoint packs.
+
+Run under test-supervisor --command. No model downloads or full model loads.
+"""
+import os
+import sys
+
+print("weight-fusion reference: importing", file=sys.stderr, flush=True)
+import hashlib
+import json
+from types import SimpleNamespace
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+from mtplx.models import qwen4_exp as qwen
+
+
+def checkpoint_metadata(root):
+    """Header/lazy-graph audit of the installed checkpoint, never mx.eval."""
+    import glob
+    from pathlib import Path
+    from contextlib import redirect_stdout
+    config = json.loads((Path(root) / "config.json").read_text())["text_config"]
+    os.environ.update(MTPLX_FUSED_GATE_UP="1", MTPLX_FUSED_GDN_INPROJ="1", MTPLX_FUSED_QSA_QKV="1")
+    weights = {}
+    for path in glob.glob(str(Path(root) / "model*.safetensors")):
+        print(f"weight-fusion headers: {Path(path).name}", file=sys.stderr, flush=True)
+        weights.update(mx.load(path))
+    before = len(weights)
+    layers = []
+    for _ in range(config["num_hidden_layers"]):
+        attention = nn.Module()
+        attention.indexer = nn.Module() if config["indexer_n_heads"] else None
+        layers.append(SimpleNamespace(linear_attn=nn.Module(), self_attn=attention,
+            mlp=SimpleNamespace(switch_mlp=SimpleNamespace(down_proj=None), shared_expert=SimpleNamespace(down_proj=None))))
+    model = SimpleNamespace(layers=layers)
+    with redirect_stdout(sys.stderr):
+        weights = qwen._fuse_gate_up_sanitize(model, weights)
+        weights = qwen._fuse_gdn_in_proj_sanitize(model, weights)
+        weights = qwen._fuse_qsa_qkv_sanitize(model, weights)
+    rows = [[name, list(a.shape), str(a.dtype).removeprefix("mlx.core.")] for name, a in sorted(weights.items())]
+    fusions = {}
+    for i, layer in enumerate(layers):
+        base = f"language_model.model.layers.{i}"
+        for suffix, module in (("mlp.switch_mlp", layer.mlp.switch_mlp), ("mlp.shared_expert", layer.mlp.shared_expert),
+                ("linear_attn.in_proj_fused", getattr(layer.linear_attn, "in_proj_fused", None)),
+                ("self_attn.qkv_fused", getattr(layer.self_attn, "qkv_fused", None))):
+            if module is not None and hasattr(module, "bits"):
+                fusions[f"{base}.{suffix}"] = dict(bits=module.bits, group=module.group_size, splits=getattr(module, "_splits", []))
+    print(json.dumps(dict(before=before, after=len(weights), fusions=fusions,
+        metadata_sha256=hashlib.sha256(json.dumps(rows, separators=(",", ":")).encode()).hexdigest()), separators=(",", ":")), flush=True)
+
+
+def main():
+    assert mx.__version__ == "0.32.2"
+    profiles = [(mask, "normal") for mask in range(8)] + [
+        (7, name) for name in ("qsa_bias", "indexer_mixed", "qkv_mixed",
+        "gdn_mixed", "gdn_missing", "indexer_missing", "no_indexer",
+        "gate_missing_bias", "shared_only", "gdn_2bit")]
+    for mask, profile in profiles:
+        print(f"weight-fusion reference: {mask}/{profile}", 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))
+        weights = {}
+        specs = [("language_model.model.layers.0.mlp.switch_mlp.gate_proj", [2, 4], 320),
+            ("language_model.model.layers.0.mlp.switch_mlp.up_proj", [2, 4], 320),
+            ("language_model.model.layers.0.mlp.shared_expert.gate_proj", [4], 320),
+            ("language_model.model.layers.0.mlp.shared_expert.up_proj", [4], 320)]
+        specs += [(f"language_model.model.layers.0.linear_attn.in_proj_{sub}", [n],
+            160 if profile == "gdn_2bit" else (640 if profile == "gdn_mixed" and sub == "a" else 320))
+            for sub, n in (("qkv", 5), ("z", 3), ("b", 1), ("a", 1))]
+        specs += [(f"language_model.model.layers.1.self_attn.{sub}", [n],
+            640 if (profile == "indexer_mixed" and sub.startswith("indexer")) or
+                (profile == "qkv_mixed" and sub == "k_proj") else 320)
+            for sub, n in (("q_proj", 5), ("k_proj", 2), ("v_proj", 2), ("indexer.index_qk_proj", 3))]
+        for index, (base, prefix, columns) in enumerate(specs):
+            for part, width, dtype in (("weight", columns, mx.uint32), ("scales", 40, mx.bfloat16), ("biases", 40, mx.bfloat16)):
+                shape = (*prefix, width)
+                data = np.arange(np.prod(shape), dtype=np.int32) % 17 + index
+                weights[f"{base}.{part}"] = mx.array(data, dtype=dtype).reshape(shape)
+        remove = {
+            "gdn_missing": ["language_model.model.layers.0.linear_attn.in_proj_a.biases"],
+            "gate_missing_bias": ["language_model.model.layers.0.mlp.switch_mlp.up_proj.biases"],
+            "shared_only": ["language_model.model.layers.0.mlp.switch_mlp.gate_proj.weight"],
+        }.get(profile, [])
+        if profile == "indexer_missing":
+            remove += [f"language_model.model.layers.1.self_attn.indexer.index_qk_proj.{p}" for p in ("weight", "scales", "biases")]
+        for name in remove:
+            del weights[name]
+        if profile == "qsa_bias":
+            weights["language_model.model.layers.1.self_attn.q_proj.bias"] = mx.zeros((5,), mx.bfloat16)
+        # Only the attributes read/replaced by the original sanitize routines.
+        mlp = SimpleNamespace(switch_mlp=SimpleNamespace(down_proj=None), shared_expert=SimpleNamespace(down_proj=None))
+        gdn, attention = nn.Module(), nn.Module()
+        attention.indexer = None if profile == "no_indexer" else nn.Module()
+        model = SimpleNamespace(layers=[SimpleNamespace(mlp=mlp, linear_attn=gdn), SimpleNamespace(self_attn=attention)])
+        # Original routines print human progress on stdout; preserve JSONL there.
+        from contextlib import redirect_stdout
+        with redirect_stdout(sys.stderr):
+            weights = qwen._fuse_gate_up_sanitize(model, weights)
+            weights = qwen._fuse_gdn_in_proj_sanitize(model, weights)
+            weights = qwen._fuse_qsa_qkv_sanitize(model, weights)
+        metadata = {}
+        for name, module in (("language_model.model.layers.0.mlp.switch_mlp", mlp.switch_mlp),
+                ("language_model.model.layers.0.mlp.shared_expert", mlp.shared_expert),
+                ("language_model.model.layers.0.linear_attn.in_proj_fused", getattr(gdn, "in_proj_fused", None)),
+                ("language_model.model.layers.1.self_attn.qkv_fused", getattr(attention, "qkv_fused", None))):
+            if module is not None and hasattr(module, "bits"):
+                metadata[name] = dict(bits=module.bits, group=module.group_size, splits=getattr(module, "_splits", []))
+        ordered = sorted(weights.items())
+        mx.eval(*[a for _, a in ordered])
+        print(json.dumps(dict(mask=mask, profile=profile, fusions=metadata, tensors=[
+            dict(name=name, shape=a.shape, sha256=hashlib.sha256(np.asarray(a.astype(mx.float32)).astype('<f4').tobytes()).hexdigest())
+            for name, a in ordered]), separators=(",", ":")), flush=True)
+
+
+if __name__ == "__main__":
+    if len(sys.argv) == 3 and sys.argv[1] == "--metadata-only":
+        checkpoint_metadata(sys.argv[2])
+    else:
+        assert len(sys.argv) == 1
+        main()

tools/qwen-chat-reference.py

@@ -0,0 +1,83 @@
+"""Generate ongoing-chat token fixtures through MTPLX's actual server encoder.
+
+Use MTPLX's installed Python environment. Reads only local tokenizer files and
+the supplied README; no inference, server, tools or model downloads are started.
+"""
+
+import argparse
+import hashlib
+import json
+import os
+from pathlib import Path
+
+os.environ["HF_HUB_OFFLINE"] = "1"
+os.environ["TRANSFORMERS_OFFLINE"] = "1"
+
+from mtplx.runtime import _load_tokenizer_resilient
+from mtplx.server import openai as server
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("model", type=Path)
+    parser.add_argument("readme", type=Path)
+    parser.add_argument("--source-run", type=Path,
+                        help="Use exact prompts and generated history from a completed model-eval JSONL")
+    args = parser.parse_args()
+    tokenizer = _load_tokenizer_resilient(
+        args.model.resolve(), json.loads((args.model / "config.json").read_text()),
+    )
+    readme = args.readme.read_text()
+    prompts = [
+        "Give a summary of the following text:\n\n" + readme,
+        "Tell me a complete short fictional story. Choose the setting yourself; do not ask questions.",
+        "Return one Python code block defining has_close_elements(numbers: list[float], threshold: float) -> bool.",
+    ]
+    replies = [
+        {"role": "assistant", "content": text,
+         "reasoning_content": "**Plan:** answer the request directly."}
+        for text in [
+            "DS4Server runs local language models on Apple silicon.",
+            "A lighthouse keeper found a letter in an empty bottle. She wrote back, and the sea carried her reply home.",
+            "```python\ndef has_close_elements(numbers, threshold):\n    return any(abs(a-b) < threshold for i, a in enumerate(numbers) for b in numbers[i+1:])\n```",
+        ]
+    ]
+    if args.source_run:
+        records = [json.loads(line) for line in args.source_run.read_text().splitlines()]
+        start = next(record for record in records if record["event"] == "start")
+        results = [record for record in records if record["event"] == "result"]
+        if not (start["plain_chat"] and not start["system_prompt"]
+                and start["settings"]["reasoning"] == "low"
+                and len(start["prompts"]) == 3
+                and [record["turn"] for record in results] == [1, 2, 3]
+                and all(record["ok"] and record["finish_reason"] == "stop" for record in results)
+                and readme in start["prompts"][0]):
+            parser.error("source must be a complete three-turn Low plain chat using the supplied README")
+        prompts = start["prompts"]
+        replies = [{"role": "assistant", "content": record["content"],
+                    "reasoning_content": record["reasoning"] or ""} for record in results]
+    history = []
+    cases = []
+    for index, prompt in enumerate(prompts):
+        history.append({"role": "user", "content": prompt})
+        ids = server._encode_messages_uncached(
+            tokenizer, [server.ChatMessage(**message) for message in history],
+            enable_thinking=True, reasoning_effort="low",
+            preserve_reasoning_history=True, tools=None,
+        )
+        cases.append({"messages": list(history), "token_ids": ids,
+                      "rendered": tokenizer.decode(ids, skip_special_tokens=False)})
+        history.append(replies[index])
+    source = Path(server.__file__).resolve()
+    print(json.dumps({"reference": "MTPLX server _encode_messages_uncached",
+                      "source": str(source),
+                      "source_sha256": hashlib.sha256(source.read_bytes()).hexdigest(),
+                      "tokenizer_sha256": hashlib.sha256((args.model / "tokenizer.json").read_bytes()).hexdigest(),
+                      "chat_template_sha256": hashlib.sha256(tokenizer.chat_template.encode()).hexdigest(),
+                      "readme_sha256": hashlib.sha256(readme.encode()).hexdigest(),
+                      "source_run_sha256": hashlib.sha256(args.source_run.read_bytes()).hexdigest() if args.source_run else None,
+                      "cases": cases}))
+
+
+if __name__ == "__main__":
+    main()

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()

tools/qwen-qmm-reference.py

@@ -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()

tools/test-supervisor.rs

@@ -0,0 +1,359 @@
+//! External watchdog for a prebuilt Rust test or a direct reference worker.
+//! Output is test progress, not proof of GPU progress. No total-runtime limit.
+#[cfg(target_os = "macos")]
+#[path = "../src/process_resources.rs"]
+mod process_resources;
+
+#[cfg(target_os = "macos")]
+mod supervisor {
+    use super::process_resources::{ResourceUsage, resource_usage};
+    use serde_json::json;
+    use std::io::{self, Read, Write};
+    use std::os::unix::process::CommandExt;
+    use std::process::{Child, Command, Stdio};
+    use std::sync::{Arc, Mutex};
+    use std::thread;
+    use std::time::{Duration, Instant};
+
+    const SAMPLE_INTERVAL: Duration = Duration::from_millis(100);
+
+    pub(super) struct Limits {
+        pub(super) bytes: u64,
+        pub(super) start: Duration,
+        pub(super) idle: Duration,
+    }
+
+    struct Worker(Child);
+
+    impl Drop for Worker {
+        fn drop(&mut self) {
+            if !matches!(self.0.try_wait(), Ok(Some(_))) {
+                // This process group was created for this child alone. Also
+                // clean up helpers on error; memory accounting covers the test
+                // process, so this launcher is not a process-tree supervisor.
+                unsafe { libc::kill(-(self.0.id() as i32), libc::SIGKILL) };
+                let _ = self.0.kill();
+                let _ = self.0.wait();
+            }
+        }
+    }
+
+    #[derive(Default)]
+    struct Output {
+        last: Option<Instant>,
+        failed: bool,
+    }
+
+    // Worker-side probes are observations, not inference progress. Buffer JSON
+    // across pipe reads so a split heartbeat cannot keep a stuck worker alive.
+    // Plain output (including Rust test banners without a newline) still counts.
+    fn useful_output(pending: &mut Vec<u8>, bytes: &[u8]) -> bool {
+        pending.extend_from_slice(bytes);
+        let mut progressed = false;
+        while let Some(end) = pending.iter().position(|&b| b == b'\n') {
+            let line = &pending[..end];
+            let observation = serde_json::from_slice::<serde_json::Value>(line)
+                .ok()
+                .is_some_and(|value| {
+                    matches!(
+                        value["event"].as_str(),
+                        Some(
+                            "gpu_canary_sample"
+                                | "gpu_canary_summary"
+                                | "reference_canary_summary"
+                                | "test_resource_sample"
+                        )
+                    )
+                });
+            progressed |= !observation && line.iter().any(|b| !b.is_ascii_whitespace());
+            pending.drain(..=end);
+        }
+        if pending.len() > 64 * 1024
+            || pending
+                .iter()
+                .find(|b| !b.is_ascii_whitespace())
+                .is_some_and(|&b| b != b'{')
+        {
+            progressed = true;
+            pending.clear();
+        }
+        progressed
+    }
+
+    fn forward(
+        mut input: impl Read + Send + 'static,
+        mut output: impl Write + Send + 'static,
+        progress: Arc<Mutex<Output>>,
+    ) -> thread::JoinHandle<()> {
+        thread::spawn(move || {
+            let mut bytes = [0; 4096];
+            let mut pending = Vec::new();
+            loop {
+                match input.read(&mut bytes) {
+                    Ok(0) => return,
+                    Ok(n) => {
+                        if output
+                            .write_all(&bytes[..n])
+                            .and_then(|()| output.flush())
+                            .is_err()
+                        {
+                            progress.lock().unwrap().failed = true;
+                            return;
+                        }
+                        if useful_output(&mut pending, &bytes[..n]) {
+                            progress.lock().unwrap().last = Some(Instant::now());
+                        }
+                    }
+                    Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
+                    Err(_) => {
+                        progress.lock().unwrap().failed = true;
+                        return;
+                    }
+                }
+            }
+        })
+    }
+
+    pub(super) fn run(command: &mut Command, limits: Limits) -> Result<(), String> {
+        if limits.bytes == 0 || limits.start.is_zero() || limits.idle.is_zero() {
+            return Err("memory and progress limits must be positive".into());
+        }
+        // Fail before spawning if process accounting is unavailable.
+        resource_usage(std::process::id()).ok_or("process accounting unavailable")?;
+        let mut worker = Worker(
+            command
+                .process_group(0)
+                .stdin(Stdio::null())
+                .stdout(Stdio::piped())
+                .stderr(Stdio::piped())
+                .spawn()
+                .map_err(|e| format!("could not start test: {e}"))?,
+        );
+        let progress = Arc::new(Mutex::new(Output::default()));
+        let readers = [
+            forward(
+                worker.0.stdout.take().unwrap(),
+                io::stdout(),
+                progress.clone(),
+            ),
+            forward(
+                worker.0.stderr.take().unwrap(),
+                io::stderr(),
+                progress.clone(),
+            ),
+        ];
+        let started = Instant::now();
+        let mut peak = ResourceUsage::default();
+        eprintln!(
+            "{}",
+            json!({"event":"test_supervisor_started", "pid":worker.0.id(),
+            "limit_bytes":limits.bytes, "start_seconds":limits.start.as_secs_f64(),
+            "idle_seconds":limits.idle.as_secs_f64(), "sample_ms":100})
+        );
+        let outcome = loop {
+            if let Some(status) = worker.0.try_wait().map_err(|e| e.to_string())? {
+                break if status.success() {
+                    Ok(())
+                } else {
+                    Err(format!("test exited with {status}"))
+                };
+            }
+            let now = Instant::now();
+            let Some(sample) = resource_usage(worker.0.id()) else {
+                // Exit can race the sample; missing telemetry on a live worker
+                // is a failure, never permission to keep running unmonitored.
+                if let Some(status) = worker.0.try_wait().map_err(|e| e.to_string())? {
+                    break if status.success() {
+                        Ok(())
+                    } else {
+                        Err(format!("test exited with {status}"))
+                    };
+                }
+                break Err("process accounting lost".into());
+            };
+            peak.include(sample);
+            eprintln!(
+                "{}",
+                json!({"event":"test_resource_sample",
+                "elapsed_ms":started.elapsed().as_millis(), "usage":sample})
+            );
+            if sample.physical_bytes > limits.bytes || sample.peak_physical_bytes > limits.bytes {
+                break Err("memory_limit".into());
+            }
+            let output = progress.lock().map_err(|e| e.to_string())?;
+            if output.failed {
+                break Err("output monitoring failed".into());
+            }
+            let (last, timeout, reason) = output
+                .last
+                .map_or((started, limits.start, "start_timeout"), |last| {
+                    (last, limits.idle, "continuation_timeout")
+                });
+            if now.saturating_duration_since(last) >= timeout {
+                break Err(reason.into());
+            }
+            drop(output);
+            thread::sleep(SAMPLE_INTERVAL);
+        };
+        drop(worker); // Kill and reap on every failure, including output errors.
+        for reader in readers {
+            reader.join().map_err(|_| "output reader panicked")?;
+        }
+        eprintln!(
+            "{}",
+            json!({"event":"test_resource_summary",
+            "elapsed_ms":started.elapsed().as_millis(), "peak":peak,
+            "error":outcome.as_ref().err()})
+        );
+        if progress.lock().map_err(|e| e.to_string())?.failed {
+            return Err("output monitoring failed".into());
+        }
+        outcome
+    }
+
+    pub(super) fn cli() -> Result<(), String> {
+        let args = std::env::args_os().skip(1).collect::<Vec<_>>();
+        if args.len() < 5 || (args[3] != "--command" && args.len() != 5) {
+            return Err("usage: test-supervisor MEMORY_MIB START_SECONDS IDLE_SECONDS TEST_BINARY TEST_FILTER\nOr: test-supervisor MEMORY_MIB START_SECONDS IDLE_SECONDS --command EXECUTABLE [ARG ...]\nMonitors the direct worker only; do not pass cargo or a spawning shell.".into());
+        }
+        let number = |i: usize| {
+            args[i]
+                .to_str()
+                .and_then(|s| s.parse::<u64>().ok())
+                .filter(|n| *n > 0)
+                .ok_or("limits must be positive integers")
+        };
+        let limits = Limits {
+            bytes: number(0)?
+                .checked_mul(1024 * 1024)
+                .ok_or("memory limit overflow")?,
+            start: Duration::from_secs(number(1)?),
+            idle: Duration::from_secs(number(2)?),
+        };
+        let mut command = worker_command(&args[3..])?;
+        run(&mut command, limits)
+    }
+
+    fn worker_command(args: &[std::ffi::OsString]) -> Result<Command, String> {
+        if args.len() < 2 || args[1].is_empty() {
+            return Err("worker executable or test filter must not be empty".into());
+        }
+        if args[0] == "--command" {
+            let mut command = Command::new(&args[1]);
+            command.args(&args[2..]);
+            Ok(command)
+        } else if args.len() == 2 {
+            let mut command = Command::new(&args[0]);
+            command
+                .arg(&args[1])
+                .args(["--include-ignored", "--test-threads=1", "--nocapture"]);
+            Ok(command)
+        } else {
+            Err("unexpected test arguments".into())
+        }
+    }
+
+    #[cfg(test)]
+    mod tests {
+        use super::*;
+
+        #[test]
+        fn direct_worker_arguments_are_not_rust_test_arguments() {
+            let args = |a: &[&str]| a.iter().map(std::ffi::OsString::from).collect::<Vec<_>>();
+            let direct = worker_command(&args(&[
+                "--command",
+                "/usr/bin/python3",
+                "-u",
+                "reference.py",
+            ]))
+            .unwrap();
+            assert_eq!(direct.get_program(), "/usr/bin/python3");
+            assert_eq!(
+                direct.get_args().collect::<Vec<_>>(),
+                ["-u", "reference.py"]
+            );
+            let test = worker_command(&args(&["test-binary", "fixture"])).unwrap();
+            assert_eq!(
+                test.get_args().collect::<Vec<_>>(),
+                [
+                    "fixture",
+                    "--include-ignored",
+                    "--test-threads=1",
+                    "--nocapture"
+                ]
+            );
+            assert!(worker_command(&args(&["--command", ""])).is_err());
+            assert!(worker_command(&args(&["test-binary", "fixture", "extra"])).is_err());
+        }
+
+        #[test]
+        fn watchdog_stops_memory_and_silence_but_not_total_runtime() {
+            let limits = || Limits {
+                bytes: 256 * 1024 * 1024,
+                start: Duration::from_millis(300),
+                idle: Duration::from_millis(300),
+            };
+            let mut sleeping = Command::new("/bin/sleep");
+            sleeping.arg("5");
+            assert_eq!(run(&mut sleeping, limits()).unwrap_err(), "start_timeout");
+            let mut started = Command::new("/bin/sh");
+            started.args(["-c", "printf ready; exec /bin/sleep 5"]);
+            assert_eq!(
+                run(&mut started, limits()).unwrap_err(),
+                "continuation_timeout"
+            );
+            for (prefix, expected) in [
+                ("", "start_timeout"),
+                ("printf 'ready\\n'; ", "continuation_timeout"),
+            ] {
+                let mut heartbeat = Command::new("/bin/sh");
+                heartbeat.args(["-c", &format!("{prefix}while :; do printf '{{\"event\":\"gpu_canary_sample\",\"ok\":true}}\\n'; sleep 0.05; done")]);
+                assert_eq!(run(&mut heartbeat, limits()).unwrap_err(), expected);
+            }
+            let mut bounded = limits();
+            bounded.bytes = 1;
+            assert_eq!(run(&mut sleeping, bounded).unwrap_err(), "memory_limit");
+            let mut progressing = Command::new("/bin/sh");
+            progressing.args([
+                "-c",
+                "for i in 1 2 3 4 5 6 7 8; do printf progress; sleep 0.1; done",
+            ]);
+            let before = Instant::now();
+            run(&mut progressing, limits()).unwrap();
+            assert!(before.elapsed() >= Duration::from_millis(600));
+            let mut failure = Command::new("/bin/sh");
+            failure.args(["-c", "exit 7"]);
+            assert!(
+                run(&mut failure, limits())
+                    .unwrap_err()
+                    .contains("exit status: 7")
+            );
+        }
+
+        #[test]
+        fn split_probe_records_do_not_reset_progress() {
+            let mut pending = Vec::new();
+            assert!(!useful_output(&mut pending, b"{\"event\":\"gpu_canary_"));
+            assert!(!useful_output(&mut pending, b"sample\",\"ok\":true}\n"));
+            assert!(useful_output(
+                &mut pending,
+                b"{\"event\":\"mtp_tokens\",\"ids\":[42]}\n"
+            ));
+            assert!(useful_output(&mut pending, b"progress"));
+            assert!(pending.is_empty());
+        }
+    }
+}
+
+fn main() {
+    #[cfg(target_os = "macos")]
+    if let Err(error) = supervisor::cli() {
+        eprintln!("Test supervisor: {error}");
+        std::process::exit(1);
+    }
+    #[cfg(not(target_os = "macos"))]
+    {
+        eprintln!("test-supervisor requires macOS process accounting");
+        std::process::exit(1);
+    }
+}

tools/test_mtplx_execution_reference.py

@@ -0,0 +1,62 @@
+"""Model-free checks for the matched-reference input contract."""
+import copy
+import contextlib
+import io
+import json
+import importlib.util
+from pathlib import Path
+import unittest
+from unittest.mock import patch
+
+spec = importlib.util.spec_from_file_location("oracle", Path(__file__).with_name("mtplx-execution-reference.py"))
+oracle = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(oracle)
+
+
+class SourceContract(unittest.TestCase):
+    def test_progress_requires_new_tokens_and_limits_output_only(self):
+        output = io.StringIO()
+        with contextlib.redirect_stderr(output), patch.object(oracle.time, "monotonic", side_effect=[0, .4, 1]):
+            report = oracle.token_report("chat_tokens", "progress", step=0)
+            report([])
+            report([10])
+            report([11, 12])
+            report([13])
+        events = [json.loads(line) for line in output.getvalue().splitlines()]
+        self.assertEqual([e["tokens"] for e in events], [1, 4])
+        self.assertTrue(all(e["event"] == "chat_decode_progress" and "ids" not in e for e in events))
+
+    def test_requires_matching_complete_warm_chat(self):
+        start = dict(event="start", model="qwen3.8-flash-next", plain_chat=True,
+            system_prompt="", canary=True, prompts=["summary", "story", "python"],
+            settings=dict(reasoning="low", power_percent=100, prefill_chunk=2048,
+                min_p=0, quality=False, seed=42, acceleration=dict(kind="mtp", enabled=True)),
+            warmup=dict(enabled=True))
+        records = [start, dict(event="warmup_result", ok=True)] + [
+            dict(event="result", turn=n, ok=True, finish_reason="stop") for n in (1, 2, 3)]
+        self.assertEqual(oracle.source_chat(records, True, True)[0], start)
+        for records_bad, mtp, canary in [
+            (records[:-1], True, True),
+            (records + [start], True, True),
+            (records, False, True),
+            (records, True, False),
+        ]:
+            with self.assertRaises(ValueError):
+                oracle.source_chat(records_bad, mtp, canary)
+        for location, name, value in [
+            ("settings", "power_percent", 75),
+            ("settings", "seed", None),
+            ("warmup", "enabled", False),
+        ]:
+            bad = copy.deepcopy(records)
+            bad[0][location][name] = value
+            with self.assertRaises(ValueError):
+                oracle.source_chat(bad, True, True)
+        bad = copy.deepcopy(records)
+        bad[-1]["finish_reason"] = "length"
+        with self.assertRaises(ValueError):
+            oracle.source_chat(bad, True, True)
+
+
+if __name__ == "__main__":
+    unittest.main()