attention: add 8-KV-group diag harness (CCL pattern + cube_start sub-meshes)

New ``test_attention_8kv_groups_diag.py`` probes the path from
validation scale to the GQA Llama-70B 1-Q-head-per-cube headline
(64 cubes = 8 KV-groups x 8 cubes/group on a 4-SIP topology) in
three incremental steps:

  step_1_single_kv_group_at_full_breadth
    ONE 2x4 multi_user_decode launch (8 cubes) via the 2D mesh-mlo
    kernel. Verifies the per-KV-group 2D AllReduce works at full
    breadth.

  step_2_four_kv_groups_one_per_sip
    Four sequential 2x4 launches, one per SIP. Uses the CCL
    milestone's set_device pattern (milestone_1h_ccl.py:283-292):
    ONE run_bench with ``target_device="all"`` and
    ``ctx.ahbm.set_device(sip)`` between launches — not four
    separate run_bench calls with ``DeviceSelector("sip:N")``
    (which would misuse DeviceSelector and hit a
    ``DPPolicy x target_device`` allocator mismatch).

  step_3_eight_kv_groups_two_per_sip
    Two 2x4 launches per SIP x 4 SIPs = 8 KV-groups (64 cubes
    total). Second launch per SIP uses ``cube_start=8`` to address
    cubes 8..15 — the disjoint sub-mesh that was unreachable before
    ``DPPolicy.cube_start`` landed. Demonstrates the headline-enabling
    use of cube_start end-to-end (DPPolicy + kernel kwarg).

All three steps pass at validation dims (S_q=1, S_kv=16, h=1,
d_head=64). Headline-scale dims, multi_user_prefill 2D, and the B=8
"Batch on batch" PE parallelism remain follow-up work.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

src/kernbench/benches/_attention_mesh_mlo_2d.py

