diff --git a/bench/bench_subm.py b/bench/bench_subm.py new file mode 100644 index 0000000..e768ba0 --- /dev/null +++ b/bench/bench_subm.py @@ -0,0 +1,128 @@ +"""Benchmark SubMConv3d across the fleet. Needs only mlx + numpy. + +Submanifold conv is the hot op in the V-VAE, so its scaling across Apple Silicon +generations is what decides which box should run LATO.2 once the port lands. Two +distinct costs are timed separately because they scale differently: + + indice map — coordinate hashing + binary search, runs on CPU in numpy, cacheable + per coordinate set (upstream reuses it via `indice_key`) + conv — 27 gathers + 27 matmuls on the GPU, paid every layer + +Occupancy is set to ~10% of the grid, roughly what a voxelised surface gives you. +""" + +from __future__ import annotations + +import argparse +import json +import platform +import subprocess +import sys +import time +from pathlib import Path + +import mlx.core as mx +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from lato_mlx.sparse.conv import SubMConv3d, build_indice_map # noqa: E402 +from lato_mlx.sparse.tensor import SparseTensor # noqa: E402 + + +def machine() -> dict: + def sysctl(k): + try: + return subprocess.run( + ["sysctl", "-n", k], capture_output=True, text=True, timeout=5 + ).stdout.strip() + except Exception: # noqa: BLE001 + return "?" + + mem = sysctl("hw.memsize") + return { + "host": platform.node().split(".")[0], + "chip": sysctl("machdep.cpu.brand_string"), + "ram_gb": round(int(mem) / 1e9) if mem.isdigit() else "?", + "mlx": mx.__version__, + } + + +def make_sparse(res: int, channels: int, occupancy: float, seed: int = 0): + rng = np.random.default_rng(seed) + n_total = res**3 + n = max(64, int(n_total * occupancy)) + flat = rng.choice(n_total, size=n, replace=False) + flat.sort() + z, rem = np.divmod(flat, res * res) + y, x = np.divmod(rem, res) + coords = np.stack([np.zeros_like(z), z, y, x], axis=1).astype(np.int32) + feats = rng.standard_normal((n, channels)).astype(np.float32) + return mx.array(coords), mx.array(feats) + + +def time_it(fn, warmup=1, iters=3): + for _ in range(warmup): + fn() + ts = [] + for _ in range(iters): + t0 = time.perf_counter() + fn() + ts.append(time.perf_counter() - t0) + return min(ts) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--out", default=None, help="write JSON result here") + ap.add_argument("--occupancy", type=float, default=0.10) + ap.add_argument("--iters", type=int, default=3) + a = ap.parse_args() + + info = machine() + print(f"== {info['host']} {info['chip']} {info['ram_gb']}GB mlx {info['mlx']}") + print(f"{'res':>5} {'chan':>5} {'voxels':>9} {'imap_ms':>9} {'conv_ms':>9} {'Mvox/s':>8}") + + rows = [] + for res, ch in [(32, 64), (64, 64), (64, 128), (128, 128)]: + try: + coords, feats = make_sparse(res, ch, a.occupancy) + n = coords.shape[0] + conv = SubMConv3d(ch, ch, 3, bias=True, indice_key=f"b{res}") + + t_imap = time_it(lambda: build_indice_map(coords, 3), warmup=0, iters=1) + + x = SparseTensor(feats, coords) + _ = conv(x) # populate the indice cache so we time conv alone + mx.eval(_.feats) + + def run(): + mx.eval(conv(x).feats) + + t_conv = time_it(run, iters=a.iters) + rows.append( + { + "res": res, + "channels": ch, + "voxels": n, + "imap_ms": t_imap * 1e3, + "conv_ms": t_conv * 1e3, + "mvox_s": n / t_conv / 1e6, + } + ) + print( + f"{res:>5} {ch:>5} {n:>9,} {t_imap*1e3:>9.1f} {t_conv*1e3:>9.1f} " + f"{n/t_conv/1e6:>8.2f}" + ) + except Exception as e: # noqa: BLE001 + print(f"{res:>5} {ch:>5} FAILED: {type(e).__name__}: {e}") + rows.append({"res": res, "channels": ch, "error": f"{type(e).__name__}: {e}"}) + + result = {"machine": info, "occupancy": a.occupancy, "rows": rows} + if a.out: + Path(a.out).parent.mkdir(parents=True, exist_ok=True) + Path(a.out).write_text(json.dumps(result, indent=2)) + print(f"\nwrote {a.out}") + + +if __name__ == "__main__": + main() diff --git a/bench/run_fleet.sh b/bench/run_fleet.sh new file mode 100755 index 0000000..d77882a --- /dev/null +++ b/bench/run_fleet.sh @@ -0,0 +1,58 @@ +#!/bin/zsh +# Run the SubMConv3d benchmark on every Apple Silicon box in the fleet. +# Ships code only (no weights); each box gets a small mlx+numpy venv under ~/lato-bench. +set -u +SRC=/Users/m3ultra/Documents/lato.2_mrp_mlx +OUT=$SRC/bench +SSH_OPTS=(-o ConnectTimeout=10 -o BatchMode=yes -o StrictHostKeyChecking=accept-new) +SSH_STR="ssh -o ConnectTimeout=10 -o BatchMode=yes -o StrictHostKeyChecking=accept-new" + +# host:label (m3ultra runs locally, already done) +HOSTS=( + "m1max@100.92.78.24:m1max" + "m2max@100.120.83.110:m2max" + "100.69.21.128:m4pro" + "johnking@100.91.239.7:m1ultra" +) + +hb(){ mkdir -p ~/.jobs; echo "$(date '+%F %T') | $1 | $2" > ~/.jobs/lato-fleet-bench.status; } +hb "starting" "${#HOSTS[@]} hosts" + +for entry in $HOSTS; do + host=${entry%%:*}; label=${entry##*:} + echo "===== $label ($host)" + hb "$label" "connecting" + + if ! ssh "${SSH_OPTS[@]}" "$host" 'true' 2>/dev/null; then + echo " SKIP unreachable"; continue + fi + + # m1max is disk-critical (~20GB free per fleet notes) - refuse rather than fill it + free_mb=$(ssh "${SSH_OPTS[@]}" "$host" "df -m / | tail -1 | awk '{print \$4}'" 2>/dev/null) + if [ -n "$free_mb" ] && [ "$free_mb" -lt 3000 ]; then + echo " SKIP only ${free_mb}MB free - not installing a venv here"; continue + fi + + ssh "${SSH_OPTS[@]}" "$host" 'mkdir -p ~/lato-bench' 2>/dev/null + rsync -a -e "$SSH_STR" --delete \ + "$SRC/lato_mlx" "$SRC/bench" "$host:~/lato-bench/" 2>&1 | tail -1 + + hb "$label" "installing venv" + ssh "${SSH_OPTS[@]}" "$host" ' + cd ~/lato-bench || exit 1 + UV=$(command -v uv || echo /opt/homebrew/bin/uv) + if [ ! -x "$UV" ]; then echo " no uv on this host"; exit 3; fi + [ -d .venv ] || "$UV" venv --python 3.12 .venv >/dev/null 2>&1 + VIRTUAL_ENV=.venv "$UV" pip install -q mlx numpy >/dev/null 2>&1 + .venv/bin/python -c "import mlx.core" 2>/dev/null || { echo " mlx install failed"; exit 4; } + ' || { echo " SKIP setup failed (rc=$?)"; continue; } + + hb "$label" "benchmarking" + ssh "${SSH_OPTS[@]}" "$host" \ + 'cd ~/lato-bench && .venv/bin/python bench/bench_subm.py --out bench/result.json' 2>&1 | tail -8 + rsync -a -e "$SSH_STR" "$host:~/lato-bench/bench/result.json" "$OUT/$label.json" 2>/dev/null \ + && echo " -> $OUT/$label.json" || echo " (no result pulled)" +done + +hb "DONE" "results in $OUT" +echo "===== done" diff --git a/lato_mlx/sparse/conv.py b/lato_mlx/sparse/conv.py index 8f1d520..38daee7 100644 --- a/lato_mlx/sparse/conv.py +++ b/lato_mlx/sparse/conv.py @@ -130,16 +130,45 @@ class SubMConv3d(nn.Module): imap = self._indice_map(x) # One appended zero row: absent neighbours index it and contribute nothing, - # which avoids a per-offset boolean mask in the accumulation loop. + # which avoids a per-offset boolean mask. feats_pad = mx.concatenate( [x.feats, mx.zeros((1, self.in_channels), dtype=x.feats.dtype)], axis=0 ) - idx = np.where(imap == _MISSING, n, imap) + idx = np.where(imap == _MISSING, n, imap) # [K^3, N] + k3 = imap.shape[0] - out = mx.zeros((n, self.out_channels), dtype=x.feats.dtype) - for i in range(imap.shape[0]): - gathered = mx.take(feats_pad, mx.array(idx[i]), axis=0) - out = out + gathered @ self.weight[i].astype(x.feats.dtype) + # Fuse the K^3 taps into ONE gather + ONE matmul. + # + # The obvious implementation loops over the K^3 offsets accumulating + # `gather(i) @ W[i]`, but that issues 2*K^3 tiny GPU dispatches with Python + # between them, and each is far too small to fill the machine. Fleet + # benchmarking made this unmistakable: the 80-core M3 Ultra came in SLOWER + # than a 38-core M2 Max (81.6ms vs 58.3ms at 128^3/128ch), i.e. throughput was + # anti-correlated with core count — the signature of launch-latency binding + # rather than compute binding. + # + # Concatenating neighbours along the channel axis turns the whole thing into a + # single [N, K^3*Cin] x [K^3*Cin, Cout] matmul, which is one dispatch big + # enough to actually occupy the GPU. + # + # That buffer is N*K^3*Cin floats, so it is chunked over rows to keep peak + # memory bounded (~256MB/chunk) — at 128^3 and 128ch the unchunked form alone + # would be ~2.9GB, which is fine on a Studio and not fine on an 8GB mini. + w_flat = self.weight.reshape(k3 * self.in_channels, self.out_channels).astype( + x.feats.dtype + ) + idx_t = mx.array(idx.T) # [N, K^3] + + bytes_per_row = k3 * self.in_channels * 4 + chunk = max(1, min(n, (256 << 20) // max(bytes_per_row, 1))) + + outs = [] + for start in range(0, n, chunk): + stop = min(start + chunk, n) + g = mx.take(feats_pad, idx_t[start:stop].reshape(-1), axis=0) + g = g.reshape(stop - start, k3 * self.in_channels) + outs.append(g @ w_flat) + out = outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=0) if hasattr(self, "bias"): out = out + self.bias.astype(x.feats.dtype)