"""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._gqa_attention_decode_cube_repl_pe_sp import ( gqa_attention_decode_cube_repl_pe_sp_kernel, ) from kernbench.benches._gqa_attention_decode_cube_repl_pe_tp import ( gqa_attention_decode_cube_repl_pe_tp_kernel, ) from kernbench.benches.milestone_gqa_headline import ( _ccl_cfg, _run_decode_panel, _summarize_op_log, ) from kernbench.benches.registry import bench from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip from kernbench.policy.placement.dp import DPPolicy _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 ★ optimal "single_kv_group_decode_gqa_cube_repl_pe_tp", # Case 2 (no comm; 8× memory) "single_kv_group_decode_gqa_cube_repl_pe_sp", # Case 3 (intra-CUBE AR only; 8× memory) # Case 1 to be added by 5C.A ) # Each entry: (kind, panel-specific params). # 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, }), "single_kv_group_decode_gqa_cube_repl_pe_tp": ("decode_cube_repl_pe_tp", { # Case 2: K, V replicated everywhere (8× memory waste); PEs TP # on batch. For B=1 only one rank works (slide-11 PE-TP waste). # No inter-rank communication. "C": 8, "P": 8, "T_q": 1, "S_kv": 131_072, "d_head": 128, "h_q": 8, "h_kv": 1, }), "single_kv_group_decode_gqa_cube_repl_pe_sp": ("decode_cube_repl_pe_sp", { # Case 3: K, V replicated per cube (8× memory); PEs SP on S_kv # within each cube. Intra-CUBE 8-way reduce; no inter-CUBE comm # (every cube ends with full answer; designated writer = cube 0). "C": 8, "P": 8, "T_q": 1, "S_kv": 131_072, "d_head": 128, "h_q": 8, "h_kv": 1, }), } # ── Per-panel runner ───────────────────────────────────────────────── def _run_decode_panel_cube_repl_pe_tp( ctx, *, panel: str, C: int, P: int, T_q: int, S_kv: int, d_head: int, h_q: int, h_kv: int, ) -> None: """Case 2 runner: K, V replicated everywhere; B=1 single-rank work. DPPolicy models the cluster-wide memory waste — every rank holds full K, V in its HBM region. Only PE 0 of CUBE 0 computes (the kernel early-returns on every other rank), so only one rank reads from its HBM copy and writes the output. """ configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) dp_repl = DPPolicy(cube="replicate", pe="replicate", num_cubes=C, num_pes=P) q = ctx.zeros((T_q, h_q * d_head), dtype="f16", dp=dp_repl, name=f"{panel}_q") k = ctx.zeros((S_kv, h_kv * d_head), dtype="f16", dp=dp_repl, name=f"{panel}_k") v = ctx.zeros((S_kv, h_kv * d_head), dtype="f16", dp=dp_repl, name=f"{panel}_v") o = ctx.empty((T_q, h_q * d_head), dtype="f16", dp=dp_repl, name=f"{panel}_o") ctx.launch( panel, gqa_attention_decode_cube_repl_pe_tp_kernel, q, k, v, o, T_q, S_kv, h_q, h_kv, d_head, C, P, _auto_dim_remap=False, ) def _run_decode_panel_cube_repl_pe_sp( ctx, *, panel: str, C: int, P: int, T_q: int, S_kv: int, d_head: int, h_q: int, h_kv: int, ) -> None: """Case 3 runner: K, V replicated per cube; PEs SP on S_kv within cube. DPPolicy models the cluster-wide 8× memory waste — every cube holds full K, V in its HBM region, then splits the S_kv axis row_wise across its 8 PEs. The kernel does an intra-CUBE 8-way reduce on (m, ℓ, O); only cube 0's PE 0 writes the output. """ 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="replicate", pe="row_wise", num_cubes=C, num_pes=P) q = ctx.zeros((T_q, h_q * d_head), dtype="f16", dp=dp_full, name=f"{panel}_q") k = ctx.zeros((S_kv, h_kv * d_head), dtype="f16", dp=dp_kv, name=f"{panel}_k") v = ctx.zeros((S_kv, h_kv * d_head), dtype="f16", dp=dp_kv, name=f"{panel}_v") o = ctx.empty((T_q, h_q * d_head), dtype="f16", dp=dp_full, name=f"{panel}_o") ctx.launch( panel, gqa_attention_decode_cube_repl_pe_sp_kernel, q, k, v, o, T_q, S_kv, h_q, h_kv, d_head, C, P, _auto_dim_remap=False, ) 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) elif kind == "decode_cube_repl_pe_tp": _run_decode_panel_cube_repl_pe_tp(ctx, panel=panel, **params) elif kind == "decode_cube_repl_pe_sp": _run_decode_panel_cube_repl_pe_sp(ctx, panel=panel, **params) else: raise RuntimeError( f"milestone-gqa-decode-4cases panel {panel!r} has " f"unsupported kind={kind!r}." ) 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}" )