paper(gemm): per-PE HBM BW + TFLOPS evaluation plots; unify peak to 8 TFLOP/s

Two new evaluation plots for §3 GEMM, both using the composite window
as the denominator (the up-front tl.load of A in load_ref is therefore
excluded — only HBM traffic and compute inside the composite count):

- gemm_hbm_bw_util.png: per-PE HBM bandwidth utilization against the
  256 GB/s per-PE ceiling. Answers "is this workload memory-bound on
  this configuration?" — small/single-tile shapes sit at 22–46%
  (pipeline-fill-limited), large or output-write-dominated shapes
  saturate (≥85%).

- gemm_per_pe_tflops.png: per-PE GEMM throughput against the 8 TFLOP/s
  engine peak. The deep-K K=3072 load_ref case lands at
  7.18 TFLOP/s (~90% of peak) — the configuration's clean win.
  ref_ref at the same shape drops to 3.83 TFLOP/s because the second
  operand doubles HBM pressure and the kernel hits the BW ceiling.
  The variant gap on this chart is the operational cost of NOT
  pre-staging the activation.

Both charts compare two operand-staging variants on every shape:
load_ref ("activation pre-staged", weight-only HBM streaming) and
ref_ref (both A and W streamed from HBM). load_load is omitted —
with both operands pre-staged the composite carries no HBM traffic
and the BW metric collapses to 0.

Analytic / measured peak unified to the single hardware spec
(8.0 TFLOP/s = 8000 flops/ns). T_STAGE is now derived as
tile_flops / peak (= 16.384 ns for the 32×64×32 tile) rather than
the previous hardcoded 16 ns, so MAC efficiency and per-PE TFLOPS%
share the same denominator and never disagree. Side effect: max
analytic-vs-measured gap tightens from 2.2 ppt to 1.4 ppt across
the sweep.

§3 Results restructured: lead with the two new evaluation plots
(BW saturation, achieved throughput), then the MAC-utilization
analytic-vs-measured chart serves as the validation step, then the
per-stage engine wall-clock as the supporting diagnostic. §3
Analysis updated to use the load_ref / ref_ref contrast as the
hardware-software boundary line that motivates the GQA kernel of §5.

§2.4 Accuracy claim updated to match: GEMM analytic-vs-measured
agreement is now "within 1.4 ppt across every swept shape", from
~7.7% at single-tile up to ~90% at K=3072. Subtitle/title vertical
stacking glitch in the matplotlib charts fixed in passing.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

src/kernbench/benches/milestone_1h_gemm.py

