yuhuabillchen/justfit-mlx/benchmarking

#!/usr/bin/env bash
# <feature> perf: upstream main vs fork, random-init tiny <model>, swept over <param>.
set -euo pipefail
W=$(mktemp -d)
git clone -q --depth 1 https://github.com/Blaizzy/mlx-vlm "$W/main"
git clone -q --depth 1 https://github.com/<you>/mlx-vlm   "$W/fork"
 
cat > "$W/b.py" <<'PY'
import time, numpy as np, mlx.core as mx
from mlx_vlm.models.<model> import Model, ModelConfig
cfg = ModelConfig(...tiny config...)          # small enough to fit in RAM
m = Model(cfg); m.eval(); mx.eval(m.parameters())
def ms(T):
    c = m.make_cache(); p = 0
    while p < T:                               # chunked prefill avoids OOM
        n = min(256, T - p)
        mx.eval(m(mx.array(np.random.randint(0, cfg.vocab_size, (1, n))), cache=c).logits); p += n
    for _ in range(4): mx.eval(m(mx.array([[0]]), cache=c).logits)   # warmup
    t = time.perf_counter()
    for _ in range(20): mx.eval(m(mx.array([[0]]), cache=c).logits)
    return (time.perf_counter() - t) / 20 * 1e3
for T in (16384, 32768, 65536):
    print(T, round(ms(T), 2), flush=True)
PY
 
run() { uv venv -q "$1/.venv"; uv pip install -q -e "$1" --python "$1/.venv/bin/python"; "$1/.venv/bin/python" "$W/b.py"; }
run "$W/main" > "$W/m.txt"
run "$W/fork" > "$W/f.txt"
echo; printf "| param | main | fork | speedup |\n|--:|--:|--:|--:|\n"
join "$W/m.txt" "$W/f.txt" | awk '{printf "| %d | %.2f ms | %.2f ms | %.2fx |\n", $1, $2, $3, $2/$3}'

import time, statistics, mlx.core as mx
def bench(fn, it=30, warm=5):
    for _ in range(warm): mx.eval(fn())
    mx.synchronize()
    if hasattr(mx, "reset_peak_memory"): mx.reset_peak_memory()
    s = []
    for _ in range(it):
        t0 = time.perf_counter(); mx.eval(fn()); s.append((time.perf_counter() - t0) * 1e3)
    return statistics.median(s), mx.get_peak_memory() / 1e9
# build inputs; assert new == baseline in the degenerate limit; then print a swept markdown table.