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

gqa: ADR-0060/0062/0063/0064 unified GQA kernels + CPU cost model

Browse Source 
Land the new GQA fused-attention kernels (ADR-0060) for prefill/decode
across long and short context, the TL discipline primitives they depend
on (ADR-0062 lazy load, ADR-0063 scratch_scope + copy_to), and the
per-op-type CPU issue cost model (ADR-0064). Remove the pre-ADR-0060
mesh-attention baseline now that the unified kernels supersede it.

ADR-0060 (long context)
- _gqa_decode.py: M-fold + 2-level chain reduce-to-root (Level-2
  intra-CUBE row-then-col + Level-1 inter-CUBE) — root-only output.
- _gqa_prefill.py: head-parallel + Ring KV rotation around C CUBEs,
  online-softmax merge per ring step, per-CUBE distributed output.
- Each merge stage wraps in scratch_scope() and persists running
  (m, l, O) via copy_to() to lift the 1 MiB scratch ceiling.

ADR-0060 §B.split.2 (short context, kv_per_cube in {1,2,4,8})
- _gqa_decode_short.py / _gqa_prefill_short.py: no cube-SP; each CUBE
  owns whole KV heads; PE-parallel heads with intra-group chain
  reduce. Prefill has no Ring KV (each head fully resident).

ADR-0062 (lazy tl.load): future-bearing TensorHandle, auto-wait at
first consuming op (dot/MATH/store/send/copy_to/composite).

ADR-0063 (tl.scratch_scope + tl.copy_to): scoped per-tile arena with
copy_to writeback primitive for persistent running state.

ADR-0064 (CPU issue cost model)
- common/cpu_issue_cost.py: per-op-type table (composite=40 ns,
  primitives=5 ns); ratios are load-bearing per D1.
- TLContext: issue_cost_table param; _emit_dispatch_overhead(kind)
  consults table with dispatch_cycles fallback (ADR-0046 §D6
  back-compat).
- Live PE_CPU paths (greenlet + legacy) construct TLContext with
  DEFAULT_CPU_ISSUE_COST so saturation lever (ADR-0060 §1) is
  measurable end-to-end.

P7 headline bench: milestone-gqa-headline writes per-panel
op_log_summary to 1H_milestone_output/gqa_headline/sweep.json. No
figure renderers yet (deferred).

Removals (pre-ADR-0060 baseline now superseded):
- benches: _attention_mesh_kv.py, _attention_mesh_mlo.py,
  _attention_mesh_mlo_2d.py, milestone_gqa_llama70b.py
- tests: test_attention_*, test_mesh_*, test_milestone_gqa_llama70b
- topology: llama70b_4sip.yaml (only consumer was the deleted diag)
- artifacts: 1H_milestone_output/gqa/ (sweep.json + 5 PNGs)
- tests/gqa/ plot helper + test (broken on Windows Tcl/Tkinter)
- ADR-0060/0061 references to deleted file paths cleaned up
  (EN + KO kept in sync).

