gqa: fold decode_long into the Case 4 kernel; drop 1D-chain dead code

Option B fold: the lrab math previously in
``_gqa_attention_decode_long.py`` (used only as a thin re-export by
the Case 4 wrapper and as the import source for ``_merge_running``
in Cases 1 / 3) now lives inline in the Case 4 kernel file. ``sub_w``
is hardcoded to 4 (the 4×2 cube sub-mesh geometry; root cube 6);
the legacy ``sub_w=0`` 1D-chain backward-compat path is removed
along with its dedicated tests — the dropped ``single_user_*`` /
``multi_user_*`` panels that exercised it are already gone.

Production changes:
  - inline lrab math into _gqa_attention_decode_long_ctx_cube_sp_pe_sp.py
    (drops sub_w param; _ROOT_CUBE=6 baked in)
  - inline _merge_running into Cases 1 (cube_sp_pe_tp) and 3
    (cube_repl_pe_sp) so they no longer depend on the deleted file
  - delete src/kernbench/benches/_gqa_attention_decode_long.py
  - remove dead _run_decode_panel / _DECODE_* constants /
    decode-side _PANEL_DISPATCH / _make_bench_fn decode branch
    from milestone_gqa_headline.py (only the prefill panel remains)
  - update _gqa_attention_decode_opt2.py docstring reference

Test changes:
  - delete 7 legacy test_gqa_*.py files that pre-dated the 4-cases
    architectural split (coverage now subsumed by the 16 4-cases tests)
  - remove test_opt2_matches_opt3_data_mode + _run_decode_data helper
    from test_gqa_decode_opt2.py (the parity check required sub_w=0
    which no longer exists; opt2's other 4 tests preserved)

20/20 tests pass (16 4-cases + 4 opt2 smoke/dispatch).

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

scripts/paper/paper_plot_gqa.py

@@ -24,10 +24,8 @@ _FIG_DIR = Path(__file__).resolve().parents[2] / "docs" / "report" / "1H-codesig
 _IN_JSON = _FIG_DIR / "gqa_latency.json"
 
 _LABELS = {
-    "single_user_prefill_gqa": "prefill\nC=1",
+    "single_kv_group_prefill_gqa_c8_p8":
-    "multi_user_prefill_gqa": "prefill\nC=4 (Ring KV)",
+        "prefill\nC=8, P=8\n(single KV group)",
-    "single_user_decode_gqa": "decode\nC=1, P=8",
-    "multi_user_decode_gqa": "decode\nC=4, P=8",
 }
 
 

src/kernbench/benches/1H_milestone_output/gqa_headline/sweep.json

@@ -1,10 +1,7 @@
 {
   "version": 1,
   "panels": [
-    "single_user_prefill_gqa",
+    "single_kv_group_prefill_gqa_c8_p8"
-    "multi_user_prefill_gqa",
-    "single_user_decode_gqa",
-    "multi_user_decode_gqa"
   ],
   "config": {
     "T_q_prefill": 4,
@@ -16,53 +13,18 @@
   },
   "rows": [
     {
-      "panel": "single_user_prefill_gqa",
+      "panel": "single_kv_group_prefill_gqa_c8_p8",
       "kind": "prefill",
-      "C": 1,
+      "C": 8,
-      "S_kv": 16,
-      "op_log_summary": {
-        "gemm_count": 2,
-        "ipcq_copy_count": 0,
-        "dma_read_count": 3,
-        "dma_write_count": 1
-      }
-    },
-    {
-      "panel": "multi_user_prefill_gqa",
-      "kind": "prefill",
-      "C": 4,
-      "S_kv": 16,
-      "op_log_summary": {
-        "gemm_count": 32,
-        "ipcq_copy_count": 24,
-        "dma_read_count": 12,
-        "dma_write_count": 4
-      }
-    },
-    {
-      "panel": "single_user_decode_gqa",
-      "kind": "decode",
-      "C": 1,
       "P": 8,
-      "S_kv": 64,
+      "T_q": 1024,
+      "S_kv": 1024,
+      "d_head": 128,
       "op_log_summary": {
-        "gemm_count": 16,
+        "gemm_count": 1024,
-        "ipcq_copy_count": 21,
+        "ipcq_copy_count": 896,
-        "dma_read_count": 24,
+        "dma_read_count": 192,
-        "dma_write_count": 1
+        "dma_write_count": 64
-      }
-    },
-    {
-      "panel": "multi_user_decode_gqa",
-      "kind": "decode",
-      "C": 4,
-      "P": 8,
-      "S_kv": 128,
-      "op_log_summary": {
-        "gemm_count": 64,
-        "ipcq_copy_count": 93,
-        "dma_read_count": 96,
-        "dma_write_count": 1
       }
     }
   ]

src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_repl_pe_sp.py

@@ -1,23 +1,27 @@
-"""GQA fused-attention decode kernel — long context (ADR-0060).
+"""GQA decode kernel — Case 3 (Cube-Repl × PE-SP).
 
-Each rank holds an ``S_local = S_kv / (C·P)`` slice of K, V and the full
+Per GQA_full_deck.pptx slide 11:
-Q (replicated). The local attention is computed via an S_kv-axis tile
+  - K, V replicated across all 8 cubes (the 8× memory waste).
-sweep (ADR-0063 §A.2) so per-rank scratch is bounded by ``TILE_S_KV``
+  - PEs SP on S_kv inside each cube: each PE attends to its
-regardless of ``S_local``. The partial ``(m, ℓ, O)`` is then reduced
+    ``S_local = S_kv / P`` slice.
-to PE 0 of CUBE 0 via a 2-level chain (intra-CUBE row+col, then
+  - Intra-CUBE 8-way reduce on the partial ``(m, ℓ, O)`` triple
-inter-CUBE), and the root writes the final output.
+    (row chain + col bridge over the 2×4 PE grid, same structural
+    pattern as Case 4's intra-CUBE phase).
+  - NO inter-CUBE communication — every cube ends with the full
+    answer (8× redundant compute is the inherent cost of Case 3).
+  - Designated writer: only PE 0 of CUBE 0 stores ``O`` to avoid
+    8 redundant DMA writes.
+
+Tensor layout:
+  Q : (T_q, h_q · d_head) replicated on every rank.
+  K : (S_kv, h_kv · d_head) with cube=replicate, pe=row_wise —
+      each PE owns its (S_local, h_kv·d_head) shard within its cube.
+  V : same as K.
+  O : (T_q, h_q · d_head) — only PE 0 of CUBE 0 stores.
 
 Topology / SFR:
-  - Requires ``configure_sfr_intercube_multisip`` when ``P > 1`` or
+  - ``configure_sfr_intercube_multisip`` provides the intra_* /
-    ``C > 1`` (provides disjoint ``intra_*`` and ``E/W/N/S`` namespaces).
+    E/W/N/S namespaces; only the intra_* lanes are exercised here.
-  - Intra-CUBE PEs are arranged as a 2×4 grid (no wrap).
-  - Inter-CUBE CUBEs are arranged as a 1D row (no wrap).
-
-Layout caveats:
-  - GEMMs use ``tl.dot`` (no composite epilogue / ``softmax_scale``).
-  - K is loaded as ``(d_head, S_local)`` via byte-conserving reshape of
-    the deployed ``(S_local, h_kv·d_head)`` slice — correct for zero /
-    symmetric inputs (ADR-0060 §3 reshape-not-transpose caveat).
 """
 from __future__ import annotations
 
@@ -35,7 +39,7 @@ def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
     return m_new, l_new, O_new
 
 
-def gqa_attention_decode_long_kernel(
+def gqa_attention_decode_long_ctx_cube_repl_pe_sp_kernel(
     q_ptr: int,
     k_ptr: int,
     v_ptr: int,
@@ -50,20 +54,9 @@ def gqa_attention_decode_long_kernel(
     *,
     tl,
 ) -> None:
-    """GQA decode with M-fold + S_kv tile sweep + 2-level chain reduce-to-root.
+    """Case-3 decode: PE-SP attention; replicated KV per cube; intra-CUBE AR only."""
-
-    Tensor layout:
-      Q : (T_q, h_q · d_head) replicated on every rank; loaded as
-          (G·T_q, d_head) — byte-conserving and math-correct for T_q=1.
-      K : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank
-          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.
-    """
     G = h_q // h_kv
-    n_ranks = C * P
+    S_local = S_kv // P  # each PE owns S_kv/P (cube=replicate, pe=row_wise)
-    S_local = S_kv // n_ranks
     pe_id = tl.program_id(axis=0)
     cube_id = tl.program_id(axis=1)
 
@@ -72,18 +65,6 @@ def gqa_attention_decode_long_kernel(
     n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
     KV_ROW_BYTES = d_head * 2  # f16
 
-    # Tile 0: establishes persistent (m_local, l_local, O_local).
-    #
-    # Cannot be folded into the Tiles 1..N loop (kernbench-only limitation):
-    #   - persistent (m, ℓ, O) must live OUTSIDE ``tl.scratch_scope``,
-    #     otherwise scope teardown discards them before the next tile's
-    #     merge can read them;
-    #   - kernbench has no scratch-backed initializer — ``tl.zeros`` /
-    #     ``tl.full`` return addr=0 handles with no backing storage, so
-    #     they cannot be overwritten via ``tl.copy_to`` to seed (-inf, 0, 0).
-    # So Tile 0 computes the initial running state directly; Tiles 1..N
-    # fold into it. Triton port: limitation does not apply (SSA tensors
-    # stay live across iterations) — a single unified loop suffices.
     tile_s0 = min(TILE_S_KV, S_local)
     K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
     V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
@@ -94,9 +75,6 @@ def gqa_attention_decode_long_kernel(
     l_local = tl.sum(exp_scores, axis=-1)
     O_local = tl.dot(exp_scores, V)
 
-    # Tiles 1..n_tiles-1: fold into running state via online-softmax merge.
-    # Triton port: drop the ``with tl.scratch_scope():`` line and replace
-    # each ``copy_to`` with a Python rebind.
     for tile_idx in range(1, n_tiles):
         tile_start = tile_idx * TILE_S_KV
         tile_s = min(TILE_S_KV, S_local - tile_start)
@@ -118,14 +96,13 @@ def gqa_attention_decode_long_kernel(
             tl.copy_to(l_local, l_new)
             tl.copy_to(O_local, O_new)
 
-    # ── Communication: chain reduce-to-root at PE 0 of CUBE 0 ──
+    # ── Intra-CUBE 8-way reduce-to-PE0 (row chain + col bridge) ──
     PE_GRID_COLS = 4
     pe_col = pe_id % PE_GRID_COLS
     pe_row = pe_id // PE_GRID_COLS
     pe_cols_used = min(PE_GRID_COLS, P)
     pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS
 
-    # Level-2 row chain (intra-CUBE, along intra_W, leftward).
     if pe_cols_used > 1:
         if pe_col < pe_cols_used - 1:
             with tl.scratch_scope():
@@ -143,7 +120,6 @@ def gqa_attention_decode_long_kernel(
             tl.send(dir="intra_W", src=l_local)
             tl.send(dir="intra_W", src=O_local)
 
-    # Level-2 col bridge (intra-CUBE, along intra_N, row-1 → row-0).
     if pe_col == 0 and pe_rows_used > 1:
         if pe_row < pe_rows_used - 1:
             with tl.scratch_scope():
@@ -161,25 +137,7 @@ 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).
+    # ── Final normalise + store (designated writer: cube 0, PE 0) ──
-    if pe_id == 0 and C > 1:
-        if cube_id < C - 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)
-        if cube_id > 0:
-            tl.send(dir="W", src=m_local)
-            tl.send(dir="W", src=l_local)
-            tl.send(dir="W", src=O_local)
-
-    # ── Final normalise + store (root only) ──
     if pe_id == 0 and cube_id == 0:
         O_final = O_local / l_local
         tl.store(o_ptr, O_final)

src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_repl_pe_tp.py

@@ -0,0 +1,101 @@
+"""GQA decode kernel — Case 2 (Cube-Repl × PE-TP).
+
+Per GQA_full_deck.pptx slide 11:
+  - K, V replicated across all 8 cubes × 8 PEs (the 8 KB/tok/PE
+    memory waste — this is the inherent cost Case 2 demonstrates).
+  - PEs nominally split on the batch dim (PE-TP). For B=1 (single-
+    user decode, the slide-11 default), only PE 0 of CUBE 0 has
+    work; the other 63 ranks idle.
+  - NO inter-rank communication (each rank has full KV — slide 11
+    lists comm cost as "none").
+
+This kernel is the simplest of the 4 cases by design: one active
+rank does the full attention locally; everyone else early-returns.
+The DPPolicy at the call site models the cluster-wide 8× memory
+waste even though only one rank reads from HBM.
+
+Tensor layout (B=1):
+  Q : (T_q, h_q · d_head) replicated on every rank; loaded as
+      (G · T_q, d_head) on the active rank.
+  K : (S_kv, h_kv · d_head) replicated on every rank.
+  V : (S_kv, h_kv · d_head) replicated on every rank.
+  O : (T_q, h_q · d_head) — only PE 0 of CUBE 0 stores.
+
+Topology / SFR:
+  - configure_sfr_intercube_multisip is fine but not strictly required
+    (no inter-rank sends/recvs happen).
+"""
+from __future__ import annotations
+
+
+TILE_S_KV = 1024  # match decode_long — per-tile S_kv width (ADR-0063 §A.2).
+
+
+def gqa_attention_decode_long_ctx_cube_repl_pe_tp_kernel(
+    q_ptr: int,
+    k_ptr: int,
+    v_ptr: int,
+    o_ptr: int,
+    T_q: int,
+    S_kv: int,
+    h_q: int,
+    h_kv: int,
+    d_head: int,
+    C: int,
+    P: int,
+    *,
+    tl,
+) -> None:
+    """Case-2 decode: single-rank attention; full KV per rank; no comm."""
+    pe_id = tl.program_id(axis=0)
+    cube_id = tl.program_id(axis=1)
+    # B=1 single-user decode + PE-TP: only one rank has work.
+    # Slide-11 acknowledges this PE-TP waste at B=1.
+    if pe_id != 0 or cube_id != 0:
+        return
+
+    G = h_q // h_kv
+    n_tiles = (S_kv + TILE_S_KV - 1) // TILE_S_KV
+    KV_ROW_BYTES = d_head * 2  # f16
+
+    # ── Load Q (full; replicated; M-folded for GQA reuse) ──
+    Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
+
+    # ── Tile 0: bootstrap persistent (m, ℓ, O) ──
+    tile_s0 = min(TILE_S_KV, S_kv)
+    K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
+    V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
+    scores = tl.dot(Q, K_T)
+    m_local = tl.max(scores, axis=-1)
+    centered = scores - m_local
+    exp_scores = tl.exp(centered)
+    l_local = tl.sum(exp_scores, axis=-1)
+    O_local = tl.dot(exp_scores, V)
+
+    # ── Tiles 1..N: fold via online-softmax merge in scratch_scope ──
+    for tile_idx in range(1, n_tiles):
+        tile_start = tile_idx * TILE_S_KV
+        tile_s = min(TILE_S_KV, S_kv - tile_start)
+        with tl.scratch_scope():
+            K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
+                            shape=(d_head, tile_s), dtype="f16")
+            V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
+                          shape=(tile_s, d_head), dtype="f16")
+            scores_t = tl.dot(Q, K_T_t)
+            m_tile = tl.max(scores_t, axis=-1)
+            centered_t = scores_t - m_tile
+            exp_scores_t = tl.exp(centered_t)
+            l_tile = tl.sum(exp_scores_t, axis=-1)
+            O_tile = tl.dot(exp_scores_t, V_t)
+            m_new = tl.maximum(m_local, m_tile)
+            scale_old = tl.exp(m_local - m_new)
+            scale_new = tl.exp(m_tile - m_new)
+            l_new = l_local * scale_old + l_tile * scale_new
+            O_new = O_local * scale_old + O_tile * scale_new
+            tl.copy_to(m_local, m_new)
+            tl.copy_to(l_local, l_new)
+            tl.copy_to(O_local, O_new)
+
+    # ── Final normalise + store ──
+    O_final = O_local / l_local
+    tl.store(o_ptr, O_final)

src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_sp_pe_sp.py

@@ -0,0 +1,295 @@
+"""GQA decode kernel — Case 4 (Cube-SP × PE-SP). ★ slide-11 optimal.
+
+Per GQA_full_deck.pptx slide 11:
+  - K, V split 64-way (cube=row_wise, pe=row_wise); each rank owns
+    ``S_local = S_kv / (C·P)``.
+  - Local attention via per-rank S_kv-axis tile sweep (ADR-0063 §A.2)
+    keeps scratch bounded by ``TILE_S_KV``.
+  - Two-level reduce on the partial ``(m, ℓ, O)``:
+      • intra-CUBE 8-way (row chain along intra_W + col bridge along
+        intra_N) over the 2×4 PE grid;
+      • inter-CUBE 2-phase lrab over the 4×2 CUBE sub-mesh (ADR-0060
+        §4.2 lrab-adapted center-root reduce): Phase 1 row reduce
+        converges at root_col=2; Phase 2 col reduce on root_col
+        converges at root_row=1; root cube id = 6.
+  - PE 0 of CUBE 6 stores ``O``.
+
+Deviation from slide 13: slide prescribes AllReduce (every rank has
+the answer); the kernel does reduce-to-root (only cube 6 has it)
+per ADR-0060 §4. Treated as the kernbench Case-4 baseline.
+
+Tensor layout:
+  Q : (T_q, h_q · d_head) replicated; loaded as (G·T_q, d_head).
+  K : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank
+      loads its (d_head, S_local) slice via byte-conserving reshape
+      of (S_local, h_kv·d_head) — correct for zero / symmetric
+      inputs (ADR-0060 §3 reshape-not-transpose caveat).
+  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 6 stores.
+
+Topology / SFR:
+  - Requires ``configure_sfr_intercube_multisip`` (provides disjoint
+    ``intra_*`` and ``E/W/N/S`` namespaces).
+  - Intra-CUBE PEs arranged as a 2×4 grid (no wrap).
+  - Inter-CUBE: 4×2 CUBE sub-mesh starting at origin (0, 0).
+"""
+from __future__ import annotations
+
+
+TILE_S_KV = 1024  # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
+
+# lrab geometry for the C=8 single-KV-head group (4×2 cube sub-mesh).
+_SUB_W = 4
+_SUB_H = 2
+_ROOT_COL = _SUB_W // 2   # 2
+_ROOT_ROW = _SUB_H // 2   # 1
+_ROOT_CUBE = _ROOT_ROW * _SUB_W + _ROOT_COL  # 6
+
+
+def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
+    """Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
+    m_new = tl.maximum(m_local, m_other)
+    scale_old = tl.exp(m_local - m_new)
+    scale_new = tl.exp(m_other - m_new)
+    l_new = l_local * scale_old + l_other * scale_new
+    O_new = O_local * scale_old + O_other * scale_new
+    return m_new, l_new, O_new
+
+
+def gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel(
+    q_ptr: int,
+    k_ptr: int,
+    v_ptr: int,
+    o_ptr: int,
+    T_q: int,
+    S_kv: int,
+    h_q: int,
+    h_kv: int,
+    d_head: int,
+    C: int,
+    P: int,
+    *,
+    tl,
+) -> None:
+    """Case-4 decode: Cube-SP × PE-SP, lrab center-root at cube 6."""
+    G = h_q // h_kv
+    n_ranks = C * P
+    S_local = S_kv // n_ranks
+    pe_id = tl.program_id(axis=0)
+    cube_id = tl.program_id(axis=1)
+
+    # ── Local attention (S_kv-axis tile sweep, ADR-0063 §A.2) ──
+    Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
+    n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
+    KV_ROW_BYTES = d_head * 2  # f16
+
+    # Tile 0: establishes persistent (m_local, l_local, O_local). Cannot
+    # fold into Tiles 1..N loop (kernbench-only): persistent tensors
+    # must live outside tl.scratch_scope or scope teardown discards
+    # them; there's no scratch-backed (-inf, 0, 0) initializer.
+    tile_s0 = min(TILE_S_KV, S_local)
+    K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
+    V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
+    scores = tl.dot(Q, K_T)
+    m_local = tl.max(scores, axis=-1)
+    centered = scores - m_local
+    exp_scores = tl.exp(centered)
+    l_local = tl.sum(exp_scores, axis=-1)
+    O_local = tl.dot(exp_scores, V)
+
+    for tile_idx in range(1, n_tiles):
+        tile_start = tile_idx * TILE_S_KV
+        tile_s = min(TILE_S_KV, S_local - tile_start)
+        with tl.scratch_scope():
+            K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
+                            shape=(d_head, tile_s), dtype="f16")
+            V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
+                          shape=(tile_s, d_head), dtype="f16")
+            scores_t = tl.dot(Q, K_T_t)
+            m_tile = tl.max(scores_t, axis=-1)
+            centered_t = scores_t - m_tile
+            exp_scores_t = tl.exp(centered_t)
+            l_tile = tl.sum(exp_scores_t, axis=-1)
+            O_tile = tl.dot(exp_scores_t, V_t)
+            m_new, l_new, O_new = _merge_running(
+                m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
+            )
+            tl.copy_to(m_local, m_new)
+            tl.copy_to(l_local, l_new)
+            tl.copy_to(O_local, O_new)
+
+    # ── Intra-CUBE reduce: row chain (intra_W) + col bridge (intra_N) ──
+    PE_GRID_COLS = 4
+    pe_col = pe_id % PE_GRID_COLS
+    pe_row = pe_id // PE_GRID_COLS
+    pe_cols_used = min(PE_GRID_COLS, P)
+    pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS
+
+    if pe_cols_used > 1:
+        if pe_col < pe_cols_used - 1:
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="intra_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)
+        if pe_col > 0:
+            tl.send(dir="intra_W", src=m_local)
+            tl.send(dir="intra_W", src=l_local)
+            tl.send(dir="intra_W", src=O_local)
+
+    if pe_col == 0 and pe_rows_used > 1:
+        if pe_row < pe_rows_used - 1:
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="intra_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)
+        if pe_row > 0:
+            tl.send(dir="intra_N", src=m_local)
+            tl.send(dir="intra_N", src=l_local)
+            tl.send(dir="intra_N", src=O_local)
+
+    # ── Inter-CUBE lrab-adapted center-root reduce (ADR-0060 §4.2) ──
+    # Only PE 0 of each CUBE participates. Adapts Phases 1-2 of
+    # lrab_hierarchical_allreduce.py: bidirectional row reduce converges
+    # at root_col; bidirectional col reduce on root_col converges at
+    # root_row. Plain ``+`` replaced by log-sum-exp ``_merge_running``.
+    if pe_id == 0:
+        row = cube_id // _SUB_W
+        col = cube_id % _SUB_W
+
+        # Phase 1: row reduce — converge at col == _ROOT_COL.
+        if 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:
+            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)
+            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:
+            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:
+                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:
+                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: PE 0 of CUBE 6) ──
+    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_decode_long_ctx_cube_sp_pe_tp.py

