201 lines
9.7 KiB
Python
201 lines
9.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Plot the benchmark matrix into publication-style figures (PNG)."""
|
|
import argparse, json, os, re
|
|
import matplotlib
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
|
|
# Cohesive figure style shared by every chart.
|
|
PALETTE = {"baseline": "#c0392b", "mooncake": "#2e6fdb", "reuser": "#27ae60",
|
|
"accent": "#7a5fb8", "muted": "#9aa0a6"}
|
|
plt.rcParams.update({
|
|
"figure.dpi": 140,
|
|
"font.size": 12,
|
|
"axes.titlesize": 15,
|
|
"axes.titleweight": "bold",
|
|
"axes.labelsize": 12,
|
|
"axes.spines.top": False,
|
|
"axes.spines.right": False,
|
|
"axes.grid": True,
|
|
"grid.color": "#e6e6e6",
|
|
"grid.linewidth": 0.8,
|
|
"axes.axisbelow": True,
|
|
"legend.frameon": False,
|
|
"figure.autolayout": True,
|
|
})
|
|
|
|
|
|
def _human(name):
|
|
"""A2_7b_30k_6ag -> '7B / 30k / 6 agents'."""
|
|
m = re.search(r"(\d+\.?\d*)b_(\d+)k?_(\d+)ag", name, re.I)
|
|
if not m:
|
|
return name
|
|
size, ctx, ag = m.groups()
|
|
return f"{size}B / {ctx}k / {ag} agents"
|
|
|
|
|
|
def fig_ttft(cells, out):
|
|
names = [_human(c["name"]) for c in cells]
|
|
off = [c["ttft_off"] for c in cells]
|
|
on = [c["ttft_on"] for c in cells]
|
|
reuser = [c.get("reuser_ttft_on", c["ttft_on"]) for c in cells]
|
|
x = np.arange(len(names)); w = 0.27
|
|
fig, ax = plt.subplots(figsize=(max(7, 2.0 * len(names)), 4.4))
|
|
ax.bar(x - w, off, w, label="no sharing (baseline)", color=PALETTE["baseline"])
|
|
ax.bar(x, on, w, label="Mooncake (all agents)", color=PALETTE["mooncake"])
|
|
ax.bar(x + w, reuser, w, label="Mooncake (reusing agents only)", color=PALETTE["reuser"])
|
|
for i, (a, b) in enumerate(zip(off, on)):
|
|
ax.text(i, max(a, b) * 1.02, f"-{100*(1-b/a):.0f}%", ha="center", fontsize=10, fontweight="bold", color="#222")
|
|
ax.set_ylabel("mean TTFT (ms)")
|
|
ax.set_title("Time-to-first-token: KV reuse via Mooncake")
|
|
ax.set_xticks(x); ax.set_xticklabels(names); ax.legend()
|
|
fig.savefig(out); print("wrote", out); plt.close(fig)
|
|
|
|
|
|
def fig_recompute(cells, out):
|
|
names = [_human(c["name"]) for c in cells]
|
|
off = [c["recompute_off_pct"] for c in cells]
|
|
on = [c["recompute_on_pct"] for c in cells]
|
|
x = np.arange(len(names)); w = 0.38
|
|
fig, ax = plt.subplots(figsize=(max(7, 2.0 * len(names)), 4.4))
|
|
ax.bar(x - w/2, off, w, label="baseline", color=PALETTE["baseline"])
|
|
bars = ax.bar(x + w/2, on, w, label="Mooncake", color=PALETTE["mooncake"])
|
|
for b, v in zip(bars, on):
|
|
ax.text(b.get_x() + b.get_width()/2, v + 2, f"{v:.0f}%", ha="center", fontsize=10, color="#222")
|
|
ax.set_ylabel("% of prompt tokens re-prefilled"); ax.set_ylim(0, 108)
|
|
ax.set_title("Redundant prefill eliminated"); ax.set_xticks(x)
|
|
ax.set_xticklabels(names); ax.legend()
|
|
fig.savefig(out); print("wrote", out); plt.close(fig)
|
|
|
|
|
|
def fig_throughput(cells, out):
|
|
names = [_human(c["name"]) for c in cells]
|
|
thru = [c["thru_x"] for c in cells]
|
|
x = np.arange(len(names))
|
|
fig, ax = plt.subplots(figsize=(max(7, 2.0 * len(names)), 4.2))
|
|
bars = ax.bar(x, thru, 0.5, color=PALETTE["accent"])
|
|
ax.axhline(1.0, ls="--", color=PALETTE["muted"], label="parity")
|
|
for b, t in zip(bars, thru):
|
|
ax.text(b.get_x() + b.get_width()/2, t + 0.02, f"{t:.2f}x", ha="center", fontsize=11, fontweight="bold")
|
|
ax.set_ylabel("throughput speedup (x)"); ax.set_title("Aggregate swarm throughput")
|
|
ax.set_xticks(x); ax.set_xticklabels(names); ax.legend()
|
|
fig.savefig(out); print("wrote", out); plt.close(fig)
|
|
|
|
|
|
def fig_per_agent(swarm_json, out):
|
|
"""Prompt tokens recomputed per agent: the pioneer pays the full prefill,
|
|
every later agent restores and recomputes almost nothing. This is the
|
|
'only agent 0 prefills' story, shown directly rather than via TTFT (which
|
|
on the Stage-1 file path can make a reuser look slower than the pioneer)."""
|
|
d = json.load(open(swarm_json))
|
|
pa = d["per_agent_cached"]
|
|
idx = [r["idx"] for r in pa]
|
|
recomputed = [r["prompt_n"] for r in pa]
|
|
total = [r["tokens"] for r in pa]
|
|
colors = [PALETTE["baseline"]] + [PALETTE["reuser"]] * (len(pa) - 1)
|
|
fig, ax = plt.subplots(figsize=(max(7.5, 1.0 * len(pa)), 4.6))
|
|
# faint full-prompt reference so the saving is obvious
|
|
ax.bar(idx, total, color="#ececec", width=0.74, label="prompt size (tokens)")
|
|
bars = ax.bar(idx, recomputed, color=colors, width=0.74, label="tokens re-prefilled")
|
|
for i, (b, rc) in enumerate(zip(bars, recomputed)):
|
|
tag = f"{rc}" if rc < 1000 else f"{rc/1000:.1f}k"
|
|
ax.text(b.get_x() + b.get_width() / 2, rc + max(total) * 0.02, tag,
|
|
ha="center", fontsize=10, fontweight="bold", color="#222")
|
|
ax.annotate("pioneer pays\nthe full prefill", xy=(0, total[0]), xytext=(0.6, total[0] * 0.82),
|
|
fontsize=10, color=PALETTE["baseline"],
|
|
arrowprops=dict(arrowstyle="->", color=PALETTE["baseline"]))
|
|
saved = sum(total) - sum(recomputed)
|
|
ax.set_xlabel("agent index (0 = pioneer that prefills + stores)")
|
|
ax.set_ylabel("prompt tokens")
|
|
ax.set_xticks(idx)
|
|
ax.set_title(f"Only the pioneer prefills ({saved//1000}k tokens saved)")
|
|
ax.legend(loc="center right")
|
|
fig.savefig(out); print("wrote", out); plt.close(fig)
|
|
|
|
|
|
def fig_scaling(scaling_json, out):
|
|
"""Baseline work grows linearly with the swarm; Mooncake stays almost flat.
|
|
scaling_json is a list of {agents, baseline_recompute_tokens,
|
|
cached_recompute_tokens, baseline_wall_s, cached_wall_s}."""
|
|
rows = sorted(json.load(open(scaling_json)), key=lambda r: r["agents"])
|
|
n = [r["agents"] for r in rows]
|
|
base_tok = [r["baseline_recompute_tokens"] / 1e3 for r in rows]
|
|
cache_tok = [r["cached_recompute_tokens"] / 1e3 for r in rows]
|
|
fig, ax = plt.subplots(figsize=(7.5, 4.6))
|
|
ax.plot(n, base_tok, "o-", color=PALETTE["baseline"], lw=2.4, ms=7,
|
|
label="no sharing (re-prefill every agent)")
|
|
ax.plot(n, cache_tok, "o-", color=PALETTE["mooncake"], lw=2.4, ms=7,
|
|
label="Mooncake (prefill once, restore the rest)")
|
|
ax.fill_between(n, cache_tok, base_tok, color=PALETTE["mooncake"], alpha=0.10)
|
|
# endpoint callouts at the widest swarm
|
|
ax.annotate(f"{base_tok[-1]:.0f}k", xy=(n[-1], base_tok[-1]), xytext=(-4, 6),
|
|
textcoords="offset points", ha="right", fontsize=11,
|
|
fontweight="bold", color=PALETTE["baseline"])
|
|
ax.annotate(f"{cache_tok[-1]:.0f}k (flat)", xy=(n[-1], cache_tok[-1]), xytext=(-4, 8),
|
|
textcoords="offset points", ha="right", fontsize=11,
|
|
fontweight="bold", color=PALETTE["mooncake"])
|
|
ax.text(n[len(n)//2], (base_tok[-1] + cache_tok[0]) / 2,
|
|
f"{base_tok[-1]-cache_tok[-1]:.0f}k tokens\nnever re-prefilled",
|
|
ha="center", va="center", fontsize=10, color="#555", style="italic")
|
|
ax.set_xlabel("agents sharing the context")
|
|
ax.set_ylabel("prompt tokens prefilled (thousands)")
|
|
ax.set_title("Redundant prefill work vs swarm size (7B, 30k ctx)")
|
|
ax.set_xticks(n); ax.set_ylim(bottom=0); ax.legend(loc="upper left")
|
|
fig.savefig(out); print("wrote", out); plt.close(fig)
|
|
|
|
|
|
def fig_arbiter(out, points=None):
|
|
"""Restore cost vs recompute cost per (model, context); the arbiter picks the
|
|
cheaper. Demonstrates the loss-free property: it restores only when restoring
|
|
is actually faster than recomputing. Values are the arbiter's own estimates
|
|
from measured online rates.
|
|
|
|
points: list of {label, restore_ms, recompute_ms, decision}. Defaults to the
|
|
two measured regimes (7B restores, 1.5B declines)."""
|
|
if points is None:
|
|
points = [
|
|
{"label": "7B / 30k ctx", "restore_ms": 2160, "recompute_ms": 4034, "decision": "RESTORE"},
|
|
{"label": "1.5B / 8k ctx", "restore_ms": 697, "recompute_ms": 311, "decision": "RECOMPUTE"},
|
|
]
|
|
labels = [p["label"] for p in points]
|
|
restore = [p["restore_ms"] for p in points]
|
|
recompute = [p["recompute_ms"] for p in points]
|
|
x = np.arange(len(labels)); w = 0.36
|
|
fig, ax = plt.subplots(figsize=(max(6.5, 2.6 * len(labels)), 4.4))
|
|
ax.bar(x - w / 2, restore, w, label="restore from store", color=PALETTE["mooncake"])
|
|
ax.bar(x + w / 2, recompute, w, label="recompute (prefill)", color=PALETTE["baseline"])
|
|
for i, p in enumerate(points):
|
|
chosen = min(p["restore_ms"], p["recompute_ms"])
|
|
ax.annotate(f"chooses\n{p['decision']}",
|
|
xy=(i, chosen), xytext=(i, chosen + max(restore + recompute) * 0.08),
|
|
ha="center", fontsize=10, fontweight="bold", color="#1a7d32")
|
|
ax.set_ylabel("estimated cost (ms)")
|
|
ax.set_title("Loss-free arbiter: restore only when it's cheaper")
|
|
ax.set_xticks(x); ax.set_xticklabels(labels); ax.legend(loc="upper left")
|
|
fig.savefig(out); print("wrote", out); plt.close(fig)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--matrix", default="run/matrix_results.json")
|
|
ap.add_argument("--per-agent", default="", help="a swarm cell json for the per-agent figure")
|
|
ap.add_argument("--scaling", default="", help="a scaling json for the agents-scaling figure")
|
|
ap.add_argument("--outdir", default="docs/figures")
|
|
args = ap.parse_args()
|
|
os.makedirs(args.outdir, exist_ok=True)
|
|
m = json.load(open(args.matrix))
|
|
cells = m["cells"]
|
|
fig_ttft(cells, os.path.join(args.outdir, "ttft.png"))
|
|
fig_recompute(cells, os.path.join(args.outdir, "recompute.png"))
|
|
fig_throughput(cells, os.path.join(args.outdir, "throughput.png"))
|
|
if args.per_agent and os.path.exists(args.per_agent):
|
|
fig_per_agent(args.per_agent, os.path.join(args.outdir, "per_agent.png"))
|
|
if args.scaling and os.path.exists(args.scaling):
|
|
fig_scaling(args.scaling, os.path.join(args.outdir, "scaling_agents.png"))
|
|
fig_arbiter(os.path.join(args.outdir, "arbiter.png"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|