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>

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