@@ -0,0 +1,246 @@
+"""GQA decode kernel — Case 1 (Cube-SP × PE-TP).
+
+Per GQA_full_deck.pptx slide 11:
+  - K, V split S_kv-wise across the 8 cubes (cube=row_wise);
+    replicated within each cube (pe=replicate). Each active rank
+    owns S_local = S_kv / C.
+  - PEs nominally split on the batch dim (PE-TP). At B=1 (the
+    slide-11 single-user default) only PE 0 of each cube has work;
+    PEs 1-7 of every cube idle (PE-TP-at-B=1 waste).
+  - NO intra-CUBE communication (only PE 0 holds a valid partial).
+  - Inter-CUBE 8-way reduce on (m, ℓ, O) via the lrab-adapted
+    center-root pattern (ADR-0060 §4.2): Phase 1 row reduce
+    converges at root_col; Phase 2 col reduce on root_col converges
+    at root_row. For sub_w=4, sub_h=2: root_col=2, root_row=1,
+    root_cube=6 (lrab geometric center).
+
+Tensor layout:
+  Q : (T_q, h_q · d_head) replicated on every rank.
+  K : (S_kv, h_kv · d_head) with cube=row_wise, pe=replicate —
+      each cube's PE 0 owns the full (S_local, h_kv·d_head) shard.
+  V : same as K.
+  O : (T_q, h_q · d_head) — only PE 0 of CUBE 6 stores.
+
+Topology / SFR:
+  - Requires ``configure_sfr_intercube_multisip`` for the E/W/N/S
+    inter-CUBE lanes used by the lrab reduce.
+"""
+from __future__ import annotations
+
+
+TILE_S_KV = 1024  # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
+
+
+def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
+    """Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
+    m_new = tl.maximum(m_local, m_other)
+    scale_old = tl.exp(m_local - m_new)
+    scale_new = tl.exp(m_other - m_new)
+    l_new = l_local * scale_old + l_other * scale_new
+    O_new = O_local * scale_old + O_other * scale_new
+    return m_new, l_new, O_new
+
+
+def gqa_attention_decode_long_ctx_cube_sp_pe_tp_kernel(
+    q_ptr: int,
+    k_ptr: int,
+    v_ptr: int,
+    o_ptr: int,
+    T_q: int,
+    S_kv: int,
+    h_q: int,
+    h_kv: int,
+    d_head: int,
+    C: int,
+    P: int,
+    *,
+    tl,
+) -> None:
+    """Case-1 decode: PE 0 per cube does attention; inter-CUBE lrab AR only."""
+    # B=1 single-user decode + PE-TP: only PE 0 per cube has work.
+    pe_id = tl.program_id(axis=0)
+    cube_id = tl.program_id(axis=1)
+    if pe_id != 0:
+        return
+
+    # Cube-SP geometry: 4×2 sub-mesh, lrab center-root cube = 6.
+    sub_w = 4
+    sub_h = C // sub_w
+    if sub_w < 2 or sub_h < 2 or sub_w * sub_h != C:
+        raise ValueError(
+            f"Case 1 requires C decomposable as sub_w=4, sub_h>=2; got C={C}"
+        )
+    root_col = sub_w // 2
+    root_row = sub_h // 2
+    root_cube = root_row * sub_w + root_col
+
+    G = h_q // h_kv
+    S_local = S_kv // C  # cube=row_wise, pe=replicate ⇒ S_local per cube
+    KV_ROW_BYTES = d_head * 2  # f16
+
+    # ── Local attention (S_kv-axis tile sweep, ADR-0063 §A.2) ──
+    Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
+    n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
+
+    tile_s0 = min(TILE_S_KV, S_local)
+    K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
+    V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
+    scores = tl.dot(Q, K_T)
+    m_local = tl.max(scores, axis=-1)
+    centered = scores - m_local
+    exp_scores = tl.exp(centered)
+    l_local = tl.sum(exp_scores, axis=-1)
+    O_local = tl.dot(exp_scores, V)
+
+    for tile_idx in range(1, n_tiles):
+        tile_start = tile_idx * TILE_S_KV
+        tile_s = min(TILE_S_KV, S_local - tile_start)
+        with tl.scratch_scope():
+            K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
+                            shape=(d_head, tile_s), dtype="f16")
+            V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
+                          shape=(tile_s, d_head), dtype="f16")
+            scores_t = tl.dot(Q, K_T_t)
+            m_tile = tl.max(scores_t, axis=-1)
+            centered_t = scores_t - m_tile
+            exp_scores_t = tl.exp(centered_t)
+            l_tile = tl.sum(exp_scores_t, axis=-1)
+            O_tile = tl.dot(exp_scores_t, V_t)
+            m_new, l_new, O_new = _merge_running(
+                m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
+            )
+            tl.copy_to(m_local, m_new)
+            tl.copy_to(l_local, l_new)
+            tl.copy_to(O_local, O_new)
+
+    # ── Inter-CUBE lrab-adapted center-root reduce (ADR-0060 §4.2) ──
+    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: PE 0 of cube 6) ──
+    if cube_id == root_cube:
+        O_final = O_local / l_local
+        tl.store(o_ptr, O_final)

src/kernbench/benches/_gqa_attention_decode_opt2.py

@@ -37,9 +37,10 @@ def gqa_attention_decode_opt2_kernel(
 ) -> None:
     """Single-rank (C=P=1) GQA decode using the opt2 two-composite form.
 
