diff --git a/src/kernbench/benches/_attention_mesh_mlo_2d.py b/src/kernbench/benches/_attention_mesh_mlo_2d.py new file mode 100644 index 0000000..d4e1bba --- /dev/null +++ b/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) diff --git a/src/kernbench/benches/milestone_gqa_llama70b.py b/src/kernbench/benches/milestone_gqa_llama70b.py index b3a38cb..f7edcc4 100644 --- a/src/kernbench/benches/milestone_gqa_llama70b.py +++ b/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,13 +155,25 @@ 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. - ctx.launch( - f"{panel}_mesh", kernel, - q, k, v, o, - S_q, _S_KV_PER_RANK, _H_Q, _H_KV, _D_HEAD, n_ranks, - rank_axis, - _auto_dim_remap=False, - ) + if mesh_shape is None: + ctx.launch( + f"{panel}_mesh", kernel, + q, k, v, o, + S_q, _S_KV_PER_RANK, _H_Q, _H_KV, _D_HEAD, n_ranks, + 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, diff --git a/tests/attention/test_mesh_mlo_2d_correctness.py b/tests/attention/test_mesh_mlo_2d_correctness.py new file mode 100644 index 0000000..803717e --- /dev/null +++ b/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}" + )