modelbeast/scripts/perfcheck/bench.py

"""PERFCHECK — per-node GPU performance canary. Prints one JSON object to stdout.

Runs on ``MODELBEAST/venvs/rmbg/bin/python``, which carries an identical torch on every node
in the fleet — so numbers are comparable machine-to-machine without installing anything new
(nothing lands on the disk-critical m1max).

Probes (~20s, sized to survive the 8GB M1 mini):

  matmul_fp16 / matmul_fp32   GFLOP/s   — raw compute
  bandwidth                   GB/s      — unified-memory read/write
  sdpa_64 / sdpa_56           ms        — attention at a fast-path head_dim (64) vs the
                                          awkward one CorridorKey's Hiera actually uses (56)

The sdpa pair is the point. CorridorKey is slow on Apple Silicon because head_dim=56 misses
the fused attention path; padding to 64 buys 1.3x-4.9x depending on the machine. ``sdpa_ratio``
(56/64) tracks that cliff per node, so if a future torch closes it we find out — instead of
carrying the padding workaround forever without noticing it became unnecessary.

MEASUREMENT NOTES — these are load-bearing, do not "simplify" them:

* **min, not mean/median.** These nodes serve jobs. A concurrent trellis_mac run on the M3 Ultra
  dragged a median-of-5 matmul from ~24500 down to ~2150 GFLOP/s — a 5x error that looks
  exactly like a catastrophic regression. The minimum time is the least-interfered sample and
  the only estimate of the machine's actual capability that survives a noisy neighbour.
* **Sizes are chosen for stability, not for speed.** n=4096 reproduced to within 0.1% across
  runs *while the GPU was contended*; n=2048 swung 48% (11700 vs 17361) because it is
  dispatch-bound rather than compute-bound. Shrinking these to make the bench faster silently
  turns it back into a random number generator.
* **Every probe syncs MPS before stopping the clock.** MPS dispatch is async; an unsynced timer
  measures queue-submission, not work.
* ``load1`` is reported so a consumer can distrust a sample taken on a busy node. Under
  contention the mins mostly hold up, but a hard-loaded node can still read low.
"""

from __future__ import annotations

import json
import os
import platform
import subprocess
import sys
import time

os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")

import torch  # noqa: E402

# 2: matmul 2048->4096, sdpa 8x1024 -> 16x2048, median->min, added gpu_busy_pre.
SCHEMA = 2


def _sync() -> None:
    torch.mps.synchronize()


def _best(fn, runs: int = 7, warmup: int = 3) -> float:
    """Fastest of ``runs`` seconds-per-call, MPS-synced. See MEASUREMENT NOTES on why min."""
    for _ in range(warmup):
        fn()
    _sync()
    best = float("inf")
    for _ in range(runs):
        t0 = time.perf_counter()
        fn()
        _sync()
        best = min(best, time.perf_counter() - t0)
    return best


def bench_matmul(dtype: torch.dtype, n: int = 4096) -> float:
    """GFLOP/s for an n^3 matmul. n=4096 is the smallest size that is compute- not
    dispatch-bound here (~67MB/tensor fp32 — fine on the 8GB mini)."""
    a = torch.randn(n, n, device="mps", dtype=dtype)
    b = torch.randn(n, n, device="mps", dtype=dtype)
    secs = _best(lambda: a @ b)
    return (2.0 * n**3) / secs / 1e9


def bench_bandwidth(mb: int = 512) -> float:
    """GB/s. x.add_(1) reads and writes the whole buffer -> 2x bytes moved.

    runs=25 rather than the default: each call is only ~1.6ms, so extra samples are nearly
    free and they reject neighbour interference. At runs=9 this spread 16% on a contended
    node (479 vs 570) — enough to trip a 20% drift alarm on its own. At runs=25: 0.3%.
    """
    n = (mb * 1024 * 1024) // 4
    x = torch.randn(n, device="mps", dtype=torch.float32)
    secs = _best(lambda: x.add_(1.0), runs=25)
    return (2.0 * n * 4) / secs / 1e9