-    Layout mirrors ``gqa_attention_decode_long_kernel`` (M-fold Q, K loaded
+    Layout mirrors ``gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel``
-    as ``(d_head, S_local)``). The KV slice is split into two sub-tiles so
+    (M-fold Q, K loaded as ``(d_head, S_local)``). The KV slice is split
-    the second one drives the ``softmax_merge`` recipe composite.
+    into two sub-tiles so the second one drives the ``softmax_merge``
+    recipe composite.
     """
     G = h_q // h_kv
     S_local = S_kv // (C * P)

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:
-    if pe_id != 0:
+        # Head-parallel only: PE 0 of each CUBE participates.
-        return
+        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_decode_long_ctx_4cases.py

@@ -0,0 +1,323 @@
+"""milestone-gqa-decode-long-ctx-4cases: long-context decode 4-cases study.
+
+Per GQA_full_deck.pptx slides 11-17: 4 KV-cache sharding strategies on
+the LLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs) at long
+context (S_kv = 128K, T_q = 1):
+
+    Case 1  Cube-SP / PE-TP  → KV split by S_kv across cubes; PEs TP on batch
+    Case 2  Cube-Repl / PE-TP → full KV per cube; PEs TP on batch
+    Case 3  Cube-Repl / PE-SP → full KV per cube; PEs SP on S_kv (intra-cube AR)
+    Case 4  Cube-SP / PE-SP  → KV split 64-way; 2-phase AR on (m,ℓ,O)  ★ optimal
+
+Each case is a separate panel. The bench drives all panels in one
+invocation and writes per-panel op_log_summary to sweep.json so the
+comparative analysis (latency, GEMM/MAC util, comm volume) can be
+generated from a single sweep.
+
+Deviation from slide 13: slide prescribes AllReduce (every rank has
+the answer); the kernel does reduce-to-root (only the lrab center
+cube has it) per ADR-0060 §4. Treated as the kernbench Case-4 baseline.
+
+Gated by ``GQA_DECODE_LONG_CTX_4CASES_RUN=1``.
+"""
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+from kernbench.benches._gqa_attention_decode_long_ctx_cube_repl_pe_sp import (
+    gqa_attention_decode_long_ctx_cube_repl_pe_sp_kernel,
+)
+from kernbench.benches._gqa_attention_decode_long_ctx_cube_repl_pe_tp import (
+    gqa_attention_decode_long_ctx_cube_repl_pe_tp_kernel,
+)
+from kernbench.benches._gqa_attention_decode_long_ctx_cube_sp_pe_sp import (
+    gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel,
+)
+from kernbench.benches._gqa_attention_decode_long_ctx_cube_sp_pe_tp import (
+    gqa_attention_decode_long_ctx_cube_sp_pe_tp_kernel,
+)
+from kernbench.benches.milestone_gqa_headline import (
+    _ccl_cfg,
+    _summarize_op_log,
+)
+from kernbench.benches.registry import bench
+from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
+from kernbench.policy.placement.dp import DPPolicy
+
+_OUTPUT_DIR = (
+    Path(__file__).resolve().parent
+    / "1H_milestone_output"
+    / "gqa_decode_long_ctx_4cases"
+)
+_SWEEP_JSON = _OUTPUT_DIR / "sweep.json"
+
+
+# ── Panel registry ───────────────────────────────────────────────────
+
+
+_PANELS = (
+    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp",      # Case 4 ★ optimal
+    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp",    # Case 2 (no comm; 8× memory)
+    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp",    # Case 3 (intra-CUBE AR only; 8× memory)
+    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp",      # Case 1 (inter-CUBE lrab only; PE-TP B=1 waste)
+)
+
+# Each entry: (kind, panel-specific params).
+# LLaMA-3.1-70B single-KV-head group target:
+#   1 KV head, h_q = 8 (G = 8 group), d_head = 128
+#   8 cubes (head-parallel group), 8 PEs/cube
+#   S_kv = 128K (long-context decode), T_q = 1 (one new token per pass)
+_PANEL_DISPATCH: dict[str, tuple[str, dict]] = {
+    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp": ("decode_long_ctx_cube_sp_pe_sp", {
+        # Case 4: KV split 64-way (Cube-SP × PE-SP), 2-level reduce.
+        # The sub_w=4/sub_h=2 lrab center-root geometry (root cube 6) is
+        # baked into the Case-4 wrapper kernel.
+        "C": 8, "P": 8,
+        "T_q": 1, "S_kv": 131_072,
+        "d_head": 128, "h_q": 8, "h_kv": 1,
+    }),
+    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp": ("decode_long_ctx_cube_repl_pe_tp", {
+        # Case 2: K, V replicated everywhere (8× memory waste); PEs TP
+        # on batch. For B=1 only one rank works (slide-11 PE-TP waste).
+        # No inter-rank communication.
+        "C": 8, "P": 8,
+        "T_q": 1, "S_kv": 131_072,
+        "d_head": 128, "h_q": 8, "h_kv": 1,
+    }),
+    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp": ("decode_long_ctx_cube_repl_pe_sp", {
+        # Case 3: K, V replicated per cube (8× memory); PEs SP on S_kv
+        # within each cube. Intra-CUBE 8-way reduce; no inter-CUBE comm
+        # (every cube ends with full answer; designated writer = cube 0).
+        "C": 8, "P": 8,
+        "T_q": 1, "S_kv": 131_072,
+        "d_head": 128, "h_q": 8, "h_kv": 1,
+    }),
+    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp": ("decode_long_ctx_cube_sp_pe_tp", {
+        # Case 1: K, V split across cubes (S_local = S_kv/C per cube);
+        # PEs TP on batch — at B=1 only PE 0 of each cube works (PE-TP
+        # waste). Inter-CUBE lrab AR (root = cube 6); no intra-CUBE comm.
+        "C": 8, "P": 8,
+        "T_q": 1, "S_kv": 131_072,
+        "d_head": 128, "h_q": 8, "h_kv": 1,
+    }),
+}
+
+
+# ── Per-panel runner ─────────────────────────────────────────────────
+
+
+def _run_decode_panel_long_ctx_cube_repl_pe_tp(
+    ctx, *, panel: str, C: int, P: int,
+    T_q: int, S_kv: int,
+    d_head: int, h_q: int, h_kv: int,
+) -> None:
+    """Case 2 runner: K, V replicated everywhere; B=1 single-rank work.
+
+    DPPolicy models the cluster-wide memory waste — every rank holds
+    full K, V in its HBM region. Only PE 0 of CUBE 0 computes (the
+    kernel early-returns on every other rank), so only one rank reads
+    from its HBM copy and writes the output.
+    """
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_repl = DPPolicy(cube="replicate", pe="replicate",
+                       num_cubes=C, num_pes=P)
+    q = ctx.zeros((T_q, h_q * d_head),
+                  dtype="f16", dp=dp_repl, name=f"{panel}_q")
+    k = ctx.zeros((S_kv, h_kv * d_head),
+                  dtype="f16", dp=dp_repl, name=f"{panel}_k")
+    v = ctx.zeros((S_kv, h_kv * d_head),
+                  dtype="f16", dp=dp_repl, name=f"{panel}_v")
+    o = ctx.empty((T_q, h_q * d_head),
+                  dtype="f16", dp=dp_repl, name=f"{panel}_o")
+    ctx.launch(
+        panel, gqa_attention_decode_long_ctx_cube_repl_pe_tp_kernel,
+        q, k, v, o,
+        T_q, S_kv, h_q, h_kv, d_head, C, P,
+        _auto_dim_remap=False,
+    )
+
+
+def _run_decode_panel_long_ctx_cube_repl_pe_sp(
+    ctx, *, panel: str, C: int, P: int,
+    T_q: int, S_kv: int,
+    d_head: int, h_q: int, h_kv: int,
+) -> None:
+    """Case 3 runner: K, V replicated per cube; PEs SP on S_kv within cube.
+
+    DPPolicy models the cluster-wide 8× memory waste — every cube
+    holds full K, V in its HBM region, then splits the S_kv axis
+    row_wise across its 8 PEs. The kernel does an intra-CUBE 8-way
+    reduce on (m, ℓ, O); only cube 0's PE 0 writes the output.
+    """
+    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="replicate", pe="row_wise",
+                     num_cubes=C, num_pes=P)
+    q = ctx.zeros((T_q, h_q * d_head),
+                  dtype="f16", dp=dp_full, name=f"{panel}_q")
+    k = ctx.zeros((S_kv, h_kv * d_head),
+                  dtype="f16", dp=dp_kv, name=f"{panel}_k")
+    v = ctx.zeros((S_kv, h_kv * d_head),
+                  dtype="f16", dp=dp_kv, name=f"{panel}_v")
+    o = ctx.empty((T_q, h_q * d_head),
+                  dtype="f16", dp=dp_full, name=f"{panel}_o")
+    ctx.launch(
+        panel, gqa_attention_decode_long_ctx_cube_repl_pe_sp_kernel,
+        q, k, v, o,
+        T_q, S_kv, h_q, h_kv, d_head, C, P,
+        _auto_dim_remap=False,
+    )
+
+
+def _run_decode_panel_long_ctx_cube_sp_pe_tp(
+    ctx, *, panel: str, C: int, P: int,
+    T_q: int, S_kv: int,
+    d_head: int, h_q: int, h_kv: int,
+) -> None:
+    """Case 1 runner: K, V split across cubes; PE-TP single-rank work at B=1.
+
+    DPPolicy: K, V are cube=row_wise (S_local = S_kv/C per cube),
+    pe=replicate (full S_local within each cube). At B=1, the kernel
+    runs only on PE 0 of every cube; PEs 1-7 early-return. PE 0 per
+    cube does local attention on S_local, then participates in the
+    inter-CUBE lrab AR; PE 0 of cube 6 (lrab center) writes O.
+    """
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_full = DPPolicy(cube="replicate", pe="replicate",
+                       num_cubes=C, num_pes=P)
+    dp_kv = DPPolicy(cube="row_wise", pe="replicate",
+                     num_cubes=C, num_pes=P)
+    q = ctx.zeros((T_q, h_q * d_head),
+                  dtype="f16", dp=dp_full, name=f"{panel}_q")
+    k = ctx.zeros((S_kv, h_kv * d_head),
+                  dtype="f16", dp=dp_kv, name=f"{panel}_k")
+    v = ctx.zeros((S_kv, h_kv * d_head),
+                  dtype="f16", dp=dp_kv, name=f"{panel}_v")
+    o = ctx.empty((T_q, h_q * d_head),
+                  dtype="f16", dp=dp_full, name=f"{panel}_o")
+    ctx.launch(
+        panel, gqa_attention_decode_long_ctx_cube_sp_pe_tp_kernel,
+        q, k, v, o,
+        T_q, S_kv, h_q, h_kv, d_head, C, P,
+        _auto_dim_remap=False,
+    )
+
+
+def _run_decode_panel_long_ctx_cube_sp_pe_sp(
+    ctx, *, panel: str, C: int, P: int,
+    T_q: int, S_kv: int,
+    d_head: int, h_q: int, h_kv: int,
+) -> None:
+    """Case 4 runner: K, V split 64-way (Cube-SP × PE-SP).
+
+    DPPolicy: K, V are cube=row_wise, pe=row_wise — each rank owns
+    S_local = S_kv/(C·P). The wrapper kernel bakes in the lrab
+    sub_w=4 / sub_h=2 geometry (root cube 6, ADR-0060 §4.2).
+    """
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_full = DPPolicy(cube="replicate", pe="replicate",
+                       num_cubes=C, num_pes=P)
+    dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
+                     num_cubes=C, num_pes=P)
+    q = ctx.zeros((T_q, h_q * d_head),
+                  dtype="f16", dp=dp_full, name=f"{panel}_q")
+    k = ctx.zeros((S_kv, h_kv * d_head),
+                  dtype="f16", dp=dp_kv, name=f"{panel}_k")
+    v = ctx.zeros((S_kv, h_kv * d_head),
+                  dtype="f16", dp=dp_kv, name=f"{panel}_v")
+    o = ctx.empty((T_q, h_q * d_head),
+                  dtype="f16", dp=dp_full, name=f"{panel}_o")
+    ctx.launch(
+        panel, gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel,
+        q, k, v, o,
+        T_q, S_kv, h_q, h_kv, d_head, C, P,
+        _auto_dim_remap=False,
+    )
+
+
+def _make_bench_fn(panel: str):
+    kind, params = _PANEL_DISPATCH[panel]
+
+    def _bench_fn(ctx):
+        if kind == "decode_long_ctx_cube_sp_pe_sp":
+            _run_decode_panel_long_ctx_cube_sp_pe_sp(ctx, panel=panel, **params)
+        elif kind == "decode_long_ctx_cube_repl_pe_tp":
+            _run_decode_panel_long_ctx_cube_repl_pe_tp(ctx, panel=panel, **params)
+        elif kind == "decode_long_ctx_cube_repl_pe_sp":
+            _run_decode_panel_long_ctx_cube_repl_pe_sp(ctx, panel=panel, **params)
+        elif kind == "decode_long_ctx_cube_sp_pe_tp":
+            _run_decode_panel_long_ctx_cube_sp_pe_tp(ctx, panel=panel, **params)
+        else:
+            raise RuntimeError(
+                f"milestone-gqa-decode-long-ctx-4cases panel {panel!r} has "
+                f"unsupported kind={kind!r}."
+            )
+    return _bench_fn
+
+
+def _run_panel(panel: str, topology: str) -> dict:
+    from kernbench.runtime_api.bench_runner import run_bench
+    from kernbench.runtime_api.types import resolve_device
+    from kernbench.sim_engine.engine import GraphEngine
+    from kernbench.topology.builder import resolve_topology
+
+    topo = resolve_topology(topology)
+    result = run_bench(
+        topology=topo, bench_fn=_make_bench_fn(panel),
+        device=resolve_device(None),
+        engine_factory=lambda t, d: GraphEngine(
+            getattr(t, "topology_obj", t), enable_data=True,
+        ),
+    )
+    if not result.completion.ok:
+        raise RuntimeError(
+            f"milestone-gqa-decode-long-ctx-4cases panel {panel!r} failed: "
+            f"{result.completion}"
+        )
+    kind, params = _PANEL_DISPATCH[panel]
+    return {
+        "panel": panel,
+        "kind": kind,
+        **params,
+        "op_log_summary": _summarize_op_log(result.engine.op_log),
+    }
+
+
+# ── Bench entry ──────────────────────────────────────────────────────
+
+
+@bench(
+    name="milestone-gqa-decode-long-ctx-4cases",
+    description=(
+        "Long-context decode 4-cases comparative study on the "
+        "LLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs)."
+    ),
+)
+def run(torch) -> None:
+    """Drive the registered decode case panels; write sweep.json.
+
+    Gated by GQA_DECODE_LONG_CTX_4CASES_RUN=1.
+    """
+    _OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    if not os.environ.get("GQA_DECODE_LONG_CTX_4CASES_RUN"):
+        raise RuntimeError(
+            "milestone-gqa-decode-long-ctx-4cases needs "
+            "GQA_DECODE_LONG_CTX_4CASES_RUN=1."
+        )
+
+    topology = os.environ.get(
+        "GQA_DECODE_LONG_CTX_4CASES_TOPOLOGY", "topology.yaml",
+    )
+    rows = [_run_panel(panel, topology) for panel in _PANELS]
+    sweep = {
+        "version": 1,
+        "panels": list(_PANELS),
+        "rows": rows,
+    }
+    _SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
+    print(
+        f"  milestone-gqa-decode-long-ctx-4cases: {len(rows)} rows -> {_SWEEP_JSON}"
+    )

src/kernbench/benches/milestone_gqa_headline.py

@@ -1,24 +1,26 @@
-"""milestone-gqa-headline bench: real GQA + 2-level SP + Ring KV.
+"""milestone-gqa-headline bench: real GQA prefill — single-KV-group target.
 
-Wires the new ``_gqa_decode`` and ``_gqa_prefill`` kernels through 4
+Drives the LLaMA-3.1-70B single-KV-head-group prefill panel through
-panels with real GQA (h_q = G·h_kv, G > 1 on the decode side), writing
+``_gqa_attention_prefill_long`` (C=8 snake Ring KV + intra-CUBE PE-SP,
-per-panel ``op_log_summary`` into ``sweep.json``. Independent from the
+all 64 ranks), writing ``op_log_summary`` into ``sweep.json``.
-existing ``milestone-gqa-llama70b`` validation-scale bench (which stays
+Independent from the existing ``milestone-gqa-llama70b`` validation-scale
-on the legacy baseline kernels).
+bench (which stays on the legacy baseline kernels).
 
-Restrictions (P7 first cut):
+Decode panels live in ``milestone_gqa_decode_long_ctx_4cases.py`` (the
-  - C ≤ 4 (single-row inter-CUBE ring SFR; multi-row deferred)
+4-cases long-context decode comparative study).
+
+Restrictions:
   - Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP
-    headline deferred)
+    headline deferred per ADR-0060 §B-item-1)
+  - T_q = S_kv = 1024 (scratch-limited; 32K headline awaits Q-axis
+    kernel tiling)
   - No figure renderers (defer to a separate cycle)
 
 Panels:
-  single_user_prefill_gqa : prefill C=1, T_q=4, S_kv=16
+  single_kv_group_prefill_gqa_c8_p8: prefill C=8 snake Ring KV +
-  multi_user_prefill_gqa  : prefill C=4 Ring KV, T_q=4, S_kv=16
+                                     intra-CUBE PE-SP, T_q=S_kv=1K,
-  single_user_decode_gqa  : decode C=1, P=8, h_q=8, h_kv=1, S_kv=64
+                                     d_head=128 — the LLaMA-3.1-70B
-                            (M-fold + intra-cube row-then-col chain)
+                                     single-KV-group prefill target.
-  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``.
 """
@@ -28,14 +30,10 @@ import json
 import os
 from pathlib import Path
 
-from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel
 from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel
 from kernbench.benches.registry import bench
 from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
