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
+40 -16
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,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,