gqa-decode: topology-agnostic reduce + multi-model attention bench

Make reduce_mlo derive its submesh dimensions (sub_w, sub_h, root_col,
root_row, root_cube) from C at call time, with peer-existence guards on
every inter-cube send/recv so any C ≥ 1 completes cleanly (previously
hardcoded to C=8's 4×2 mesh; non-rectangular C like 7 or 12 hit
IpcqDeadlock at the root cube's east receive). Byte-equal at C=8 —
existing composite digest tests still pass 8/8.

Callsites in the four Case-6 kernels (primitive / primitive-tiled /
composite / composite_extended) updated to pass C into reduce_mlo and
use root_cube_for(C) for the final tl.store gate.

Add the multi-model attention bench: sweeps the composite kernel at
S_kv=128K across six GQA models (Gemma-2 27B, LLaMA-3 8B, Qwen 2.5 7B,
LLaMA-3 70B, Qwen 2.5 72B, Command R+), with per-model topology
C = h_q (cubes per KV group = query heads per KV group). Captures
per-op-kind occupancy (matmul GEMM+MATH, comm DMA) alongside latency
and PE_CPU dispatch. Three-panel plot writes to bench-output and
paper-figures dir.

Headline result: same-h_q-different-family models land within 0.2 µs
(model-agnostic given fixed h_kv=1 per KV group + d_head=128); latency
climbs 263 → 492 µs as C climbs 2 → 12, driven by reduce depth
(comm/matmul ratio 0.4 → 0.6 across the sweep).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

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()

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
+    }
+  ]
+}

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) ──
+    # ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
-    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)

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) ──
+    # ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
-    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)

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) ──
+    # ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
-    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)

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)

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
+Extracted from the Cube-SP × PE-SP decode kernel so the four command-form
-command-form variants (primitive / composite / composite_extended)
+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
+the reduce does not. Byte-equal behavior at C=8 is guarded by
-reduce (guarded by tests/attention/test_gqa_decode_long_ctx_composite.py).
+``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
+  • inter-CUBE 2-phase lrab-adapted center-root reduce over a
-    lrab-adapted center-root reduce): Phase 1 row reduce converges at
+    ``sub_w × sub_h`` CUBE sub-mesh (ADR-0060 §4.2): Phase 1 does a
-    root_col=2; Phase 2 col reduce on root_col converges at root_row=1;
+    bidirectional row reduce converging at ``root_col = sub_w // 2``;
-    root cube id = 6.
+    Phase 2 does a bidirectional col reduce on ``root_col`` converging
-Plain ``+`` is replaced by log-sum-exp ``_merge_running``.
+    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):
+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 CUBE 6). In place.
+    """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``;
+    Updates ``m_local`` / ``l_local`` / ``O_local`` in place via
-    after this call only PE 0 of CUBE ``_ROOT_CUBE`` holds the full result.
+    ``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,21 +127,31 @@ 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
+    # at ``root_col``; bidirectional col reduce on ``root_col`` converges
-    # root_row. Plain ``+`` replaced by log-sum-exp ``_merge_running``.
+    # 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
+        row = cube_id // sub_w
-        col = 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:
+            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:
+        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")
@@ -109,10 +162,12 @@ 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)
+            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:
+        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")
@@ -123,6 +178,7 @@ 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)
+            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")
@@ -133,7 +189,8 @@ 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_COL < col < _SUB_W - 1:
+        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")
@@ -144,21 +201,25 @@ 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)
+            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:
+        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.
+        # Phase 2: col reduce on col == root_col — converge at row == root_row.
-        if col == _ROOT_COL:
+        if col == root_col:
             if row == 0:
+                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:
+            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")
@@ -169,10 +230,12 @@ 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)
+                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:
+            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")
@@ -183,7 +246,7 @@ 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)
-                if _SUB_H - 1 > _ROOT_ROW:
+                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,7 +257,8 @@ 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:
+            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")
@@ -205,10 +269,12 @@ 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)
+                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:
+            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)

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()