@@ -0,0 +1,217 @@
+"""Mesh-native 2D row-then-col AllReduce-mlo attention — decode (ADR-0059 extension).
+
+Each cube holds the full Q (replicated) and 1/(mesh_rows * mesh_cols) of
+KV (sequence-sharded across the 2D cube sub-mesh). The kernel decomposes
+the AllReduce-mlo into a two-stage reduction:
+
+  Stage 1 — Row reduce (E/W edges, ``mesh_cols - 1`` steps)
+     Bidirectional ring within each row. After this stage every cube in
+     row ``r`` holds the partial ``(m, ℓ, o)`` over the ``mesh_cols`` KV
+     chunks in row ``r``.
+
+  Stage 2 — Col reduce (N/S edges, ``mesh_rows - 1`` steps)
+     Bidirectional ring within each column. After this stage every cube
+     holds the partial over all ``mesh_rows × mesh_cols`` KV chunks —
+     the AllReduce result.
+
+The online-softmax mlo merge is associative, so row-then-col partitioning
+of the reduction is mathematically equivalent to a 1D ring AllReduce-mlo
+over all ``mesh_rows × mesh_cols`` cubes. The 2D form takes
+``(mesh_cols - 1) + (mesh_rows - 1)`` steps instead of
+``mesh_rows × mesh_cols - 1`` (e.g. 4 vs 7 at 2×4; 6 vs 15 at 4×4).
+
+Designed to run on hardware wired by
+``configure_sfr_intercube_multisip``, which installs both E/W and N/S
+intra-SIP cube-mesh edges (``sfr_config.py:135-143``). The 1D
+``_attention_mesh_mlo.py`` remains for the single_user PE-ring case;
+this 2D variant supersedes it for multi_user_decode where the per-KV-group
+cube count crosses a row boundary in the 4×4 cube mesh.
+
+``mesh_rows = 1`` is supported as a degenerate row-only case so the
+validation config (``_N_RANKS_MULTI_USER = 4`` → ``(1, 4)``) reduces to
+the 1D ring's step count without behavioral change.
+"""
+from __future__ import annotations
+
+from kernbench.common.pe_commands import TensorHandle
+
+
+def _view(handle: TensorHandle, new_shape: tuple[int, ...]) -> TensorHandle:
+    """Reshape — metadata only, no command emitted (cf. ``tl.trans``)."""
+    return TensorHandle(
+        id=handle.id,
+        addr=handle.addr,
+        shape=new_shape,
+        dtype=handle.dtype,
+        nbytes=handle.nbytes,
+        data=handle.data,
+        space=handle.space,
+        pinned=handle.pinned,
+    )
+
+
+def _bidir_allreduce_mlo(
+    m: TensorHandle,
+    ell: TensorHandle,
+    o: TensorHandle,
+    rank: int,
+    n_ranks: int,
+    dir_pos: str,
+    dir_neg: str,
+    *,
+    tl,
+) -> tuple[TensorHandle, TensorHandle, TensorHandle]:
+    """One bidirectional AllReduce-mlo ring along ``(dir_pos, dir_neg)``.
+
+    Mirrors the 1D ``_attention_mesh_mlo.py`` algorithm but parameterized
+    on direction labels so the 2D kernel can call it once with ``("E", "W")``
+    for the row reduce and once with ``("S", "N")`` for the col reduce.
+    Forwards the received triplets in subsequent steps so chunk ``c_i``
+    reaches rank ``j`` at step ``|i - j|``.
+
+    Returns the running ``(m, ℓ, o)`` after ``n_ranks - 1`` steps. Degenerate
+    cases (``n_ranks <= 1``) are no-ops — the for-loop body simply does not
+    execute.
+    """
+    has_pos = rank < n_ranks - 1
+    has_neg = rank > 0
+
+    to_send_pos_m: TensorHandle | None = m
+    to_send_pos_ell: TensorHandle | None = ell
+    to_send_pos_o: TensorHandle | None = o
+    to_send_neg_m: TensorHandle | None = m
+    to_send_neg_ell: TensorHandle | None = ell
+    to_send_neg_o: TensorHandle | None = o
+
+    for step in range(1, n_ranks):
+        if has_pos and to_send_pos_m is not None:
+            tl.send(dir=dir_pos, src=to_send_pos_m)
+            tl.send(dir=dir_pos, src=to_send_pos_ell)
+            tl.send(dir=dir_pos, src=to_send_pos_o)
+        if has_neg and to_send_neg_m is not None:
+            tl.send(dir=dir_neg, src=to_send_neg_m)
+            tl.send(dir=dir_neg, src=to_send_neg_ell)
+            tl.send(dir=dir_neg, src=to_send_neg_o)
+
+        m_from_neg: TensorHandle | None = None
+        ell_from_neg: TensorHandle | None = None
+        o_from_neg: TensorHandle | None = None
+        if has_neg and (rank - step) >= 0:
+            m_from_neg = tl.recv(dir=dir_neg, shape=m.shape, dtype="f16")
+            ell_from_neg = tl.recv(dir=dir_neg, shape=ell.shape, dtype="f16")
+            o_from_neg = tl.recv(dir=dir_neg, shape=o.shape, dtype="f16")
+            m_combined = tl.maximum(m, m_from_neg)
+            scale_old = tl.exp(m - m_combined)
+            scale_new = tl.exp(m_from_neg - m_combined)
+            ell = ell * scale_old + ell_from_neg * scale_new
+            o = o * scale_old + o_from_neg * scale_new
+            m = m_combined
+
+        m_from_pos: TensorHandle | None = None
+        ell_from_pos: TensorHandle | None = None
+        o_from_pos: TensorHandle | None = None
+        if has_pos and (rank + step) < n_ranks:
+            m_from_pos = tl.recv(dir=dir_pos, shape=m.shape, dtype="f16")
+            ell_from_pos = tl.recv(dir=dir_pos, shape=ell.shape, dtype="f16")
+            o_from_pos = tl.recv(dir=dir_pos, shape=o.shape, dtype="f16")
+            m_combined = tl.maximum(m, m_from_pos)
+            scale_old = tl.exp(m - m_combined)
+            scale_new = tl.exp(m_from_pos - m_combined)
+            ell = ell * scale_old + ell_from_pos * scale_new
+            o = o * scale_old + o_from_pos * scale_new
+            m = m_combined
+
+        to_send_pos_m = m_from_neg
+        to_send_pos_ell = ell_from_neg
+        to_send_pos_o = o_from_neg
+        to_send_neg_m = m_from_pos
+        to_send_neg_ell = ell_from_pos
+        to_send_neg_o = o_from_pos
+
+    return m, ell, o
+
+
+def attention_mesh_mlo_2d_kernel(
+    q_ptr: int,
+    k_ptr: int,
+    v_ptr: int,
+    o_ptr: int,
+    S_q: int,
+    S_kv_per_rank: int,
+    h_q: int,
+    h_kv: int,
+    d_head: int,
+    mesh_rows: int,
+    mesh_cols: int,
+    rank_axis: int = 0,
+    cube_start: int = 0,
+    *,
+    tl,
+) -> None:
+    """2D row-then-col AllReduce-mlo decode kernel — see module docstring.
+
+    ``rank_axis`` selects which program-id dimension carries the cube
+    rank (matches the 1D kernel convention):
+
+      0 — single_user_* (TL/BL): rank == tl.program_id(axis=0) (PE id).
+          Not used at headline scale — single_user uses the 1D intra-cube
+          PE ring (``_attention_mesh_mlo``). Kept here so the signature
+          mirrors the 1D kernel.
+
+      1 — multi_user_* (TR/BR): rank == tl.program_id(axis=1) (cube id).
+          KV is split @ cubes inter-cube; the ring runs over the
+          ``mesh_rows × mesh_cols`` cubes of one KV-group. The kernel
+          gates ``pe_id != 0`` to return early — same v1 simplification
+          as ``_attention_mesh_mlo`` (validation B=1).
+
+    ``cube_start`` matches the value passed to ``DPPolicy.cube_start`` for
+    the launch's tensor placement. kernbench's ``tl.program_id(axis=1)``
+    returns the physical cube id (ADR-0022), so when the launch is
+    offset within the SIP (e.g. cube_start=8 placing the second 2×4
+    KV-group on cubes 8..15), the kernel must subtract ``cube_start``
+    to recover the launch-local rank for ring arithmetic. Default 0
+    preserves the cube_start=0 launches unchanged.
+    """
+    # For multi_user (rank_axis=1) only PE 0 in each cube runs the ring.
+    if rank_axis != 0 and tl.program_id(axis=0) != 0:
+        return
+
+    rank = tl.program_id(axis=rank_axis)
+    if rank_axis != 0:
+        rank = rank - cube_start
+    my_row = rank // mesh_cols
+    my_col = rank % mesh_cols
+
+    # Q is replicated on every cube — loaded once.
+    Q = tl.load(q_ptr, shape=(S_q, h_q * d_head), dtype="f16")
+
+    # Local KV chunk (sequence-sharded across the 2D sub-mesh).
+    K = tl.load(k_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")
+    V = tl.load(v_ptr, shape=(S_kv_per_rank, h_kv, d_head), dtype="f16")
+
+    # ── One-shot local partial attention ──────────────────────────
+    K_2d_T = _view(K, (h_q * d_head, S_kv_per_rank))
+    V_2d = _view(V, (S_kv_per_rank, h_q * d_head))
+    scores = tl.dot(Q, K_2d_T)
+    m = tl.max(scores, axis=-1)
+    P = tl.softmax(scores, axis=-1)
+    scores_centered = scores - m
+    exp_scores = tl.exp(scores_centered)
+    ell = tl.sum(exp_scores, axis=-1)
+    o = tl.dot(P, V_2d)
+
+    # ── Stage 1: row AllReduce (E/W, mesh_cols - 1 steps) ─────────
+    m, ell, o = _bidir_allreduce_mlo(
+        m, ell, o, my_col, mesh_cols, "E", "W", tl=tl,
+    )
+
+    # ── Stage 2: col AllReduce (N/S, mesh_rows - 1 steps) ─────────
+    # ``dir_pos="S"`` matches the SFR convention: ``S`` goes to higher
+    # row (configure_sfr_intercube_multisip:140).
+    m, ell, o = _bidir_allreduce_mlo(
+        m, ell, o, my_row, mesh_rows, "S", "N", tl=tl,
+    )
+
+    # Final normalize: O := o / ℓ.
+    O_final = o / ell
+    tl.store(o_ptr, O_final)

src/kernbench/benches/milestone_gqa_llama70b.py

@@ -52,6 +52,7 @@ from typing import Any
 
 from kernbench.benches._attention_mesh_kv import attention_mesh_kv_kernel
 from kernbench.benches._attention_mesh_mlo import attention_mesh_mlo_kernel
+from kernbench.benches._attention_mesh_mlo_2d import attention_mesh_mlo_2d_kernel
 from kernbench.benches.registry import bench
 from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
 from kernbench.ccl.sfr_config import (
@@ -82,23 +83,32 @@ _PANELS_V1 = (
     "multi_user_decode",
 )
 
-# Panel → (kernel, SFR install, S_q, n_ranks, rank_axis)
-_PANEL_DISPATCH: dict[str, tuple[Any, Any, int, int, int]] = {
+# Panel → (kernel, SFR install, S_q, n_ranks, rank_axis, mesh_shape)
+# ``mesh_shape`` is ``None`` for 1D-ring kernels and ``(rows, cols)`` for the
+# 2D row-then-col AllReduce-mlo kernel (multi_user_decode); when set, the
+# launch passes ``(mesh_rows, mesh_cols)`` instead of ``n_ranks``.
+_PANEL_DISPATCH: dict[
+    str, tuple[Any, Any, int, int, int, tuple[int, int] | None]
+] = {
     "single_user_prefill": (
         attention_mesh_kv_kernel, configure_sfr_intracube_pe_ring,
-        _S_Q_PREFILL, _N_RANKS_SINGLE_USER, 0,
+        _S_Q_PREFILL, _N_RANKS_SINGLE_USER, 0, None,
     ),
     "multi_user_prefill": (
         attention_mesh_kv_kernel, configure_sfr_intercube_multisip,
-        _S_Q_PREFILL, _N_RANKS_MULTI_USER, 1,
+        _S_Q_PREFILL, _N_RANKS_MULTI_USER, 1, None,
     ),
     "single_user_decode": (
         attention_mesh_mlo_kernel, configure_sfr_intracube_pe_ring,
-        _S_Q_DECODE, _N_RANKS_SINGLE_USER, 0,
+        _S_Q_DECODE, _N_RANKS_SINGLE_USER, 0, None,
     ),
+    # multi_user_decode uses the C2 2D AllReduce-mlo kernel. (1, 4)
+    # degenerates to a row-only AllReduce equivalent to the prior 1D ring
+    # at n_ranks=4 — no op_log_summary regression. Headline 8-cube
+    # KV-groups land at (2, 4).
     "multi_user_decode": (
-        attention_mesh_mlo_kernel, configure_sfr_intercube_multisip,
-        _S_Q_DECODE, _N_RANKS_MULTI_USER, 1,
+        attention_mesh_mlo_2d_kernel, configure_sfr_intercube_multisip,
+        _S_Q_DECODE, _N_RANKS_MULTI_USER, 1, (1, _N_RANKS_MULTI_USER),
     ),
 }
 
@@ -107,7 +117,9 @@ _PANEL_DISPATCH: dict[str, tuple[Any, Any, int, int, int]] = {
 
 
 def _make_bench_fn(panel: str):
-    kernel, sfr_install, S_q, n_ranks, rank_axis = _PANEL_DISPATCH[panel]
+    kernel, sfr_install, S_q, n_ranks, rank_axis, mesh_shape = (
+        _PANEL_DISPATCH[panel]
+    )
     is_multi_user = panel.startswith("multi_user_")
 
     def _bench_fn(ctx):
@@ -143,6 +155,7 @@ def _make_bench_fn(panel: str):
                       dtype=_DTYPE, dp=dp_full, name=f"{panel}_o")
         # rank_axis is a positional arg; _auto_dim_remap=False keeps
         # d_head=64 from colliding with the multi_user K's global M=64.
+        if mesh_shape is None:
             ctx.launch(
                 f"{panel}_mesh", kernel,
                 q, k, v, o,
@@ -150,6 +163,17 @@ def _make_bench_fn(panel: str):
                 rank_axis,
                 _auto_dim_remap=False,
             )
+        else:
+            mesh_rows, mesh_cols = mesh_shape
+            ctx.launch(
+                f"{panel}_mesh", kernel,
+                q, k, v, o,
+                S_q, _S_KV_PER_RANK, _H_Q, _H_KV, _D_HEAD,
+                mesh_rows, mesh_cols,
+                rank_axis,
+                0,  # cube_start=0: this panel's launch starts at cube 0
+                _auto_dim_remap=False,
+            )
 
     return _bench_fn
 
@@ -210,7 +234,7 @@ def _run_panel(panel: str, topology: str) -> dict:
         raise RuntimeError(
             f"milestone-gqa-llama70b panel {panel!r} failed: {result.completion}"
         )
-    _, _, _, n_ranks, _ = _PANEL_DISPATCH[panel]
+    _, _, _, n_ranks, _, _ = _PANEL_DISPATCH[panel]
     return {
         "panel": panel,
         "n_ranks": n_ranks,

src/kernbench/policy/placement/dp.py

@@ -33,12 +33,20 @@ class DPPolicy:
     Optional overrides (``None`` = use topology dimensions):
       - num_pes: override PEs per cube
       - num_cubes: override cubes per SIP
+
+    Cube range:
+      - cube_start: index of the first cube in the SIP to place shards
+        on. Defaults to 0. Enables disjoint sub-meshes within one SIP
+        (e.g. cubes 8..15 alongside cubes 0..7) — required by the
+        GQA Llama-70B 8-KV-group headline target (two 2×4 KV-groups
+        per SIP).
     """
 
     cube: Literal["replicate", "column_wise", "row_wise"] = "replicate"
     pe: Literal["replicate", "column_wise", "row_wise"] = "replicate"
     num_pes: int | None = None
     num_cubes: int | None = None
+    cube_start: int = 0
 
 
 @dataclass(frozen=True)
@@ -125,7 +133,7 @@ def resolve_dp_policy(
         for ls in local_shards:
             all_shards.append(ShardSpec(
                 sip=target_sip,
-                cube=cube_id,
+                cube=policy.cube_start + cube_id,
                 pe=ls.local_pe,
                 offset_bytes=cube_offset + ls.offset_bytes,
                 nbytes=ls.nbytes,

tests/attention/test_attention_8kv_groups_diag.py

@@ -0,0 +1,286 @@
+"""Diagnostic harness for the Llama-70B "1 Q-head per cube" target.
+
+Per the GQA Llama-70B sharding study at
+``llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py``,
+the 1 Q-head/cube baseline uses 64 cubes (4 SIPs × 16 cubes/SIP) organized
+into 8 KV-groups of 8 cubes each. Each KV-group occupies a ``2×4``
+sub-mesh within a SIP's ``4×4`` cube grid and runs the C2 2D row-then-col
+AllReduce-mlo (ADR-0059 extension). This harness probes the gap between
+validation and headline in three incrementally-larger steps:
+
+  step_1_single_kv_group_at_full_breadth
+     ONE multi_user_decode launch on a 2×4 sub-mesh (8 cubes) of the
+     4-SIP topology, via the 2D mesh-mlo kernel. Verifies the per-KV-group
+     2D AllReduce works at full breadth. Smallest dim possible
+     (S_q=1, S_kv=16, h=1, d_head=64) to keep wall time bounded.
+
+  step_2_four_kv_groups_one_per_sip
+     Four sequential multi_user_decode launches, each targeting a
+     different SIP. Verifies that per-SIP isolation works (each SIP holds
+     its own 2×4 KV-group; the SFR install only writes intra-SIP
+     E/W + N/S edges so the 4 groups don't see each other).
+
+  step_3_eight_kv_groups_two_per_sip
+     The actual study target: 8 KV-groups, two per SIP (cubes 0..7 vs
+     cubes 8..15 within each SIP). Expected to FAIL with current infra —
+     DPPolicy doesn't take a cube offset and target_device is SIP-level,
+     so back-to-back launches both land on cubes 0..7 of their target SIP.
+     Captures what's needed to lift the 4-group cap to 8.
+
+Each step prints what it observed; the test asserts only the documented
+expected outcomes so we can land it, watch CI, and iterate. Steps that
+are *expected* to fail (step 3) are marked xfail with a precise reason.
+"""
+from __future__ import annotations
+
+import traceback
+from pathlib import Path
+
+import pytest
+
+from kernbench.benches._attention_mesh_mlo_2d import attention_mesh_mlo_2d_kernel
+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 DeviceSelector, resolve_device
+from kernbench.sim_engine.engine import GraphEngine
+from kernbench.topology.builder import resolve_topology
+
+TOPOLOGY_4SIP = (
+    Path(__file__).resolve().parents[2] / "topologies" / "llama70b_4sip.yaml"
+)
+TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
+
+S_Q_DECODE = 1
+S_KV_PER_RANK = 16
+H_Q = 1
+H_KV = 1
+D_HEAD = 64
+# 2×4 sub-mesh per KV-group (study: 8 cubes per KV-group at Q/cube=1).
+MESH_ROWS = 2
+MESH_COLS = 4
+N_CUBES_PER_KV_GROUP = MESH_ROWS * MESH_COLS
+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 _make_one_kv_group_bench(mesh_rows: int, mesh_cols: int):
+    """Return a bench_fn that runs ONE multi_user_decode kernel on a
+    ``mesh_rows × mesh_cols`` sub-mesh."""
+    n_cubes = mesh_rows * mesh_cols
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp_full = DPPolicy(cube="replicate", pe="replicate",
+                           num_cubes=n_cubes, num_pes=8)
+        dp_kv = DPPolicy(cube="row_wise", pe="replicate",
+                         num_cubes=n_cubes, num_pes=8)
+        q = ctx.zeros((S_Q_DECODE, H_Q * D_HEAD),
+                      dtype=DTYPE, dp=dp_full, name="q")
+        k = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name="k")
+        v = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name="v")
+        o = ctx.empty((S_Q_DECODE, H_Q * D_HEAD),
+                      dtype=DTYPE, dp=dp_full, name="o")
+        ctx.launch(
+            f"single_kv_group_{mesh_rows}x{mesh_cols}",
+            attention_mesh_mlo_2d_kernel,
+            q, k, v, o,
+            S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD,
+            mesh_rows, mesh_cols,
+            1,                               # rank_axis=1 → cube-level ring
+            0,                               # cube_start=0 — single sub-mesh launch
+            _auto_dim_remap=False,
+        )
+
+    return _bench_fn
+
+
+def _run_one_kv_group(topology_path: Path, mesh_rows: int, mesh_cols: int,
+                      target_device=None):
+    topo = resolve_topology(str(topology_path))
+    captured: dict = {"engine": None}
+
+    def factory(t, d):
+        eng = _engine_factory(t, d)
+        captured["engine"] = eng
+        return eng
+
+    exc = None
+    result = None
+    try:
+        result = run_bench(
+            topology=topo,
+            bench_fn=_make_one_kv_group_bench(mesh_rows, mesh_cols),
+            device=target_device or resolve_device(None),
+            engine_factory=factory,
+        )
+    except BaseException as e:  # noqa: BLE001
+        exc = e
+    return exc, result, captured["engine"]
+
+
+# ── Step 1 — single KV-group at the study's full breadth ──────────
+
+
+def test_step_1_single_kv_group_at_full_breadth():
+    """One multi_user_decode launch on a 2×4 sub-mesh, 4-SIP topology.
+
+    Uses the C2 2D row-then-col AllReduce-mlo kernel: stage 1 reduces
+    across cols (E/W) within each row, stage 2 reduces across rows (N/S).
+    Expected to PASS — N/S edges are wired by
+    ``configure_sfr_intercube_multisip`` and the 2D fan-out avoids the
+    row-boundary IpcqInvalidDirection that the 1D kernel hit at cube 4.
+    """
+    if not TOPOLOGY_4SIP.exists():
+        pytest.skip(f"4-SIP topology missing: {TOPOLOGY_4SIP}")
+    exc, result, engine = _run_one_kv_group(
+        TOPOLOGY_4SIP, mesh_rows=MESH_ROWS, mesh_cols=MESH_COLS,
+    )
+    if exc is not None:
+        oplog_len = len(getattr(engine, "op_log", []) or []) if engine else 0
+        print(f"\nstep_1 FAIL — op_log records before crash: {oplog_len}")
+        traceback.print_exception(type(exc), exc, exc.__traceback__)
+        raise AssertionError(f"step_1 failed: {exc}") from exc
+    assert result is not None and result.completion.ok, (
+        f"step_1: completion not ok — {result.completion if result else None}"
+    )
+
+
+# ── Step 2 — 4 KV-groups, one per SIP, sequential launches ────────
+
+
+def _make_multi_sip_bench_fn(sip_groups: list[tuple[int, str, int]]):
+    """One bench_fn that does one 2×4 multi_user_decode launch per item.
+
+    Each ``(sip, tag, cube_start)`` tuple becomes one launch:
+      - ``ctx.ahbm.set_device(sip)`` switches allocations to that SIP
+        (mirrors ``milestone_1h_ccl.py:283-292``).
+      - ``cube_start`` selects which 8-cube sub-mesh within the SIP:
+        ``0`` → cubes 0..7 (rows 0..1), ``8`` → cubes 8..15 (rows 2..3).
+      - ``tag`` disambiguates tensor names so launches in the same
+        run_bench don't collide on the allocator namespace.
+    """
+    n_cubes = MESH_ROWS * MESH_COLS
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        for sip, tag, cube_start in sip_groups:
+            ctx.ahbm.set_device(sip)
+            dp_full = DPPolicy(cube="replicate", pe="replicate",
+                               num_cubes=n_cubes, num_pes=8,
+                               cube_start=cube_start)
+            dp_kv = DPPolicy(cube="row_wise", pe="replicate",
+                             num_cubes=n_cubes, num_pes=8,
+                             cube_start=cube_start)
+            q = ctx.zeros((S_Q_DECODE, H_Q * D_HEAD),
+                          dtype=DTYPE, dp=dp_full, name=f"q_{tag}")
+            k = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
+                          dtype=DTYPE, dp=dp_kv, name=f"k_{tag}")
+            v = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
+                          dtype=DTYPE, dp=dp_kv, name=f"v_{tag}")
+            o = ctx.empty((S_Q_DECODE, H_Q * D_HEAD),
+                          dtype=DTYPE, dp=dp_full, name=f"o_{tag}")
+            ctx.launch(
+                f"kv_group_{tag}", attention_mesh_mlo_2d_kernel,
+                q, k, v, o,
+                S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD,
+                MESH_ROWS, MESH_COLS,
+                1,                               # rank_axis=1 → cube-level ring
+                cube_start,                      # converts physical id → launch-local rank
+                _auto_dim_remap=False,
+            )
+
+    return _bench_fn
+
+
+def _run_multi_sip(sip_groups: list[tuple[int, str, int]]):
+    """Run a single run_bench call covering all (sip, tag) groups."""
+    topo = resolve_topology(str(TOPOLOGY_4SIP))
+    captured: dict = {"engine": None}
+
+    def factory(t, d):
+        eng = _engine_factory(t, d)
+        captured["engine"] = eng
+        return eng
+
+    exc = None
+    result = None
+    try:
+        result = run_bench(
+            topology=topo,
+            bench_fn=_make_multi_sip_bench_fn(sip_groups),
+            device=resolve_device(None),  # "all" SIPs in scope
+            engine_factory=factory,
+        )
+    except BaseException as e:  # noqa: BLE001
+        exc = e
+    return exc, result, captured["engine"]
+
+
+def test_step_2_four_kv_groups_one_per_sip():
+    """Four multi_user_decode launches, one per SIP, in ONE run_bench call.
+
+    Uses the CCL milestone pattern (``milestone_1h_ccl.py:283-292``):
+    ``target_device="all"`` scopes the runtime to every SIP; then
+    ``ctx.ahbm.set_device(sip)`` before each ``ctx.zeros``/``launch``
+    switches which SIP the next allocation+launch lands on. This is the
+    canonical sequential per-SIP pattern in the codebase — four separate
+    ``run_bench`` calls with ``DeviceSelector("sip:N")`` is a misuse.
+
+    Expected to PASS — the SFR install draws intra-SIP edges only, so the
+    4 KV-groups can't see each other.
+    """
+    if not TOPOLOGY_4SIP.exists():
+        pytest.skip(f"4-SIP topology missing: {TOPOLOGY_4SIP}")
+    sip_groups = [(sip, f"sip{sip}", 0) for sip in range(4)]
+    exc, result, engine = _run_multi_sip(sip_groups)
+    if exc is not None:
+        oplog_len = len(getattr(engine, "op_log", []) or []) if engine else 0
+        print(f"\nstep_2 FAIL — op_log records before crash: {oplog_len}")
+        traceback.print_exception(type(exc), exc, exc.__traceback__)
+        raise AssertionError(f"step_2 failed: {exc}") from exc
+    assert result is not None and result.completion.ok, (
+        f"step_2: completion not ok — {result.completion if result else None}"
+    )
+
+
+# ── Step 3 — 8 KV-groups, two per SIP (study target) ──────────────
+
+
+def test_step_3_eight_kv_groups_two_per_sip():
+    """Two launches per SIP × 4 SIPs = 8 KV-groups total in one run_bench.
+
+    The headline target: 64 cubes serving 8 KV-groups, two disjoint 2×4
+    sub-meshes per SIP. ``cube_start=0`` puts the first KV-group on
+    cubes 0..7 (rows 0..1); ``cube_start=8`` puts the second on cubes
+    8..15 (rows 2..3). This is the use case ``DPPolicy.cube_start`` was
+    added to enable.
+
+    Expected to PASS — the SFR install draws intra-SIP edges only, so
+    the 8 KV-groups can't see each other; ``cube_start`` ensures the
+    two halves of each SIP land on disjoint cubes.
+    """
+    if not TOPOLOGY_4SIP.exists():
+        pytest.skip(f"4-SIP topology missing: {TOPOLOGY_4SIP}")
+    sip_groups = [
+        (sip, f"sip{sip}_half{half}", half * N_CUBES_PER_KV_GROUP)
+        for sip in range(4) for half in (0, 1)
+    ]
+    exc, result, engine = _run_multi_sip(sip_groups)
+    if exc is not None:
+        raise AssertionError(f"step_3 failed: {exc}") from exc
+    assert result is not None and result.completion.ok, (
+        f"step_3: completion not ok — {result.completion if result else None}"
+    )

tests/attention/test_mesh_mlo_2d_correctness.py

@@ -0,0 +1,142 @@
+"""Phase 1 spec test for the 2D row-then-col AllReduce-mlo decode kernel.
+
+The 2D kernel decomposes a ``(mesh_rows × mesh_cols)`` cube sub-mesh into a
+two-stage AllReduce-mlo: stage 1 reduces across columns within each row
+(E/W edges), stage 2 reduces across rows within each column (N/S edges).
+After both stages every cube holds the same final ``(m, ℓ, o)``.
+
+This module is Phase 1 of C2 (see CLAUDE.md change protocol): it pins
+the kernel's interface and observable behavior. Production code for the
+kernel lands in Phase 2; until then this file fails to import.
+
+Test shapes (run on default ``topology.yaml`` — 2 SIPs × 4×4 cube_mesh):
+
+  1×4 sub-mesh (4 cubes, row 0 only)
+     Degenerates to a row-only AllReduce — equivalent in step count to
+     the existing 1D kernel at n_ranks=4. Verifies the kernel reduces
+     correctly when mesh_rows=1 (stage 2 collapses to no-op).
+
+  2×4 sub-mesh (8 cubes, rows 0+1)
+     The 8-KV-group target. Verifies that cubes 4..7 use ``dir="N"``
+     (not ``dir="W"``) to reach row 0 — surfacing the IpcqInvalidDirection
+     bug that the 1D kernel hit at cube 4 (rank 4, no W neighbor).
+
+  4×4 sub-mesh (16 cubes, full SIP)
+     Full-SIP scale. Verifies the algorithm fans out over (cols-1)=3
+     row steps followed by (rows-1)=3 col steps.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+from kernbench.benches._attention_mesh_mlo_2d import attention_mesh_mlo_2d_kernel
+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"
+
+S_Q = 1
+S_KV_PER_RANK = 16
+H_Q = 1
+H_KV = 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_2d(mesh_rows: int, mesh_cols: int, cube_start: int = 0):
+    """Build a bench_fn and run it on the default topology."""
+    n_cubes = mesh_rows * mesh_cols
+    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=n_cubes, num_pes=8,
+                           cube_start=cube_start)
+        dp_kv = DPPolicy(cube="row_wise", pe="replicate",
+                         num_cubes=n_cubes, num_pes=8,
+                         cube_start=cube_start)
+        q = ctx.zeros((S_Q, H_Q * D_HEAD),
+                      dtype=DTYPE, dp=dp_full, name="q")
+        k = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name="k")
+        v = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name="v")
+        o = ctx.empty((S_Q, H_Q * D_HEAD),
+                      dtype=DTYPE, dp=dp_full, name="o")
+        ctx.launch(
+            f"mesh_mlo_2d_{mesh_rows}x{mesh_cols}_start{cube_start}",
+            attention_mesh_mlo_2d_kernel,
+            q, k, v, o,
+            S_Q, S_KV_PER_RANK, H_Q, H_KV, D_HEAD,
+            mesh_rows, mesh_cols,
+            1,  # rank_axis=1 → cube-level ring
+            cube_start,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo,
+        bench_fn=_bench_fn,
+        device=resolve_device(None),
+        engine_factory=_engine_factory,
+    )
+
+
+def test_2d_kernel_1x4_row_only():
+    """1×4 sub-mesh: row-only AllReduce, stage 2 collapses to no-op."""
+    result = _run_2d(mesh_rows=1, mesh_cols=4)
+    assert result.completion.ok, (
+        f"1x4: completion not ok - {result.completion}"
+    )
+
+
+def test_2d_kernel_2x4_eight_cubes():
+    """2×4 sub-mesh: the 8-KV-group target.
+
+    Verifies that cube 4 (row 1, col 0) uses ``dir="N"`` to reach cube 0
+    (row 0, col 0) for stage 2, not ``dir="W"`` — the 1D kernel hit
+    IpcqInvalidDirection here.
+    """
+    result = _run_2d(mesh_rows=2, mesh_cols=4)
+    assert result.completion.ok, (
+        f"2x4: completion not ok - {result.completion}"
+    )
+
+
+def test_2d_kernel_4x4_full_sip():
+    """4×4 sub-mesh: full-SIP scale (16 cubes)."""
+    result = _run_2d(mesh_rows=4, mesh_cols=4)
+    assert result.completion.ok, (
+        f"4x4: completion not ok - {result.completion}"
+    )
+
+
+def test_2d_kernel_2x4_at_cube_start_eight():
+    """2×4 sub-mesh at cube_start=8 (cubes 8..15, rows 2..3).
+
+    The second KV-group per SIP in the 8-KV-group Llama-70B headline.
+    Verifies the kernel converts ``program_id(axis=1)`` (physical cube
+    id) back to launch-local rank via ``cube_start`` — without that
+    subtraction, cube 8 would compute my_row=2 (out of sub-mesh bounds)
+    and deadlock waiting on cube 4 which isn't in the launch.
+    """
+    result = _run_2d(mesh_rows=2, mesh_cols=4, cube_start=8)
+    assert result.completion.ok, (
+        f"2x4 @ cube_start=8: completion not ok - {result.completion}"
+    )

tests/test_dppolicy_cube_start.py

@@ -0,0 +1,126 @@
+"""Phase 1 spec test for ``DPPolicy.cube_start``.
+
+``cube_start`` is an optional override that shifts the cube indices in
+``ShardSpec`` generated by ``resolve_dp_policy``. It enables addressing a
+disjoint cube sub-mesh within one SIP (e.g. cubes 8..15 in addition to
+0..7), which the GQA Llama-70B 8-KV-group headline target requires
+(two 2×4 KV-groups per SIP × 4 SIPs = 64 cubes).
+
+Default ``cube_start=0`` preserves every existing call-site bit-for-bit:
+``DPPolicy(num_cubes=8)`` still resolves to ``ShardSpec.cube ∈ [0, 8)``.
+The CCL milestone bench's full-SIP ``DPPolicy(num_cubes=16)`` likewise
+continues to produce cubes 0..15.
+
+Phase 1: production code is unchanged → these tests FAIL until the
+Phase 2 diff lands. Phase 2 makes all of them pass.
+"""
+from __future__ import annotations
+
+from kernbench.policy.placement.dp import DPPolicy, resolve_dp_policy
+
+
+# ── Default ``cube_start=0`` preserves existing behavior ─────────────
+
+
+def test_dppolicy_default_cube_start_is_zero():
+    """``DPPolicy(num_cubes=8)`` sets ``cube_start=0`` by default."""
+    dp = DPPolicy(cube="row_wise", num_cubes=8)
+    assert dp.cube_start == 0
+
+
+def test_resolve_default_yields_cubes_zero_through_seven():
+    """Backward-compat: default ``cube_start=0`` → cubes 0..num_cubes-1."""
+    dp = DPPolicy(cube="row_wise", pe="replicate", num_cubes=8, num_pes=1)
+    shards = resolve_dp_policy(
+        dp, shape=(32, 64), itemsize=2,
+        num_pe=1, num_cubes=8, target_sip=0,
+    )
+    cube_ids = sorted({s.cube for s in shards})
+    assert cube_ids == list(range(8))
+
+
+def test_resolve_full_sip_default_yields_cubes_zero_through_fifteen():
+    """CCL pattern: ``num_cubes=16, cube_start=0`` → cubes 0..15 unchanged."""
+    dp = DPPolicy(cube="row_wise", pe="replicate", num_cubes=16, num_pes=1)
+    shards = resolve_dp_policy(
+        dp, shape=(32, 64), itemsize=2,
+        num_pe=1, num_cubes=16, target_sip=0,
+    )
+    cube_ids = sorted({s.cube for s in shards})
+    assert cube_ids == list(range(16))
+
+
+# ── ``cube_start=N`` shifts cube indices ─────────────────────────────
+
+
+def test_resolve_cube_start_eight_yields_cubes_eight_through_fifteen():
+    """``cube_start=8, num_cubes=8`` → ``ShardSpec.cube ∈ [8, 16)``.
+
+    This is the headline-enabling case: the second 2×4 KV-group per SIP
+    in the 8-KV-group Llama-70B target lands on cubes 8..15.
+    """
+    dp = DPPolicy(
+        cube="row_wise", pe="replicate",
+        num_cubes=8, num_pes=1, cube_start=8,
+    )
+    shards = resolve_dp_policy(
+        dp, shape=(32, 64), itemsize=2,
+        num_pe=1, num_cubes=8, target_sip=0,
+    )
+    cube_ids = sorted({s.cube for s in shards})
+    assert cube_ids == list(range(8, 16))
+
+
+def test_resolve_cube_start_preserves_shard_count():
+    """Shifting ``cube_start`` does not change the total shard count."""
+    dp_default = DPPolicy(
+        cube="row_wise", pe="replicate", num_cubes=8, num_pes=1,
+    )
+    dp_shifted = DPPolicy(
+        cube="row_wise", pe="replicate",
+        num_cubes=8, num_pes=1, cube_start=8,
+    )
+    shards_def = resolve_dp_policy(
+        dp_default, shape=(32, 64), itemsize=2,
+        num_pe=1, num_cubes=8, target_sip=0,
+    )
+    shards_shf = resolve_dp_policy(
+        dp_shifted, shape=(32, 64), itemsize=2,
+        num_pe=1, num_cubes=8, target_sip=0,
+    )
+    assert len(shards_def) == len(shards_shf)
+
+
+# ── ``cube_start`` works with ``cube="replicate"`` too ──────────────
+
+
+def test_resolve_cube_start_with_replicate_policy():
+    """``cube="replicate", cube_start=8`` → every cube in [8, 16) gets a copy."""
+    dp = DPPolicy(
+        cube="replicate", pe="replicate",
+        num_cubes=8, num_pes=1, cube_start=8,
+    )
+    shards = resolve_dp_policy(
+        dp, shape=(32, 64), itemsize=2,
+        num_pe=1, num_cubes=8, target_sip=0,
+    )
+    cube_ids = sorted({s.cube for s in shards})
+    assert cube_ids == list(range(8, 16))
+
+
+# ── No duplicates regardless of ``cube_start`` ───────────────────────
+
+
+def test_resolve_cube_start_yields_distinct_cube_ids():
+    """For ``pe="replicate", num_pes=1``, every shard has a distinct cube."""
+    dp = DPPolicy(
+        cube="row_wise", pe="replicate",
+        num_cubes=8, num_pes=1, cube_start=8,
+    )
+    shards = resolve_dp_policy(
+        dp, shape=(32, 64), itemsize=2,
+        num_pe=1, num_cubes=8, target_sip=0,
+    )
+    cube_ids = [s.cube for s in shards]
+    assert len(cube_ids) == len(set(cube_ids))
+    assert all(8 <= c < 16 for c in cube_ids)