kernbench2/scripts/paper/paper_plot_gqa_prefill_compute_bound.py

"""Comparative figure for the compute-bound prefill composite study.

Reads sweep_prefill_compute_bound.json (emitted by milestone-1h-gqa,
sweep ``prefill_cb``) and writes one two-panel PNG:

  gqa_prefill_compute_bound.png
    Left  — end-to-end prefill latency (µs) vs context length.
    Right — MAC utilization (achieved / 8 TFLOP·s⁻¹ per-PE peak) vs context.

Unlike memory-bound decode (where command form is latency-neutral), in
compute-bound prefill the composite command keeps the MAC array fed by
streaming DMA↔compute per HW tile, so it wins on both latency and
utilization — and the margin grows with context (deeper P·V reduction =
more tiles to pipeline).

Run (after the bench):
    GQA_1H_RUN=1 GQA_1H_SWEEPS=prefill_cb python -m kernbench.cli.main run \\
        --bench milestone-1h-gqa --topology topology.yaml
    python scripts/paper/paper_plot_gqa_prefill_compute_bound.py
"""
from __future__ import annotations

import json
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

_REPO_ROOT = Path(__file__).resolve().parents[2]
_FIG_DIR = (
    _REPO_ROOT / "src" / "kernbench" / "benches"
    / "1H_milestone_output" / "gqa" / "long_ctx"
)
_SWEEP_JSON = _FIG_DIR / "sweep_prefill_compute_bound.json"
_PAPER_FIG_DIR = (
    _REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
)

_VARIANT_STYLE = {
    "primitive":          ("primitive (tl.dot, hand-tiled)", "#c0504d", "o"),
    "composite":          ("composite GEMM",                 "#3b6ea5", "s"),
    "composite_extended": ("composite + softmax_merge",      "#4f8a4f", "^"),
}
_ORDER = ("primitive", "composite", "composite_extended")


def _ctx_label(c: int) -> str:
    return f"{c // 1024}K" if c >= 1024 else str(c)


def _series(rows, variant, key):
    pts = sorted(((r["ctx_len"], r[key]) for r in rows
                  if r["variant"] == variant), key=lambda t: t[0])
    return [p[0] for p in pts], [p[1] for p in pts]


def main() -> None:
    sweep = json.loads(_SWEEP_JSON.read_text())
    rows = sweep["rows"]
    ctxs = sweep["ctx_points"]

    fig, (ax_lat, ax_util) = plt.subplots(1, 2, figsize=(13.0, 4.8))

    for v in _ORDER:
        label, color, marker = _VARIANT_STYLE[v]
        xs, lat = _series(rows, v, "latency_ns")
        ax_lat.plot(xs, [y / 1e3 for y in lat], marker=marker,
                    color=color, label=label, lw=2)
        xs, util = _series(rows, v, "mac_util")
        ax_util.plot(xs, [u * 100 for u in util], marker=marker,
                     color=color, label=label, lw=2)

    for ax in (ax_lat, ax_util):
        ax.set_xscale("log", base=2)
        ax.set_xticks(ctxs)
        ax.set_xticklabels([_ctx_label(c) for c in ctxs])
        ax.set_xlabel(r"context length (= $T_q$ = $S_{kv}$)")
        ax.grid(True, ls=":", alpha=0.5)
        ax.legend(fontsize=9)

    ax_lat.set_ylabel("end-to-end prefill latency (µs)")
    ax_lat.set_title("Compute-bound prefill latency per command form")
    ax_util.set_ylabel("MAC utilization (% of 8 TFLOP·s⁻¹ peak)")
    ax_util.set_title(
        "MAC utilization — composite keeps the array fed; primitive starves"
    )
    ax_util.axhline(100, color="#888", ls="--", lw=1, alpha=0.6)

    fig.suptitle(
        "Compute-bound prefill attention — use of composite commands\n"
        "single-rank, GQA single-KV-head group ($h_q{=}8$, $d_{\\text{head}}"
        "{=}128$); $M{=}8T_q$ tile-filling",
        fontsize=11,
    )
    fig.tight_layout(rect=(0, 0, 1, 0.92))

    out = _FIG_DIR / "gqa_prefill_compute_bound.png"
    fig.savefig(out, dpi=150)
    plt.close(fig)
    print(f"wrote {out}")

    if _PAPER_FIG_DIR.is_dir():
        dst = _PAPER_FIG_DIR / out.name
        dst.write_bytes(out.read_bytes())
        print(f"copied {dst}")


if __name__ == "__main__":
    main()