attention: add 2D row-then-col AllReduce-mlo decode kernel (C2)

New ``_attention_mesh_mlo_2d.py`` decomposes a
``(mesh_rows x mesh_cols)`` cube sub-mesh into two stages of
bidirectional AllReduce-mlo:

  Stage 1 — row reduce (E/W edges, mesh_cols-1 steps)
  Stage 2 — col reduce (N/S edges, mesh_rows-1 steps)

After both stages every cube holds the same final ``(m, l, o)`` and
writes the normalized output. The online-softmax mlo merge is
associative, so row-then-col partitioning is mathematically
equivalent to a 1D ring AllReduce-mlo over all
``mesh_rows * mesh_cols`` cubes but uses fewer hops:

  - 2x4 (8 cubes / KV-group):  4 steps vs 7 (1.75x faster)
  - 4x4 (16 cubes / full SIP): 6 steps vs 15 (2.5x faster)

Motivation: the original 1D ring kernel ``_attention_mesh_mlo.py``
hit ``IpcqInvalidDirection`` at cube 4 when ``n_ranks=8`` on the
4x4 cube mesh — cube 4 has no W neighbor at the row 0/1 boundary.
N/S edges are already installed by ``configure_sfr_intercube_multisip``
so the 2D kernel runs on existing wiring without SFR changes.

The kernel accepts ``cube_start: int = 0`` and subtracts it from
``program_id(axis=1)`` so the ring math uses launch-local rank. This
matters because kernbench's ``program_id(axis=1)`` returns the
physical cube id (ADR-0022), so a launch starting at cube 8 would
otherwise compute ``my_row = 8//4 = 2`` (out of sub-mesh bounds) and
deadlock. Default ``cube_start=0`` keeps the existing
multi_user_decode validation behavior bit-for-bit.

Bench dispatch: ``multi_user_decode`` in milestone-gqa-llama70b now
uses the 2D kernel via a new ``mesh_shape`` column in
``_PANEL_DISPATCH``. At validation ``N_RANKS_MULTI_USER=4``, the
shape is ``(1, 4)`` — a degenerate single-row mesh, equivalent in
step count and op_log structure to the prior 1D ring at n_ranks=4.
The other three panels keep their 1D kernels.

Tests: 4 new unit tests in ``test_mesh_mlo_2d_correctness.py`` —
1x4 (degenerate row), 2x4 (8-KV-group target), 4x4 (full SIP), and
2x4 at cube_start=8 (the second sub-mesh per SIP). Existing
milestone (12 tests) and mesh-kernels-rank-axis (7 tests) suites
stay green — no regression.

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 → (kernel, SFR install, S_q, n_ranks, rank_axis, mesh_shape)
-_PANEL_DISPATCH: dict[str, tuple[Any, Any, int, int, int]] = {
+# ``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,
+        attention_mesh_mlo_2d_kernel, configure_sfr_intercube_multisip,
-        _S_Q_DECODE, _N_RANKS_MULTI_USER, 1,
+        _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,

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}"
+    )