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>
This commit is contained in:
@@ -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)
|
||||
@@ -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,
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
Reference in New Issue
Block a user