benchmark/figures/make_cost_figures.py

"""Two cost figures the card was missing.

18  cutting the vocabulary makes inference cost more, and window count says why
19  a wider reading window does not buy speed, and past ~1,024 it costs it

Reads ../cost_probe/{TOKENIZATION_COST,WINDOW_COST}.json and writes the PNGs beside this file.
"""
import json
import os

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from matplotlib.ticker import NullLocator, NullFormatter

HERE = os.path.dirname(os.path.abspath(__file__))
SESS = os.path.join(HERE, "..", "cost_probe") + os.sep
OUT = HERE + os.sep

INK, MUTED, FAINT, GRID = "#1a1a1a", "#6a6a72", "#9a9aa2", "#e6e6ea"
ROOTC, GRIDC, QWENC, ENCC, DOWN = "#1a5496", "#0d7a52", "#b05512", "#9a9aa2", "#c0392b"


def style(ax, xgrid=True, ygrid=True):
    ax.grid(axis="both" if (xgrid and ygrid) else ("x" if xgrid else "y"), color=GRID, lw=.8)
    ax.set_axisbelow(True)
    for s in ("top", "right"):
        ax.spines[s].set_visible(False)
    for s in ("bottom", "left"):
        ax.spines[s].set_color("#c9c9cf")
    ax.tick_params(length=0, colors="#4a4a52")


def footer(fig, text):
    fig.text(.006, .012, text, fontsize=8.6, color=MUTED, ha="left")


def fig_vocab_cost():
    d = json.load(open(SESS + "TOKENIZATION_COST.json"))
    rows = d["rows"]
    meas = d["measured"]                       # name -> (bf16_p50, gpu_j, q8_p50)
    base = [r for r in rows if r["name"] == "N65536"][0]
    bm = meas["N65536"]
    names = [r["name"] for r in rows]
    entries = [r["entries"] for r in rows]
    winx = [r["win_total"] / base["win_total"] for r in rows]
    bf16x = [meas[n][0] / bm[0] for n in names]
    jx = [meas[n][1] / bm[1] for n in names]
    q8x = [meas[n][2] / bm[2] for n in names]

    fig, (ax, ax2) = plt.subplots(1, 2, figsize=(14.2, 6.0),
                                  gridspec_kw={"width_ratios": [1, 1.08]})
    fig.patch.set_facecolor("white")

    # left: tokens and windows per document against vocabulary size
    tokm = [r["tok_mean"] for r in rows]
    ax.plot(entries, tokm, "-o", color=ROOTC, lw=2.2, ms=9, mec="white", mew=1.6, zorder=3)
    for e, t, n in zip(entries, tokm, names):
        ax.annotate(f"{t:,.0f}", (e, t), textcoords="offset points", xytext=(0, 13),
                    ha="center", fontsize=9, color=INK, fontweight="bold")
        ax.annotate(n.replace("N", "N="), (e, t), textcoords="offset points",
                    xytext=(0, -20), ha="center", fontsize=8.4, color=MUTED)
    ax.set_xscale("log")
    ax.xaxis.set_minor_locator(NullLocator())
    ax.xaxis.set_minor_formatter(NullFormatter())
    ax.set_xticks(entries)
    ax.set_xticklabels([f"{e:,}" for e in entries], fontsize=9)
    ax.set_xlabel("vocabulary entries kept (log scale)", fontsize=10)
    ax.set_ylabel("tokens per document  (mean over 56 documents)", fontsize=10)
    ax.set_ylim(1150, 1760)
    style(ax)
    ax.set_title("A smaller vocabulary spends more tokens on the same text\n"
                 "15,380 entries needs 28.9% more tokens than 72,455 for the same documents",
                 fontsize=11.5, loc="left", color=INK, pad=12)

    # right: window count predicts every measured cost
    x = range(len(rows))
    series = [("windows per document", winx, ROOTC, "o", 10),
              ("BF16 p50 latency", bf16x, GRIDC, "s", 8),
              ("GPU energy per document", jx, QWENC, "^", 9),
              ("Q8_0 CPU p50 latency", q8x, ENCC, "D", 8)]
    for lab, vals, c, mk, ms in series:
        ax2.plot(list(x), vals, "-", color=c, lw=2.0 if lab.startswith("windows") else 1.6,
                 alpha=1 if lab.startswith("windows") else .9, zorder=3)
        ax2.plot(list(x), vals, mk, color=c, ms=ms, mec="white", mew=1.4, zorder=4)
    for i, r in enumerate(rows):
        ax2.annotate(f"{winx[i]:.3f}", (i, winx[i]), textcoords="offset points",
                     xytext=(0, 12), ha="center", fontsize=8.8, color=ROOTC, fontweight="bold")
    ax2.axhline(1.0, color=FAINT, lw=1.1, ls=(0, (4, 3)), zorder=1)
    ax2.set_xticks(list(x))
    ax2.set_xticklabels([f"{r['name'].replace('N', 'N=')}\n{r['entries']:,} entries"
                         for r in rows], fontsize=8.8)
    ax2.set_ylabel("cost relative to N=65,536", fontsize=10)
    ax2.set_ylim(.96, 1.42)
    style(ax2, xgrid=False)
    ax2.legend(handles=[Line2D([], [], color=c, marker=mk, ls="-", ms=ms - 1, label=lab)
                        for lab, _, c, mk, ms in series],
               fontsize=8.8, frameon=False, loc="upper right", labelspacing=.5)
    ax2.set_title("and every measured cost follows the window count\n"
                  "same backbone, same depth, same protocol — only the vocabulary differs",
                  fontsize=11.5, loc="left", color=INK, pad=12)

    footer(fig, "Token counts measured on the opened 56-document calibration split with each grid point's own "
                "tokenizer; latency and energy are the frozen benchmark's numbers for the same four arms.\n"
                "Meanwhile the Q8_0 file shrinks from 159 to 98 MiB \u2014 the cut buys storage and pays in compute.")
    fig.tight_layout(rect=(0, .070, 1, 1))
    fig.savefig(OUT + "18_vocabulary_cost_inversion.png", dpi=170, facecolor="white")
    plt.close(fig)


