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>

tests/test_intercube_snake_ring.py

@@ -0,0 +1,273 @@
+"""Tests for snake/serpentine ring extension of configure_sfr_intercube_ring.
+
+Verifies the multi-row snake ring SFR wiring used for ADR-0060 §5.5
+prefill Ring KV at C=8 on a 2×4 sub-mesh of the 4×4 CUBE mesh.
+
+The snake path for ``submesh_shape=(2, 4), origin=(0, 0)`` is::
+
+    (0,0)→(0,1)→(0,2)→(0,3)→(1,3)→(1,2)→(1,1)→(1,0)→wrap to (0,0)
+
+In cube indices on a mesh_w=4 grid: ``[0, 1, 2, 3, 7, 6, 5, 4]``.
+
+Every consecutive pair (including the wrap) is a 1-hop CUBE NOC neighbour.
+The kernel sees a 1D logical E/W ring; the snake is invisible to it.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
+from kernbench.ccl.sfr_config import configure_sfr_intercube_ring
+from kernbench.sim_engine.engine import GraphEngine
+from kernbench.topology.builder import resolve_topology
+
+TOPOLOGY_PATH = Path(__file__).parent.parent / "topology.yaml"
+
+N_CUBES = 16          # 4×4 SIP CUBE mesh
+PES_PER_CUBE = 8
+MESH_W = 4
+
+# Canonical snake path for (2,4) sub-mesh at origin (0,0).
+SNAKE_2X4_O00 = [0, 1, 2, 3, 7, 6, 5, 4]
+
+# Canonical snake path for (2,4) sub-mesh at origin (1,0) (rows 1-2).
+SNAKE_2X4_O10 = [4, 5, 6, 7, 11, 10, 9, 8]
+
+
+def _engine_and_spec():
+    topo = resolve_topology(str(TOPOLOGY_PATH))
+    engine = GraphEngine(topo.topology_obj, enable_data=True)
+    return engine, topo.topology_obj.spec
+
+
+def _merged_cfg():
+    cfg = load_ccl_config()
+    return resolve_algorithm_config(cfg, name="lrab_hierarchical_allreduce")
+
+
+def _qp(engine, sip: int, cube: int, pe: int):
+    return engine._components[f"sip{sip}.cube{cube}.pe{pe}.pe_ipcq"].queue_pairs
+
+
+def _row_col(cube: int) -> tuple[int, int]:
+    return cube // MESH_W, cube % MESH_W
+
+
+def _l1(c1: int, c2: int) -> int:
+    r1, col1 = _row_col(c1)
+    r2, col2 = _row_col(c2)
+    return abs(r1 - r2) + abs(col1 - col2)
+
+
+class TestSnakeRing2x4Origin00:
+    """Snake ring through the top 2×4 sub-mesh of the 4×4 SIP."""
+
+    def test_snake_ring_2x4_world_size(self):
+        """All PEs on all CUBEs of all SIPs still enumerated (snake only
+        restricts which CUBEs have ring links, not the world)."""
+        engine, spec = _engine_and_spec()
+        cfg = _merged_cfg()
+        plan = configure_sfr_intercube_ring(
+            engine, spec, cfg, submesh_shape=(2, 4),
+        )
+        n_sips = int(spec["system"]["sips"]["count"])
+        assert plan["world_size"] == n_sips * N_CUBES * PES_PER_CUBE
+        assert len(plan["rank_to_pe"]) == plan["world_size"]
+
+    def test_snake_ring_2x4_path_corners_E(self):
+        """The four 'interesting' E hops on the snake path:
+
+          cube0.E → cube1   (start of row 0)
+          cube3.E → cube7   (row-bridge: top-right corner down)
+          cube7.E → cube6   (row 1, going leftward)
+          cube4.E → cube0   (wrap: bottom-left to top-left)
+        """
+        engine, spec = _engine_and_spec()
+        cfg = _merged_cfg()
+        configure_sfr_intercube_ring(
+            engine, spec, cfg, submesh_shape=(2, 4),
+        )
+        assert _qp(engine, 0, 0, 0)["E"]["peer"].cube == 1
+        assert _qp(engine, 0, 3, 0)["E"]["peer"].cube == 7
+        assert _qp(engine, 0, 7, 0)["E"]["peer"].cube == 6
+        assert _qp(engine, 0, 4, 0)["E"]["peer"].cube == 0
+
+    def test_snake_ring_2x4_path_corners_W(self):
+        """Symmetric W hops (W is reverse of E along the snake):
+
+          cube0.W → cube4   (wrap reverse)
+          cube7.W → cube3   (row-bridge reverse: up)
+          cube6.W → cube7   (row 1 reverse: rightward)
+          cube4.W → cube5   (row-1 leftmost, W goes right within row 1)
+
+        Note: under the snake path [0,1,2,3,7,6,5,4], W(cube_i) is the
+        predecessor on that path. So W(4)=5 (previous on path).
+        """
+        engine, spec = _engine_and_spec()
+        cfg = _merged_cfg()
+        configure_sfr_intercube_ring(
+            engine, spec, cfg, submesh_shape=(2, 4),
+        )
+        assert _qp(engine, 0, 0, 0)["W"]["peer"].cube == 4
+        assert _qp(engine, 0, 7, 0)["W"]["peer"].cube == 3
+        assert _qp(engine, 0, 6, 0)["W"]["peer"].cube == 7
+        assert _qp(engine, 0, 4, 0)["W"]["peer"].cube == 5
+
+    def test_snake_ring_2x4_interior(self):
+        """Interior hops along each row (not at row-bridge or wrap)."""
+        engine, spec = _engine_and_spec()
+        cfg = _merged_cfg()
+        configure_sfr_intercube_ring(
+            engine, spec, cfg, submesh_shape=(2, 4),
+        )
+        # Row 0 interior: cube1 between cube0 and cube2.
+        assert _qp(engine, 0, 1, 0)["E"]["peer"].cube == 2
+        assert _qp(engine, 0, 1, 0)["W"]["peer"].cube == 0
+        # Row 1 interior: cube5 between cube6 and cube4 along the snake.
+        assert _qp(engine, 0, 5, 0)["E"]["peer"].cube == 4
+        assert _qp(engine, 0, 5, 0)["W"]["peer"].cube == 6
+
+    def test_snake_ring_2x4_all_hops_are_1_hop(self):
+        """Every E hop on the snake is L1-distance 1 in the cube mesh.
+
+        Defensive — holds by construction (snake = Hamiltonian cycle
+        on the 2×4 grid graph), but explicit guards against future
+        path-builder regressions.
+        """
+        engine, spec = _engine_and_spec()
+        cfg = _merged_cfg()
+        configure_sfr_intercube_ring(
+            engine, spec, cfg, submesh_shape=(2, 4),
+        )
+        for cube in SNAKE_2X4_O00:
+            qp = _qp(engine, 0, cube, 0)
+            e_peer = qp["E"]["peer"].cube
+            w_peer = qp["W"]["peer"].cube
+            assert _l1(cube, e_peer) == 1, (
+                f"snake E hop cube{cube}→cube{e_peer} is not 1-hop"
+            )
+            assert _l1(cube, w_peer) == 1, (
+                f"snake W hop cube{cube}→cube{w_peer} is not 1-hop"
+            )
+
+    def test_snake_ring_2x4_off_path_cubes_have_no_ring_links(self):
+        """CUBEs not on the snake path (rows 2-3, i.e. cubes 8..15)
+        get no E/W ring entries — consistent with the current
+        ``if cube < ring_size`` behaviour."""
+        engine, spec = _engine_and_spec()
+        cfg = _merged_cfg()
+        configure_sfr_intercube_ring(
+            engine, spec, cfg, submesh_shape=(2, 4),
+        )
+        for cube in range(8, 16):
+            qp = _qp(engine, 0, cube, 0)
+            assert "E" not in qp, (
+                f"sip0.cube{cube}.pe0 has E but is off the snake path"
+            )
+            assert "W" not in qp
+
+    def test_snake_ring_2x4_intra_namespace_unchanged(self):
+        """Snake doesn't disturb the intra-cube 2×4 PE grid wiring
+        (intra_N/S/E/W) — it only changes cube-level E/W."""
+        from kernbench.ccl.sfr_config import _intra_cube_neighbors
+
+        engine, spec = _engine_and_spec()
+        cfg = _merged_cfg()
+        configure_sfr_intercube_ring(
+            engine, spec, cfg, submesh_shape=(2, 4),
+        )
+        # Spot-check a few cubes (including off-path) and all 8 PEs.
+        for cube in (0, 3, 4, 7, 10, 15):
+            for pe in range(PES_PER_CUBE):
+                qp = _qp(engine, 0, cube, pe)
+                expected = _intra_cube_neighbors(pe)
+                for d, expected_pe in expected.items():
+                    assert d in qp, f"cube{cube}.pe{pe} missing {d}"
+                    assert qp[d]["peer"].pe == expected_pe
+                    assert qp[d]["peer"].cube == cube
+
+
+class TestSnakeRing2x4OriginShift:
+    """Snake ring through rows 1-2 (origin=(1, 0)) — proves the
+    origin parameter shifts the sub-mesh."""
+
+    def test_snake_ring_origin_shift_path(self):
+        """For origin=(1, 0), snake path is [4,5,6,7,11,10,9,8]:
+
+          row 1 left→right:  cube4 → cube5 → cube6 → cube7
+                                                       ↓
+          row 2 right→left:  cube11 → cube10 → cube9 → cube8
+                                                          ↺ wraps to cube4
+        """
+        engine, spec = _engine_and_spec()
+        cfg = _merged_cfg()
+        configure_sfr_intercube_ring(
+            engine, spec, cfg, submesh_shape=(2, 4), submesh_origin=(1, 0),
+        )
+        # Path corners.
+        assert _qp(engine, 0, 4, 0)["E"]["peer"].cube == 5
+        assert _qp(engine, 0, 7, 0)["E"]["peer"].cube == 11
+        assert _qp(engine, 0, 11, 0)["E"]["peer"].cube == 10
+        assert _qp(engine, 0, 8, 0)["E"]["peer"].cube == 4  # wrap
+
+    def test_snake_ring_origin_shift_off_path_no_links(self):
+        """For origin=(1, 0), cubes 0..3 (row 0) and 12..15 (row 3)
+        get no ring links."""
+        engine, spec = _engine_and_spec()
+        cfg = _merged_cfg()
+        configure_sfr_intercube_ring(
+            engine, spec, cfg, submesh_shape=(2, 4), submesh_origin=(1, 0),
+        )
+        for cube in list(range(0, 4)) + list(range(12, 16)):
+            qp = _qp(engine, 0, cube, 0)
+            assert "E" not in qp
+            assert "W" not in qp
+
+
+class TestSnakeRingBackwardCompat:
+    """The existing 1D-row API (no submesh_shape) must behave identically
+    to before the snake extension lands."""
+
+    def test_snake_ring_backward_compat_1d_default(self):
+        """ring_size=4 (single-row) without submesh_shape — identical
+        behaviour to today's 1D-row ring: cube0..3 with E/W wrap;
+        cubes 4..15 have no ring links."""
+        engine, spec = _engine_and_spec()
+        cfg = _merged_cfg()
+        configure_sfr_intercube_ring(
+            engine, spec, cfg, ring_size=4,
+        )
+        # Row-0 1D ring with wrap.
+        assert _qp(engine, 0, 0, 0)["E"]["peer"].cube == 1
+        assert _qp(engine, 0, 0, 0)["W"]["peer"].cube == 3  # wrap
+        assert _qp(engine, 0, 3, 0)["E"]["peer"].cube == 0  # wrap
+        # Off-row cubes get no ring links.
+        for cube in range(4, 16):
+            qp = _qp(engine, 0, cube, 0)
+            assert "E" not in qp
+            assert "W" not in qp
+
+
+class TestSnakeRingValidation:
+    """Input validation of the submesh_shape / ring_size combinations."""
+
+    def test_snake_ring_invalid_submesh_overflow(self):
+        """submesh_shape that doesn't fit the cube mesh is rejected."""
+        engine, spec = _engine_and_spec()
+        cfg = _merged_cfg()
+        with pytest.raises(ValueError, match=r"sub.?mesh"):
+            configure_sfr_intercube_ring(
+                engine, spec, cfg, submesh_shape=(2, 5),  # col overflow
+            )
+
+    def test_snake_ring_ring_size_mismatch(self):
+        """submesh_shape and ring_size disagree → ValueError."""
+        engine, spec = _engine_and_spec()
+        cfg = _merged_cfg()
+        with pytest.raises(ValueError, match=r"ring_size"):
+            configure_sfr_intercube_ring(
+                engine, spec, cfg,
+                submesh_shape=(2, 4), ring_size=4,
+            )