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

gqa: fold decode_long into the Case 4 kernel; drop 1D-chain dead code

Browse Source 
Option B fold: the lrab math previously in
``_gqa_attention_decode_long.py`` (used only as a thin re-export by
the Case 4 wrapper and as the import source for ``_merge_running``
in Cases 1 / 3) now lives inline in the Case 4 kernel file. ``sub_w``
is hardcoded to 4 (the 4×2 cube sub-mesh geometry; root cube 6);
the legacy ``sub_w=0`` 1D-chain backward-compat path is removed
along with its dedicated tests — the dropped ``single_user_*`` /
``multi_user_*`` panels that exercised it are already gone.

Production changes:
  - inline lrab math into _gqa_attention_decode_long_ctx_cube_sp_pe_sp.py
    (drops sub_w param; _ROOT_CUBE=6 baked in)
  - inline _merge_running into Cases 1 (cube_sp_pe_tp) and 3
    (cube_repl_pe_sp) so they no longer depend on the deleted file
  - delete src/kernbench/benches/_gqa_attention_decode_long.py
  - remove dead _run_decode_panel / _DECODE_* constants /
    decode-side _PANEL_DISPATCH / _make_bench_fn decode branch
    from milestone_gqa_headline.py (only the prefill panel remains)
  - update _gqa_attention_decode_opt2.py docstring reference

