91 lines
2.9 KiB
Python
91 lines
2.9 KiB
Python
"""MLX fleet probe — one JSON line of capability + micro-bench per machine.
|
|
|
|
Benches (each ~seconds, sized for 16GB boxes):
|
|
matmul fp16/bf16 4096^2 -> TFLOPS (bf16-on-M1 question)
|
|
elementwise add 256MB -> GB/s (memory bandwidth proxy)
|
|
sdpa head_dim 56 vs 64 -> ratio (fused fast-path check, corridorkey lesson)
|
|
"""
|
|
import json
|
|
import platform
|
|
import subprocess
|
|
import time
|
|
|
|
R = {"host": platform.node().split(".")[0]}
|
|
try:
|
|
R["chip"] = subprocess.check_output(
|
|
["sysctl", "-n", "machdep.cpu.brand_string"], text=True).strip()
|
|
R["ram_gb"] = round(int(subprocess.check_output(
|
|
["sysctl", "-n", "hw.memsize"], text=True)) / 2**30)
|
|
R["macos"] = subprocess.check_output(
|
|
["sw_vers", "-productVersion"], text=True).strip()
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
import mlx.core as mx
|
|
R["mlx"] = getattr(mx, "__version__", "?")
|
|
except Exception as e:
|
|
R["mlx"] = None
|
|
R["error"] = f"mlx import failed: {e}"
|
|
print(json.dumps(R))
|
|
raise SystemExit(0)
|
|
|
|
try:
|
|
di = mx.metal.device_info() if hasattr(mx, "metal") else {}
|
|
R["gpu"] = {k: di[k] for k in ("architecture", "max_recommended_working_set_size",
|
|
"memory_size") if k in di}
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def timeit(fn, warmup=3, iters=10):
|
|
for _ in range(warmup):
|
|
mx.eval(fn())
|
|
t0 = time.perf_counter()
|
|
for _ in range(iters):
|
|
mx.eval(fn())
|
|
return (time.perf_counter() - t0) / iters
|
|
|
|
|
|
N = 4096
|
|
flops = 2 * N * N * N
|
|
for dt, name in ((mx.float16, "fp16"), (mx.bfloat16, "bf16")):
|
|
try:
|
|
a = mx.random.normal((N, N)).astype(dt)
|
|
b = mx.random.normal((N, N)).astype(dt)
|
|
mx.eval(a, b)
|
|
t = timeit(lambda: a @ b)
|
|
R[f"matmul_{name}_tflops"] = round(flops / t / 1e12, 2)
|
|
except Exception as e:
|
|
R[f"matmul_{name}_tflops"] = f"fail: {e}"
|
|
|
|
try:
|
|
M = 64 * 1024 * 1024 # 64M floats = 256MB per array
|
|
x = mx.random.normal((M,))
|
|
y = mx.random.normal((M,))
|
|
mx.eval(x, y)
|
|
t = timeit(lambda: x + y, warmup=2, iters=8)
|
|
R["bandwidth_gbs"] = round(3 * M * 4 / t / 1e9) # 2 reads + 1 write, fp32
|
|
except Exception as e:
|
|
R["bandwidth_gbs"] = f"fail: {e}"
|
|
|
|
try:
|
|
B, H, L = 1, 8, 2048
|
|
sdpa = {}
|
|
for D in (56, 64):
|
|
q = mx.random.normal((B, H, L, D)).astype(mx.float16)
|
|
k = mx.random.normal((B, H, L, D)).astype(mx.float16)
|
|
v = mx.random.normal((B, H, L, D)).astype(mx.float16)
|
|
mx.eval(q, k, v)
|
|
t = timeit(lambda: mx.fast.scaled_dot_product_attention(
|
|
q, k, v, scale=D ** -0.5), warmup=3, iters=20)
|
|
sdpa[f"d{D}_us"] = round(t * 1e6)
|
|
sdpa["d56_vs_d64"] = round(sdpa["d56_us"] / max(sdpa["d64_us"], 1), 2)
|
|
R["sdpa"] = sdpa
|
|
except Exception as e:
|
|
R["sdpa"] = f"fail: {e}"
|
|
|
|
R["peak_mem_gb"] = round(mx.get_peak_memory() / 2**30, 2) \
|
|
if hasattr(mx, "get_peak_memory") else None
|
|
print(json.dumps(R))
|