diff --git a/docs/report/1H-codesign-paper/figures/gqa_decode_models.png b/docs/report/1H-codesign-paper/figures/gqa_decode_models.png new file mode 100644 index 0000000..5e767fb Binary files /dev/null and b/docs/report/1H-codesign-paper/figures/gqa_decode_models.png differ diff --git a/scripts/paper/paper_plot_gqa_decode_models.py b/scripts/paper/paper_plot_gqa_decode_models.py new file mode 100644 index 0000000..89f30a3 --- /dev/null +++ b/scripts/paper/paper_plot_gqa_decode_models.py @@ -0,0 +1,192 @@ +"""Bar plot for the multi-model GQA composite bench. + +Topology per model: cubes per KV group C = h_q, P = 8 PEs / cube. + +Reads ``sweep_decode_models.json`` (produced by +``gqa_decode_long_ctx_models.py``) and writes a two-panel PNG +comparing end-to-end latency and PE_CPU dispatch across six GQA models +where each model uses a topology sized to match its G value +(cubes per KV group = h_q per KV group; PEs per cube = 8): + + Left panel — end-to-end decode latency (µs), one bar per model, sorted + by G. Bar labels show C, N, and latency. + Right panel — PE_CPU command count per model. + +Run (after the bench): + python scripts/paper/paper_plot_gqa_decode_models.py +""" +from __future__ import annotations + +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +import numpy as np # noqa: E402 + +_REPO_ROOT = Path(__file__).resolve().parents[2] +_FIG_DIR = ( + _REPO_ROOT / "src" / "kernbench" / "benches" + / "1H_milestone_output" / "gqa" / "long_ctx" +) +_SWEEP_JSON = _FIG_DIR / "sweep_decode_models.json" +_PAPER_FIG_DIR = ( + _REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures" +) + +# Model display labels +_MODEL_LABELS = { + "gemma2-27b": "Gemma 2 27B\n(G=2)", + "llama3-8b": "LLaMA-3 8B\n(G=4)", + "qwen2.5-7b": "Qwen 2.5 7B\n(G=7)", + "llama3-70b": "LLaMA-3 70B\n(G=8)", + "qwen2.5-72b": "Qwen 2.5 72B\n(G=8)", + "command-r-plus": "Command R+\n(G=12)", +} + +# Bar color per model family +_FAMILY_COLOUR = { + "Google": "#4f8a4f", + "Meta": "#3b6ea5", + "Alibaba": "#c86432", + "Cohere": "#8b5a2b", +} + + +def main() -> None: + sweep = json.loads(_SWEEP_JSON.read_text()) + rows_by_model = {r["model"]: r for r in sweep["rows"]} + # Sort models by G (ascending) so the trend is left-to-right. + ordered = sorted(_MODEL_LABELS.keys(), key=lambda m: rows_by_model[m]["G"]) + + fig, (ax_lat, ax_cmd, ax_brk) = plt.subplots(1, 3, figsize=(18.5, 5.0)) + xs = np.arange(len(ordered)) + + # ── Left: latency ───────────────────────────────────────────── + latencies = [rows_by_model[m]["latency_ns"] / 1e3 for m in ordered] + colours = [_FAMILY_COLOUR[rows_by_model[m]["family"]] for m in ordered] + bars = ax_lat.bar(xs, latencies, 0.62, + color=colours, edgecolor="#222", linewidth=0.7) + for i, m in enumerate(ordered): + r = rows_by_model[m] + # Latency label above bar + ax_lat.text(i, latencies[i] + max(latencies) * 0.015, + f"{latencies[i]:.1f} µs", + ha="center", va="bottom", fontsize=9, + fontweight="bold", color="#222") + # Topology label inside bar + ax_lat.text(i, latencies[i] / 2, + f"cubes = {r['C']}", + ha="center", va="center", fontsize=9, color="white") + + ax_lat.set_xticks(xs) + ax_lat.set_xticklabels([_MODEL_LABELS[m] for m in ordered], fontsize=9) + ax_lat.set_ylabel("end-to-end decode latency (µs)") + ax_lat.set_title( + f"Composite Case-6 decode latency at $S_{{kv}}=128$K\n" + "topology per model: cubes per KV group = h_q of the KV group", + fontsize=11, + ) + ax_lat.grid(True, axis="y", ls=":", alpha=0.5) + ax_lat.set_ylim(0, max(latencies) * 1.18) + + # Legend (family) + from matplotlib.patches import Patch + seen_families = [] + handles = [] + for m in ordered: + fam = rows_by_model[m]["family"] + if fam not in seen_families: + seen_families.append(fam) + handles.append(Patch(facecolor=_FAMILY_COLOUR[fam], + edgecolor="#222", label=fam)) + ax_lat.legend(handles=handles, loc="upper left", fontsize=9, + title="family") + + # ── Right: PE_CPU dispatch ───────────────────────────────────── + cmds = [rows_by_model[m]["pe_cpu_cmd_count"] for m in ordered] + bars2 = ax_cmd.bar(xs, cmds, 0.62, + color=colours, edgecolor="#222", linewidth=0.7) + for i, m in enumerate(ordered): + ax_cmd.text(i, cmds[i] + max(cmds) * 0.02, + f"{cmds[i]}", + ha="center", va="bottom", fontsize=9, + fontweight="bold", color="#222") + r = rows_by_model[m] + ax_cmd.text(i, cmds[i] / 2, + f"cubes = {r['C']}", + ha="center", va="center", fontsize=9, color="white") + + ax_cmd.set_xticks(xs) + ax_cmd.set_xticklabels([_MODEL_LABELS[m] for m in ordered], fontsize=9) + ax_cmd.set_ylabel("PE_CPU commands issued") + ax_cmd.set_title( + "PE_CPU dispatch (composite, per KV group)\n" + "growing sub-mesh width or height → more reduce hops", + fontsize=11, + ) + ax_cmd.grid(True, axis="y", ls=":", alpha=0.5) + ax_cmd.set_ylim(0, max(cmds) * 1.18) + + # ── Third: matmul vs comm (op-kind occupancy) ───────────────────── + # Only the two kinds that carry useful work — matmul (GEMM + MATH + # occupancy summed across all engines) and comm (DMA occupancy). + # These are op-log sums across components, not critical-path + # attribution; parallelism means they don't sum to wall-clock + # latency. Bars are absolute µs so the reduce-cost growth with C + # is visible. + matmul_us = [rows_by_model[m].get("matmul_ns", 0.0) / 1e3 for m in ordered] + comm_us = [rows_by_model[m].get("comm_ns", 0.0) / 1e3 for m in ordered] + max_v = max(max(matmul_us), max(comm_us)) + width = 0.35 + + ax_brk.bar(xs - width / 2, matmul_us, width, + color="#3b6ea5", edgecolor="#222", linewidth=0.6, + label="matmul (GEMM + MATH)") + ax_brk.bar(xs + width / 2, comm_us, width, + color="#c0504d", edgecolor="#222", linewidth=0.6, + label="communication (DMA)") + for i in range(len(ordered)): + ax_brk.text(xs[i] - width / 2, matmul_us[i] + max_v * 0.015, + f"{matmul_us[i]:.1f}", + ha="center", va="bottom", fontsize=8, + color="#222", fontweight="bold") + ax_brk.text(xs[i] + width / 2, comm_us[i] + max_v * 0.015, + f"{comm_us[i]:.1f}", + ha="center", va="bottom", fontsize=8, + color="#222", fontweight="bold") + + ax_brk.set_xticks(xs) + ax_brk.set_xticklabels([_MODEL_LABELS[m] for m in ordered], fontsize=9) + ax_brk.set_ylabel("op-log occupancy (µs, summed across engines)") + ax_brk.set_title( + "Matmul vs communication engine work per model\n" + "sum of GEMM/MATH and DMA occupancy — not wall-clock (overlap)", + fontsize=11, + ) + ax_brk.grid(True, axis="y", ls=":", alpha=0.5) + ax_brk.set_ylim(0, max_v * 1.20) + ax_brk.legend(loc="upper left", fontsize=9) + + fig.suptitle( + "Multi-model attention comparison — Case-6 composite kernel" + " ($S_{kv}=128$K, $T_q=1$, P=8 PEs / cube)", + fontsize=12, fontweight="bold", + ) + fig.tight_layout(rect=(0, 0, 1, 0.93)) + + out = _FIG_DIR / "gqa_decode_models.png" + fig.savefig(out, dpi=150) + plt.close(fig) + print(f"wrote {out}") + + if _PAPER_FIG_DIR.is_dir(): + dst = _PAPER_FIG_DIR / out.name + dst.write_bytes(out.read_bytes()) + print(f"copied {dst}") + + +if __name__ == "__main__": + main() diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_models.png b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_models.png new file mode 100644 index 0000000..5e767fb Binary files /dev/null and b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_decode_models.png differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode_models.json b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode_models.json new file mode 100644 index 0000000..db9452e --- /dev/null +++ b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/sweep_decode_models.json @@ -0,0 +1,173 @@ +{ + "version": 1, + "variant": "composite", + "S_kv": 131072, + "P": 8, + "note": "C = G (G-matched topology per model)", + "models": [ + "llama3-8b", + "llama3-70b", + "qwen2.5-7b", + "qwen2.5-72b", + "gemma2-27b", + "command-r-plus" + ], + "rows": [ + { + "model": "llama3-8b", + "family": "Meta", + "params_b": 8, + "h_q": 4, + "h_kv": 1, + "G": 4, + "d_head": 128, + "full_h_q": 32, + "full_h_kv": 8, + "hidden": 4096, + "layers": 32, + "C": 4, + "P": 8, + "N": 32, + "T_q": 1, + "S_kv": 131072, + "S_local": 4096, + "variant": "composite", + "pe_cpu_cmd_count": 44, + "pe_cpu_dispatch_cycles": 494, + "latency_ns": 394539.4071250012, + "matmul_ns": 4783.8393749997485, + "comm_ns": 1812.7050000057789, + "other_ns": 1684807746.5613055 + }, + { + "model": "llama3-70b", + "family": "Meta", + "params_b": 70, + "h_q": 8, + "h_kv": 1, + "G": 8, + "d_head": 128, + "full_h_q": 64, + "full_h_kv": 8, + "hidden": 8192, + "layers": 80, + "C": 8, + "P": 8, + "N": 64, + "T_q": 1, + "S_kv": 131072, + "S_local": 2048, + "variant": "composite", + "pe_cpu_cmd_count": 94, + "pe_cpu_dispatch_cycles": 978, + "latency_ns": 460651.3170000027, + "matmul_ns": 10069.639499999117, + "comm_ns": 4287.750000016997, + "other_ns": 2079259822.423066 + }, + { + "model": "qwen2.5-7b", + "family": "Alibaba", + "params_b": 7, + "h_q": 7, + "h_kv": 1, + "G": 7, + "d_head": 128, + "full_h_q": 28, + "full_h_kv": 4, + "hidden": 3584, + "layers": 28, + "C": 7, + "P": 8, + "N": 56, + "T_q": 1, + "S_kv": 131072, + "S_local": 2340, + "variant": "composite", + "pe_cpu_cmd_count": 77, + "pe_cpu_dispatch_cycles": 813, + "latency_ns": 456482.3358750015, + "matmul_ns": 8809.787499998813, + "comm_ns": 3486.4500000112457, + "other_ns": 2032890493.459298 + }, + { + "model": "qwen2.5-72b", + "family": "Alibaba", + "params_b": 72, + "h_q": 8, + "h_kv": 1, + "G": 8, + "d_head": 128, + "full_h_q": 64, + "full_h_kv": 8, + "hidden": 8192, + "layers": 80, + "C": 8, + "P": 8, + "N": 64, + "T_q": 1, + "S_kv": 131072, + "S_local": 2048, + "variant": "composite", + "pe_cpu_cmd_count": 94, + "pe_cpu_dispatch_cycles": 978, + "latency_ns": 460651.3170000027, + "matmul_ns": 10069.639499999117, + "comm_ns": 4287.750000016997, + "other_ns": 2079259822.423066 + }, + { + "model": "gemma2-27b", + "family": "Google", + "params_b": 27, + "h_q": 2, + "h_kv": 1, + "G": 2, + "d_head": 128, + "full_h_q": 32, + "full_h_kv": 16, + "hidden": 4608, + "layers": 46, + "C": 2, + "P": 8, + "N": 16, + "T_q": 1, + "S_kv": 131072, + "S_local": 8192, + "variant": "composite", + "pe_cpu_cmd_count": 60, + "pe_cpu_dispatch_cycles": 648, + "latency_ns": 263522.86775000195, + "matmul_ns": 2326.0945000193315, + "comm_ns": 907.5250000030501, + "other_ns": 1235149167.507364 + }, + { + "model": "command-r-plus", + "family": "Cohere", + "params_b": 104, + "h_q": 12, + "h_kv": 1, + "G": 12, + "d_head": 128, + "full_h_q": 96, + "full_h_kv": 8, + "hidden": 12288, + "layers": 64, + "C": 12, + "P": 8, + "N": 96, + "T_q": 1, + "S_kv": 131072, + "S_local": 1365, + "variant": "composite", + "pe_cpu_cmd_count": 111, + "pe_cpu_dispatch_cycles": 1143, + "latency_ns": 491928.65900000196, + "matmul_ns": 15894.392999998527, + "comm_ns": 9904.505000027013, + "other_ns": 2372909985.3171477 + } + ] +} \ No newline at end of file diff --git a/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp.py b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp.py index a1ca935..b43c4a5 100644 --- a/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp.py +++ b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp.py @@ -37,9 +37,9 @@ Topology / SFR: from __future__ import annotations from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import ( - _ROOT_CUBE, _merge_running, reduce_mlo, + root_cube_for, ) TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width). @@ -108,9 +108,9 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel( tl.copy_to(O_local, O_new) # ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ── - reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl) + reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl) - # ── Final normalise + store (root only: PE 0 of CUBE 6) ── - if pe_id == 0 and cube_id == _ROOT_CUBE: + # ── Final normalise + store (root only: PE 0 of the C-derived root cube) ── + if pe_id == 0 and cube_id == root_cube_for(C): 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_sp_pe_sp_composite.py b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite.py index 83c70e7..13f3c09 100644 --- a/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite.py +++ b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite.py @@ -18,8 +18,8 @@ under the ADR-0064 Rev2 structural model (the CPU-offload win). from __future__ import annotations from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import ( - _ROOT_CUBE, reduce_mlo, + root_cube_for, ) @@ -63,9 +63,9 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel( tl.composite(op="gemm", a=exp_scores, b=V, out=O_local) # P·V, one command # ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ── - reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl) + reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl) - # ── Final normalise + store (root only: PE 0 of CUBE 6) ── - if pe_id == 0 and cube_id == _ROOT_CUBE: + # ── Final normalise + store (root only: PE 0 of the C-derived root cube) ── + if pe_id == 0 and cube_id == root_cube_for(C): 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_sp_pe_sp_composite_ext.py b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext.py index 0dc8069..4c4d1b0 100644 --- a/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext.py +++ b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext.py @@ -27,8 +27,8 @@ numeric parity of the recipe's 8 MATH ops is a separate follow-up from __future__ import annotations from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import ( - _ROOT_CUBE, reduce_mlo, + root_cube_for, ) # Small primitive seed slice that establishes the running (m, ℓ, O) the @@ -86,9 +86,9 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_kernel( ) # ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ── - reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl) + reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl) - # ── Final normalise + store (root only: PE 0 of CUBE 6) ── - if pe_id == 0 and cube_id == _ROOT_CUBE: + # ── Final normalise + store (root only: PE 0 of the C-derived root cube) ── + if pe_id == 0 and cube_id == root_cube_for(C): 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_sp_pe_sp_hand_tiled_16x16x16.py b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16.py index 85622a7..c7993cf 100644 --- a/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16.py +++ b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16.py @@ -37,9 +37,9 @@ from math import ceil from kernbench.common.pe_commands import GemmCmd from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import ( - _ROOT_CUBE, _merge_running, reduce_mlo, + root_cube_for, ) TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width). @@ -208,8 +208,8 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel( tl.copy_to(l_local, l_new) tl.copy_to(O_local, O_new) - reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl) + reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl) - if pe_id == 0 and cube_id == _ROOT_CUBE: + if pe_id == 0 and cube_id == root_cube_for(C): O_final = O_local / l_local tl.store(o_ptr, O_final) diff --git a/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_mlo_reduce.py b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_mlo_reduce.py index e7db131..613f554 100644 --- a/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_mlo_reduce.py +++ b/src/kernbench/benches/gqa_helpers/long_ctx/_gqa_mlo_reduce.py @@ -1,33 +1,66 @@ """Shared (m, ℓ, O) reduce for the Case-6 long-context decode kernels. -Extracted from the Cube-SP × PE-SP decode kernel so the three -command-form variants (primitive / composite / composite_extended) +Extracted from the Cube-SP × PE-SP decode kernel so the four command-form +variants (primitive / primitive-tiled / composite / composite_extended) share one identical reduce. The local attention differs per variant; -the reduce does not. Behavior is byte-equal to the pre-extraction inline -reduce (guarded by tests/attention/test_gqa_decode_long_ctx_composite.py). +the reduce does not. Byte-equal behavior at C=8 is guarded by +``tests/attention/test_gqa_decode_long_ctx_composite.py``. The reduce is two-level: • intra-CUBE 8-way (row chain along intra_W + col bridge along intra_N) over the 2×4 PE grid; - • inter-CUBE 2-phase lrab over the 4×2 CUBE sub-mesh (ADR-0060 §4.2 - lrab-adapted center-root reduce): Phase 1 row reduce converges at - root_col=2; Phase 2 col reduce on root_col converges at root_row=1; - root cube id = 6. -Plain ``+`` is replaced by log-sum-exp ``_merge_running``. + • inter-CUBE 2-phase lrab-adapted center-root reduce over a + ``sub_w × sub_h`` CUBE sub-mesh (ADR-0060 §4.2): Phase 1 does a + bidirectional row reduce converging at ``root_col = sub_w // 2``; + Phase 2 does a bidirectional col reduce on ``root_col`` converging + at ``root_row = sub_h // 2``. + +The submesh dimensions ``(sub_w, sub_h)`` and root are **derived from +C** (number of cubes launched per KV group), not hardcoded, so the +same reduce works for any ``C ≥ 1``. Peer-existence guards on every +send/recv elide operations that would target un-launched neighbors +(needed for non-rectangular C like 7 or 12; safe no-ops when the +neighbor exists). For ``C = 8`` this reproduces the previous +hardcoded ``{sub_w=4, sub_h=2, root_col=2, root_row=1, root_cube=6}`` +exactly. + +Plain ``+`` in the tree is replaced by log-sum-exp ``_merge_running``. """ from __future__ import annotations -# 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 - PE_GRID_COLS = 4 +def _submesh_for(C): + """Derive ``(sub_w, sub_h, root_col, root_row, root_cube)`` for a + C-cube inter-cube reduce. + + Submesh is at most 4 columns wide (matching the SIP's 4×4 cube + mesh); rows fill left-to-right, top-to-bottom. For the Case-6 + baseline C=8, this returns ``(4, 2, 2, 1, 6)`` — identical to + the previously hardcoded values. + """ + sub_w = min(C, 4) + sub_h = (C + sub_w - 1) // sub_w + root_col = sub_w // 2 + root_row = sub_h // 2 + root_cube = root_row * sub_w + root_col + return sub_w, sub_h, root_col, root_row, root_cube + + +def root_cube_for(C): + """Root cube id for a C-cube reduce. Callers use this to gate the + final ``tl.store`` in the kernel epilogue.""" + _, _, _, _, root_cube = _submesh_for(C) + return root_cube + + +# Backward-compat alias: preserves the Case-6 baseline value (6) for +# any external code that still imports the module-level constant. +_ROOT_CUBE = root_cube_for(8) + + 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) @@ -38,13 +71,23 @@ def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl): return m_new, l_new, O_new -def reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, *, tl): - """Two-level (m, ℓ, O) reduce-to-root (PE 0 of CUBE 6). In place. +def reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, *, tl): + """Two-level (m, ℓ, O) reduce-to-root (PE 0 of ``root_cube_for(C)``). + In place. - Updates ``m_local``/``l_local``/``O_local`` in place via ``copy_to``; - after this call only PE 0 of CUBE ``_ROOT_CUBE`` holds the full result. + Updates ``m_local`` / ``l_local`` / ``O_local`` in place via + ``copy_to``; after this call only PE 0 of ``root_cube_for(C)`` + holds the fully-merged result. + + Peer-existence guards on every inter-cube send/recv ensure the + reduce completes for any launched C (rectangular or not); a + send/recv toward a non-launched neighbor is skipped. For C = 8 + every neighbor exists → guards are all True → byte-equal to the + previously hardcoded Case-6 reduce. """ - # ── Intra-CUBE reduce: row chain (intra_W) + col bridge (intra_N) ── + sub_w, sub_h, root_col, root_row, _root_cube = _submesh_for(C) + + # ── Intra-CUBE reduce (unchanged — depends only on P) ───────────── pe_col = pe_id % PE_GRID_COLS pe_row = pe_id // PE_GRID_COLS pe_cols_used = min(PE_GRID_COLS, P) @@ -84,106 +127,126 @@ def reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, *, tl): tl.send(dir="intra_N", src=l_local) tl.send(dir="intra_N", src=O_local) - # ── Inter-CUBE lrab-adapted center-root reduce (ADR-0060 §4.2) ── + # ── Inter-CUBE lrab-adapted center-root reduce ──────────────────── # Only PE 0 of each CUBE participates. Adapts Phases 1-2 of # lrab_hierarchical_allreduce.py: bidirectional row reduce converges - # at root_col; bidirectional col reduce on root_col converges at - # root_row. Plain ``+`` replaced by log-sum-exp ``_merge_running``. + # at ``root_col``; bidirectional col reduce on ``root_col`` converges + # at ``root_row``. Plain ``+`` replaced by log-sum-exp + # ``_merge_running``. Peer-existence guards elide sends/recvs whose + # neighbor cube was not launched. if pe_id == 0: - row = cube_id // _SUB_W - col = cube_id % _SUB_W + row = cube_id // sub_w + col = cube_id % sub_w - # Phase 1: row reduce — converge at col == _ROOT_COL. + # Peer existence within the launched set of C cubes. + east_exists = (col + 1 < sub_w) and (cube_id + 1 < C) + west_exists = col > 0 + south_exists = cube_id + sub_w < C + north_exists = row > 0 + + # 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) + if east_exists: + 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: + if west_exists: + 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) + if east_exists: + 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: + if west_exists: + 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) + if east_exists: + 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: + if east_exists: + 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) + if west_exists: + 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: + if west_exists: + 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: + # 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: + if south_exists: + 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: + if north_exists: + 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 south_exists: + 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: + if north_exists: + 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 and south_exists: 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") @@ -194,21 +257,24 @@ def reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, *, 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) + elif root_row < row < sub_h - 1: + if south_exists: + 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) + if north_exists: + 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: + if north_exists: + tl.send(dir="N", src=m_local) + tl.send(dir="N", src=l_local) + tl.send(dir="N", src=O_local) diff --git a/src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_models.py b/src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_models.py new file mode 100644 index 0000000..fdc7b3e --- /dev/null +++ b/src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_models.py @@ -0,0 +1,262 @@ +"""GQA composite kernel across models — per-KV-group multi-model comparison. + +The topology per model uses C = h_q (cubes per KV group = query heads +per KV group in the model). P = 8 PEs / cube is fixed to the physical +layout. So Gemma-2 27B (G=2) runs on a 2-cube slice; Command R+ (G=12) +runs on 12 cubes; LLaMA-3-70B (G=8) matches the Case-6 baseline +exactly. + +Companion to ``gqa_decode_long_ctx_models.py``. That bench sweeps a +fixed set of topology sizes (N ∈ {16, 32, 64}) across six GQA models +and showed that latency is model-agnostic when only G varies. + +This bench takes a different slice: it fixes the topology to match +each model's G ratio — one cube per Q head in the KV group: + + C = G (cubes per KV group = h_q of the KV group) + P = 8 (fixed, physical PEs per cube) + N = C · P = 8·G + +Six models, one topology per model, one context length: + + Gemma 2 27B G=2 → C=2 → N=16 + LLaMA-3 8B G=4 → C=4 → N=32 + Qwen 2.5 7B G=7 → C=7 → N=56 (2 KV groups per SIP, 2 cubes idle) + LLaMA-3 70B G=8 → C=8 → N=64 + Qwen 2.5 72B G=8 → C=8 → N=64 + Command R+ G=12 → C=12 → N=96 (1 KV group per SIP, 4 cubes idle) + +Physical SIP is 4×4 = 16 cubes (topology.yaml sip.cube_mesh: {w:4, h:4}), +so all six configurations fit in one SIP. + +Total: 6 engine runs at S_kv = 131 072 (~40 min wall-clock). +""" +from __future__ import annotations + +import json +from pathlib import Path + +from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite import ( # noqa: E501 + gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_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 + +_OUTPUT_DIR = ( + Path(__file__).resolve().parents[2] + / "1H_milestone_output" / "gqa" / "long_ctx" +) +_SWEEP_JSON = _OUTPUT_DIR / "sweep_decode_models.json" + + +# ── Model registry (per-KV-group shapes) ───────────────────────────── + + +# Kernel operates on ONE KV group at a time. h_q / h_kv here are the +# per-group counts (not the model's totals); G = h_q / h_kv. +# Full-model context (full_h_q, full_h_kv, hidden, layers) is metadata +# for annotation only — it does not enter the kernel launch. +_MODELS = { + "llama3-8b": dict( + family="Meta", params_b=8, + h_q=4, h_kv=1, d_head=128, + full_h_q=32, full_h_kv=8, hidden=4096, layers=32, + ), + "llama3-70b": dict( + family="Meta", params_b=70, + h_q=8, h_kv=1, d_head=128, + full_h_q=64, full_h_kv=8, hidden=8192, layers=80, + ), + "qwen2.5-7b": dict( + family="Alibaba", params_b=7, + h_q=7, h_kv=1, d_head=128, + full_h_q=28, full_h_kv=4, hidden=3584, layers=28, + ), + "qwen2.5-72b": dict( + family="Alibaba", params_b=72, + h_q=8, h_kv=1, d_head=128, + full_h_q=64, full_h_kv=8, hidden=8192, layers=80, + ), + "gemma2-27b": dict( + family="Google", params_b=27, + h_q=2, h_kv=1, d_head=128, + full_h_q=32, full_h_kv=16, hidden=4608, layers=46, + ), + "command-r-plus": dict( + family="Cohere", params_b=104, + h_q=12, h_kv=1, d_head=128, + full_h_q=96, full_h_kv=8, hidden=12288, layers=64, + ), +} +_MODEL_ORDER = tuple(_MODELS.keys()) + +_S_KV = 131_072 # 128 K context +_T_Q = 1 +_P = 8 # fixed physical PEs per cube + + +def _c_for(model_key: str) -> int: + m = _MODELS[model_key] + return m["h_q"] // m["h_kv"] + + +def _end_to_end_ns(op_log) -> float: + if not op_log: + return 0.0 + return max(r.t_end for r in op_log) - min(r.t_start for r in op_log) + + +def _op_breakdown_ns(op_log) -> dict[str, float]: + """Sum occupancy per op_kind across all components (not + critical-path — ops overlap on independent components). First-order + view of where time is spent inside the engine.""" + totals = {"matmul": 0.0, "comm": 0.0, "other": 0.0} + for r in op_log: + dur = r.t_end - r.t_start + if r.op_kind in ("gemm", "math"): + totals["matmul"] += dur + elif r.op_kind == "memory": + totals["comm"] += dur + else: + totals["other"] += dur + return totals + + +def _run_panel_fn(model_key: str, C: int, P: int, S_kv: int): + m = _MODELS[model_key] + panel = f"decode_models_{model_key}_c{C}p{P}_s{S_kv}" + + def _bench_fn(ctx): + 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", pe="row_wise", + num_cubes=C, num_pes=P) + q = ctx.zeros((_T_Q, m["h_q"] * m["d_head"]), + dtype="f16", dp=dp_full, name=f"{panel}_q") + k = ctx.zeros((S_kv, m["h_kv"] * m["d_head"]), + dtype="f16", dp=dp_kv, name=f"{panel}_k") + v = ctx.zeros((S_kv, m["h_kv"] * m["d_head"]), + dtype="f16", dp=dp_kv, name=f"{panel}_v") + o = ctx.empty((_T_Q, m["h_q"] * m["d_head"]), + dtype="f16", dp=dp_full, name=f"{panel}_o") + ctx.launch( + panel, + gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel, + q, k, v, o, + _T_Q, S_kv, m["h_q"], m["h_kv"], m["d_head"], C, P, + _auto_dim_remap=False, + ) + + return _bench_fn + + +def _emit_dispatch(model_key: str, C: int, P: int, + S_kv: int) -> tuple[int, float]: + from kernbench.common.pe_commands import PeCpuOverheadCmd + from kernbench.common.pe_cost_model import DEFAULT_PE_COST_MODEL + from kernbench.triton_emu.tl_context import TLContext, run_kernel + + m = _MODELS[model_key] + cube_id = min(6, C - 1) + tl = TLContext( + pe_id=0, num_programs=P, cost_model=DEFAULT_PE_COST_MODEL, + cube_id=cube_id, num_cubes=C, scratch_base=1 << 61, + scratch_size=1 << 20, + ) + run_kernel( + gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel, tl, + 0x1000, 0x2000, 0x3000, 0x4000, + _T_Q, S_kv, m["h_q"], m["h_kv"], m["d_head"], C, P, + ) + cmds = [c for c in tl.commands if isinstance(c, PeCpuOverheadCmd)] + return len(cmds), sum(c.cycles for c in cmds) + + +def _engine_run(model_key: str, C: int, P: int, + S_kv: int, topology: str): + """Run the engine sim; return (latency_ns, op_kind_breakdown).""" + 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=_run_panel_fn(model_key, C, P, S_kv), + 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"gqa-decode-models {model_key} C={C} P={P} S={S_kv} " + f"failed: {result.completion}" + ) + op_log = result.engine.op_log + latency = _end_to_end_ns(op_log) + breakdown = _op_breakdown_ns(op_log) + return latency, breakdown + + +def run_sweep(topology: str = "topology.yaml") -> int: + """One row per model at its per-KV-group topology (C = h_q). Writes + sweep_decode_models.json.""" + _OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + rows = [] + for model_key in _MODEL_ORDER: + m = _MODELS[model_key] + C = _c_for(model_key) + N = C * _P + n_cmds, cycles = _emit_dispatch(model_key, C, _P, _S_KV) + latency, breakdown = _engine_run(model_key, C, _P, _S_KV, topology) + rows.append({ + "model": model_key, + "family": m["family"], + "params_b": m["params_b"], + "h_q": m["h_q"], + "h_kv": m["h_kv"], + "G": C, + "d_head": m["d_head"], + "full_h_q": m["full_h_q"], + "full_h_kv": m["full_h_kv"], + "hidden": m["hidden"], + "layers": m["layers"], + "C": C, "P": _P, "N": N, + "T_q": _T_Q, + "S_kv": _S_KV, + "S_local": _S_KV // N, + "variant": "composite", + "pe_cpu_cmd_count": n_cmds, + "pe_cpu_dispatch_cycles": cycles, + "latency_ns": latency, + # Op-kind occupancy (summed across all components — not + # critical-path; first-order view of where engine time goes) + "matmul_ns": breakdown["matmul"], + "comm_ns": breakdown["comm"], + "other_ns": breakdown["other"], + }) + print( + f" {model_key:<18} G={C:>2} C={C:>2} P={_P} N={N:>3} " + f"latency={latency/1e3:>7.1f} µs cmds={n_cmds:>4} " + f"matmul={breakdown['matmul']/1e3:>7.1f} " + f"comm={breakdown['comm']/1e3:>7.1f}" + ) + sweep = { + "version": 1, + "variant": "composite", + "S_kv": _S_KV, + "P": _P, + "note": "C = h_q per model (cubes per KV group = query heads per KV group)", + "models": list(_MODEL_ORDER), + "rows": rows, + } + _SWEEP_JSON.write_text(json.dumps(sweep, indent=2)) + print(f" gqa-decode-models: {len(rows)} rows -> {_SWEEP_JSON}") + return len(rows) + + +if __name__ == "__main__": + run_sweep()