gqa: single-KV-group LLaMA-3.1-70B prefill milestone (Increments 1-5)

End-to-end wires C=8 P=8 d_head=128 prefill with snake-ring inter-CUBE
SFR, intra-CUBE PE-SP (all 64 ranks active), and the milestone bench
panel. Decode kernel gains lrab-adapted center-root reduce for the
2×4 sub-mesh per ADR-0060 §4.2.

Increment 1 — SFR multi-row snake
  src/kernbench/ccl/sfr_config.py: configure_sfr_intercube_ring gains
  submesh_shape / submesh_origin kwargs; installs a Hamiltonian snake
  ring through a rectangular sub-mesh (every hop is 1-hop physical
  neighbour). Backward-compat: 1D-row behaviour preserved when
  submesh_shape is None.
  tests/test_intercube_snake_ring.py (12 tests)

Increment 2 — Decode lrab-adapted center-root reduce
  src/kernbench/benches/_gqa_attention_decode_long.py: new sub_w param
  (default 0 = existing 1D-chain). sub_w >= 2 selects the ADR-0060
  §4.2 prescribed lrab-adapted Phase 1+2 reduce (bidirectional row +
  bidirectional col converge to the center cube), with log-sum-exp
  _merge_running replacing the plain + of lrab.
  tests/attention/test_gqa_decode_long_2d_reduce.py (4 tests)

Increment 3 — Prefill kernel at C=8 (no production change)
  Verified by inspection that the existing prefill_long kernel +
  Increment 1's snake SFR already work at C=8 without any kernel
  edit. The kernel speaks logical W/E; the snake routes it.
  tests/attention/test_gqa_prefill_long_c8_snake.py (3 tests)

Increment 4 — Intra-CUBE PE-SP in prefill (all 64 ranks)
  src/kernbench/benches/_gqa_attention_prefill_long.py: new P param
  (default 1 = existing PE-0-only). P > 1 splits T_q query-axis-wise
  across the P PEs of each CUBE; output rows are disjoint per PE so
  no intra-CUBE reduce is needed; each PE drives its own same-lane
  ring (P parallel rings).
  tests/attention/test_gqa_prefill_long_pe_sp.py (5 tests)

Increment 5 — LLaMA-scale milestone bench panel
  src/kernbench/benches/milestone_gqa_headline.py: new panel
  single_kv_group_prefill_gqa_c8_p8 (C=8, P=8, T_q=S_kv=32K,
  d_head=128). _run_prefill_panel extended with P/T_q/d_head
  defaults; routes snake SFR when C > mesh_w.
  tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py (3 tests)

Total: 4 production files modified, 5 new test files, 27 new tests.
Followed the Phase 1/2 protocol per CLAUDE.md throughout.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

src/kernbench/benches/milestone_gqa_headline.py

@@ -6,19 +6,24 @@ per-panel ``op_log_summary`` into ``sweep.json``. Independent from the
 existing ``milestone-gqa-llama70b`` validation-scale bench (which stays
 on the legacy baseline kernels).
 
-Restrictions (P7 first cut):
-  - C ≤ 4 (single-row inter-CUBE ring SFR; multi-row deferred)
+Restrictions:
+  - Decode side capped at C ≤ 4 (1D chain reduce; 2D mesh wiring on
+    the decode panels deferred to the 4-cases decode comparative study)
   - Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP
-    headline deferred)
+    headline deferred per ADR-0060 §B-item-1)
   - No figure renderers (defer to a separate cycle)
 
 Panels:
