gqa: reorganize benches into gqa_helpers/ subpackage; drop legacy headline
Splits the GQA helpers into a dedicated subpackage to make room for the
prefill 4-cases study (next commit) and a single umbrella bench
(milestone-1h-gqa, after that).
Layout:
benches/gqa_helpers/
long_ctx/ — decode 4-cases kernels + sweep runner
short_ctx/ — prefill/decode short-context kernels
shared/ — _gqa_panel_helpers + decode_opt2 (context-agnostic)
The registry audit now skips subpackages so gqa_helpers/ (without a
leading underscore) doesn't get audited for @bench decorators.
Also drops the legacy milestone-gqa-headline bench, its
_gqa_attention_prefill_long kernel, 6 dependent prefill tests, the
paper_gqa_latency.py report harness, and the 3 stale headline-derived
PNGs the §6 wire-up referenced (paper will re-pull from the new
1H_milestone_output/gqa/long_ctx/ once §6 is updated).
The _ccl_cfg and _summarize_op_log helpers used to live in the
headline bench; extracted them to gqa_helpers/shared/_gqa_panel_helpers.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -1,101 +0,0 @@
|
||||
"""Report harness: GQA end-to-end latency + op/engine breakdown.
|
||||
|
||||
Isolated under ``scripts/paper/`` for the 1H codesign report only — it is
|
||||
NOT a registered bench and does not touch other people's benches. It
|
||||
reuses the existing GQA headline panels (the real GQA kernels wired in
|
||||
``milestone_gqa_headline``) but, unlike that milestone (which records only
|
||||
op-counts), it also harvests per-panel end-to-end latency and per-engine
|
||||
occupancy from ``result.engine.op_log``.
|
||||
|
||||
Latency definition (same window convention as ``milestone_1h_gemm``):
|
||||
end_to_end_ns = max(r.t_end) - min(r.t_start) over all op_log records.
|
||||
|
||||
Output: docs/report/1H-codesign-paper/figures/gqa_latency.json
|
||||
|
||||
Run:
|
||||
python scripts/paper/paper_gqa_latency.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from kernbench.benches.milestone_gqa_headline import (
|
||||
_PANEL_DISPATCH,
|
||||
_PANELS,
|
||||
_make_bench_fn,
|
||||
_summarize_op_log,
|
||||
)
|
||||
|
||||
_REPORT_DIR = Path(__file__).resolve().parents[2] / "docs" / "report" / "1H-codesign-paper"
|
||||
_FIG_DIR = _REPORT_DIR / "figures"
|
||||
_OUT_JSON = _FIG_DIR / "gqa_latency.json"
|
||||
|
||||
# PE engine component suffixes whose occupancy we break out.
|
||||
_ENGINES = ("pe_gemm", "pe_math", "pe_dma", "pe_fetch_store", "pe_ipcq", "pe_cpu")
|
||||
|
||||
|
||||
def _occupancy_ns(op_log, suffix: str) -> float:
|
||||
return sum(
|
||||
r.t_end - r.t_start
|
||||
for r in op_log
|
||||
if r.component_id.endswith("." + suffix)
|
||||
)
|
||||
|
||||
|
||||
def _end_to_end_ns(op_log) -> float:
|
||||
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 _run_panel(panel: str, topology: str) -> dict:
|
||||
from kernbench.runtime_api.bench_runner import run_bench
|
||||
from kernbench.runtime_api.types import resolve_device
|
||||
from kernbench.sim_engine.engine import GraphEngine
|
||||
from kernbench.topology.builder import resolve_topology
|
||||
|
||||
topo = resolve_topology(topology)
|
||||
result = run_bench(
|
||||
topology=topo,
|
||||
bench_fn=_make_bench_fn(panel),
|
||||
device=resolve_device(None),
|
||||
engine_factory=lambda t, d: GraphEngine(
|
||||
getattr(t, "topology_obj", t), enable_data=True,
|
||||
),
|
||||
)
|
||||
if not result.completion.ok:
|
||||
raise RuntimeError(f"panel {panel!r} failed: {result.completion}")
|
||||
|
||||
op_log = result.engine.op_log
|
||||
kind, params = _PANEL_DISPATCH[panel]
|
||||
return {
|
||||
"panel": panel,
|
||||
"kind": kind,
|
||||
**params,
|
||||
"latency_ns": _end_to_end_ns(op_log),
|
||||
"op_log_summary": _summarize_op_log(op_log),
|
||||
"engine_occupancy_ns": {
|
||||
eng: _occupancy_ns(op_log, eng) for eng in _ENGINES
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
topology = "topology.yaml"
|
||||
rows = [_run_panel(panel, topology) for panel in _PANELS]
|
||||
_FIG_DIR.mkdir(parents=True, exist_ok=True)
|
||||
out = {"version": 1, "panels": list(_PANELS), "rows": rows}
|
||||
_OUT_JSON.write_text(json.dumps(out, indent=2))
|
||||
print(f"wrote {_OUT_JSON}")
|
||||
for r in rows:
|
||||
s = r["op_log_summary"]
|
||||
print(
|
||||
f" {r['panel']:24s} latency={r['latency_ns']:10.1f} ns "
|
||||
f"gemm={s['gemm_count']:3d} ipcq={s['ipcq_copy_count']:3d} "
|
||||
f"dma_rd={s['dma_read_count']:3d} dma_wr={s['dma_write_count']:2d}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,12 +1,13 @@
|
||||
"""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
|
||||
milestone-gqa-decode-long-ctx-4cases``) and writes four 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
|
||||
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 per-PE KV bytes per case
|
||||
gqa_decode_long_ctx_4cases_parallelism.png per-PE S_local (compute work)
|
||||
|
||||
Run (after the bench):
|
||||
GQA_DECODE_LONG_CTX_4CASES_RUN=1 python -m kernbench.cli.main run \\
|
||||
@@ -24,11 +25,12 @@ 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 = (
|
||||
# Sweep JSON + PNGs live together under the bench output dir.
|
||||
_FIG_DIR = (
|
||||
_REPO_ROOT / "src" / "kernbench" / "benches"
|
||||
/ "1H_milestone_output" / "gqa_decode_long_ctx_4cases" / "sweep.json"
|
||||
/ "1H_milestone_output" / "gqa" / "long_ctx"
|
||||
)
|
||||
_SWEEP_JSON = _FIG_DIR / "sweep_decode.json"
|
||||
|
||||
# Panel name → (short label, case ordinal for left-to-right plot order).
|
||||
_CASE_INFO = {
|
||||
@@ -98,39 +100,63 @@ def _plot_traffic(rows: list[dict]) -> Path:
|
||||
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.
|
||||
def _s_local_per_pe(panel: str, *, S_kv: int, C: int, P: int) -> int:
|
||||
"""S_local (token count) each PE attends over locally.
|
||||
|
||||
Encodes the cube/pe sharding axes from the panel name:
|
||||
cube_sp_pe_tp (Case 1): S_kv / C (pe=replicate within cube)
|
||||
cube_repl_pe_tp (Case 2): S_kv (pe=replicate; only 1 PE works)
|
||||
cube_repl_pe_sp (Case 3): S_kv / P (pe=row_wise within cube)
|
||||
cube_sp_pe_sp (Case 4): S_kv / (C·P) (★ 64-way split)
|
||||
"""
|
||||
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)
|
||||
return S_per_cube // P if "pe_sp" in panel else S_per_cube
|
||||
|
||||
|
||||
def _active_pe_count(panel: str, *, C: int, P: int) -> int:
|
||||
"""Number of PEs doing non-idle attention work.
|
||||
|
||||
cube_sp_pe_tp (Case 1): C (PE 0 of each cube; 7 PEs idle per cube)
|
||||
cube_repl_pe_tp (Case 2): 1 (only PE 0 of CUBE 0)
|
||||
cube_repl_pe_sp (Case 3): C·P (all PEs busy, but cubes are redundant)
|
||||
cube_sp_pe_sp (Case 4): C·P (all 64 PEs doing unique work)
|
||||
"""
|
||||
if "cube_repl" in panel and "pe_tp" in panel:
|
||||
return 1
|
||||
if "cube_sp" in panel and "pe_tp" in panel:
|
||||
return C
|
||||
return C * P
|
||||
|
||||
|
||||
def _kv_bytes_per_pe(panel: str, *, S_kv: int, h_kv: int,
|
||||
d_head: int, C: int, P: int) -> int:
|
||||
"""KV bytes a single PE references (K + V, f16, 2 B/elem)."""
|
||||
s_local = _s_local_per_pe(panel, S_kv=S_kv, C=C, P=P)
|
||||
return 2 * s_local * h_kv * d_head * 2
|
||||
|
||||
|
||||
def _plot_memory(rows: list[dict]) -> Path:
|
||||
"""Per-PE KV bytes — Case 4 wins (64-way split)."""
|
||||
rows = _sorted_by_case(rows)
|
||||
labels = [_CASE_INFO[r["panel"]][0] for r in rows]
|
||||
mib_per_cube = [
|
||||
_kv_bytes_per_cube(
|
||||
mib_per_pe = [
|
||||
_kv_bytes_per_pe(
|
||||
r["panel"], S_kv=r["S_kv"], h_kv=r["h_kv"],
|
||||
d_head=r["d_head"], C=r["C"],
|
||||
d_head=r["d_head"], C=r["C"], P=r["P"],
|
||||
) / (1024 * 1024)
|
||||
for r in rows
|
||||
]
|
||||
# SP = blue (efficient); Repl = red (wasteful).
|
||||
colors = ["#3b6ea5", "#c0504d", "#c0504d", "#3b6ea5"]
|
||||
colors = ["#888", "#c0504d", "#888", "#3b6ea5"] # 4 highlighted, 2 marked red
|
||||
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)")
|
||||
bars = ax.bar(labels, mib_per_pe, color=colors, width=0.6)
|
||||
ax.set_ylabel("KV bytes per PE (MiB, K + V, f16)")
|
||||
ax.set_title(
|
||||
"Long-context decode 4-cases — KV memory per cube\n"
|
||||
"Long-context decode 4-cases — KV memory per PE\n"
|
||||
"(one KV-head group; per-layer, per-token state)"
|
||||
)
|
||||
ax.bar_label(bars, fmt="%.1f", padding=3, fontsize=9)
|
||||
ax.bar_label(bars, fmt="%.3f", padding=3, fontsize=9)
|
||||
ax.grid(axis="y", ls=":", alpha=0.5)
|
||||
ax.set_ylim(0, max(mib_per_cube) * 1.15)
|
||||
ax.set_ylim(0, max(mib_per_pe) * 1.15)
|
||||
fig.tight_layout()
|
||||
out = _FIG_DIR / "gqa_decode_long_ctx_4cases_memory.png"
|
||||
fig.savefig(out, dpi=150)
|
||||
@@ -138,15 +164,44 @@ def _plot_memory(rows: list[dict]) -> Path:
|
||||
return out
|
||||
|
||||
|
||||
def _plot_parallelism(rows: list[dict]) -> Path:
|
||||
"""Total active PE-token compute load — exposes Case 3's redundancy."""
|
||||
rows = _sorted_by_case(rows)
|
||||
labels = [_CASE_INFO[r["panel"]][0] for r in rows]
|
||||
total_work = [
|
||||
_active_pe_count(r["panel"], C=r["C"], P=r["P"])
|
||||
* _s_local_per_pe(r["panel"], S_kv=r["S_kv"], C=r["C"], P=r["P"])
|
||||
for r in rows
|
||||
]
|
||||
colors = ["#888", "#888", "#c0504d", "#3b6ea5"] # 4 highlighted, 3 marked red
|
||||
fig, ax = plt.subplots(figsize=(8.0, 4.5))
|
||||
bars = ax.bar(labels, total_work, color=colors, width=0.6)
|
||||
ax.set_ylabel("active-PE × S_local (PE-tokens; lower ⇒ less wasted work)")
|
||||
ax.set_title(
|
||||
"Long-context decode 4-cases — total compute load across active PEs\n"
|
||||
"(Case 3 replicates the full K/V across 8 cubes ⇒ 8× wasted PE-tokens)"
|
||||
)
|
||||
ax.bar_label(bars, fmt="%d", padding=3, fontsize=9)
|
||||
ax.grid(axis="y", ls=":", alpha=0.5)
|
||||
ax.set_ylim(0, max(total_work) * 1.15)
|
||||
fig.tight_layout()
|
||||
out = _FIG_DIR / "gqa_decode_long_ctx_4cases_parallelism.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)
|
||||
p4 = _plot_parallelism(rows)
|
||||
print(f"wrote {p1}")
|
||||
print(f"wrote {p2}")
|
||||
print(f"wrote {p3}")
|
||||
print(f"wrote {p4}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user