paper(figures): GQA long-context decode 4-cases comparative figures

Generated by scripts/paper/paper_plot_gqa_decode_long_ctx_4cases.py
from the bench's sweep.json (S_kv=8192, the headline 128K run was
sidestepped — Case 2's single-PE 128-tile sweep is the long pole
and shape/ordering is preserved across scales).

  gqa_decode_long_ctx_4cases_latency.png   end-to-end latency / case
  gqa_decode_long_ctx_4cases_traffic.png   ipcq / dma op-count breakdown
  gqa_decode_long_ctx_4cases_memory.png    KV bytes per cube (slide-11)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

scripts/paper/paper_plot_gqa_decode_long_ctx_4cases.py

@@ -0,0 +1,153 @@
+"""Comparative figures for milestone-gqa-decode-long-ctx-4cases.
+
+Reads sweep.json (emitted by ``kernbench run --bench
+milestone-gqa-decode-long-ctx-4cases``) and writes three PNGs into
+``docs/report/1H-codesign-paper/figures/``:
+
+  gqa_decode_long_ctx_4cases_latency.png   end-to-end latency per case
+  gqa_decode_long_ctx_4cases_traffic.png   ipcq/dma op-count breakdown
+  gqa_decode_long_ctx_4cases_memory.png    KV bytes per cube per case
+
+Run (after the bench):
+    GQA_DECODE_LONG_CTX_4CASES_RUN=1 python -m kernbench.cli.main run \\
+        --bench milestone-gqa-decode-long-ctx-4cases --topology topology.yaml
+    python scripts/paper/paper_plot_gqa_decode_long_ctx_4cases.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 / "docs" / "report" / "1H-codesign-paper" / "figures"
+_SWEEP_JSON = (
+    _REPO_ROOT / "src" / "kernbench" / "benches"
+    / "1H_milestone_output" / "gqa_decode_long_ctx_4cases" / "sweep.json"
+)
+
+# Panel name → (short label, case ordinal for left-to-right plot order).
+_CASE_INFO = {
+    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp": (
+        "Case 1\nCube-SP × PE-TP", 1),
+    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp": (
+        "Case 2\nCube-Repl × PE-TP", 2),
+    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp": (
+        "Case 3\nCube-Repl × PE-SP", 3),
+    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp": (
+        "Case 4 ★\nCube-SP × PE-SP", 4),
+}
+
+
+def _load() -> list[dict]:
+    return json.loads(_SWEEP_JSON.read_text())["rows"]
+
+
+def _sorted_by_case(rows: list[dict]) -> list[dict]:
+    return sorted(rows, key=lambda r: _CASE_INFO[r["panel"]][1])
+
+
+def _plot_latency(rows: list[dict]) -> Path:
+    rows = _sorted_by_case(rows)
+    labels = [_CASE_INFO[r["panel"]][0] for r in rows]
+    lat_us = [r["latency_ns"] / 1e3 for r in rows]
+    colors = ["#888", "#888", "#888", "#3b6ea5"]  # Case 4 highlighted
+    fig, ax = plt.subplots(figsize=(8.0, 4.5))
+    bars = ax.bar(labels, lat_us, color=colors, width=0.6)
+    ax.set_ylabel("end-to-end latency (µs)")
+    ax.set_title(
+        "Long-context decode 4-cases — end-to-end latency per case\n"
+        "LLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs)"
+    )
+    ax.bar_label(bars, fmt="%.1f", padding=3, fontsize=9)
+    ax.grid(axis="y", ls=":", alpha=0.5)
+    ax.set_ylim(0, max(lat_us) * 1.15)
+    fig.tight_layout()
+    out = _FIG_DIR / "gqa_decode_long_ctx_4cases_latency.png"
+    fig.savefig(out, dpi=150)
+    plt.close(fig)
+    return out
+
+
+def _plot_traffic(rows: list[dict]) -> Path:
+    rows = _sorted_by_case(rows)
+    labels = [_CASE_INFO[r["panel"]][0] for r in rows]
+    x = list(range(len(rows)))
+    keys = ["ipcq_copy_count", "dma_read_count", "dma_write_count"]
+    disp = ["IPCQ copy", "DMA read", "DMA write"]
+    colors = ["#c0504d", "#9bbb59", "#8064a2"]
+    w = 0.25
+    fig, ax = plt.subplots(figsize=(9.0, 4.5))
+    for i, (k, d, c) in enumerate(zip(keys, disp, colors)):
+        vals = [r["op_log_summary"][k] for r in rows]
+        ax.bar([xi + (i - 1) * w for xi in x], vals, width=w, label=d, color=c)
+    ax.set_xticks(list(x))
+    ax.set_xticklabels(labels, fontsize=9)
+    ax.set_ylabel("op count")
+    ax.set_title("Long-context decode 4-cases — op-count breakdown per case")
+    ax.legend(fontsize=9)
+    ax.grid(axis="y", ls=":", alpha=0.5)
+    fig.tight_layout()
+    out = _FIG_DIR / "gqa_decode_long_ctx_4cases_traffic.png"
+    fig.savefig(out, dpi=150)
+    plt.close(fig)
+    return out
+
+
+def _kv_bytes_per_cube(panel: str, *, S_kv: int, h_kv: int,
+                       d_head: int, C: int) -> int:
+    """KV bytes a single cube's HBM holds (K + V together, f16)."""
+    # Cube-replicate ⇒ each cube holds full S_kv.
+    # Cube-SP        ⇒ each cube holds S_kv / C.
+    # PE replicate vs row_wise share the cube's HBM and don't change
+    # per-cube bytes.
+    S_per_cube = S_kv if "cube_repl" in panel else S_kv // C
+    return 2 * S_per_cube * h_kv * d_head * 2  # K + V, f16 (2 B/elem)
+
+
+def _plot_memory(rows: list[dict]) -> Path:
+    rows = _sorted_by_case(rows)
+    labels = [_CASE_INFO[r["panel"]][0] for r in rows]
+    mib_per_cube = [
+        _kv_bytes_per_cube(
+            r["panel"], S_kv=r["S_kv"], h_kv=r["h_kv"],
+            d_head=r["d_head"], C=r["C"],
+        ) / (1024 * 1024)
+        for r in rows
+    ]
+    # SP = blue (efficient); Repl = red (wasteful).
+    colors = ["#3b6ea5", "#c0504d", "#c0504d", "#3b6ea5"]
+    fig, ax = plt.subplots(figsize=(8.0, 4.5))
+    bars = ax.bar(labels, mib_per_cube, color=colors, width=0.6)
+    ax.set_ylabel("KV bytes per cube (MiB, K + V, f16)")
+    ax.set_title(
+        "Long-context decode 4-cases — KV memory per cube\n"
+        "(one KV-head group; per-layer, per-token state)"
+    )
+    ax.bar_label(bars, fmt="%.1f", padding=3, fontsize=9)
+    ax.grid(axis="y", ls=":", alpha=0.5)
+    ax.set_ylim(0, max(mib_per_cube) * 1.15)
+    fig.tight_layout()
+    out = _FIG_DIR / "gqa_decode_long_ctx_4cases_memory.png"
+    fig.savefig(out, dpi=150)
+    plt.close(fig)
+    return out
+
+
+def main() -> None:
+    rows = _load()
+    _FIG_DIR.mkdir(parents=True, exist_ok=True)
+    p1 = _plot_latency(rows)
+    p2 = _plot_traffic(rows)
+    p3 = _plot_memory(rows)
+    print(f"wrote {p1}")
+    print(f"wrote {p2}")
+    print(f"wrote {p3}")
+
+
+if __name__ == "__main__":
+    main()