Test changes:
  - delete 7 legacy test_gqa_*.py files that pre-dated the 4-cases
    architectural split (coverage now subsumed by the 16 4-cases tests)
  - remove test_opt2_matches_opt3_data_mode + _run_decode_data helper
    from test_gqa_decode_opt2.py (the parity check required sub_w=0
    which no longer exists; opt2's other 4 tests preserved)

20/20 tests pass (16 4-cases + 4 opt2 smoke/dispatch).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-06-15 16:32:37 -07:00
parent ddee28a499
commit 756680f4e6
14 changed files with 315 additions and 1810 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/attention/test_gqa_decode.py
-143
View File
@@ -1,143 +0,0 @@
"""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_attention_decode_long import gqa_attention_decode_long_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_attention_decode_long_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_long_2d_reduce.py
-191
View File
@@ -1,191 +0,0 @@
"""Tests for lrab-adapted center-root inter-CUBE reduce in decode_long.
Verifies the ADR-0060 §4.2 prescribed Level-1 collective for the
single-KV-group LLaMA-3.1-70B target (C=8 over a 2×4 sub-mesh, root at
the geometric center CUBE).
Activation contract:
- ``sub_w == 0`` (default; omitted from launch args) → existing 1D
inter-CUBE chain that converges at CUBE 0. Byte-for-byte unchanged.
- ``sub_w >= 2`` with ``sub_h = C // sub_w >= 2`` → lrab-adapted
center-root mesh reduce. Root cube is
``(sub_h//2)*sub_w + (sub_w//2)``.
Degenerate (``sub_h < 2`` or ``sub_w < 2``) combinations are rejected at
launch time per ADR-0060 §4.2 + CLAUDE.md "Simplicity First": those
configurations belong to the 1D-chain code path, not lrab.
Phase 1 (this commit): tests only — production code lands in Phase 2.
T1, T2, T4 fail today (kernel signature does not accept ``sub_w``).
T3 passes today as the backward-compat anchor for the existing
1D-chain path.
"""
from __future__ import annotations
import re
from pathlib import Path
import pytest
from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_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"
D_HEAD = 64
DTYPE = "f16"
_CUBE_RE = re.compile(r"\bcube(\d+)\b")
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 _dma_write_cubes(op_log) -> list[int]:
"""Return the list of CUBE ids that emitted a ``dma_write``.
Decode's final-store path uses one ``tl.store`` from the root rank.
Multiple HBM-channel-level write records may correspond to the same
logical store; this just reports the source CUBE for each.
"""
cubes: list[int] = []
for r in op_log:
if r.op_name != "dma_write":
continue
m = _CUBE_RE.search(r.component_id)
if m is not None:
cubes.append(int(m.group(1)))
return cubes
def _run_decode_long(
*, C: int, P: int, S_kv: int, sub_w: int | None,
):
"""Drive a decode_long launch. ``sub_w=None`` means omit the arg
entirely (exercises the kernel's current default behaviour)."""
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" if C > 1 else "replicate",
pe="row_wise", num_cubes=C, num_pes=P)
T_q = 1
h_q = 8
h_kv = 1
ctx.zeros((T_q, h_q * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"q_dec_2d_c{C}")
k = ctx.zeros((S_kv, h_kv * D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"k_dec_2d_c{C}")
v = ctx.zeros((S_kv, h_kv * D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"v_dec_2d_c{C}")
o = ctx.empty((T_q, h_q * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"o_dec_2d_c{C}")
q = ctx.zeros((T_q, h_q * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"q_dec_2d_c{C}_2")
launch_args = [q, k, v, o, T_q, S_kv, h_q, h_kv, D_HEAD, C, P]
if sub_w is not None:
launch_args.append(sub_w)
ctx.launch(
f"gqa_decode_long_2d_c{C}_sw{sub_w}",
gqa_attention_decode_long_kernel,
*launch_args,
_auto_dim_remap=False,
)
return run_bench(
topology=topo, bench_fn=_bench_fn,
device=resolve_device(None), engine_factory=_engine_factory,
)
# ── T1: C=8 lrab-adapted center-root completes ───────────────────────
def test_decode_2d_sub_w4_sub_h2_completes():
"""ADR-0060 §4.2 prescribed center-root mesh reduce at the
milestone target: C=8 over a 2×4 sub-mesh (sub_w=4, sub_h=2),
P=8 PEs per CUBE. With zero inputs, output is trivially zero;
the test guards that the kernel reaches completion without
deadlock or scratch overflow.
"""
result = _run_decode_long(C=8, P=8, S_kv=2048, sub_w=4)
assert result.completion.ok, (
f"decode at C=8, sub_w=4 (lrab-adapted) must complete; "
f"got {result.completion}"
)
# ── T2: root lives at the geometric center CUBE (cube 6) ─────────────
def test_decode_2d_sub_w4_sub_h2_root_at_center_cube_6():
"""For sub_w=4, sub_h=2: root_col=2, root_row=1, root_cube=6.
Only the root CUBE issues the final ``tl.store`` — every other
rank short-circuits. The dma_write records must come exclusively
from CUBE 6.
"""
result = _run_decode_long(C=8, P=8, S_kv=2048, sub_w=4)
assert result.completion.ok
cubes = _dma_write_cubes(result.engine.op_log)
assert cubes, "expected at least one dma_write for the final O store"
distinct = set(cubes)
assert distinct == {6}, (
f"final dma_write must come exclusively from CUBE 6 (sub_w=4, "
f"sub_h=2 root); got cubes={sorted(distinct)}"
)
# ── T3: backward-compat — sub_w omitted → existing 1D chain at C=4 ───
def test_decode_2d_backward_compat_sub_w0_default():
"""When ``sub_w`` is omitted from the launch, the kernel must
behave identically to today: 1D inter-CUBE chain along W,
converging at CUBE 0. This serves as a regression anchor — Phase 2
must not change the existing C=4 panel behaviour.
Passes today as well as after Phase 2.
"""
result = _run_decode_long(C=4, P=8, S_kv=2048, sub_w=None)
assert result.completion.ok, (
f"decode at C=4 with default sub_w (1D chain) must complete; "
f"got {result.completion}"
)
cubes = _dma_write_cubes(result.engine.op_log)
assert cubes, "expected at least one dma_write for the final O store"
distinct = set(cubes)
assert distinct == {0}, (
f"1D-chain root must be CUBE 0; got cubes={sorted(distinct)}"
)
# ── T4: degenerate sub_w configurations are rejected ─────────────────
def test_decode_2d_invalid_sub_w_rejected_when_sub_h_lt_2():
"""ADR-0060 §4.2 + CLAUDE.md "Simplicity First": only sub-meshes
with both sub_w >= 2 and sub_h >= 2 use lrab. C=4 with sub_w=4
would give sub_h=1 (degenerate; Phase 2 of lrab becomes a no-op).
Callers must use the 1D-chain path (sub_w=0/omitted) for that case.
The kernel must reject the degenerate combination with a clear
error rather than silently producing a 1D-chain result at the
wrong root location.
"""
with pytest.raises((ValueError, AssertionError), match=r"sub_[wh]"):
_run_decode_long(C=4, P=8, S_kv=2048, sub_w=4)
tests/attention/test_gqa_decode_mc.py
-182
View File
@@ -1,182 +0,0 @@
"""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_attention_decode_long import gqa_attention_decode_long_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_attention_decode_long_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_opt2.py
-52
View File
@@ -181,55 +181,3 @@ def test_opt2_bench_completes_oplog_mode():
assert result.completion.ok, f"opt2 decode failed: {result.completion}"
# ── opt2 ↔ opt3 numeric parity in data mode (ADR-0065 N4) ─────────────
def _run_decode_data(kernel, name, q_d, k_d, v_d, *, S_kv=16):
"""Run a decode kernel (C=P=1) in data mode with seeded Q/K/V; return the
HBM output array."""
from pathlib import Path
import numpy as np
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.data_executor import DataExecutor
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import resolve_topology
topo = resolve_topology(str(Path(__file__).resolve().parents[2] / "topology.yaml"))
cap: dict = {}
def bench_fn(ctx):
dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1)
q = ctx.zeros((1, 8 * D), dtype="f16", dp=dp, name=name + "q")
k = ctx.zeros((S_kv, D), dtype="f16", dp=dp, name=name + "k")
v = ctx.zeros((S_kv, D), dtype="f16", dp=dp, name=name + "v")
o = ctx.empty((1, 8 * D), dtype="f16", dp=dp, name=name + "o")
q.copy_(ctx.from_numpy(q_d)); k.copy_(ctx.from_numpy(k_d)); v.copy_(ctx.from_numpy(v_d))
cap["o"] = o
ctx.launch(name, kernel, q, k, v, o, 1, S_kv, 8, 1, D, 1, 1, _auto_dim_remap=False)
r = run_bench(topology=topo, bench_fn=bench_fn, device=resolve_device(None),
engine_factory=lambda t, d: GraphEngine(getattr(t, "topology_obj", t),
enable_data=True))
assert r.completion.ok
DataExecutor(r.engine.op_log, r.engine.memory_store).run()
o = cap["o"]
return r.engine.memory_store.read("hbm", o._handle.va_base, shape=(1, 8 * D), dtype="f16")
def test_opt2_matches_opt3_data_mode():
import numpy as np
from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel
from kernbench.benches._gqa_attention_decode_opt2 import gqa_attention_decode_opt2_kernel
rng = np.random.default_rng(0)
q_d = (0.1 * rng.standard_normal((1, 8 * D))).astype(np.float16)
k_d = (0.1 * rng.standard_normal((16, D))).astype(np.float16)
v_d = (0.1 * rng.standard_normal((16, D))).astype(np.float16)
o3 = _run_decode_data(gqa_attention_decode_long_kernel, "p3", q_d, k_d, v_d)
o2 = _run_decode_data(gqa_attention_decode_opt2_kernel, "p2", q_d, k_d, v_d)
assert np.allclose(np.asarray(o2, np.float32), np.asarray(o3, np.float32),
rtol=5e-2, atol=5e-2), f"opt2 {np.asarray(o2).ravel()[:4]} vs opt3 {np.asarray(o3).ravel()[:4]}"
tests/attention/test_gqa_decode_sp.py
-153
View File
@@ -1,153 +0,0 @@
"""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_attention_decode_long import gqa_attention_decode_long_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_attention_decode_long_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
@@ -1,157 +0,0 @@
"""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_attention_decode_long import gqa_attention_decode_long_kernel # noqa: F401
from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_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_attention_decode_long_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_attention_prefill_long_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_scoped_writeback.py
-198
View File
@@ -1,198 +0,0 @@
"""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_attention_decode_long import gqa_attention_decode_long_kernel # noqa: F401
from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_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_attention_decode_long_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_attention_prefill_long_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_tile_sweep.py
-290
View File
@@ -1,290 +0,0 @@
"""Phase 1 spec test for P3b: S_kv tile sweep in the GQA kernels
(ADR-0063 §A.2 + ADR-0060 §B "long/short context split").
The local one-shot partial in three kernels — ``_gqa_attention_decode_long.py``,
``_gqa_attention_decode_short.py``, ``_gqa_attention_prefill_short.py`` — loads its rank's
entire ``(d_head, S_local)`` / ``(S_local, d_head)`` KV slice in one
shot, so per-rank scratch grows linearly with ``S_local``. The
``_attention_*`` baselines hit a TCM ceiling once ``S_local`` exceeds
the 1 MiB pool divided by the per-rank intermediate footprint.
P3b replaces the one-shot partial with a tile sweep:
- Module-level ``TILE_S_KV = 1024`` per kernel file.
- Tile 0 establishes the persistent ``(m_local, l_local, O_local)``.
- Tiles 1..n_tiles-1 wrap their intermediates (K_T tile, V tile,
scores, centered, exp_scores, partials) in ``tl.scratch_scope()``
and persist the merged ``(m, ℓ, O)`` to the persistent handles via
``tl.copy_to`` (ADR-0063 §D3 / §D3.1).
When ``S_local ≤ TILE_S_KV`` the loop body never runs and the op_log
is structurally identical to today's one-shot path — every existing
validation-scale test continues to pass.
``_gqa_attention_prefill_long.py`` is **out of scope** for P3b. Its inner step
is IPCQ partial-recv (Ring KV rotation), not an HBM load; tiling it
intersects the ring-step structure and is a separate phase.
Phase 1 (this commit): tests only — production code lands in Phase 2.
The four multi-tile tests fail today because the kernels never call
``copy_to`` outside the chain-reduce merges (so isolating to a
chain-free config gives ``copy_to == 0`` today). The 128K test fails
today with a ``TLContext`` scratch overflow.
"""
from __future__ import annotations
from pathlib import Path
from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel # noqa: F401
from kernbench.benches._gqa_attention_decode_short import gqa_attention_decode_short_kernel # noqa: F401
from kernbench.benches._gqa_attention_prefill_short import gqa_attention_prefill_short_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"
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_long runner (chain-free configs isolate tile-sweep behavior) ──
def _run_decode_long(*, C: int, P: int, S_kv: int, h_q: int = 1, h_kv: int = 1):
"""Run the long-context decode kernel.
For tile-sweep isolation, callers pass C=1, P=1 — this leaves both
chain-reduce levels inactive so any ``copy_to`` in op_log comes
solely from the tile-merge body.
"""
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" if C > 1 else "replicate",
pe="row_wise" if P > 1 else "replicate",
num_cubes=C, num_pes=P)
q = ctx.zeros((1, h_q * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"q_tl_c{C}_p{P}_s{S_kv}")
k = ctx.zeros((S_kv, h_kv * D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"k_tl_c{C}_p{P}_s{S_kv}")
v = ctx.zeros((S_kv, h_kv * D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"v_tl_c{C}_p{P}_s{S_kv}")
o = ctx.empty((1, h_q * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"o_tl_c{C}_p{P}_s{S_kv}")
ctx.launch(
f"gqa_decode_long_tile_c{C}_p{P}_s{S_kv}",
gqa_attention_decode_long_kernel,
q, k, v, o,
1, 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,
)
# ── decode_short / prefill_short runners ─────────────────────────────
def _run_decode_short(*, kv_per_cube: int, C: int, P: int, S_kv: int,
h_q: int = 8, h_kv: int = 8):
"""For tile-sweep isolation: kv_per_cube=P → group_size=1 (no 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=C, num_pes=P)
dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
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_dsh_s{S_kv}")
k = ctx.zeros((h_kv * S_kv, D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"k_dsh_s{S_kv}")
v = ctx.zeros((h_kv * S_kv, D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"v_dsh_s{S_kv}")
o = ctx.empty((1, h_q * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"o_dsh_s{S_kv}")
ctx.launch(
f"gqa_decode_short_tile_s{S_kv}",
gqa_attention_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 _run_prefill_short(*, kv_per_cube: int, C: int, P: int,
T_q: int, S_kv: int, h_kv: int = 8):
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" if C > 1 else "replicate",
pe="row_wise", num_cubes=C, num_pes=P)
q = ctx.zeros((T_q, h_kv * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"q_psh_s{S_kv}")
k = ctx.zeros((h_kv * S_kv, D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"k_psh_s{S_kv}")
v = ctx.zeros((h_kv * S_kv, D_HEAD),
dtype=DTYPE, dp=dp_kv, name=f"v_psh_s{S_kv}")
o = ctx.empty((T_q, h_kv * D_HEAD),
dtype=DTYPE, dp=dp_full, name=f"o_psh_s{S_kv}")
ctx.launch(
f"gqa_prefill_short_tile_s{S_kv}",
gqa_attention_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,
)
# ── T1: single-tile path is op_log-stable (regression guard) ─────────
def test_decode_long_single_tile_path_unchanged():
"""ADR-0063 §A.2: when S_local <= TILE_S_KV the tile-sweep loop
body never runs and op_log is structurally identical to today's
one-shot path.
Passes today (one-shot path) AND after Phase 2 (n_tiles=1 skips
the merge loop). The witness: with chain reduce inactive (C=1,
P=1), there must be zero ``copy_to`` entries — neither today nor
after Phase 2 should the kernel emit tile-merge writebacks here.
"""
result = _run_decode_long(C=1, P=1, S_kv=64)
assert result.completion.ok, (
f"single-tile decode_long must complete; got {result.completion}"
)
n_copy = _count(result.engine.op_log, "copy")
assert n_copy == 0, (
f"S_local <= TILE_S_KV must not emit tile-merge copy_to; got {n_copy}"
)
# ── T2: decode_long multi-tile emits tile-merge copy_to ──────────────
def test_decode_long_multi_tile_emits_copy_to_merges():
"""ADR-0063 §A.2 + §D3.1: when S_local > TILE_S_KV the kernel must
wrap each subsequent tile's intermediates in ``scratch_scope`` and
persist the merged ``(m, ℓ, O)`` via ``tl.copy_to``.
Chain-free config (C=1, P=1) isolates the tile-merge body — any
``copy_to`` in op_log comes from the tile sweep, not chain reduce.
S_kv=2048, C=1, P=1 → S_local=2048 → 2 tiles → 1 merge × 3 handles
(m, ℓ, O) ⇒ 3 ``copy`` entries.
Currently 0 because the kernel performs a one-shot partial.
"""
result = _run_decode_long(C=1, P=1, S_kv=2048)
assert result.completion.ok, (
f"multi-tile decode_long must complete; got {result.completion}"
)
n_copy = _count(result.engine.op_log, "copy")
assert n_copy >= 3, (
f"decode_long multi-tile must emit >= 3 copy_to entries "
f"(1 merge × 3 handles); got {n_copy}"
)
# ── T3: decode_short multi-tile emits tile-merge copy_to ─────────────
def test_decode_short_multi_tile_emits_copy_to_merges():
"""Same property as T2 for the short decode kernel.
kv_per_cube=8, P=8, C=1 → group_size=1 (no chain reduce); S_kv=2048
→ S_local=2048 → 2 tiles per PE → 3 ``copy`` entries per PE × 8 PEs
= 24 total.
Currently 0 because the kernel performs a one-shot partial.
"""
result = _run_decode_short(kv_per_cube=8, C=1, P=8, S_kv=2048)
assert result.completion.ok, (
f"multi-tile decode_short must complete; got {result.completion}"
)
n_copy = _count(result.engine.op_log, "copy")
assert n_copy >= 3 * 8, (
f"decode_short multi-tile must emit >= 24 copy_to entries "
f"(1 merge × 3 handles × 8 PEs); got {n_copy}"
)
# ── T4: prefill_short multi-tile emits tile-merge copy_to ────────────
def test_prefill_short_multi_tile_emits_copy_to_merges():
"""Same property as T3 for the short prefill kernel.
kv_per_cube=8, P=8, C=1, T_q=4, S_kv=2048 → group_size=1 (no chain),
S_local=2048 → 2 tiles per PE → 3 ``copy`` entries per PE × 8 PEs
= 24 total.
Currently 0 because the kernel performs a one-shot partial.
"""
result = _run_prefill_short(
kv_per_cube=8, C=1, P=8, T_q=4, S_kv=2048,
)
assert result.completion.ok, (
f"multi-tile prefill_short must complete; got {result.completion}"
)
n_copy = _count(result.engine.op_log, "copy")
assert n_copy >= 3 * 8, (
f"prefill_short multi-tile must emit >= 24 copy_to entries "
f"(1 merge × 3 handles × 8 PEs); got {n_copy}"
)
# ── T5: decode_long at S_kv=256K completes (TCM ceiling lifted) ──────
def test_decode_long_context_256k_completes():
"""ADR-0063 §A.2 (test req 3): headline ceiling-lift. C=1 P=8
decode at S_kv=256K → S_local=32K. Today the per-rank score stack
(3 ×M·S_local·2 bytes, M=G·T_q=8) is ~1.5 MB → exceeds the 1 MiB
scratch pool → TLContext scratch overflow. After Phase 2 the tile
sweep bounds per-tile score stack at 3·M·TILE_S_KV·2 ≈ 48 KB and
the run completes regardless of S_kv.
"""
result = _run_decode_long(C=1, P=8, S_kv=262144, h_q=8, h_kv=1)
assert result.completion.ok, (
f"decode_long at S_kv=256K must complete after tile sweep; "
f"got {result.completion}"
)