-from kernbench.ccl.sfr_config import (
+from kernbench.ccl.sfr_config import configure_sfr_intercube_ring
-    configure_sfr_intercube_multisip,
-    configure_sfr_intercube_ring,
-)
 from kernbench.policy.placement.dp import DPPolicy
 
 _OUTPUT_DIR = Path(__file__).resolve().parent / "1H_milestone_output" / "gqa_headline"
@@ -46,24 +44,26 @@ _SWEEP_JSON = _OUTPUT_DIR / "sweep.json"
 _DTYPE = "f16"
 _D_HEAD = 64
 _T_Q_PREFILL = 4
-_T_Q_DECODE = 1
 _S_KV_PREFILL = 16
-_H_Q_DECODE = 8        # real GQA: G = H_Q_DECODE / H_KV_DECODE = 8
-_H_KV_DECODE = 1
 
 _PANELS = (
-    "single_user_prefill_gqa",
+    "single_kv_group_prefill_gqa_c8_p8",
-    "multi_user_prefill_gqa",
-    "single_user_decode_gqa",
-    "multi_user_decode_gqa",
 )
 
 # Each entry: (kind, panel-specific params)
 _PANEL_DISPATCH: dict[str, tuple[str, dict]] = {
-    "single_user_prefill_gqa": ("prefill", {"C": 1, "S_kv": _S_KV_PREFILL}),
+    "single_kv_group_prefill_gqa_c8_p8": ("prefill", {
-    "multi_user_prefill_gqa":  ("prefill", {"C": 4, "S_kv": _S_KV_PREFILL}),
+        "C": 8, "P": 8,
-    "single_user_decode_gqa":  ("decode",  {"C": 1, "P": 8, "S_kv": 64}),
+        # T_q = S_kv = 1024 (one-shot long-context prefill, scratch-limited).
-    "multi_user_decode_gqa":   ("decode",  {"C": 4, "P": 8, "S_kv": 128}),
+        # The bootstrap section of the prefill kernel leaves K_t/V_t/scores/
+        # exp_scores persistent (outside tl.scratch_scope), inflating the
+        # baseline. At T_q=2048 the peak hits ~1.05 MB vs 1.0 MB budget; at
+        # 1024 it fits comfortably. The true LLaMA 32K headline awaits a
+        # future increment to (a) add Q-axis tiling and (b) move the
+        # bootstrap into scratch discipline.
+        "T_q": 1_024, "S_kv": 1_024,
+        "d_head": 128,
+    }),
 }
 
 
@@ -76,63 +76,61 @@ 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(
+        mesh_w = int(ctx.spec["sip"]["cube_mesh"]["w"])
-            ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C,
+        if C <= mesh_w:
-        )
+            configure_sfr_intercube_ring(
-    dp_q = DPPolicy(cube="replicate", pe="replicate",
+                ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C,
-                    num_cubes=C, num_pes=1)
+            )
+        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)
+                    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,
-    )
-
-
-def _run_decode_panel(ctx, *, panel: str, C: int, P: int, S_kv: int) -> None:
-    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)
-    q = ctx.zeros((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD),
-                  dtype=_DTYPE, dp=dp_full, name=f"{panel}_q")
-    k = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD),
-                  dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k")
-    v = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD),
-                  dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v")
-    o = ctx.empty((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD),
-                  dtype=_DTYPE, dp=dp_full, name=f"{panel}_o")
-    ctx.launch(
-        panel, gqa_attention_decode_long_kernel,
-        q, k, v, o,
-        _T_Q_DECODE, S_kv, _H_Q_DECODE, _H_KV_DECODE, _D_HEAD, C, P,
         _auto_dim_remap=False,
     )
 
 
 def _make_bench_fn(panel: str):
     kind, params = _PANEL_DISPATCH[panel]
+    assert kind == "prefill", (
+        f"milestone-gqa-headline only registers prefill panels; got {kind!r}"
+    )
 
     def _bench_fn(ctx):
-        if kind == "prefill":
+        _run_prefill_panel(ctx, panel=panel, **params)
-            _run_prefill_panel(ctx, panel=panel, **params)
-        else:
-            _run_decode_panel(ctx, panel=panel, **params)
 
     return _bench_fn
 
@@ -197,10 +195,10 @@ def _run_panel(panel: str, topology: str) -> dict:
 
 @bench(
     name="milestone-gqa-headline",
-    description="Headline GQA milestone — real GQA h_q=8/h_kv=1 + 2-level SP (decode) + Ring KV (prefill).",
+    description="Headline GQA prefill milestone — LLaMA-3.1-70B single-KV-group target (C=8, P=8).",
 )
 def run(torch) -> None:
