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/_gqa_attention_decode_long.py

@@ -4,14 +4,29 @@ Each rank holds an ``S_local = S_kv / (C·P)`` slice of K, V and the full
 Q (replicated). The local attention is computed via an S_kv-axis tile
 sweep (ADR-0063 §A.2) so per-rank scratch is bounded by ``TILE_S_KV``
 regardless of ``S_local``. The partial ``(m, ℓ, O)`` is then reduced
-to PE 0 of CUBE 0 via a 2-level chain (intra-CUBE row+col, then
+via a 2-level pattern (intra-CUBE row+col, then inter-CUBE), and the
-inter-CUBE), and the root writes the final output.
+root writes the final output.
+
+Inter-CUBE reduce (selected by the ``sub_w`` launch arg):
+  - ``sub_w == 0`` (default): 1D chain along ``W``, root at CUBE 0.
+    Used for ``C ∈ {1, 4}`` panels.
+  - ``sub_w >= 2``: ADR-0060 §4.2 prescribed **lrab-adapted center-root
+    mesh reduce** over a ``sub_w × sub_h`` sub-mesh (``sub_h = C //
+    sub_w``). Phase 1 row reduce converges at ``root_col = sub_w//2``;
+    Phase 2 col reduce on ``root_col`` converges at
+    ``root_row = sub_h//2``. Root cube id:
+    ``root_row * sub_w + root_col``. Used for the C=8 single-KV-group
+    LLaMA-3.1-70B target (sub_w=4, sub_h=2, root=cube 6).
+    Requires ``sub_w >= 2 and sub_h >= 2`` — degenerate 1×N / N×1
+    layouts must use the 1D-chain path with ``sub_w=0``.
 
 Topology / SFR:
   - Requires ``configure_sfr_intercube_multisip`` when ``P > 1`` or
     ``C > 1`` (provides disjoint ``intra_*`` and ``E/W/N/S`` namespaces).
   - Intra-CUBE PEs are arranged as a 2×4 grid (no wrap).
-  - Inter-CUBE CUBEs are arranged as a 1D row (no wrap).
+  - Inter-CUBE CUBEs: 1D row (``sub_w == 0``) or rectangular sub-mesh
+    of the SIP's 4×4 CUBE mesh starting at origin (0, 0) with
+    ``sub_w == mesh_w`` (``sub_w >= 2``).
 
 Layout caveats:
   - GEMMs use ``tl.dot`` (no composite epilogue / ``softmax_scale``).
@@ -47,10 +62,11 @@ def gqa_attention_decode_long_kernel(
     d_head: int,
     C: int,
     P: int,
+    sub_w: int = 0,
     *,
     tl,
 ) -> None:
-    """GQA decode with M-fold + S_kv tile sweep + 2-level chain reduce-to-root.
+    """GQA decode with M-fold + S_kv tile sweep + 2-level reduce-to-root.
 
     Tensor layout:
       Q : (T_q, h_q · d_head) replicated on every rank; loaded as