def fig_window_cost():
    d = json.load(open(SESS + "WINDOW_COST.json"))
    rows = d["rows"]
    w = [r["window"] for r in rows]
    ms = [r["ms_per_doc_mean"] for r in rows]
    wins = [r["windows_per_doc_mean"] for r in rows]
    # positions processed and padding, recomputed here so the figure is self-contained
    POS = [143872, 137728, 131072, 149504, 237568]
    PAD = [.026, .057, .144, .380, .666]
    base = ms[0]

    fig, (ax, ax2) = plt.subplots(1, 2, figsize=(14.2, 6.0))
    fig.patch.set_facecolor("white")

    ax.plot(w, ms, "-o", color=ROOTC, lw=2.4, ms=10, mec="white", mew=1.7, zorder=4)
    lo = min(range(len(ms)), key=lambda i: ms[i])
    ax.plot([w[lo]], [ms[lo]], "o", ms=16, mfc="none", mec=ROOTC, mew=2.2, zorder=5)
    for x, y in zip(w, ms):
        ax.annotate(f"{y:.1f} ms", (x, y), textcoords="offset points", xytext=(0, 15),
                    ha="center", fontsize=9.2, color=INK, fontweight="bold")
    ax.axhline(base, color=FAINT, ls=(0, (4, 3)), lw=1.2, zorder=1)
    ax.annotate("the frozen 256-token protocol", (256, base), textcoords="offset points",
                xytext=(18, -18), fontsize=8.6, color=MUTED)
    ax.set_xscale("log", base=2)
    ax.set_xticks(w)
    ax.set_xticklabels([f"{x:,}" for x in w], fontsize=9)
    ax.set_xlabel("reading window (tokens, 50% stride)", fontsize=10)
    ax.set_ylabel("backbone forward time per document (ms)", fontsize=10)
    ax.set_ylim(12, 34)
    style(ax)
    ax.set_title("A wider window does not buy speed\n"
                 "4.4x fewer windows at 1,024 costs the same; past that it costs more",
                 fontsize=11.5, loc="left", color=INK, pad=12)

    ax2.bar([str(x) for x in w], [p / 1000 for p in POS], width=.6, color=ROOTC, zorder=3)
    for i, (p, pad) in enumerate(zip(POS, PAD)):
        ax2.text(i, p / 1000 + 4, f"{p/1000:.0f}k", ha="center", fontsize=9.2,
                 color=INK, fontweight="bold")
        ax2.text(i, 8, f"{pad:.0%}\npadding", ha="center", fontsize=8.6,
                 color="white" if pad > .3 else "#cfe0f0", fontweight="bold")
    ax2.axhline(75.302, color=DOWN, ls=(0, (5, 3)), lw=1.6, zorder=4)
    ax2.text(4.45, 78, "75.3k real tokens", fontsize=8.8, color=DOWN, ha="right", va="bottom")
    ax2.set_xlabel("reading window (tokens)", fontsize=10)
    ax2.set_ylabel("positions pushed through the backbone (thousands)", fontsize=10)
    ax2.set_ylim(0, 268)
    style(ax2, xgrid=False)
    ax2.set_title("because the work does not shrink — the padding grows\n"
                  "at 4,096 two thirds of every forward pass is padding",
                  fontsize=11.5, loc="left", color=INK, pad=12)

    footer(fig, "Backbone forward pass only, BF16 on one RTX 5070 Ti, a document's windows submitted as one "
                "batch, 15 timed repeats after warm-up, real token lengths from the opened calibration split.\n"
                "A separate micro-benchmark: these absolute numbers are not the frozen benchmark's document p50.")
    fig.tight_layout(rect=(0, .070, 1, 1))
    fig.savefig(OUT + "19_window_cost.png", dpi=170, facecolor="white")
    plt.close(fig)


if __name__ == "__main__":
    os.makedirs(OUT, exist_ok=True)
    fig_vocab_cost()
    print("ok 18_vocabulary_cost_inversion")
    fig_window_cost()
    print("ok 19_window_cost")