diff --git a/src/kernbench/benches/milestone_gqa_decode_4cases.py b/src/kernbench/benches/milestone_gqa_decode_4cases.py new file mode 100644 index 0000000..002b152 --- /dev/null +++ b/src/kernbench/benches/milestone_gqa_decode_4cases.py @@ -0,0 +1,152 @@ +"""milestone-gqa-decode-4cases: comparative study of 4 decode sharding cases. + +Per GQA_full_deck.pptx slides 11-17: 4 KV-cache sharding strategies on +the LLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs): + + Case 1 Cube-SP / PE-TP → KV split by S_kv across cubes; PEs TP on batch + Case 2 Cube-Repl / PE-TP → full KV per cube; PEs TP on batch + Case 3 Cube-Repl / PE-SP → full KV per cube; PEs SP on S_kv (intra-cube AR) + Case 4 Cube-SP / PE-SP → KV split 64-way; 2-phase AR on (m,ℓ,O) ★ optimal + +Each case is a separate panel. The bench drives all panels in one +invocation and writes per-panel op_log_summary to sweep.json so the +comparative analysis (latency, GEMM/MAC util, comm volume) can be +generated from a single sweep. + +Status (initial commit, 5C.D): + - Case 4 panel implemented (uses _gqa_attention_decode_long with + sub_w=4 — ADR-0060 §4.2 lrab-adapted center-root reduce; that is + structurally Case 4 per slide 11 with the reduce-to-root variant + of the AR pattern). + - Cases 1-3 panels: TBD in subsequent sub-increments (5C.A/B/C). + +Deviation from slide 13: slide prescribes AllReduce (every rank has +the answer); the kernel does reduce-to-root (only the lrab center +cube has it) per ADR-0060 §4. Treated as the kernbench Case-4 baseline. + +Gated by ``GQA_DECODE_4CASES_RUN=1``. +""" +from __future__ import annotations + +import json +import os +from pathlib import Path + +from kernbench.benches.milestone_gqa_headline import ( + _ccl_cfg, + _run_decode_panel, + _summarize_op_log, +) +from kernbench.benches.registry import bench + +_OUTPUT_DIR = ( + Path(__file__).resolve().parent + / "1H_milestone_output" + / "gqa_decode_4cases" +) +_SWEEP_JSON = _OUTPUT_DIR / "sweep.json" + + +# ── Panel registry ─────────────────────────────────────────────────── + + +_PANELS = ( + "single_kv_group_decode_gqa_cube_sp_pe_sp", # Case 4 + # Cases 1-3 to be added by 5C.A/B/C +) + +# Each entry: (kind, panel-specific params for _run_decode_panel). +# LLaMA-3.1-70B single-KV-head group target: +# 1 KV head, h_q = 8 (G = 8 group), d_head = 128 +# 8 cubes (head-parallel group), 8 PEs/cube +# S_kv = 128K (long-context decode), T_q = 1 (one new token per pass) +_PANEL_DISPATCH: dict[str, tuple[str, dict]] = { + "single_kv_group_decode_gqa_cube_sp_pe_sp": ("decode", { + # Case 4: KV split 64-way (Cube-SP × PE-SP), 2-level reduce. + # sub_w=4 ⇒ sub_h=2 ⇒ lrab center-root cube = (1,2) = cube 6. + "C": 8, "P": 8, "sub_w": 4, + "T_q": 1, "S_kv": 131_072, + "d_head": 128, "h_q": 8, "h_kv": 1, + }), +} + + +# ── Per-panel runner ───────────────────────────────────────────────── + + +def _make_bench_fn(panel: str): + kind, params = _PANEL_DISPATCH[panel] + + def _bench_fn(ctx): + if kind == "decode": + _run_decode_panel(ctx, panel=panel, **params) + else: + raise RuntimeError( + f"milestone-gqa-decode-4cases panel {panel!r} has " + f"unsupported kind={kind!r}; only 'decode' is allowed." + ) + return _bench_fn + + +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"milestone-gqa-decode-4cases 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-decode-4cases", + description=( + "Comparative decode study of 4 KV-cache sharding cases 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_4CASES_RUN=1. + """ + _OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + if not os.environ.get("GQA_DECODE_4CASES_RUN"): + raise RuntimeError( + "milestone-gqa-decode-4cases needs GQA_DECODE_4CASES_RUN=1." + ) + + topology = os.environ.get( + "GQA_DECODE_4CASES_TOPOLOGY", "topology.yaml", + ) + rows = [_run_panel(panel, topology) for panel in _PANELS] + sweep = { + "version": 1, + "panels": list(_PANELS), + "rows": rows, + } + _SWEEP_JSON.write_text(json.dumps(sweep, indent=2)) + print( + f" milestone-gqa-decode-4cases: {len(rows)} rows -> {_SWEEP_JSON}" + ) diff --git a/src/kernbench/benches/milestone_gqa_headline.py b/src/kernbench/benches/milestone_gqa_headline.py index ae045a2..3fc4789 100644 --- a/src/kernbench/benches/milestone_gqa_headline.py +++ b/src/kernbench/benches/milestone_gqa_headline.py @@ -143,24 +143,31 @@ def _run_prefill_panel( ) -def _run_decode_panel(ctx, *, panel: str, C: int, P: int, S_kv: int) -> None: +def _run_decode_panel( + ctx, *, panel: str, C: int, P: int, S_kv: int, + sub_w: int = 0, + T_q: int = _T_Q_DECODE, + d_head: int = _D_HEAD, + h_q: int = _H_Q_DECODE, + h_kv: int = _H_KV_DECODE, +) -> None: configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) dp_full = DPPolicy(cube="replicate", pe="replicate", num_cubes=C, num_pes=P) dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate", pe="row_wise", num_cubes=C, num_pes=P) - q = ctx.zeros((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD), + q = ctx.zeros((T_q, h_q * d_head), dtype=_DTYPE, dp=dp_full, name=f"{panel}_q") - k = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD), + k = ctx.zeros((S_kv, h_kv * d_head), dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k") - v = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD), + v = ctx.zeros((S_kv, h_kv * d_head), dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v") - o = ctx.empty((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD), + o = ctx.empty((T_q, h_q * d_head), dtype=_DTYPE, dp=dp_full, name=f"{panel}_o") ctx.launch( panel, gqa_attention_decode_long_kernel, q, k, v, o, - _T_Q_DECODE, S_kv, _H_Q_DECODE, _H_KV_DECODE, _D_HEAD, C, P, + T_q, S_kv, h_q, h_kv, d_head, C, P, sub_w, _auto_dim_remap=False, ) diff --git a/tests/attention/test_milestone_gqa_decode_4cases.py b/tests/attention/test_milestone_gqa_decode_4cases.py new file mode 100644 index 0000000..2655531 --- /dev/null +++ b/tests/attention/test_milestone_gqa_decode_4cases.py @@ -0,0 +1,188 @@ +"""Tests for the decode 4-cases comparative-study bench. + +Per ``GQA_full_deck.pptx`` slides 11-17, the 4 cases differ in how KV +cache is sharded across the 8 cubes and 8 PEs of a single KV-head +group on LLaMA-3.1-70B GQA: + + Case 1 Cube-SP / PE-TP → KV split by S_kv across cubes; PEs TP on batch + Case 2 Cube-Repl / PE-TP → full KV per cube; PEs TP on batch + Case 3 Cube-Repl / PE-SP → full KV per cube; PEs SP on S_kv (intra-cube AR) + Case 4 Cube-SP / PE-SP → KV split 64-way; 2-phase AR on (m,ℓ,O) ★ optimal + +This file grows as each case lands. Phase 1 of 5C.D adds Case 4 first, +which is structurally the existing ``_gqa_attention_decode_long.py`` at +``sub_w=4`` (Increment 2's lrab-adapted center-root reduce). The Phase 2 +production change introduces a new bench file +``src/kernbench/benches/milestone_gqa_decode_4cases.py`` housing the 4 +case panels under a single ``milestone-gqa-decode-4cases`` entry, and +extends ``_run_decode_panel`` in ``milestone_gqa_headline`` to accept +``sub_w``/``d_head``/``h_q``/``h_kv`` overrides. + +Deviation from slide 13: slide prescribes AllReduce on (m,ℓ,O); the +kernel does reduce-to-root (only the lrab center cube has the answer) +per ADR-0060 §4. Treated as the kernbench Case-4 baseline. + +Phase 1: tests only. T1, T2, T3, T4 fail today (new bench file does +not yet exist; ``_run_decode_panel`` does not yet accept ``sub_w``). +""" +from __future__ import annotations + +import re +from pathlib import Path + +from kernbench.benches.milestone_gqa_headline import _run_decode_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" + +_CASE4_PANEL = "single_kv_group_decode_gqa_cube_sp_pe_sp" +_CUBE_RE = re.compile(r"\bcube(\d+)\b") + + +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 _dma_write_cubes(op_log) -> list[int]: + cubes: list[int] = [] + for r in op_log: + if r.op_name != "dma_write": + continue + m = _CUBE_RE.search(r.component_id) + if m is not None: + cubes.append(int(m.group(1))) + return cubes + + +def _run_case4_smoke(*, S_kv: int): + """Drive the Case 4 decode panel via ``_run_decode_panel``. + + Uses ``S_kv=8192`` (smoke) to keep test time bounded; the headline + ``S_kv=128K`` runs come from ``kernbench run --bench + milestone-gqa-decode-4cases``, not pytest. + """ + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + _run_decode_panel( + ctx, panel=_CASE4_PANEL, + C=8, P=8, sub_w=4, + S_kv=S_kv, + d_head=128, h_q=8, h_kv=1, + ) + + return run_bench( + topology=topo, bench_fn=_bench_fn, + device=resolve_device(None), + engine_factory=_engine_factory, + ) + + +# ── T1: Case 4 panel is registered in the new bench ────────────────── + + +def test_case4_panel_registered(): + """The Case 4 panel must be in the new bench's ``_PANELS`` + + ``_PANEL_DISPATCH`` with the expected LLaMA-3.1-70B target dims. + + Headline config: + C = 8 (head-parallel CUBE Group) + P = 8 (intra-CUBE PE-SP) + sub_w = 4 (lrab-adapted center-root reduce; root cube 6) + T_q = 1 (decode: one new token per pass) + S_kv = 131_072 (LLaMA long-context decode target) + d_head = 128, h_q = 8, h_kv = 1 + """ + from kernbench.benches.milestone_gqa_decode_4cases import ( + _PANEL_DISPATCH, + _PANELS, + ) + assert _CASE4_PANEL in _PANELS, ( + f"{_CASE4_PANEL!r} not in _PANELS; got {_PANELS}" + ) + assert _CASE4_PANEL in _PANEL_DISPATCH + kind, params = _PANEL_DISPATCH[_CASE4_PANEL] + assert kind == "decode", f"kind={kind!r}, expected 'decode'" + assert params.get("C") == 8 + assert params.get("P") == 8 + assert params.get("sub_w") == 4 + assert params.get("T_q") == 1 + assert params.get("S_kv") == 131_072 + assert params.get("d_head") == 128 + assert params.get("h_q") == 8 + assert params.get("h_kv") == 1 + + +# ── T2: Case 4 runner drives the kernel to completion ─────────────── + + +def test_case4_runner_smoke(): + """``_run_decode_panel`` must accept the new ``sub_w``, + ``d_head``, ``h_q``, ``h_kv`` kwargs and launch the kernel at + ``(C, P, sub_w) = (8, 8, 4)`` (the Case 4 / lrab path). + + Smoke uses ``S_kv=8192`` so the simulation completes quickly. The + headline 128K dims run via ``kernbench run --bench + milestone-gqa-decode-4cases``. + """ + result = _run_case4_smoke(S_kv=8192) + assert result.completion.ok, ( + f"Case 4 decode smoke at C=8 P=8 sub_w=4 must complete; " + f"got {result.completion}" + ) + + +# ── T3: reduce-to-root lands at the lrab center cube (cube 6) ─────── + + +def test_case4_root_at_center_cube_6(): + """For ``sub_w=4, sub_h=2``: root_col=2, root_row=1, root_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) + 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" + distinct = set(cubes) + assert distinct == {6}, ( + f"Case 4 root must be the lrab center cube 6; " + f"got cubes={sorted(distinct)}" + ) + + +# ── T4: 2-phase AR ipcq pattern matches the predicted Case-4 traffic ─ + + +def test_case4_two_level_ar_ipcq_pattern(): + """Total ipcq_copy for the Case 4 reduce at (C, P, sub_w) = + (8, 8, 4): + + Intra-CUBE (per CUBE = 8 PEs in a 2×4 grid): + row chain along intra_W: cols 1,2,3 each row × 2 rows × + 3 tensors = 18 + col bridge along intra_N: pe4 only × 3 tensors = 3 + per-CUBE intra total = 21 + × 8 CUBEs = 168 + + Inter-CUBE lrab (sub_w=4, sub_h=2): + Phase 1 row reduce — 3 sends/row × 3 tensors × 2 rows = 18 + Phase 2 col reduce — cube 2 → S × 3 tensors = 3 + inter-CUBE total = 21 + + Grand total: 168 + 21 = 189 + """ + result = _run_case4_smoke(S_kv=8192) + assert result.completion.ok + n_copy = _count(result.engine.op_log, "ipcq_copy") + assert n_copy == 189, ( + f"Case 4 expected 189 ipcq_copy " + f"(168 intra-CUBE + 21 inter-CUBE lrab); got {n_copy}" + ) diff --git a/tests/attention/test_milestone_gqa_headline.py b/tests/attention/test_milestone_gqa_headline.py index 77b89d8..d728710 100644 --- a/tests/attention/test_milestone_gqa_headline.py +++ b/tests/attention/test_milestone_gqa_headline.py @@ -37,6 +37,7 @@ BENCH_NAME = "milestone-gqa-headline" PANELS = ( "single_user_prefill_gqa", "multi_user_prefill_gqa", + "single_kv_group_prefill_gqa_c8_p8", "single_user_decode_gqa", "multi_user_decode_gqa", ) @@ -88,12 +89,12 @@ def test_validation_run_completes_ok(monkeypatch): # ── sweep.json shape ────────────────────────────────────────────────── -def test_sweep_json_has_four_panels(monkeypatch): +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"]) == 4 + assert len(data["rows"]) == len(PANELS) assert {r["panel"] for r in data["rows"]} == set(PANELS)