@@ -59,8 +75,26 @@ def gqa_attention_decode_long_kernel(
           loads its (d_head, S_local) slice via byte-conserving reshape.
       V : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank
           loads its (S_local, d_head) slice.
-      O : (T_q, h_q · d_head) — only PE 0 of CUBE 0 stores.
+      O : (T_q, h_q · d_head) — only PE 0 of the root CUBE stores
+          (CUBE 0 for ``sub_w=0``; geometric center cube for ``sub_w>=2``).
     """
+    if sub_w > 0:
+        sub_h = C // sub_w
+        if sub_w < 2 or sub_h < 2 or sub_w * sub_h != C:
+            raise ValueError(
+                f"sub_w={sub_w} requires sub_w>=2 and sub_h>=2 and "
+                f"sub_w*sub_h==C; got sub_h={sub_h}, C={C}. "
+                "Use sub_w=0 for the 1D-chain path."
+            )
+        root_col = sub_w // 2
+        root_row = sub_h // 2
+        root_cube = root_row * sub_w + root_col
+    else:
+        sub_h = 0
+        root_col = 0
+        root_row = 0
+        root_cube = 0
+
     G = h_q // h_kv
     n_ranks = C * P
     S_local = S_kv // n_ranks
@@ -161,8 +195,9 @@ def gqa_attention_decode_long_kernel(
             tl.send(dir="intra_N", src=l_local)
             tl.send(dir="intra_N", src=O_local)
 
-    # Level-1 inter-CUBE chain (along W, leftward; only PE 0 of each CUBE).
+    # Level-1 inter-CUBE reduce (only PE 0 of each CUBE participates).
-    if pe_id == 0 and C > 1:
+    if pe_id == 0 and C > 1 and sub_w == 0:
+        # 1D chain along W, leftward; root at CUBE 0.
         if cube_id < C - 1:
             with tl.scratch_scope():
                 m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
@@ -178,8 +213,141 @@ def gqa_attention_decode_long_kernel(
             tl.send(dir="W", src=m_local)
             tl.send(dir="W", src=l_local)
             tl.send(dir="W", src=O_local)
+    elif pe_id == 0 and sub_w > 0:
+        # ── ADR-0060 §4.2 lrab-adapted center-root mesh reduce ──
+        # 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. Reduce-only
+        # (Phases 3-5 dropped: no inter-SIP, no broadcast — attention
+        # needs the answer at one rank). Plain ``+`` replaced with the
+        # log-sum-exp ``_merge_running`` for (m, ℓ, O).
+        row = cube_id // sub_w
+        col = cube_id % sub_w
+
+        # Phase 1: row reduce — converge at col == root_col.
+        if col == 0 and root_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:
+            if root_col > 0:
+                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 sub_w - 1 > root_col:
+                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 and sub_w - 1 > root_col:
+            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:
+            if row == 0 and root_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:
+                if root_row > 0:
+                    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:
+                    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)
+            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)
 
     # ── Final normalise + store (root only) ──
-    if pe_id == 0 and cube_id == 0:
+    if pe_id == 0 and cube_id == root_cube:
         O_final = O_local / l_local
         tl.store(o_ptr, O_final)

src/kernbench/benches/_gqa_attention_prefill_long.py

@@ -16,10 +16,17 @@ persistent scratch is ``(m, ℓ, O)`` only (~1 KB); per-tile in-scope
 scratch is bounded by ``TILE_S_KV`` regardless of ``S_local``.
 
 Topology / SFR:
-  - Requires ``configure_sfr_intercube_ring(ring_size=C)`` (1D ring of
+  - Requires ``configure_sfr_intercube_ring`` — either ``ring_size=C``
-    C CUBEs with wrap at the CUBE level).
+    (1D row, ``C ≤ mesh_w``) or ``submesh_shape=(rows, cols)`` (snake
-  - Only PE 0 of each CUBE participates (head-parallel; intra-CUBE PE
+    Hamiltonian cycle through a rectangular sub-mesh, for ``C > mesh_w``).
-    parallelism is a separate phase).
+  - ``P == 1`` (default): only PE 0 of each CUBE participates;
+    intra-CUBE PE parallelism is disabled.
+  - ``P > 1``: all P PEs of each CUBE participate via query-axis
+    split (ADR-0060 §5.5 last bullet + §B-item-3) — each PE owns
+    ``T_q/P`` disjoint query rows. Output rows are disjoint per PE,
+    so no intra-CUBE reduce is needed. Each PE drives its own
+    same-lane ring via independent IPCQ channels (P parallel rings).
+    Requires ``T_q % P == 0``.
 
 Layout caveats:
   - GEMMs use ``tl.dot`` (no composite epilogue / ``softmax_scale``).
@@ -55,36 +62,56 @@ def gqa_attention_prefill_long_kernel(
     S_kv: int,
     d_head: int,
     C: int,
+    P: int = 1,
     *,
     tl,
 ) -> None:
     """Head-parallel prefill attention with tile-granular Ring KV (ADR-0060 §5.5).
 
-    Tensor layout:
+    Tensor layout (per-rank shapes):
-      Q : (T_q, d_head)           one head per CUBE; replicated.
+      Q : (T_q_local, d_head)     one head per CUBE; replicated within
+          a CUBE for P=1, or sharded ``pe="row_wise"`` for P>1 so each
+          PE owns ``T_q_local = T_q // P`` query rows.
       K : (S_kv, d_head)          sharded cube_row_wise → each CUBE owns
           (S_local, d_head). Tiles loaded as (d_head, tile_s) via
           byte-conserving reshape.
       V : (S_kv, d_head)          sharded cube_row_wise → each CUBE owns
           (S_local, d_head). Tiles loaded as (tile_s, d_head).
       O : (T_q * C, d_head)       sharded cube_row_wise → each CUBE
-          writes its own (T_q, d_head) slice. NO reduce.
+          writes its own ``T_q`` rows; for P>1 each PE writes a
+          disjoint ``(T_q_local, d_head)`` slice (no intra-CUBE
+          reduce — disjoint output rows). NO inter-CUBE reduce.
 
     Algorithm: nested loop over (ring_step k, tile_idx t). At k=0 each
-    CUBE loads its own block's tiles from HBM; at k > 0 it receives
+    rank loads its own block's tiles from HBM; at k > 0 it receives
     tiles from its E neighbour. Tiles are forwarded W to the next
     ring step. Each tile's partial is folded into the running
     (m, ℓ, O) via online-softmax merge.
     """
     pe_id = tl.program_id(axis=0)
-    # Head-parallel: only PE 0 of each CUBE participates.
+    if P == 1:
+        # Head-parallel only: PE 0 of each CUBE participates.
         if pe_id != 0:
             return
+        T_q_local = T_q
+    else:
+        # Intra-CUBE PE-SP (ADR-0060 §5.5 last bullet + §B-item-3):
+        # query-axis split across P PEs of each CUBE. Output rows are
+        # disjoint per PE ⇒ no intra-CUBE reduce. Each PE drives its
+        # own same-lane ring via independent IPCQ channels.
+        if pe_id >= P:
+            return
+        if T_q % P != 0:
+            raise ValueError(
+                f"T_q={T_q} must be divisible by P={P} when P > 1; "
+                "use P=1 for the PE-0-only fallback path."
+            )
+        T_q_local = T_q // P
 
     S_local = S_kv // C
     n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
     KV_ROW_BYTES = d_head * 2  # f16
-    Q = tl.load(q_ptr, shape=(T_q, d_head), dtype="f16")
+    Q = tl.load(q_ptr, shape=(T_q_local, d_head), dtype="f16")
 
     # ── Bootstrap: (t=0, k=0) — load my own tile 0, establish (m, ℓ, O) ──
     #

src/kernbench/benches/milestone_gqa_headline.py

@@ -6,15 +6,20 @@ 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):
+Restrictions:
-  - C ≤ 4 (single-row inter-CUBE ring SFR; multi-row deferred)
+  - 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_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
@@ -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:
+        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,
             )
-    dp_q = DPPolicy(cube="replicate", pe="replicate",
+        else:
-                    num_cubes=C, num_pes=1)
+            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)
+                    pe=o_pe, num_cubes=C, num_pes=P)
-    q = ctx.zeros((_T_Q_PREFILL, _D_HEAD),
+    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,
     )
 

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
+        ring_size: number of CUBEs in the ring. Defaults to the full
-            0..ring_size-1). Defaults to the full cube_mesh count.
+            cube_mesh count for the 1D-row case, or ``rows*cols`` for
-            Must be ≤ mesh_w (single row); multi-row rings span
+            the snake case. If passed alongside ``submesh_shape``, must
-            non-neighbour boundaries.
+            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,6 +301,31 @@ 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 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:
@@ -288,6 +333,10 @@ def configure_sfr_intercube_ring(
                 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) ──
+        # ── Cube ring (E/W along resolved ring_path, with wrap) ──
-        if cube < ring_size:
+        pos = ring_pos.get(cube)
-            nbrs["E"] = _pe_idx(sip, (cube + 1) % ring_size, pe)
+        if pos is not None:
-            nbrs["W"] = _pe_idx(sip, (cube - 1) % ring_size, pe)
+            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:

tests/attention/test_gqa_decode_long_2d_reduce.py

@@ -0,0 +1,191 @@
+"""Tests for lrab-adapted center-root inter-CUBE reduce in decode_long.
+
+Verifies the ADR-0060 §4.2 prescribed Level-1 collective for the
+single-KV-group LLaMA-3.1-70B target (C=8 over a 2×4 sub-mesh, root at
+the geometric center CUBE).
+
+Activation contract:
+  - ``sub_w == 0`` (default; omitted from launch args) → existing 1D
+    inter-CUBE chain that converges at CUBE 0. Byte-for-byte unchanged.
+  - ``sub_w >= 2`` with ``sub_h = C // sub_w >= 2`` → lrab-adapted
+    center-root mesh reduce. Root cube is
+    ``(sub_h//2)*sub_w + (sub_w//2)``.
+
+Degenerate (``sub_h < 2`` or ``sub_w < 2``) combinations are rejected at
+launch time per ADR-0060 §4.2 + CLAUDE.md "Simplicity First": those
+configurations belong to the 1D-chain code path, not lrab.
+
+Phase 1 (this commit): tests only — production code lands in Phase 2.
+T1, T2, T4 fail today (kernel signature does not accept ``sub_w``).
+T3 passes today as the backward-compat anchor for the existing
+1D-chain path.
+"""
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+import pytest
+
+from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel  # noqa: F401
+from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
+from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
+from kernbench.policy.placement.dp import DPPolicy
+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
+
+TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
+
+D_HEAD = 64
+DTYPE = "f16"
+_CUBE_RE = re.compile(r"\bcube(\d+)\b")
+
+
+def _ccl_cfg():
+    return resolve_algorithm_config(
+        load_ccl_config(), name="lrab_hierarchical_allreduce",
+    )
+
+
+def _engine_factory(t, d):
+    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
+
+
+def _dma_write_cubes(op_log) -> list[int]:
+    """Return the list of CUBE ids that emitted a ``dma_write``.
+
+    Decode's final-store path uses one ``tl.store`` from the root rank.
+    Multiple HBM-channel-level write records may correspond to the same
+    logical store; this just reports the source CUBE for each.
+    """
+    cubes: list[int] = []
+    for r in op_log:
+        if r.op_name != "dma_write":
+            continue
+        m = _CUBE_RE.search(r.component_id)
+        if m is not None:
+            cubes.append(int(m.group(1)))
+    return cubes
+
+
+def _run_decode_long(
+    *, C: int, P: int, S_kv: int, sub_w: int | None,
+):
+    """Drive a decode_long launch. ``sub_w=None`` means omit the arg
+    entirely (exercises the kernel's current default behaviour)."""
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    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" if C > 1 else "replicate",
+                         pe="row_wise", num_cubes=C, num_pes=P)
+        T_q = 1
+        h_q = 8
+        h_kv = 1
+        ctx.zeros((T_q, h_q * D_HEAD),
+                  dtype=DTYPE, dp=dp_full, name=f"q_dec_2d_c{C}")
+        k = ctx.zeros((S_kv, h_kv * D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"k_dec_2d_c{C}")
+        v = ctx.zeros((S_kv, h_kv * D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"v_dec_2d_c{C}")
+        o = ctx.empty((T_q, h_q * D_HEAD),
+                      dtype=DTYPE, dp=dp_full, name=f"o_dec_2d_c{C}")
+        q = ctx.zeros((T_q, h_q * D_HEAD),
+                      dtype=DTYPE, dp=dp_full, name=f"q_dec_2d_c{C}_2")
+        launch_args = [q, k, v, o, T_q, S_kv, h_q, h_kv, D_HEAD, C, P]
+        if sub_w is not None:
+            launch_args.append(sub_w)
+        ctx.launch(
+            f"gqa_decode_long_2d_c{C}_sw{sub_w}",
+            gqa_attention_decode_long_kernel,
+            *launch_args,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None), engine_factory=_engine_factory,
+    )
+
+
+# ── T1: C=8 lrab-adapted center-root completes ───────────────────────
+
+
+def test_decode_2d_sub_w4_sub_h2_completes():
+    """ADR-0060 §4.2 prescribed center-root mesh reduce at the
+    milestone target: C=8 over a 2×4 sub-mesh (sub_w=4, sub_h=2),
+    P=8 PEs per CUBE. With zero inputs, output is trivially zero;
+    the test guards that the kernel reaches completion without
+    deadlock or scratch overflow.
+    """
+    result = _run_decode_long(C=8, P=8, S_kv=2048, sub_w=4)
+    assert result.completion.ok, (
+        f"decode at C=8, sub_w=4 (lrab-adapted) must complete; "
+        f"got {result.completion}"
+    )
+
+
+# ── T2: root lives at the geometric center CUBE (cube 6) ─────────────
+
+
+def test_decode_2d_sub_w4_sub_h2_root_at_center_cube_6():
+    """For sub_w=4, sub_h=2: root_col=2, root_row=1, root_cube=6.
+
+    Only the root CUBE issues the final ``tl.store`` — every other
+    rank short-circuits. The dma_write records must come exclusively
+    from CUBE 6.
+    """
+    result = _run_decode_long(C=8, P=8, S_kv=2048, sub_w=4)
+    assert result.completion.ok
+    cubes = _dma_write_cubes(result.engine.op_log)
+    assert cubes, "expected at least one dma_write for the final O store"
+    distinct = set(cubes)
+    assert distinct == {6}, (
+        f"final dma_write must come exclusively from CUBE 6 (sub_w=4, "
+        f"sub_h=2 root); got cubes={sorted(distinct)}"
+    )
+
+
+# ── T3: backward-compat — sub_w omitted → existing 1D chain at C=4 ───
+
+
+def test_decode_2d_backward_compat_sub_w0_default():
+    """When ``sub_w`` is omitted from the launch, the kernel must
+    behave identically to today: 1D inter-CUBE chain along W,
+    converging at CUBE 0. This serves as a regression anchor — Phase 2
+    must not change the existing C=4 panel behaviour.
+
+    Passes today as well as after Phase 2.
+    """
+    result = _run_decode_long(C=4, P=8, S_kv=2048, sub_w=None)
+    assert result.completion.ok, (
+        f"decode at C=4 with default sub_w (1D chain) must complete; "
+        f"got {result.completion}"
+    )
+    cubes = _dma_write_cubes(result.engine.op_log)
+    assert cubes, "expected at least one dma_write for the final O store"
+    distinct = set(cubes)
+    assert distinct == {0}, (
+        f"1D-chain root must be CUBE 0; got cubes={sorted(distinct)}"
+    )
+
+
+# ── T4: degenerate sub_w configurations are rejected ─────────────────
+
+
+def test_decode_2d_invalid_sub_w_rejected_when_sub_h_lt_2():
+    """ADR-0060 §4.2 + CLAUDE.md "Simplicity First": only sub-meshes
+    with both sub_w >= 2 and sub_h >= 2 use lrab. C=4 with sub_w=4
+    would give sub_h=1 (degenerate; Phase 2 of lrab becomes a no-op).
+    Callers must use the 1D-chain path (sub_w=0/omitted) for that case.
+
+    The kernel must reject the degenerate combination with a clear
+    error rather than silently producing a 1D-chain result at the
+    wrong root location.
+    """
+    with pytest.raises((ValueError, AssertionError), match=r"sub_[wh]"):
+        _run_decode_long(C=4, P=8, S_kv=2048, sub_w=4)

tests/attention/test_gqa_prefill_long_c8_snake.py

@@ -0,0 +1,154 @@
+"""Tests for prefill_long Ring KV at C=8 over a 2×4 snake sub-mesh.
+
+ADR-0060 §5.5 design target for the single-KV-group LLaMA-3.1-70B
+configuration: ``C = G = 8`` (one Q head per CUBE), Ring KV rotates
+the 8 KV slices around all 8 CUBEs.
+
+Increment 1 wired ``configure_sfr_intercube_ring(submesh_shape=(2, 4))``
+to map the kernel's logical E/W to a snake/serpentine 1-hop path
+through the top 2×4 sub-mesh of the 4×4 SIP CUBE mesh. The kernel
+itself uses only logical ``dir="W"`` / ``dir="E"`` — the snake is
+transparent at kernel level.
+
+This file verifies the assembly works end-to-end at C=8:
+  T1  kernel completes (no scratch overflow, no deadlock)
+  T2  per-CUBE head-parallel output (8 dma_writes, one per CUBE)
+  T3  tile-granular ring traffic at C=8 follows the same
+      ``(C-1)·n_tiles·2·C`` formula as the existing tile-ring tests
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel  # noqa: F401
+from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
+from kernbench.ccl.sfr_config import configure_sfr_intercube_ring
+from kernbench.policy.placement.dp import DPPolicy
+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
+
+TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
+
+D_HEAD = 64
+DTYPE = "f16"
+
+
+def _ccl_cfg():
+    return resolve_algorithm_config(
+        load_ccl_config(), name="lrab_hierarchical_allreduce",
+    )
+
+
+def _engine_factory(t, d):
+    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
+
+
+def _count(op_log, name: str) -> int:
+    return sum(1 for r in op_log if r.op_name == name)
+
+
+def _run_prefill_c8_snake(*, T_q: int, S_kv: int):
+    """Head-parallel prefill at C=8 with snake-mapped Ring KV over the
+    top 2×4 sub-mesh of the 4×4 SIP CUBE mesh."""
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+    C = 8
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_ring(
+            ctx.engine, ctx.spec, _ccl_cfg(),
+            submesh_shape=(2, 4),
+        )
+        dp_q = DPPolicy(cube="replicate", pe="replicate",
+                        num_cubes=C, num_pes=1)
+        dp_kv = DPPolicy(cube="row_wise", pe="replicate",
+                         num_cubes=C, num_pes=1)
+        dp_o = DPPolicy(cube="row_wise", pe="replicate",
+                        num_cubes=C, num_pes=1)
+        q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp_q,
+                      name=f"q_c8_snake_t{T_q}_s{S_kv}")
+        k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
+                      name=f"k_c8_snake_t{T_q}_s{S_kv}")
+        v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
+                      name=f"v_c8_snake_t{T_q}_s{S_kv}")
+        o = ctx.empty((T_q * C, D_HEAD), dtype=DTYPE, dp=dp_o,
+                      name=f"o_c8_snake_t{T_q}_s{S_kv}")
+        ctx.launch(
+            f"gqa_prefill_long_c8_snake_t{T_q}_s{S_kv}",
+            gqa_attention_prefill_long_kernel,
+            q, k, v, o,
+            T_q, S_kv, D_HEAD, C,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None),
+        engine_factory=_engine_factory,
+    )
+
+
+# ── T1: kernel completes end-to-end at C=8 ───────────────────────────
+
+
+def test_prefill_long_c8_snake_completes():
+    """ADR-0060 §5.5 design target: prefill Ring KV at C=G=8 over the
+    2×4 snake sub-mesh. With zero inputs, output is trivially zero;
+    this test guards that the kernel reaches completion without
+    deadlock, scratch overflow, or routing failure.
+
+    Configuration: T_q=4, S_kv=8192 → S_local=1024, n_tiles=1
+    (TILE_S_KV=1024). Bounded scratch.
+    """
+    result = _run_prefill_c8_snake(T_q=4, S_kv=8192)
+    assert result.completion.ok, (
+        f"prefill at C=8 with snake ring must complete; "
+        f"got {result.completion}"
+    )
+
+
+# ── T2: 8 dma_writes (head-parallel, one head per CUBE) ──────────────
+
+
+def test_prefill_long_c8_snake_distributed_output_count():
+    """Head-parallel: each CUBE owns one Q head and writes its own
+    head's output (T_q, d_head) rows. No inter-CUBE reduce on the
+    output side — so ``dma_write_count == C == 8``.
+
+    This is the C=8 analogue of
+    ``test_prefill_long_tile_ring_dma_write_count`` at C=4.
+    """
+    result = _run_prefill_c8_snake(T_q=4, S_kv=8192)
+    assert result.completion.ok
+    n_writes = _count(result.engine.op_log, "dma_write")
+    assert n_writes == 8, (
+        f"head-parallel C=8: expected 8 dma_writes (one per CUBE); "
+        f"got {n_writes}"
+    )
+
+
+# ── T3: tile-granular ring ipcq_copy count scales as (C-1)·n_tiles·2·C
+
+
+def test_prefill_long_c8_snake_tile_ipcq_count():
+    """ADR-0060 §5.5.1 tile-granular ring: per-CUBE send count is
+    ``2·n_tiles·(C-1)`` (K + V per tile per ring step). Aggregated
+    across C CUBEs, total ipcq_copy = ``(C-1)·n_tiles·2·C``.
+
+    Configuration: T_q=4, S_kv=8192, C=8 → S_local=1024, n_tiles=1
+    (TILE_S_KV=1024). Expected total = ``7·1·2·8 = 112``.
+
+    Verifies that the snake-mapped 1-hop physical links carry the
+    same logical ring traffic as the existing C=4 1D-row ring tests.
+    """
+    C = 8
+    n_tiles = 1  # S_local=1024 / TILE_S_KV=1024
+    result = _run_prefill_c8_snake(T_q=4, S_kv=8192)
+    assert result.completion.ok
+    n_copy = _count(result.engine.op_log, "ipcq_copy")
+    expected = (C - 1) * n_tiles * 2 * C
+    assert n_copy == expected, (
+        f"tile-granular ring at C=8: expected {expected} ipcq_copy "
+        f"((C-1)·n_tiles·2·C = {C - 1}·{n_tiles}·2·{C}); got {n_copy}"
+    )

tests/attention/test_gqa_prefill_long_pe_sp.py

@@ -0,0 +1,233 @@
+"""Tests for intra-CUBE PE-SP in prefill_long (query-axis split).
+
+ADR-0060 §5.5 last bullet + §B-item-3: split the head's query rows
+``[T_q, d_head]`` across the ``P`` PEs of each CUBE so all P PEs work
+in parallel. Output rows are disjoint across PEs ⇒ no intra-CUBE
+reduce needed.
+
+Same-lane SFR wiring: PE ``i`` in CUBE A has its own E/W ring link to
+PE ``i`` in CUBE B (the snake's prev/next). All P PEs of a CUBE see
+the same K, V (HBM-resident, ``pe="replicate"``) ⇒ P parallel rings
+run in lockstep, each PE rotating its own K/V copies via its own IPCQ
+channels.
+
+Activation contract:
+  - ``P == 1`` (default; omitted from launch args) → existing
+    PE-0-only behavior. Byte-for-byte unchanged.
+  - ``P > 1`` → all P PEs active; each handles ``T_q // P`` query
+    rows. Requires ``T_q % P == 0`` (degenerate T_q < P is rejected
+    — caller must use ``P=1`` for that workload, ADR-0060 §B-item-3).
+
+Phase 1: tests only — production code lands in Phase 2.
+T1, T2, T3, T5 fail today (TypeError: kernel signature has no P).
+T4 passes today as the backward-compat anchor.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel  # noqa: F401
+from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
+from kernbench.ccl.sfr_config import configure_sfr_intercube_ring
+from kernbench.policy.placement.dp import DPPolicy
+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
+
+TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
+
+D_HEAD = 64
+DTYPE = "f16"
+
+
+def _ccl_cfg():
+    return resolve_algorithm_config(
+        load_ccl_config(), name="lrab_hierarchical_allreduce",
+    )
+
+
+def _engine_factory(t, d):
+    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
+
+
+def _count(op_log, name: str) -> int:
+    return sum(1 for r in op_log if r.op_name == name)
+
+
+def _run_prefill_pe_sp(
+    *, T_q: int, S_kv: int, C: int, P: int | None,
+    snake: bool,
+):
+    """Drive a prefill_long launch with optional intra-CUBE PE-SP.
+
+    ``P=None`` → omit P from the launch args (exercises the kernel's
+    default behaviour).
+    ``snake=True`` → install the 2×4 snake ring SFR (Increment 1);
+    else install the 1D-row ring at ``ring_size=C``.
+    """
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        if snake:
+            configure_sfr_intercube_ring(
+                ctx.engine, ctx.spec, _ccl_cfg(),
+                submesh_shape=(2, 4),
+            )
+        else:
+            configure_sfr_intercube_ring(
+                ctx.engine, ctx.spec, _ccl_cfg(),
+                ring_size=C,
+            )
+        num_pes = P if (P is not None and P > 1) else 1
+        # When PE-SP is active, Q and O are split row-wise across PEs;
+        # K, V remain replicated within a CUBE.
+        q_pe = "row_wise" if num_pes > 1 else "replicate"
+        o_pe = "row_wise" if num_pes > 1 else "replicate"
+        dp_q = DPPolicy(cube="replicate", pe=q_pe,
+                        num_cubes=C, num_pes=num_pes)
+        dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
+                         pe="replicate", num_cubes=C, num_pes=num_pes)
+        dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate",
+                        pe=o_pe, num_cubes=C, num_pes=num_pes)
+        suffix = f"c{C}_p{P}_t{T_q}_s{S_kv}"
+        q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp_q,
+                      name=f"q_pe_sp_{suffix}")
+        k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
+                      name=f"k_pe_sp_{suffix}")
+        v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
+                      name=f"v_pe_sp_{suffix}")
+        o = ctx.empty((T_q * C, D_HEAD), dtype=DTYPE, dp=dp_o,
+                      name=f"o_pe_sp_{suffix}")
+        launch_args = [q, k, v, o, T_q, S_kv, D_HEAD, C]
+        if P is not None:
+            launch_args.append(P)
+        ctx.launch(
+            f"gqa_prefill_long_pe_sp_{suffix}",
+            gqa_attention_prefill_long_kernel,
+            *launch_args,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None),
+        engine_factory=_engine_factory,
+    )
+
+
+# ── T1: C=8 P=8 completes end-to-end ─────────────────────────────────
+
+
+def test_prefill_c8_p8_pe_sp_completes():
+    """ADR-0060 §5.5 + §B-item-3 design target: 64 ranks active
+    (8 CUBEs × 8 PEs), query-axis split across PEs, snake-mapped
+    Ring KV. Guards no scratch overflow, no deadlock across the
+    P parallel same-lane rings.
+    """
+    result = _run_prefill_pe_sp(
+        T_q=8, S_kv=8192, C=8, P=8, snake=True,
+    )
+    assert result.completion.ok, (
+        f"prefill at C=8, P=8 (PE-SP) must complete; "
+        f"got {result.completion}"
+    )
+
+
+# ── T2: 64 dma_writes (one per PE, disjoint query rows) ──────────────
+
+
+def test_prefill_c8_p8_pe_sp_64_dma_writes():
+    """With query-axis split, each PE owns ``T_q/P = 1`` query row
+    and stores its own ``(1, d_head)`` slice. Across all 64 ranks
+    (8 CUBEs × 8 PEs), ``dma_write_count == 64``.
+
+    This is the structural signal that PE-SP is actually wired:
+    PE-0-only would give 8 dma_writes (one per CUBE).
+    """
+    result = _run_prefill_pe_sp(
+        T_q=8, S_kv=8192, C=8, P=8, snake=True,
+    )
+    assert result.completion.ok
+    n_writes = _count(result.engine.op_log, "dma_write")
+    assert n_writes == 64, (
+        f"PE-SP at C=8, P=8: expected 64 dma_writes "
+        f"(8 cubes × 8 PEs, each writes its T_q/P=1 query row); "
+        f"got {n_writes}"
+    )
+
+
+# ── T3: P parallel rings — ipcq_copy scales by P ─────────────────────
+
+
+def test_prefill_c8_p8_pe_sp_ring_ipcq_count():
+    """Per-PE same-lane rings: each PE runs its own ring traffic via
+    its own IPCQ channels. Total inter-CUBE ipcq_copy:
+    ``(C-1) · n_tiles · 2 · C · P`` (= existing C=8 formula × P).
+
+    Configuration: T_q=8, S_kv=8192, C=8 → S_local=1024, n_tiles=1
+    (TILE_S_KV=1024). Expected = ``7·1·2·8·8`` = **896**.
+    """
+    C = 8
+    P = 8
+    n_tiles = 1
+    result = _run_prefill_pe_sp(
+        T_q=8, S_kv=8192, C=C, P=P, snake=True,
+    )
+    assert result.completion.ok
+    n_copy = _count(result.engine.op_log, "ipcq_copy")
+    expected = (C - 1) * n_tiles * 2 * C * P
+    assert n_copy == expected, (
+        f"PE-SP parallel rings at C=8, P=8: expected {expected} "
+        f"ipcq_copy ((C-1)·n_tiles·2·C·P = "
+        f"{C - 1}·{n_tiles}·2·{C}·{P}); got {n_copy}"
+    )
+
+
+# ── T4: backward-compat — P omitted → existing PE-0-only behaviour ───
+
+
+def test_prefill_p_default_1_backward_compat():
+    """When ``P`` is omitted from the launch args, the kernel must
+    behave identically to today: only PE 0 of each CUBE participates;
+    one head per CUBE; one dma_write per CUBE.
+
+    At C=4 this gives 4 dma_writes (matches the existing
+    ``test_prefill_long_tile_ring_dma_write_count``). Phase 2 must
+    not regress this path.
+
+    Passes today AND after Phase 2.
+    """
+    C = 4
+    result = _run_prefill_pe_sp(
+        T_q=4, S_kv=8192, C=C, P=None, snake=False,
+    )
+    assert result.completion.ok, (
+        f"prefill at C=4 with default P (PE-0-only) must complete; "
+        f"got {result.completion}"
+    )
+    n_writes = _count(result.engine.op_log, "dma_write")
+    assert n_writes == C, (
+        f"default P=1 (PE-0-only): expected {C} dma_writes "
+        f"(one per CUBE); got {n_writes}"
+    )
+
+
+# ── T5: validation — T_q must be divisible by P when P > 1 ───────────
+
+
+def test_prefill_pe_sp_rejects_non_divisible_t_q():
+    """ADR-0060 §B-item-3 fallback (KV-block split + intra-CUBE
+    reduce for T_q < P) is deferred. Callers must request ``P=1`` for
+    workloads where T_q < P or T_q % P != 0.
+
+    C=4, P=8, T_q=4 violates ``T_q % P == 0`` (and also T_q < P).
+    The kernel must raise ValueError with a clear error message
+    rather than silently producing wrong results.
+    """
+    with pytest.raises((ValueError, AssertionError), match=r"T_q"):
+        _run_prefill_pe_sp(
+            T_q=4, S_kv=8192, C=4, P=8, snake=False,
+        )

tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py

@@ -0,0 +1,146 @@
+"""Tests for the single-KV-group prefill panel in milestone_gqa_headline.
+
+The single-KV-group (LLaMA-3.1-70B target) prefill panel wires together
+all Increment 1–4 work in a single bench panel:
+  - C=8 head-parallel + Ring KV via snake-mapped 2×4 sub-mesh (Inc 1, 3)
+  - P=8 intra-CUBE PE-SP (Inc 4)
+  - d_head=128 (LLaMA-3.1-70B), one-shot prefill T_q = S_kv = 32K
+
+This file verifies the bench-config wiring:
+  T1  the new ``single_kv_group_prefill_gqa_c8_p8`` panel is registered
+      with the expected dims in ``_PANELS`` + ``_PANEL_DISPATCH``.
+  T2  ``_run_prefill_panel`` accepts the new ``P``, ``T_q``, ``d_head``
+      kwargs and drives the prefill kernel to completion at the
+      milestone-target ``(C, P) = (8, 8)``. Uses smaller T_q/S_kv to
+      keep test time bounded — full 32K runs come from
+      ``kernbench run milestone-gqa-headline``.
+  T3  Existing prefill panels still work when ``_run_prefill_panel``
+      is called without the new kwargs (backward compat anchor).
+
+Phase 1: tests only — production code lands in Phase 2.
+T1 fails today (panel not registered).
+T2 fails today (helper doesn't accept the new kwargs).
+T3 passes today as the backward-compat anchor.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+from kernbench.benches.milestone_gqa_headline import (  # noqa: F401
+    _PANEL_DISPATCH,
+    _PANELS,
+    _run_prefill_panel,
+)
+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
+
+TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
+
+
+def _engine_factory(t, d):
+    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
+
+
+# ── T1: new panel registered ─────────────────────────────────────────
+
+
+def test_single_kv_group_prefill_panel_registered():
+    """The ``single_kv_group_prefill_gqa_c8_p8`` panel must be in ``_PANELS`` and
+    its dispatch entry must have the expected LLaMA-3.1-70B params.
+
+    Headline config:
+      C = 8        (head-parallel, snake sub-mesh)
+      P = 8        (intra-CUBE PE-SP)
+      T_q = 32_768 (one-shot long-context prefill)
+      S_kv = 32_768
+      d_head = 128 (LLaMA-3.1-70B)
+    """
+    panel_name = "single_kv_group_prefill_gqa_c8_p8"
+    assert panel_name in _PANELS, (
+        f"{panel_name!r} not in _PANELS; got {_PANELS}"
+    )
+    assert panel_name in _PANEL_DISPATCH, (
+        f"{panel_name!r} not in _PANEL_DISPATCH"
+    )
+    kind, params = _PANEL_DISPATCH[panel_name]
+    assert kind == "prefill", f"kind={kind!r}, expected 'prefill'"
+    assert params.get("C") == 8, f"C={params.get('C')}, expected 8"
+    assert params.get("P") == 8, f"P={params.get('P')}, expected 8"
+    assert params.get("T_q") == 32_768, (
+        f"T_q={params.get('T_q')}, expected 32_768"
+    )
+    assert params.get("S_kv") == 32_768, (
+        f"S_kv={params.get('S_kv')}, expected 32_768"
+    )
+    assert params.get("d_head") == 128, (
+        f"d_head={params.get('d_head')}, expected 128"
+    )
+
+
+# ── T2: helper drives the C=8, P=8 prefill kernel to completion ──────
+
+
+def test_single_kv_group_prefill_panel_runner_smoke():
+    """``_run_prefill_panel`` must accept the new ``P``, ``T_q``,
+    ``d_head`` kwargs and successfully launch the prefill kernel at
+    ``(C, P) = (8, 8)`` with the snake-mapped 2×4 SFR.
+
+    Uses **smaller** T_q/S_kv than the headline panel so the test
+    completes quickly. The headline 32K dims are exercised via
+    ``kernbench run milestone-gqa-headline``, not pytest.
+    """
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        _run_prefill_panel(
+            ctx,
+            panel="single_kv_group_prefill_gqa_c8_p8",
+            C=8, P=8,
+            T_q=8, S_kv=8192,   # smoke dims, not the 32K headline
+            d_head=128,
+        )
+
+    result = run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None),
+        engine_factory=_engine_factory,
+    )
+    assert result.completion.ok, (
+        f"single_kv_group prefill panel smoke at C=8 P=8 must complete; "
+        f"got {result.completion}"
+    )
+
+
+# ── T3: existing panel calls still work (backward compat) ────────────
+
+
+def test_existing_prefill_panel_runner_backward_compat():
+    """Existing callers of ``_run_prefill_panel`` (no ``P``/``T_q``/
+    ``d_head`` overrides) must continue to work after the signature
+    extension. Exercises the ``multi_user_prefill_gqa`` panel dims
+    from ``_PANEL_DISPATCH``.
+
+    Passes today AND after Phase 2.
+    """
+    panel_name = "multi_user_prefill_gqa"
+    _kind, params = _PANEL_DISPATCH[panel_name]
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        _run_prefill_panel(
+            ctx,
+            panel=panel_name,
+            C=params["C"], S_kv=params["S_kv"],
+        )
+
+    result = run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None),
+        engine_factory=_engine_factory,
+    )
+    assert result.completion.ok, (
+        f"existing prefill panel {panel_name!r} must still run with "
+        f"the original kwargs; got {result.completion}"
+    )

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,
+            )