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/ccl/sfr_config.py

@@ -248,14 +248,29 @@ def configure_sfr_intercube_ring(
     cfg: dict,
     *,
     ring_size: int | None = None,
+    submesh_shape: tuple[int, int] | None = None,
+    submesh_origin: tuple[int, int] = (0, 0),
 ) -> dict[str, Any]:
-    """Install intra-cube PE grid + a 1D CUBE-level ring with wrap.
+    """Install intra-cube PE grid + a CUBE-level ring with wrap.
+
+    Two ring layouts:
+
+      - **1D row** (default; ``submesh_shape=None``): cubes
+        ``0..ring_size-1`` form a 1D ring with wrap. ``ring_size`` must
+        be ≤ ``mesh_w`` so every hop is a 1-hop CUBE NOC neighbour.
+      - **Snake/serpentine** (``submesh_shape=(rows, cols)``,
+        ADR-0060 §5.5 prefill Ring KV at C=G=8 on a 2×4 sub-mesh):
+        a boustrophedon Hamiltonian cycle through the ``rows × cols``
+        sub-mesh rooted at ``submesh_origin``. Every consecutive pair
+        on the snake (including the wrap) is a 1-hop CUBE NOC
+        neighbour, so the kernel sees a 1D logical E/W ring without
+        being aware of the underlying 2D layout.
 
     Direction namespaces (disjoint, same as
     ``configure_sfr_intercube_multisip``):
 
       - ``intra_N/S/E/W`` : 2×4 PE grid within each cube (no wrap)
-      - ``E/W``           : 1D ring of cubes 0..ring_size-1 WITH WRAP
+      - ``E/W``           : ring along the resolved path WITH WRAP
                             (symmetric to ``configure_sfr_intracube_pe_ring``
                              at PE level — wrap applied at CUBE level here)
       - ``global_*``      : SIP topology (same as multisip)
@@ -264,10 +279,15 @@ def configure_sfr_intercube_ring(
     ``configure_sfr_intercube_multisip`` for the full 4×4 cube mesh.
 
     Args:
-        ring_size: number of CUBEs in the ring (wrap applies to cubes
-            0..ring_size-1). Defaults to the full cube_mesh count.
-            Must be ≤ mesh_w (single row); multi-row rings span
-            non-neighbour boundaries.
+        ring_size: number of CUBEs in the ring. Defaults to the full
+            cube_mesh count for the 1D-row case, or ``rows*cols`` for
+            the snake case. If passed alongside ``submesh_shape``, must
+            equal ``rows*cols``.
+        submesh_shape: ``(rows, cols)`` of the snake sub-mesh. If
+            given, ring follows a boustrophedon path through that
+            rectangle. If ``None``, falls back to the 1D-row layout.
+        submesh_origin: ``(row, col)`` top-left of the sub-mesh inside
+            the cube mesh. Defaults to ``(0, 0)``.
     """
     cm = spec["sip"]["cube_mesh"]
     mesh_w = int(cm["w"])
@@ -281,13 +301,42 @@ def configure_sfr_intercube_ring(
     sip_w = int(sip_w) if sip_w is not None else None
     sip_h = int(sip_h) if sip_h is not None else None
 
-    if ring_size is None:
-        ring_size = n_cubes
-    if ring_size > mesh_w:
-        raise ValueError(
-            f"intercube_ring ring_size={ring_size} > mesh_w={mesh_w}; "
-            "multi-row rings cross non-neighbour boundaries"
-        )
+    if submesh_shape is not None:
+        sub_h, sub_w = submesh_shape
+        origin_row, origin_col = submesh_origin
+        if (sub_h <= 0 or sub_w <= 0
+                or origin_row < 0 or origin_col < 0
+                or origin_row + sub_h > mesh_h
+                or origin_col + sub_w > mesh_w):
+            raise ValueError(
+                f"submesh_shape={submesh_shape} at origin={submesh_origin} "
+                f"does not fit cube_mesh ({mesh_h}x{mesh_w})"
+            )
+        expected_size = sub_h * sub_w
+        if ring_size is not None and ring_size != expected_size:
+            raise ValueError(
+                f"ring_size={ring_size} inconsistent with "
+                f"submesh_shape={submesh_shape} (expected {expected_size})"
+            )
+        ring_size = expected_size
+        # Boustrophedon: even rows L→R, odd rows R→L.
+        ring_path: list[int] = []
+        for r in range(sub_h):
+            cols = range(sub_w) if r % 2 == 0 else range(sub_w - 1, -1, -1)
+            for c in cols:
+                ring_path.append((origin_row + r) * mesh_w + (origin_col + c))
+    else:
+        if ring_size is None:
+            ring_size = n_cubes
+        if ring_size > mesh_w:
+            raise ValueError(
+                f"intercube_ring ring_size={ring_size} > mesh_w={mesh_w}; "
+                "multi-row rings cross non-neighbour boundaries"
+            )
+        ring_path = list(range(ring_size))
+
+    ring_pos: dict[int, int] = {c: i for i, c in enumerate(ring_path)}
+    ring_len = len(ring_path)
 
     if sip_topology not in _TOPO_BUILTINS:
         raise ValueError(
@@ -329,10 +378,11 @@ def configure_sfr_intercube_ring(
         for d, peer_pe in _intra_cube_neighbors(pe).items():
             nbrs[d] = _pe_idx(sip, cube, peer_pe)
 
-        # ── Cube ring (E/W with wrap for cubes 0..ring_size-1) ──
-        if cube < ring_size:
-            nbrs["E"] = _pe_idx(sip, (cube + 1) % ring_size, pe)
-            nbrs["W"] = _pe_idx(sip, (cube - 1) % ring_size, pe)
+        # ── Cube ring (E/W along resolved ring_path, with wrap) ──
+        pos = ring_pos.get(cube)
+        if pos is not None:
+            nbrs["E"] = _pe_idx(sip, ring_path[(pos + 1) % ring_len], pe)
+            nbrs["W"] = _pe_idx(sip, ring_path[(pos - 1) % ring_len], pe)
 
         # ── Inter-SIP same-(cube, pe) (global_*) ──
         if n_sips > 1: