1#!/usr/bin/env python32"""Generate every figure of the paper from logged measurements (benchmarks/raw, benchmarks/processed). PDF outputs in3paper/figures/ and figures/. Nothing here is hand-drawn from numbers typed in."""4import glob, json, os, statistics, sys5import matplotlib6matplotlib.use("Agg")7import matplotlib.pyplot as plt8ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))9FIG = os.path.join(ROOT, "paper", "figures"); FIG2 = os.path.join(ROOT, "figures")10os.makedirs(FIG, exist_ok=True); os.makedirs(FIG2, exist_ok=True)11plt.rcParams.update({"font.size": 9, "axes.spines.top": False, "axes.spines.right": False, "figure.dpi": 150})12PAL = ["#2E5A88", "#C8553D", "#4F8A5B", "#8E6C8A", "#D9A441", "#5B5B5B", "#7FB2D6"]1314def save(fig, name):15 for d in (FIG, FIG2):16 fig.savefig(os.path.join(d, name + ".pdf"), bbox_inches="tight")17 plt.close(fig); print("wrote", name)1819def fig_sim_hit_vs_capacity():20 p = os.path.join(ROOT, "benchmarks/processed/gateC/sim_olmoe_v2_ssd043_m8.json")21 if not os.path.exists(p): return22 d = json.load(open(p)); rows = [r for r in d["results"] if r["split"] == "heldout" and r["policy"] == "lru"]23 preds = []24 for r in rows:25 if r["predictor"] not in preds: preds.append(r["predictor"])26 fig, axes = plt.subplots(1, 2, figsize=(7.2, 2.8))27 for i, pr in enumerate(preds):28 rs = sorted([r for r in rows if r["predictor"] == pr], key=lambda r: r["capacity"])29 axes[0].plot([r["capacity"] for r in rs], [r["hit_rate"] for r in rs], marker="o", ms=3, color=PAL[i % len(PAL)], label=pr)30 axes[1].plot([r["capacity"] for r in rs], [r["stall_ms_per_token"] for r in rs], marker="o", ms=3, color=PAL[i % len(PAL)], label=pr)31 axes[0].set_xlabel("slots per layer $C$"); axes[0].set_ylabel("expert hit rate (held-out)"); axes[0].set_ylim(0.4, 1.0)32 axes[1].set_xlabel("slots per layer $C$"); axes[1].set_ylabel("modelled stall (ms / token)")33 axes[1].legend(frameon=False, fontsize=7, ncol=2)34 fig.suptitle(f"OLMoE-1B-7B, trace-driven simulation (t_read {d['t_read_ms']} ms, {d['loaders']} loader)", fontsize=9)35 save(fig, "sim_hit_vs_capacity")3637def fig_capacity_sweep():38 paths = [os.path.join(ROOT, "benchmarks/raw/gateB/capacity_sweep_laptop/sweep.jsonl"),39 os.path.join(ROOT, "benchmarks/raw/gateC/prefetch_online_laptop/smoke_lazyinstall.jsonl"),40 os.path.join(ROOT, "benchmarks/raw/gateC/cold_nvme_laptop/sweep.jsonl")]41 recs = []42 for p in paths:43 if os.path.exists(p):44 recs += [dict(json.loads(l), src=os.path.basename(os.path.dirname(p))) for l in open(p) if l.strip()]45 recs = [r for r in recs if not r.get("warmup")]46 if not recs: return47 groups = {}48 for r in recs:49 key = (r["src"], r["backend"], r.get("capacity"), r.get("predictor") or "none", r.get("m"), bool(r.get("nocache")), r.get("tag", ""))50 groups.setdefault(key, []).append(r["gen_tps"])51 labels, vals, errs, cols = [], [], [], []52 for k, v in sorted(groups.items(), key=lambda kv: (kv[0][5], -(kv[0][2] or 999), kv[0][3])):53 src, be, cap, pr, m, cold, tag = k54 if be == "mlx_lm": lab = "mlx-lm (all resident)"55 else: lab = f"C={cap} {pr}{' m='+str(m) if pr!='none' else ''}{' cold' if cold else ' page-cache'} [{src.split('_')[0]}{(' '+tag) if tag else ''}]"56 labels.append(lab); vals.append(statistics.median(v)); errs.append((max(v) - min(v)) / 2); cols.append(PAL[1] if cold else PAL[0])57 fig, ax = plt.subplots(figsize=(7.2, 0.22 * len(labels) + 1.2))58 y = range(len(labels)); ax.barh(y, vals, xerr=errs, color=cols, height=0.7, error_kw={"lw": 0.8})59 ax.set_yticks(list(y)); ax.set_yticklabels(labels, fontsize=7); ax.invert_yaxis(); ax.set_xlabel("decode tok/s (median; bar = range)")60 ax.set_title("OLMoE-1B-7B 4-bit on M5 Max: residency, tier and prefetch", fontsize=9)61 save(fig, "capacity_sweep")6263def fig_two_node():64 p = os.path.join(ROOT, "benchmarks/raw/gateE/two_node_m3u96a_m3u96b.jsonl")65 if not os.path.exists(p): return66 recs = [json.loads(l) for l in open(p) if l.strip()]; recs = [r for r in recs if not r["warmup"]]67 groups = {}68 for r in recs: groups.setdefault((r["capacity"], r["placement"]), []).append(r)69 fig, ax = plt.subplots(figsize=(4.2, 2.6))70 labels = []; i = 071 for (cap, pl), v in sorted(groups.items()):72 tps = [x["gen_tps"] for x in v]; rw = [x["remote"]["remote_wait_s"] / max(1, x["gen_tokens"]) * 1e3 if x.get("remote") else 0 for x in v]73 ax.bar(i, statistics.median(tps), color=PAL[0 if pl == "none" else 1]); labels.append(f"C={cap}\n{pl}"); i += 174 ax.set_xticks(range(i)); ax.set_xticklabels(labels, fontsize=7); ax.set_ylabel("decode tok/s (median)")75 ax.set_title("OLMoE: M3U96a local vs odd experts on M3U96b (1 GbE)", fontsize=8)76 save(fig, "two_node")7778if __name__ == "__main__":79 fig_sim_hit_vs_capacity(); fig_capacity_sweep(); fig_two_node()80