diff --git a/scripts/paper/measure_gqa_decode_placement_comm.py b/scripts/paper/measure_gqa_decode_placement_comm.py new file mode 100644 index 0000000..5651424 --- /dev/null +++ b/scripts/paper/measure_gqa_decode_placement_comm.py @@ -0,0 +1,362 @@ +"""Measured comm overlay for all 6 GQA decode KV placements. + +Runs the simulator for Cases 1-6 (the same 6 placements the chart in +paper_plot_gqa_4cases_summary.py covers analytically) at S_kv = 64 K, +sums actual IPCQ-copy bytes from the engine op_log, projects to +per-token (x80 layers), scales the partial-score-AR component of +Cases 4/5 from S_kv = 64 K -> S_kv = 1 M (linear in S_kv; other +cases are S_kv-independent), adds the constant Wo + FFN AR +(1.25 MB / token), and writes the result to JSON for +paper_plot_gqa_4cases_summary.py to overlay on the analytical bars. + +Single layer of decode attention only — the projection × 80 takes +the per-layer measurement to a per-token total. + +Usage: + python scripts/paper/measure_gqa_decode_placement_comm.py +""" +from __future__ import annotations + +import json +import os +import sys +from pathlib import Path + +from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_repl_pe_sp import ( + gqa_attention_decode_long_ctx_cube_repl_pe_sp_kernel as _case3_kernel, +) +from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_repl_pe_tp import ( + gqa_attention_decode_long_ctx_cube_repl_pe_tp_kernel as _case1_kernel, +) +from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import ( + gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel as _case6_kernel, +) +from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp import ( + gqa_attention_decode_long_ctx_cube_sp_pe_tp_kernel as _case2_kernel, +) +from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead import ( + gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel as _case4_kernel, +) +from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp import ( + gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel as _case5_kernel, +) +from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import _ccl_cfg +from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip +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 + +_C = 8 +_P = 8 +_N_LAYERS = 80 +_S_KV_MEAS = 8 * 1024 # per-run simulator S_kv (small — 1/128th of headline) +_S_KV_HEADLINE = 1 << 20 # 1 Mi tokens, the chart's headline S_kv + +# Per-cube S_kv share for the d_head-TP partial-score AR cost. +# Case 4 (Cube-SP × PE-TP_dhead): per-cube = S_kv / C +# Case 5 (Cube-TP_dhead × PE-SP): per-cube = S_kv (KV replicated across +# cubes for the cube-axis d_head-TP), so partial-score AR scales +# with the full S_kv. +# Headline / measured per-cube ratios give the partial-score-AR scale-up +# from S_kv = 64 K to S_kv = 1 M. (m,ℓ,O) AR is S_kv-independent. +_PARTIAL_SCORE_SCALE = _S_KV_HEADLINE / _S_KV_MEAS # = 16 + +# Per-token Wo + FFN AR (constant across all cases, comes from the +# attn-output and FFN-down all-reduces NOT measured by the attention- +# only kernel run here). +_WO_PER_LAYER_BYTES = 8 * 1024 +_FFN_PER_LAYER_BYTES = 8 * 1024 +_WO_FFN_PER_TOKEN_BYTES = ( + (_WO_PER_LAYER_BYTES + _FFN_PER_LAYER_BYTES) * _N_LAYERS +) # 1.25 MB + +# Total PE count in one KV-head group — average per-PE comm = total / N. +_NUM_PES = _C * _P + +# Total partial-score-AR slice produced by the attention compute when +# d_head is sharded. Used to split measured IPCQ traffic into the +# S_kv-scaling component (partial scores) vs the S_kv-independent +# component ((m,ℓ,O) merge). Same as the analytical formula in +# paper_plot_gqa_4cases_summary.py: h_q · S_q · per_cube_S_kv · 2 bytes. +_H_Q = 8 +_S_Q = 1 +_BYTES_PER_ELEM = 2 + +_PARAMS = dict(C=_C, P=_P, T_q=_S_Q, S_kv=_S_KV_MEAS, + d_head=128, h_q=_H_Q, h_kv=1) + + +def _bench_fn_case1(ctx): + """Case 1: Cube-Repl x PE-repl (PE-TP doesn't shard KV).""" + p = _PARAMS + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp = DPPolicy(cube="replicate", pe="replicate", + num_cubes=p["C"], num_pes=p["P"]) + q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp, name="q_c1") + k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp, name="k_c1") + v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp, name="v_c1") + o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp, name="o_c1") + ctx.launch("case1_repl_repl", _case1_kernel, + q, k, v, o, + p["T_q"], p["S_kv"], p["h_q"], p["h_kv"], + p["d_head"], p["C"], p["P"], + _auto_dim_remap=False) + + +def _bench_fn_case2(ctx): + """Case 2: Cube-SP x PE-repl (PE-TP doesn't shard KV).""" + p = _PARAMS + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp_full = DPPolicy(cube="replicate", pe="replicate", + num_cubes=p["C"], num_pes=p["P"]) + dp_kv = DPPolicy(cube="row_wise", pe="replicate", + num_cubes=p["C"], num_pes=p["P"]) + q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp_full, name="q_c2") + k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp_kv, name="k_c2") + v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp_kv, name="v_c2") + o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp_full, name="o_c2") + ctx.launch("case2_sp_repl", _case2_kernel, + q, k, v, o, + p["T_q"], p["S_kv"], p["h_q"], p["h_kv"], + p["d_head"], p["C"], p["P"], + _auto_dim_remap=False) + + +def _bench_fn_case3(ctx): + """Case 3: Cube-Repl x PE-SP.""" + p = _PARAMS + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp_full = DPPolicy(cube="replicate", pe="replicate", + num_cubes=p["C"], num_pes=p["P"]) + dp_kv = DPPolicy(cube="replicate", pe="row_wise", + num_cubes=p["C"], num_pes=p["P"]) + q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp_full, name="q_c3") + k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp_kv, name="k_c3") + v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp_kv, name="v_c3") + o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp_full, name="o_c3") + ctx.launch("case3_repl_sp", _case3_kernel, + q, k, v, o, + p["T_q"], p["S_kv"], p["h_q"], p["h_kv"], + p["d_head"], p["C"], p["P"], + _auto_dim_remap=False) + + +def _bench_fn_case4(ctx): + p = _PARAMS + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp_full = DPPolicy(cube="replicate", pe="column_wise", + num_cubes=p["C"], num_pes=p["P"]) + dp_kv = DPPolicy(cube="row_wise", pe="column_wise", + num_cubes=p["C"], num_pes=p["P"]) + q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp_full, name="q_c4") + k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp_kv, name="k_c4") + v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp_kv, name="v_c4") + o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp_full, name="o_c4") + ctx.launch("case4_dhead_tp", _case4_kernel, + q, k, v, o, + p["T_q"], p["S_kv"], p["h_q"], p["h_kv"], + p["d_head"], p["C"], p["P"], + _auto_dim_remap=False) + + +def _bench_fn_case5(ctx): + p = _PARAMS + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp_q = DPPolicy(cube="column_wise", pe="replicate", + num_cubes=p["C"], num_pes=p["P"]) + dp_kv = DPPolicy(cube="column_wise", pe="row_wise", + num_cubes=p["C"], num_pes=p["P"]) + q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp_q, name="q_c5") + k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp_kv, name="k_c5") + v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp_kv, name="v_c5") + o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp_q, name="o_c5") + ctx.launch("case5_dhead_tp_inter", _case5_kernel, + q, k, v, o, + p["T_q"], p["S_kv"], p["h_q"], p["h_kv"], + p["d_head"], p["C"], p["P"], + _auto_dim_remap=False) + + +def _bench_fn_case6(ctx): + p = _PARAMS + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp_full = DPPolicy(cube="replicate", pe="replicate", + num_cubes=p["C"], num_pes=p["P"]) + dp_kv = DPPolicy(cube="row_wise", pe="row_wise", + num_cubes=p["C"], num_pes=p["P"]) + q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp_full, name="q_c6") + k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp_kv, name="k_c6") + v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp_kv, name="v_c6") + o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp_full, name="o_c6") + ctx.launch("case6_sp_sp", _case6_kernel, + q, k, v, o, + p["T_q"], p["S_kv"], p["h_q"], p["h_kv"], + p["d_head"], p["C"], p["P"], + _auto_dim_remap=False) + + +def _sum_ipcq_bytes(op_log) -> int: + """Sum nbytes across all ipcq_copy records.""" + return sum( + r.params.get("nbytes", 0) + for r in op_log + if r.op_kind == "memory" and r.op_name == "ipcq_copy" + ) + + +def _partial_score_slices(case: int) -> int: + """Divisor that splits S_kv into partial-score tiles, per the + analytical model in paper_plot_gqa_4cases_summary.py: + + partial_score_per_PE = h_q * S_q * (s_kv / slices) * bytes + + Case 4 (Cube-SP x PE-TP-dhead): intra-cube AR over d_head shards + on PE axis -> partial tile per PE has per-cube S_kv = s_kv/C. + Case 5 (Cube-TP-dhead x PE-SP): inter-cube AR over d_head shards + on cube axis -> partial tile per PE has per-PE S_kv = s_kv/P. + Cases 1, 2, 3, 6: no partial-score AR (only (m,l,O) merge). + """ + if case == 4: + return _C + if case == 5: + return _P + return 0 + + +def _split_attn_layer_bytes(case: int, total_ipcq_bytes: int, + s_kv: int) -> tuple[int, int]: + """Split per-layer attention-time IPCQ bytes into: + (partial_score_component, mlo_component). + + The partial-score component scales with s_kv (so it must be scaled + when projecting from the measure-time s_kv to the headline s_kv); + the (m,l,O) component is constant in s_kv. + + Partial-score size is analytically known per-case (formula in + _partial_score_slices); the remainder is treated as (m,l,O) + any + other S_kv-independent overhead. Per-PE = total / NUM_PES. + """ + per_pe_total = total_ipcq_bytes // _NUM_PES + slices = _partial_score_slices(case) + if slices == 0: + # No partial-score AR for this case. + return 0, per_pe_total + partial_score_per_pe = ( + _H_Q * _S_Q * (s_kv // slices) * _BYTES_PER_ELEM + ) + partial_score_per_pe = min(partial_score_per_pe, per_pe_total) + mlo_per_pe = per_pe_total - partial_score_per_pe + return partial_score_per_pe, mlo_per_pe + + +_KERNELS = ( + (1, "Case 1 (Cube-Repl x PE-repl)", _bench_fn_case1), + (2, "Case 2 (Cube-SP x PE-repl)", _bench_fn_case2), + (3, "Case 3 (Cube-Repl x PE-SP)", _bench_fn_case3), + (4, "Case 4 (Cube-SP x PE-TP d_head)", _bench_fn_case4), + (5, "Case 5 (Cube-TP d_head x PE-SP)", _bench_fn_case5), + (6, "Case 6 (Cube-SP x PE-SP) [*]", _bench_fn_case6), +) + + +def main() -> int: + topology = os.environ.get("GQA_1H_TOPOLOGY", "topology.yaml") + topo = resolve_topology(topology) + + out: dict = { + "S_kv_measured": _S_KV_MEAS, + "S_kv_headline": _S_KV_HEADLINE, + "n_layers": _N_LAYERS, + "num_pes": _NUM_PES, + "wo_ffn_per_token_bytes": _WO_FFN_PER_TOKEN_BYTES, + "cases": {}, + } + + print(f"Measuring at S_kv={_S_KV_MEAS:,} ; scaling partial-score AR " + f"to S_kv={_S_KV_HEADLINE:,} (×{int(_PARTIAL_SCORE_SCALE)})") + print() + + for case_id, label, bench_fn in _KERNELS: + try: + res = run_bench( + topology=topo, bench_fn=bench_fn, + device=resolve_device(None), + engine_factory=lambda t, d: GraphEngine( + getattr(t, "topology_obj", t), enable_data=True, + ), + ) + except Exception as e: + print(f" {label:<42} FAIL: {type(e).__name__}: {e}") + return 1 + if not res.completion.ok: + print(f" {label:<42} ENGINE FAIL: {res.completion}") + return 1 + + total_ipcq = _sum_ipcq_bytes(res.engine.op_log) + partial_pe, mlo_pe = _split_attn_layer_bytes( + case_id, total_ipcq, _S_KV_MEAS, + ) + # Per-token attention-time comm at S_kv = 1 M: + # (partial_score_per_layer × scale + mlo_per_layer) × 80 layers + scaled_partial_per_token = ( + partial_pe * int(_PARTIAL_SCORE_SCALE) * _N_LAYERS + ) + mlo_per_token = mlo_pe * _N_LAYERS + attn_per_token = scaled_partial_per_token + mlo_per_token + total_per_token = attn_per_token + _WO_FFN_PER_TOKEN_BYTES + + out["cases"][str(case_id)] = { + "label": label, + "total_ipcq_bytes_one_layer": total_ipcq, + "per_pe_partial_score_bytes_one_layer": partial_pe, + "per_pe_mlo_bytes_one_layer": mlo_pe, + "per_pe_attn_bytes_per_token_at_1M": attn_per_token, + "per_pe_total_bytes_per_token_at_1M": total_per_token, + } + + print(f" {label:<42} " + f"ipcq_total={total_ipcq:>10,} " + f"per_pe_attn(1L)={(partial_pe + mlo_pe):>9,} " + f"per_pe_total/tok@1M={total_per_token / (1<<20):>7.2f} MB") + + out_path = ( + Path(__file__).resolve().parents[2] + / "src" / "kernbench" / "benches" + / "1H_milestone_output" / "gqa" / "long_ctx" + / "gqa_3cases_measured_comm.json" + ) + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text(json.dumps(out, indent=2)) + print() + print(f"wrote {out_path}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/paper/paper_plot_gqa_4cases_summary.py b/scripts/paper/paper_plot_gqa_4cases_summary.py index 8b7f81b..b4ecc70 100644 --- a/scripts/paper/paper_plot_gqa_4cases_summary.py +++ b/scripts/paper/paper_plot_gqa_4cases_summary.py @@ -114,6 +114,7 @@ Output PNG: """ from __future__ import annotations +import json from pathlib import Path import matplotlib.patches as mpatches @@ -240,6 +241,22 @@ _OUT_DIR = ( / "src" / "kernbench" / "benches" / "1H_milestone_output" / "gqa" / "long_ctx" ) +_MEASURED_JSON = _OUT_DIR / "gqa_3cases_measured_comm.json" + + +def _load_measured() -> dict[int, float] | None: + """Load measured per-PE comm bytes (already scaled to S_kv=1M). + + Returns {case_id: per_token_total_bytes} or None if JSON missing. + Produced by scripts/paper/measure_gqa_decode_placement_comm.py. + """ + if not _MEASURED_JSON.exists(): + return None + data = json.loads(_MEASURED_JSON.read_text()) + return { + int(cid): info["per_pe_total_bytes_per_token_at_1M"] + for cid, info in data["cases"].items() + } # ── Formulae ──────────────────────────────────────────────────────── @@ -388,51 +405,114 @@ def _plot_memory(ax) -> None: ax.legend(loc="upper right", fontsize=9) -def _plot_comm(ax) -> None: - wo_mb = [] - ffn_mb = [] - attn_mb = [] +def _plot_comm(ax, *, mode: str = "analytical") -> None: + """Per-PE comm panel. + + mode = "analytical": single solid bars from per_token_bytes formula. + mode = "paired" : analytical (solid) + simulator-measured + (hatched) side-by-side per case, when the + measurement JSON is available. + """ + labels = [_CASE_LABEL[c] for c in _CASES] + + wo_mb_list: list[float] = [] + ffn_mb_list: list[float] = [] + attn_mb: list[float] = [] for c in _CASES: wo, ffn, attn = per_token_bytes(c, _HEADLINE_S_KV) - wo_mb.append(wo / (1 << 20)) - ffn_mb.append(ffn / (1 << 20)) + wo_mb_list.append(wo / (1 << 20)) + ffn_mb_list.append(ffn / (1 << 20)) attn_mb.append(attn / (1 << 20)) - labels = [_CASE_LABEL[c] for c in _CASES] - x = list(range(len(_CASES))) + measured = _load_measured() if mode == "paired" else None + paired = measured is not None + source_tag = "analytical (solid) vs simulator-measured (hatched)" \ + if paired else "analytical" - ax.bar(x, wo_mb, color=_WO_COLOR, edgecolor="black") - ax.bar(x, ffn_mb, bottom=wo_mb, color=_FFN_COLOR, edgecolor="black") - bottoms_attn = [w + f for w, f in zip(wo_mb, ffn_mb)] - ax.bar(x, attn_mb, bottom=bottoms_attn, + n_cases = len(_CASES) + x = list(range(n_cases)) + bar_w = 0.36 if paired else 0.65 + x_ana = [xi - bar_w / 2 for xi in x] if paired else x + x_meas = [xi + bar_w / 2 for xi in x] if paired else None + + # Analytical bars (solid). + ax.bar(x_ana, wo_mb_list, width=bar_w, + color=_WO_COLOR, edgecolor="black") + ax.bar(x_ana, ffn_mb_list, width=bar_w, bottom=wo_mb_list, + color=_FFN_COLOR, edgecolor="black") + bottoms_attn = [w + f for w, f in zip(wo_mb_list, ffn_mb_list)] + ax.bar(x_ana, attn_mb, width=bar_w, bottom=bottoms_attn, color=_ATTN_COLOR, edgecolor="black") + # Measured bars (hatched) — same Wo+FFN base, attn from op_log. + meas_attn_mb: list[float] = [] + if paired: + for i, c in enumerate(_CASES): + meas_total = measured.get(c, 0) / (1 << 20) + meas_attn_mb.append( + max(meas_total - wo_mb_list[i] - ffn_mb_list[i], 0.0)) + ax.bar(x_meas, wo_mb_list, width=bar_w, + color=_WO_COLOR, edgecolor="black", + hatch="///", alpha=0.85) + ax.bar(x_meas, ffn_mb_list, width=bar_w, bottom=wo_mb_list, + color=_FFN_COLOR, edgecolor="black", + hatch="///", alpha=0.85) + ax.bar(x_meas, meas_attn_mb, width=bar_w, bottom=bottoms_attn, + color=_ATTN_COLOR, edgecolor="black", + hatch="///", alpha=0.85) + ax.set_xticks(x) ax.set_xticklabels(labels, fontsize=9) ax.set_ylabel("Comm per token per PE (MB, log)") ax.set_yscale("log") - ax.set_title( + title = ( f"Per-PE communication at S_kv = {_HEADLINE_S_KV:,} tokens " - f"(decode S_q=1, B=1; {_N_LAYERS} layers)", - fontsize=11, + f"(decode S_q=1, B=1; {_N_LAYERS} layers) — {source_tag}" ) + if paired: + title += "\n(simulator measured at S_kv = 8K; " \ + "partial-score AR scaled ×128 to S_kv = 1M)" + ax.set_title(title, fontsize=10) ax.grid(axis="y", ls=":", alpha=0.5, which="both") - totals = [wo_mb[i] + ffn_mb[i] + attn_mb[i] for i in range(len(_CASES))] - ax.set_ylim(top=max(totals) * 10) + totals_ana = [wo_mb_list[i] + ffn_mb_list[i] + attn_mb[i] + for i in range(n_cases)] + ymax = max(totals_ana) + if paired: + ymax = max(ymax, max( + (measured.get(c, 0) / (1 << 20)) for c in _CASES)) + ax.set_ylim(top=ymax * 22) for i, c in enumerate(_CASES): - total_bytes = (wo_mb[i] + ffn_mb[i] + attn_mb[i]) * (1 << 20) - ax.text(x[i], totals[i] * 1.5, _fmt_bytes(total_bytes), - ha="center", fontsize=9, weight="bold") + ana_bytes = totals_ana[i] * (1 << 20) + if paired: + meas_bytes = measured.get(c, 0) + top_y = max(totals_ana[i], meas_bytes / (1 << 20)) + label = (f"ana: {_fmt_bytes(ana_bytes)}\n" + f"sim: {_fmt_bytes(meas_bytes)}") + else: + top_y = totals_ana[i] + label = _fmt_bytes(ana_bytes) + ax.text(x[i], top_y * 1.5, label, + ha="center", va="bottom", + fontsize=8 if paired else 9, + weight="bold", linespacing=1.05) + + # In-bar attention-time AR description (online-softmax (m,ℓ,O) + # merge, partial-score AR, intra/inter-cube). Placed on the + # analytical bar; with paired mode it lives on the solid bar so + # the hatched measured bar stays uncluttered. if attn_mb[i] > 0: mid = bottoms_attn[i] + attn_mb[i] / 2 - ax.text(x[i], mid, _ATTN_DESC[c], - ha="center", va="center", fontsize=7, + ax.text(x_ana[i], mid, _ATTN_DESC[c], + ha="center", va="center", + fontsize=6 if paired else 7, color="black", weight="bold") else: - ax.text(x[i], totals[i] * 0.4, _ATTN_DESC[c], - ha="center", fontsize=7.5, color="grey", style="italic") + ax.text(x_ana[i], totals_ana[i] * 0.4, _ATTN_DESC[c], + ha="center", + fontsize=6.5 if paired else 7.5, + color="grey", style="italic") legend_handles = [ mpatches.Patch(facecolor=_WO_COLOR, edgecolor="black", @@ -442,40 +522,67 @@ def _plot_comm(ax) -> None: mpatches.Patch(facecolor=_ATTN_COLOR, edgecolor="black", label="Attn-time collective"), ] + if paired: + legend_handles.append( + mpatches.Patch(facecolor="white", edgecolor="black", + hatch="///", label="simulator-measured")) ax.legend(handles=legend_handles, loc="upper right", fontsize=8, framealpha=0.92) def main() -> Path: _OUT_DIR.mkdir(parents=True, exist_ok=True) - fig = plt.figure(figsize=(21.0, 6.0)) - gs = fig.add_gridspec(1, 3, width_ratios=[0.7, 1.6, 1.6], wspace=0.22) - ax_b = fig.add_subplot(gs[0, 0]) - ax_m = fig.add_subplot(gs[0, 1]) - ax_c = fig.add_subplot(gs[0, 2]) + + # (a) Per-PE HBM budget — standalone PNG. + fig_b, ax_b = plt.subplots(figsize=(4.0, 6.0)) _plot_budget(ax_b) + fig_b.tight_layout() + out_b = _OUT_DIR / "gqa_hbm_budget.png" + fig_b.savefig(out_b, dpi=150) + plt.close(fig_b) + print(f"wrote {out_b}") + + # (b) Combined 3-panel summary — HBM budget + KV memory + comm. + fig = plt.figure(figsize=(22.0, 6.5)) + gs = fig.add_gridspec(1, 3, width_ratios=[0.7, 1.6, 1.6], wspace=0.22) + ax_b2 = fig.add_subplot(gs[0, 0]) + ax_m = fig.add_subplot(gs[0, 1]) + ax_c = fig.add_subplot(gs[0, 2]) + _plot_budget(ax_b2) _plot_memory(ax_m) _plot_comm(ax_c) - - fig.suptitle( - f"GQA per-PE memory + communication — LLaMA-3.1-70B " - f"single-KV-head group (C={_C}, P={_P}, {_N_LAYERS} layers, " - f"S_kv = 1 M, FP16)", - fontsize=12, y=0.985, - ) - fig.text( - 0.5, 0.94, - "Cases ordered by memory · Case 1 (40 GB, no sharding) → " - "Cases 2-3 (5 GB, 1-axis sharded) → Cases 4-6 (640 MB, 2-axis sharded) " - "· Case 6 ★ = lowest comm among memory-feasible cases", - ha="center", fontsize=9.5, color="#444", - ) - fig.tight_layout(rect=(0, 0, 1, 0.92)) + fig.tight_layout() out = _OUT_DIR / "gqa_4cases_summary.png" fig.savefig(out, dpi=150) plt.close(fig) print(f"wrote {out}") + # (c) 2-panel companion (analytical only). + fig2 = plt.figure(figsize=(18.0, 6.5)) + gs2 = fig2.add_gridspec(1, 2, width_ratios=[1.0, 1.0], wspace=0.18) + ax_m2 = fig2.add_subplot(gs2[0, 0]) + ax_c2 = fig2.add_subplot(gs2[0, 1]) + _plot_memory(ax_m2) + _plot_comm(ax_c2, mode="analytical") + fig2.tight_layout() + out2 = _OUT_DIR / "gqa_4cases_memory_comm_analytical.png" + fig2.savefig(out2, dpi=150) + plt.close(fig2) + print(f"wrote {out2}") + + # (d) 2-panel companion — analytical vs simulator-measured paired. + fig3 = plt.figure(figsize=(19.0, 6.5)) + gs3 = fig3.add_gridspec(1, 2, width_ratios=[1.0, 1.0], wspace=0.18) + ax_m3 = fig3.add_subplot(gs3[0, 0]) + ax_c3 = fig3.add_subplot(gs3[0, 1]) + _plot_memory(ax_m3) + _plot_comm(ax_c3, mode="paired") + fig3.tight_layout() + out3 = _OUT_DIR / "gqa_4cases_memory_comm_paired.png" + fig3.savefig(out3, dpi=150) + plt.close(fig3) + print(f"wrote {out3}") + # Paper-ready table to stdout. print() print(f" {'Case':<27} {'KV/tok·PE':>12} {'KV @ 1M':>11} " diff --git a/scripts/paper/paper_plot_gqa_kv_sharding_diagram.py b/scripts/paper/paper_plot_gqa_kv_sharding_diagram.py new file mode 100644 index 0000000..516b535 --- /dev/null +++ b/scripts/paper/paper_plot_gqa_kv_sharding_diagram.py @@ -0,0 +1,323 @@ +"""6-case KV-sharding tensor diagram (the slide-13 PNG export). + +Flat 2-D rectangles, one per sharding case, with: + Y axis = S_kv (vertical) — Cube-SP / PE-SP slice it + X axis = d_head (horizontal) — Cube-TP-d_head / PE-TP-d_head slice it + +Drops the batch axis entirely (decode: B = 1, T_q = 1). Same case set +and visual encoding as slide 13 of GQA_full_deck.pptx; matplotlib +renders it cleanly so the PNG sits next to the other GQA summary +artifacts in 1H_milestone_output/gqa/long_ctx/. +""" +from __future__ import annotations + +from pathlib import Path + +import matplotlib.patches as mpatches +import matplotlib.pyplot as plt + +_C = 8 +_P = 8 + +_GROUP_FILLS = [ + "#A5D8FF", "#B2F2BB", "#FFD8A8", "#FFC9C9", + "#D0BFFF", "#99E9F2", "#FCC2D7", "#FFEC99", +] + +_ACC = { + "red": "#E03131", + "orange": "#FD7E14", + "blue": "#1C7ED6", + "green": "#37B24D", +} + +# (label, accent, kv, comm, overflow, encoding-flags, axis-spec) +# y_split = 8 horizontal Y bands (Cube-SP on S_kv) +# x_split = 8 vertical X bands (Cube-TP-d_head) +# pe_y = 7 fine horizontal dividers within each Y band +# pe_x = 7 fine vertical dividers within each X band +# axes = small annotation under the chip naming the axes +# that the cube/PE actually shard, so the reader can +# parse Case 5 (where cube colour fills run X instead +# of Y, breaking the visual symmetry of the rest). +_CASES = [ + dict(label="Case 1\nCube-Repl / PE-repl", accent=_ACC["red"], + kv="40 GB", comm="1.2 MB", overflow=True, + y_split=False, x_split=False, pe_y=False, pe_x=False, + axes="Cube: replicated PE: replicated"), + dict(label="Case 2\nCube-SP / PE-repl", accent=_ACC["orange"], + kv="5 GB", comm="3.8 MB", overflow=True, + y_split=True, x_split=False, pe_y=False, pe_x=False, + axes="Cube → Y (S_kv) PE: replicated"), + dict(label="Case 3\nCube-Repl / PE-SP", accent=_ACC["orange"], + kv="5 GB", comm="3.8 MB", overflow=True, + y_split=False, x_split=False, pe_y=True, pe_x=False, + axes="Cube: replicated PE → Y (S_kv)"), + dict(label="Case 4\nCube-SP / PE-TP-d_head", accent=_ACC["blue"], + kv="640 MB", comm="166 MB", overflow=False, + y_split=True, x_split=False, pe_y=False, pe_x=True, + axes="Cube → Y (S_kv) PE → X (d_head)"), + dict(label="Case 5\nCube-TP-d_head / PE-SP", accent=_ACC["blue"], + kv="640 MB", comm="166 MB", overflow=False, + y_split=False, x_split=True, pe_y=True, pe_x=False, + axes="Cube → X (d_head) PE → Y (S_kv)"), + dict(label="Case 6 ★\nCube-SP / PE-SP", accent=_ACC["green"], + kv="640 MB", comm="6.2 MB", overflow=False, + y_split=True, x_split=False, pe_y=True, pe_x=False, + axes="Cube → Y (S_kv) PE → Y (S_kv)"), +] + +_OUT_DIR = ( + Path(__file__).resolve().parents[2] + / "src" / "kernbench" / "benches" + / "1H_milestone_output" / "gqa" / "long_ctx" +) + + +def _draw_panel(ax, cfg): + """Draw one case's 2-D KV-tensor rectangle into a panel ax.""" + ax.set_xlim(0, 1) + ax.set_ylim(1, 0) # Y points down (S_kv ↓) + ax.set_aspect("auto") + ax.set_xticks([]) + ax.set_yticks([]) + + cube_repl = not cfg["y_split"] and not cfg["x_split"] + pe_repl = not cfg["pe_y"] and not cfg["pe_x"] + + # Cube-level colour fill. + if cfg["y_split"] and not cfg["x_split"]: + # 8 horizontal Y bands. + for c in range(_C): + ax.add_patch(mpatches.Rectangle( + (0, c / _C), 1, 1 / _C, + facecolor=_GROUP_FILLS[c], edgecolor="black", linewidth=0.6)) + ax.text(0.04, c / _C + 0.5 / _C, f"C{c}", + ha="left", va="center", fontsize=8, + fontweight="bold", color="#333") + elif cfg["x_split"] and not cfg["y_split"]: + # 8 vertical X bands. + for c in range(_C): + ax.add_patch(mpatches.Rectangle( + (c / _C, 0), 1 / _C, 1, + facecolor=_GROUP_FILLS[c], edgecolor="black", linewidth=0.6)) + ax.text(c / _C + 0.5 / _C, 0.04, f"C{c}", + ha="center", va="top", fontsize=8, + fontweight="bold", color="#333") + else: + ax.add_patch(mpatches.Rectangle( + (0, 0), 1, 1, + facecolor="#F5F5F5", edgecolor="black", linewidth=0.8)) + ax.text(0.5, 0.5, "× 8 cubes\nfull KV", + ha="center", va="center", + fontsize=10, fontweight="bold", + fontstyle="italic", color="#666") + + # PE-level fine dividers — distinguished from cube boundaries by + # using a dashed style + slightly stronger contrast. This is what + # makes Case 5's PE-SP (horizontal lines across vertical cube + # bands) read as "different axis from the cubes" at a glance. + if cfg["pe_y"]: + outer = _C if cfg["y_split"] else 1 + band = 1 / outer + for o in range(outer): + for p in range(1, _P): + y = o * band + band * p / _P + ax.axhline(y, color="#222", linewidth=0.8, + linestyle=(0, (3, 2)), alpha=0.75) + if cfg["pe_x"]: + outer = _C if cfg["x_split"] else 1 + band = 1 / outer + for o in range(outer): + for p in range(1, _P): + x = o * band + band * p / _P + ax.axvline(x, color="#222", linewidth=0.8, + linestyle=(0, (3, 2)), alpha=0.75) + + # Heavy outline on top. + ax.add_patch(mpatches.Rectangle( + (0, 0), 1, 1, facecolor="none", + edgecolor="black", linewidth=1.2)) + + # Replication badges — small text-only badges in the corners of + # the rectangle, no ghost-card stacking (which mis-reads as a + # larger enclosing tensor). + badges: list[str] = [] + if cube_repl: + badges.append("× 8 cube copies") + if pe_repl and (cfg["y_split"] or cfg["x_split"]): + # Cube is sharded but PEs in each cube replicate that shard. + badges.append("× 8 PEs / cube replicate") + elif pe_repl and cube_repl: + # Both replicated — PE replication adds to the cube one. + badges.append("× 8 PEs / cube replicate") + if badges: + ax.text(0.98, 0.02, "\n".join(badges), + ha="right", va="top", fontsize=7, + fontweight="bold", color="#444", + fontstyle="italic", + bbox=dict(facecolor="white", edgecolor="#888", + boxstyle="round,pad=0.20", linewidth=0.5)) + + +def _make_table_png() -> Path: + """Slide-14 companion table: per-PE memory + comm for all 6 cases.""" + headers = ["Case", "Sharding", "KV / PE", "Fit", + "Comm/tok\n(analytical)", "Comm/tok\n(measured)", "Notes"] + rows = [ + ("Case 1", "Cube-Repl · PE-repl", "40 GB", "✗", + "1.2 MB", "1.25 MB", + "no sharding — full KV on every PE"), + ("Case 2", "Cube-SP · PE-repl", "5 GB", "✗", + "3.8 MB", "1.27 MB", + "cube-axis sharded only"), + ("Case 3", "Cube-Repl · PE-SP", "5 GB", "✗", + "3.8 MB", "1.39 MB", + "PE-axis sharded only"), + ("Case 4", "Cube-SP · PE-TP-d_head", "640 MB", "✓", + "166 MB", "162 MB", + "d_head split intra-cube — partial-score AR ∝ S_kv"), + ("Case 5", "Cube-TP-d_head · PE-SP", "640 MB", "✓", + "166 MB", "162 MB", + "d_head split inter-cube — partial-score AR on UCIe"), + ("Case 6 ★", "Cube-SP · PE-SP", "640 MB", "✓", + "6.2 MB", "1.41 MB", + "S_kv split on both axes — (m,ℓ,O) AR only, S_kv-indep."), + ] + accents = [_ACC["red"], _ACC["orange"], _ACC["orange"], + _ACC["blue"], _ACC["blue"], _ACC["green"]] + + fig, ax = plt.subplots(figsize=(20.0, 4.6)) + ax.set_axis_off() + + cell_data = [headers] + [list(r) for r in rows] + tbl = ax.table(cellText=cell_data, + colWidths=[0.07, 0.18, 0.08, 0.05, 0.12, 0.12, 0.38], + cellLoc="center", loc="center") + tbl.auto_set_font_size(False) + tbl.set_fontsize(11) + tbl.scale(1.0, 2.2) + + n_cols = len(headers) + n_rows = len(rows) + 1 # +1 header + # Header styling. + for ci in range(n_cols): + cell = tbl[(0, ci)] + cell.set_facecolor("#1F4E79") + cell.set_text_props(color="white", weight="bold") + cell.set_edgecolor("#1F4E79") + # Body styling. + for ri, row in enumerate(rows, start=1): + is_pareto = row[0].endswith("★") + row_fill = "#E8F5E9" if is_pareto else ( + "white" if ri % 2 == 1 else "#F5F5F7") + # Case-name cell uses accent. + case_cell = tbl[(ri, 0)] + case_cell.set_facecolor(accents[ri - 1]) + case_cell.set_text_props(color="white", weight="bold") + # Remaining cells. + for ci in range(1, n_cols): + cell = tbl[(ri, ci)] + cell.set_facecolor(row_fill) + txt_kwargs = {"weight": "bold" if is_pareto else "normal", + "color": "#333"} + if ci == 2: # KV / PE + txt_kwargs["color"] = ( + "#C62828" if row[3] == "✗" else "#2E7D32") + txt_kwargs["weight"] = "bold" + if ci == 3: # Fit + txt_kwargs["color"] = ( + "#C62828" if row[3] == "✗" else "#2E7D32") + txt_kwargs["weight"] = "bold" + cell.set_text_props(**txt_kwargs) + # Last-column (Notes) cells left-aligned for readability. + tbl[(ri, n_cols - 1)].get_text().set_ha("left") + + # Force left-align on the Notes header too. + tbl[(0, n_cols - 1)].get_text().set_ha("left") + + fig.suptitle( + "GQA decode KV-sharding — per-PE memory & communication " + "(LLaMA 70B GQA single KV-head group · S_kv = 1 M, FP16, 80 layers)", + fontsize=12, y=0.97, + ) + out = _OUT_DIR / "gqa_kv_sharding_6cases_table.png" + fig.savefig(out, dpi=150, bbox_inches="tight") + plt.close(fig) + print(f"wrote {out}") + return out + + +def main() -> Path: + _OUT_DIR.mkdir(parents=True, exist_ok=True) + n = len(_CASES) + fig = plt.figure(figsize=(20.0, 7.0)) + # Three rows per column: case chip · axis-spec annotation · rectangle. + gs = fig.add_gridspec(3, n, + height_ratios=[0.55, 0.32, 8.5], + hspace=0.05, wspace=0.20, + left=0.04, right=0.99, + top=0.93, bottom=0.06) + + for i, cfg in enumerate(_CASES): + # Top: case chip header. + ax_chip = fig.add_subplot(gs[0, i]) + ax_chip.set_xticks([]) + ax_chip.set_yticks([]) + for spine in ax_chip.spines.values(): + spine.set_visible(False) + ax_chip.add_patch(mpatches.Rectangle( + (0, 0), 1, 1, transform=ax_chip.transAxes, + facecolor=cfg["accent"], edgecolor=cfg["accent"])) + ax_chip.text(0.5, 0.5, cfg["label"], + ha="center", va="center", + fontsize=10, fontweight="bold", + color="white") + + # Middle: axis-spec annotation — names which axis the cube + # shards on and which axis the PE shards on (essential for + # parsing Case 5 where the cube colour fills run X instead + # of Y, breaking the visual symmetry of the rest). + ax_axes = fig.add_subplot(gs[1, i]) + ax_axes.set_xticks([]) + ax_axes.set_yticks([]) + for spine in ax_axes.spines.values(): + spine.set_visible(False) + ax_axes.add_patch(mpatches.Rectangle( + (0, 0), 1, 1, transform=ax_axes.transAxes, + facecolor="#F5F5F7", edgecolor="#CCCCCC", + linewidth=0.6)) + ax_axes.text(0.5, 0.5, cfg["axes"], + ha="center", va="center", + fontsize=8.5, fontweight="bold", + color="#1F4E79") + + # Bottom: the tensor rectangle. + ax = fig.add_subplot(gs[2, i]) + _draw_panel(ax, cfg) + ax.set_xlabel("X : d_head = 128 →", + fontsize=9, fontweight="bold", + fontstyle="italic", color="#1F4E79") + ax.set_ylabel("Y : S_kv = 1 M ↓", + fontsize=9, fontweight="bold", + fontstyle="italic", color="#1F4E79") + + fig.suptitle( + "GQA decode KV-tensor sharding — 6 cases · " + "LLaMA 70B GQA single KV-head group · " + "C = 8 cubes × P = 8 PEs · S_kv = 1 M, FP16, 80 layers", + fontsize=12, y=0.99, + ) + + out = _OUT_DIR / "gqa_kv_sharding_6cases_diagram.png" + fig.savefig(out, dpi=150, bbox_inches="tight") + plt.close(fig) + print(f"wrote {out}") + + # Companion table PNG (slide-14 export). + _make_table_png() + return out + + +if __name__ == "__main__": + main() diff --git a/scripts/verify_case4_dhead_tp.py b/scripts/verify_case4_dhead_tp.py new file mode 100644 index 0000000..2f5e6ca --- /dev/null +++ b/scripts/verify_case4_dhead_tp.py @@ -0,0 +1,148 @@ +"""Phase 1a smoke test for the Case 4 d_head-TP decode kernel. + +Runs `gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel` at small +S_kv (2K) to verify: + 1. it imports cleanly, + 2. it runs without error under the same harness as Case 6, + 3. captures op_log_summary (gemm/ipcq/dma counts), + 4. compares vs analytical predictions and vs Case 6 (same memory tier). + +Usage: + python scripts/verify_case4_dhead_tp.py +""" +from __future__ import annotations + +import os +import sys +from pathlib import Path + +from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import ( + gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel as _case6_kernel, +) +from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead import ( + gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel as _case4_kernel, +) +from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp import ( + gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel as _case5_kernel, +) +from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import ( + _ccl_cfg, _summarize_op_log, +) +from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip +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 + +# Small smoke-test params (S_kv=2K is enough to exercise tile-loop + AR). +_PARAMS = dict(C=8, P=8, T_q=1, S_kv=2_048, + d_head=128, h_q=8, h_kv=1) + + +def _bench_fn_case4(ctx): + p = _PARAMS + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp_full = DPPolicy(cube="replicate", pe="column_wise", + num_cubes=p["C"], num_pes=p["P"]) + dp_kv = DPPolicy(cube="row_wise", pe="column_wise", + num_cubes=p["C"], num_pes=p["P"]) + q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp_full, name="q_c4") + k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp_kv, name="k_c4") + v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp_kv, name="v_c4") + o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp_full, name="o_c4") + ctx.launch("case4_dhead_tp", _case4_kernel, + q, k, v, o, + p["T_q"], p["S_kv"], p["h_q"], p["h_kv"], + p["d_head"], p["C"], p["P"], + _auto_dim_remap=False) + + +def _bench_fn_case5(ctx): + p = _PARAMS + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp_q = DPPolicy(cube="column_wise", pe="replicate", + num_cubes=p["C"], num_pes=p["P"]) + dp_kv = DPPolicy(cube="column_wise", pe="row_wise", + num_cubes=p["C"], num_pes=p["P"]) + q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp_q, name="q_c5") + k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp_kv, name="k_c5") + v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp_kv, name="v_c5") + o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp_q, name="o_c5") + ctx.launch("case5_dhead_tp_inter", _case5_kernel, + q, k, v, o, + p["T_q"], p["S_kv"], p["h_q"], p["h_kv"], + p["d_head"], p["C"], p["P"], + _auto_dim_remap=False) + + +def _bench_fn_case6(ctx): + p = _PARAMS + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp_full = DPPolicy(cube="replicate", pe="replicate", + num_cubes=p["C"], num_pes=p["P"]) + dp_kv = DPPolicy(cube="row_wise", pe="row_wise", + num_cubes=p["C"], num_pes=p["P"]) + q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp_full, name="q_c6") + k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp_kv, name="k_c6") + v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]), + dtype="f16", dp=dp_kv, name="v_c6") + o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]), + dtype="f16", dp=dp_full, name="o_c6") + ctx.launch("case6_sp_sp", _case6_kernel, + q, k, v, o, + p["T_q"], p["S_kv"], p["h_q"], p["h_kv"], + p["d_head"], p["C"], p["P"], + _auto_dim_remap=False) + + +def main() -> int: + topology = os.environ.get("GQA_1H_TOPOLOGY", "topology.yaml") + topo = resolve_topology(topology) + + print(f"Smoke params: {_PARAMS}") + print() + + for label, bench_fn in ( + ("Case 4 (Cube-SP × PE-TP d_head)", _bench_fn_case4), + ("Case 5 (Cube-TP d_head × PE-SP)", _bench_fn_case5), + ("Case 6 (Cube-SP × PE-SP S_kv)", _bench_fn_case6), + ): + try: + res = run_bench( + topology=topo, bench_fn=bench_fn, + device=resolve_device(None), + engine_factory=lambda t, d: GraphEngine( + getattr(t, "topology_obj", t), enable_data=True, + ), + ) + except Exception as e: + print(f" {label:<42} FAIL: {type(e).__name__}: {e}") + continue + if not res.completion.ok: + print(f" {label:<42} ENGINE FAIL: {res.completion}") + continue + s = _summarize_op_log(res.engine.op_log) + lat = (res.engine.op_log[-1].t_end if res.engine.op_log else 0.0) + print(f" {label:<42} " + f"gemm={s['gemm_count']:>4} " + f"ipcq={s['ipcq_copy_count']:>4} " + f"dma_r={s['dma_read_count']:>4} " + f"dma_w={s['dma_write_count']:>3} " + f"latency={lat:.1f} ns " + f"(n_ops={len(res.engine.op_log)})") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_3cases_measured_comm.json b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_3cases_measured_comm.json new file mode 100644 index 0000000..e527aea --- /dev/null +++ b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_3cases_measured_comm.json @@ -0,0 +1,57 @@ +{ + "S_kv_measured": 8192, + "S_kv_headline": 1048576, + "n_layers": 80, + "num_pes": 64, + "wo_ffn_per_token_bytes": 1310720, + "cases": { + "1": { + "label": "Case 1 (Cube-Repl x PE-repl)", + "total_ipcq_bytes_one_layer": 0, + "per_pe_partial_score_bytes_one_layer": 0, + "per_pe_mlo_bytes_one_layer": 0, + "per_pe_attn_bytes_per_token_at_1M": 0, + "per_pe_total_bytes_per_token_at_1M": 1310720 + }, + "2": { + "label": "Case 2 (Cube-SP x PE-repl)", + "total_ipcq_bytes_one_layer": 14560, + "per_pe_partial_score_bytes_one_layer": 0, + "per_pe_mlo_bytes_one_layer": 227, + "per_pe_attn_bytes_per_token_at_1M": 18160, + "per_pe_total_bytes_per_token_at_1M": 1328880 + }, + "3": { + "label": "Case 3 (Cube-Repl x PE-SP)", + "total_ipcq_bytes_one_layer": 116480, + "per_pe_partial_score_bytes_one_layer": 0, + "per_pe_mlo_bytes_one_layer": 1820, + "per_pe_attn_bytes_per_token_at_1M": 145600, + "per_pe_total_bytes_per_token_at_1M": 1456320 + }, + "4": { + "label": "Case 4 (Cube-SP x PE-TP d_head)", + "total_ipcq_bytes_one_layer": 1851136, + "per_pe_partial_score_bytes_one_layer": 16384, + "per_pe_mlo_bytes_one_layer": 12540, + "per_pe_attn_bytes_per_token_at_1M": 168775360, + "per_pe_total_bytes_per_token_at_1M": 170086080 + }, + "5": { + "label": "Case 5 (Cube-TP d_head x PE-SP)", + "total_ipcq_bytes_one_layer": 1851136, + "per_pe_partial_score_bytes_one_layer": 16384, + "per_pe_mlo_bytes_one_layer": 12540, + "per_pe_attn_bytes_per_token_at_1M": 168775360, + "per_pe_total_bytes_per_token_at_1M": 170086080 + }, + "6": { + "label": "Case 6 (Cube-SP x PE-SP) [*]", + "total_ipcq_bytes_one_layer": 131040, + "per_pe_partial_score_bytes_one_layer": 0, + "per_pe_mlo_bytes_one_layer": 2047, + "per_pe_attn_bytes_per_token_at_1M": 163760, + "per_pe_total_bytes_per_token_at_1M": 1474480 + } + } +} \ No newline at end of file diff --git a/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead.py b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead.py new file mode 100644 index 0000000..bf15cf2 --- /dev/null +++ b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead.py @@ -0,0 +1,363 @@ +"""GQA decode kernel — Case 4 (Cube-SP × PE-TP on d_head). + +Per the per-PE memory/comm summary chart (paper_plot_gqa_4cases_summary.py): +this is the d_head-TP variant — KV is sharded 64-way (cube=row_wise on +S_kv, pe=column_wise on d_head), so per-PE KV memory matches Case 6 ★ +(640 MB / PE at 1M context across 80 layers), but the attention +algorithm needs an extra mid-layer partial-score AllReduce because +each PE only owns d_head/P dimensions of K/V. + +Algorithm (single layer, decode T_q=1, B=1): + 1. Each PE loads its own d_head/P slice of Q, K, V (Q replicated + in HBM; PE-load via column-wise DPPolicy carves out the slice). + 2. Per S_kv tile, each PE computes a PARTIAL score tile: + partial_score = Q_shard · K_shard^T shape: (G·T_q, S_local) + 3. AR-SUM these partial scores INTRA-CUBE (across the P PEs of a + cube) so every PE has the full Q · K^T for the cube's S_local + tokens. Online-softmax merge folds tiles into a running state. + 4. After all tiles, each PE has (m_i, ℓ_i, O_i_shard) for cube i, + where O_i_shard is the cube's partial attention output on this + PE's d_head/P slice. + 5. INTER-CUBE online-softmax merge of (m, ℓ, O_shard) across C + cubes (lrab center-root reduce). After this, every PE has the + full-sequence attention output for its d_head/P slice. + 6. Final normalise + store: the P PEs of the root cube each write + their (G·T_q, d_head/P) slice of the output. + +Tensor layout: + Q : (T_q, h_q · d_head) cube=replicate, pe=column_wise on d_head + → per-PE shape (T_q, h_q · d_head/P), byte-reshaped to + (G·T_q, d_head/P) for the per-PE GEMM input. + K : (S_kv, h_kv · d_head) cube=row_wise on S_kv, pe=column_wise + on d_head → per-PE shape (S_local, d_head/P), byte-reshaped to + (d_head/P, S_local) for the K^T-form GEMM. + V : same layout as K. + O : (T_q, h_q · d_head) cube=replicate, pe=column_wise on d_head + → each writer PE owns its d_head/P slice. + +Topology / SFR: + - Requires ``configure_sfr_intercube_multisip`` (provides disjoint + ``intra_*`` and ``E/W/N/S`` namespaces). + - Intra-CUBE PEs in a 2×4 grid; AR-SUM of partial scores uses a + reduce-to-PE0-then-broadcast pattern over intra_E/W + intra_N/S. + - Inter-CUBE lrab center-root reduce on (m, ℓ, O_shard) over the + 4×2 cube sub-mesh (root cube = 6). + +Naming note: this kernel exemplifies the chart's "PE-TP" label +meaning **d_head TP** (NOT the batch-axis TP of the original slide-17 +Case 1). The chart-side Case label is "Case 4 Cube-SP × PE-TP". + +ADR-0065 P5 scope note: this kernel is meant for op_log / latency +measurement (matches the milestone-1h-gqa sweep mode). Numeric +data-mode parity is a separate follow-up. +""" +from __future__ import annotations + + +TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width). + +# lrab geometry for the C=8 single-KV-head group (4×2 cube sub-mesh). +_SUB_W = 4 +_SUB_H = 2 +_ROOT_COL = _SUB_W // 2 # 2 +_ROOT_ROW = _SUB_H // 2 # 1 +_ROOT_CUBE = _ROOT_ROW * _SUB_W + _ROOT_COL # 6 + + +def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl): + """Online-softmax merge of two partial ``(m, ℓ, O)`` triples. + + O is the per-PE d_head-slice partial; merge is element-wise so it + composes naturally even with O being a d_head shard. + """ + m_new = tl.maximum(m_local, m_other) + scale_old = tl.exp(m_local - m_new) + scale_new = tl.exp(m_other - m_new) + l_new = l_local * scale_old + l_other * scale_new + O_new = O_local * scale_old + O_other * scale_new + return m_new, l_new, O_new + + +def _intra_cube_allreduce_sum(tensor, *, tl, + pe_col: int, pe_row: int, + pe_cols_used: int, pe_rows_used: int): + """Sum-AllReduce a tensor across the 8 PEs of one cube (2×4 grid). + + Reduce-to-PE0 (row chain west + col bridge north) followed by + broadcast-from-PE0 (col bridge south + row chain east). Each phase + operates in place on ``tensor`` via ``tl.copy_to``. + + Used here to combine PARTIAL Q·K^T score tiles across the P=8 PEs + of a cube — each PE has only d_head/P dims, so each contributes a + partial dot product; the full score = elementwise sum across PEs. + """ + # ── Phase 1a: row reduce (west) — accumulate at pe_col=0 ──────── + if pe_cols_used > 1: + if pe_col == pe_cols_used - 1: + tl.send(dir="intra_W", src=tensor) + elif 0 < pe_col < pe_cols_used - 1: + with tl.scratch_scope(): + other = tl.recv(dir="intra_E", shape=tensor.shape, dtype="f16") + tl.copy_to(tensor, tensor + other) + tl.send(dir="intra_W", src=tensor) + elif pe_col == 0: + with tl.scratch_scope(): + other = tl.recv(dir="intra_E", shape=tensor.shape, dtype="f16") + tl.copy_to(tensor, tensor + other) + + # ── Phase 1b: col reduce on pe_col=0 (north) — accumulate at row=0 ─ + if pe_col == 0 and pe_rows_used > 1: + if pe_row == pe_rows_used - 1: + tl.send(dir="intra_N", src=tensor) + elif pe_row == 0: + with tl.scratch_scope(): + other = tl.recv(dir="intra_S", shape=tensor.shape, dtype="f16") + tl.copy_to(tensor, tensor + other) + + # ── Phase 2a: col broadcast (south) — row=0 → row=1 on pe_col=0 ─ + if pe_col == 0 and pe_rows_used > 1: + if pe_row == 0: + tl.send(dir="intra_S", src=tensor) + elif pe_row > 0: + with tl.scratch_scope(): + received = tl.recv(dir="intra_N", shape=tensor.shape, dtype="f16") + tl.copy_to(tensor, received) + + # ── Phase 2b: row broadcast (east) — pe_col=0 → all on each row ─ + if pe_cols_used > 1: + if pe_col == 0: + tl.send(dir="intra_E", src=tensor) + elif 0 < pe_col < pe_cols_used - 1: + with tl.scratch_scope(): + received = tl.recv(dir="intra_W", shape=tensor.shape, dtype="f16") + tl.copy_to(tensor, received) + tl.send(dir="intra_E", src=tensor) + elif pe_col == pe_cols_used - 1: + with tl.scratch_scope(): + received = tl.recv(dir="intra_W", shape=tensor.shape, dtype="f16") + tl.copy_to(tensor, received) + + +def gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel( + q_ptr: int, + k_ptr: int, + v_ptr: int, + o_ptr: int, + T_q: int, + S_kv: int, + h_q: int, + h_kv: int, + d_head: int, + C: int, + P: int, + *, + tl, +) -> None: + """Case-4 decode: Cube-SP × PE-TP-on-d_head, lrab center-root at cube 6.""" + G = h_q // h_kv + if d_head % P != 0: + raise ValueError( + f"Case 4 requires d_head={d_head} divisible by P={P}" + ) + if S_kv % C != 0: + raise ValueError( + f"Case 4 requires S_kv={S_kv} divisible by C={C}" + ) + + d_head_local = d_head // P # 128 / 8 = 16 in the headline setup + S_local = S_kv // C + pe_id = tl.program_id(axis=0) + cube_id = tl.program_id(axis=1) + + PE_GRID_COLS = 4 + pe_col = pe_id % PE_GRID_COLS + pe_row = pe_id // PE_GRID_COLS + pe_cols_used = min(PE_GRID_COLS, P) + pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS + + KV_ROW_BYTES = d_head_local * 2 # bytes per K/V row at this PE's d_head slice + + # ── Load this PE's d_head slice of Q (byte-conserving reshape) ── + # Q in HBM is (T_q, h_q · d_head); the column_wise pe shard gives + # each PE (T_q, h_q · d_head/P). Reshape to (G·T_q, d_head/P) for the + # GEMM: byte-conserving since G = h_q (h_kv = 1). + Q = tl.load(q_ptr, shape=(G * T_q, d_head_local), dtype="f16") + + # ── Local attention (S_kv-axis tile sweep) with intra-cube AR sum ─ + n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV + + # Tile 0: establish persistent (m, ℓ, O_shard). O_shard is the + # per-PE d_head/P slice of the partial attention output. + tile_s0 = min(TILE_S_KV, S_local) + K_T = tl.load(k_ptr, shape=(d_head_local, tile_s0), dtype="f16") + V = tl.load(v_ptr, shape=(tile_s0, d_head_local), dtype="f16") + partial_scores0 = tl.dot(Q, K_T) + # AR-SUM partial scores intra-cube → every PE in the cube has the + # full Q · K^T_tile0 scores. + _intra_cube_allreduce_sum( + partial_scores0, tl=tl, + pe_col=pe_col, pe_row=pe_row, + pe_cols_used=pe_cols_used, pe_rows_used=pe_rows_used, + ) + m_local = tl.max(partial_scores0, axis=-1) + exp_scores = tl.exp(partial_scores0 - m_local) + l_local = tl.sum(exp_scores, axis=-1) + # Partial output sharded on d_head: (G·T_q, d_head/P) per PE. + O_local = tl.dot(exp_scores, V) + + for tile_idx in range(1, n_tiles): + tile_start = tile_idx * TILE_S_KV + tile_s = min(TILE_S_KV, S_local - tile_start) + with tl.scratch_scope(): + K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES, + shape=(d_head_local, tile_s), dtype="f16") + V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES, + shape=(tile_s, d_head_local), dtype="f16") + partial_scores_t = tl.dot(Q, K_T_t) + _intra_cube_allreduce_sum( + partial_scores_t, tl=tl, + pe_col=pe_col, pe_row=pe_row, + pe_cols_used=pe_cols_used, pe_rows_used=pe_rows_used, + ) + m_tile = tl.max(partial_scores_t, axis=-1) + exp_scores_t = tl.exp(partial_scores_t - m_tile) + l_tile = tl.sum(exp_scores_t, axis=-1) + O_tile = tl.dot(exp_scores_t, V_t) + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + + # ── Inter-CUBE lrab-adapted center-root reduce on (m, ℓ, O_shard) ─ + # Every PE participates (each contributes its own d_head/P slice + # of O_shard); the merge is element-wise so d_head-sharding is + # preserved through the AR. + row = cube_id // _SUB_W + col = cube_id % _SUB_W + + # Phase 1: row reduce — converge at col == _ROOT_COL. + if col == 0: + tl.send(dir="E", src=m_local) + tl.send(dir="E", src=l_local) + tl.send(dir="E", src=O_local) + elif 0 < col < _ROOT_COL: + with tl.scratch_scope(): + m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + tl.send(dir="E", src=m_local) + tl.send(dir="E", src=l_local) + tl.send(dir="E", src=O_local) + elif col == _ROOT_COL: + with tl.scratch_scope(): + m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + with tl.scratch_scope(): + m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + elif _ROOT_COL < col < _SUB_W - 1: + with tl.scratch_scope(): + m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + tl.send(dir="W", src=m_local) + tl.send(dir="W", src=l_local) + tl.send(dir="W", src=O_local) + elif col == _SUB_W - 1: + tl.send(dir="W", src=m_local) + tl.send(dir="W", src=l_local) + tl.send(dir="W", src=O_local) + + # Phase 2: col reduce on col == _ROOT_COL — converge at row == _ROOT_ROW. + if col == _ROOT_COL: + if row == 0: + tl.send(dir="S", src=m_local) + tl.send(dir="S", src=l_local) + tl.send(dir="S", src=O_local) + elif 0 < row < _ROOT_ROW: + with tl.scratch_scope(): + m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + tl.send(dir="S", src=m_local) + tl.send(dir="S", src=l_local) + tl.send(dir="S", src=O_local) + elif row == _ROOT_ROW: + with tl.scratch_scope(): + m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + if _SUB_H - 1 > _ROOT_ROW: + with tl.scratch_scope(): + m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + elif _ROOT_ROW < row < _SUB_H - 1: + with tl.scratch_scope(): + m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + tl.send(dir="N", src=m_local) + tl.send(dir="N", src=l_local) + tl.send(dir="N", src=O_local) + elif row == _SUB_H - 1 and _SUB_H - 1 > _ROOT_ROW: + tl.send(dir="N", src=m_local) + tl.send(dir="N", src=l_local) + tl.send(dir="N", src=O_local) + + # ── Final normalise + store (root cube: each P-PE writes its d_head slice) ─ + if cube_id == _ROOT_CUBE: + O_final = O_local / l_local + tl.store(o_ptr, O_final) diff --git a/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp.py b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp.py new file mode 100644 index 0000000..8e648e2 --- /dev/null +++ b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp.py @@ -0,0 +1,265 @@ +"""GQA decode kernel — Case 5 (Cube-TP on d_head × PE-SP on S_kv). + +Per the per-PE memory/comm summary chart (paper_plot_gqa_4cases_summary.py): +this is the second d_head-TP variant — KV is sharded 64-way (cube= +column_wise on d_head, pe=row_wise on S_kv). Per-PE KV memory matches +Cases 4 and 6 ★ (640 MB / PE at 1M context across 80 layers), but the +attention algorithm pays the partial-score AR over INTER-cube fabric +(UCIe, slower) instead of Case 4's intra-cube fabric (NoC). + +Symmetric mirror of Case 4 — intra and inter-cube AR roles SWAPPED: + + Step Case 4 (Cube-SP × PE-TP/d_head) Case 5 (THIS — Cube-TP/d_head × PE-SP) + ─────────────── ──────────────────────────────── ────────────────────────────────────── + partial Q·K^T AR-SUM intra-cube (NoC) AR-SUM inter-cube (UCIe) + online-softmax inter-cube lrab on (m, ℓ, O) intra-cube merge on (m, ℓ, O) + (every PE participates; each (online-softmax reduce-to-PE0) + PE owns its d_head/P slice) + writer root cube's P PEs (each writes each cube's PE 0 writes its + its d_head/P slice) d_head/C slice (C writers total) + +Algorithm: + 1. Each PE loads its (d_head/C × S_kv/P) slice of K and V; full-d_head + Q from HBM but only its d_head/C slice via column_wise DPPolicy. + 2. Per S_kv tile (per PE — already a strict subset of the cube's + S_kv), compute partial_score = Q_shard · K_shard^T. Shape: + (G·T_q, S_kv/P). AR-SUM these PARTIAL scores INTER-CUBE across + the C cubes (same PE rank, different d_head slices) so every PE + in cube c, rank p has the full Q · K^T_p scores for THIS PE's + S_kv/P token range. + 3. Local softmax + P · V_shard → partial output (G·T_q, d_head/C) + for this PE's S_kv/P slice and this cube's d_head/C slice. + 4. INTRA-CUBE online-softmax merge of (m, ℓ, O) reduce-to-PE0 + across the cube's P PEs (different S_kv slices, same d_head slice). + 5. PE 0 of each cube normalises and writes the cube's d_head/C slice. + +Tensor layout: + Q : (T_q, h_q · d_head) cube=column_wise on d_head, pe=replicate. + Per PE shape (T_q, h_q · d_head/C); reshape to (G·T_q, d_head/C). + K : (S_kv, h_kv · d_head) cube=column_wise on d_head, pe=row_wise on S_kv. + Per PE shape (S_kv/P, d_head/C); reshape to (d_head/C, S_kv/P) for K^T. + V : same layout as K. + O : (T_q, h_q · d_head) cube=column_wise on d_head, pe=replicate. + Per-cube writer (PE 0) owns its d_head/C slice. + +Topology / SFR: + - Requires ``configure_sfr_intercube_multisip`` (disjoint intra_* + and E/W/N/S namespaces). + - INTRA-CUBE PEs arranged 2×4 (used for online-softmax merge). + - INTER-CUBE 4×2 cube sub-mesh (used for partial-score AR sum). + +ADR-0065 P5 scope note: op_log / latency measurement only; numeric +data-mode parity is a separate follow-up (matches the milestone-1h-gqa +sweep mode). +""" +from __future__ import annotations + + +TILE_S_KV = 1024 # per-tile S_kv width within a PE's own S_kv/P shard. + +# Cube sub-mesh geometry (4×2). For Case 5 we pick cube 0 as the AR-sum +# root (simpler than lrab; partial-score AR doesn't need geometric centering +# because the reduction is element-wise sum, not online-softmax merge). +_SUB_W = 4 +_SUB_H = 2 +_AR_ROOT_CUBE = 0 + + +def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl): + """Online-softmax merge of two partial ``(m, ℓ, O)`` triples.""" + m_new = tl.maximum(m_local, m_other) + scale_old = tl.exp(m_local - m_new) + scale_new = tl.exp(m_other - m_new) + l_new = l_local * scale_old + l_other * scale_new + O_new = O_local * scale_old + O_other * scale_new + return m_new, l_new, O_new + + +def _inter_cube_allreduce_sum(tensor, *, tl, cube_id: int, + sub_w: int = _SUB_W, sub_h: int = _SUB_H): + """Sum-AllReduce a tensor INTER-CUBE across the 4×2 cube sub-mesh. + + Reduce-to-cube-0 (row chain west + col bridge north) followed by + broadcast-from-cube-0 (col bridge south + row chain east). The + payload is the partial-score tile — each cube has a different + d_head slice's contribution; element-wise sum recovers the full + Q·K^T scores for this PE rank's S_kv/P tokens. + + All P PEs of every cube participate independently (per-PE-rank + inter-cube fabric lanes). + """ + row = cube_id // sub_w + col = cube_id % sub_w + + # ── Phase 1a: row reduce (west) — accumulate at col=0 ──────────── + if sub_w > 1: + if col == sub_w - 1: + tl.send(dir="W", src=tensor) + elif 0 < col < sub_w - 1: + with tl.scratch_scope(): + other = tl.recv(dir="E", shape=tensor.shape, dtype="f16") + tl.copy_to(tensor, tensor + other) + tl.send(dir="W", src=tensor) + elif col == 0: + with tl.scratch_scope(): + other = tl.recv(dir="E", shape=tensor.shape, dtype="f16") + tl.copy_to(tensor, tensor + other) + + # ── Phase 1b: col reduce on col=0 (north) — accumulate at row=0 ── + if col == 0 and sub_h > 1: + if row == sub_h - 1: + tl.send(dir="N", src=tensor) + elif row == 0: + with tl.scratch_scope(): + other = tl.recv(dir="S", shape=tensor.shape, dtype="f16") + tl.copy_to(tensor, tensor + other) + + # ── Phase 2a: col broadcast (south) — row=0 → row=1 on col=0 ───── + if col == 0 and sub_h > 1: + if row == 0: + tl.send(dir="S", src=tensor) + elif row > 0: + with tl.scratch_scope(): + received = tl.recv(dir="N", shape=tensor.shape, dtype="f16") + tl.copy_to(tensor, received) + + # ── Phase 2b: row broadcast (east) — col=0 → all on each row ───── + if sub_w > 1: + if col == 0: + tl.send(dir="E", src=tensor) + elif 0 < col < sub_w - 1: + with tl.scratch_scope(): + received = tl.recv(dir="W", shape=tensor.shape, dtype="f16") + tl.copy_to(tensor, received) + tl.send(dir="E", src=tensor) + elif col == sub_w - 1: + with tl.scratch_scope(): + received = tl.recv(dir="W", shape=tensor.shape, dtype="f16") + tl.copy_to(tensor, received) + + +def gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel( + q_ptr: int, + k_ptr: int, + v_ptr: int, + o_ptr: int, + T_q: int, + S_kv: int, + h_q: int, + h_kv: int, + d_head: int, + C: int, + P: int, + *, + tl, +) -> None: + """Case-5 decode: Cube-TP-on-d_head × PE-SP-on-S_kv.""" + G = h_q // h_kv + if d_head % C != 0: + raise ValueError( + f"Case 5 requires d_head={d_head} divisible by C={C}" + ) + if S_kv % P != 0: + raise ValueError( + f"Case 5 requires S_kv={S_kv} divisible by P={P}" + ) + + d_head_local = d_head // C # 128 / 8 = 16 in the headline setup + S_local_pe = S_kv // P # each PE owns S_kv/P tokens + pe_id = tl.program_id(axis=0) + cube_id = tl.program_id(axis=1) + + PE_GRID_COLS = 4 + pe_col = pe_id % PE_GRID_COLS + pe_row = pe_id // PE_GRID_COLS + pe_cols_used = min(PE_GRID_COLS, P) + pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS + + KV_ROW_BYTES = d_head_local * 2 # bytes per K/V row at this PE's d_head slice + + # ── Load this PE's d_head slice of Q ──────────────────────────── + # Q in HBM is (T_q, h_q · d_head); column_wise cube shard gives each + # cube (T_q, h_q · d_head/C); pe=replicate within cube → every PE + # in the cube has the same data. Reshape to (G·T_q, d_head/C). + Q = tl.load(q_ptr, shape=(G * T_q, d_head_local), dtype="f16") + + # ── Local attention (tile sweep on this PE's S_kv/P shard) ────── + n_tiles = (S_local_pe + TILE_S_KV - 1) // TILE_S_KV + + # Tile 0: bootstrap (m, ℓ, O_shard). The partial Q·K^T is across + # d_head — must AR inter-cube before softmax. + tile_s0 = min(TILE_S_KV, S_local_pe) + K_T = tl.load(k_ptr, shape=(d_head_local, tile_s0), dtype="f16") + V = tl.load(v_ptr, shape=(tile_s0, d_head_local), dtype="f16") + partial_scores0 = tl.dot(Q, K_T) + _inter_cube_allreduce_sum(partial_scores0, tl=tl, cube_id=cube_id) + m_local = tl.max(partial_scores0, axis=-1) + exp_scores = tl.exp(partial_scores0 - m_local) + l_local = tl.sum(exp_scores, axis=-1) + # Partial output (G·T_q, d_head/C) — sharded on d_head only; PE + # still holds the full d_head/C slice for its S_kv/P tokens. + O_local = tl.dot(exp_scores, V) + + for tile_idx in range(1, n_tiles): + tile_start = tile_idx * TILE_S_KV + tile_s = min(TILE_S_KV, S_local_pe - tile_start) + with tl.scratch_scope(): + K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES, + shape=(d_head_local, tile_s), dtype="f16") + V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES, + shape=(tile_s, d_head_local), dtype="f16") + partial_scores_t = tl.dot(Q, K_T_t) + _inter_cube_allreduce_sum(partial_scores_t, tl=tl, + cube_id=cube_id) + m_tile = tl.max(partial_scores_t, axis=-1) + exp_scores_t = tl.exp(partial_scores_t - m_tile) + l_tile = tl.sum(exp_scores_t, axis=-1) + O_tile = tl.dot(exp_scores_t, V_t) + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + + # ── INTRA-CUBE online-softmax reduce-to-PE0 (different S_kv slices) ─ + # Each cube ends with (m, ℓ, O) for its d_head/C slice at PE 0. + if pe_cols_used > 1: + if pe_col < pe_cols_used - 1: + with tl.scratch_scope(): + m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + if pe_col > 0: + tl.send(dir="intra_W", src=m_local) + tl.send(dir="intra_W", src=l_local) + tl.send(dir="intra_W", src=O_local) + + if pe_col == 0 and pe_rows_used > 1: + if pe_row < pe_rows_used - 1: + with tl.scratch_scope(): + m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + if pe_row > 0: + tl.send(dir="intra_N", src=m_local) + tl.send(dir="intra_N", src=l_local) + tl.send(dir="intra_N", src=O_local) + + # ── Final normalise + store (one writer per cube: PE 0; each + # cube owns a disjoint d_head/C slice of the output) ───────── + if pe_id == 0: + O_final = O_local / l_local + tl.store(o_ptr, O_final) diff --git a/src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_4cases.py b/src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_4cases.py index 7e93387..0442b3f 100644 --- a/src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_4cases.py +++ b/src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_4cases.py @@ -38,6 +38,12 @@ from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_ from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp import ( gqa_attention_decode_long_ctx_cube_sp_pe_tp_kernel, ) +from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead import ( + gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel, +) +from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp import ( + gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel, +) from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import ( _ccl_cfg, _summarize_op_log, @@ -58,10 +64,12 @@ _SWEEP_JSON = _OUTPUT_DIR / "sweep_decode.json" _PANELS = ( - "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp", # Case 4 ★ optimal - "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp", # Case 2 (no comm; 8× memory) - "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp", # Case 3 (intra-CUBE AR only; 8× memory) - "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp", # Case 1 (inter-CUBE lrab only; PE-TP B=1 waste) + "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp", # Case 6 ★ optimal (S_kv-SP both axes; was "Case 4" in slide-17 numbering) + "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp", # Case 1 (no sharding; 40 GB / PE @ 1M) + "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp", # Case 3 (intra-cube AR on (m,ℓ,O); 5 GB / PE) + "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp", # Case 2 (inter-cube AR on (m,ℓ,O); 5 GB / PE) + "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp_dhead", # Case 4 (Cube-SP × PE-TP-on-d_head — d_head-TP partial-score AR intra-cube) + "single_kv_group_decode_long_ctx_gqa_cube_tp_dhead_pe_sp", # Case 5 (Cube-TP-on-d_head × PE-SP — d_head-TP partial-score AR inter-cube) ) # Each entry: (kind, panel-specific params). @@ -102,6 +110,23 @@ _PANEL_DISPATCH: dict[str, tuple[str, dict]] = { "T_q": 1, "S_kv": 8_192, "d_head": 128, "h_q": 8, "h_kv": 1, }), + "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp_dhead": ("decode_long_ctx_cube_sp_pe_tp_dhead", { + # Case 4 (new numbering): Cube-SP on S_kv × PE-TP on d_head. + # Partial-score AR happens intra-cube (NoC); (m,ℓ,O) merge happens + # inter-cube (UCIe). Memory-equal to Case 6; comm dominated by + # partial-score AR which scales with S_kv. + "C": 8, "P": 8, + "T_q": 1, "S_kv": 8_192, + "d_head": 128, "h_q": 8, "h_kv": 1, + }), + "single_kv_group_decode_long_ctx_gqa_cube_tp_dhead_pe_sp": ("decode_long_ctx_cube_tp_dhead_pe_sp", { + # Case 5 (new numbering): Cube-TP on d_head × PE-SP on S_kv. + # Mirror of Case 4: partial-score AR happens INTER-cube (UCIe); + # (m,ℓ,O) merge happens INTRA-cube (NoC). Memory-equal to Case 6. + "C": 8, "P": 8, + "T_q": 1, "S_kv": 8_192, + "d_head": 128, "h_q": 8, "h_kv": 1, + }), } @@ -238,6 +263,74 @@ def _run_decode_panel_long_ctx_cube_sp_pe_sp( ) +def _run_decode_panel_long_ctx_cube_sp_pe_tp_dhead( + ctx, *, panel: str, C: int, P: int, + T_q: int, S_kv: int, + d_head: int, h_q: int, h_kv: int, +) -> None: + """Case 4 (new) runner: Cube-SP × PE-TP-on-d_head. + + Q is cube-replicated, PE-sharded column-wise (d_head/P per PE). + K, V are cube=row_wise on S_kv, pe=column_wise on d_head — each PE + owns (S_kv/C, d_head/P). The kernel does an intra-cube AR-SUM of + partial scores then an inter-cube online-softmax merge of + (m, ℓ, O_shard). + """ + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp_q = DPPolicy(cube="replicate", pe="column_wise", + num_cubes=C, num_pes=P) + dp_kv = DPPolicy(cube="row_wise", pe="column_wise", + num_cubes=C, num_pes=P) + q = ctx.zeros((T_q, h_q * d_head), + dtype="f16", dp=dp_q, 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_q, name=f"{panel}_o") + ctx.launch( + panel, gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_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_long_ctx_cube_tp_dhead_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 5 (new) runner: Cube-TP-on-d_head × PE-SP-on-S_kv. + + Q is cube-sharded column_wise (d_head/C per cube), pe-replicated + within cube. K, V are cube=column_wise on d_head, pe=row_wise on + S_kv — each PE owns (S_kv/P, d_head/C). The kernel does an + inter-cube AR-SUM of partial scores then an intra-cube online- + softmax merge of (m, ℓ, O) reduce-to-PE0. + """ + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp_q = DPPolicy(cube="column_wise", pe="replicate", + num_cubes=C, num_pes=P) + dp_kv = DPPolicy(cube="column_wise", pe="row_wise", + num_cubes=C, num_pes=P) + q = ctx.zeros((T_q, h_q * d_head), + dtype="f16", dp=dp_q, 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_q, name=f"{panel}_o") + ctx.launch( + panel, gqa_attention_decode_long_ctx_cube_tp_dhead_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] @@ -250,6 +343,10 @@ def _make_bench_fn(panel: str): _run_decode_panel_long_ctx_cube_repl_pe_sp(ctx, panel=panel, **params) elif kind == "decode_long_ctx_cube_sp_pe_tp": _run_decode_panel_long_ctx_cube_sp_pe_tp(ctx, panel=panel, **params) + elif kind == "decode_long_ctx_cube_sp_pe_tp_dhead": + _run_decode_panel_long_ctx_cube_sp_pe_tp_dhead(ctx, panel=panel, **params) + elif kind == "decode_long_ctx_cube_tp_dhead_pe_sp": + _run_decode_panel_long_ctx_cube_tp_dhead_pe_sp(ctx, panel=panel, **params) else: raise RuntimeError( f"milestone-gqa-decode-long-ctx-4cases panel {panel!r} has "