-  single_user_prefill_gqa : prefill C=1, T_q=4, S_kv=16
-  multi_user_prefill_gqa  : prefill C=4 Ring KV, T_q=4, S_kv=16
-  single_user_decode_gqa  : decode C=1, P=8, h_q=8, h_kv=1, S_kv=64
-                            (M-fold + intra-cube row-then-col chain)
-  multi_user_decode_gqa   : decode C=4, P=8, h_q=8, h_kv=1, S_kv=128
-                            (M-fold + 2-level chain reduce-to-root)
+  single_user_prefill_gqa         : prefill C=1, T_q=4, S_kv=16
+  multi_user_prefill_gqa          : prefill C=4 Ring KV, T_q=4, S_kv=16
+  single_kv_group_prefill_gqa_c8_p8: prefill C=8 snake Ring KV +
+                                     intra-CUBE PE-SP (all 64 ranks),
+                                     T_q=S_kv=32K, d_head=128 — the
+                                     LLaMA-3.1-70B single-KV-group target
+  single_user_decode_gqa          : decode C=1, P=8, h_q=8, h_kv=1, S_kv=64
+                                    (M-fold + intra-cube row-then-col chain)
+  multi_user_decode_gqa           : decode C=4, P=8, h_q=8, h_kv=1, S_kv=128
+                                    (M-fold + 2-level chain reduce-to-root)
 
 Gated by ``GQA_HEADLINE_RUN=1``.
 """
@@ -54,6 +59,7 @@ _H_KV_DECODE = 1
 _PANELS = (
     "single_user_prefill_gqa",
     "multi_user_prefill_gqa",
+    "single_kv_group_prefill_gqa_c8_p8",
     "single_user_decode_gqa",
     "multi_user_decode_gqa",
 )
@@ -62,6 +68,11 @@ _PANELS = (
 _PANEL_DISPATCH: dict[str, tuple[str, dict]] = {
     "single_user_prefill_gqa": ("prefill", {"C": 1, "S_kv": _S_KV_PREFILL}),
     "multi_user_prefill_gqa":  ("prefill", {"C": 4, "S_kv": _S_KV_PREFILL}),
+    "single_kv_group_prefill_gqa_c8_p8": ("prefill", {
+        "C": 8, "P": 8,
+        "T_q": 32_768, "S_kv": 32_768,
+        "d_head": 128,
+    }),
     "single_user_decode_gqa":  ("decode",  {"C": 1, "P": 8, "S_kv": 64}),
     "multi_user_decode_gqa":   ("decode",  {"C": 4, "P": 8, "S_kv": 128}),
 }
@@ -76,29 +87,49 @@ def _ccl_cfg():
 # ── Per-kind launch helpers ──────────────────────────────────────────
 
 
-def _run_prefill_panel(ctx, *, panel: str, C: int, S_kv: int) -> None:
+def _run_prefill_panel(
+    ctx, *, panel: str, C: int, S_kv: int,
+    P: int = 1,
+    T_q: int = _T_Q_PREFILL,
+    d_head: int = _D_HEAD,
+) -> None:
     if C > 1:
-        configure_sfr_intercube_ring(
-            ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C,
-        )
-    dp_q = DPPolicy(cube="replicate", pe="replicate",
-                    num_cubes=C, num_pes=1)
+        mesh_w = int(ctx.spec["sip"]["cube_mesh"]["w"])
+        if C <= mesh_w:
+            configure_sfr_intercube_ring(
+                ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C,
+            )
+        else:
+            if C % mesh_w != 0:
+                raise ValueError(
+                    f"C={C} > mesh_w={mesh_w} requires C divisible by mesh_w"
+                )
+            configure_sfr_intercube_ring(
+                ctx.engine, ctx.spec, _ccl_cfg(),
+                submesh_shape=(C // mesh_w, mesh_w),
+            )
+    # Q and O switch to pe="row_wise" when intra-CUBE PE-SP is active
+    # (ADR-0060 §5.5 + §B-item-3: disjoint query-row split across PEs).
+    q_pe = "row_wise" if P > 1 else "replicate"
+    o_pe = "row_wise" if P > 1 else "replicate"
+    dp_q = DPPolicy(cube="replicate", pe=q_pe,
+                    num_cubes=C, num_pes=P)
     dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
-                     pe="replicate", num_cubes=C, num_pes=1)
+                     pe="replicate", num_cubes=C, num_pes=P)
     dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate",
-                    pe="replicate", num_cubes=C, num_pes=1)
-    q = ctx.zeros((_T_Q_PREFILL, _D_HEAD),
+                    pe=o_pe, num_cubes=C, num_pes=P)
+    q = ctx.zeros((T_q, d_head),
                   dtype=_DTYPE, dp=dp_q, name=f"{panel}_q")
-    k = ctx.zeros((S_kv, _D_HEAD),
+    k = ctx.zeros((S_kv, d_head),
                   dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k")
-    v = ctx.zeros((S_kv, _D_HEAD),
+    v = ctx.zeros((S_kv, d_head),
                   dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v")
-    o = ctx.empty((_T_Q_PREFILL * C, _D_HEAD),
+    o = ctx.empty((T_q * C, d_head),
                   dtype=_DTYPE, dp=dp_o, name=f"{panel}_o")
     ctx.launch(
         panel, gqa_attention_prefill_long_kernel,
         q, k, v, o,
-        _T_Q_PREFILL, S_kv, _D_HEAD, C,
+        T_q, S_kv, d_head, C, P,
         _auto_dim_remap=False,
     )