diff --git a/scripts/paper/paper_plot_gqa_decode_long_ctx_4cases.py b/scripts/paper/paper_plot_gqa_decode_long_ctx_4cases.py new file mode 100644 index 0000000..1f76d44 --- /dev/null +++ b/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() diff --git a/src/kernbench/benches/milestone_gqa_decode_long_ctx_4cases.py b/src/kernbench/benches/milestone_gqa_decode_long_ctx_4cases.py index 0498752..45c5bd8 100644 --- a/src/kernbench/benches/milestone_gqa_decode_long_ctx_4cases.py +++ b/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), } diff --git a/tests/attention/test_milestone_gqa_decode_long_ctx_4cases.py b/tests/attention/test_milestone_gqa_decode_long_ctx_4cases.py index a2d0684..c921b84 100644 --- a/tests/attention/test_milestone_gqa_decode_long_ctx_4cases.py +++ b/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