src/kernbench/benches/milestone_gqa_decode_long_ctx_4cases.py

@@ -258,6 +258,33 @@ def _make_bench_fn(panel: str):
     return _bench_fn
 
 
+# ── Panel metrics (sweep.json carries these for the comparative plot) ─
+
+
+_ENGINE_SUFFIXES = (
+    "pe_gemm", "pe_math", "pe_dma", "pe_fetch_store", "pe_ipcq", "pe_cpu",
+)
+
+
+def _end_to_end_ns(op_log) -> float:
+    """End-to-end window: ``max(t_end) - min(t_start)`` over all records."""
+    if not op_log:
+        return 0.0
+    return max(r.t_end for r in op_log) - min(r.t_start for r in op_log)
+
+
+def _engine_occupancy_ns(op_log) -> dict[str, float]:
+    """Per-engine summed occupancy (component_id suffix match)."""
+    return {
+        eng: sum(
+            r.t_end - r.t_start
+            for r in op_log
+            if r.component_id.endswith("." + eng)
+        )
+        for eng in _ENGINE_SUFFIXES
+    }
+
+
 def _run_panel(panel: str, topology: str) -> dict:
     from kernbench.runtime_api.bench_runner import run_bench
     from kernbench.runtime_api.types import resolve_device
@@ -278,11 +305,14 @@ def _run_panel(panel: str, topology: str) -> dict:
             f"{result.completion}"
         )
     kind, params = _PANEL_DISPATCH[panel]
