diff --git a/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_latency.png b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_latency.png deleted file mode 100644 index 69a179d..0000000 Binary files a/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_latency.png and /dev/null differ diff --git a/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_memory.png b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_memory.png deleted file mode 100644 index 9e40551..0000000 Binary files a/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_memory.png and /dev/null differ diff --git a/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_traffic.png b/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_traffic.png deleted file mode 100644 index 12fafe2..0000000 Binary files a/docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_4cases_traffic.png and /dev/null differ diff --git a/scripts/paper/paper_gqa_latency.py b/scripts/paper/paper_gqa_latency.py deleted file mode 100644 index 1c0a0ad..0000000 --- a/scripts/paper/paper_gqa_latency.py +++ /dev/null @@ -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() diff --git a/scripts/paper/paper_plot_gqa_decode_long_ctx_4cases.py b/scripts/paper/paper_plot_gqa_decode_long_ctx_4cases.py index 1f76d44..dd11408 100644 --- a/scripts/paper/paper_plot_gqa_decode_long_ctx_4cases.py +++ b/scripts/paper/paper_plot_gqa_decode_long_ctx_4cases.py @@ -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__": diff --git a/src/kernbench/benches/1H_milestone_output/gqa_headline/sweep.json b/src/kernbench/benches/1H_milestone_output/gqa_headline/sweep.json deleted file mode 100644 index e73d9cf..0000000 --- a/src/kernbench/benches/1H_milestone_output/gqa_headline/sweep.json +++ /dev/null @@ -1,31 +0,0 @@ -{ - "version": 1, - "panels": [ - "single_kv_group_prefill_gqa_c8_p8" - ], - "config": { - "T_q_prefill": 4, - "T_q_decode": 1, - "S_kv_prefill": 16, - "h_q_decode": 8, - "h_kv_decode": 1, - "d_head": 64 - }, - "rows": [ - { - "panel": "single_kv_group_prefill_gqa_c8_p8", - "kind": "prefill", - "C": 8, - "P": 8, - "T_q": 1024, - "S_kv": 1024, - "d_head": 128, - "op_log_summary": { - "gemm_count": 1024, - "ipcq_copy_count": 896, - "dma_read_count": 192, - "dma_write_count": 64 - } - } - ] -} \ No newline at end of file diff --git a/src/kernbench/benches/_gqa_attention_prefill_long.py b/src/kernbench/benches/_gqa_attention_prefill_long.py deleted file mode 100644 index 77970b8..0000000 --- a/src/kernbench/benches/_gqa_attention_prefill_long.py +++ /dev/null @@ -1,190 +0,0 @@ -"""GQA fused-attention prefill kernel — long context (ADR-0060 §5.5). - -Head-parallel: each CUBE owns one Q head and one KV slice. Over C ring -steps the KV blocks rotate around the C CUBEs (W-send / E-recv) so every -CUBE sees every block; the online-softmax merge folds each step into the -running ``(m, ℓ, O)``. No inter-CUBE reduce — each CUBE writes its own -head's output. - -The ring is **tile-granular** (ADR-0060 §5.5 + ADR-0063 §A.2): each -ring step transmits ``n_tiles`` tiles of size ``TILE_S_KV`` rather than -one full ``S_local`` slice. The kernel loop is nested over -``(ring_step, tile_idx)``, with the bootstrap at ``(k=0, t=0)`` -establishing the persistent ``(m, ℓ, O)`` and every subsequent -iteration folding its tile in inside a ``tl.scratch_scope``. Per-rank -persistent scratch is ``(m, ℓ, O)`` only (~1 KB); per-tile in-scope -scratch is bounded by ``TILE_S_KV`` regardless of ``S_local``. - -Topology / SFR: - - Requires ``configure_sfr_intercube_ring`` — either ``ring_size=C`` - (1D row, ``C ≤ mesh_w``) or ``submesh_shape=(rows, cols)`` (snake - Hamiltonian cycle through a rectangular sub-mesh, for ``C > mesh_w``). - - ``P == 1`` (default): only PE 0 of each CUBE participates; - intra-CUBE PE parallelism is disabled. - - ``P > 1``: all P PEs of each CUBE participate via query-axis - split (ADR-0060 §5.5 last bullet + §B-item-3) — each PE owns - ``T_q/P`` disjoint query rows. Output rows are disjoint per PE, - so no intra-CUBE reduce is needed. Each PE drives its own - same-lane ring via independent IPCQ channels (P parallel rings). - Requires ``T_q % P == 0``. - -Layout caveats: - - GEMMs use ``tl.dot`` (no composite epilogue / ``softmax_scale``). - - K loaded as ``(d_head, tile_s)`` via byte-conserving reshape of - the deployed ``(tile_s, d_head)`` shard (ADR-0060 §3 - reshape-not-transpose caveat). - - No causal masking / step-skip; blocking ``tl.recv``. - - ``TILE_S_KV`` is assumed to divide ``S_local`` evenly; last-tile - padding is not modelled. -""" -from __future__ import annotations - - -TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width). - - -def _merge_running(m, l, O, m_step, l_step, O_step, *, tl): - """Online-softmax merge of two partial ``(m, ℓ, O)`` triples.""" - m_new = tl.maximum(m, m_step) - scale_old = tl.exp(m - m_new) - scale_step = tl.exp(m_step - m_new) - l_new = l * scale_old + l_step * scale_step - O_new = O * scale_old + O_step * scale_step - return m_new, l_new, O_new - - -def gqa_attention_prefill_long_kernel( - q_ptr: int, - k_ptr: int, - v_ptr: int, - o_ptr: int, - T_q: int, - S_kv: int, - d_head: int, - C: int, - P: int = 1, - *, - tl, -) -> None: - """Head-parallel prefill attention with tile-granular Ring KV (ADR-0060 §5.5). - - Tensor layout (per-rank shapes): - Q : (T_q_local, d_head) one head per CUBE; replicated within - a CUBE for P=1, or sharded ``pe="row_wise"`` for P>1 so each - PE owns ``T_q_local = T_q // P`` query rows. - K : (S_kv, d_head) sharded cube_row_wise → each CUBE owns - (S_local, d_head). Tiles loaded as (d_head, tile_s) via - byte-conserving reshape. - V : (S_kv, d_head) sharded cube_row_wise → each CUBE owns - (S_local, d_head). Tiles loaded as (tile_s, d_head). - O : (T_q * C, d_head) sharded cube_row_wise → each CUBE - writes its own ``T_q`` rows; for P>1 each PE writes a - disjoint ``(T_q_local, d_head)`` slice (no intra-CUBE - reduce — disjoint output rows). NO inter-CUBE reduce. - - Algorithm: nested loop over (ring_step k, tile_idx t). At k=0 each - rank loads its own block's tiles from HBM; at k > 0 it receives - tiles from its E neighbour. Tiles are forwarded W to the next - ring step. Each tile's partial is folded into the running - (m, ℓ, O) via online-softmax merge. - """ - pe_id = tl.program_id(axis=0) - if P == 1: - # Head-parallel only: PE 0 of each CUBE participates. - if pe_id != 0: - return - T_q_local = T_q - else: - # Intra-CUBE PE-SP (ADR-0060 §5.5 last bullet + §B-item-3): - # query-axis split across P PEs of each CUBE. Output rows are - # disjoint per PE ⇒ no intra-CUBE reduce. Each PE drives its - # own same-lane ring via independent IPCQ channels. - if pe_id >= P: - return - if T_q % P != 0: - raise ValueError( - f"T_q={T_q} must be divisible by P={P} when P > 1; " - "use P=1 for the PE-0-only fallback path." - ) - T_q_local = T_q // P - - S_local = S_kv // C - n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV - KV_ROW_BYTES = d_head * 2 # f16 - Q = tl.load(q_ptr, shape=(T_q_local, d_head), dtype="f16") - - # ── Bootstrap: (t=0, k=0) — load my own tile 0, establish (m, ℓ, O) ── - # - # Cannot be folded into the (t, k) loop (kernbench-only limitation): - # - persistent (m, ℓ, O) must live OUTSIDE ``tl.scratch_scope``, - # otherwise scope teardown discards them before the next tile's - # merge can read them; - # - kernbench has no scratch-backed initializer — ``tl.zeros`` / - # ``tl.full`` return addr=0 handles with no backing storage, so - # they cannot be overwritten via ``tl.copy_to`` to seed (-inf, 0, 0); - # - the ring-step ordering also requires k=0 to send its loaded - # tile W before any k>0 iteration can recv from E, so the very - # first (t=0, k=0) step has to run outside the scoped loop body - # where the send is conditional on ``C > 1``. - # Triton port: limitation does not apply (SSA tensors stay live - # across iterations); the bootstrap collapses into the loop. - tile_s = min(TILE_S_KV, S_local) - K_t = tl.load(k_ptr, shape=(d_head, tile_s), dtype="f16") - V_t = tl.load(v_ptr, shape=(tile_s, d_head), dtype="f16") - if C > 1: - tl.send(dir="W", src=K_t) - tl.send(dir="W", src=V_t) - scores = tl.dot(Q, K_t) - m = tl.max(scores, axis=-1) - exp_scores = tl.exp(scores - m) - l = tl.sum(exp_scores, axis=-1) - O = tl.dot(exp_scores, V_t) - - # ── Nested loop: outer tile, inner ring step ── - # Each tile propagates all the way through the ring before the next - # tile starts; IPCQ in-flight depth stays at 1 per direction. - # - # Per outer ``t``: - # k=0 : load my own tile ``t`` from HBM - # k=1..C-1 : receive tile ``t`` of a peer's block from E - # (sent by my E neighbour during their previous k step) - # k dict: } -# ── Bench entry ────────────────────────────────────────────────────── +# ── Sweep entry (called by the umbrella ``milestone_1h_gqa``) ──────── -@bench( - name="milestone-gqa-decode-long-ctx-4cases", - description=( - "Long-context decode 4-cases comparative study on the " - "LLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs)." - ), -) -def run(torch) -> None: - """Drive the registered decode case panels; write sweep.json. - - Gated by GQA_DECODE_LONG_CTX_4CASES_RUN=1. - """ +def run_sweep(topology: str = "topology.yaml") -> int: + """Drive all decode case panels; write sweep.json. Returns row count.""" _OUTPUT_DIR.mkdir(parents=True, exist_ok=True) - if not os.environ.get("GQA_DECODE_LONG_CTX_4CASES_RUN"): - raise RuntimeError( - "milestone-gqa-decode-long-ctx-4cases needs " - "GQA_DECODE_LONG_CTX_4CASES_RUN=1." - ) - - topology = os.environ.get( - "GQA_DECODE_LONG_CTX_4CASES_TOPOLOGY", "topology.yaml", - ) rows = [_run_panel(panel, topology) for panel in _PANELS] sweep = { "version": 1, @@ -348,6 +329,5 @@ def run(torch) -> None: "rows": rows, } _SWEEP_JSON.write_text(json.dumps(sweep, indent=2)) - print( - f" milestone-gqa-decode-long-ctx-4cases: {len(rows)} rows -> {_SWEEP_JSON}" - ) + print(f" gqa-decode-long-ctx-4cases: {len(rows)} rows -> {_SWEEP_JSON}") + return len(rows) diff --git a/src/kernbench/benches/gqa_helpers/shared/__init__.py b/src/kernbench/benches/gqa_helpers/shared/__init__.py new file mode 100644 index 0000000..ae7dc03 --- /dev/null +++ b/src/kernbench/benches/gqa_helpers/shared/__init__.py @@ -0,0 +1,5 @@ +"""GQA helpers shared across long-ctx and short-ctx variants. + +Includes the panel-metrics helpers (``_gqa_panel_helpers``) and the +context-agnostic decode-opt2 kernel. +""" diff --git a/src/kernbench/benches/_gqa_attention_decode_opt2.py b/src/kernbench/benches/gqa_helpers/shared/_gqa_attention_decode_opt2.py similarity index 100% rename from src/kernbench/benches/_gqa_attention_decode_opt2.py rename to src/kernbench/benches/gqa_helpers/shared/_gqa_attention_decode_opt2.py diff --git a/src/kernbench/benches/gqa_helpers/shared/_gqa_panel_helpers.py b/src/kernbench/benches/gqa_helpers/shared/_gqa_panel_helpers.py new file mode 100644 index 0000000..1878725 --- /dev/null +++ b/src/kernbench/benches/gqa_helpers/shared/_gqa_panel_helpers.py @@ -0,0 +1,40 @@ +"""Shared helpers for the GQA 4-cases milestone benches. + +Used by both ``milestone_gqa_decode_long_ctx_4cases`` and +``milestone_gqa_prefill_long_ctx_4cases``. Underscore-prefixed so the +benches package auto-discovery (which skips ``_*.py``) doesn't try to +audit it for a ``@bench`` decorator. +""" +from __future__ import annotations + +from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config + + +def _ccl_cfg() -> dict: + """Resolve the lrab hierarchical AllReduce CCL config.""" + return resolve_algorithm_config( + load_ccl_config(), name="lrab_hierarchical_allreduce", + ) + + +def _summarize_op_log(op_log) -> dict[str, int]: + """Per-panel op_log counts (gemm, ipcq_copy, dma_read, dma_write).""" + gemm_count = 0 + ipcq_copy_count = 0 + dma_read_count = 0 + dma_write_count = 0 + for r in op_log: + if r.op_kind == "gemm": + gemm_count += 1 + elif r.op_name == "dma_read": + dma_read_count += 1 + elif r.op_name == "dma_write": + dma_write_count += 1 + elif r.op_name == "ipcq_copy": + ipcq_copy_count += 1 + return { + "gemm_count": gemm_count, + "ipcq_copy_count": ipcq_copy_count, + "dma_read_count": dma_read_count, + "dma_write_count": dma_write_count, + } diff --git a/src/kernbench/benches/gqa_helpers/short_ctx/__init__.py b/src/kernbench/benches/gqa_helpers/short_ctx/__init__.py new file mode 100644 index 0000000..9731757 --- /dev/null +++ b/src/kernbench/benches/gqa_helpers/short_ctx/__init__.py @@ -0,0 +1 @@ +"""GQA short-context kernels.""" diff --git a/src/kernbench/benches/_gqa_attention_decode_short.py b/src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_decode_short.py similarity index 100% rename from src/kernbench/benches/_gqa_attention_decode_short.py rename to src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_decode_short.py diff --git a/src/kernbench/benches/_gqa_attention_prefill_short.py b/src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_prefill_short.py similarity index 100% rename from src/kernbench/benches/_gqa_attention_prefill_short.py rename to src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_prefill_short.py diff --git a/src/kernbench/benches/milestone_gqa_headline.py b/src/kernbench/benches/milestone_gqa_headline.py deleted file mode 100644 index d213838..0000000 --- a/src/kernbench/benches/milestone_gqa_headline.py +++ /dev/null @@ -1,231 +0,0 @@ -"""milestone-gqa-headline bench: real GQA prefill — single-KV-group target. - -Drives the LLaMA-3.1-70B single-KV-head-group prefill panel through -``_gqa_attention_prefill_long`` (C=8 snake Ring KV + intra-CUBE PE-SP, -all 64 ranks), writing ``op_log_summary`` into ``sweep.json``. -Independent from the existing ``milestone-gqa-llama70b`` validation-scale -bench (which stays on the legacy baseline kernels). - -Decode panels live in ``milestone_gqa_decode_long_ctx_4cases.py`` (the -4-cases long-context decode comparative study). - -Restrictions: - - Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP - headline deferred per ADR-0060 §B-item-1) - - T_q = S_kv = 1024 (scratch-limited; 32K headline awaits Q-axis - kernel tiling) - - No figure renderers (defer to a separate cycle) - -Panels: - single_kv_group_prefill_gqa_c8_p8: prefill C=8 snake Ring KV + - intra-CUBE PE-SP, T_q=S_kv=1K, - d_head=128 — the LLaMA-3.1-70B - single-KV-group prefill target. - -Gated by ``GQA_HEADLINE_RUN=1``. -""" -from __future__ import annotations - -import json -import os -from pathlib import Path - -from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel -from kernbench.benches.registry import bench -from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config -from kernbench.ccl.sfr_config import configure_sfr_intercube_ring -from kernbench.policy.placement.dp import DPPolicy - -_OUTPUT_DIR = Path(__file__).resolve().parent / "1H_milestone_output" / "gqa_headline" -_SWEEP_JSON = _OUTPUT_DIR / "sweep.json" - -# ── Panel configs ──────────────────────────────────────────────────── - -_DTYPE = "f16" -_D_HEAD = 64 -_T_Q_PREFILL = 4 -_S_KV_PREFILL = 16 - -_PANELS = ( - "single_kv_group_prefill_gqa_c8_p8", -) - -# Each entry: (kind, panel-specific params) -_PANEL_DISPATCH: dict[str, tuple[str, dict]] = { - "single_kv_group_prefill_gqa_c8_p8": ("prefill", { - "C": 8, "P": 8, - # T_q = S_kv = 1024 (one-shot long-context prefill, scratch-limited). - # The bootstrap section of the prefill kernel leaves K_t/V_t/scores/ - # exp_scores persistent (outside tl.scratch_scope), inflating the - # baseline. At T_q=2048 the peak hits ~1.05 MB vs 1.0 MB budget; at - # 1024 it fits comfortably. The true LLaMA 32K headline awaits a - # future increment to (a) add Q-axis tiling and (b) move the - # bootstrap into scratch discipline. - "T_q": 1_024, "S_kv": 1_024, - "d_head": 128, - }), -} - - -def _ccl_cfg(): - return resolve_algorithm_config( - load_ccl_config(), name="lrab_hierarchical_allreduce", - ) - - -# ── Per-kind launch helpers ────────────────────────────────────────── - - -def _run_prefill_panel( - ctx, *, panel: str, C: int, S_kv: int, - P: int = 1, - T_q: int = _T_Q_PREFILL, - d_head: int = _D_HEAD, -) -> None: - if C > 1: - mesh_w = int(ctx.spec["sip"]["cube_mesh"]["w"]) - if C <= mesh_w: - configure_sfr_intercube_ring( - ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C, - ) - else: - if C % mesh_w != 0: - raise ValueError( - f"C={C} > mesh_w={mesh_w} requires C divisible by mesh_w" - ) - configure_sfr_intercube_ring( - ctx.engine, ctx.spec, _ccl_cfg(), - submesh_shape=(C // mesh_w, mesh_w), - ) - # Q and O switch to pe="row_wise" when intra-CUBE PE-SP is active - # (ADR-0060 §5.5 + §B-item-3: disjoint query-row split across PEs). - q_pe = "row_wise" if P > 1 else "replicate" - o_pe = "row_wise" if P > 1 else "replicate" - dp_q = DPPolicy(cube="replicate", pe=q_pe, - num_cubes=C, num_pes=P) - dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate", - pe="replicate", num_cubes=C, num_pes=P) - dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate", - pe=o_pe, num_cubes=C, num_pes=P) - q = ctx.zeros((T_q, d_head), - dtype=_DTYPE, dp=dp_q, name=f"{panel}_q") - k = ctx.zeros((S_kv, d_head), - dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k") - v = ctx.zeros((S_kv, d_head), - dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v") - o = ctx.empty((T_q * C, d_head), - dtype=_DTYPE, dp=dp_o, name=f"{panel}_o") - ctx.launch( - panel, gqa_attention_prefill_long_kernel, - q, k, v, o, - T_q, S_kv, d_head, C, P, - _auto_dim_remap=False, - ) - - -def _make_bench_fn(panel: str): - kind, params = _PANEL_DISPATCH[panel] - assert kind == "prefill", ( - f"milestone-gqa-headline only registers prefill panels; got {kind!r}" - ) - - def _bench_fn(ctx): - _run_prefill_panel(ctx, panel=panel, **params) - - return _bench_fn - - -# ── Op-log summary ────────────────────────────────────────────────── - - -def _summarize_op_log(op_log) -> dict[str, int]: - """Per-panel op_log counts (gemm, ipcq_copy, dma_read, dma_write).""" - gemm_count = 0 - ipcq_copy_count = 0 - dma_read_count = 0 - dma_write_count = 0 - for r in op_log: - if r.op_kind == "gemm": - gemm_count += 1 - elif r.op_name == "dma_read": - dma_read_count += 1 - elif r.op_name == "dma_write": - dma_write_count += 1 - elif r.op_name == "ipcq_copy": - ipcq_copy_count += 1 - return { - "gemm_count": gemm_count, - "ipcq_copy_count": ipcq_copy_count, - "dma_read_count": dma_read_count, - "dma_write_count": dma_write_count, - } - - -def _run_panel(panel: str, topology: str) -> dict: - """Run one panel in a fresh engine; return its row 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"milestone-gqa-headline panel {panel!r} failed: " - f"{result.completion}" - ) - kind, params = _PANEL_DISPATCH[panel] - return { - "panel": panel, - "kind": kind, - **params, - "op_log_summary": _summarize_op_log(result.engine.op_log), - } - - -# ── Bench entry ────────────────────────────────────────────────────── - - -@bench( - name="milestone-gqa-headline", - description="Headline GQA prefill milestone — LLaMA-3.1-70B single-KV-group target (C=8, P=8).", -) -def run(torch) -> None: - """Drive the headline prefill panel; write sweep.json. - - Gated by GQA_HEADLINE_RUN=1. - """ - _OUTPUT_DIR.mkdir(parents=True, exist_ok=True) - if not os.environ.get("GQA_HEADLINE_RUN"): - raise RuntimeError( - "milestone-gqa-headline needs GQA_HEADLINE_RUN=1." - ) - - topology = os.environ.get("GQA_HEADLINE_TOPOLOGY", "topology.yaml") - rows = [_run_panel(panel, topology) for panel in _PANELS] - sweep = { - "version": 1, - "panels": list(_PANELS), - "config": { - "T_q_prefill": _T_Q_PREFILL, - "S_kv_prefill": _S_KV_PREFILL, - "d_head": _D_HEAD, - }, - "rows": rows, - } - _SWEEP_JSON.write_text(json.dumps(sweep, indent=2)) - print(f" milestone-gqa-headline: {len(rows)} rows -> {_SWEEP_JSON}") - - # Sentinel tensor (ADR-0045 D4 / ADR-0054 D2 carve-out). - torch.zeros( - (1, 1), dtype="f16", - dp=DPPolicy(cube="row_wise", pe="replicate", num_cubes=1, num_pes=1), - name="milestone_gqa_headline_sentinel", - ) diff --git a/src/kernbench/benches/registry.py b/src/kernbench/benches/registry.py index 58d8af8..02e3c09 100644 --- a/src/kernbench/benches/registry.py +++ b/src/kernbench/benches/registry.py @@ -99,7 +99,11 @@ def _audit_modules(imported: list[str], registered: set[str]) -> None: def _eager_import_and_audit(pkg_path: list[str], pkg_name: str) -> None: imported: list[str] = [] for m in iter_modules(pkg_path): - if m.name == "registry" or m.name.startswith("_"): + # Skip helper-name modules, the registry itself, and any + # subpackages (audit policy applies to ``.py`` bench files + # only — subpackages are explicitly imported by the benches + # that need them). + if m.ispkg or m.name == "registry" or m.name.startswith("_"): continue mod = import_module(f"{pkg_name}.{m.name}") imported.append(mod.__name__) diff --git a/tests/attention/test_gqa_decode_opt2.py b/tests/attention/test_gqa_decode_opt2.py index ecfbc2d..b7efee4 100644 --- a/tests/attention/test_gqa_decode_opt2.py +++ b/tests/attention/test_gqa_decode_opt2.py @@ -151,7 +151,7 @@ def test_k_before_v_in_opt2_plan(): def test_opt2_bench_completes_oplog_mode(): from pathlib import Path - from kernbench.benches._gqa_attention_decode_opt2 import ( # noqa: F401 + from kernbench.benches.gqa_helpers.shared._gqa_attention_decode_opt2 import ( # noqa: F401 gqa_attention_decode_opt2_kernel, ) from kernbench.policy.placement.dp import DPPolicy diff --git a/tests/attention/test_gqa_prefill.py b/tests/attention/test_gqa_prefill.py deleted file mode 100644 index 090b479..0000000 --- a/tests/attention/test_gqa_prefill.py +++ /dev/null @@ -1,118 +0,0 @@ -"""Phase 1 spec test for P6a GQA prefill kernel (head-parallel, C=1 baseline). - -P6a introduces ``_gqa_attention_prefill_long.py`` with the head-parallel structure (one -Q head per CUBE, per-CUBE distributed output, no reduce). C=1 is the -degenerate case — no Ring KV, no IPCQ traffic. Validates kernel -structure and T_q > 1 attention. - -P6b (deferred) adds the Ring KV rotation for C > 1, which needs either -a new SFR install function (intra_* + wrapped E/W at CUBE level) or a -topology-specific config — separate design call. - -The prefill kernel differs from decode (P1a/P2a/P2b) in three ways -(ADR-0060 §5.5 / TL;DR): - 1. Q has T_q > 1 rows (not just decode's single timestep). - 2. Head-parallel placement: each CUBE owns ONE Q head — no M-fold. - 3. Each CUBE writes its own head's output — NO reduce. - -Phase 1 (this commit): tests only — production code lands in Phase 2. -""" -from __future__ import annotations - -from pathlib import Path - -from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel # noqa: F401 -from kernbench.policy.placement.dp import DPPolicy -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 - -TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" - -D_HEAD = 64 -DTYPE = "f16" - - -def _engine_factory(t, d): - return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) - - -def _run_prefill(*, T_q: int, S_kv: int, C: int = 1): - """C=1 head-parallel prefill: single CUBE owns the one head + full KV.""" - topo = resolve_topology(str(TOPOLOGY_DEFAULT)) - - def _bench_fn(ctx): - dp = DPPolicy(cube="replicate", pe="replicate", - num_cubes=C, num_pes=1) - # Q: (T_q, d_head) — one head per CUBE (head-parallel; for C=1 - # only one head total). 2D layout matches what the kernel loads. - # P6b will use a 3D (h_q, T_q, d_head) Q with cube_row_wise - # sharding so each CUBE owns its head. - q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp, - name=f"q_t{T_q}_c{C}") - # K, V: full local for C=1 (no ring). Kernel loads K as - # (d_head, S_kv) via byte-conserving reshape. - k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp, - name=f"k_t{T_q}_c{C}") - v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp, - name=f"v_t{T_q}_c{C}") - # O: (T_q, d_head) — per-CUBE distributed output. - o = ctx.empty((T_q, D_HEAD), dtype=DTYPE, dp=dp, - name=f"o_t{T_q}_c{C}") - ctx.launch( - f"gqa_prefill_p6a_t{T_q}_s{S_kv}_c{C}", - gqa_attention_prefill_long_kernel, - q, k, v, o, - T_q, S_kv, D_HEAD, C, - _auto_dim_remap=False, - ) - - return run_bench( - topology=topo, bench_fn=_bench_fn, - device=resolve_device(None), - engine_factory=_engine_factory, - ) - - -def _count(op_log, name: str) -> int: - return sum(1 for r in op_log if r.op_name == name) - - -def test_prefill_c_one_t_q_one_completes(): - """C=1, T_q=1: smallest workload (decode-like).""" - result = _run_prefill(T_q=1, S_kv=16, C=1) - assert result.completion.ok, ( - f"prefill C=1 T_q=1 failed: {result.completion}" - ) - - -def test_prefill_c_one_t_q_four_completes(): - """C=1, T_q=4: real prefill (Q has multiple rows) — distinguishes - prefill from decode (T_q=1).""" - result = _run_prefill(T_q=4, S_kv=16, C=1) - assert result.completion.ok, ( - f"prefill C=1 T_q=4 failed: {result.completion}" - ) - - -def test_prefill_c_one_no_ipcq_traffic(): - """C=1: no ring step, no IPCQ traffic.""" - result = _run_prefill(T_q=4, S_kv=16, C=1) - assert result.completion.ok - n_copy = _count(result.engine.op_log, "ipcq_copy") - assert n_copy == 0, ( - f"C=1 must have no IPCQ traffic (no ring); got {n_copy}" - ) - - -def test_prefill_c_one_one_dma_write(): - """C=1, one head: exactly one dma_write (per-CUBE distributed output; - no reduce). For C > 1 in P6b this becomes dma_write_count == C.""" - result = _run_prefill(T_q=4, S_kv=16, C=1) - assert result.completion.ok - n_writes = _count(result.engine.op_log, "dma_write") - assert n_writes == 1, ( - f"C=1 prefill: expected 1 dma_write (one head per cube); " - f"got {n_writes}" - ) diff --git a/tests/attention/test_gqa_prefill_long_c8_snake.py b/tests/attention/test_gqa_prefill_long_c8_snake.py deleted file mode 100644 index 58ca62e..0000000 --- a/tests/attention/test_gqa_prefill_long_c8_snake.py +++ /dev/null @@ -1,154 +0,0 @@ -"""Tests for prefill_long Ring KV at C=8 over a 2×4 snake sub-mesh. - -ADR-0060 §5.5 design target for the single-KV-group LLaMA-3.1-70B -configuration: ``C = G = 8`` (one Q head per CUBE), Ring KV rotates -the 8 KV slices around all 8 CUBEs. - -Increment 1 wired ``configure_sfr_intercube_ring(submesh_shape=(2, 4))`` -to map the kernel's logical E/W to a snake/serpentine 1-hop path -through the top 2×4 sub-mesh of the 4×4 SIP CUBE mesh. The kernel -itself uses only logical ``dir="W"`` / ``dir="E"`` — the snake is -transparent at kernel level. - -This file verifies the assembly works end-to-end at C=8: - T1 kernel completes (no scratch overflow, no deadlock) - T2 per-CUBE head-parallel output (8 dma_writes, one per CUBE) - T3 tile-granular ring traffic at C=8 follows the same - ``(C-1)·n_tiles·2·C`` formula as the existing tile-ring tests -""" -from __future__ import annotations - -from pathlib import Path - -from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel # noqa: F401 -from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config -from kernbench.ccl.sfr_config import configure_sfr_intercube_ring -from kernbench.policy.placement.dp import DPPolicy -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 - -TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" - -D_HEAD = 64 -DTYPE = "f16" - - -def _ccl_cfg(): - return resolve_algorithm_config( - load_ccl_config(), name="lrab_hierarchical_allreduce", - ) - - -def _engine_factory(t, d): - return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) - - -def _count(op_log, name: str) -> int: - return sum(1 for r in op_log if r.op_name == name) - - -def _run_prefill_c8_snake(*, T_q: int, S_kv: int): - """Head-parallel prefill at C=8 with snake-mapped Ring KV over the - top 2×4 sub-mesh of the 4×4 SIP CUBE mesh.""" - topo = resolve_topology(str(TOPOLOGY_DEFAULT)) - C = 8 - - def _bench_fn(ctx): - configure_sfr_intercube_ring( - ctx.engine, ctx.spec, _ccl_cfg(), - submesh_shape=(2, 4), - ) - dp_q = DPPolicy(cube="replicate", pe="replicate", - num_cubes=C, num_pes=1) - dp_kv = DPPolicy(cube="row_wise", pe="replicate", - num_cubes=C, num_pes=1) - dp_o = DPPolicy(cube="row_wise", pe="replicate", - num_cubes=C, num_pes=1) - q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp_q, - name=f"q_c8_snake_t{T_q}_s{S_kv}") - k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, - name=f"k_c8_snake_t{T_q}_s{S_kv}") - v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, - name=f"v_c8_snake_t{T_q}_s{S_kv}") - o = ctx.empty((T_q * C, D_HEAD), dtype=DTYPE, dp=dp_o, - name=f"o_c8_snake_t{T_q}_s{S_kv}") - ctx.launch( - f"gqa_prefill_long_c8_snake_t{T_q}_s{S_kv}", - gqa_attention_prefill_long_kernel, - q, k, v, o, - T_q, S_kv, D_HEAD, C, - _auto_dim_remap=False, - ) - - return run_bench( - topology=topo, bench_fn=_bench_fn, - device=resolve_device(None), - engine_factory=_engine_factory, - ) - - -# ── T1: kernel completes end-to-end at C=8 ─────────────────────────── - - -def test_prefill_long_c8_snake_completes(): - """ADR-0060 §5.5 design target: prefill Ring KV at C=G=8 over the - 2×4 snake sub-mesh. With zero inputs, output is trivially zero; - this test guards that the kernel reaches completion without - deadlock, scratch overflow, or routing failure. - - Configuration: T_q=4, S_kv=8192 → S_local=1024, n_tiles=1 - (TILE_S_KV=1024). Bounded scratch. - """ - result = _run_prefill_c8_snake(T_q=4, S_kv=8192) - assert result.completion.ok, ( - f"prefill at C=8 with snake ring must complete; " - f"got {result.completion}" - ) - - -# ── T2: 8 dma_writes (head-parallel, one head per CUBE) ────────────── - - -def test_prefill_long_c8_snake_distributed_output_count(): - """Head-parallel: each CUBE owns one Q head and writes its own - head's output (T_q, d_head) rows. No inter-CUBE reduce on the - output side — so ``dma_write_count == C == 8``. - - This is the C=8 analogue of - ``test_prefill_long_tile_ring_dma_write_count`` at C=4. - """ - result = _run_prefill_c8_snake(T_q=4, S_kv=8192) - assert result.completion.ok - n_writes = _count(result.engine.op_log, "dma_write") - assert n_writes == 8, ( - f"head-parallel C=8: expected 8 dma_writes (one per CUBE); " - f"got {n_writes}" - ) - - -# ── T3: tile-granular ring ipcq_copy count scales as (C-1)·n_tiles·2·C - - -def test_prefill_long_c8_snake_tile_ipcq_count(): - """ADR-0060 §5.5.1 tile-granular ring: per-CUBE send count is - ``2·n_tiles·(C-1)`` (K + V per tile per ring step). Aggregated - across C CUBEs, total ipcq_copy = ``(C-1)·n_tiles·2·C``. - - Configuration: T_q=4, S_kv=8192, C=8 → S_local=1024, n_tiles=1 - (TILE_S_KV=1024). Expected total = ``7·1·2·8 = 112``. - - Verifies that the snake-mapped 1-hop physical links carry the - same logical ring traffic as the existing C=4 1D-row ring tests. - """ - C = 8 - n_tiles = 1 # S_local=1024 / TILE_S_KV=1024 - result = _run_prefill_c8_snake(T_q=4, S_kv=8192) - assert result.completion.ok - n_copy = _count(result.engine.op_log, "ipcq_copy") - expected = (C - 1) * n_tiles * 2 * C - assert n_copy == expected, ( - f"tile-granular ring at C=8: expected {expected} ipcq_copy " - f"((C-1)·n_tiles·2·C = {C - 1}·{n_tiles}·2·{C}); got {n_copy}" - ) diff --git a/tests/attention/test_gqa_prefill_long_pe_sp.py b/tests/attention/test_gqa_prefill_long_pe_sp.py deleted file mode 100644 index 8c8b775..0000000 --- a/tests/attention/test_gqa_prefill_long_pe_sp.py +++ /dev/null @@ -1,233 +0,0 @@ -"""Tests for intra-CUBE PE-SP in prefill_long (query-axis split). - -ADR-0060 §5.5 last bullet + §B-item-3: split the head's query rows -``[T_q, d_head]`` across the ``P`` PEs of each CUBE so all P PEs work -in parallel. Output rows are disjoint across PEs ⇒ no intra-CUBE -reduce needed. - -Same-lane SFR wiring: PE ``i`` in CUBE A has its own E/W ring link to -PE ``i`` in CUBE B (the snake's prev/next). All P PEs of a CUBE see -the same K, V (HBM-resident, ``pe="replicate"``) ⇒ P parallel rings -run in lockstep, each PE rotating its own K/V copies via its own IPCQ -channels. - -Activation contract: - - ``P == 1`` (default; omitted from launch args) → existing - PE-0-only behavior. Byte-for-byte unchanged. - - ``P > 1`` → all P PEs active; each handles ``T_q // P`` query - rows. Requires ``T_q % P == 0`` (degenerate T_q < P is rejected - — caller must use ``P=1`` for that workload, ADR-0060 §B-item-3). - -Phase 1: tests only — production code lands in Phase 2. -T1, T2, T3, T5 fail today (TypeError: kernel signature has no P). -T4 passes today as the backward-compat anchor. -""" -from __future__ import annotations - -from pathlib import Path - -import pytest - -from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel # noqa: F401 -from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config -from kernbench.ccl.sfr_config import configure_sfr_intercube_ring -from kernbench.policy.placement.dp import DPPolicy -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 - -TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" - -D_HEAD = 64 -DTYPE = "f16" - - -def _ccl_cfg(): - return resolve_algorithm_config( - load_ccl_config(), name="lrab_hierarchical_allreduce", - ) - - -def _engine_factory(t, d): - return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) - - -def _count(op_log, name: str) -> int: - return sum(1 for r in op_log if r.op_name == name) - - -def _run_prefill_pe_sp( - *, T_q: int, S_kv: int, C: int, P: int | None, - snake: bool, -): - """Drive a prefill_long launch with optional intra-CUBE PE-SP. - - ``P=None`` → omit P from the launch args (exercises the kernel's - default behaviour). - ``snake=True`` → install the 2×4 snake ring SFR (Increment 1); - else install the 1D-row ring at ``ring_size=C``. - """ - topo = resolve_topology(str(TOPOLOGY_DEFAULT)) - - def _bench_fn(ctx): - if snake: - configure_sfr_intercube_ring( - ctx.engine, ctx.spec, _ccl_cfg(), - submesh_shape=(2, 4), - ) - else: - configure_sfr_intercube_ring( - ctx.engine, ctx.spec, _ccl_cfg(), - ring_size=C, - ) - num_pes = P if (P is not None and P > 1) else 1 - # When PE-SP is active, Q and O are split row-wise across PEs; - # K, V remain replicated within a CUBE. - q_pe = "row_wise" if num_pes > 1 else "replicate" - o_pe = "row_wise" if num_pes > 1 else "replicate" - dp_q = DPPolicy(cube="replicate", pe=q_pe, - num_cubes=C, num_pes=num_pes) - dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate", - pe="replicate", num_cubes=C, num_pes=num_pes) - dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate", - pe=o_pe, num_cubes=C, num_pes=num_pes) - suffix = f"c{C}_p{P}_t{T_q}_s{S_kv}" - q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp_q, - name=f"q_pe_sp_{suffix}") - k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, - name=f"k_pe_sp_{suffix}") - v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, - name=f"v_pe_sp_{suffix}") - o = ctx.empty((T_q * C, D_HEAD), dtype=DTYPE, dp=dp_o, - name=f"o_pe_sp_{suffix}") - launch_args = [q, k, v, o, T_q, S_kv, D_HEAD, C] - if P is not None: - launch_args.append(P) - ctx.launch( - f"gqa_prefill_long_pe_sp_{suffix}", - gqa_attention_prefill_long_kernel, - *launch_args, - _auto_dim_remap=False, - ) - - return run_bench( - topology=topo, bench_fn=_bench_fn, - device=resolve_device(None), - engine_factory=_engine_factory, - ) - - -# ── T1: C=8 P=8 completes end-to-end ───────────────────────────────── - - -def test_prefill_c8_p8_pe_sp_completes(): - """ADR-0060 §5.5 + §B-item-3 design target: 64 ranks active - (8 CUBEs × 8 PEs), query-axis split across PEs, snake-mapped - Ring KV. Guards no scratch overflow, no deadlock across the - P parallel same-lane rings. - """ - result = _run_prefill_pe_sp( - T_q=8, S_kv=8192, C=8, P=8, snake=True, - ) - assert result.completion.ok, ( - f"prefill at C=8, P=8 (PE-SP) must complete; " - f"got {result.completion}" - ) - - -# ── T2: 64 dma_writes (one per PE, disjoint query rows) ────────────── - - -def test_prefill_c8_p8_pe_sp_64_dma_writes(): - """With query-axis split, each PE owns ``T_q/P = 1`` query row - and stores its own ``(1, d_head)`` slice. Across all 64 ranks - (8 CUBEs × 8 PEs), ``dma_write_count == 64``. - - This is the structural signal that PE-SP is actually wired: - PE-0-only would give 8 dma_writes (one per CUBE). - """ - result = _run_prefill_pe_sp( - T_q=8, S_kv=8192, C=8, P=8, snake=True, - ) - assert result.completion.ok - n_writes = _count(result.engine.op_log, "dma_write") - assert n_writes == 64, ( - f"PE-SP at C=8, P=8: expected 64 dma_writes " - f"(8 cubes × 8 PEs, each writes its T_q/P=1 query row); " - f"got {n_writes}" - ) - - -# ── T3: P parallel rings — ipcq_copy scales by P ───────────────────── - - -def test_prefill_c8_p8_pe_sp_ring_ipcq_count(): - """Per-PE same-lane rings: each PE runs its own ring traffic via - its own IPCQ channels. Total inter-CUBE ipcq_copy: - ``(C-1) · n_tiles · 2 · C · P`` (= existing C=8 formula × P). - - Configuration: T_q=8, S_kv=8192, C=8 → S_local=1024, n_tiles=1 - (TILE_S_KV=1024). Expected = ``7·1·2·8·8`` = **896**. - """ - C = 8 - P = 8 - n_tiles = 1 - result = _run_prefill_pe_sp( - T_q=8, S_kv=8192, C=C, P=P, snake=True, - ) - assert result.completion.ok - n_copy = _count(result.engine.op_log, "ipcq_copy") - expected = (C - 1) * n_tiles * 2 * C * P - assert n_copy == expected, ( - f"PE-SP parallel rings at C=8, P=8: expected {expected} " - f"ipcq_copy ((C-1)·n_tiles·2·C·P = " - f"{C - 1}·{n_tiles}·2·{C}·{P}); got {n_copy}" - ) - - -# ── T4: backward-compat — P omitted → existing PE-0-only behaviour ─── - - -def test_prefill_p_default_1_backward_compat(): - """When ``P`` is omitted from the launch args, the kernel must - behave identically to today: only PE 0 of each CUBE participates; - one head per CUBE; one dma_write per CUBE. - - At C=4 this gives 4 dma_writes (matches the existing - ``test_prefill_long_tile_ring_dma_write_count``). Phase 2 must - not regress this path. - - Passes today AND after Phase 2. - """ - C = 4 - result = _run_prefill_pe_sp( - T_q=4, S_kv=8192, C=C, P=None, snake=False, - ) - assert result.completion.ok, ( - f"prefill at C=4 with default P (PE-0-only) must complete; " - f"got {result.completion}" - ) - n_writes = _count(result.engine.op_log, "dma_write") - assert n_writes == C, ( - f"default P=1 (PE-0-only): expected {C} dma_writes " - f"(one per CUBE); got {n_writes}" - ) - - -# ── T5: validation — T_q must be divisible by P when P > 1 ─────────── - - -def test_prefill_pe_sp_rejects_non_divisible_t_q(): - """ADR-0060 §B-item-3 fallback (KV-block split + intra-CUBE - reduce for T_q < P) is deferred. Callers must request ``P=1`` for - workloads where T_q < P or T_q % P != 0. - - C=4, P=8, T_q=4 violates ``T_q % P == 0`` (and also T_q < P). - The kernel must raise ValueError with a clear error message - rather than silently producing wrong results. - """ - with pytest.raises((ValueError, AssertionError), match=r"T_q"): - _run_prefill_pe_sp( - T_q=4, S_kv=8192, C=4, P=8, snake=False, - ) diff --git a/tests/attention/test_gqa_prefill_long_tile_ring.py b/tests/attention/test_gqa_prefill_long_tile_ring.py deleted file mode 100644 index 93833a6..0000000 --- a/tests/attention/test_gqa_prefill_long_tile_ring.py +++ /dev/null @@ -1,162 +0,0 @@ -"""Phase 1 spec test for P3c: tile-granular Ring KV in prefill_long -(ADR-0060 §5.5 amendment + ADR-0063 §A.2). - -Today's prefill_long ring sends and receives full ``(d_head, S_local)`` -KV slices per step. Step 0's local-attention intermediates also live -outside any ``scratch_scope`` (they're loaded as full-slice ``Kc``, -``Vc`` and feed ``scores``, ``exp_scores`` as persistent allocations). -At larger ``S_local`` both the step-0 leak and the ring step's -in-scope intermediates grow linearly with ``S_local``, and at -``S_local = 32K`` (S_kv=128K, C=4, T_q=4) the peak overflows the -1 MiB pool. - -P3c converts the ring to **tile-granular**: a nested loop -``for k in range(C): for t in range(n_tiles): ...`` where each -iteration sends/recvs one ``(d_head, TILE_S_KV)`` tile (and its V -counterpart). The persistent state shrinks to ``(m, ℓ, O)`` only -(~1 KB); per-tile in-scope scratch is bounded by -``TILE_S_KV`` regardless of ``S_local``. Ceiling lifted. - -Trade-off: the per-CUBE send count grows from ``2·(C-1)`` to -``2·n_tiles·(C-1)``. At ``n_tiles=1`` (small ``S_local``) the count is -unchanged, so the existing ``test_prefill_ring_c_*`` tests at -``S_kv ∈ {16, 32}`` still pass. - -Phase 1 (this commit): tests only — production code lands in Phase 2. -T1 fails today with a ``TLContext scratch overflow``; T2 fails today -with the slice-granular ipcq_copy count; T5 passes today and is a -regression guard for the per-CUBE output write count. -""" -from __future__ import annotations - -from pathlib import Path - -from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel # noqa: F401 -from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config -from kernbench.ccl.sfr_config import configure_sfr_intercube_ring -from kernbench.policy.placement.dp import DPPolicy -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 - -TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" - -D_HEAD = 64 -DTYPE = "f16" - - -def _ccl_cfg(): - return resolve_algorithm_config( - load_ccl_config(), name="lrab_hierarchical_allreduce", - ) - - -def _engine_factory(t, d): - return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) - - -def _count(op_log, name: str) -> int: - return sum(1 for r in op_log if r.op_name == name) - - -def _run_prefill_ring(*, T_q: int, S_kv: int, C: int): - """Head-parallel prefill with Ring KV across C CUBEs.""" - topo = resolve_topology(str(TOPOLOGY_DEFAULT)) - - def _bench_fn(ctx): - configure_sfr_intercube_ring( - ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C, - ) - dp_q = DPPolicy(cube="replicate", pe="replicate", - num_cubes=C, num_pes=1) - dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate", - pe="replicate", num_cubes=C, num_pes=1) - dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate", - pe="replicate", num_cubes=C, num_pes=1) - q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp_q, - name=f"q_tlr_t{T_q}_c{C}_s{S_kv}") - k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, - name=f"k_tlr_t{T_q}_c{C}_s{S_kv}") - v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, - name=f"v_tlr_t{T_q}_c{C}_s{S_kv}") - o = ctx.empty((T_q * C, D_HEAD), dtype=DTYPE, dp=dp_o, - name=f"o_tlr_t{T_q}_c{C}_s{S_kv}") - ctx.launch( - f"gqa_prefill_long_tile_ring_t{T_q}_c{C}_s{S_kv}", - gqa_attention_prefill_long_kernel, - q, k, v, o, - T_q, S_kv, D_HEAD, C, - _auto_dim_remap=False, - ) - - return run_bench( - topology=topo, bench_fn=_bench_fn, - device=resolve_device(None), - engine_factory=_engine_factory, - ) - - -# ── T1: 128K ceiling lift ──────────────────────────────────────────── - - -def test_prefill_long_context_128k_completes(): - """ADR-0063 §A.2 headline ceiling lift. At S_kv=128K with C=4, - S_local=32K. Today: step-0 score-stack (~768 KB persistent) + ring - scope (~768 KB) → 1.5 MB peak → TLContext scratch overflow. After - P3c the persistent state is just ``(m, ℓ, O)`` (≈ 1 KB) and the - per-tile in-scope scratch is bounded by TILE_S_KV. - """ - result = _run_prefill_ring(T_q=4, S_kv=131_072, C=4) - assert result.completion.ok, ( - f"prefill_long at S_kv=128K must complete after tile-granular " - f"ring lands; got {result.completion}" - ) - - -# ── T2: tile-granular ipcq_copy count ──────────────────────────────── - - -def test_prefill_long_tile_granular_ipcq_count(): - """ADR-0060 §5.5 amendment: with tile-granular sends, the per-CUBE - send count grows from ``2·(C-1)`` to ``2·n_tiles·(C-1)``. Aggregated - across all C CUBEs the total ipcq_copy becomes - ``2·n_tiles·(C-1)·C``. - - Config: T_q=4, S_kv=4096, C=2 → S_local=2048, n_tiles=2 (with - TILE_S_KV=1024). Today: slice-granular total = - ``(C-1)·2·C = 4``. After P3c: tile-granular total = - ``(C-1)·n_tiles·2·C = 8``. - """ - C = 2 - n_tiles = 2 # S_local=2048 / TILE_S_KV=1024 - result = _run_prefill_ring(T_q=4, S_kv=4096, C=C) - assert result.completion.ok, ( - f"prefill_long multi-tile ring must complete; got {result.completion}" - ) - n_copy = _count(result.engine.op_log, "ipcq_copy") - expected = (C - 1) * n_tiles * 2 * C - assert n_copy == expected, ( - f"tile-granular ring: expected {expected} ipcq_copy " - f"((C-1)·n_tiles·2·C = {C - 1}·{n_tiles}·2·{C}); got {n_copy}" - ) - - -# ── T5: per-CUBE distributed output unchanged (regression guard) ───── - - -def test_prefill_long_tile_ring_dma_write_count(): - """ADR-0060 §5.5: per-CUBE distributed output must hold under the - tile-granular ring rewrite. Each CUBE still writes its own head's - output (no inter-CUBE reduce); dma_write_count == C. - - Today passes; must continue to pass after P3c. - """ - C = 4 - result = _run_prefill_ring(T_q=4, S_kv=4096, C=C) - assert result.completion.ok - n_writes = _count(result.engine.op_log, "dma_write") - assert n_writes == C, ( - f"per-CUBE distributed output: expected {C} dma_writes (one per " - f"CUBE); got {n_writes}" - ) diff --git a/tests/attention/test_gqa_prefill_ring.py b/tests/attention/test_gqa_prefill_ring.py deleted file mode 100644 index 1ce602c..0000000 --- a/tests/attention/test_gqa_prefill_ring.py +++ /dev/null @@ -1,161 +0,0 @@ -"""Phase 1 spec test for P6b GQA prefill Ring KV (head-parallel, C>1). - -P6b adds the Ring KV rotation (ADR-0060 §5.5) to the prefill kernel. -Each CUBE owns one Q head + its KV slice; over C ring steps the KV -blocks rotate around the C CUBEs so every CUBE sees every block. The -running ``(m, ℓ, O)`` is folded inside the loop. No reduce — each CUBE -writes its own head's output. - -Requires a new SFR install ``configure_sfr_intercube_ring(ring_size=C)`` -that wires: - - ``intra_*`` : 2×4 PE grid within a CUBE (same as multisip) - - ``E/W`` : 1D ring of CUBEs 0..ring_size-1 WITH WRAP - (symmetric to ``configure_sfr_intracube_pe_ring``, - applied at CUBE level) - - ``global_*`` : SIP-level (same as multisip) - -The kernel ring body: - for step in range(1, C): - tl.send(dir="W", src=Kc) - tl.send(dir="W", src=Vc) - Kc = tl.recv(dir="E", ...) - Vc = tl.recv(dir="E", ...) - ... local partial + online-softmax merge into (m, ℓ, O) ... - -Per CUBE per step: 2 sends (K, V) → 2 ``ipcq_copy``. Across all CUBEs: -``(C-1) * 2 * C`` total ipcq_copy. - -Restriction in P6b first cut: ``C ∈ {1, 2, 4}`` (single row of the 4×4 -cube mesh). C=8 ring would span rows — follow-on. - -Phase 1 (this commit): tests only — production code lands in Phase 2. -""" -from __future__ import annotations - -from pathlib import Path - -from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel # noqa: F401 -from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config -from kernbench.ccl.sfr_config import configure_sfr_intercube_ring # noqa: F401 (Phase 2) -from kernbench.policy.placement.dp import DPPolicy -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 - -TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" - -D_HEAD = 64 -DTYPE = "f16" - - -def _ccl_cfg(): - return resolve_algorithm_config( - load_ccl_config(), name="lrab_hierarchical_allreduce", - ) - - -def _engine_factory(t, d): - return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) - - -def _run_prefill_ring(*, T_q: int, S_kv: int, C: int): - """Head-parallel prefill with Ring KV across C CUBEs.""" - topo = resolve_topology(str(TOPOLOGY_DEFAULT)) - - def _bench_fn(ctx): - # P6b: new SFR install with cube-level ring wrap. - configure_sfr_intercube_ring( - ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C, - ) - # Q replicated on every CUBE (zeros for testing; per-CUBE head - # indexing is implicit). KV sequence-sharded by CUBE. O - # distributed — each CUBE writes its slice of (T_q*C, d_head). - dp_q = DPPolicy(cube="replicate", pe="replicate", - num_cubes=C, num_pes=1) - dp_kv = DPPolicy(cube="row_wise", pe="replicate", - num_cubes=C, num_pes=1) - dp_o = DPPolicy(cube="row_wise", pe="replicate", - num_cubes=C, num_pes=1) - q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp_q, - name=f"q_t{T_q}_c{C}_ring") - k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, - name=f"k_t{T_q}_c{C}_ring") - v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, - name=f"v_t{T_q}_c{C}_ring") - # O: (T_q * C, d_head), each CUBE writes (T_q, d_head) at its - # slice. dma_write_count = C (per-CUBE distributed output). - o = ctx.empty((T_q * C, D_HEAD), dtype=DTYPE, dp=dp_o, - name=f"o_t{T_q}_c{C}_ring") - ctx.launch( - f"gqa_prefill_ring_t{T_q}_s{S_kv}_c{C}", - gqa_attention_prefill_long_kernel, - q, k, v, o, - T_q, S_kv, D_HEAD, C, - _auto_dim_remap=False, - ) - - return run_bench( - topology=topo, bench_fn=_bench_fn, - device=resolve_device(None), - engine_factory=_engine_factory, - ) - - -def _count(op_log, name: str) -> int: - return sum(1 for r in op_log if r.op_name == name) - - -# ── C=2 Ring KV ──────────────────────────────────────────────────────── - - -def test_prefill_ring_c_two_completes(): - """C=2: 1 ring step rotates KV between the 2 CUBEs.""" - result = _run_prefill_ring(T_q=4, S_kv=16, C=2) - assert result.completion.ok, ( - f"prefill ring C=2 failed: {result.completion}" - ) - - -def test_prefill_ring_c_two_distributed_output(): - """C=2: per-CUBE distributed output, no reduce → dma_write_count == 2.""" - result = _run_prefill_ring(T_q=4, S_kv=16, C=2) - assert result.completion.ok - n_writes = _count(result.engine.op_log, "dma_write") - assert n_writes == 2, ( - f"C=2 prefill: expected 2 dma_write (one per CUBE); got {n_writes}" - ) - - -def test_prefill_ring_c_two_ipcq_count(): - """C=2: 1 ring step × 2 handles (K, V) × 2 CUBEs = 4 ipcq_copy.""" - result = _run_prefill_ring(T_q=4, S_kv=16, C=2) - assert result.completion.ok - n_copy = _count(result.engine.op_log, "ipcq_copy") - expected = (2 - 1) * 2 * 2 - assert n_copy == expected, ( - f"C=2 ring: expected {expected} ipcq_copy " - f"((C-1)·2·C = 1·2·2); got {n_copy}" - ) - - -# ── C=4 Ring KV (combined assertions) ───────────────────────────────── - - -def test_prefill_ring_c_four_combined(): - """C=4: 3 ring steps rotate KV around 4 CUBEs. - Expected: completes; 4 dma_writes; (4-1)·2·4 = 24 ipcq_copy.""" - result = _run_prefill_ring(T_q=4, S_kv=32, C=4) - assert result.completion.ok, ( - f"prefill ring C=4 failed: {result.completion}" - ) - n_writes = _count(result.engine.op_log, "dma_write") - assert n_writes == 4, ( - f"C=4 prefill: expected 4 dma_write (one per CUBE); got {n_writes}" - ) - n_copy = _count(result.engine.op_log, "ipcq_copy") - expected = (4 - 1) * 2 * 4 - assert n_copy == expected, ( - f"C=4 ring: expected {expected} ipcq_copy " - f"((C-1)·2·C = 3·2·4); got {n_copy}" - ) diff --git a/tests/attention/test_gqa_short_context.py b/tests/attention/test_gqa_short_context.py index b1191fd..bdc85a2 100644 --- a/tests/attention/test_gqa_short_context.py +++ b/tests/attention/test_gqa_short_context.py @@ -76,7 +76,7 @@ def _run_decode_short(*, h_q: int, h_kv: int, kv_per_cube: int, O: replicated; each group root writes the full byte-conserving (h_q·T_q, D_HEAD) result. """ - from kernbench.benches._gqa_attention_decode_short import gqa_attention_decode_short_kernel # Phase 2 + from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short import gqa_attention_decode_short_kernel # Phase 2 topo = resolve_topology(str(TOPOLOGY_DEFAULT)) @@ -196,7 +196,7 @@ def _run_prefill_short(*, h_kv: int, kv_per_cube: int, Layout: same head-stacked K/V scheme as decode short, with Q/O replicated and byte-conserving reshape inside the kernel. """ - from kernbench.benches._gqa_attention_prefill_short import gqa_attention_prefill_short_kernel # Phase 2 + from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short import gqa_attention_prefill_short_kernel # Phase 2 topo = resolve_topology(str(TOPOLOGY_DEFAULT)) 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 c921b84..65ce1f6 100644 --- a/tests/attention/test_milestone_gqa_decode_long_ctx_4cases.py +++ b/tests/attention/test_milestone_gqa_decode_long_ctx_4cases.py @@ -68,7 +68,7 @@ def _run_case4_smoke(*, S_kv: int): ``S_kv=128K`` runs come from ``kernbench run --bench milestone-gqa-decode-long-ctx-4cases``, not pytest. """ - from kernbench.benches.milestone_gqa_decode_long_ctx_4cases import ( + from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import ( _run_decode_panel_long_ctx_cube_sp_pe_sp, ) topo = resolve_topology(str(TOPOLOGY_DEFAULT)) @@ -99,13 +99,13 @@ def test_case4_panel_registered(): C = 8 (head-parallel CUBE Group) P = 8 (intra-CUBE PE-SP) T_q = 1 (decode: one new token per pass) - S_kv = 131_072 (LLaMA long-context decode target) + S_kv = 8_192 (LLaMA long-context decode target) d_head = 128, h_q = 8, h_kv = 1 The lrab sub_w=4 / sub_h=2 geometry is baked into the Case 4 wrapper kernel; it is not a panel parameter. """ - from kernbench.benches.milestone_gqa_decode_long_ctx_4cases import ( + from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import ( _PANEL_DISPATCH, _PANELS, ) @@ -120,7 +120,7 @@ def test_case4_panel_registered(): assert params.get("C") == 8 assert params.get("P") == 8 assert params.get("T_q") == 1 - assert params.get("S_kv") == 131_072 + assert params.get("S_kv") == 8_192 assert params.get("d_head") == 128 assert params.get("h_q") == 8 assert params.get("h_kv") == 1 @@ -197,7 +197,7 @@ def _run_case2_smoke(*, S_kv: int): slide-11 memory waste); for B=1 only one rank does the work; no inter-rank comm. """ - from kernbench.benches.milestone_gqa_decode_long_ctx_4cases import ( + from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import ( _run_decode_panel_long_ctx_cube_repl_pe_tp, ) topo = resolve_topology(str(TOPOLOGY_DEFAULT)) @@ -229,7 +229,7 @@ def test_case2_panel_registered(): For B=1 only one rank works (PEs 1-7 idle — slide-11 calls out this PE-TP waste). """ - from kernbench.benches.milestone_gqa_decode_long_ctx_4cases import ( + from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import ( _PANEL_DISPATCH, _PANELS, ) @@ -242,7 +242,7 @@ def test_case2_panel_registered(): assert params.get("C") == 8 assert params.get("P") == 8 assert params.get("T_q") == 1 - assert params.get("S_kv") == 131_072 + assert params.get("S_kv") == 8_192 assert params.get("d_head") == 128 assert params.get("h_q") == 8 assert params.get("h_kv") == 1 @@ -303,7 +303,7 @@ def _run_case3_smoke(*, S_kv: int): 8-way across PEs within each cube. Intra-CUBE 8-way reduce on (m, ℓ, O); no inter-CUBE comm (every cube ends with full answer). """ - from kernbench.benches.milestone_gqa_decode_long_ctx_4cases import ( + from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import ( _run_decode_panel_long_ctx_cube_repl_pe_sp, ) topo = resolve_topology(str(TOPOLOGY_DEFAULT)) @@ -333,7 +333,7 @@ def test_case3_panel_registered(): Case 3: Cube-Repl × PE-SP. K, V replicated per cube; PEs SP on S_kv. Intra-CUBE 8-way AR; no inter-CUBE comm. """ - from kernbench.benches.milestone_gqa_decode_long_ctx_4cases import ( + from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import ( _PANEL_DISPATCH, _PANELS, ) @@ -346,7 +346,7 @@ def test_case3_panel_registered(): assert params.get("C") == 8 assert params.get("P") == 8 assert params.get("T_q") == 1 - assert params.get("S_kv") == 131_072 + assert params.get("S_kv") == 8_192 assert params.get("d_head") == 128 assert params.get("h_q") == 8 assert params.get("h_kv") == 1 @@ -423,7 +423,7 @@ def _run_case1_smoke(*, S_kv: int): each cube has work; PEs 1-7 idle. Inter-CUBE 8-way reduce via the lrab-adapted center-root pattern (root cube 6); no intra-CUBE comm. """ - from kernbench.benches.milestone_gqa_decode_long_ctx_4cases import ( + from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import ( _run_decode_panel_long_ctx_cube_sp_pe_tp, ) topo = resolve_topology(str(TOPOLOGY_DEFAULT)) @@ -454,7 +454,7 @@ def test_case1_panel_registered(): batch. At B=1 only PE 0 of each cube works (PE-TP waste). Inter-CUBE lrab AR; no intra-CUBE comm. """ - from kernbench.benches.milestone_gqa_decode_long_ctx_4cases import ( + from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import ( _PANEL_DISPATCH, _PANELS, ) @@ -467,7 +467,7 @@ def test_case1_panel_registered(): assert params.get("C") == 8 assert params.get("P") == 8 assert params.get("T_q") == 1 - assert params.get("S_kv") == 131_072 + assert params.get("S_kv") == 8_192 assert params.get("d_head") == 128 assert params.get("h_q") == 8 assert params.get("h_kv") == 1 @@ -548,7 +548,7 @@ def test_panel_metrics_helpers_present_and_correct(): 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 ( + from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import ( _end_to_end_ns, _engine_occupancy_ns, ) @@ -579,7 +579,7 @@ def test_run_panel_returns_latency_and_engine_occupancy(monkeypatch): 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 + import kernbench.benches.gqa_helpers.long_ctx.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} diff --git a/tests/attention/test_milestone_gqa_headline.py b/tests/attention/test_milestone_gqa_headline.py deleted file mode 100644 index 5fb6b9a..0000000 --- a/tests/attention/test_milestone_gqa_headline.py +++ /dev/null @@ -1,89 +0,0 @@ -"""Tests for the milestone-gqa-headline bench. - -The bench wires the new ``_gqa_attention_prefill_long`` kernel into a -single-KV-group milestone panel (LLaMA-3.1-70B target, C=8 P=8). The -4 legacy C=1 / C=4 panels (single_user_*, multi_user_*) were removed -once pytest regression covered those configurations comprehensively -and the comparative decode work moved into the -``milestone-gqa-decode-4cases`` bench. - -Single SIP scope (default ``topology.yaml`` 4×4 cube mesh; 4-SIP -headline deferred per ADR-0060 §B-item-1). -""" -from __future__ import annotations - -import json - -import kernbench.benches.milestone_gqa_headline as bench_mod # noqa: F401 (Phase 2) -from kernbench.benches.registry import resolve -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 - -BENCH_NAME = "milestone-gqa-headline" - -PANELS = ( - "single_kv_group_prefill_gqa_c8_p8", -) - - -def _run_validation(): - topo = resolve_topology("topology.yaml") - return run_bench( - topology=topo, - bench_fn=resolve(BENCH_NAME).run, - device=resolve_device(None), - engine_factory=lambda t, d: GraphEngine( - getattr(t, "topology_obj", t), enable_data=True, - ), - ) - - -def _sweep_json(monkeypatch) -> dict: - monkeypatch.setenv("GQA_HEADLINE_RUN", "1") - out = bench_mod._OUTPUT_DIR / "sweep.json" - if not out.exists(): - result = _run_validation() - assert result.completion.ok, result.completion - assert out.exists(), f"missing {out}" - return json.loads(out.read_text()) - - -# ── Registration ─────────────────────────────────────────────────────── - - -def test_bench_registered(): - spec = resolve(BENCH_NAME) - assert spec.name == BENCH_NAME - assert callable(spec.run) - assert spec.description.strip(), "description must be non-empty" - - -# ── Validation run completes ────────────────────────────────────────── - - -def test_validation_run_completes_ok(monkeypatch): - monkeypatch.setenv("GQA_HEADLINE_RUN", "1") - result = _run_validation() - assert result.completion.ok, ( - f"headline validation run failed: {result.completion}" - ) - - -# ── sweep.json shape ────────────────────────────────────────────────── - - -def test_sweep_json_has_expected_panels(monkeypatch): - data = _sweep_json(monkeypatch) - assert set(data["panels"]) == set(PANELS), ( - f"panels mismatch: expected {set(PANELS)}, got {set(data['panels'])}" - ) - assert len(data["rows"]) == len(PANELS) - assert {r["panel"] for r in data["rows"]} == set(PANELS) - - -# Per-panel architectural assertions for the surviving single-KV-group -# panel live in tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py. -# Decode architectural assertions live in -# tests/attention/test_milestone_gqa_decode_4cases.py. diff --git a/tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py b/tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py deleted file mode 100644 index 9935254..0000000 --- a/tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py +++ /dev/null @@ -1,121 +0,0 @@ -"""Tests for the single-KV-group prefill panel in milestone_gqa_headline. - -The single-KV-group (LLaMA-3.1-70B target) prefill panel wires together -all Increment 1–4 work in a single bench panel: - - C=8 head-parallel + Ring KV via snake-mapped 2×4 sub-mesh (Inc 1, 3) - - P=8 intra-CUBE PE-SP (Inc 4) - - d_head=128 (LLaMA-3.1-70B), one-shot prefill T_q = S_kv = 1K - (scratch-limited; LLaMA 32K headline awaits Q-axis kernel tiling) - -This file verifies the bench-config wiring: - T1 the new ``single_kv_group_prefill_gqa_c8_p8`` panel is registered - with the expected dims in ``_PANELS`` + ``_PANEL_DISPATCH``. - T2 ``_run_prefill_panel`` accepts the new ``P``, ``T_q``, ``d_head`` - kwargs and drives the prefill kernel to completion at the - milestone-target ``(C, P) = (8, 8)``. Uses smaller T_q/S_kv to - keep test time bounded — full 32K runs come from - ``kernbench run milestone-gqa-headline``. - T3 Existing prefill panels still work when ``_run_prefill_panel`` - is called without the new kwargs (backward compat anchor). - -Phase 1: tests only — production code lands in Phase 2. -T1 fails today (panel not registered). -T2 fails today (helper doesn't accept the new kwargs). -T3 passes today as the backward-compat anchor. -""" -from __future__ import annotations - -from pathlib import Path - -from kernbench.benches.milestone_gqa_headline import ( # noqa: F401 - _PANEL_DISPATCH, - _PANELS, - _run_prefill_panel, -) -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 - -TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" - - -def _engine_factory(t, d): - return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) - - -# ── T1: new panel registered ───────────────────────────────────────── - - -def test_single_kv_group_prefill_panel_registered(): - """The ``single_kv_group_prefill_gqa_c8_p8`` panel must be in ``_PANELS`` and - its dispatch entry must have the expected LLaMA-3.1-70B params. - - Headline config (scratch-limited; LLaMA 32K headline awaits - Q-axis kernel tiling): - C = 8 (head-parallel, snake sub-mesh) - P = 8 (intra-CUBE PE-SP) - T_q = 1_024 (one-shot long-context prefill) - S_kv = 1_024 - (scratch-limited; LLaMA 32K headline awaits Q-axis kernel tiling) - d_head = 128 (LLaMA-3.1-70B) - """ - panel_name = "single_kv_group_prefill_gqa_c8_p8" - assert panel_name in _PANELS, ( - f"{panel_name!r} not in _PANELS; got {_PANELS}" - ) - assert panel_name in _PANEL_DISPATCH, ( - f"{panel_name!r} not in _PANEL_DISPATCH" - ) - kind, params = _PANEL_DISPATCH[panel_name] - assert kind == "prefill", f"kind={kind!r}, expected 'prefill'" - assert params.get("C") == 8, f"C={params.get('C')}, expected 8" - assert params.get("P") == 8, f"P={params.get('P')}, expected 8" - assert params.get("T_q") == 1_024, ( - f"T_q={params.get('T_q')}, expected 1_024" - ) - assert params.get("S_kv") == 1_024, ( - f"S_kv={params.get('S_kv')}, expected 1_024" - ) - assert params.get("d_head") == 128, ( - f"d_head={params.get('d_head')}, expected 128" - ) - - -# ── T2: helper drives the C=8, P=8 prefill kernel to completion ────── - - -def test_single_kv_group_prefill_panel_runner_smoke(): - """``_run_prefill_panel`` must accept the new ``P``, ``T_q``, - ``d_head`` kwargs and successfully launch the prefill kernel at - ``(C, P) = (8, 8)`` with the snake-mapped 2×4 SFR. - - Uses **smaller** T_q/S_kv than the headline panel so the test - completes quickly. The headline 32K dims are exercised via - ``kernbench run milestone-gqa-headline``, not pytest. - """ - topo = resolve_topology(str(TOPOLOGY_DEFAULT)) - - def _bench_fn(ctx): - _run_prefill_panel( - ctx, - panel="single_kv_group_prefill_gqa_c8_p8", - C=8, P=8, - T_q=8, S_kv=8192, # smoke dims, not the 32K headline - d_head=128, - ) - - result = run_bench( - topology=topo, bench_fn=_bench_fn, - device=resolve_device(None), - engine_factory=_engine_factory, - ) - assert result.completion.ok, ( - f"single_kv_group prefill panel smoke at C=8 P=8 must complete; " - f"got {result.completion}" - ) - - -# Backward-compat test removed when the legacy multi_user_prefill_gqa -# panel was dropped; signature-extension contract is now exercised -# implicitly by the live single_kv_group_prefill_gqa_c8_p8 panel.