@@ -277,12 +277,35 @@ STAGE_COLORS = {
 # MAC-utilization analytic model constants.
 # Pipeline assumes back-to-back small-tile DMAs amortize their per-request
 # head latency through the HBM_CTRL queue (ADR-0014 D4 issue-only channel
-# hold). Steady-state cycle is BW-bound at one tile / T_STAGE.
-_T_STAGE = 16.0                  # steady-state cycle (ns/tile)
+# hold). Steady-state cycle is BW-bound at one tile / T_STAGE, where
+# T_STAGE is derived from the hardware peak so that analytic and measured
+# share the same per-PE GEMM-throughput denominator
+# (\\SI{8}{\\tera\\flop\\per\\second} = 8000 flops/ns) — see
+# `_t_stage_ns_from_tile()` below.
 _T_PIPELINE_FILL = 60.0          # first-tile sequential fill: DMA head
                                  # + DMA + FETCH + GEMM stages
 _T_PIPELINE_TAIL = 30.0          # tail: last STORE + DMA_WRITE + drain
 
+# Per-PE GEMM peak. Single source of truth for both the analytic
+# `compute_total = tiles * T_STAGE` term and the measured-eff
+# denominator. Matches topology.yaml pe_gemm.peak_tflops_f16.
+_GEMM_PEAK_TFLOPS = 8.0
+_GEMM_PEAK_FLOPS_PER_NS = _GEMM_PEAK_TFLOPS * 1000.0   # = 8000 flops/ns
+
+
+def _t_stage_ns_from_tile(tile_flops: int) -> float:
+    """Per-tile bottleneck cycle (ns) implied by the hardware peak.
+
+    The composite pipeline's steady-state stage cycle equals the time the
+    GEMM engine spends on one tile at peak throughput. Deriving T_STAGE
+    from tile_flops / peak (rather than hardcoding 16 ns) keeps the
+    analytic model self-consistent with the measured-eff denominator: a
+    tile with tile_flops = 131072 and an \\SI{8}{\\tera\\flop\\per\\second}
+    engine takes exactly 131072 / 8000 = 16.384 ns of GEMM time.
+    """
+    return tile_flops / _GEMM_PEAK_FLOPS_PER_NS
+
+
 _PLOT_VARIANT = "load_ref"
 
 
@@ -342,10 +365,15 @@ def _grouped_bar_png(
         _xtick_labels(shape_labels, tile_counts, flagged), fontsize=8,
     )
     ax.set_ylabel(y_label)
-    ax.set_title(title, fontsize=13, fontweight="bold")
+    # Title with optional subtitle stacked above (subtitle in lighter ink,
+    # title in bold below it). Coordinates are in axes-fraction so they
+    # survive tight_layout without colliding with the data area.
     if subtitle:
-        ax.text(0.5, 1.01, subtitle, transform=ax.transAxes, ha="center",
+        ax.text(0.5, 1.10, subtitle, transform=ax.transAxes, ha="center",
                 va="bottom", fontsize=8, color="#475569")
+        ax.set_title(title, fontsize=13, fontweight="bold", pad=22)
+    else:
+        ax.set_title(title, fontsize=13, fontweight="bold")
     if threshold is not None:
         ax.axhline(threshold, ls="--", color="gray", lw=1.0)
     ax.legend(fontsize=8, loc="upper right")
@@ -420,6 +448,7 @@ def emit_mac_utilization_measured(
     flagged = [_under_tile(*k, TILE_M, TILE_K, TILE_N) for k in shapes]
     tile_counts = [by_shape[k]["tile_count_expected"] for k in shapes]
 
+    t_stage_ns = _t_stage_ns_from_tile(tile_flops)
     gemm_util, useful_eff = [], []
     for k in shapes:
         r = by_shape[k]
@@ -428,7 +457,7 @@ def emit_mac_utilization_measured(
         tiles = r["tile_count_expected"]
         gu = useful / (tile_flops * tiles) * 100
         gemm_util.append(gu)
-        compute_total = tiles * _T_STAGE
+        compute_total = tiles * t_stage_ns
         wall = _T_PIPELINE_FILL + compute_total + _T_PIPELINE_TAIL
         ueff = (compute_total * (gu / 100.0) / wall) * 100 if wall > 0 else 0.0
         useful_eff.append(ueff)
@@ -458,7 +487,7 @@ def emit_mac_utilization_theoretical_vs_measured(
     tile = data["tile_sizes"]
     TILE_M, TILE_K, TILE_N = tile["M"], tile["K"], tile["N"]
     tile_flops = 2 * TILE_M * TILE_K * TILE_N
-    peak_per_ns = tile_flops / _T_STAGE
+    t_stage_ns = _t_stage_ns_from_tile(tile_flops)
 
     by_shape = {(r["M"], r["K"], r["N"]): r
                 for r in rows if r["variant"] == _PLOT_VARIANT}
@@ -479,11 +508,11 @@ def emit_mac_utilization_theoretical_vs_measured(
         gu_t.append(gut * 100)
         rec = r.get("stages", {}).get("GEMM", {}).get("record_count", 0) or tiles
         gu_m.append((useful / (tile_flops * rec) * 100) if rec else 0.0)
-        compute_total = tiles * _T_STAGE
+        compute_total = tiles * t_stage_ns
         wall_t = _T_PIPELINE_FILL + compute_total + _T_PIPELINE_TAIL
         eff_t.append((compute_total * gut / wall_t * 100) if wall_t > 0 else 0.0)
         cw = r.get("composite_window_ns", 0.0) or 0.0
-        eff_m.append((useful / cw / peak_per_ns * 100) if cw > 0 else 0.0)
+        eff_m.append((useful / cw / _GEMM_PEAK_FLOPS_PER_NS * 100) if cw > 0 else 0.0)
 
     series = {
         "GEMM util % (theoretical)": gu_t,
@@ -507,6 +536,202 @@ def emit_mac_utilization_theoretical_vs_measured(
     )
 
 
+# ── per-PE composite-pipeline plots: HBM BW util & GEMM TFLOPS ─────────
+#
+# These two plots evaluate the composite GEMM kernel against the
+# per-PE HBM bandwidth (256 GB/s) and per-PE GEMM peak (8 TFLOP/s).
+# Both metrics use composite_window_ns as the denominator — the
+# composite-internal pipeline window, with any pre-composite tl.load
+# of operand A excluded by construction. The two variants on each plot
+# answer different staging questions:
+#
+#   - load_ref: activation A pre-staged on chip, only weight W streamed
+#     from HBM during the composite. ("Weight-only HBM" case.)
+#   - ref_ref:  both A and W streamed from HBM during the composite.
+#     ("Both A and W from HBM" case.)
+#
+# load_load is omitted from these plots — with both operands pre-staged
+# the composite carries no HBM traffic and the BW metric collapses to 0,
+# which is uninformative against the saturation question.
+
+# Per-tile DMA payload sizes (one TILE_M x TILE_K or TILE_K x TILE_N
+# read; one TILE_M x TILE_N write per output tile).
+_PIPELINE_VARIANTS = ("load_ref", "ref_ref")
+_VARIANT_LABEL = {
+    "load_ref": "Activation pre-staged",
+    "ref_ref":  "A and W both from HBM",
+}
+_VARIANT_COLOR = {
+    "load_ref": "#3B82F6",   # blue
+    "ref_ref":  "#F59E0B",   # orange
+}
+
+# Per-PE HBM peak BW (GB/s); matches topology.yaml hbm_to_router_bw_gbs
+# and (8 PCs / PE) x (32 GB/s / PC).
+_HBM_BW_PER_PE_GBS = 256.0
+# Per-PE GEMM peak (TFLOP/s); matches topology.yaml pe_gemm.peak_tflops_f16.
+_GEMM_PEAK_TFLOPS = 8.0
+
+
+def _shape_summary(rows: list[dict], tile: dict):
+    """Common shape-axis state for the BW / TFLOPS plots.
+
+    Returns (shape_labels, tile_counts, flagged, shape_keys) — only
+    shapes that actually appear in `rows` for any of the pipeline
+    variants, in their original order.
+    """
+    seen: list[tuple[int, int, int]] = []
+    by_key: dict[tuple[int, int, int], dict] = {}
+    for r in rows:
+        key = (r["M"], r["K"], r["N"])
+        if key not in by_key:
+            by_key[key] = r
+            seen.append(key)
+    shape_labels = [_shape_label(by_key[k]) for k in seen]
+    flagged = [_under_tile(*k, tile["M"], tile["K"], tile["N"]) for k in seen]
+    tile_counts = [by_key[k]["tile_count_expected"] for k in seen]
+    return shape_labels, tile_counts, flagged, seen
+
+
+def _composite_bytes_through_hbm(r: dict, tile: dict) -> float:
+    """Total bytes transferred via HBM during the composite window.
+
+    Derived from the simulator's per-stage record counts so it
+    automatically reflects the variant's pre-staging choices:
+    load_ref skips A reads, ref_ref reads both A and B, load_load
+    reads neither. Each DMA record carries a full hardware tile
+    (TILE_M x TILE_K read for A, TILE_K x TILE_N read for B,
+    TILE_M x TILE_N write for the output), regardless of how much
+    of that tile is "useful" — which is the actual byte traffic the
+    HBM controller sees.
+    """
+    bpe = 2  # f16; matches matmul_composite.py default
+    a_tile_bytes = tile["M"] * tile["K"] * bpe
+    b_tile_bytes = tile["K"] * tile["N"] * bpe
+    out_tile_bytes = tile["M"] * tile["N"] * bpe
+    # The schedule emits one DMA_READ per (A,B) per non-pinned operand.
+    # Without per-record operand tags in the aggregated JSON, treat all
+    # DMA_READ records as B-shaped reads on load_ref (B is the only
+    # streamed operand) and average the two tile shapes on ref_ref.
+    # In the current topology both tile shapes are 4 KiB so this
+    # distinction is moot — kept symbolic for future tile-size sweeps.
+    variant = r.get("variant")
+    n_read = r["stages"]["DMA_READ"]["record_count"]
+    n_write = r["stages"]["DMA_WRITE"]["record_count"]
+    if variant == "ref_ref":
+        # A reads + B reads, roughly half each.
+        read_bytes = (n_read // 2) * a_tile_bytes \
+                   + (n_read - n_read // 2) * b_tile_bytes
+    else:
+        read_bytes = n_read * b_tile_bytes
+    write_bytes = n_write * out_tile_bytes
+    return float(read_bytes + write_bytes)
+
+
+def _by_variant(rows: list[dict], variant: str) -> dict:
+    return {(r["M"], r["K"], r["N"]): r
+            for r in rows if r.get("variant") == variant}
+
+
+def emit_hbm_bw_utilization(
+    sweep_json: Path | str = DEFAULT_SWEEP_JSON,
+    out_dir: Path | str = DEFAULT_PLOTS_DIR,
+) -> str | None:
+    """Per-PE HBM bandwidth utilization during the composite window.
+
+    Two grouped bars per shape: load_ref ("activation pre-staged")
+    vs ref_ref ("both A and W from HBM"). Y-axis is BW
+    util % against the per-PE HBM ceiling (\\SI{256}{\\giga\\byte\\per
+    \\second}).
+    """
+    data = _load_sweep_data(sweep_json)
+    rows = [r for r in data["rows"] if r.get("variant") in _PIPELINE_VARIANTS]
+    if not rows:
+        return None
+    tile = data["tile_sizes"]
+    shape_labels, tile_counts, flagged, shape_keys = _shape_summary(rows, tile)
+
+    series: dict = {}
+    for variant in _PIPELINE_VARIANTS:
+        by = _by_variant(rows, variant)
+        vals = []
+        for k in shape_keys:
+            r = by.get(k)
+            if r is None:
+                vals.append(0.0)
+                continue
+            cw = r.get("composite_window_ns", 0.0) or 0.0
+            if cw <= 0:
+                vals.append(0.0)
+                continue
+            bytes_hbm = _composite_bytes_through_hbm(r, tile)
+            # bytes_hbm [B] / cw [ns]  =  GB/s  (since 1 B/ns = 1 GB/s).
+            bw_gbs = bytes_hbm / cw
+            vals.append(bw_gbs / _HBM_BW_PER_PE_GBS * 100.0)
+        series[_VARIANT_LABEL[variant]] = vals
+    colors = {_VARIANT_LABEL[v]: _VARIANT_COLOR[v] for v in _PIPELINE_VARIANTS}
+
+    return _grouped_bar_png(
+        "gemm_hbm_bw_util.png", out_dir=Path(out_dir),
+        title="GEMM HBM bandwidth utilization (per PE, composite window)",
+        subtitle=(f"BW util = bytes through HBM / composite window "
+                  f"/ {_HBM_BW_PER_PE_GBS:.0f} GB/s. "
+                  f"Saturation → memory-bound regime."),
+        shape_labels=shape_labels, tile_counts=tile_counts, flagged=flagged,
+        series=series, colors=colors, y_label="HBM BW utilization (%)",
+        threshold=100.0,
+    )
+
+
+def emit_per_pe_tflops(
+    sweep_json: Path | str = DEFAULT_SWEEP_JSON,
+    out_dir: Path | str = DEFAULT_PLOTS_DIR,
+) -> str | None:
+    """Per-PE GEMM throughput in TFLOP/s during the composite window.
+
+    Two grouped bars per shape (load_ref vs ref_ref). Reference line
+    at the per-PE GEMM engine peak (\\SI{8}{\\tera\\flop\\per\\second}).
+    """
+    data = _load_sweep_data(sweep_json)
+    rows = [r for r in data["rows"] if r.get("variant") in _PIPELINE_VARIANTS]
+    if not rows:
+        return None
+    tile = data["tile_sizes"]
+    shape_labels, tile_counts, flagged, shape_keys = _shape_summary(rows, tile)
+
+    series: dict = {}
+    for variant in _PIPELINE_VARIANTS:
+        by = _by_variant(rows, variant)
+        vals = []
+        for k in shape_keys:
+            r = by.get(k)
+            if r is None:
+                vals.append(0.0)
+                continue
+            cw = r.get("composite_window_ns", 0.0) or 0.0
+            if cw <= 0:
+                vals.append(0.0)
+                continue
+            M, K, N = r["M"], r["K"], r["N"]
+            useful = 2 * M * K * N
+            # flops / ns  -> GFLOP/s; / 1000  -> TFLOP/s.
+            tflops = useful / cw / 1000.0
+            vals.append(tflops)
+        series[_VARIANT_LABEL[variant]] = vals
+    colors = {_VARIANT_LABEL[v]: _VARIANT_COLOR[v] for v in _PIPELINE_VARIANTS}
+
+    return _grouped_bar_png(
+        "gemm_per_pe_tflops.png", out_dir=Path(out_dir),
+        title="GEMM achieved throughput (per PE, composite window)",
+        subtitle=(f"TFLOP/s = 2 · M · K · N / composite "
+                  f"window. Peak per PE = "
+                  f"{_GEMM_PEAK_TFLOPS:.0f} TFLOP/s (dashed line)."),
+        shape_labels=shape_labels, tile_counts=tile_counts, flagged=flagged,
+        series=series, colors=colors, y_label="Per-PE TFLOP/s",
+        threshold=_GEMM_PEAK_TFLOPS,
+    )
+
+
 def emit_all_gemm_plots(
     sweep_json: Path | str = DEFAULT_SWEEP_JSON,
     out_dir: Path | str = DEFAULT_PLOTS_DIR,
@@ -515,7 +740,9 @@ def emit_all_gemm_plots(
     paths = []
     for fn in (emit_stage_breakdown,
                emit_mac_utilization_measured,
-               emit_mac_utilization_theoretical_vs_measured):
+               emit_mac_utilization_theoretical_vs_measured,
+               emit_hbm_bw_utilization,
+               emit_per_pe_tflops):
         p = fn(sweep_json, out_dir)
         if p:
             paths.append(p)