Tests: 124/124 focused regression green (attention + Phase E + TL
discipline + triton_emu + pe_components). Full regression: 764 pass,
2 pre-existing test_bench_registry failures (stale EXPECTED_NAMES
across multiple benches, not introduced here).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-06-09 18:15:59 -07:00
parent b3730a33eb
commit d282144339
52 changed files with 3952 additions and 2526 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/attention/test_attention_8kv_groups_diag.py
-286
View File
@@ -1,286 +0,0 @@
"""Diagnostic harness for the Llama-70B "1 Q-head per cube" target.
Per the GQA Llama-70B sharding study at
``llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py``,
the 1 Q-head/cube baseline uses 64 cubes (4 SIPs × 16 cubes/SIP) organized
into 8 KV-groups of 8 cubes each. Each KV-group occupies a ``2×4``
sub-mesh within a SIP's ``4×4`` cube grid and runs the C2 2D row-then-col
AllReduce-mlo (ADR-0059 extension). This harness probes the gap between
validation and headline in three incrementally-larger steps:
step_1_single_kv_group_at_full_breadth
ONE multi_user_decode launch on a 2×4 sub-mesh (8 cubes) of the
4-SIP topology, via the 2D mesh-mlo kernel. Verifies the per-KV-group
2D AllReduce works at full breadth. Smallest dim possible
(S_q=1, S_kv=16, h=1, d_head=64) to keep wall time bounded.
step_2_four_kv_groups_one_per_sip
Four sequential multi_user_decode launches, each targeting a
different SIP. Verifies that per-SIP isolation works (each SIP holds
its own 2×4 KV-group; the SFR install only writes intra-SIP
E/W + N/S edges so the 4 groups don't see each other).
step_3_eight_kv_groups_two_per_sip
The actual study target: 8 KV-groups, two per SIP (cubes 0..7 vs
cubes 8..15 within each SIP). Expected to FAIL with current infra —
DPPolicy doesn't take a cube offset and target_device is SIP-level,
so back-to-back launches both land on cubes 0..7 of their target SIP.
Captures what's needed to lift the 4-group cap to 8.
Each step prints what it observed; the test asserts only the documented
expected outcomes so we can land it, watch CI, and iterate. Steps that
are *expected* to fail (step 3) are marked xfail with a precise reason.
"""
from __future__ import annotations
import traceback
from pathlib import Path
import pytest
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 DeviceSelector, resolve_device
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import resolve_topology
TOPOLOGY_4SIP = (
Path(__file__).resolve().parents[2] / "topologies" / "llama70b_4sip.yaml"
)
TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
S_Q_DECODE = 1
S_KV_PER_RANK = 16
H_Q = 1
H_KV = 1
D_HEAD = 64
# 2×4 sub-mesh per KV-group (study: 8 cubes per KV-group at Q/cube=1).
MESH_ROWS = 2
MESH_COLS = 4
N_CUBES_PER_KV_GROUP = MESH_ROWS * MESH_COLS
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 _make_one_kv_group_bench(mesh_rows: int, mesh_cols: int):
"""Return a bench_fn that runs ONE multi_user_decode kernel on a
``mesh_rows × mesh_cols`` sub-mesh."""
n_cubes = mesh_rows * mesh_cols
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)
dp_kv = DPPolicy(cube="row_wise", pe="replicate",
num_cubes=n_cubes, num_pes=8)
q = ctx.zeros((S_Q_DECODE, 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_DECODE, H_Q * D_HEAD),
dtype=DTYPE, dp=dp_full, name="o")
ctx.launch(
f"single_kv_group_{mesh_rows}x{mesh_cols}",
attention_mesh_mlo_2d_kernel,
q, k, v, o,
S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD,
mesh_rows, mesh_cols,
1, # rank_axis=1 → cube-level ring
0, # cube_start=0 — single sub-mesh launch
_auto_dim_remap=False,
)
return _bench_fn
def _run_one_kv_group(topology_path: Path, mesh_rows: int, mesh_cols: int,
target_device=None):
topo = resolve_topology(str(topology_path))
captured: dict = {"engine": None}
def factory(t, d):
eng = _engine_factory(t, d)
captured["engine"] = eng
return eng
exc = None
result = None
try:
result = run_bench(
topology=topo,
bench_fn=_make_one_kv_group_bench(mesh_rows, mesh_cols),
device=target_device or resolve_device(None),
engine_factory=factory,
)
except BaseException as e: # noqa: BLE001
exc = e
return exc, result, captured["engine"]
# ── Step 1 — single KV-group at the study's full breadth ──────────
def test_step_1_single_kv_group_at_full_breadth():
"""One multi_user_decode launch on a 2×4 sub-mesh, 4-SIP topology.
Uses the C2 2D row-then-col AllReduce-mlo kernel: stage 1 reduces
across cols (E/W) within each row, stage 2 reduces across rows (N/S).
Expected to PASS — N/S edges are wired by
``configure_sfr_intercube_multisip`` and the 2D fan-out avoids the
row-boundary IpcqInvalidDirection that the 1D kernel hit at cube 4.
"""
if not TOPOLOGY_4SIP.exists():
pytest.skip(f"4-SIP topology missing: {TOPOLOGY_4SIP}")
exc, result, engine = _run_one_kv_group(
TOPOLOGY_4SIP, mesh_rows=MESH_ROWS, mesh_cols=MESH_COLS,
)
if exc is not None:
oplog_len = len(getattr(engine, "op_log", []) or []) if engine else 0
print(f"\nstep_1 FAIL — op_log records before crash: {oplog_len}")
traceback.print_exception(type(exc), exc, exc.__traceback__)
raise AssertionError(f"step_1 failed: {exc}") from exc
assert result is not None and result.completion.ok, (
f"step_1: completion not ok — {result.completion if result else None}"
)
# ── Step 2 — 4 KV-groups, one per SIP, sequential launches ────────
def _make_multi_sip_bench_fn(sip_groups: list[tuple[int, str, int]]):
"""One bench_fn that does one 2×4 multi_user_decode launch per item.
Each ``(sip, tag, cube_start)`` tuple becomes one launch:
- ``ctx.ahbm.set_device(sip)`` switches allocations to that SIP
(mirrors ``milestone_1h_ccl.py:283-292``).
- ``cube_start`` selects which 8-cube sub-mesh within the SIP:
``0`` → cubes 0..7 (rows 0..1), ``8`` → cubes 8..15 (rows 2..3).
- ``tag`` disambiguates tensor names so launches in the same
run_bench don't collide on the allocator namespace.
"""
n_cubes = MESH_ROWS * MESH_COLS
def _bench_fn(ctx):
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
for sip, tag, cube_start in sip_groups:
ctx.ahbm.set_device(sip)
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_DECODE, H_Q * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"q_{tag}")
k = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"k_{tag}")
v = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"v_{tag}")
o = ctx.empty((S_Q_DECODE, H_Q * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"o_{tag}")
ctx.launch(
f"kv_group_{tag}", attention_mesh_mlo_2d_kernel,
q, k, v, o,
S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD,
MESH_ROWS, MESH_COLS,
1, # rank_axis=1 → cube-level ring
cube_start, # converts physical id → launch-local rank
_auto_dim_remap=False,
)
return _bench_fn
def _run_multi_sip(sip_groups: list[tuple[int, str, int]]):
"""Run a single run_bench call covering all (sip, tag) groups."""
topo = resolve_topology(str(TOPOLOGY_4SIP))
captured: dict = {"engine": None}
def factory(t, d):
eng = _engine_factory(t, d)
captured["engine"] = eng
return eng
exc = None
result = None
try:
result = run_bench(
topology=topo,
bench_fn=_make_multi_sip_bench_fn(sip_groups),
device=resolve_device(None), # "all" SIPs in scope
engine_factory=factory,
)
except BaseException as e: # noqa: BLE001
exc = e
return exc, result, captured["engine"]
def test_step_2_four_kv_groups_one_per_sip():
"""Four multi_user_decode launches, one per SIP, in ONE run_bench call.
Uses the CCL milestone pattern (``milestone_1h_ccl.py:283-292``):
``target_device="all"`` scopes the runtime to every SIP; then
``ctx.ahbm.set_device(sip)`` before each ``ctx.zeros``/``launch``
switches which SIP the next allocation+launch lands on. This is the
canonical sequential per-SIP pattern in the codebase — four separate
``run_bench`` calls with ``DeviceSelector("sip:N")`` is a misuse.
Expected to PASS — the SFR install draws intra-SIP edges only, so the
4 KV-groups can't see each other.
"""
if not TOPOLOGY_4SIP.exists():
pytest.skip(f"4-SIP topology missing: {TOPOLOGY_4SIP}")
sip_groups = [(sip, f"sip{sip}", 0) for sip in range(4)]
exc, result, engine = _run_multi_sip(sip_groups)
if exc is not None:
oplog_len = len(getattr(engine, "op_log", []) or []) if engine else 0
print(f"\nstep_2 FAIL — op_log records before crash: {oplog_len}")
traceback.print_exception(type(exc), exc, exc.__traceback__)
raise AssertionError(f"step_2 failed: {exc}") from exc
assert result is not None and result.completion.ok, (
f"step_2: completion not ok — {result.completion if result else None}"
)
# ── Step 3 — 8 KV-groups, two per SIP (study target) ──────────────
def test_step_3_eight_kv_groups_two_per_sip():
"""Two launches per SIP × 4 SIPs = 8 KV-groups total in one run_bench.
The headline target: 64 cubes serving 8 KV-groups, two disjoint 2×4
sub-meshes per SIP. ``cube_start=0`` puts the first KV-group on
cubes 0..7 (rows 0..1); ``cube_start=8`` puts the second on cubes
8..15 (rows 2..3). This is the use case ``DPPolicy.cube_start`` was
added to enable.
Expected to PASS — the SFR install draws intra-SIP edges only, so
the 8 KV-groups can't see each other; ``cube_start`` ensures the
two halves of each SIP land on disjoint cubes.
"""
if not TOPOLOGY_4SIP.exists():
pytest.skip(f"4-SIP topology missing: {TOPOLOGY_4SIP}")
sip_groups = [
(sip, f"sip{sip}_half{half}", half * N_CUBES_PER_KV_GROUP)
for sip in range(4) for half in (0, 1)
]
exc, result, engine = _run_multi_sip(sip_groups)
if exc is not None:
raise AssertionError(f"step_3 failed: {exc}") from exc
assert result is not None and result.completion.ok, (
f"step_3: completion not ok — {result.completion if result else None}"
)
tests/attention/test_attention_mesh_panels_diag.py
-198
View File
@@ -1,198 +0,0 @@
"""End-to-end engine drives for the four GQA Llama-70B panels (sub-cycle 4c step 2).
Mirrors the existing single_user_decode diag harness across all four panels
of the milestone-gqa-llama70b sweep (ADR-0057):
single_user_prefill ring-K/V kernel, intracube PE ring (8 PEs / 1 cube)
single_user_decode allreduce-mlo kernel, intracube PE ring
multi_user_prefill ring-K/V kernel, intercube multisip (4 cubes)
multi_user_decode allreduce-mlo kernel, intercube multisip
Each test runs the panel through ``run_bench`` with ``enable_data=True``
and asserts ``result.completion.ok``. Failures dump the engine's op_log
tail and the exception, mirroring the decode-diag harness format.
Validation-scale config matches ADR-0057 D4:
S_q_prefill=16, S_kv_per_rank=16, h_q=h_kv=1, d_head=64
n_ranks_single_user=8, n_ranks_multi_user=4
"""
from __future__ import annotations
import traceback
from pathlib import Path
import pytest
from kernbench.benches._attention_mesh_kv import attention_mesh_kv_kernel
from kernbench.benches._attention_mesh_mlo import attention_mesh_mlo_kernel
from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
from kernbench.ccl.sfr_config import (
configure_sfr_intercube_multisip,
configure_sfr_intracube_pe_ring,
)
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_PATH = Path(__file__).resolve().parents[2] / "topology.yaml"
S_Q_PREFILL = 16
S_Q_DECODE = 1
S_KV_PER_RANK = 16
H_Q = 1
H_KV = 1
D_HEAD = 64
N_RANKS_SINGLE_USER = 8
N_RANKS_MULTI_USER = 4
DTYPE = "f16"
# ── Helpers ──────────────────────────────────────────────────────
def _engine_factory(t, d):
return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
def _run_panel(bench_fn):
"""Drive a panel through run_bench; return (exc, result, engine)."""
topo = resolve_topology(str(TOPOLOGY_PATH))
captured: dict = {"engine": None}
def factory(t, d):
eng = _engine_factory(t, d)
captured["engine"] = eng
return eng
exc = None
result = None
try:
result = run_bench(
topology=topo, bench_fn=bench_fn,
device=resolve_device(None), engine_factory=factory,
)
except BaseException as e: # noqa: BLE001
exc = e
return exc, result, captured["engine"]
def _assert_ok(name: str, exc, result, engine) -> None:
if exc is not None:
oplog_len = len(getattr(engine, "op_log", []) or []) if engine else 0
print(f"\n========== {name} FAIL ==========")
print(f"op_log records before crash: {oplog_len}")
print(f"{type(exc).__name__}: {exc}")
traceback.print_exception(type(exc), exc, exc.__traceback__)
raise AssertionError(
f"{name} failed at runtime: {exc}"
) from exc
assert result is not None, f"{name}: no result"
assert result.completion.ok, f"{name}: completion not ok — {result.completion}"
# ── Panel bench fns ──────────────────────────────────────────────
def _bench_fn_single_user_prefill(ctx):
configure_sfr_intracube_pe_ring(
ctx.engine, ctx.spec,
resolve_algorithm_config(load_ccl_config(), name="lrab_hierarchical_allreduce"),
)
n = N_RANKS_SINGLE_USER
dp_full = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=n)
dp_kv = DPPolicy(cube="replicate", pe="row_wise", num_cubes=1, num_pes=n)
q = ctx.zeros((S_Q_PREFILL, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="q")
k = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="k")
v = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="v")
o = ctx.empty((S_Q_PREFILL, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="o")
ctx.launch(
"single_user_prefill_mesh", attention_mesh_kv_kernel,
q, k, v, o,
S_Q_PREFILL, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, n,
)
def _bench_fn_single_user_decode(ctx):
configure_sfr_intracube_pe_ring(
ctx.engine, ctx.spec,
resolve_algorithm_config(load_ccl_config(), name="lrab_hierarchical_allreduce"),
)
n = N_RANKS_SINGLE_USER
dp_full = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=n)
dp_kv = DPPolicy(cube="replicate", pe="row_wise", num_cubes=1, num_pes=n)
q = ctx.zeros((S_Q_DECODE, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="q")
k = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="k")
v = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="v")
o = ctx.empty((S_Q_DECODE, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="o")
ctx.launch(
"single_user_decode_mesh", attention_mesh_mlo_kernel,
q, k, v, o,
S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, n,
)
def _bench_fn_multi_user_prefill(ctx):
configure_sfr_intercube_multisip(
ctx.engine, ctx.spec,
resolve_algorithm_config(load_ccl_config(), name="lrab_hierarchical_allreduce"),
)
n = N_RANKS_MULTI_USER
dp_full = DPPolicy(cube="replicate", pe="replicate", num_cubes=n, num_pes=8)
dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=n, num_pes=8)
q = ctx.zeros((S_Q_PREFILL, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="q")
k = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="k")
v = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="v")
o = ctx.empty((S_Q_PREFILL, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="o")
ctx.launch(
"multi_user_prefill_mesh", attention_mesh_kv_kernel,
q, k, v, o,
S_Q_PREFILL, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, n,
1, # rank_axis=1 → ring at cube level (ADR-0059 multi_user)
_auto_dim_remap=False,
)
def _bench_fn_multi_user_decode(ctx):
configure_sfr_intercube_multisip(
ctx.engine, ctx.spec,
resolve_algorithm_config(load_ccl_config(), name="lrab_hierarchical_allreduce"),
)
n = N_RANKS_MULTI_USER
dp_full = DPPolicy(cube="replicate", pe="replicate", num_cubes=n, num_pes=8)
dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=n, num_pes=8)
q = ctx.zeros((S_Q_DECODE, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="q")
k = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="k")
v = ctx.zeros((S_KV_PER_RANK * n, H_KV * D_HEAD), dtype=DTYPE, dp=dp_kv, name="v")
o = ctx.empty((S_Q_DECODE, H_Q * D_HEAD), dtype=DTYPE, dp=dp_full, name="o")
ctx.launch(
"multi_user_decode_mesh", attention_mesh_mlo_kernel,
q, k, v, o,
S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, n,
1, # rank_axis=1 → ring at cube level (ADR-0059 multi_user)
_auto_dim_remap=False,
)
# ── Tests ────────────────────────────────────────────────────────
def test_single_user_prefill_through_engine():
exc, result, engine = _run_panel(_bench_fn_single_user_prefill)
_assert_ok("single_user_prefill", exc, result, engine)
def test_single_user_decode_through_engine():
exc, result, engine = _run_panel(_bench_fn_single_user_decode)
_assert_ok("single_user_decode", exc, result, engine)
def test_multi_user_prefill_through_engine():
exc, result, engine = _run_panel(_bench_fn_multi_user_prefill)
_assert_ok("multi_user_prefill", exc, result, engine)
def test_multi_user_decode_through_engine():
exc, result, engine = _run_panel(_bench_fn_multi_user_decode)
_assert_ok("multi_user_decode", exc, result, engine)
tests/attention/test_gqa_decode.py
+143
View File
@@ -0,0 +1,143 @@
"""Phase 1 spec test for P1a GQA decode kernel (real GQA via M-fold).
P1a is the first phase of the DDD-0060 plan, split out of the original
P1 (the composite-hybrid swap is P1b, deferred until the tl.composite
output-handle question is decided). P1a is the *correctness* unlock:
the kernel processes ONE KV head at a time and folds the G query heads
into the matmul M (row) dimension so a single Q·Kᵀ serves all G heads
sharing one K (ADR-0060 §5.2). This lifts the baseline's
``h_q == h_kv == 1`` cap pinned at
``tests/attention/test_milestone_gqa_llama70b.py:137-142``.
P1a stays inside the existing ``tl`` API: the two attention GEMMs use the
blocking ``tl.dot`` so the chain ``Q·Kᵀ → softmax → P·V → store`` fits
without any composite-output chaining. The composite swap (P1b) will
revisit this once the API for feeding a composite's output into a
downstream MATH op is settled.
Phase 1 (this commit): tests only — production code lands in Phase 2.
Tests fail at import in Phase 1 with ModuleNotFoundError; Phase 2 makes
them pass.
"""
from __future__ import annotations
from pathlib import Path
from kernbench.benches._gqa_decode import gqa_decode_kernel # noqa: F401 (Phase 2 deliverable)
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"
# Decode shapes — P1a is one-shot, no tiling, single rank. P3 will tile.
T_Q = 1
D_HEAD = 64
S_KV = 16
DTYPE = "f16"
def _engine_factory(t, d):
return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
def _run_decode_p1(*, h_q: int, h_kv: int):
"""One-shot GQA decode on a single PE in a single CUBE (P=1, no SP).
Tensor layout (natural K — same as the P2a SP tests; the kernel
reshapes K to (d_head, S_local) via byte-conserving load):
Q : (T_q, h_q · d_head) — natural Q layout; kernel reshapes to
(G·T_q, d_head) — byte-conserving and math-correct because
T_q=1 collapses axis ordering.
K : (S_kv, h_kv · d_head) — natural K layout; kernel loads as
(d_head, S_local) — reshape-as-transpose caveat (ADR-0060 §3
/ §B item 2), correct for zero inputs used here.
V : (S_kv, h_kv · d_head) — natural V layout.
O : (T_q, h_q · d_head) — same shape as Q.
"""
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
def _bench_fn(ctx):
dp = DPPolicy(cube="replicate", pe="replicate",
num_cubes=1, num_pes=1)
q = ctx.zeros((T_Q, h_q * D_HEAD),
dtype=DTYPE, dp=dp, name=f"q_h{h_q}_kv{h_kv}")
k = ctx.zeros((S_KV, h_kv * D_HEAD),
dtype=DTYPE, dp=dp, name=f"k_h{h_q}_kv{h_kv}")
v = ctx.zeros((S_KV, h_kv * D_HEAD),
dtype=DTYPE, dp=dp, name=f"v_h{h_q}_kv{h_kv}")
o = ctx.empty((T_Q, h_q * D_HEAD),
dtype=DTYPE, dp=dp, name=f"o_h{h_q}_kv{h_kv}")
ctx.launch(
f"gqa_decode_p1_h{h_q}_kv{h_kv}",
gqa_decode_kernel,
q, k, v, o,
T_Q, S_KV, h_q, h_kv, D_HEAD,
1, 1, # C=1, P=1 (no SP, degenerate)
_auto_dim_remap=False,
)
return run_bench(
topology=topo,
bench_fn=_bench_fn,
device=resolve_device(None),
engine_factory=_engine_factory,
)
def _dma_read_count(op_log) -> int:
return sum(1 for r in op_log if r.op_name == "dma_read")
# ── Headline unlock: real GQA (h_q = G·h_kv) runs in data mode ──────────
def test_real_gqa_h_q_eight_h_kv_one_completes_in_data_mode():
"""ADR-0060 §A.1 headline unlock — the baseline raises
``ValueError: Shape mismatch …`` in MemoryStore at h_q=8, h_kv=1
because ``_view(K, (h_q·d, S_kv))`` only byte-conserves when h_q==h_kv.
M-fold processes one KV head at a time using only byte-conserving
reshapes, so this completes.
"""
result = _run_decode_p1(h_q=8, h_kv=1)
assert result.completion.ok, (
f"real GQA (h_q=8, h_kv=1) decode failed: {result.completion}"
)
# ── M-fold property: K/V loads do not scale with G ─────────────────────
def test_kv_dma_read_count_independent_of_g():
"""ADR-0060 TL;DR / §5.2: M-fold loads K and V once per KV head and
folds the G query heads into the GEMM M dim. The dma_read_count must
therefore be identical between (G=1, h_kv=1) and (G=8, h_kv=1) — both
issue exactly 3 reads (Q + K + V). This pins the GQA-reuse property
that the rest of the plan (composite streaming in P4, etc.) builds on.
"""
g1 = _run_decode_p1(h_q=1, h_kv=1)
g8 = _run_decode_p1(h_q=8, h_kv=1)
n_g1 = _dma_read_count(g1.engine.op_log)
n_g8 = _dma_read_count(g8.engine.op_log)
assert n_g1 == 3, f"G=1 dma_read_count must be 3 (Q+K+V); got {n_g1}"
assert n_g8 == 3, f"G=8 dma_read_count must be 3 (Q+K+V); got {n_g8}"
assert n_g1 == n_g8, (
f"K/V dma_read_count must be independent of G; "
f"got G=1 -> {n_g1}, G=8 -> {n_g8}"
)
# ── Backward-compat: degenerate G=1 still works ────────────────────────
def test_degenerate_g_equals_one_still_works():
"""G=1 (h_q == h_kv == 1) is the baseline-compatible config. M-fold
degenerates to (T_q, d) = (1, 64) — the same shape the baseline
already exercises — so this proves no regression on that path.
"""
result = _run_decode_p1(h_q=1, h_kv=1)
assert result.completion.ok, (
f"degenerate G=1 decode failed: {result.completion}"
)
tests/attention/test_gqa_decode_mc.py
+182
View File
@@ -0,0 +1,182 @@
"""Phase 1 spec test for P2b GQA decode multi-cube SP (both Level-2 + Level-1).
P2b extends P2a to multiple CUBEs in one CUBE Group. The kernel uses the
canonical full SFR install (``configure_sfr_intercube_multisip``) which
provides disjoint direction namespaces:
- ``intra_E / intra_W / intra_N / intra_S`` — PE↔PE within a CUBE
(logical 2×4 grid, no wrap)
- ``E / W / N / S`` — CUBE↔CUBE inter-CUBE
(mesh, no wrap)
Reduce pattern (chain reduce-to-root, ADR-0060 §A.2 spirit, §4 chain
deviation noted):
Level-2 (intra-CUBE, 2×4 grid):
row-then-col chain — each row reduces leftward along ``intra_W`` to
its col-0 PE, then PE 4 sends to PE 0 along ``intra_N``. 7 chain
steps per CUBE × 3 handles each = 21 ``ipcq_copy`` per CUBE.
Level-1 (inter-CUBE):
only PE 0 of each CUBE participates. Chain leftward along ``W``.
(C-1) chain steps × 3 handles each.
Final store at PE 0 of CUBE 0 only.
Phase 1 (this commit): tests only — production code lands in Phase 2.
Phase 2 also updates ``test_gqa_decode.py`` (add C=1) and
``test_gqa_decode_sp.py`` (switch SFR + add C=1).
"""
from __future__ import annotations
from pathlib import Path
from kernbench.benches._gqa_decode import gqa_decode_kernel # noqa: F401
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"
T_Q = 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_decode_mc(*, h_q: int, h_kv: int, C: int, P: int, S_kv: int):
"""Multi-CUBE SP decode: C cubes × P PEs each share the work."""
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=C, num_pes=P)
dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
num_cubes=C, num_pes=P)
# Total KV split across C×P ranks; each rank sees S_kv/(C·P) rows.
q = ctx.zeros((T_Q, h_q * D_HEAD),
dtype=DTYPE, dp=dp_full,
name=f"q_h{h_q}_kv{h_kv}_c{C}_p{P}")
k = ctx.zeros((S_kv, h_kv * D_HEAD),
dtype=DTYPE, dp=dp_kv,
name=f"k_h{h_q}_kv{h_kv}_c{C}_p{P}")
v = ctx.zeros((S_kv, h_kv * D_HEAD),
dtype=DTYPE, dp=dp_kv,
name=f"v_h{h_q}_kv{h_kv}_c{C}_p{P}")
o = ctx.empty((T_Q, h_q * D_HEAD),
dtype=DTYPE, dp=dp_full,
name=f"o_h{h_q}_kv{h_kv}_c{C}_p{P}")
ctx.launch(
f"gqa_decode_mc_h{h_q}_kv{h_kv}_c{C}_p{P}",
gqa_decode_kernel,
q, k, v, o,
T_Q, S_kv, h_q, h_kv, D_HEAD,
C, P,
_auto_dim_remap=False,
)
return run_bench(
topology=topo,
bench_fn=_bench_fn,
device=resolve_device(None),
engine_factory=_engine_factory,
)
def _count(op_log, name: str) -> int:
return sum(1 for r in op_log if r.op_name == name)
# ── Degenerate C=1 P=1 ────────────────────────────────────────────────
def test_mc_c_one_p_one_degenerate():
"""C=1, P=1: single rank, no SP. No IPCQ traffic; one dma_write."""
result = _run_decode_mc(h_q=1, h_kv=1, C=1, P=1, S_kv=16)
assert result.completion.ok, (
f"C=1 P=1 degenerate failed: {result.completion}"
)
assert _count(result.engine.op_log, "ipcq_copy") == 0
assert _count(result.engine.op_log, "dma_write") == 1
# ── Intra-CUBE only (single CUBE, P=8 on 2×4 grid) ────────────────────
def test_mc_c_one_p_eight_intracube_grid():
"""C=1, P=8: intra-cube row-then-col chain on the 2×4 grid.
7 chain steps × 3 handles (m, ℓ, O) = 21 ipcq_copy."""
result = _run_decode_mc(h_q=1, h_kv=1, C=1, P=8, S_kv=64)
assert result.completion.ok, (
f"C=1 P=8 intra-cube failed: {result.completion}"
)
assert _count(result.engine.op_log, "dma_write") == 1
n_copy = _count(result.engine.op_log, "ipcq_copy")
assert n_copy == 21, (
f"C=1 P=8: expected 21 ipcq_copy (7 chain × 3 handles); got {n_copy}"
)
# ── Multi-CUBE root-only write ────────────────────────────────────────
def test_mc_c_two_p_eight_root_only_writes_o():
"""C=2, P=8: 16 ranks total. Only PE 0 of CUBE 0 writes O."""
result = _run_decode_mc(h_q=1, h_kv=1, C=2, P=8, S_kv=128)
assert result.completion.ok, (
f"C=2 P=8 multi-cube failed: {result.completion}"
)
n_writes = _count(result.engine.op_log, "dma_write")
assert n_writes == 1, (
f"root-only write must hold for C=2 P=8 (16 ranks); got {n_writes}"
)
# ── Multi-CUBE total IPCQ chain count ─────────────────────────────────
def test_mc_c_two_p_eight_total_ipcq_count():
"""C=2, P=8: 21 intra-cube ipcq_copy per CUBE × 2 CUBEs + 3 inter-cube
chain ipcq_copy (C-1=1 step × 3 handles) = 45 total."""
result = _run_decode_mc(h_q=1, h_kv=1, C=2, P=8, S_kv=128)
assert result.completion.ok, (
f"C=2 P=8 multi-cube failed: {result.completion}"
)
n_copy = _count(result.engine.op_log, "ipcq_copy")
expected = 21 * 2 + (2 - 1) * 3
assert n_copy == expected, (
f"C=2 P=8: expected {expected} ipcq_copy (21 intra × 2 CUBEs + 3 "
f"inter); got {n_copy}"
)
# ── Real GQA × multi-CUBE SP combined ─────────────────────────────────
def test_mc_real_gqa_c_two_p_eight():
"""Headline: real GQA (h_q = G·h_kv with G=8) AND multi-CUBE SP
(C=2, P=8) together — the case the original baseline can express
neither part of."""
result = _run_decode_mc(h_q=8, h_kv=1, C=2, P=8, S_kv=128)
assert result.completion.ok, (
f"real GQA + multi-CUBE SP combined failed: {result.completion}"
)
n_writes = _count(result.engine.op_log, "dma_write")
assert n_writes == 1, (
f"root-only write must hold under M-fold + multi-CUBE; "
f"got {n_writes}"
)
tests/attention/test_gqa_decode_sp.py
+153
View File
@@ -0,0 +1,153 @@
"""Phase 1 spec test for P2a GQA decode SP (chain reduce-to-root, Level-2 only).
P2a is the first half of DDD-0060 P2: the kernel becomes multi-PE within
one CUBE and reduces to root (PE 0) using a chain over the 1D intra-cube
ring (W direction). This **replaces the baseline's bidirectional O(N)
fan-out** where every rank ends with O — ADR-0060 §A.2's headline.
Deviation from DDD-0060 §7 P2 gate: the gate text asks for
``⌈log₂ P⌉`` reduce rounds. The intra-cube SFR install
(``configure_sfr_intracube_pe_ring``) wires only a 1D E/W ring, so a
true tree would require either multi-hop forwarding or a new SFR install
(future ADR). P2a uses a **chain reduce-to-root**: ``P-1`` rounds along
W. The architectural property the ADR cares about
(root-only output vs every-rank-has-O) is preserved; the logarithmic
collective is deferred.
P2b (deferred) covers Level-1 inter-CUBE center-mesh reduce (C>1).
Phase 1 (this commit): tests only — production code lands in Phase 2.
"""
from __future__ import annotations
from pathlib import Path
from kernbench.benches._gqa_decode import gqa_decode_kernel # noqa: F401
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"
T_Q = 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_decode_sp(*, h_q: int, h_kv: int, P: int, S_kv: int):
"""Single-CUBE SP decode: P PEs share the work along the intra-cube ring."""
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=1, num_pes=P)
dp_kv = DPPolicy(cube="replicate", pe="row_wise",
num_cubes=1, num_pes=P)
q = ctx.zeros((T_Q, h_q * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"q_h{h_q}_kv{h_kv}_p{P}")
# KV: total S_kv split across P PEs along axis 0 (row_wise sharding).
# Each PE sees (S_kv/P, h_kv·D_HEAD).
k = ctx.zeros((S_kv, h_kv * D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"k_h{h_q}_kv{h_kv}_p{P}")
v = ctx.zeros((S_kv, h_kv * D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"v_h{h_q}_kv{h_kv}_p{P}")
o = ctx.empty((T_Q, h_q * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"o_h{h_q}_kv{h_kv}_p{P}")
ctx.launch(
f"gqa_decode_sp_h{h_q}_kv{h_kv}_p{P}",
gqa_decode_kernel,
q, k, v, o,
T_Q, S_kv, h_q, h_kv, D_HEAD,
1, P, # C=1, P=P (single-CUBE SP)
_auto_dim_remap=False,
)
return run_bench(
topology=topo,
bench_fn=_bench_fn,
device=resolve_device(None),
engine_factory=_engine_factory,
)
def _count(op_log, name: str) -> int:
return sum(1 for r in op_log if r.op_name == name)
# ── Root-only write ────────────────────────────────────────────────────
def test_sp_chain_reduce_root_only_writes_o():
"""ADR-0060 §A.2: only the root rank (PE 0) writes O. Baseline today
has every rank write the full final O (bidirectional fan-out)."""
result = _run_decode_sp(h_q=1, h_kv=1, P=8, S_kv=64)
assert result.completion.ok, f"P=8 chain reduce failed: {result.completion}"
n_writes = _count(result.engine.op_log, "dma_write")
assert n_writes == 1, (
f"reduce-to-root must produce exactly 1 dma_write (PE 0); "
f"got {n_writes}"
)
# ── Chain step count ───────────────────────────────────────────────────
def test_sp_chain_reduce_p_minus_one_ipcq_pairs():
"""Chain reduce-to-root has P-1 send→recv pairs along the W chain;
each pair logs one ``ipcq_copy`` (inbound DMA, per
``milestone_gqa_llama70b._summarize_op_log``). Each chain step ships
the triplet (m, ℓ, O) → 3 handles per step → 7 steps × 3 = 21."""
result = _run_decode_sp(h_q=1, h_kv=1, P=8, S_kv=64)
assert result.completion.ok, f"P=8 chain reduce failed: {result.completion}"
n_copy = _count(result.engine.op_log, "ipcq_copy")
expected = (8 - 1) * 3
assert n_copy == expected, (
f"chain reduce: expected {expected} ipcq_copy (P-1=7 steps × "
f"3 handles m/ℓ/O); got {n_copy}"
)
# ── Real GQA × SP combined ─────────────────────────────────────────────
def test_sp_real_gqa_h_q_eight_h_kv_one_p_eight():
"""The combined unlock: real GQA (h_q=G·h_kv with G=8) AND SP
(P=8) together — neither expressible by the baseline."""
result = _run_decode_sp(h_q=8, h_kv=1, P=8, S_kv=64)
assert result.completion.ok, (
f"real GQA + SP combined run failed: {result.completion}"
)
n_writes = _count(result.engine.op_log, "dma_write")
assert n_writes == 1, (
f"root-only write must hold under M-fold too; got {n_writes}"
)
# ── Degenerate P=1 ────────────────────────────────────────────────────
def test_sp_p_one_degenerate_no_ipcq_traffic():
"""P=1: SP degenerates to a single rank. No IPCQ traffic; one dma_write."""
result = _run_decode_sp(h_q=8, h_kv=1, P=1, S_kv=16)
assert result.completion.ok, f"P=1 degenerate failed: {result.completion}"
n_send = _count(result.engine.op_log, "ipcq_send")
n_recv = _count(result.engine.op_log, "ipcq_recv")
assert n_send == 0, f"P=1 must have no ipcq_send; got {n_send}"
assert n_recv == 0, f"P=1 must have no ipcq_recv; got {n_recv}"
n_writes = _count(result.engine.op_log, "dma_write")
assert n_writes == 1, f"P=1: one dma_write; got {n_writes}"
tests/attention/test_gqa_long_context.py
+157
View File
@@ -0,0 +1,157 @@
"""Phase 1 spec test for Phase C: long-context regression for the GQA
kernels (ADR-0060 §B "long/short context split" + ADR-0063 §A.2).
The headline panel today caps prefill at S_kv=16 / decode at S_kv≤128 —
NOT because the algorithm fails, but because the bump allocator never
recycles. ADR-0063 §A.2 (test req 3) requires a sweep at an S that
would overflow 1 MiB without recycling to complete after scope
discipline lands.
Scratch-budget estimate at C=4, P=8 (current 1 MiB pool):
decode: per-rank S_local = S_kv / 32; intermediates ≲ 500 KB at
S_kv=32K (fits today — used as regression guard).
prefill: per-rank S_local = S_kv / 4; ~16·S_local bytes per ring
step × 4 steps ≈ 64·S_local bytes. Overflows at
~64 K tokens (64 × 16K > 1 MiB).
Phase 1 (this commit): tests only — production code lands in Phase 2.
The prefill test fails today with a RuntimeError("TLContext scratch
overflow"). After Phase 2 (scratch_scope + copy_to discipline) it
completes.
"""
from __future__ import annotations
from pathlib import Path
from kernbench.benches._gqa_decode import gqa_decode_kernel # noqa: F401
from kernbench.benches._gqa_prefill import gqa_prefill_kernel # noqa: F401
from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
from kernbench.ccl.sfr_config import (
configure_sfr_intercube_multisip,
configure_sfr_intercube_ring,
)
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"
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)
# ── Decode at moderate-long context (regression guard) ───────────────
def test_decode_long_context_32k_completes():
"""Decode at S_kv=32K (C=1, P=8) — per-rank S_local=4K. Should
complete with current scratch usage (~few KB intermediates) and
must continue to complete after the rewrite.
This is a regression guard: scratch discipline shouldn't break
decode's existing long-context capability.
"""
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
def _bench_fn(ctx):
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
P = 8
S_kv = 32_768
dp_full = DPPolicy(cube="replicate", pe="replicate",
num_cubes=1, num_pes=P)
dp_kv = DPPolicy(cube="replicate", pe="row_wise",
num_cubes=1, num_pes=P)
ctx.zeros((1, 8 * D_HEAD), dtype=DTYPE, dp=dp_full, name="q_long_dec")
k = ctx.zeros((S_kv, D_HEAD),
dtype=DTYPE, dp=dp_kv, name="k_long_dec")
v = ctx.zeros((S_kv, D_HEAD),
dtype=DTYPE, dp=dp_kv, name="v_long_dec")
o = ctx.empty((1, 8 * D_HEAD),
dtype=DTYPE, dp=dp_full, name="o_long_dec")
q = ctx.zeros((1, 8 * D_HEAD),
dtype=DTYPE, dp=dp_full, name="q_long_dec_2")
ctx.launch(
"gqa_decode_long_32k",
gqa_decode_kernel,
q, k, v, o,
1, S_kv, 8, 1, D_HEAD,
1, P,
_auto_dim_remap=False,
)
result = run_bench(
topology=topo, bench_fn=_bench_fn,
device=resolve_device(None), engine_factory=_engine_factory,
)
assert result.completion.ok, (
f"decode at S_kv=32K must complete; got {result.completion}"
)
# ── Prefill at long context overflows without scope discipline ───────
def test_prefill_long_context_completes_after_scope_discipline():
"""ADR-0063 §A.2 Test Req 3: a sweep at an ``S`` that would overflow
1 MiB without recycling must complete after scope discipline lands.
With C=4 and S_kv chosen so per-rank S_local·8·4 > 1 MiB
(~16K tokens per rank ⇒ S_kv ≥ 64K), the current kernel overflows
the 1 MiB scratch pool. After Phase 2 (scratch_scope wraps each
ring step's intermediates; copy_to persists running state), it
completes.
This is the headline test that proves the S-ceiling is gone for
prefill.
"""
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
def _bench_fn(ctx):
configure_sfr_intercube_ring(
ctx.engine, ctx.spec, _ccl_cfg(), ring_size=4,
)
T_q = 4
S_kv = 65_536 # 64 K — per-rank 16 K, ~2 MB scratch w/o scope
C = 4
dp_q = DPPolicy(cube="replicate", pe="replicate",
num_cubes=C, num_pes=1)
dp_kv = DPPolicy(cube="row_wise", pe="replicate",
num_cubes=C, num_pes=1)
dp_o = DPPolicy(cube="row_wise", pe="replicate",
num_cubes=C, num_pes=1)
q = ctx.zeros((T_q, D_HEAD),
dtype=DTYPE, dp=dp_q, name="q_long_pre")
k = ctx.zeros((S_kv, D_HEAD),
dtype=DTYPE, dp=dp_kv, name="k_long_pre")
v = ctx.zeros((S_kv, D_HEAD),
dtype=DTYPE, dp=dp_kv, name="v_long_pre")
o = ctx.empty((T_q * C, D_HEAD),
dtype=DTYPE, dp=dp_o, name="o_long_pre")
ctx.launch(
"gqa_prefill_long_64k",
gqa_prefill_kernel,
q, k, v, o,
T_q, S_kv, D_HEAD, C,
_auto_dim_remap=False,
)
result = run_bench(
topology=topo, bench_fn=_bench_fn,
device=resolve_device(None), engine_factory=_engine_factory,
)
assert result.completion.ok, (
f"prefill at S_kv=64K must complete after scope discipline lands; "
f"got {result.completion}"
)
tests/attention/test_gqa_prefill.py
+118
View File
@@ -0,0 +1,118 @@
"""Phase 1 spec test for P6a GQA prefill kernel (head-parallel, C=1 baseline).
P6a introduces ``_gqa_prefill.py`` with the head-parallel structure (one
Q head per CUBE, per-CUBE distributed output, no reduce). C=1 is the
degenerate case — no Ring KV, no IPCQ traffic. Validates kernel
structure and T_q > 1 attention.
P6b (deferred) adds the Ring KV rotation for C > 1, which needs either
a new SFR install function (intra_* + wrapped E/W at CUBE level) or a
topology-specific config — separate design call.
The prefill kernel differs from decode (P1a/P2a/P2b) in three ways
(ADR-0060 §5.5 / TL;DR):
1. Q has T_q > 1 rows (not just decode's single timestep).
2. Head-parallel placement: each CUBE owns ONE Q head — no M-fold.
3. Each CUBE writes its own head's output — NO reduce.
Phase 1 (this commit): tests only — production code lands in Phase 2.
"""
from __future__ import annotations
from pathlib import Path
from kernbench.benches._gqa_prefill import gqa_prefill_kernel # noqa: F401
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"
D_HEAD = 64
DTYPE = "f16"
def _engine_factory(t, d):
return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
def _run_prefill(*, T_q: int, S_kv: int, C: int = 1):
"""C=1 head-parallel prefill: single CUBE owns the one head + full KV."""
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
def _bench_fn(ctx):
dp = DPPolicy(cube="replicate", pe="replicate",
num_cubes=C, num_pes=1)
# Q: (T_q, d_head) — one head per CUBE (head-parallel; for C=1
# only one head total). 2D layout matches what the kernel loads.
# P6b will use a 3D (h_q, T_q, d_head) Q with cube_row_wise
# sharding so each CUBE owns its head.
q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp,
name=f"q_t{T_q}_c{C}")
# K, V: full local for C=1 (no ring). Kernel loads K as
# (d_head, S_kv) via byte-conserving reshape.
k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp,
name=f"k_t{T_q}_c{C}")
v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp,
name=f"v_t{T_q}_c{C}")
# O: (T_q, d_head) — per-CUBE distributed output.
o = ctx.empty((T_q, D_HEAD), dtype=DTYPE, dp=dp,
name=f"o_t{T_q}_c{C}")
ctx.launch(
f"gqa_prefill_p6a_t{T_q}_s{S_kv}_c{C}",
gqa_prefill_kernel,
q, k, v, o,
T_q, S_kv, D_HEAD, C,
_auto_dim_remap=False,
)
return run_bench(
topology=topo, bench_fn=_bench_fn,
device=resolve_device(None),
engine_factory=_engine_factory,
)
def _count(op_log, name: str) -> int:
return sum(1 for r in op_log if r.op_name == name)
def test_prefill_c_one_t_q_one_completes():
"""C=1, T_q=1: smallest workload (decode-like)."""
result = _run_prefill(T_q=1, S_kv=16, C=1)
assert result.completion.ok, (
f"prefill C=1 T_q=1 failed: {result.completion}"
)
def test_prefill_c_one_t_q_four_completes():
"""C=1, T_q=4: real prefill (Q has multiple rows) — distinguishes
prefill from decode (T_q=1)."""
result = _run_prefill(T_q=4, S_kv=16, C=1)
assert result.completion.ok, (
f"prefill C=1 T_q=4 failed: {result.completion}"
)
def test_prefill_c_one_no_ipcq_traffic():
"""C=1: no ring step, no IPCQ traffic."""
result = _run_prefill(T_q=4, S_kv=16, C=1)
assert result.completion.ok
n_copy = _count(result.engine.op_log, "ipcq_copy")
assert n_copy == 0, (
f"C=1 must have no IPCQ traffic (no ring); got {n_copy}"
)
def test_prefill_c_one_one_dma_write():
"""C=1, one head: exactly one dma_write (per-CUBE distributed output;
no reduce). For C > 1 in P6b this becomes dma_write_count == C."""
result = _run_prefill(T_q=4, S_kv=16, C=1)
assert result.completion.ok
n_writes = _count(result.engine.op_log, "dma_write")
assert n_writes == 1, (
f"C=1 prefill: expected 1 dma_write (one head per cube); "
f"got {n_writes}"
)
tests/attention/test_gqa_prefill_ring.py
+161
View File
@@ -0,0 +1,161 @@
"""Phase 1 spec test for P6b GQA prefill Ring KV (head-parallel, C>1).
P6b adds the Ring KV rotation (ADR-0060 §5.5) to the prefill kernel.
Each CUBE owns one Q head + its KV slice; over C ring steps the KV
blocks rotate around the C CUBEs so every CUBE sees every block. The
running ``(m, ℓ, O)`` is folded inside the loop. No reduce — each CUBE
writes its own head's output.
Requires a new SFR install ``configure_sfr_intercube_ring(ring_size=C)``
that wires:
- ``intra_*`` : 2×4 PE grid within a CUBE (same as multisip)
- ``E/W`` : 1D ring of CUBEs 0..ring_size-1 WITH WRAP
(symmetric to ``configure_sfr_intracube_pe_ring``,
applied at CUBE level)
- ``global_*`` : SIP-level (same as multisip)
The kernel ring body:
for step in range(1, C):
tl.send(dir="W", src=Kc)
tl.send(dir="W", src=Vc)
Kc = tl.recv(dir="E", ...)
Vc = tl.recv(dir="E", ...)
... local partial + online-softmax merge into (m, ℓ, O) ...
Per CUBE per step: 2 sends (K, V) → 2 ``ipcq_copy``. Across all CUBEs:
``(C-1) * 2 * C`` total ipcq_copy.
Restriction in P6b first cut: ``C ∈ {1, 2, 4}`` (single row of the 4×4
cube mesh). C=8 ring would span rows — follow-on.
Phase 1 (this commit): tests only — production code lands in Phase 2.
"""
from __future__ import annotations
from pathlib import Path
from kernbench.benches._gqa_prefill import gqa_prefill_kernel # noqa: F401
from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
from kernbench.ccl.sfr_config import configure_sfr_intercube_ring # noqa: F401 (Phase 2)
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"
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_prefill_ring(*, T_q: int, S_kv: int, C: int):
"""Head-parallel prefill with Ring KV across C CUBEs."""
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
def _bench_fn(ctx):
# P6b: new SFR install with cube-level ring wrap.
configure_sfr_intercube_ring(
ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C,
)
# Q replicated on every CUBE (zeros for testing; per-CUBE head
# indexing is implicit). KV sequence-sharded by CUBE. O
# distributed — each CUBE writes its slice of (T_q*C, d_head).
dp_q = DPPolicy(cube="replicate", pe="replicate",
num_cubes=C, num_pes=1)
dp_kv = DPPolicy(cube="row_wise", pe="replicate",
num_cubes=C, num_pes=1)
dp_o = DPPolicy(cube="row_wise", pe="replicate",
num_cubes=C, num_pes=1)
q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp_q,
name=f"q_t{T_q}_c{C}_ring")
k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
name=f"k_t{T_q}_c{C}_ring")
v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
name=f"v_t{T_q}_c{C}_ring")
# O: (T_q * C, d_head), each CUBE writes (T_q, d_head) at its
# slice. dma_write_count = C (per-CUBE distributed output).
o = ctx.empty((T_q * C, D_HEAD), dtype=DTYPE, dp=dp_o,
name=f"o_t{T_q}_c{C}_ring")
ctx.launch(
f"gqa_prefill_ring_t{T_q}_s{S_kv}_c{C}",
gqa_prefill_kernel,
q, k, v, o,
T_q, S_kv, D_HEAD, C,
_auto_dim_remap=False,
)
return run_bench(
topology=topo, bench_fn=_bench_fn,
device=resolve_device(None),
engine_factory=_engine_factory,
)
def _count(op_log, name: str) -> int:
return sum(1 for r in op_log if r.op_name == name)
# ── C=2 Ring KV ────────────────────────────────────────────────────────
def test_prefill_ring_c_two_completes():
"""C=2: 1 ring step rotates KV between the 2 CUBEs."""
result = _run_prefill_ring(T_q=4, S_kv=16, C=2)
assert result.completion.ok, (
f"prefill ring C=2 failed: {result.completion}"
)
def test_prefill_ring_c_two_distributed_output():
"""C=2: per-CUBE distributed output, no reduce → dma_write_count == 2."""
result = _run_prefill_ring(T_q=4, S_kv=16, C=2)
assert result.completion.ok
n_writes = _count(result.engine.op_log, "dma_write")
assert n_writes == 2, (
f"C=2 prefill: expected 2 dma_write (one per CUBE); got {n_writes}"
)
def test_prefill_ring_c_two_ipcq_count():
"""C=2: 1 ring step × 2 handles (K, V) × 2 CUBEs = 4 ipcq_copy."""
result = _run_prefill_ring(T_q=4, S_kv=16, C=2)
assert result.completion.ok
n_copy = _count(result.engine.op_log, "ipcq_copy")
expected = (2 - 1) * 2 * 2
assert n_copy == expected, (
f"C=2 ring: expected {expected} ipcq_copy "
f"((C-1)·2·C = 1·2·2); got {n_copy}"
)
# ── C=4 Ring KV (combined assertions) ─────────────────────────────────
def test_prefill_ring_c_four_combined():
"""C=4: 3 ring steps rotate KV around 4 CUBEs.
Expected: completes; 4 dma_writes; (4-1)·2·4 = 24 ipcq_copy."""
result = _run_prefill_ring(T_q=4, S_kv=32, C=4)
assert result.completion.ok, (
f"prefill ring C=4 failed: {result.completion}"
)
n_writes = _count(result.engine.op_log, "dma_write")
assert n_writes == 4, (
f"C=4 prefill: expected 4 dma_write (one per CUBE); got {n_writes}"
)
n_copy = _count(result.engine.op_log, "ipcq_copy")
expected = (4 - 1) * 2 * 4
assert n_copy == expected, (
f"C=4 ring: expected {expected} ipcq_copy "
f"((C-1)·2·C = 3·2·4); got {n_copy}"
)
tests/attention/test_gqa_scoped_writeback.py
+198
View File
@@ -0,0 +1,198 @@
"""Phase 1 spec test for Phase C: scratch_scope + tl.copy_to discipline
in the GQA kernels (ADR-0060 §5.2 / §5.5 + ADR-0063 §D3 / §D3.1).
ADR-0060 §5.2 (decode pseudocode line 75 / §5.5 (prefill pseudocode line
96) both wrap per-tile / per-ring-step intermediates in
``with tl.scratch_scope():`` and persist the merged running ``(m, ℓ, O)``
to a persistent arena allocated outside the scope. The original ADR-0063
§D3 specifies the two-arena pattern; §D3.1 specifies the
``tl.copy_to(dst, src)`` writeback primitive used to persist scoped
results.
Currently neither kernel uses ``scratch_scope`` or ``copy_to``; their
chain-merge / ring-merge bodies allocate every intermediate from the
bump cursor and never recycle. Result: op_log has 0 ``copy`` entries.
After Phase 2: each per-tile / per-step merge writes the new running
``(m, ℓ, O)`` via ``copy_to`` to the persistent arena. Per merge step
→ 3 copy ops (m, ℓ, O).
Phase 1 (this commit): tests only — production code lands in Phase 2.
"""
from __future__ import annotations
from pathlib import Path
from kernbench.benches._gqa_decode import gqa_decode_kernel # noqa: F401
from kernbench.benches._gqa_prefill import gqa_prefill_kernel # noqa: F401
from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
from kernbench.ccl.sfr_config import (
configure_sfr_intercube_multisip,
configure_sfr_intercube_ring,
)
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"
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 _count(op_log, name: str) -> int:
return sum(1 for r in op_log if r.op_name == name)
# ── Decode chain merges must use scratch_scope + copy_to ─────────────
def _run_decode_sp(*, h_q: int, h_kv: int, P: int, S_kv: int):
"""Single-CUBE SP decode (C=1, P PEs along intra-cube chain)."""
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=1, num_pes=P)
dp_kv = DPPolicy(cube="replicate", pe="row_wise",
num_cubes=1, num_pes=P)
q = ctx.zeros((1, h_q * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"q_sc_{P}")
k = ctx.zeros((S_kv, h_kv * D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"k_sc_{P}")
v = ctx.zeros((S_kv, h_kv * D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"v_sc_{P}")
o = ctx.empty((1, h_q * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"o_sc_{P}")
ctx.launch(
f"gqa_decode_scoped_{P}",
gqa_decode_kernel,
q, k, v, o,
1, S_kv, h_q, h_kv, D_HEAD,
1, P,
_auto_dim_remap=False,
)
return run_bench(
topology=topo, bench_fn=_bench_fn,
device=resolve_device(None), engine_factory=_engine_factory,
)
def test_decode_chain_merges_emit_copy_to_writeback():
"""ADR-0060 §5.2 + ADR-0063 §D3.1: each chain-merge step must wrap
its intermediates in ``scratch_scope`` and persist the new running
``(m, ℓ, O)`` via ``tl.copy_to``.
For (C=1, P=8): 7 intra-cube chain merges × 3 handles (m, ℓ, O) per
merge ⇒ 21 ``copy`` entries.
Currently 0 because the kernel never calls ``tl.copy_to``.
"""
result = _run_decode_sp(h_q=8, h_kv=1, P=8, S_kv=64)
assert result.completion.ok, f"decode SP failed: {result.completion}"
n_copy = _count(result.engine.op_log, "copy")
assert n_copy > 0, (
f"decode kernel must emit copy_to writeback per merge step "
f"(ADR-0060 §5.2 + ADR-0063 §D3.1); got 0 ``copy`` entries"
)
# ── Prefill Ring KV merges must use scratch_scope + copy_to ──────────
def _run_prefill_ring(*, T_q: int, S_kv: int, C: int):
"""Head-parallel prefill with Ring KV across C CUBEs."""
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
def _bench_fn(ctx):
configure_sfr_intercube_ring(
ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C,
)
dp_q = DPPolicy(cube="replicate", pe="replicate",
num_cubes=C, num_pes=1)
dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
pe="replicate", num_cubes=C, num_pes=1)
dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate",
pe="replicate", num_cubes=C, num_pes=1)
q = ctx.zeros((T_q, D_HEAD),
dtype=DTYPE, dp=dp_q, name=f"q_ring_{C}")
k = ctx.zeros((S_kv, D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"k_ring_{C}")
v = ctx.zeros((S_kv, D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"v_ring_{C}")
o = ctx.empty((T_q * C, D_HEAD),
dtype=DTYPE, dp=dp_o, name=f"o_ring_{C}")
ctx.launch(
f"gqa_prefill_scoped_{C}",
gqa_prefill_kernel,
q, k, v, o,
T_q, S_kv, D_HEAD, C,
_auto_dim_remap=False,
)
return run_bench(
topology=topo, bench_fn=_bench_fn,
device=resolve_device(None), engine_factory=_engine_factory,
)
def test_prefill_ring_step_merges_emit_copy_to_writeback():
"""ADR-0060 §5.5 + ADR-0063 §D3.1: each Ring KV step's online-softmax
merge must wrap its intermediates in ``scratch_scope`` and persist
the new running ``(m, ℓ, O)`` via ``tl.copy_to``.
For C=4: 3 ring-step merges (steps 1..C-1) × 3 handles (m, ℓ, O) per
merge ⇒ 9 ``copy`` entries per participating CUBE. Aggregated across
C CUBEs: ⇒ 36 ``copy`` entries.
Currently 0 because the kernel never calls ``tl.copy_to``.
"""
result = _run_prefill_ring(T_q=4, S_kv=16, C=4)
assert result.completion.ok, f"prefill ring failed: {result.completion}"
n_copy = _count(result.engine.op_log, "copy")
assert n_copy > 0, (
f"prefill ring kernel must emit copy_to writeback per merge step "
f"(ADR-0060 §5.5 + ADR-0063 §D3.1); got 0 ``copy`` entries"
)
# ── Scoped kernels still produce the same op_log shape (regression) ──
def test_decode_scoped_still_has_root_only_write():
"""ADR-0060 §A.2 root-only output must hold under the rewrite:
adding scratch_scope + copy_to should not change the reduce
topology; only the per-PE scratch usage. Single PE 0 writes O."""
result = _run_decode_sp(h_q=8, h_kv=1, P=8, S_kv=64)
assert result.completion.ok
n_writes = _count(result.engine.op_log, "dma_write")
assert n_writes == 1, (
f"scoped decode must still produce 1 dma_write (root-only); "
f"got {n_writes}"
)
def test_prefill_scoped_still_has_per_cube_distributed_output():
"""ADR-0060 §5.5 per-CUBE distributed output must hold under the
rewrite: scoped prefill still writes one O slice per CUBE."""
result = _run_prefill_ring(T_q=4, S_kv=16, C=4)
assert result.completion.ok
n_writes = _count(result.engine.op_log, "dma_write")
assert n_writes == 4, (
f"scoped prefill must write one O per CUBE (C=4 → 4 dma_write); "
f"got {n_writes}"
)
tests/attention/test_gqa_short_context.py
+260
View File
@@ -0,0 +1,260 @@
"""Phase 1 spec test for Phase D: short-context GQA kernels
(ADR-0060 §B "Items from the long/short context split", item B.split.2).
The long-context kernels (§5.2 decode / §5.5 prefill) shard each KV
head row-wise across all CUBEs. At short context (S_kv < 256K), that
shard is too thin to feed the engines and the cube-level collective
overhead dominates. The short-context kernels drop cube-SP entirely:
each CUBE owns ``kv_per_cube`` *whole* KV heads, no S_kv sharding across
CUBEs.
Design (per AskUserQuestion answers in this session):
- PE-parallel heads: P PEs split into ``kv_per_cube`` groups, each
group does PE-SP across (P/kv_per_cube) PEs for one owned head.
- scratch_scope + tl.copy_to discipline mirrors the long kernels.
- One short kernel handles kv_per_cube ∈ {1, 2, 4, 8} via a parameter.
Group layout on the 2×4 PE grid:
kv_per_cube=1, group=8 PEs (full 2×4): row chain + col bridge.
kv_per_cube=2, group=4 PEs (one row): row chain only.
kv_per_cube=4, group=2 PEs (adj cols): 1-step chain.
kv_per_cube=8, group=1 PE: no chain.
After chain reduce-to-group-root, the group's root PE writes its
owned head's output. No inter-CUBE reduce.
Phase 1 (this commit): tests only — production code lands in Phase 2.
All tests fail because the short kernels (``_gqa_decode_short.py`` and
``_gqa_prefill_short.py``) do not exist yet.
"""
from __future__ import annotations
from pathlib import Path
import pytest
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"
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 _count(op_log, name: str) -> int:
return sum(1 for r in op_log if r.op_name == name)
# ── Decode short-context kernel ──────────────────────────────────────
def _run_decode_short(*, h_q: int, h_kv: int, kv_per_cube: int,
C: int, P: int, S_kv: int):
"""Run the short-context decode kernel with PE-parallel heads.
Layout (after design iteration — see Phase D failure-recovery):
Q: (T_q, h_q·D_HEAD) replicated; kernel reshapes byte-conservingly.
K, V: (h_kv·S_kv, D_HEAD) head-stacked, ``cube=row_wise, pe=row_wise``
so each PE gets a contiguous (S_local, D_HEAD) chunk at its own
addressable shard (no partial reads needed).
O: replicated; each group root writes the full byte-conserving
(h_q·T_q, D_HEAD) result.
"""
from kernbench.benches._gqa_decode_short import gqa_decode_short_kernel # Phase 2
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=C, num_pes=P)
# Head-stacked KV with row_wise sharding so each PE's chunk is
# contiguous and exactly (S_local, D_HEAD) addressable.
dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
num_cubes=C, num_pes=P)
q = ctx.zeros((1, h_q * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"q_short_{kv_per_cube}_{C}")
k = ctx.zeros((h_kv * S_kv, D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"k_short_{kv_per_cube}_{C}")
v = ctx.zeros((h_kv * S_kv, D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"v_short_{kv_per_cube}_{C}")
o = ctx.empty((1, h_q * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"o_short_{kv_per_cube}_{C}")
ctx.launch(
f"gqa_decode_short_{kv_per_cube}_{C}",
gqa_decode_short_kernel,
q, k, v, o,
1, S_kv, h_q, h_kv, D_HEAD, C, P, kv_per_cube,
_auto_dim_remap=False,
)
return run_bench(
topology=topo, bench_fn=_bench_fn,
device=resolve_device(None), engine_factory=_engine_factory,
)
def test_short_decode_smoke_kv_per_cube_2_C_4():
"""ADR-0060 §B.split.2 headline: kv_per_cube=2, C=4 — each CUBE
owns 2 heads (8 heads / 4 CUBEs). PE-parallel heads splits the 8
PEs into 2 groups of 4, each group does PE-SP for one owned head.
Smoke: kernel completes."""
result = _run_decode_short(
h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
)
assert result.completion.ok, (
f"short decode kv_per_cube=2 must complete; got {result.completion}"
)
def test_short_decode_smoke_kv_per_cube_4_C_2():
"""kv_per_cube=4, C=2 — half the CUBEs participate, each owns 4
heads, 4 PE groups of 2 PEs each."""
result = _run_decode_short(
h_q=8, h_kv=8, kv_per_cube=4, C=2, P=8, S_kv=64,
)
assert result.completion.ok, (
f"short decode kv_per_cube=4 must complete; got {result.completion}"
)
def test_short_decode_smoke_kv_per_cube_8_C_1():
"""kv_per_cube=8, C=1 — all heads on one CUBE. 8 PE groups of 1 PE
each → no chain reduce, each PE writes its head's output."""
result = _run_decode_short(
h_q=8, h_kv=8, kv_per_cube=8, C=1, P=8, S_kv=64,
)
assert result.completion.ok, (
f"short decode kv_per_cube=8 must complete; got {result.completion}"
)
def test_short_decode_dma_write_count_equals_h_kv():
"""ADR-0060 §B.split.2: each owned head produces exactly one
output (no inter-CUBE reduce). Total dma_writes = h_kv across all
participating CUBEs and groups.
For h_kv=8, kv_per_cube=2, C=4: each CUBE writes 2 outputs →
4 × 2 = 8 dma_writes total.
"""
result = _run_decode_short(
h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
)
assert result.completion.ok
n_writes = _count(result.engine.op_log, "dma_write")
assert n_writes == 8, (
f"short decode: expected 8 dma_writes (h_kv); got {n_writes}"
)
def test_short_decode_no_inter_cube_traffic():
"""ADR-0060 §B.split.2: each head is fully owned by one CUBE → no
inter-CUBE reduce. The kernel must not invoke CUBE-level E/W IPCQ.
Today's long-context kernel at C=4 emits ~12 inter-CUBE ipcq_copy
via direction "E"/"W". The short kernel must emit zero of those,
keeping all IPCQ traffic on the ``intra_*`` directions.
"""
result = _run_decode_short(
h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
)
assert result.completion.ok
# Count IPCQ traffic that targeted the CUBE-level "E"/"W" directions.
inter_cube_ipcq = sum(
1 for r in result.engine.op_log
if r.op_name == "ipcq_copy"
and r.params.get("direction") in ("E", "W")
)
assert inter_cube_ipcq == 0, (
f"short decode must have no inter-CUBE E/W IPCQ; got {inter_cube_ipcq}"
)
# ── Prefill short-context kernel ─────────────────────────────────────
def _run_prefill_short(*, h_kv: int, kv_per_cube: int,
C: int, P: int, T_q: int, S_kv: int):
"""Run the short-context prefill kernel with PE-parallel heads.
Layout: same head-stacked K/V scheme as decode short, with Q/O
replicated and byte-conserving reshape inside the kernel.
"""
from kernbench.benches._gqa_prefill_short import gqa_prefill_short_kernel # Phase 2
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
def _bench_fn(ctx):
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
dp_q = DPPolicy(cube="replicate", pe="replicate",
num_cubes=C, num_pes=P)
dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
num_cubes=C, num_pes=P)
dp_o = DPPolicy(cube="replicate", pe="replicate",
num_cubes=C, num_pes=P)
q = ctx.zeros((T_q, h_kv * D_HEAD),
dtype=DTYPE, dp=dp_q, name=f"q_pre_short_{kv_per_cube}_{C}")
k = ctx.zeros((h_kv * S_kv, D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"k_pre_short_{kv_per_cube}_{C}")
v = ctx.zeros((h_kv * S_kv, D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"v_pre_short_{kv_per_cube}_{C}")
o = ctx.empty((T_q, h_kv * D_HEAD),
dtype=DTYPE, dp=dp_o, name=f"o_pre_short_{kv_per_cube}_{C}")
ctx.launch(
f"gqa_prefill_short_{kv_per_cube}_{C}",
gqa_prefill_short_kernel,
q, k, v, o,
T_q, S_kv, h_kv, D_HEAD, C, P, kv_per_cube,
_auto_dim_remap=False,
)
return run_bench(
topology=topo, bench_fn=_bench_fn,
device=resolve_device(None), engine_factory=_engine_factory,
)
def test_short_prefill_smoke_kv_per_cube_2_C_4():
"""Short prefill kv_per_cube=2, C=4. Smoke: completes."""
result = _run_prefill_short(
h_kv=8, kv_per_cube=2, C=4, P=8, T_q=4, S_kv=64,
)
assert result.completion.ok, (
f"short prefill kv_per_cube=2 must complete; got {result.completion}"
)
def test_short_prefill_no_ring_KV_traffic():
"""ADR-0060 §B.split.2: short prefill DOES NOT use Ring KV — each
CUBE owns its KV heads fully, no rotation. The kernel must not
emit any KV-rotation IPCQ traffic at the CUBE level.
"""
result = _run_prefill_short(
h_kv=8, kv_per_cube=2, C=4, P=8, T_q=4, S_kv=64,
)
assert result.completion.ok
inter_cube_ipcq = sum(
1 for r in result.engine.op_log
if r.op_name == "ipcq_copy"
and r.params.get("direction") in ("E", "W")
)
assert inter_cube_ipcq == 0, (
f"short prefill must have no Ring KV (no inter-CUBE E/W IPCQ); "
f"got {inter_cube_ipcq}"
)
tests/attention/test_mesh_kernels_rank_axis.py
-172
View File
@@ -1,172 +0,0 @@
"""Phase 1 spec test for ``rank_axis`` parameter on the two mesh kernels.
ADR-0059's mesh kernels currently hard-code ``rank = tl.program_id(axis=0)``,
which only works for single_user_* panels (rank == pe_id within cube).
For multi_user_* panels the ring is at the cube level — rank should be
``cube_id`` (axis=1), and the 7 non-rank-leader PEs in each cube should
not run the ring (they only hold KV replicas).
This test pins the desired ``rank_axis`` kwarg semantics:
rank_axis = 0 (default, single_user)
rank = tl.program_id(axis=0). Every PE in the cube runs the ring.
Existing behavior — no change.
rank_axis = 1 (multi_user)
if tl.program_id(axis=0) != 0: return. (7/8 PEs early-exit.)
rank = tl.program_id(axis=1).
Phase 1 expectation: tests fail today (kernels don't accept the kwarg).
Phase 2 lands the parameter on both kernels; tests turn green and the
multi_user_* diag harness clears its first send.
"""
from __future__ import annotations
from kernbench.common.ipcq_types import IpcqRecvCmd, IpcqSendCmd
from kernbench.common.pe_commands import GemmCmd
from kernbench.triton_emu.tl_context import TLContext, run_kernel
from kernbench.benches._attention_mesh_kv import attention_mesh_kv_kernel
from kernbench.benches._attention_mesh_mlo import attention_mesh_mlo_kernel
S_Q_PREFILL = 16
S_Q_DECODE = 1
S_KV_PER_RANK = 16
H_Q = 1
H_KV = 1
D_HEAD = 64
N_RANKS_MULTI = 4
PES_PER_CUBE = 8
Q_PTR = 0x10000
K_PTR = 0x20000
V_PTR = 0x30000
O_PTR = 0x40000
def _tl(pe_id: int, cube_id: int, num_pes: int, num_cubes: int) -> TLContext:
return TLContext(
pe_id=pe_id,
num_programs=num_pes,
cube_id=cube_id,
num_cubes=num_cubes,
dispatch_cycles=0,
scratch_base=0x80000,
scratch_size=1 << 20,
)
# ── Default rank_axis=0 backward-compat ──────────────────────────
def test_mlo_kernel_default_rank_axis_zero_emits_commands_on_all_pes():
"""rank_axis defaults to 0 → kernel uses pe_id as rank, runs on every
PE. Verify by running rank=3 (interior PE) in a single-cube 8-rank
setup and asserting at least one GEMM and at least one IPCQ send
are emitted (interior ranks send in both directions)."""
tl = _tl(pe_id=3, cube_id=0, num_pes=8, num_cubes=1)
run_kernel(
attention_mesh_mlo_kernel, tl,
Q_PTR, K_PTR, V_PTR, O_PTR,
S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, 8,
)
assert any(isinstance(c, GemmCmd) for c in tl.commands), \
"default rank_axis=0 must run the kernel (≥1 GEMM)"
assert any(isinstance(c, IpcqSendCmd) for c in tl.commands), \
"interior rank must emit ≥1 IpcqSendCmd"
def test_kv_kernel_default_rank_axis_zero_emits_commands_on_all_pes():
tl = _tl(pe_id=3, cube_id=0, num_pes=8, num_cubes=1)
run_kernel(
attention_mesh_kv_kernel, tl,
Q_PTR, K_PTR, V_PTR, O_PTR,
S_Q_PREFILL, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, 8,
)
assert any(isinstance(c, GemmCmd) for c in tl.commands)
assert any(isinstance(c, IpcqSendCmd) for c in tl.commands)
# ── rank_axis=1 multi_user semantics ─────────────────────────────
def test_mlo_kernel_rank_axis_one_gates_non_zero_pe_to_no_commands():
"""rank_axis=1 + pe_id != 0 → kernel must early-return; no GEMM,
no DMA, no IPCQ. The 7 non-rank-leader PEs in a multi_user cube
must stay completely silent so the cube-level SFR install isn't
asked to route sends from PEs that have no neighbors installed."""
tl = _tl(pe_id=2, cube_id=1, num_pes=PES_PER_CUBE, num_cubes=N_RANKS_MULTI)
run_kernel(
attention_mesh_mlo_kernel, tl,
Q_PTR, K_PTR, V_PTR, O_PTR,
S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, N_RANKS_MULTI,
rank_axis=1,
)
assert not any(isinstance(c, GemmCmd) for c in tl.commands), \
"pe_id=2 with rank_axis=1 must not emit GEMMs"
assert not any(isinstance(c, IpcqSendCmd) for c in tl.commands), \
"pe_id=2 with rank_axis=1 must not emit IpcqSendCmd"
assert not any(isinstance(c, IpcqRecvCmd) for c in tl.commands), \
"pe_id=2 with rank_axis=1 must not emit IpcqRecvCmd"
def test_kv_kernel_rank_axis_one_gates_non_zero_pe_to_no_commands():
tl = _tl(pe_id=2, cube_id=1, num_pes=PES_PER_CUBE, num_cubes=N_RANKS_MULTI)
run_kernel(
attention_mesh_kv_kernel, tl,
Q_PTR, K_PTR, V_PTR, O_PTR,
S_Q_PREFILL, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, N_RANKS_MULTI,
rank_axis=1,
)
assert not any(isinstance(c, GemmCmd) for c in tl.commands)
assert not any(isinstance(c, IpcqSendCmd) for c in tl.commands)
assert not any(isinstance(c, IpcqRecvCmd) for c in tl.commands)
def test_mlo_kernel_rank_axis_one_pe_zero_uses_cube_id_as_rank():
"""rank_axis=1 + pe_id == 0 → kernel runs the ring with rank=cube_id.
For cube_id=1 in a 4-cube ring, rank=1 is an interior rank: has_E=True
AND has_W=True → IPCQ sends emitted in both E and W directions.
"""
tl = _tl(pe_id=0, cube_id=1, num_pes=PES_PER_CUBE, num_cubes=N_RANKS_MULTI)
run_kernel(
attention_mesh_mlo_kernel, tl,
Q_PTR, K_PTR, V_PTR, O_PTR,
S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, N_RANKS_MULTI,
rank_axis=1,
)
sends = [c for c in tl.commands if isinstance(c, IpcqSendCmd)]
assert any(s.direction == "E" for s in sends), \
"cube_id=1 (interior) must emit ≥1 E-send"
assert any(s.direction == "W" for s in sends), \
"cube_id=1 (interior) must emit ≥1 W-send"
def test_kv_kernel_rank_axis_one_pe_zero_uses_cube_id_as_rank():
tl = _tl(pe_id=0, cube_id=1, num_pes=PES_PER_CUBE, num_cubes=N_RANKS_MULTI)
run_kernel(
attention_mesh_kv_kernel, tl,
Q_PTR, K_PTR, V_PTR, O_PTR,
S_Q_PREFILL, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, N_RANKS_MULTI,
rank_axis=1,
)
sends = [c for c in tl.commands if isinstance(c, IpcqSendCmd)]
assert any(s.direction == "E" for s in sends)
assert any(s.direction == "W" for s in sends)
def test_mlo_kernel_rank_axis_one_west_edge_cube_no_west_sends():
"""cube_id=0 (west edge) with rank_axis=1: rank=0, has_W=False → no
W-direction IPCQ sends. has_E=True → ≥1 E-direction send."""
tl = _tl(pe_id=0, cube_id=0, num_pes=PES_PER_CUBE, num_cubes=N_RANKS_MULTI)
run_kernel(
attention_mesh_mlo_kernel, tl,
Q_PTR, K_PTR, V_PTR, O_PTR,
S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD, N_RANKS_MULTI,
rank_axis=1,
)
sends = [c for c in tl.commands if isinstance(c, IpcqSendCmd)]
assert any(s.direction == "E" for s in sends), \
"west-edge cube_id=0 must still emit ≥1 E-send"
assert not any(s.direction == "W" for s in sends), \
"west-edge cube_id=0 must NOT emit any W-send (no W neighbor)"
tests/attention/test_mesh_mlo_2d_correctness.py
-142
View File
@@ -1,142 +0,0 @@
"""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}"
)
tests/attention/test_milestone_gqa_headline.py
+148
View File
@@ -0,0 +1,148 @@
"""Phase 1 spec test for P7: headline milestone-gqa bench with real GQA.
P7 wires the new ``_gqa_decode`` and ``_gqa_prefill`` kernels into a
new 4-panel milestone bench (independent from the existing
``milestone-gqa-llama70b`` which still covers the baseline kernels).
Real GQA (``h_q > h_kv`` with G=8) runs end-to-end through a
milestone-style sweep + sweep.json output.
Restriction in P7 first cut:
- C ≤ 4 (single-row inter-CUBE ring SFR; ADR-0060 §B leaves
multi-row rings for follow-on)
- Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP
headline left to follow-on)
- No figure renderers (defer to a separate cycle)
Panels (4 total):
single_user_prefill_gqa: C=1, T_q=4, S_kv=16 (no ring)
multi_user_prefill_gqa : C=4, T_q=4, S_kv=16 (Ring KV, 3 steps)
single_user_decode_gqa : C=1, P=8, h_q=8, h_kv=1, S_kv=64 (M-fold + intra-cube chain)
multi_user_decode_gqa : C=4, P=8, h_q=8, h_kv=1, S_kv=128 (M-fold + 2-level chain)
Phase 1 (this commit): tests only — bench module lands in Phase 2.
"""
from __future__ import annotations
import json
import kernbench.benches.milestone_gqa_headline as bench_mod # noqa: F401 (Phase 2)
from kernbench.benches.registry import resolve
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
BENCH_NAME = "milestone-gqa-headline"
PANELS = (
"single_user_prefill_gqa",
"multi_user_prefill_gqa",
"single_user_decode_gqa",
"multi_user_decode_gqa",
)
def _run_validation():
topo = resolve_topology("topology.yaml")
return run_bench(
topology=topo,
bench_fn=resolve(BENCH_NAME).run,
device=resolve_device(None),
engine_factory=lambda t, d: GraphEngine(
getattr(t, "topology_obj", t), enable_data=True,
),
)
def _sweep_json(monkeypatch) -> dict:
monkeypatch.setenv("GQA_HEADLINE_RUN", "1")
out = bench_mod._OUTPUT_DIR / "sweep.json"
if not out.exists():
result = _run_validation()
assert result.completion.ok, result.completion
assert out.exists(), f"missing {out}"
return json.loads(out.read_text())
# ── Registration ───────────────────────────────────────────────────────
def test_bench_registered():
spec = resolve(BENCH_NAME)
assert spec.name == BENCH_NAME
assert callable(spec.run)
assert spec.description.strip(), "description must be non-empty"
# ── Validation run completes ──────────────────────────────────────────
def test_validation_run_completes_ok(monkeypatch):
monkeypatch.setenv("GQA_HEADLINE_RUN", "1")
result = _run_validation()
assert result.completion.ok, (
f"headline validation run failed: {result.completion}"
)
# ── sweep.json shape ──────────────────────────────────────────────────
def test_sweep_json_has_four_panels(monkeypatch):
data = _sweep_json(monkeypatch)
assert set(data["panels"]) == set(PANELS), (
f"panels mismatch: expected {set(PANELS)}, got {set(data['panels'])}"
)
assert len(data["rows"]) == 4
assert {r["panel"] for r in data["rows"]} == set(PANELS)
def test_decode_panels_use_real_gqa(monkeypatch):
"""ADR-0060 §A.1 headline: decode panels must use h_q = G·h_kv with G>1."""
data = _sweep_json(monkeypatch)
cfg = data["config"]
assert cfg["h_q_decode"] > cfg["h_kv_decode"], (
f"decode must use real GQA (h_q > h_kv); got "
f"h_q={cfg['h_q_decode']}, h_kv={cfg['h_kv_decode']}"
)
# ── Per-panel architectural assertions ────────────────────────────────
def _row(rows, panel: str) -> dict:
for r in rows:
if r["panel"] == panel:
return r
raise AssertionError(f"missing row for panel {panel!r}")
def test_prefill_ring_panel_has_ipcq_traffic(monkeypatch):
"""multi_user_prefill_gqa uses Ring KV → IPCQ traffic > 0."""
data = _sweep_json(monkeypatch)
row = _row(data["rows"], "multi_user_prefill_gqa")
n_copy = row["op_log_summary"].get("ipcq_copy_count", 0)
assert n_copy > 0, (
f"multi_user_prefill_gqa must have Ring KV traffic; got "
f"ipcq_copy_count={n_copy}"
)
def test_decode_reduce_panel_writes_once(monkeypatch):
"""multi_user_decode_gqa: chain reduce-to-root → exactly 1 dma_write."""
data = _sweep_json(monkeypatch)
row = _row(data["rows"], "multi_user_decode_gqa")
n_writes = row["op_log_summary"]["dma_write_count"]
assert n_writes == 1, (
f"multi_user_decode_gqa root-only write: expected 1; got {n_writes}"
)
def test_prefill_panel_distributes_output(monkeypatch):
"""multi_user_prefill_gqa: per-CUBE distributed output → dma_write_count == C."""
data = _sweep_json(monkeypatch)
row = _row(data["rows"], "multi_user_prefill_gqa")
n_writes = row["op_log_summary"]["dma_write_count"]
assert n_writes == 4, (
f"multi_user_prefill_gqa per-CUBE distributed: expected 4; got {n_writes}"
)
tests/attention/test_milestone_gqa_llama70b.py
-222
View File
@@ -1,222 +0,0 @@
"""Phase 1 spec test for ``milestone-gqa-llama70b`` bench (sub-cycle 4a, all 4 panels).
ADR-0057 (Proposed) defines an eval bench that drives both attention kernels
(ADR-0055 ring-K/V, ADR-0056 allreduce-mlo) and emits per-panel op_log
summaries into ``src/kernbench/benches/1H_milestone_output/gqa/sweep.json``.
v1 (sub-cycle 4a) covers ALL FOUR panels:
Panel name in JSON / test Study label SFR install used
─────────────────────────────────────────────────────────────────────────────
single_user_prefill TL configure_sfr_intracube_pe_ring
multi_user_prefill TR configure_sfr_intercube_multisip
single_user_decode BL configure_sfr_intracube_pe_ring
multi_user_decode BR configure_sfr_intercube_multisip
single_user_* panels became runnable after sub-cycle 4-pre delivered the
new SFR install function (ADR-0058).
In Phase 1 the bench module does not exist; pytest collection fails with
``ModuleNotFoundError``. Once Phase 2 lands the bench module, every
assertion below must pass.
Assertions:
- Bench is registered as ``milestone-gqa-llama70b``.
- A validation run (``GQA_VALIDATION=1``) completes ok via run_bench.
- sweep.json conforms to ADR-0057 D7 (v1 schema).
- All four panel rows present with sane op_log summaries.
- Both decode panels have gemm_count = 2 × n_ranks (one-shot per rank).
- Both prefill panels have gemm_count = 2 × n_ranks² (per-step GEMMs).
"""
from __future__ import annotations
import json
from kernbench.benches.registry import resolve
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
# Production module (Phase 2 deliverable; absent in Phase 1).
import kernbench.benches.milestone_gqa_llama70b as gqa_bench
BENCH_NAME = "milestone-gqa-llama70b"
PANELS_V1 = (
"single_user_prefill",
"multi_user_prefill",
"single_user_decode",
"multi_user_decode",
)
SINGLE_USER_PANELS = ("single_user_prefill", "single_user_decode")
MULTI_USER_PANELS = ("multi_user_prefill", "multi_user_decode")
PREFILL_PANELS = ("single_user_prefill", "multi_user_prefill")
DECODE_PANELS = ("single_user_decode", "multi_user_decode")
def _run_validation():
"""Drive the bench through run_bench at validation scale."""
topo = resolve_topology("topology.yaml")
return run_bench(
topology=topo,
bench_fn=resolve(BENCH_NAME).run,
device=resolve_device(None),
engine_factory=lambda t, d: GraphEngine(
getattr(t, "topology_obj", t), enable_data=True,
),
)
# ── Registration ────────────────────────────────────────────────────────
def test_bench_registered():
spec = resolve(BENCH_NAME)
assert spec.name == BENCH_NAME
assert callable(spec.run)
assert spec.description.strip(), "description must be non-empty"
# ── Validation run end-to-end ────────────────────────────────────────────
def test_validation_run_completes_ok(monkeypatch):
monkeypatch.setenv("GQA_VALIDATION", "1")
result = _run_validation()
assert result.completion.ok, (
f"validation run failed: {result.completion}"
)
# ── JSON shape (ADR-0057 D7 amended for 4 panels) ──────────────────────
def _sweep_json(monkeypatch) -> dict:
"""Run the bench (if needed) and return the parsed sweep.json."""
monkeypatch.setenv("GQA_VALIDATION", "1")
out = gqa_bench._OUTPUT_DIR / "sweep.json"
if not out.exists():
result = _run_validation()
assert result.completion.ok, result.completion
assert out.exists(), f"missing {out}"
return json.loads(out.read_text())
def test_sweep_json_has_v1_schema(monkeypatch):
data = _sweep_json(monkeypatch)
assert data["version"] == 1
assert data["validation_scale"] is True
assert isinstance(data["panels"], list)
assert isinstance(data["config"], dict)
assert isinstance(data["rows"], list)
def test_sweep_json_panels_are_all_four(monkeypatch):
"""v1 covers all four panels — single_user_{prefill,decode} +
multi_user_{prefill,decode}. Q/cube sweep deferred to 4b."""
data = _sweep_json(monkeypatch)
assert set(data["panels"]) == set(PANELS_V1)
def test_sweep_json_config_matches_adr0057_d4(monkeypatch):
"""Validation-scale config per ADR-0057 D4 (amended for 4 panels + scratch budget).
S_q_prefill and S_kv_per_rank are deliberately small (16 each) so the
simulator's 1 MB per-PE TCM kernel scratch (topology.yaml
``pe_tcm.kernel_scratch_mb: 1``) is not exhausted by the
bump-allocated handle outputs of softmax/exp/dot/sum chains over
n_ranks ring steps. Headline-scale runs in 4c will lift these into a
config-driven sweep.
"""
data = _sweep_json(monkeypatch)
cfg = data["config"]
assert cfg["S_q_prefill"] == 16
assert cfg["S_kv_per_rank"] == 16
# v1 uses h_q == h_kv == 1 to avoid ADR-0055 D3's GQA broadcast view
# (which is symbolic and does not survive MemoryStore's nbytes check
# under simulator data execution). Real GQA (h_q > h_kv) is deferred
# to sub-cycle 4c (headline scale).
assert cfg["h_q"] == 1
assert cfg["h_kv"] == 1
assert cfg["d_head"] == 64
# single_user_* uses the 8 PEs in one cube as ring ranks.
assert cfg["n_ranks_single_user"] == 8
# multi_user_* uses cube-level ring; validation uses 4 cubes.
assert cfg["n_ranks_multi_user"] == 4
def test_sweep_json_has_one_row_per_panel(monkeypatch):
data = _sweep_json(monkeypatch)
assert len(data["rows"]) == 4
panels_in_rows = {r["panel"] for r in data["rows"]}
assert panels_in_rows == set(PANELS_V1)
# ── Per-row op_log summary sanity (ADR-0057 D7) ─────────────────────────
def _row(rows: list, panel: str) -> dict:
matches = [r for r in rows if r["panel"] == panel]
assert len(matches) == 1, f"expected exactly one {panel} row; got {len(matches)}"
return matches[0]
def _assert_sane_summary(row: dict) -> None:
s = row["op_log_summary"]
panel = row["panel"]
assert s["gemm_count"] > 0, f"{panel} must run GEMMs"
assert s["ipcq_send_count"] > 0, f"{panel} must send (ring/allreduce phase)"
assert s["ipcq_recv_count"] > 0, f"{panel} must recv"
assert s["dma_read_count"] >= 3, f"{panel}: Q + K + V loads"
assert s["dma_write_count"] >= 1, f"{panel}: final O store"
def test_single_user_prefill_row_has_sane_op_log_summary(monkeypatch):
data = _sweep_json(monkeypatch)
_assert_sane_summary(_row(data["rows"], "single_user_prefill"))
def test_multi_user_prefill_row_has_sane_op_log_summary(monkeypatch):
data = _sweep_json(monkeypatch)
_assert_sane_summary(_row(data["rows"], "multi_user_prefill"))
def test_single_user_decode_row_has_sane_op_log_summary(monkeypatch):
data = _sweep_json(monkeypatch)
_assert_sane_summary(_row(data["rows"], "single_user_decode"))
def test_multi_user_decode_row_has_sane_op_log_summary(monkeypatch):
data = _sweep_json(monkeypatch)
_assert_sane_summary(_row(data["rows"], "multi_user_decode"))
# ── Architectural invariant: decode = one-shot per rank ─────────────────
def test_single_user_decode_gemm_count_is_exactly_2_per_rank(monkeypatch):
"""ADR-0056 D3: decode kernel does ONE local partial-attention pass per
rank → exactly 2 GEMMs per rank (Q·K^T + S·V). With n_ranks ranks the
total = 2 × n_ranks. This distinguishes decode from prefill where each
ring step has 2 GEMMs and the total scales as 2 × n_ranks²."""
data = _sweep_json(monkeypatch)
row = _row(data["rows"], "single_user_decode")
n_ranks = row["n_ranks"]
assert row["op_log_summary"]["gemm_count"] == 2 * n_ranks, (
f"single_user_decode gemm_count must be 2 × n_ranks = {2 * n_ranks}; "
f"got {row['op_log_summary']['gemm_count']}"
)
def test_multi_user_decode_gemm_count_is_exactly_2_per_rank(monkeypatch):
"""Same one-shot invariant as single_user_decode — the kernel is the
same; what differs is who the ranks are (cubes vs PEs)."""
data = _sweep_json(monkeypatch)
row = _row(data["rows"], "multi_user_decode")
n_ranks = row["n_ranks"]
assert row["op_log_summary"]["gemm_count"] == 2 * n_ranks, (
f"multi_user_decode gemm_count must be 2 × n_ranks = {2 * n_ranks}; "
f"got {row['op_log_summary']['gemm_count']}"
)