-    """Drive 4 headline panels through the new GQA kernels; write sweep.json.
+    """Drive the headline prefill panel; write sweep.json.
 
     Gated by GQA_HEADLINE_RUN=1.
     """
@@ -217,10 +215,7 @@ def run(torch) -> None:
         "panels": list(_PANELS),
         "config": {
             "T_q_prefill": _T_Q_PREFILL,
-            "T_q_decode": _T_Q_DECODE,
             "S_kv_prefill": _S_KV_PREFILL,
-            "h_q_decode": _H_Q_DECODE,
-            "h_kv_decode": _H_KV_DECODE,
             "d_head": _D_HEAD,
         },
         "rows": rows,

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,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:
+    if submesh_shape is not None:
-        ring_size = n_cubes
+        sub_h, sub_w = submesh_shape
-    if ring_size > mesh_w:
+        origin_row, origin_col = submesh_origin
-        raise ValueError(
+        if (sub_h <= 0 or sub_w <= 0
-            f"intercube_ring ring_size={ring_size} > mesh_w={mesh_w}; "
+                or origin_row < 0 or origin_col < 0
-            "multi-row rings cross non-neighbour boundaries"
+                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) ──
+        # ── 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.py

@@ -1,143 +0,0 @@
-"""Phase 1 spec test for P1a GQA decode kernel (real GQA via M-fold).
-
-P1a is the first phase of the DDD-0060 plan, split out of the original
-P1 (the composite-hybrid swap is P1b, deferred until the tl.composite
-output-handle question is decided). P1a is the *correctness* unlock:
-the kernel processes ONE KV head at a time and folds the G query heads
-into the matmul M (row) dimension so a single Q·Kᵀ serves all G heads
-sharing one K (ADR-0060 §5.2). This lifts the baseline's
-``h_q == h_kv == 1`` cap pinned at
-``tests/attention/test_milestone_gqa_llama70b.py:137-142``.
-
-P1a stays inside the existing ``tl`` API: the two attention GEMMs use the
-blocking ``tl.dot`` so the chain ``Q·Kᵀ → softmax → P·V → store`` fits
-without any composite-output chaining. The composite swap (P1b) will
-revisit this once the API for feeding a composite's output into a
-downstream MATH op is settled.
-
-Phase 1 (this commit): tests only — production code lands in Phase 2.
-Tests fail at import in Phase 1 with ModuleNotFoundError; Phase 2 makes
-them pass.
-"""
-from __future__ import annotations
-
-from pathlib import Path
-
-from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel  # noqa: F401  (Phase 2 deliverable)
-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"
-
-# Decode shapes — P1a is one-shot, no tiling, single rank. P3 will tile.
-T_Q = 1
-D_HEAD = 64
-S_KV = 16
-DTYPE = "f16"
-
-
-def _engine_factory(t, d):
-    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
-
-
-def _run_decode_p1(*, h_q: int, h_kv: int):
-    """One-shot GQA decode on a single PE in a single CUBE (P=1, no SP).
-
-    Tensor layout (natural K — same as the P2a SP tests; the kernel
-    reshapes K to (d_head, S_local) via byte-conserving load):
-      Q : (T_q, h_q · d_head)   — natural Q layout; kernel reshapes to
-          (G·T_q, d_head) — byte-conserving and math-correct because
-          T_q=1 collapses axis ordering.
-      K : (S_kv, h_kv · d_head) — natural K layout; kernel loads as
-          (d_head, S_local) — reshape-as-transpose caveat (ADR-0060 §3
-          / §B item 2), correct for zero inputs used here.
-      V : (S_kv, h_kv · d_head) — natural V layout.
-      O : (T_q, h_q · d_head)   — same shape as Q.
-    """
-    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
-
-    def _bench_fn(ctx):
-        dp = DPPolicy(cube="replicate", pe="replicate",
-                      num_cubes=1, num_pes=1)
-        q = ctx.zeros((T_Q, h_q * D_HEAD),
-                      dtype=DTYPE, dp=dp, name=f"q_h{h_q}_kv{h_kv}")
-        k = ctx.zeros((S_KV, h_kv * D_HEAD),
-                      dtype=DTYPE, dp=dp, name=f"k_h{h_q}_kv{h_kv}")
-        v = ctx.zeros((S_KV, h_kv * D_HEAD),
-                      dtype=DTYPE, dp=dp, name=f"v_h{h_q}_kv{h_kv}")
-        o = ctx.empty((T_Q, h_q * D_HEAD),
-                      dtype=DTYPE, dp=dp, name=f"o_h{h_q}_kv{h_kv}")
-        ctx.launch(
-            f"gqa_decode_p1_h{h_q}_kv{h_kv}",
-            gqa_attention_decode_long_kernel,
-            q, k, v, o,
-            T_Q, S_KV, h_q, h_kv, D_HEAD,
-            1, 1,                               # C=1, P=1 (no SP, degenerate)
-            _auto_dim_remap=False,
-        )
-
-    return run_bench(
-        topology=topo,
-        bench_fn=_bench_fn,
-        device=resolve_device(None),
-        engine_factory=_engine_factory,
-    )
-
-
-def _dma_read_count(op_log) -> int:
-    return sum(1 for r in op_log if r.op_name == "dma_read")
-
-
-# ── Headline unlock: real GQA (h_q = G·h_kv) runs in data mode ──────────
-
-
-def test_real_gqa_h_q_eight_h_kv_one_completes_in_data_mode():
-    """ADR-0060 §A.1 headline unlock — the baseline raises
-    ``ValueError: Shape mismatch …`` in MemoryStore at h_q=8, h_kv=1
-    because ``_view(K, (h_q·d, S_kv))`` only byte-conserves when h_q==h_kv.
-    M-fold processes one KV head at a time using only byte-conserving
-    reshapes, so this completes.
-    """
-    result = _run_decode_p1(h_q=8, h_kv=1)
-    assert result.completion.ok, (
-        f"real GQA (h_q=8, h_kv=1) decode failed: {result.completion}"
-    )
-
-
-# ── M-fold property: K/V loads do not scale with G ─────────────────────
-
-
-def test_kv_dma_read_count_independent_of_g():
-    """ADR-0060 TL;DR / §5.2: M-fold loads K and V once per KV head and
-    folds the G query heads into the GEMM M dim. The dma_read_count must
-    therefore be identical between (G=1, h_kv=1) and (G=8, h_kv=1) — both
-    issue exactly 3 reads (Q + K + V). This pins the GQA-reuse property
-    that the rest of the plan (composite streaming in P4, etc.) builds on.
-    """
-    g1 = _run_decode_p1(h_q=1, h_kv=1)
-    g8 = _run_decode_p1(h_q=8, h_kv=1)
-    n_g1 = _dma_read_count(g1.engine.op_log)
-    n_g8 = _dma_read_count(g8.engine.op_log)
-    assert n_g1 == 3, f"G=1 dma_read_count must be 3 (Q+K+V); got {n_g1}"
-    assert n_g8 == 3, f"G=8 dma_read_count must be 3 (Q+K+V); got {n_g8}"
-    assert n_g1 == n_g8, (
-        f"K/V dma_read_count must be independent of G; "
-        f"got G=1 -> {n_g1}, G=8 -> {n_g8}"
-    )
-
-
-# ── Backward-compat: degenerate G=1 still works ────────────────────────
-
-
-def test_degenerate_g_equals_one_still_works():
-    """G=1 (h_q == h_kv == 1) is the baseline-compatible config. M-fold
-    degenerates to (T_q, d) = (1, 64) — the same shape the baseline
-    already exercises — so this proves no regression on that path.
-    """
-    result = _run_decode_p1(h_q=1, h_kv=1)
-    assert result.completion.ok, (
-        f"degenerate G=1 decode failed: {result.completion}"
-    )

tests/attention/test_gqa_decode_mc.py

@@ -1,182 +0,0 @@
-"""Phase 1 spec test for P2b GQA decode multi-cube SP (both Level-2 + Level-1).
-
-P2b extends P2a to multiple CUBEs in one CUBE Group. The kernel uses the
-canonical full SFR install (``configure_sfr_intercube_multisip``) which
-provides disjoint direction namespaces:
-
-  - ``intra_E / intra_W / intra_N / intra_S``  — PE↔PE within a CUBE
-    (logical 2×4 grid, no wrap)
-  - ``E / W / N / S``                          — CUBE↔CUBE inter-CUBE
-    (mesh, no wrap)
-
-Reduce pattern (chain reduce-to-root, ADR-0060 §A.2 spirit, §4 chain
-deviation noted):
-
-  Level-2 (intra-CUBE, 2×4 grid):
-    row-then-col chain — each row reduces leftward along ``intra_W`` to
-    its col-0 PE, then PE 4 sends to PE 0 along ``intra_N``. 7 chain
-    steps per CUBE × 3 handles each = 21 ``ipcq_copy`` per CUBE.
-
-  Level-1 (inter-CUBE):
-    only PE 0 of each CUBE participates. Chain leftward along ``W``.
-    (C-1) chain steps × 3 handles each.
-
-  Final store at PE 0 of CUBE 0 only.
-
-Phase 1 (this commit): tests only — production code lands in Phase 2.
-Phase 2 also updates ``test_gqa_decode.py`` (add C=1) and
-``test_gqa_decode_sp.py`` (switch SFR + add C=1).
-"""
-from __future__ import annotations
-
-from pathlib import Path
-
-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"
-
-T_Q = 1
-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 _run_decode_mc(*, h_q: int, h_kv: int, C: int, P: int, S_kv: int):
-    """Multi-CUBE SP decode: C cubes × P PEs each share the work."""
-    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", pe="row_wise",
-                         num_cubes=C, num_pes=P)
-        # Total KV split across C×P ranks; each rank sees S_kv/(C·P) rows.
-        q = ctx.zeros((T_Q, h_q * D_HEAD),
-                      dtype=DTYPE, dp=dp_full,
-                      name=f"q_h{h_q}_kv{h_kv}_c{C}_p{P}")
-        k = ctx.zeros((S_kv, h_kv * D_HEAD),
-                      dtype=DTYPE, dp=dp_kv,
-                      name=f"k_h{h_q}_kv{h_kv}_c{C}_p{P}")
-        v = ctx.zeros((S_kv, h_kv * D_HEAD),
-                      dtype=DTYPE, dp=dp_kv,
-                      name=f"v_h{h_q}_kv{h_kv}_c{C}_p{P}")
-        o = ctx.empty((T_Q, h_q * D_HEAD),
-                      dtype=DTYPE, dp=dp_full,
-                      name=f"o_h{h_q}_kv{h_kv}_c{C}_p{P}")
-        ctx.launch(
-            f"gqa_decode_mc_h{h_q}_kv{h_kv}_c{C}_p{P}",
-            gqa_attention_decode_long_kernel,
-            q, k, v, o,
-            T_Q, S_kv, h_q, h_kv, D_HEAD,
-            C, P,
-            _auto_dim_remap=False,
-        )
-
-    return run_bench(
-        topology=topo,
-        bench_fn=_bench_fn,
-        device=resolve_device(None),
-        engine_factory=_engine_factory,
-    )
-
-
-def _count(op_log, name: str) -> int:
-    return sum(1 for r in op_log if r.op_name == name)
-
-
-# ── Degenerate C=1 P=1 ────────────────────────────────────────────────
-
-
-def test_mc_c_one_p_one_degenerate():
-    """C=1, P=1: single rank, no SP. No IPCQ traffic; one dma_write."""
-    result = _run_decode_mc(h_q=1, h_kv=1, C=1, P=1, S_kv=16)
-    assert result.completion.ok, (
-        f"C=1 P=1 degenerate failed: {result.completion}"
-    )
-    assert _count(result.engine.op_log, "ipcq_copy") == 0
-    assert _count(result.engine.op_log, "dma_write") == 1
-
-
-# ── Intra-CUBE only (single CUBE, P=8 on 2×4 grid) ────────────────────
-
-
-def test_mc_c_one_p_eight_intracube_grid():
-    """C=1, P=8: intra-cube row-then-col chain on the 2×4 grid.
-    7 chain steps × 3 handles (m, ℓ, O) = 21 ipcq_copy."""
-    result = _run_decode_mc(h_q=1, h_kv=1, C=1, P=8, S_kv=64)
-    assert result.completion.ok, (
-        f"C=1 P=8 intra-cube failed: {result.completion}"
-    )
-    assert _count(result.engine.op_log, "dma_write") == 1
-    n_copy = _count(result.engine.op_log, "ipcq_copy")
-    assert n_copy == 21, (
-        f"C=1 P=8: expected 21 ipcq_copy (7 chain × 3 handles); got {n_copy}"
-    )
-
-
-# ── Multi-CUBE root-only write ────────────────────────────────────────
-
-
-def test_mc_c_two_p_eight_root_only_writes_o():
-    """C=2, P=8: 16 ranks total. Only PE 0 of CUBE 0 writes O."""
-    result = _run_decode_mc(h_q=1, h_kv=1, C=2, P=8, S_kv=128)
-    assert result.completion.ok, (
-        f"C=2 P=8 multi-cube failed: {result.completion}"
-    )
-    n_writes = _count(result.engine.op_log, "dma_write")
-    assert n_writes == 1, (
-        f"root-only write must hold for C=2 P=8 (16 ranks); got {n_writes}"
-    )
-
-
-# ── Multi-CUBE total IPCQ chain count ─────────────────────────────────
-
-
-def test_mc_c_two_p_eight_total_ipcq_count():
-    """C=2, P=8: 21 intra-cube ipcq_copy per CUBE × 2 CUBEs + 3 inter-cube
-    chain ipcq_copy (C-1=1 step × 3 handles) = 45 total."""
-    result = _run_decode_mc(h_q=1, h_kv=1, C=2, P=8, S_kv=128)
-    assert result.completion.ok, (
-        f"C=2 P=8 multi-cube failed: {result.completion}"
-    )
-    n_copy = _count(result.engine.op_log, "ipcq_copy")
-    expected = 21 * 2 + (2 - 1) * 3
-    assert n_copy == expected, (
-        f"C=2 P=8: expected {expected} ipcq_copy (21 intra × 2 CUBEs + 3 "
-        f"inter); got {n_copy}"
-    )
-
-
-# ── Real GQA × multi-CUBE SP combined ─────────────────────────────────
-
-
-def test_mc_real_gqa_c_two_p_eight():
-    """Headline: real GQA (h_q = G·h_kv with G=8) AND multi-CUBE SP
-    (C=2, P=8) together — the case the original baseline can express
-    neither part of."""
-    result = _run_decode_mc(h_q=8, h_kv=1, C=2, P=8, S_kv=128)
-    assert result.completion.ok, (
-        f"real GQA + multi-CUBE SP combined failed: {result.completion}"
-    )
-    n_writes = _count(result.engine.op_log, "dma_write")
-    assert n_writes == 1, (
-        f"root-only write must hold under M-fold + multi-CUBE; "
-        f"got {n_writes}"
-    )

tests/attention/test_gqa_decode_opt2.py

@@ -181,55 +181,3 @@ def test_opt2_bench_completes_oplog_mode():
     assert result.completion.ok, f"opt2 decode failed: {result.completion}"
 
 
-# ── opt2 ↔ opt3 numeric parity in data mode (ADR-0065 N4) ─────────────
-
-
-def _run_decode_data(kernel, name, q_d, k_d, v_d, *, S_kv=16):
-    """Run a decode kernel (C=P=1) in data mode with seeded Q/K/V; return the
-    HBM output array."""
-    from pathlib import Path
-
-    import numpy as np
-    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.data_executor import DataExecutor
-    from kernbench.sim_engine.engine import GraphEngine
-    from kernbench.topology.builder import resolve_topology
-
-    topo = resolve_topology(str(Path(__file__).resolve().parents[2] / "topology.yaml"))
-    cap: dict = {}
-
-    def bench_fn(ctx):
-        dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1)
-        q = ctx.zeros((1, 8 * D), dtype="f16", dp=dp, name=name + "q")
-        k = ctx.zeros((S_kv, D), dtype="f16", dp=dp, name=name + "k")
-        v = ctx.zeros((S_kv, D), dtype="f16", dp=dp, name=name + "v")
-        o = ctx.empty((1, 8 * D), dtype="f16", dp=dp, name=name + "o")
-        q.copy_(ctx.from_numpy(q_d)); k.copy_(ctx.from_numpy(k_d)); v.copy_(ctx.from_numpy(v_d))
-        cap["o"] = o
-        ctx.launch(name, kernel, q, k, v, o, 1, S_kv, 8, 1, D, 1, 1, _auto_dim_remap=False)
-
-    r = run_bench(topology=topo, bench_fn=bench_fn, device=resolve_device(None),
-                  engine_factory=lambda t, d: GraphEngine(getattr(t, "topology_obj", t),
-                                                          enable_data=True))
-    assert r.completion.ok
-    DataExecutor(r.engine.op_log, r.engine.memory_store).run()
-    o = cap["o"]
-    return r.engine.memory_store.read("hbm", o._handle.va_base, shape=(1, 8 * D), dtype="f16")
-
-
-def test_opt2_matches_opt3_data_mode():
-    import numpy as np
-    from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel
-    from kernbench.benches._gqa_attention_decode_opt2 import gqa_attention_decode_opt2_kernel
-
-    rng = np.random.default_rng(0)
-    q_d = (0.1 * rng.standard_normal((1, 8 * D))).astype(np.float16)
-    k_d = (0.1 * rng.standard_normal((16, D))).astype(np.float16)
-    v_d = (0.1 * rng.standard_normal((16, D))).astype(np.float16)
-
-    o3 = _run_decode_data(gqa_attention_decode_long_kernel, "p3", q_d, k_d, v_d)
-    o2 = _run_decode_data(gqa_attention_decode_opt2_kernel, "p2", q_d, k_d, v_d)
-    assert np.allclose(np.asarray(o2, np.float32), np.asarray(o3, np.float32),
-                       rtol=5e-2, atol=5e-2), f"opt2 {np.asarray(o2).ravel()[:4]} vs opt3 {np.asarray(o3).ravel()[:4]}"

tests/attention/test_gqa_decode_sp.py

@@ -1,153 +0,0 @@
-"""Phase 1 spec test for P2a GQA decode SP (chain reduce-to-root, Level-2 only).
-
-P2a is the first half of DDD-0060 P2: the kernel becomes multi-PE within
-one CUBE and reduces to root (PE 0) using a chain over the 1D intra-cube
-ring (W direction). This **replaces the baseline's bidirectional O(N)
-fan-out** where every rank ends with O — ADR-0060 §A.2's headline.
-
-Deviation from DDD-0060 §7 P2 gate: the gate text asks for
-``⌈log₂ P⌉`` reduce rounds. The intra-cube SFR install
-(``configure_sfr_intracube_pe_ring``) wires only a 1D E/W ring, so a
-true tree would require either multi-hop forwarding or a new SFR install
-(future ADR). P2a uses a **chain reduce-to-root**: ``P-1`` rounds along
-W. The architectural property the ADR cares about
-(root-only output vs every-rank-has-O) is preserved; the logarithmic
-collective is deferred.
-
-P2b (deferred) covers Level-1 inter-CUBE center-mesh reduce (C>1).
-
-Phase 1 (this commit): tests only — production code lands in Phase 2.
-"""
-from __future__ import annotations
-
-from pathlib import Path
-
-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"
-
-T_Q = 1
-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 _run_decode_sp(*, h_q: int, h_kv: int, P: int, S_kv: int):
-    """Single-CUBE SP decode: P PEs share the work along the intra-cube ring."""
-    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=1, num_pes=P)
-        dp_kv = DPPolicy(cube="replicate", pe="row_wise",
-                         num_cubes=1, num_pes=P)
-        q = ctx.zeros((T_Q, h_q * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name=f"q_h{h_q}_kv{h_kv}_p{P}")
-        # KV: total S_kv split across P PEs along axis 0 (row_wise sharding).
-        # Each PE sees (S_kv/P, h_kv·D_HEAD).
-        k = ctx.zeros((S_kv, h_kv * D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"k_h{h_q}_kv{h_kv}_p{P}")
-        v = ctx.zeros((S_kv, h_kv * D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"v_h{h_q}_kv{h_kv}_p{P}")
-        o = ctx.empty((T_Q, h_q * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name=f"o_h{h_q}_kv{h_kv}_p{P}")
-        ctx.launch(
-            f"gqa_decode_sp_h{h_q}_kv{h_kv}_p{P}",
-            gqa_attention_decode_long_kernel,
-            q, k, v, o,
-            T_Q, S_kv, h_q, h_kv, D_HEAD,
-            1, P,                                 # C=1, P=P (single-CUBE SP)
-            _auto_dim_remap=False,
-        )
-
-    return run_bench(
-        topology=topo,
-        bench_fn=_bench_fn,
-        device=resolve_device(None),
-        engine_factory=_engine_factory,
-    )
-
-
-def _count(op_log, name: str) -> int:
-    return sum(1 for r in op_log if r.op_name == name)
-
-
-# ── Root-only write ────────────────────────────────────────────────────
-
-
-def test_sp_chain_reduce_root_only_writes_o():
-    """ADR-0060 §A.2: only the root rank (PE 0) writes O. Baseline today
-    has every rank write the full final O (bidirectional fan-out)."""
-    result = _run_decode_sp(h_q=1, h_kv=1, P=8, S_kv=64)
-    assert result.completion.ok, f"P=8 chain reduce failed: {result.completion}"
-    n_writes = _count(result.engine.op_log, "dma_write")
-    assert n_writes == 1, (
-        f"reduce-to-root must produce exactly 1 dma_write (PE 0); "
-        f"got {n_writes}"
-    )
-
-
-# ── Chain step count ───────────────────────────────────────────────────
-
-
-def test_sp_chain_reduce_p_minus_one_ipcq_pairs():
-    """Chain reduce-to-root has P-1 send→recv pairs along the W chain;
-    each pair logs one ``ipcq_copy`` (inbound DMA, per
-    ``milestone_gqa_llama70b._summarize_op_log``). Each chain step ships
-    the triplet (m, ℓ, O) → 3 handles per step → 7 steps × 3 = 21."""
-    result = _run_decode_sp(h_q=1, h_kv=1, P=8, S_kv=64)
-    assert result.completion.ok, f"P=8 chain reduce failed: {result.completion}"
-    n_copy = _count(result.engine.op_log, "ipcq_copy")
-    expected = (8 - 1) * 3
-    assert n_copy == expected, (
-        f"chain reduce: expected {expected} ipcq_copy (P-1=7 steps × "
-        f"3 handles m/ℓ/O); got {n_copy}"
-    )
-
-
-# ── Real GQA × SP combined ─────────────────────────────────────────────
-
-
-def test_sp_real_gqa_h_q_eight_h_kv_one_p_eight():
-    """The combined unlock: real GQA (h_q=G·h_kv with G=8) AND SP
-    (P=8) together — neither expressible by the baseline."""
-    result = _run_decode_sp(h_q=8, h_kv=1, P=8, S_kv=64)
-    assert result.completion.ok, (
-        f"real GQA + SP combined run failed: {result.completion}"
-    )
-    n_writes = _count(result.engine.op_log, "dma_write")
-    assert n_writes == 1, (
-        f"root-only write must hold under M-fold too; got {n_writes}"
-    )
-
-
-# ── Degenerate P=1 ────────────────────────────────────────────────────
-
-
-def test_sp_p_one_degenerate_no_ipcq_traffic():
-    """P=1: SP degenerates to a single rank. No IPCQ traffic; one dma_write."""
-    result = _run_decode_sp(h_q=8, h_kv=1, P=1, S_kv=16)
-    assert result.completion.ok, f"P=1 degenerate failed: {result.completion}"
-    n_send = _count(result.engine.op_log, "ipcq_send")
-    n_recv = _count(result.engine.op_log, "ipcq_recv")
-    assert n_send == 0, f"P=1 must have no ipcq_send; got {n_send}"
-    assert n_recv == 0, f"P=1 must have no ipcq_recv; got {n_recv}"
-    n_writes = _count(result.engine.op_log, "dma_write")
-    assert n_writes == 1, f"P=1: one dma_write; got {n_writes}"

tests/attention/test_gqa_long_context.py

@@ -1,157 +0,0 @@
-"""Phase 1 spec test for Phase C: long-context regression for the GQA
-kernels (ADR-0060 §B "long/short context split" + ADR-0063 §A.2).
-
-The headline panel today caps prefill at S_kv=16 / decode at S_kv≤128 —
-NOT because the algorithm fails, but because the bump allocator never
-recycles. ADR-0063 §A.2 (test req 3) requires a sweep at an S that
-would overflow 1 MiB without recycling to complete after scope
-discipline lands.
-
-Scratch-budget estimate at C=4, P=8 (current 1 MiB pool):
-  decode:  per-rank S_local = S_kv / 32; intermediates ≲ 500 KB at
-           S_kv=32K (fits today — used as regression guard).
-  prefill: per-rank S_local = S_kv / 4; ~16·S_local bytes per ring
-           step × 4 steps ≈ 64·S_local bytes. Overflows at
-           ~64 K tokens (64 × 16K > 1 MiB).
-
-Phase 1 (this commit): tests only — production code lands in Phase 2.
-The prefill test fails today with a RuntimeError("TLContext scratch
-overflow"). After Phase 2 (scratch_scope + copy_to discipline) it
-completes.
-"""
-from __future__ import annotations
-
-from pathlib import Path
-
-from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel  # noqa: F401
-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_multisip,
-    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)
-
-
-# ── Decode at moderate-long context (regression guard) ───────────────
-
-
-def test_decode_long_context_32k_completes():
-    """Decode at S_kv=32K (C=1, P=8) — per-rank S_local=4K. Should
-    complete with current scratch usage (~few KB intermediates) and
-    must continue to complete after the rewrite.
-
-    This is a regression guard: scratch discipline shouldn't break
-    decode's existing long-context capability.
-    """
-    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
-
-    def _bench_fn(ctx):
-        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
-        P = 8
-        S_kv = 32_768
-        dp_full = DPPolicy(cube="replicate", pe="replicate",
-                           num_cubes=1, num_pes=P)
-        dp_kv = DPPolicy(cube="replicate", pe="row_wise",
-                         num_cubes=1, num_pes=P)
-        ctx.zeros((1, 8 * D_HEAD), dtype=DTYPE, dp=dp_full, name="q_long_dec")
-        k = ctx.zeros((S_kv, D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name="k_long_dec")
-        v = ctx.zeros((S_kv, D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name="v_long_dec")
-        o = ctx.empty((1, 8 * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name="o_long_dec")
-        q = ctx.zeros((1, 8 * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name="q_long_dec_2")
-        ctx.launch(
-            "gqa_decode_long_32k",
-            gqa_attention_decode_long_kernel,
-            q, k, v, o,
-            1, S_kv, 8, 1, D_HEAD,
-            1, P,
-            _auto_dim_remap=False,
-        )
-
-    result = run_bench(
-        topology=topo, bench_fn=_bench_fn,
-        device=resolve_device(None), engine_factory=_engine_factory,
-    )
-    assert result.completion.ok, (
-        f"decode at S_kv=32K must complete; got {result.completion}"
-    )
-
-
-# ── Prefill at long context overflows without scope discipline ───────
-
-
-def test_prefill_long_context_completes_after_scope_discipline():
-    """ADR-0063 §A.2 Test Req 3: a sweep at an ``S`` that would overflow
-    1 MiB without recycling must complete after scope discipline lands.
-
-    With C=4 and S_kv chosen so per-rank S_local·8·4 > 1 MiB
-    (~16K tokens per rank ⇒ S_kv ≥ 64K), the current kernel overflows
-    the 1 MiB scratch pool. After Phase 2 (scratch_scope wraps each
-    ring step's intermediates; copy_to persists running state), it
-    completes.
-
-    This is the headline test that proves the S-ceiling is gone for
-    prefill.
-    """
-    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
-
-    def _bench_fn(ctx):
-        configure_sfr_intercube_ring(
-            ctx.engine, ctx.spec, _ccl_cfg(), ring_size=4,
-        )
-        T_q = 4
-        S_kv = 65_536  # 64 K — per-rank 16 K, ~2 MB scratch w/o scope
-        C = 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="q_long_pre")
-        k = ctx.zeros((S_kv, D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name="k_long_pre")
-        v = ctx.zeros((S_kv, D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name="v_long_pre")
-        o = ctx.empty((T_q * C, D_HEAD),
-                      dtype=DTYPE, dp=dp_o, name="o_long_pre")
-        ctx.launch(
-            "gqa_prefill_long_64k",
-            gqa_attention_prefill_long_kernel,
-            q, k, v, o,
-            T_q, S_kv, D_HEAD, C,
-            _auto_dim_remap=False,
-        )
-
-    result = run_bench(
-        topology=topo, bench_fn=_bench_fn,
-        device=resolve_device(None), engine_factory=_engine_factory,
-    )
-    assert result.completion.ok, (
-        f"prefill at S_kv=64K must complete after scope discipline lands; "
-        f"got {result.completion}"
-    )

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_gqa_scoped_writeback.py

@@ -1,198 +0,0 @@
-"""Phase 1 spec test for Phase C: scratch_scope + tl.copy_to discipline
-in the GQA kernels (ADR-0060 §5.2 / §5.5 + ADR-0063 §D3 / §D3.1).
-
-ADR-0060 §5.2 (decode pseudocode line 75 / §5.5 (prefill pseudocode line
-96) both wrap per-tile / per-ring-step intermediates in
-``with tl.scratch_scope():`` and persist the merged running ``(m, ℓ, O)``
-to a persistent arena allocated outside the scope. The original ADR-0063
-§D3 specifies the two-arena pattern; §D3.1 specifies the
-``tl.copy_to(dst, src)`` writeback primitive used to persist scoped
-results.
-
-Currently neither kernel uses ``scratch_scope`` or ``copy_to``; their
-chain-merge / ring-merge bodies allocate every intermediate from the
-bump cursor and never recycle. Result: op_log has 0 ``copy`` entries.
-
-After Phase 2: each per-tile / per-step merge writes the new running
-``(m, ℓ, O)`` via ``copy_to`` to the persistent arena. Per merge step
-→ 3 copy ops (m, ℓ, O).
-
-Phase 1 (this commit): tests only — production code lands in Phase 2.
-"""
-from __future__ import annotations
-
-from pathlib import Path
-
-from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel  # noqa: F401
-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_multisip,
-    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)
-
-
-# ── Decode chain merges must use scratch_scope + copy_to ─────────────
-
-
-def _run_decode_sp(*, h_q: int, h_kv: int, P: int, S_kv: int):
-    """Single-CUBE SP decode (C=1, P PEs along intra-cube chain)."""
-    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=1, num_pes=P)
-        dp_kv = DPPolicy(cube="replicate", pe="row_wise",
-                         num_cubes=1, num_pes=P)
-        q = ctx.zeros((1, h_q * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name=f"q_sc_{P}")
-        k = ctx.zeros((S_kv, h_kv * D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"k_sc_{P}")
-        v = ctx.zeros((S_kv, h_kv * D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"v_sc_{P}")
-        o = ctx.empty((1, h_q * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name=f"o_sc_{P}")
-        ctx.launch(
-            f"gqa_decode_scoped_{P}",
-            gqa_attention_decode_long_kernel,
-            q, k, v, o,
-            1, S_kv, h_q, h_kv, D_HEAD,
-            1, P,
-            _auto_dim_remap=False,
-        )
-
-    return run_bench(
-        topology=topo, bench_fn=_bench_fn,
-        device=resolve_device(None), engine_factory=_engine_factory,
-    )
-
-
-def test_decode_chain_merges_emit_copy_to_writeback():
-    """ADR-0060 §5.2 + ADR-0063 §D3.1: each chain-merge step must wrap
-    its intermediates in ``scratch_scope`` and persist the new running
-    ``(m, ℓ, O)`` via ``tl.copy_to``.
-
-    For (C=1, P=8): 7 intra-cube chain merges × 3 handles (m, ℓ, O) per
-    merge ⇒ 21 ``copy`` entries.
-
-    Currently 0 because the kernel never calls ``tl.copy_to``.
-    """
-    result = _run_decode_sp(h_q=8, h_kv=1, P=8, S_kv=64)
-    assert result.completion.ok, f"decode SP failed: {result.completion}"
-    n_copy = _count(result.engine.op_log, "copy")
-    assert n_copy > 0, (
-        f"decode kernel must emit copy_to writeback per merge step "
-        f"(ADR-0060 §5.2 + ADR-0063 §D3.1); got 0 ``copy`` entries"
-    )
-
-
-# ── Prefill Ring KV merges must use scratch_scope + copy_to ──────────
-
-
-def _run_prefill_ring(*, T_q: int, S_kv: int, C: int):
-    """Head-parallel prefill with Ring KV across C CUBEs."""
-    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
-
-    def _bench_fn(ctx):
-        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)
-        dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
-                         pe="replicate", num_cubes=C, num_pes=1)
-        dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate",
-                        pe="replicate", num_cubes=C, num_pes=1)
-        q = ctx.zeros((T_q, D_HEAD),
-                      dtype=DTYPE, dp=dp_q, name=f"q_ring_{C}")
-        k = ctx.zeros((S_kv, D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"k_ring_{C}")
-        v = ctx.zeros((S_kv, D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"v_ring_{C}")
-        o = ctx.empty((T_q * C, D_HEAD),
-                      dtype=DTYPE, dp=dp_o, name=f"o_ring_{C}")
-        ctx.launch(
-            f"gqa_prefill_scoped_{C}",
-            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,
-    )
-
-
-def test_prefill_ring_step_merges_emit_copy_to_writeback():
-    """ADR-0060 §5.5 + ADR-0063 §D3.1: each Ring KV step's online-softmax
-    merge must wrap its intermediates in ``scratch_scope`` and persist
-    the new running ``(m, ℓ, O)`` via ``tl.copy_to``.
-
-    For C=4: 3 ring-step merges (steps 1..C-1) × 3 handles (m, ℓ, O) per
-    merge ⇒ 9 ``copy`` entries per participating CUBE. Aggregated across
-    C CUBEs: ⇒ 36 ``copy`` entries.
-
-    Currently 0 because the kernel never calls ``tl.copy_to``.
-    """
-    result = _run_prefill_ring(T_q=4, S_kv=16, C=4)
-    assert result.completion.ok, f"prefill ring failed: {result.completion}"
-    n_copy = _count(result.engine.op_log, "copy")
-    assert n_copy > 0, (
-        f"prefill ring kernel must emit copy_to writeback per merge step "
-        f"(ADR-0060 §5.5 + ADR-0063 §D3.1); got 0 ``copy`` entries"
-    )
-
-
-# ── Scoped kernels still produce the same op_log shape (regression) ──
-
-
-def test_decode_scoped_still_has_root_only_write():
-    """ADR-0060 §A.2 root-only output must hold under the rewrite:
-    adding scratch_scope + copy_to should not change the reduce
-    topology; only the per-PE scratch usage. Single PE 0 writes O."""
-    result = _run_decode_sp(h_q=8, h_kv=1, P=8, S_kv=64)
-    assert result.completion.ok
-    n_writes = _count(result.engine.op_log, "dma_write")
-    assert n_writes == 1, (
-        f"scoped decode must still produce 1 dma_write (root-only); "
-        f"got {n_writes}"
-    )
-
-
-def test_prefill_scoped_still_has_per_cube_distributed_output():
-    """ADR-0060 §5.5 per-CUBE distributed output must hold under the
-    rewrite: scoped prefill still writes one O slice per CUBE."""
-    result = _run_prefill_ring(T_q=4, S_kv=16, C=4)
-    assert result.completion.ok
-    n_writes = _count(result.engine.op_log, "dma_write")
-    assert n_writes == 4, (
-        f"scoped prefill must write one O per CUBE (C=4 → 4 dma_write); "
-        f"got {n_writes}"
-    )

tests/attention/test_gqa_tile_sweep.py

@@ -1,290 +0,0 @@
-"""Phase 1 spec test for P3b: S_kv tile sweep in the GQA kernels
-(ADR-0063 §A.2 + ADR-0060 §B "long/short context split").
-
-The local one-shot partial in three kernels — ``_gqa_attention_decode_long.py``,
-``_gqa_attention_decode_short.py``, ``_gqa_attention_prefill_short.py`` — loads its rank's
-entire ``(d_head, S_local)`` / ``(S_local, d_head)`` KV slice in one
-shot, so per-rank scratch grows linearly with ``S_local``. The
-``_attention_*`` baselines hit a TCM ceiling once ``S_local`` exceeds
-the 1 MiB pool divided by the per-rank intermediate footprint.
-
-P3b replaces the one-shot partial with a tile sweep:
-  - Module-level ``TILE_S_KV = 1024`` per kernel file.
-  - Tile 0 establishes the persistent ``(m_local, l_local, O_local)``.
-  - Tiles 1..n_tiles-1 wrap their intermediates (K_T tile, V tile,
-    scores, centered, exp_scores, partials) in ``tl.scratch_scope()``
-    and persist the merged ``(m, ℓ, O)`` to the persistent handles via
-    ``tl.copy_to`` (ADR-0063 §D3 / §D3.1).
-
-When ``S_local ≤ TILE_S_KV`` the loop body never runs and the op_log
-is structurally identical to today's one-shot path — every existing
-validation-scale test continues to pass.
-
-``_gqa_attention_prefill_long.py`` is **out of scope** for P3b. Its inner step
-is IPCQ partial-recv (Ring KV rotation), not an HBM load; tiling it
-intersects the ring-step structure and is a separate phase.
-
-Phase 1 (this commit): tests only — production code lands in Phase 2.
-The four multi-tile tests fail today because the kernels never call
-``copy_to`` outside the chain-reduce merges (so isolating to a
-chain-free config gives ``copy_to == 0`` today). The 128K test fails
-today with a ``TLContext`` scratch overflow.
-"""
-from __future__ import annotations
-
-from pathlib import Path
-
-from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel  # noqa: F401
-from kernbench.benches._gqa_attention_decode_short import gqa_attention_decode_short_kernel  # noqa: F401
-from kernbench.benches._gqa_attention_prefill_short import gqa_attention_prefill_short_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"
-
-
-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)
-
-
-# ── decode_long runner (chain-free configs isolate tile-sweep behavior) ──
-
-
-def _run_decode_long(*, C: int, P: int, S_kv: int, h_q: int = 1, h_kv: int = 1):
-    """Run the long-context decode kernel.
-
-    For tile-sweep isolation, callers pass C=1, P=1 — this leaves both
-    chain-reduce levels inactive so any ``copy_to`` in op_log comes
-    solely from the tile-merge body.
-    """
-    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" if P > 1 else "replicate",
-                         num_cubes=C, num_pes=P)
-        q = ctx.zeros((1, h_q * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name=f"q_tl_c{C}_p{P}_s{S_kv}")
-        k = ctx.zeros((S_kv, h_kv * D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"k_tl_c{C}_p{P}_s{S_kv}")
-        v = ctx.zeros((S_kv, h_kv * D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"v_tl_c{C}_p{P}_s{S_kv}")
-        o = ctx.empty((1, h_q * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name=f"o_tl_c{C}_p{P}_s{S_kv}")
-        ctx.launch(
-            f"gqa_decode_long_tile_c{C}_p{P}_s{S_kv}",
-            gqa_attention_decode_long_kernel,
-            q, k, v, o,
-            1, S_kv, h_q, h_kv, D_HEAD, C, P,
-            _auto_dim_remap=False,
-        )
-
-    return run_bench(
-        topology=topo, bench_fn=_bench_fn,
-        device=resolve_device(None), engine_factory=_engine_factory,
-    )
-
-
-# ── decode_short / prefill_short runners ─────────────────────────────
-
-
-def _run_decode_short(*, kv_per_cube: int, C: int, P: int, S_kv: int,
-                      h_q: int = 8, h_kv: int = 8):
-    """For tile-sweep isolation: kv_per_cube=P → group_size=1 (no chain)."""
-    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)
-        q = ctx.zeros((1, h_q * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name=f"q_dsh_s{S_kv}")
-        k = ctx.zeros((h_kv * S_kv, D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"k_dsh_s{S_kv}")
-        v = ctx.zeros((h_kv * S_kv, D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"v_dsh_s{S_kv}")
-        o = ctx.empty((1, h_q * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name=f"o_dsh_s{S_kv}")
-        ctx.launch(
-            f"gqa_decode_short_tile_s{S_kv}",
-            gqa_attention_decode_short_kernel,
-            q, k, v, o,
-            1, S_kv, h_q, h_kv, D_HEAD, C, P, kv_per_cube,
-            _auto_dim_remap=False,
-        )
-
-    return run_bench(
-        topology=topo, bench_fn=_bench_fn,
-        device=resolve_device(None), engine_factory=_engine_factory,
-    )
-
-
-def _run_prefill_short(*, kv_per_cube: int, C: int, P: int,
-                       T_q: int, S_kv: int, h_kv: int = 8):
-    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)
-        q = ctx.zeros((T_q, h_kv * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name=f"q_psh_s{S_kv}")
-        k = ctx.zeros((h_kv * S_kv, D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"k_psh_s{S_kv}")
-        v = ctx.zeros((h_kv * S_kv, D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"v_psh_s{S_kv}")
-        o = ctx.empty((T_q, h_kv * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name=f"o_psh_s{S_kv}")
-        ctx.launch(
-            f"gqa_prefill_short_tile_s{S_kv}",
-            gqa_attention_prefill_short_kernel,
-            q, k, v, o,
-            T_q, S_kv, h_kv, D_HEAD, C, P, kv_per_cube,
-            _auto_dim_remap=False,
-        )
-
-    return run_bench(
-        topology=topo, bench_fn=_bench_fn,
-        device=resolve_device(None), engine_factory=_engine_factory,
-    )
-
-
-# ── T1: single-tile path is op_log-stable (regression guard) ─────────
-
-
-def test_decode_long_single_tile_path_unchanged():
-    """ADR-0063 §A.2: when S_local <= TILE_S_KV the tile-sweep loop
-    body never runs and op_log is structurally identical to today's
-    one-shot path.
-
-    Passes today (one-shot path) AND after Phase 2 (n_tiles=1 skips
-    the merge loop). The witness: with chain reduce inactive (C=1,
-    P=1), there must be zero ``copy_to`` entries — neither today nor
-    after Phase 2 should the kernel emit tile-merge writebacks here.
-    """
-    result = _run_decode_long(C=1, P=1, S_kv=64)
-    assert result.completion.ok, (
-        f"single-tile decode_long must complete; got {result.completion}"
-    )
-    n_copy = _count(result.engine.op_log, "copy")
-    assert n_copy == 0, (
-        f"S_local <= TILE_S_KV must not emit tile-merge copy_to; got {n_copy}"
-    )
-
-
-# ── T2: decode_long multi-tile emits tile-merge copy_to ──────────────
-
-
-def test_decode_long_multi_tile_emits_copy_to_merges():
-    """ADR-0063 §A.2 + §D3.1: when S_local > TILE_S_KV the kernel must
-    wrap each subsequent tile's intermediates in ``scratch_scope`` and
-    persist the merged ``(m, ℓ, O)`` via ``tl.copy_to``.
-
-    Chain-free config (C=1, P=1) isolates the tile-merge body — any
-    ``copy_to`` in op_log comes from the tile sweep, not chain reduce.
-
-    S_kv=2048, C=1, P=1 → S_local=2048 → 2 tiles → 1 merge × 3 handles
-    (m, ℓ, O) ⇒ 3 ``copy`` entries.
-
-    Currently 0 because the kernel performs a one-shot partial.
-    """
-    result = _run_decode_long(C=1, P=1, S_kv=2048)
-    assert result.completion.ok, (
-        f"multi-tile decode_long must complete; got {result.completion}"
-    )
-    n_copy = _count(result.engine.op_log, "copy")
-    assert n_copy >= 3, (
-        f"decode_long multi-tile must emit >= 3 copy_to entries "
-        f"(1 merge × 3 handles); got {n_copy}"
-    )
-
-
-# ── T3: decode_short multi-tile emits tile-merge copy_to ─────────────
-
-
-def test_decode_short_multi_tile_emits_copy_to_merges():
-    """Same property as T2 for the short decode kernel.
-
-    kv_per_cube=8, P=8, C=1 → group_size=1 (no chain reduce); S_kv=2048
-    → S_local=2048 → 2 tiles per PE → 3 ``copy`` entries per PE × 8 PEs
-    = 24 total.
-
-    Currently 0 because the kernel performs a one-shot partial.
-    """
-    result = _run_decode_short(kv_per_cube=8, C=1, P=8, S_kv=2048)
-    assert result.completion.ok, (
-        f"multi-tile decode_short must complete; got {result.completion}"
-    )
-    n_copy = _count(result.engine.op_log, "copy")
-    assert n_copy >= 3 * 8, (
-        f"decode_short multi-tile must emit >= 24 copy_to entries "
-        f"(1 merge × 3 handles × 8 PEs); got {n_copy}"
-    )
-
-
-# ── T4: prefill_short multi-tile emits tile-merge copy_to ────────────
-
-
-def test_prefill_short_multi_tile_emits_copy_to_merges():
-    """Same property as T3 for the short prefill kernel.
-
-    kv_per_cube=8, P=8, C=1, T_q=4, S_kv=2048 → group_size=1 (no chain),
-    S_local=2048 → 2 tiles per PE → 3 ``copy`` entries per PE × 8 PEs
-    = 24 total.
-
-    Currently 0 because the kernel performs a one-shot partial.
-    """
-    result = _run_prefill_short(
-        kv_per_cube=8, C=1, P=8, T_q=4, S_kv=2048,
-    )
-    assert result.completion.ok, (
-        f"multi-tile prefill_short must complete; got {result.completion}"
-    )
-    n_copy = _count(result.engine.op_log, "copy")
-    assert n_copy >= 3 * 8, (
-        f"prefill_short multi-tile must emit >= 24 copy_to entries "
-        f"(1 merge × 3 handles × 8 PEs); got {n_copy}"
-    )
-
-
-# ── T5: decode_long at S_kv=256K completes (TCM ceiling lifted) ──────
-
-
-def test_decode_long_context_256k_completes():
-    """ADR-0063 §A.2 (test req 3): headline ceiling-lift. C=1 P=8
-    decode at S_kv=256K → S_local=32K. Today the per-rank score stack
-    (3 ×M·S_local·2 bytes, M=G·T_q=8) is ~1.5 MB → exceeds the 1 MiB
-    scratch pool → TLContext scratch overflow. After Phase 2 the tile
-    sweep bounds per-tile score stack at 3·M·TILE_S_KV·2 ≈ 48 KB and
-    the run completes regardless of S_kv.
-    """
-    result = _run_decode_long(C=1, P=8, S_kv=262144, h_q=8, h_kv=1)
-    assert result.completion.ok, (
-        f"decode_long at S_kv=256K must complete after tile sweep; "
-        f"got {result.completion}"
-    )

tests/attention/test_milestone_gqa_decode_long_ctx_4cases.py

@@ -0,0 +1,524 @@
+"""Tests for the long-context decode 4-cases comparative-study bench.
+
+Per ``GQA_full_deck.pptx`` slides 11-17, the 4 cases differ in how KV
+cache is sharded across the 8 cubes and 8 PEs of a single KV-head
+group on LLaMA-3.1-70B GQA:
+
+    Case 1  Cube-SP / PE-TP  → KV split by S_kv across cubes; PEs TP on batch
+    Case 2  Cube-Repl / PE-TP → full KV per cube; PEs TP on batch
+    Case 3  Cube-Repl / PE-SP → full KV per cube; PEs SP on S_kv (intra-cube AR)
+    Case 4  Cube-SP / PE-SP  → KV split 64-way; 2-phase AR on (m,ℓ,O)  ★ optimal
+
+Deviation from slide 13: slide prescribes AllReduce on (m,ℓ,O); the
+kernel does reduce-to-root (only the lrab center cube has the answer)
+per ADR-0060 §4. Treated as the kernbench Case-4 baseline.
+"""
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+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"
+
+_CASE4_PANEL = "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp"
+_CASE3_PANEL = "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp"
+_CASE2_PANEL = "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp"
+_CASE1_PANEL = "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp"
+_CUBE_RE = re.compile(r"\bcube(\d+)\b")
+
+
+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 _dma_write_cubes(op_log) -> list[int]:
+    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
+
+
+# ── Case 4 (Cube-SP × PE-SP) — ★ optimal ────────────────────────────
+
+
+def _run_case4_smoke(*, S_kv: int):
+    """Drive the Case 4 decode panel via the case-specific runner.
+
+    Uses ``S_kv=8192`` (smoke) to keep test time bounded; the headline
+    ``S_kv=128K`` runs come from ``kernbench run --bench
+    milestone-gqa-decode-long-ctx-4cases``, not pytest.
+    """
+    from kernbench.benches.milestone_gqa_decode_long_ctx_4cases import (
+        _run_decode_panel_long_ctx_cube_sp_pe_sp,
+    )
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        _run_decode_panel_long_ctx_cube_sp_pe_sp(
+            ctx, panel=_CASE4_PANEL,
+            C=8, P=8,
+            T_q=1, S_kv=S_kv,
+            d_head=128, h_q=8, h_kv=1,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None),
+        engine_factory=_engine_factory,
+    )
+
+
+# ── Case 4 — T1: panel is registered in the bench ───────────────────
+
+
+def test_case4_panel_registered():
+    """The Case 4 panel must be in the bench's ``_PANELS`` +
+    ``_PANEL_DISPATCH`` with the expected LLaMA-3.1-70B target dims.
+
+    Headline config:
+      C = 8                 (head-parallel CUBE Group)
+      P = 8                 (intra-CUBE PE-SP)
+      T_q = 1               (decode: one new token per pass)
+      S_kv = 131_072        (LLaMA long-context decode target)
+      d_head = 128, h_q = 8, h_kv = 1
+
+    The lrab sub_w=4 / sub_h=2 geometry is baked into the Case 4
+    wrapper kernel; it is not a panel parameter.
+    """
+    from kernbench.benches.milestone_gqa_decode_long_ctx_4cases import (
+        _PANEL_DISPATCH,
+        _PANELS,
+    )
+    assert _CASE4_PANEL in _PANELS, (
+        f"{_CASE4_PANEL!r} not in _PANELS; got {_PANELS}"
+    )
+    assert _CASE4_PANEL in _PANEL_DISPATCH
+    kind, params = _PANEL_DISPATCH[_CASE4_PANEL]
+    assert kind == "decode_long_ctx_cube_sp_pe_sp", (
+        f"kind={kind!r}, expected 'decode_long_ctx_cube_sp_pe_sp'"
+    )
+    assert params.get("C") == 8
+    assert params.get("P") == 8
+    assert params.get("T_q") == 1
+    assert params.get("S_kv") == 131_072
+    assert params.get("d_head") == 128
+    assert params.get("h_q") == 8
+    assert params.get("h_kv") == 1
+
+
+# ── Case 4 — T2: runner drives the kernel to completion ─────────────
+
+
+def test_case4_runner_smoke():
+    """Case 4 runner drives the new kernel to completion at smoke S_kv."""
+    result = _run_case4_smoke(S_kv=8192)
+    assert result.completion.ok, (
+        f"Case 4 decode smoke at C=8 P=8 must complete; "
+        f"got {result.completion}"
+    )
+
+
+# ── Case 4 — T3: reduce-to-root at the lrab center cube (cube 6) ────
+
+
+def test_case4_root_at_center_cube_6():
+    """For ``sub_w=4, sub_h=2``: root_col=2, root_row=1, root_cube=6.
+    The decode kernel writes the final O exclusively from PE 0 of cube
+    6 (ADR-0060 §4 reduce-to-root variant of the Case-4 AR pattern).
+    """
+    result = _run_case4_smoke(S_kv=8192)
+    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"Case 4 root must be the lrab center cube 6; "
+        f"got cubes={sorted(distinct)}"
+    )
+
+
+# ── Case 4 — T4: 2-phase AR ipcq pattern matches Case-4 traffic ─────
+
+
+def test_case4_two_level_ar_ipcq_pattern():
+    """Total ipcq_copy for the Case 4 reduce at (C, P, sub_w) =
+    (8, 8, 4):
+
+      Intra-CUBE (per CUBE = 8 PEs in a 2×4 grid):
+        row chain along intra_W: cols 1,2,3 each row × 2 rows ×
+          3 tensors                                       = 18
+        col bridge along intra_N: pe4 only × 3 tensors    =  3
+        per-CUBE intra total                              = 21
+        × 8 CUBEs                                         = 168
+
+      Inter-CUBE lrab (sub_w=4, sub_h=2):
+        Phase 1 row reduce — 3 sends/row × 3 tensors × 2 rows = 18
+        Phase 2 col reduce — cube 2 → S × 3 tensors           =  3
+        inter-CUBE total                                      = 21
+
+      Grand total: 168 + 21 = 189
+    """
+    result = _run_case4_smoke(S_kv=8192)
+    assert result.completion.ok
+    n_copy = _count(result.engine.op_log, "ipcq_copy")
+    assert n_copy == 189, (
+        f"Case 4 expected 189 ipcq_copy "
+        f"(168 intra-CUBE + 21 inter-CUBE lrab); got {n_copy}"
+    )
+
+
+# ── Case 2 (Cube-Repl × PE-TP) ──────────────────────────────────────
+
+
+def _run_case2_smoke(*, S_kv: int):
+    """Drive the Case 2 decode panel via the case-specific runner.
+
+    Case 2 = Cube-Repl × PE-TP. K, V are replicated everywhere (the
+    slide-11 memory waste); for B=1 only one rank does the work; no
+    inter-rank comm.
+    """
+    from kernbench.benches.milestone_gqa_decode_long_ctx_4cases import (
+        _run_decode_panel_long_ctx_cube_repl_pe_tp,
+    )
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        _run_decode_panel_long_ctx_cube_repl_pe_tp(
+            ctx, panel=_CASE2_PANEL,
+            C=8, P=8,
+            T_q=1, S_kv=S_kv,
+            d_head=128, h_q=8, h_kv=1,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None),
+        engine_factory=_engine_factory,
+    )
+
+
+# ── Case 2 — T1: panel registered ───────────────────────────────────
+
+
+def test_case2_panel_registered():
+    """The Case 2 panel must be in the bench's ``_PANELS`` +
+    ``_PANEL_DISPATCH`` with the expected single-KV-group dims.
+
+    Case 2: Cube-Repl × PE-TP. K, V replicated everywhere
+    (8 KB/tok/PE — slide-11 memory waste); no inter-rank comm.
+    For B=1 only one rank works (PEs 1-7 idle — slide-11 calls
+    out this PE-TP waste).
+    """
+    from kernbench.benches.milestone_gqa_decode_long_ctx_4cases import (
+        _PANEL_DISPATCH,
+        _PANELS,
+    )
+    assert _CASE2_PANEL in _PANELS, (
+        f"{_CASE2_PANEL!r} not in _PANELS; got {_PANELS}"
+    )
+    assert _CASE2_PANEL in _PANEL_DISPATCH
+    kind, params = _PANEL_DISPATCH[_CASE2_PANEL]
+    assert kind == "decode_long_ctx_cube_repl_pe_tp"
+    assert params.get("C") == 8
+    assert params.get("P") == 8
+    assert params.get("T_q") == 1
+    assert params.get("S_kv") == 131_072
+    assert params.get("d_head") == 128
+    assert params.get("h_q") == 8
+    assert params.get("h_kv") == 1
+
+
+# ── Case 2 — T2: smoke runner completes ─────────────────────────────
+
+
+def test_case2_runner_smoke():
+    """Case 2 runner drives the new kernel to completion at smoke S_kv."""
+    result = _run_case2_smoke(S_kv=8192)
+    assert result.completion.ok, (
+        f"Case 2 decode smoke at C=8 P=8 must complete; "
+        f"got {result.completion}"
+    )
+
+
+# ── Case 2 — T3: zero inter-rank comm by design ─────────────────────
+
+
+def test_case2_zero_ipcq_copy_no_comm():
+    """Case 2's defining property: full KV per rank ⇒ NO inter-rank
+    communication. Slide 11 lists comm cost as 'none'.
+    """
+    result = _run_case2_smoke(S_kv=8192)
+    assert result.completion.ok
+    n_copy = _count(result.engine.op_log, "ipcq_copy")
+    assert n_copy == 0, (
+        f"Case 2 must have zero inter-rank comm; got ipcq_copy={n_copy}"
+    )
+
+
+# ── Case 2 — T4: single dma_write from cube 0 (B=1 single-rank work) ─
+
+
+def test_case2_single_dma_write_at_cube_0():
+    """For B=1, only PE 0 of CUBE 0 does the work (the inherent PE-TP
+    waste at B=1). Exactly 1 dma_write, from cube 0.
+    """
+    result = _run_case2_smoke(S_kv=8192)
+    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 == {0}, (
+        f"Case 2 B=1 single writer must be cube 0; "
+        f"got cubes={sorted(distinct)}"
+    )
+
+
+# ── Case 3 (Cube-Repl × PE-SP) ──────────────────────────────────────
+
+
+def _run_case3_smoke(*, S_kv: int):
+    """Drive the Case 3 decode panel via the case-specific runner.
+
+    Case 3 = Cube-Repl × PE-SP. K, V replicated per cube; S_kv split
+    8-way across PEs within each cube. Intra-CUBE 8-way reduce on
+    (m, ℓ, O); no inter-CUBE comm (every cube ends with full answer).
+    """
+    from kernbench.benches.milestone_gqa_decode_long_ctx_4cases import (
+        _run_decode_panel_long_ctx_cube_repl_pe_sp,
+    )
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        _run_decode_panel_long_ctx_cube_repl_pe_sp(
+            ctx, panel=_CASE3_PANEL,
+            C=8, P=8,
+            T_q=1, S_kv=S_kv,
+            d_head=128, h_q=8, h_kv=1,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None),
+        engine_factory=_engine_factory,
+    )
+
+
+# ── Case 3 — T1: panel registered ───────────────────────────────────
+
+
+def test_case3_panel_registered():
+    """The Case 3 panel must be in the bench's ``_PANELS`` +
+    ``_PANEL_DISPATCH`` with the expected single-KV-group dims.
+
+    Case 3: Cube-Repl × PE-SP. K, V replicated per cube; PEs SP on
+    S_kv. Intra-CUBE 8-way AR; no inter-CUBE comm.
+    """
+    from kernbench.benches.milestone_gqa_decode_long_ctx_4cases import (
+        _PANEL_DISPATCH,
+        _PANELS,
+    )
+    assert _CASE3_PANEL in _PANELS, (
+        f"{_CASE3_PANEL!r} not in _PANELS; got {_PANELS}"
+    )
+    assert _CASE3_PANEL in _PANEL_DISPATCH
+    kind, params = _PANEL_DISPATCH[_CASE3_PANEL]
+    assert kind == "decode_long_ctx_cube_repl_pe_sp"
+    assert params.get("C") == 8
+    assert params.get("P") == 8
+    assert params.get("T_q") == 1
+    assert params.get("S_kv") == 131_072
+    assert params.get("d_head") == 128
+    assert params.get("h_q") == 8
+    assert params.get("h_kv") == 1
+
+
+# ── Case 3 — T2: smoke runner completes ─────────────────────────────
+
+
+def test_case3_runner_smoke():
+    """Case 3 runner drives the new kernel to completion at smoke S_kv."""
+    result = _run_case3_smoke(S_kv=8192)
+    assert result.completion.ok, (
+        f"Case 3 decode smoke at C=8 P=8 must complete; "
+        f"got {result.completion}"
+    )
+
+
+# ── Case 3 — T3: intra-CUBE-only AR (168 ipcq, no inter-CUBE) ───────
+
+
+def test_case3_intra_cube_ar_only_ipcq():
+    """Case 3 reduce pattern: per-CUBE 8-way PE-SP AR (same structural
+    cost as Case 4's intra-CUBE phase = 21 ipcq_copy per cube), and
+    NO inter-CUBE traffic (each cube has a full copy of KV).
+
+      per-CUBE intra (2×4 PE grid):
+        row chain along intra_W: cols 1,2,3 each row × 2 rows ×
+          3 tensors (m, ℓ, O)                          = 18
+        col bridge along intra_N: pe4 only × 3 tensors =  3
+        per-CUBE intra total                           = 21
+        × 8 CUBEs                                      = 168
+
+      inter-CUBE: 0 (replicated KV ⇒ no AllReduce needed).
+
+      Grand total: 168.
+    """
+    result = _run_case3_smoke(S_kv=8192)
+    assert result.completion.ok
+    n_copy = _count(result.engine.op_log, "ipcq_copy")
+    assert n_copy == 168, (
+        f"Case 3 expected 168 ipcq_copy (intra-CUBE only, no inter-CUBE); "
+        f"got {n_copy}"
+    )
+
+
+# ── Case 3 — T4: single dma_write from cube 0 (designated writer) ───
+
+
+def test_case3_single_dma_write_at_cube_0():
+    """Every cube ends with the full answer after intra-CUBE AR; only
+    the designated writer (cube 0, PE 0) stores O to avoid 8 redundant
+    DMAs.
+    """
+    result = _run_case3_smoke(S_kv=8192)
+    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 == {0}, (
+        f"Case 3 designated writer must be cube 0; "
+        f"got cubes={sorted(distinct)}"
+    )
+
+
+# ── Case 1 (Cube-SP × PE-TP) ────────────────────────────────────────
+
+
+def _run_case1_smoke(*, S_kv: int):
+    """Drive the Case 1 decode panel via the case-specific runner.
+
+    Case 1 = Cube-SP × PE-TP. K, V split S_kv-wise across the 8 cubes
+    (cube=row_wise), replicated within each cube (pe=replicate). PEs
+    nominally split on the batch dim (PE-TP); at B=1 only PE 0 of
+    each cube has work; PEs 1-7 idle. Inter-CUBE 8-way reduce via the
+    lrab-adapted center-root pattern (root cube 6); no intra-CUBE comm.
+    """
+    from kernbench.benches.milestone_gqa_decode_long_ctx_4cases import (
+        _run_decode_panel_long_ctx_cube_sp_pe_tp,
+    )
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        _run_decode_panel_long_ctx_cube_sp_pe_tp(
+            ctx, panel=_CASE1_PANEL,
+            C=8, P=8,
+            T_q=1, S_kv=S_kv,
+            d_head=128, h_q=8, h_kv=1,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None),
+        engine_factory=_engine_factory,
+    )
+
+
+# ── Case 1 — T1: panel registered ───────────────────────────────────
+
+
+def test_case1_panel_registered():
+    """The Case 1 panel must be in the bench's ``_PANELS`` +
+    ``_PANEL_DISPATCH`` with the expected single-KV-group dims.
+
+    Case 1: Cube-SP × PE-TP. K, V split across cubes; PEs TP on
+    batch. At B=1 only PE 0 of each cube works (PE-TP waste).
+    Inter-CUBE lrab AR; no intra-CUBE comm.
+    """
+    from kernbench.benches.milestone_gqa_decode_long_ctx_4cases import (
+        _PANEL_DISPATCH,
+        _PANELS,
+    )
+    assert _CASE1_PANEL in _PANELS, (
+        f"{_CASE1_PANEL!r} not in _PANELS; got {_PANELS}"
+    )
+    assert _CASE1_PANEL in _PANEL_DISPATCH
+    kind, params = _PANEL_DISPATCH[_CASE1_PANEL]
+    assert kind == "decode_long_ctx_cube_sp_pe_tp"
+    assert params.get("C") == 8
+    assert params.get("P") == 8
+    assert params.get("T_q") == 1
+    assert params.get("S_kv") == 131_072
+    assert params.get("d_head") == 128
+    assert params.get("h_q") == 8
+    assert params.get("h_kv") == 1
+
+
+# ── Case 1 — T2: smoke runner completes ─────────────────────────────
+
+
+def test_case1_runner_smoke():
+    """Case 1 runner drives the new kernel to completion at smoke S_kv."""
+    result = _run_case1_smoke(S_kv=8192)
+    assert result.completion.ok, (
+        f"Case 1 decode smoke at C=8 P=8 must complete; "
+        f"got {result.completion}"
+    )
+
+
+# ── Case 1 — T3: inter-CUBE lrab only (21 ipcq, no intra-CUBE) ──────
+
+
+def test_case1_inter_cube_lrab_only_ipcq():
+    """Case 1 reduce pattern: only PE 0 of each cube has work (PE-TP
+    at B=1), so NO intra-CUBE AR. Inter-CUBE 8-way reduce uses the
+    lrab-adapted center-root pattern (same structural cost as Case 4's
+    inter-CUBE phase).
+
+      Inter-CUBE lrab (sub_w=4, sub_h=2):
+        Phase 1 row reduce — 3 sends/row × 3 tensors × 2 rows = 18
+        Phase 2 col reduce — cube 2 → S × 3 tensors           =  3
+        inter-CUBE total                                      = 21
+
+      Intra-CUBE: 0 (PEs 1-7 idle, no partials to merge).
+
+      Grand total: 21.
+    """
+    result = _run_case1_smoke(S_kv=8192)
+    assert result.completion.ok
+    n_copy = _count(result.engine.op_log, "ipcq_copy")
+    assert n_copy == 21, (
+        f"Case 1 expected 21 ipcq_copy (inter-CUBE lrab only, no intra-CUBE); "
+        f"got {n_copy}"
+    )
+
+
+# ── Case 1 — T4: root at lrab center cube 6 ─────────────────────────
+
+
+def test_case1_root_at_center_cube_6():
+    """Case 1 uses the same lrab-adapted center-root reduce as Case 4's
+    inter-CUBE phase; the answer lands at the lrab center cube
+    (sub_w=4, sub_h=2 → root_col=2, root_row=1 → root_cube=6).
+    """
+    result = _run_case1_smoke(S_kv=8192)
+    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"Case 1 root must be the lrab center cube 6; "
+        f"got cubes={sorted(distinct)}"
+    )

tests/attention/test_milestone_gqa_headline.py

@@ -1,25 +1,14 @@
-"""Phase 1 spec test for P7: headline milestone-gqa bench with real GQA.
+"""Tests for the milestone-gqa-headline bench.
 
-P7 wires the new ``_gqa_decode`` and ``_gqa_prefill`` kernels into a
+The bench wires the new ``_gqa_attention_prefill_long`` kernel into a
-new 4-panel milestone bench (independent from the existing
+single-KV-group milestone panel (LLaMA-3.1-70B target, C=8 P=8). The
-``milestone-gqa-llama70b`` which still covers the baseline kernels).
+4 legacy C=1 / C=4 panels (single_user_*, multi_user_*) were removed
-Real GQA (``h_q > h_kv`` with G=8) runs end-to-end through a
+once pytest regression covered those configurations comprehensively
-milestone-style sweep + sweep.json output.
+and the comparative decode work moved into the
+``milestone-gqa-decode-4cases`` bench.
 
-Restriction in P7 first cut:
+Single SIP scope (default ``topology.yaml`` 4×4 cube mesh; 4-SIP
-  - C ≤ 4 (single-row inter-CUBE ring SFR; ADR-0060 §B leaves
+headline deferred per ADR-0060 §B-item-1).
-    multi-row rings for follow-on)
-  - Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP
-    headline left to follow-on)
-  - No figure renderers (defer to a separate cycle)
-
-Panels (4 total):
-  single_user_prefill_gqa: C=1, T_q=4, S_kv=16   (no ring)
-  multi_user_prefill_gqa : C=4, T_q=4, S_kv=16   (Ring KV, 3 steps)
-  single_user_decode_gqa : C=1, P=8, h_q=8, h_kv=1, S_kv=64    (M-fold + intra-cube chain)
-  multi_user_decode_gqa  : C=4, P=8, h_q=8, h_kv=1, S_kv=128   (M-fold + 2-level chain)
-
-Phase 1 (this commit): tests only — bench module lands in Phase 2.
 """
 from __future__ import annotations
 
@@ -35,10 +24,7 @@ from kernbench.topology.builder import resolve_topology
 BENCH_NAME = "milestone-gqa-headline"
 
 PANELS = (
-    "single_user_prefill_gqa",
+    "single_kv_group_prefill_gqa_c8_p8",
-    "multi_user_prefill_gqa",
-    "single_user_decode_gqa",
-    "multi_user_decode_gqa",
 )
 
 
@@ -88,61 +74,16 @@ def test_validation_run_completes_ok(monkeypatch):
 # ── sweep.json shape ──────────────────────────────────────────────────
 
 
-def test_sweep_json_has_four_panels(monkeypatch):
+def test_sweep_json_has_expected_panels(monkeypatch):
     data = _sweep_json(monkeypatch)
     assert set(data["panels"]) == set(PANELS), (
         f"panels mismatch: expected {set(PANELS)}, got {set(data['panels'])}"
     )
-    assert len(data["rows"]) == 4
+    assert len(data["rows"]) == len(PANELS)
     assert {r["panel"] for r in data["rows"]} == set(PANELS)
 
 
-def test_decode_panels_use_real_gqa(monkeypatch):
+# Per-panel architectural assertions for the surviving single-KV-group
-    """ADR-0060 §A.1 headline: decode panels must use h_q = G·h_kv with G>1."""
+# panel live in tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py.
-    data = _sweep_json(monkeypatch)
+# Decode architectural assertions live in
-    cfg = data["config"]
+# tests/attention/test_milestone_gqa_decode_4cases.py.
-    assert cfg["h_q_decode"] > cfg["h_kv_decode"], (
-        f"decode must use real GQA (h_q > h_kv); got "
-        f"h_q={cfg['h_q_decode']}, h_kv={cfg['h_kv_decode']}"
-    )
-
-
-# ── Per-panel architectural assertions ────────────────────────────────
-
-
-def _row(rows, panel: str) -> dict:
-    for r in rows:
-        if r["panel"] == panel:
-            return r
-    raise AssertionError(f"missing row for panel {panel!r}")
-
-
-def test_prefill_ring_panel_has_ipcq_traffic(monkeypatch):
-    """multi_user_prefill_gqa uses Ring KV → IPCQ traffic > 0."""
-    data = _sweep_json(monkeypatch)
-    row = _row(data["rows"], "multi_user_prefill_gqa")
-    n_copy = row["op_log_summary"].get("ipcq_copy_count", 0)
-    assert n_copy > 0, (
-        f"multi_user_prefill_gqa must have Ring KV traffic; got "
-        f"ipcq_copy_count={n_copy}"
-    )
-
-
-def test_decode_reduce_panel_writes_once(monkeypatch):
-    """multi_user_decode_gqa: chain reduce-to-root → exactly 1 dma_write."""
-    data = _sweep_json(monkeypatch)
-    row = _row(data["rows"], "multi_user_decode_gqa")
-    n_writes = row["op_log_summary"]["dma_write_count"]
-    assert n_writes == 1, (
-        f"multi_user_decode_gqa root-only write: expected 1; got {n_writes}"
-    )
-
-
-def test_prefill_panel_distributes_output(monkeypatch):
-    """multi_user_prefill_gqa: per-CUBE distributed output → dma_write_count == C."""
-    data = _sweep_json(monkeypatch)
-    row = _row(data["rows"], "multi_user_prefill_gqa")
-    n_writes = row["op_log_summary"]["dma_write_count"]
-    assert n_writes == 4, (
-        f"multi_user_prefill_gqa per-CUBE distributed: expected 4; got {n_writes}"
-    )

tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py

@@ -0,0 +1,121 @@
+"""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 = 1K
+    (scratch-limited; LLaMA 32K headline awaits Q-axis kernel tiling)
+
+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 (scratch-limited; LLaMA 32K headline awaits
+    Q-axis kernel tiling):
+      C = 8        (head-parallel, snake sub-mesh)
+      P = 8        (intra-CUBE PE-SP)
+      T_q = 1_024  (one-shot long-context prefill)
+      S_kv = 1_024
+      (scratch-limited; LLaMA 32K headline awaits Q-axis kernel tiling)
+      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") == 1_024, (
+        f"T_q={params.get('T_q')}, expected 1_024"
+    )
+    assert params.get("S_kv") == 1_024, (
+        f"S_kv={params.get('S_kv')}, expected 1_024"
+    )
+    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}"
+    )
+
+
+# Backward-compat test removed when the legacy multi_user_prefill_gqa
+# panel was dropped; signature-extension contract is now exercised
+# implicitly by the live single_kv_group_prefill_gqa_c8_p8 panel.

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