mlx/scripts/benchmark.py

"""Full-model parity and measurements on this Mac; emits no simulated hardware results."""

import argparse
import json
import time
from pathlib import Path

import mlx.core as mx
import numpy as np

from solomon_mlx import Solomon
from solomon_mlx._vendor.semantics import listed_probs, p_yes
from solomon_mlx.artifacts import digest, sha256


def probabilities(job, row):
    if job["task"] in ("boolean", "multilabel", "entity"):
        p = p_yes(row["letter_logits"])
        return np.array([p, 1 - p])
    n = job["n"] - 2 if job["head_key"].endswith("choiceR") else job["n"]
    return listed_probs(row["letter_logits"], n)


def run(model_dir, jobs_path, reference_path, output):
    output = Path(output)
    if output.exists():
        raise FileExistsError("Benchmark outputs are immutable")
    jobs = json.loads(Path(jobs_path).read_text())
    reference = json.loads(Path(reference_path).read_text())
    ref = {r["id"]: r for r in reference["rows"]}
    if set(ref) != {j["id"] for j in jobs}:
        raise ValueError("Benchmark and reference jobs differ")
    mx.reset_peak_memory()
    started = time.perf_counter()
    model = Solomon.load(model_dir)
    load_seconds = time.perf_counter() - started
    states = {}
    rows = []
    prefills = []
    try:
        for job in jobs:
            key = digest(job["parts"])
            if key not in states:
                state = model.prefill(job["parts"])
                states[key] = state
                prefills.append(
                    {
                        "document": key,
                        "tokens": state.prefix_tokens,
                        "seconds": state._data["prefill_seconds"],
                        "vision_seconds": state._data["vision_seconds"],
                        "cache_bytes": sum(c.nbytes for c in state._data["cache"]),
                    }
                )
            state = states[key]
            row = model.engine.ask(
                state._data,
                job["block"],
                job["n"],
                job["head_key"],
                execution=job.get("execution", "cached"),
                taps=job.get("taps", []),
            )
            p, q = probabilities(job, row), probabilities(job, ref[job["id"]])
            row.update(
                id=job["id"],
                decision_agrees=bool(p.argmax() == q.argmax()),
                max_probability_drift=float(np.max(np.abs(p - q))),
                max_logit_drift=float(
                    np.max(np.abs(np.array(row["letter_logits"]) - ref[job["id"]]["letter_logits"]))
                ),
                prefix_ids_exact=state._data["prefix_ids"] == ref[job["id"]]["prefix_ids"],
            )
            if "token_ids" in row:
                row["token_ids_exact"] = row["token_ids"] == ref[job["id"]]["token_ids"]
            if row.get("taps"):
                row["layer_max_hidden_drift"] = {
                    k: float(np.max(np.abs(np.array(v) - ref[job["id"]]["taps"][k])))
                    for k, v in row["taps"].items()
                }
            rows.append(row)
            print(job["id"], row["seconds"], row["decision_agrees"], flush=True)
        # Replay checks real prefill and repeated question semantics on the same binding.
        first = states[digest(jobs[0]["parts"])]
        recipe = output.with_suffix(".replay.json")
        first.save(recipe)
        with model.replay(recipe) as restored:
            b, n, h = jobs[0]["block"], jobs[0]["n"], jobs[0]["head_key"]
            replay = model.engine.ask(restored._data, b, n, h)
            replay_drift = float(np.max(np.abs(np.array(replay["letter_logits"]) - rows[0]["letter_logits"])))
    finally:
        for state in states.values():
            state.close()
    warm = [r["seconds"] for r in rows if r["reused_prefix_tokens"]]
    report = {
        "runtime": model.identity,
        "device": mx.device_info(),
        "jobs_sha256": sha256(jobs_path),
        "reference_sha256": sha256(reference_path),
        "load_seconds": load_seconds,
        "prefills": prefills,
        "rows": rows,
        "warm_question_latency_median_seconds": float(np.median(warm)),
        "questions_per_second": len(warm) / sum(warm),
        "peak_metal_bytes": mx.get_peak_memory(),
        "replay_max_logit_drift": replay_drift,
        "decision_agreement": float(np.mean([r["decision_agrees"] for r in rows])),
        "scope": "development parity fixtures; not held-out task accuracy or release qualification",
        "calibration_status": "uncalibrated",
    }
    output.write_text(json.dumps(report, indent=2))
    return report


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="models/quality")
    p.add_argument("--jobs", default="evaluations/golden-jobs.json")
    p.add_argument("--reference", default="evaluations/bf16-reference-1789901869/report.json")
    p.add_argument("--output", default="evaluations/bf16-text-benchmark.json")
    a = p.parse_args()
    run(a.model, a.jobs, a.reference, a.output)