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" and r["predictor"] != "probe"]23 for extra in ("sim_olmoe_v2_ssd043_m8_probe_fixed.json", "sim_olmoe_v2_ssd043_m8_probe_depths.json"):24 q = os.path.join(ROOT, "benchmarks/processed/gateC", extra)25 if os.path.exists(q):26 rows += [r for r in json.load(open(q))["results"] if r["split"] == "heldout" and r["policy"] == "lru" and r["predictor"] not in {x["predictor"] for x in rows}]27 preds = []28 for r in rows:29 if r["predictor"] not in preds: preds.append(r["predictor"])30 fig, axes = plt.subplots(1, 2, figsize=(7.2, 2.8))31 for i, pr in enumerate(preds):32 rs = sorted([r for r in rows if r["predictor"] == pr], key=lambda r: r["capacity"])33 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)34 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)35 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)36 axes[1].set_xlabel("slots per layer $C$"); axes[1].set_ylabel("modelled stall (ms / token)")37 axes[1].legend(frameon=False, fontsize=7, ncol=2)38 fig.suptitle(f"OLMoE-1B-7B, trace-driven simulation (t_read {d['t_read_ms']} ms, {d['loaders']} loader)", fontsize=9)39 save(fig, "sim_hit_vs_capacity")4041def fig_capacity_sweep():42 paths = [os.path.join(ROOT, "benchmarks/raw/gateB/capacity_sweep_laptop/sweep.jsonl"),43 os.path.join(ROOT, "benchmarks/raw/gateC/prefetch_online_laptop/smoke_lazyinstall.jsonl"),44 os.path.join(ROOT, "benchmarks/raw/gateC/cold_nvme_laptop/sweep.jsonl")]45 recs = []46 for p in paths:47 if os.path.exists(p):48 recs += [dict(json.loads(l), src=os.path.basename(os.path.dirname(p))) for l in open(p) if l.strip()]49 recs = [r for r in recs if not r.get("warmup")]50 if not recs: return51 groups = {}52 for r in recs:53 key = (r["src"], r["backend"], r.get("capacity"), r.get("predictor") or "none", r.get("m"), bool(r.get("nocache")), r.get("tag", ""))54 groups.setdefault(key, []).append(r["gen_tps"])55 labels, vals, errs, cols = [], [], [], []56 for k, v in sorted(groups.items(), key=lambda kv: (kv[0][5], -(kv[0][2] or 999), kv[0][3])):57 src, be, cap, pr, m, cold, tag = k58 if be == "mlx_lm": lab = "mlx-lm (all resident)"59 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 ''}]"60 labels.append(lab); vals.append(statistics.median(v)); errs.append((max(v) - min(v)) / 2); cols.append(PAL[1] if cold else PAL[0])61 fig, ax = plt.subplots(figsize=(7.2, 0.22 * len(labels) + 1.2))62 y = range(len(labels)); ax.barh(y, vals, xerr=errs, color=cols, height=0.7, error_kw={"lw": 0.8})63 ax.set_yticks(list(y)); ax.set_yticklabels(labels, fontsize=7); ax.invert_yaxis(); ax.set_xlabel("decode tok/s (median; bar = range)")64 ax.set_title("OLMoE-1B-7B 4-bit on M5 Max: residency, tier and prefetch", fontsize=9)65 save(fig, "capacity_sweep")6667def fig_two_node():68 p = os.path.join(ROOT, "benchmarks/raw/gateE/two_node_m3u96a_m3u96b.jsonl")69 if not os.path.exists(p): return70 recs = [json.loads(l) for l in open(p) if l.strip()]; recs = [r for r in recs if not r["warmup"]]71 groups = {}72 for r in recs: groups.setdefault((r["capacity"], r["placement"]), []).append(r)73 fig, ax = plt.subplots(figsize=(4.2, 2.6))74 labels = []; i = 075 for (cap, pl), v in sorted(groups.items()):76 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]77 ax.bar(i, statistics.median(tps), color=PAL[0 if pl == "none" else 1]); labels.append(f"C={cap}\n{pl}"); i += 178 ax.set_xticks(range(i)); ax.set_xticklabels(labels, fontsize=7); ax.set_ylabel("decode tok/s (median)")79 ax.set_title("OLMoE: M3U96a local vs odd experts on M3U96b (1 GbE)", fontsize=8)80 save(fig, "two_node")8182def fig_attribution():83 p = os.path.join(ROOT, "benchmarks/raw/gateC/cold_nvme_laptop/attribution.jsonl")84 if not os.path.exists(p): return85 recs = [json.loads(l) for l in open(p) if l.strip()]; recs = [r for r in recs if not r["warmup"]]86 groups = {}87 for r in recs: groups.setdefault((r["predictor"], r["m"]), []).append(r)88 order = sorted(groups, key=lambda k: (k[0] != "none", k[0], k[1]))89 comps = ["sync", "prefetch", "ensure"]; names = {"sync": "GPU sync (graph incl. installs)", "prefetch": "prefetch bookkeeping + late waits", "ensure": "demand reads (SSD)"}90 fig, ax = plt.subplots(figsize=(7.2, 3.0)); bottom = [0.0] * len(order)91 for ci, c in enumerate(comps):92 vals = []93 for k in order:94 v = groups[k]; n = sum(x["gen_tokens"] for x in v); vals.append(sum(x["host_time_s"][c] for x in v) / n * 1e3)95 ax.bar(range(len(order)), vals, bottom=bottom, color=PAL[ci], label=names[c]); bottom = [b + v for b, v in zip(bottom, vals)]96 for i, k in enumerate(order):97 v = groups[k]; ax.text(i, bottom[i] + 0.5, f"{statistics.median(x['gen_tps'] for x in v):.0f} tok/s\nhit {statistics.median(x['cache_hit_rate'] for x in v):.2f}", ha="center", fontsize=6)98 ax.set_xticks(range(len(order))); ax.set_xticklabels([f"{k[0]}\nm={k[1]}" for k in order], fontsize=6.5); ax.set_ylabel("host time per token (ms)")99 ax.set_ylim(0, max(bottom) * 1.25); ax.legend(frameon=False, fontsize=7, loc="upper left")100 ax.set_title("OLMoE, 16 slots, cold NVMe (M5 Max): where the decode step waits", fontsize=9)101 save(fig, "attribution")102103if __name__ == "__main__":104 fig_sim_hit_vs_capacity(); fig_capacity_sweep(); fig_two_node(); fig_attribution()105