Fuse the K^3 taps into one gather + matmul (2.1x on M3 Ultra)
Fleet benchmarking exposed the problem: throughput was ANTI-correlated with GPU core count. The 80-core M3 Ultra came in slowest at 128^3/128ch (81.6ms) behind a 38-core M2 Max (58.3ms), M1 Ultra (66.7ms) and even a 32-core M1 Max (69.2ms). That ordering only makes sense if the op is bound by dispatch latency rather than compute - the per-offset loop issued 2*K^3 = 54 tiny GPU ops per layer, none big enough to occupy the machine, and the Ultra's fused-die design penalises exactly that. Concatenating the K^3 neighbour taps along the channel axis collapses it to a single [N, K^3*Cin] x [K^3*Cin, Cout] matmul. Chunked over rows so peak memory stays ~256MB (the unchunked buffer is ~2.9GB at 128^3/128ch - fine on a Studio, not fine on an 8GB mini). m3ultra 128^3/128ch: 81.6ms -> 39.3ms (2.08x), 2.57 -> 5.34 Mvox/s m3ultra 64^3/128ch: 21.9ms -> 6.4ms (3.4x) 7/7 tests still pass against the naive reference.
This commit is contained in:
John
parent
97dcdfb54a
commit
6fcf677cef
128
bench/bench_subm.py
Normal file
128
bench/bench_subm.py
Normal file
@ -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()
|
||||
58
bench/run_fleet.sh
Executable file
58
bench/run_fleet.sh
Executable file
@ -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"
|
||||
@ -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)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user