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