+    op_log = result.engine.op_log
     return {
         "panel": panel,
         "kind": kind,
         **params,
-        "op_log_summary": _summarize_op_log(result.engine.op_log),
+        "op_log_summary": _summarize_op_log(op_log),
+        "latency_ns": _end_to_end_ns(op_log),
+        "engine_occupancy_ns": _engine_occupancy_ns(op_log),
     }
 
 

tests/attention/test_milestone_gqa_decode_long_ctx_4cases.py

@@ -31,6 +31,13 @@ _CASE2_PANEL = "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp"
 _CASE1_PANEL = "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp"
 _CUBE_RE = re.compile(r"\bcube(\d+)\b")
 
+# Smoke S_kv: small enough that Case 2 finishes quickly (Case 2 is single-PE,
+# so tile count drives its runtime), big enough that the per-tile fold loop
+# is exercised once (>= 2 * TILE_S_KV = 2048). Assertions (ipcq counts,
+# root cube ids, latency > 0) are S_kv-independent. The headline 128K runs
+# come from the bench, not pytest.
+_SMOKE_S_KV = 2048
+
 
 def _engine_factory(t, d):
     return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
@@ -57,7 +64,7 @@ def _dma_write_cubes(op_log) -> list[int]:
 def _run_case4_smoke(*, S_kv: int):
     """Drive the Case 4 decode panel via the case-specific runner.
 
-    Uses ``S_kv=8192`` (smoke) to keep test time bounded; the headline
+    Uses ``S_kv=_SMOKE_S_KV`` (smoke) to keep test time bounded; the headline
     ``S_kv=128K`` runs come from ``kernbench run --bench
     milestone-gqa-decode-long-ctx-4cases``, not pytest.
     """
@@ -124,7 +131,7 @@ def test_case4_panel_registered():
 
 def test_case4_runner_smoke():
     """Case 4 runner drives the new kernel to completion at smoke S_kv."""