def bench_sdpa(head_dim: int, heads: int = 16, seq: int = 2048) -> float:
    """Milliseconds for one fused-attention call at the given head_dim.

    heads=16/seq=2048 is deliberate: it lands ~1-4ms, where the ratio reproduces to 0.6%
    (4.28/4.30/4.33 across runs on a contended M3 Ultra). The obvious smaller 8x1024 is
    sub-millisecond and pure dispatch noise — it produced ratios of 0.79, 3.95 and 2.35 on
    three consecutive runs of the same machine, i.e. it would have "proven" head_dim=56 is
    *faster*, contradicting the real 8-machine result.
    """
    shape = (1, heads, seq, head_dim)
    q = torch.randn(*shape, device="mps", dtype=torch.float16)
    k = torch.randn(*shape, device="mps", dtype=torch.float16)
    v = torch.randn(*shape, device="mps", dtype=torch.float16)
    secs = _best(lambda: torch.nn.functional.scaled_dot_product_attention(q, k, v), runs=9)
    return secs * 1000.0


def gpu_busy_pre() -> int:
    """Peak GPU utilization % observed over ~1.2s, sampled BEFORE any torch work starts.

    Must run before we touch the GPU: our own matmul pins the device at ~100%, so this can
    only tell us about a *neighbour* while we are still idle. That ordering is the whole trick
    — sampling during the bench is worthless.

    ``load1`` cannot substitute: the M3 Ultra sat at load 2.45 with its GPU fully held by a
    trellis_mac job. Conversely a GPU job alternates phases (100% CPU / 0% GPU one moment,
    97% GPU the next), so a single reading misses it — hence max-of-several. The first ioreg
    read is discarded because the counter is since-last-read and reports garbage when cold (an
    idle m4mini reported 97% on its first sample, 0% thereafter).
    """
    def sample() -> int:
        try:
            out = subprocess.run(
                ["ioreg", "-r", "-d", "1", "-w", "0", "-c", "IOAccelerator"],
                capture_output=True, text=True, timeout=8,
            ).stdout
        except Exception:
            return 0
        vals = [int(s.split("=")[1]) for s in out.split(",")
                if s.strip().startswith('"Device Utilization %"')]
        return max(vals) if vals else 0

    sample()  # discard the cold read
    peak = 0
    for _ in range(4):
        time.sleep(0.3)
        peak = max(peak, sample())
    return peak


def _sysctl(key: str) -> str:
    try:
        return subprocess.run(
            ["sysctl", "-n", key], capture_output=True, text=True, timeout=5
        ).stdout.strip()
    except Exception:
        return ""


def host_facts() -> dict:
    mem = _sysctl("hw.memsize")
    return {
        "host": platform.node().split(".")[0],
        "chip": _sysctl("machdep.cpu.brand_string"),
        "cores": _sysctl("hw.ncpu"),
        "ram_gb": round(int(mem) / 1024**3) if mem.isdigit() else None,
        "macos": platform.mac_ver()[0],
        "torch": torch.__version__,
        "load1": round(os.getloadavg()[0], 2),
    }


def main() -> int:
    if not torch.backends.mps.is_available():
        print(json.dumps({"error": "MPS unavailable", **host_facts()}))
        return 1

    # Before any torch work — see gpu_busy_pre().
    busy = gpu_busy_pre()

    # BUMP `schema` whenever a probe's shape/size changes. Numbers are only comparable within
    # a schema, and run_fleet drops older samples from the median rather than averaging a
    # 2048-matmul against a 4096-matmul and calling the difference a regression.
    out = {"schema": SCHEMA, **host_facts(), "gpu_busy_pre": busy, "metrics": {}}
    m = out["metrics"]
    m["matmul_fp16_gflops"] = round(bench_matmul(torch.float16), 1)
    m["matmul_fp32_gflops"] = round(bench_matmul(torch.float32), 1)
    m["bandwidth_gbs"] = round(bench_bandwidth(), 1)
    sdpa64 = bench_sdpa(64)
    sdpa56 = bench_sdpa(56)
    m["sdpa_64_ms"] = round(sdpa64, 3)
    m["sdpa_56_ms"] = round(sdpa56, 3)
    # >1 means the awkward head_dim is slower, i.e. CorridorKey's padding win still stands.
    m["sdpa_ratio"] = round(sdpa56 / sdpa64, 2) if sdpa64 else None
    out["load1_end"] = round(os.getloadavg()[0], 2)

    print(json.dumps(out))
    return 0


if __name__ == "__main__":
    sys.exit(main())