ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Mukesh Garg
2026-06-04 12:39:48 -07:00
parent e2fe33180d
commit 4859149392
3 changed files with 399 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/benches/_attention_mesh_mlo_2d.py
+217
View File
@@ -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
+33 -9
View File
@@ -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,
tests/attention/test_mesh_mlo_2d_correctness.py
+142
View File
@@ -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}"
)