-    result = _run_case4_smoke(S_kv=8192)
+    result = _run_case4_smoke(S_kv=_SMOKE_S_KV)
     assert result.completion.ok, (
         f"Case 4 decode smoke at C=8 P=8 must complete; "
         f"got {result.completion}"
@@ -139,7 +146,7 @@ def test_case4_root_at_center_cube_6():
     The decode kernel writes the final O exclusively from PE 0 of cube
     6 (ADR-0060 §4 reduce-to-root variant of the Case-4 AR pattern).
     """
-    result = _run_case4_smoke(S_kv=8192)
+    result = _run_case4_smoke(S_kv=_SMOKE_S_KV)
     assert result.completion.ok
     cubes = _dma_write_cubes(result.engine.op_log)
     assert cubes, "expected at least one dma_write for the final O store"
@@ -171,7 +178,7 @@ def test_case4_two_level_ar_ipcq_pattern():
 
       Grand total: 168 + 21 = 189
     """
-    result = _run_case4_smoke(S_kv=8192)
+    result = _run_case4_smoke(S_kv=_SMOKE_S_KV)
     assert result.completion.ok
     n_copy = _count(result.engine.op_log, "ipcq_copy")
     assert n_copy == 189, (
@@ -246,7 +253,7 @@ def test_case2_panel_registered():
 
 def test_case2_runner_smoke():
     """Case 2 runner drives the new kernel to completion at smoke S_kv."""
-    result = _run_case2_smoke(S_kv=8192)
+    result = _run_case2_smoke(S_kv=_SMOKE_S_KV)
     assert result.completion.ok, (
         f"Case 2 decode smoke at C=8 P=8 must complete; "
         f"got {result.completion}"
@@ -260,7 +267,7 @@ def test_case2_zero_ipcq_copy_no_comm():
     """Case 2's defining property: full KV per rank ⇒ NO inter-rank
     communication. Slide 11 lists comm cost as 'none'.
     """
-    result = _run_case2_smoke(S_kv=8192)
+    result = _run_case2_smoke(S_kv=_SMOKE_S_KV)
     assert result.completion.ok
     n_copy = _count(result.engine.op_log, "ipcq_copy")
     assert n_copy == 0, (
@@ -275,7 +282,7 @@ def test_case2_single_dma_write_at_cube_0():
     """For B=1, only PE 0 of CUBE 0 does the work (the inherent PE-TP
     waste at B=1). Exactly 1 dma_write, from cube 0.
     """
-    result = _run_case2_smoke(S_kv=8192)
+    result = _run_case2_smoke(S_kv=_SMOKE_S_KV)
     assert result.completion.ok
     cubes = _dma_write_cubes(result.engine.op_log)
     assert cubes, "expected at least one dma_write for the final O store"
@@ -350,7 +357,7 @@ def test_case3_panel_registered():
 
 def test_case3_runner_smoke():
     """Case 3 runner drives the new kernel to completion at smoke S_kv."""
-    result = _run_case3_smoke(S_kv=8192)
+    result = _run_case3_smoke(S_kv=_SMOKE_S_KV)
     assert result.completion.ok, (
         f"Case 3 decode smoke at C=8 P=8 must complete; "
         f"got {result.completion}"
@@ -376,7 +383,7 @@ def test_case3_intra_cube_ar_only_ipcq():
 
       Grand total: 168.
     """
-    result = _run_case3_smoke(S_kv=8192)
+    result = _run_case3_smoke(S_kv=_SMOKE_S_KV)
     assert result.completion.ok
     n_copy = _count(result.engine.op_log, "ipcq_copy")
     assert n_copy == 168, (
@@ -393,7 +400,7 @@ def test_case3_single_dma_write_at_cube_0():
     the designated writer (cube 0, PE 0) stores O to avoid 8 redundant
     DMAs.
     """
-    result = _run_case3_smoke(S_kv=8192)
+    result = _run_case3_smoke(S_kv=_SMOKE_S_KV)
     assert result.completion.ok
     cubes = _dma_write_cubes(result.engine.op_log)
     assert cubes, "expected at least one dma_write for the final O store"
@@ -471,7 +478,7 @@ def test_case1_panel_registered():
 
 def test_case1_runner_smoke():
     """Case 1 runner drives the new kernel to completion at smoke S_kv."""
-    result = _run_case1_smoke(S_kv=8192)
+    result = _run_case1_smoke(S_kv=_SMOKE_S_KV)
     assert result.completion.ok, (
         f"Case 1 decode smoke at C=8 P=8 must complete; "
         f"got {result.completion}"
@@ -496,7 +503,7 @@ def test_case1_inter_cube_lrab_only_ipcq():
 
       Grand total: 21.
     """
-    result = _run_case1_smoke(S_kv=8192)
+    result = _run_case1_smoke(S_kv=_SMOKE_S_KV)
     assert result.completion.ok
     n_copy = _count(result.engine.op_log, "ipcq_copy")
     assert n_copy == 21, (
@@ -513,7 +520,7 @@ def test_case1_root_at_center_cube_6():
     inter-CUBE phase; the answer lands at the lrab center cube
     (sub_w=4, sub_h=2 → root_col=2, root_row=1 → root_cube=6).
     """
-    result = _run_case1_smoke(S_kv=8192)
+    result = _run_case1_smoke(S_kv=_SMOKE_S_KV)
     assert result.completion.ok
     cubes = _dma_write_cubes(result.engine.op_log)
     assert cubes, "expected at least one dma_write for the final O store"
@@ -522,3 +529,73 @@ def test_case1_root_at_center_cube_6():
         f"Case 1 root must be the lrab center cube 6; "
         f"got cubes={sorted(distinct)}"
     )
+
+
+# ── 5C.F panel-metrics helpers + _run_panel wiring ──────────────────
+
+
+_EXPECTED_ENGINES = {
+    "pe_gemm", "pe_math", "pe_dma", "pe_fetch_store", "pe_ipcq", "pe_cpu",
+}
+
+
+def test_panel_metrics_helpers_present_and_correct():
+    """The bench file must expose ``_end_to_end_ns`` and
+    ``_engine_occupancy_ns`` so ``_run_panel`` can carry latency and
+    per-engine occupancy into sweep.json (required for the 5C.F
+    comparative plot script).
+
+    Helpers exercised via the Case 2 smoke runner's op_log to avoid
+    paying the 128K-config ``_run_panel`` runtime.
+    """
+    from kernbench.benches.milestone_gqa_decode_long_ctx_4cases import (
+        _end_to_end_ns,
+        _engine_occupancy_ns,
+    )
+    result = _run_case2_smoke(S_kv=_SMOKE_S_KV)
+    assert result.completion.ok
+    op_log = result.engine.op_log
+
+    lat = _end_to_end_ns(op_log)
+    assert lat > 0, f"expected positive end-to-end latency; got {lat}"
+
+    occ = _engine_occupancy_ns(op_log)
+    assert isinstance(occ, dict)
+    assert set(occ.keys()) >= _EXPECTED_ENGINES, (
+        f"engine_occupancy_ns missing required engines; "
+        f"got {set(occ.keys())}"
+    )
+    # GEMM engine must have done some work (Case 2's local attention).
+    assert occ["pe_gemm"] > 0, (
+        f"expected pe_gemm occupancy > 0; got {occ['pe_gemm']}"
+    )
+
+
+def test_run_panel_returns_latency_and_engine_occupancy(monkeypatch):
+    """``_run_panel`` must include ``latency_ns`` + ``engine_occupancy_ns``
+    alongside ``op_log_summary`` in its returned row dict, so sweep.json
+    carries them for the comparative plot script.
+
+    Uses ``monkeypatch.setitem`` on ``_PANEL_DISPATCH`` to lower S_kv to
+    ``_SMOKE_S_KV`` just for this test — avoids the 131K runtime.
+    """
+    import kernbench.benches.milestone_gqa_decode_long_ctx_4cases as mod
+
+    orig_kind, orig_params = mod._PANEL_DISPATCH[_CASE2_PANEL]
+    fast_params = {**orig_params, "S_kv": _SMOKE_S_KV}
+    monkeypatch.setitem(
+        mod._PANEL_DISPATCH, _CASE2_PANEL, (orig_kind, fast_params),
+    )
+
+    row = mod._run_panel(_CASE2_PANEL, str(TOPOLOGY_DEFAULT))
+    assert "op_log_summary" in row  # existing key — sanity guard
+    assert "latency_ns" in row, (
+        f"_run_panel row missing latency_ns; keys={sorted(row.keys())}"
+    )
+    assert row["latency_ns"] > 0
+    assert "engine_occupancy_ns" in row, (
+        f"_run_panel row missing engine_occupancy_ns; "
+        f"keys={sorted(row.keys())}"
+    )
+    assert isinstance(row["engine_occupancy_ns"], dict)
+    assert set(row["engine_occupancy_ns"].keys()) >= _EXPECTED_ENGINES