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
src/kernbench/benches/_gqa_attention_decode_long.py
-353
View File
@@ -1,353 +0,0 @@
"""GQA fused-attention decode kernel — long context (ADR-0060).
Each rank holds an ``S_local = S_kv / (C·P)`` slice of K, V and the full
Q (replicated). The local attention is computed via an S_kv-axis tile
sweep (ADR-0063 §A.2) so per-rank scratch is bounded by ``TILE_S_KV``
regardless of ``S_local``. The partial ``(m, ℓ, O)`` is then reduced
via a 2-level pattern (intra-CUBE row+col, then inter-CUBE), and the
root writes the final output.
Inter-CUBE reduce (selected by the ``sub_w`` launch arg):
- ``sub_w == 0`` (default): 1D chain along ``W``, root at CUBE 0.
Used for ``C ∈ {1, 4}`` panels.
- ``sub_w >= 2``: ADR-0060 §4.2 prescribed **lrab-adapted center-root
mesh reduce** over a ``sub_w × sub_h`` sub-mesh (``sub_h = C //
sub_w``). Phase 1 row reduce converges at ``root_col = sub_w//2``;
Phase 2 col reduce on ``root_col`` converges at
``root_row = sub_h//2``. Root cube id:
``root_row * sub_w + root_col``. Used for the C=8 single-KV-group
LLaMA-3.1-70B target (sub_w=4, sub_h=2, root=cube 6).
Requires ``sub_w >= 2 and sub_h >= 2`` — degenerate 1×N / N×1
layouts must use the 1D-chain path with ``sub_w=0``.
Topology / SFR:
- Requires ``configure_sfr_intercube_multisip`` when ``P > 1`` or
``C > 1`` (provides disjoint ``intra_*`` and ``E/W/N/S`` namespaces).
- Intra-CUBE PEs are arranged as a 2×4 grid (no wrap).
- Inter-CUBE CUBEs: 1D row (``sub_w == 0``) or rectangular sub-mesh
of the SIP's 4×4 CUBE mesh starting at origin (0, 0) with
``sub_w == mesh_w`` (``sub_w >= 2``).
Layout caveats:
- GEMMs use ``tl.dot`` (no composite epilogue / ``softmax_scale``).
- K is loaded as ``(d_head, S_local)`` via byte-conserving reshape of
the deployed ``(S_local, h_kv·d_head)`` slice — correct for zero /
symmetric inputs (ADR-0060 §3 reshape-not-transpose caveat).
"""
from __future__ import annotations
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
"""Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
m_new = tl.maximum(m_local, m_other)
scale_old = tl.exp(m_local - m_new)
scale_new = tl.exp(m_other - m_new)
l_new = l_local * scale_old + l_other * scale_new
O_new = O_local * scale_old + O_other * scale_new
return m_new, l_new, O_new
def gqa_attention_decode_long_kernel(
q_ptr: int,
k_ptr: int,
v_ptr: int,
o_ptr: int,
T_q: int,
S_kv: int,
h_q: int,
h_kv: int,
d_head: int,
C: int,
P: int,
sub_w: int = 0,
*,
tl,
) -> None:
"""GQA decode with M-fold + S_kv tile sweep + 2-level reduce-to-root.
Tensor layout:
Q : (T_q, h_q · d_head) replicated on every rank; loaded as
(G·T_q, d_head) — byte-conserving and math-correct for T_q=1.
K : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank
loads its (d_head, S_local) slice via byte-conserving reshape.
V : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank
loads its (S_local, d_head) slice.
O : (T_q, h_q · d_head) — only PE 0 of the root CUBE stores
(CUBE 0 for ``sub_w=0``; geometric center cube for ``sub_w>=2``).
"""
if sub_w > 0:
sub_h = C // sub_w
if sub_w < 2 or sub_h < 2 or sub_w * sub_h != C:
raise ValueError(
f"sub_w={sub_w} requires sub_w>=2 and sub_h>=2 and "
f"sub_w*sub_h==C; got sub_h={sub_h}, C={C}. "
"Use sub_w=0 for the 1D-chain path."
)
root_col = sub_w // 2
root_row = sub_h // 2
root_cube = root_row * sub_w + root_col
else:
sub_h = 0
root_col = 0
root_row = 0
root_cube = 0
G = h_q // h_kv
n_ranks = C * P
S_local = S_kv // n_ranks
pe_id = tl.program_id(axis=0)
cube_id = tl.program_id(axis=1)
# ── Local attention (S_kv-axis tile sweep, ADR-0063 §A.2) ──
Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
KV_ROW_BYTES = d_head * 2 # f16
# Tile 0: establishes persistent (m_local, l_local, O_local).
#
# Cannot be folded into the Tiles 1..N loop (kernbench-only limitation):
# - persistent (m, ℓ, O) must live OUTSIDE ``tl.scratch_scope``,
# otherwise scope teardown discards them before the next tile's
# merge can read them;
# - kernbench has no scratch-backed initializer — ``tl.zeros`` /
# ``tl.full`` return addr=0 handles with no backing storage, so
# they cannot be overwritten via ``tl.copy_to`` to seed (-inf, 0, 0).
# So Tile 0 computes the initial running state directly; Tiles 1..N
# fold into it. Triton port: limitation does not apply (SSA tensors
# stay live across iterations) — a single unified loop suffices.
tile_s0 = min(TILE_S_KV, S_local)
K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
scores = tl.dot(Q, K_T)
m_local = tl.max(scores, axis=-1)
centered = scores - m_local
exp_scores = tl.exp(centered)
l_local = tl.sum(exp_scores, axis=-1)
O_local = tl.dot(exp_scores, V)
# Tiles 1..n_tiles-1: fold into running state via online-softmax merge.
# Triton port: drop the ``with tl.scratch_scope():`` line and replace
# each ``copy_to`` with a Python rebind.
for tile_idx in range(1, n_tiles):
tile_start = tile_idx * TILE_S_KV
tile_s = min(TILE_S_KV, S_local - tile_start)
with tl.scratch_scope():
K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
shape=(d_head, tile_s), dtype="f16")
V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
shape=(tile_s, d_head), dtype="f16")
scores_t = tl.dot(Q, K_T_t)
m_tile = tl.max(scores_t, axis=-1)
centered_t = scores_t - m_tile
exp_scores_t = tl.exp(centered_t)
l_tile = tl.sum(exp_scores_t, axis=-1)
O_tile = tl.dot(exp_scores_t, V_t)
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
# ── Communication: chain reduce-to-root at PE 0 of CUBE 0 ──
PE_GRID_COLS = 4
pe_col = pe_id % PE_GRID_COLS
pe_row = pe_id // PE_GRID_COLS
pe_cols_used = min(PE_GRID_COLS, P)
pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS
# Level-2 row chain (intra-CUBE, along intra_W, leftward).
if pe_cols_used > 1:
if pe_col < pe_cols_used - 1:
with tl.scratch_scope():
m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
if pe_col > 0:
tl.send(dir="intra_W", src=m_local)
tl.send(dir="intra_W", src=l_local)
tl.send(dir="intra_W", src=O_local)
# Level-2 col bridge (intra-CUBE, along intra_N, row-1 → row-0).
if pe_col == 0 and pe_rows_used > 1:
if pe_row < pe_rows_used - 1:
with tl.scratch_scope():
m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
if pe_row > 0:
tl.send(dir="intra_N", src=m_local)
tl.send(dir="intra_N", src=l_local)
tl.send(dir="intra_N", src=O_local)
# Level-1 inter-CUBE reduce (only PE 0 of each CUBE participates).
if pe_id == 0 and C > 1 and sub_w == 0:
# 1D chain along W, leftward; root at CUBE 0.
if cube_id < C - 1:
with tl.scratch_scope():
m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
if cube_id > 0:
tl.send(dir="W", src=m_local)
tl.send(dir="W", src=l_local)
tl.send(dir="W", src=O_local)
elif pe_id == 0 and sub_w > 0:
# ── ADR-0060 §4.2 lrab-adapted center-root mesh reduce ──
# Adapts Phases 1-2 of lrab_hierarchical_allreduce.py:
# bidirectional row reduce converges at root_col, bidirectional
# col reduce on root_col converges at root_row. Reduce-only
# (Phases 3-5 dropped: no inter-SIP, no broadcast — attention
# needs the answer at one rank). Plain ``+`` replaced with the
# log-sum-exp ``_merge_running`` for (m, ℓ, O).
row = cube_id // sub_w
col = cube_id % sub_w
# Phase 1: row reduce — converge at col == root_col.
if col == 0 and root_col > 0:
tl.send(dir="E", src=m_local)
tl.send(dir="E", src=l_local)
tl.send(dir="E", src=O_local)
elif 0 < col < root_col:
with tl.scratch_scope():
m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
tl.send(dir="E", src=m_local)
tl.send(dir="E", src=l_local)
tl.send(dir="E", src=O_local)
elif col == root_col:
if root_col > 0:
with tl.scratch_scope():
m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
if sub_w - 1 > root_col:
with tl.scratch_scope():
m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
elif root_col < col < sub_w - 1:
with tl.scratch_scope():
m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
tl.send(dir="W", src=m_local)
tl.send(dir="W", src=l_local)
tl.send(dir="W", src=O_local)
elif col == sub_w - 1 and sub_w - 1 > root_col:
tl.send(dir="W", src=m_local)
tl.send(dir="W", src=l_local)
tl.send(dir="W", src=O_local)
# Phase 2: col reduce on col == root_col — converge at row == root_row.
if col == root_col:
if row == 0 and root_row > 0:
tl.send(dir="S", src=m_local)
tl.send(dir="S", src=l_local)
tl.send(dir="S", src=O_local)
elif 0 < row < root_row:
with tl.scratch_scope():
m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
tl.send(dir="S", src=m_local)
tl.send(dir="S", src=l_local)
tl.send(dir="S", src=O_local)
elif row == root_row:
if root_row > 0:
with tl.scratch_scope():
m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
if sub_h - 1 > root_row:
with tl.scratch_scope():
m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
elif root_row < row < sub_h - 1:
with tl.scratch_scope():
m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
tl.send(dir="N", src=m_local)
tl.send(dir="N", src=l_local)
tl.send(dir="N", src=O_local)
elif row == sub_h - 1 and sub_h - 1 > root_row:
tl.send(dir="N", src=m_local)
tl.send(dir="N", src=l_local)
tl.send(dir="N", src=O_local)
# ── Final normalise + store (root only) ──
if pe_id == 0 and cube_id == root_cube:
O_final = O_local / l_local
tl.store(o_ptr, O_final)
src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_repl_pe_sp.py
+10 -2
View File
@@ -25,12 +25,20 @@ Topology / SFR:
""" """
from __future__ import annotations from __future__ import annotations
from kernbench.benches._gqa_attention_decode_long import _merge_running
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width). TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
"""Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
m_new = tl.maximum(m_local, m_other)
scale_old = tl.exp(m_local - m_new)
scale_new = tl.exp(m_other - m_new)
l_new = l_local * scale_old + l_other * scale_new
O_new = O_local * scale_old + O_other * scale_new
return m_new, l_new, O_new
def gqa_attention_decode_long_ctx_cube_repl_pe_sp_kernel( def gqa_attention_decode_long_ctx_cube_repl_pe_sp_kernel(
q_ptr: int, q_ptr: int,
k_ptr: int, k_ptr: int,
src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_sp_pe_sp.py
+269 -20
View File
@@ -1,26 +1,60 @@
"""GQA decode kernel — Case 4 (Cube-SP × PE-SP). ★ slide-11 optimal. """GQA decode kernel — Case 4 (Cube-SP × PE-SP). ★ slide-11 optimal.
This is a thin entry-point wrapper around the underlying decode_long
lrab-adapted center-root kernel with ``sub_w`` baked to 4 (the C=8
single-KV-head-group geometry: 4×2 cube sub-mesh, root cube 6). The
math lives in ``_gqa_attention_decode_long.py`` (which still serves
non-4-cases panels via ``sub_w=0`` and the headline GQA bench); this
file exists so the 4-cases bench refers to all four cases via
parallel ``_gqa_attention_decode_<case>.py`` kernel-file names.
Per GQA_full_deck.pptx slide 11: Per GQA_full_deck.pptx slide 11:
- K, V split 64-way (cube=row_wise, pe=row_wise); each rank owns - K, V split 64-way (cube=row_wise, pe=row_wise); each rank owns
S_local = S_kv / (C·P). ``S_local = S_kv / (C·P)``.
- Two-level reduce on the partial (m, ℓ, O): intra-CUBE 8-way - Local attention via per-rank S_kv-axis tile sweep (ADR-0063 §A.2)
(row chain + col bridge over the 2×4 PE grid) then inter-CUBE keeps scratch bounded by ``TILE_S_KV``.
2-phase lrab over the 4×2 cube sub-mesh — converges at cube 6. - Two-level reduce on the partial ``(m, ℓ, O)``:
- PE 0 of CUBE 6 stores O. • intra-CUBE 8-way (row chain along intra_W + col bridge along
intra_N) over the 2×4 PE grid;
• inter-CUBE 2-phase lrab over the 4×2 CUBE sub-mesh (ADR-0060
§4.2 lrab-adapted center-root reduce): Phase 1 row reduce
converges at root_col=2; Phase 2 col reduce on root_col
converges at root_row=1; root cube id = 6.
- PE 0 of CUBE 6 stores ``O``.
Deviation from slide 13: slide prescribes AllReduce (every rank has
the answer); the kernel does reduce-to-root (only cube 6 has it)
per ADR-0060 §4. Treated as the kernbench Case-4 baseline.
Tensor layout:
Q : (T_q, h_q · d_head) replicated; loaded as (G·T_q, d_head).
K : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank
loads its (d_head, S_local) slice via byte-conserving reshape
of (S_local, h_kv·d_head) — correct for zero / symmetric
inputs (ADR-0060 §3 reshape-not-transpose caveat).
V : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank
loads its (S_local, d_head) slice.
O : (T_q, h_q · d_head) — only PE 0 of CUBE 6 stores.
Topology / SFR:
- Requires ``configure_sfr_intercube_multisip`` (provides disjoint
``intra_*`` and ``E/W/N/S`` namespaces).
- Intra-CUBE PEs arranged as a 2×4 grid (no wrap).
- Inter-CUBE: 4×2 CUBE sub-mesh starting at origin (0, 0).
""" """
from __future__ import annotations from __future__ import annotations
from kernbench.benches._gqa_attention_decode_long import (
gqa_attention_decode_long_kernel, TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
)
# lrab geometry for the C=8 single-KV-head group (4×2 cube sub-mesh).
_SUB_W = 4
_SUB_H = 2
_ROOT_COL = _SUB_W // 2 # 2
_ROOT_ROW = _SUB_H // 2 # 1
_ROOT_CUBE = _ROOT_ROW * _SUB_W + _ROOT_COL # 6
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
"""Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
m_new = tl.maximum(m_local, m_other)
scale_old = tl.exp(m_local - m_new)
scale_new = tl.exp(m_other - m_new)
l_new = l_local * scale_old + l_other * scale_new
O_new = O_local * scale_old + O_other * scale_new
return m_new, l_new, O_new
def gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel( def gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel(
@@ -39,8 +73,223 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel(
tl, tl,
) -> None: ) -> None:
"""Case-4 decode: Cube-SP × PE-SP, lrab center-root at cube 6.""" """Case-4 decode: Cube-SP × PE-SP, lrab center-root at cube 6."""
gqa_attention_decode_long_kernel( G = h_q // h_kv
q_ptr, k_ptr, v_ptr, o_ptr, n_ranks = C * P
T_q, S_kv, h_q, h_kv, d_head, C, P, 4, # sub_w=4 baked in S_local = S_kv // n_ranks
tl=tl, pe_id = tl.program_id(axis=0)
cube_id = tl.program_id(axis=1)
# ── Local attention (S_kv-axis tile sweep, ADR-0063 §A.2) ──
Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
KV_ROW_BYTES = d_head * 2 # f16
# Tile 0: establishes persistent (m_local, l_local, O_local). Cannot
# fold into Tiles 1..N loop (kernbench-only): persistent tensors
# must live outside tl.scratch_scope or scope teardown discards
# them; there's no scratch-backed (-inf, 0, 0) initializer.
tile_s0 = min(TILE_S_KV, S_local)
K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
scores = tl.dot(Q, K_T)
m_local = tl.max(scores, axis=-1)
centered = scores - m_local
exp_scores = tl.exp(centered)
l_local = tl.sum(exp_scores, axis=-1)
O_local = tl.dot(exp_scores, V)
for tile_idx in range(1, n_tiles):
tile_start = tile_idx * TILE_S_KV
tile_s = min(TILE_S_KV, S_local - tile_start)
with tl.scratch_scope():
K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
shape=(d_head, tile_s), dtype="f16")
V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
shape=(tile_s, d_head), dtype="f16")
scores_t = tl.dot(Q, K_T_t)
m_tile = tl.max(scores_t, axis=-1)
centered_t = scores_t - m_tile
exp_scores_t = tl.exp(centered_t)
l_tile = tl.sum(exp_scores_t, axis=-1)
O_tile = tl.dot(exp_scores_t, V_t)
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
) )
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
# ── Intra-CUBE reduce: row chain (intra_W) + col bridge (intra_N) ──
PE_GRID_COLS = 4
pe_col = pe_id % PE_GRID_COLS
pe_row = pe_id // PE_GRID_COLS
pe_cols_used = min(PE_GRID_COLS, P)
pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS
if pe_cols_used > 1:
if pe_col < pe_cols_used - 1:
with tl.scratch_scope():
m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
if pe_col > 0:
tl.send(dir="intra_W", src=m_local)
tl.send(dir="intra_W", src=l_local)
tl.send(dir="intra_W", src=O_local)
if pe_col == 0 and pe_rows_used > 1:
if pe_row < pe_rows_used - 1:
with tl.scratch_scope():
m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
if pe_row > 0:
tl.send(dir="intra_N", src=m_local)
tl.send(dir="intra_N", src=l_local)
tl.send(dir="intra_N", src=O_local)
# ── Inter-CUBE lrab-adapted center-root reduce (ADR-0060 §4.2) ──
# Only PE 0 of each CUBE participates. Adapts Phases 1-2 of
# lrab_hierarchical_allreduce.py: bidirectional row reduce converges
# at root_col; bidirectional col reduce on root_col converges at
# root_row. Plain ``+`` replaced by log-sum-exp ``_merge_running``.
if pe_id == 0:
row = cube_id // _SUB_W
col = cube_id % _SUB_W
# Phase 1: row reduce — converge at col == _ROOT_COL.
if col == 0:
tl.send(dir="E", src=m_local)
tl.send(dir="E", src=l_local)
tl.send(dir="E", src=O_local)
elif 0 < col < _ROOT_COL:
with tl.scratch_scope():
m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
tl.send(dir="E", src=m_local)
tl.send(dir="E", src=l_local)
tl.send(dir="E", src=O_local)
elif col == _ROOT_COL:
with tl.scratch_scope():
m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
with tl.scratch_scope():
m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
elif _ROOT_COL < col < _SUB_W - 1:
with tl.scratch_scope():
m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
tl.send(dir="W", src=m_local)
tl.send(dir="W", src=l_local)
tl.send(dir="W", src=O_local)
elif col == _SUB_W - 1:
tl.send(dir="W", src=m_local)
tl.send(dir="W", src=l_local)
tl.send(dir="W", src=O_local)
# Phase 2: col reduce on col == _ROOT_COL — converge at row == _ROOT_ROW.
if col == _ROOT_COL:
if row == 0:
tl.send(dir="S", src=m_local)
tl.send(dir="S", src=l_local)
tl.send(dir="S", src=O_local)
elif 0 < row < _ROOT_ROW:
with tl.scratch_scope():
m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
tl.send(dir="S", src=m_local)
tl.send(dir="S", src=l_local)
tl.send(dir="S", src=O_local)
elif row == _ROOT_ROW:
with tl.scratch_scope():
m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
if _SUB_H - 1 > _ROOT_ROW:
with tl.scratch_scope():
m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
elif _ROOT_ROW < row < _SUB_H - 1:
with tl.scratch_scope():
m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16")
m_new, l_new, O_new = _merge_running(
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
)
tl.copy_to(m_local, m_new)
tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new)
tl.send(dir="N", src=m_local)
tl.send(dir="N", src=l_local)
tl.send(dir="N", src=O_local)
elif row == _SUB_H - 1 and _SUB_H - 1 > _ROOT_ROW:
tl.send(dir="N", src=m_local)
tl.send(dir="N", src=l_local)
tl.send(dir="N", src=O_local)
# ── Final normalise + store (root only: PE 0 of CUBE 6) ──
if pe_id == 0 and cube_id == _ROOT_CUBE:
O_final = O_local / l_local
tl.store(o_ptr, O_final)
src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_sp_pe_tp.py
+10 -2
View File
@@ -27,12 +27,20 @@ Topology / SFR:
""" """
from __future__ import annotations from __future__ import annotations
from kernbench.benches._gqa_attention_decode_long import _merge_running
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width). TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
"""Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
m_new = tl.maximum(m_local, m_other)
scale_old = tl.exp(m_local - m_new)
scale_new = tl.exp(m_other - m_new)
l_new = l_local * scale_old + l_other * scale_new
O_new = O_local * scale_old + O_other * scale_new
return m_new, l_new, O_new
def gqa_attention_decode_long_ctx_cube_sp_pe_tp_kernel( def gqa_attention_decode_long_ctx_cube_sp_pe_tp_kernel(
q_ptr: int, q_ptr: int,
k_ptr: int, k_ptr: int,
src/kernbench/benches/_gqa_attention_decode_opt2.py
+4 -3
View File
@@ -37,9 +37,10 @@ def gqa_attention_decode_opt2_kernel(
) -> None: ) -> None:
"""Single-rank (C=P=1) GQA decode using the opt2 two-composite form. """Single-rank (C=P=1) GQA decode using the opt2 two-composite form.
Layout mirrors ``gqa_attention_decode_long_kernel`` (M-fold Q, K loaded Layout mirrors ``gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel``
as ``(d_head, S_local)``). The KV slice is split into two sub-tiles so (M-fold Q, K loaded as ``(d_head, S_local)``). The KV slice is split
the second one drives the ``softmax_merge`` recipe composite. into two sub-tiles so the second one drives the ``softmax_merge``
recipe composite.
""" """
G = h_q // h_kv G = h_q // h_kv
S_local = S_kv // (C * P) S_local = S_kv // (C * P)
src/kernbench/benches/milestone_gqa_headline.py
+20 -62
View File
@@ -1,29 +1,26 @@
"""milestone-gqa-headline bench: real GQA + 2-level SP + Ring KV. """milestone-gqa-headline bench: real GQA prefill — single-KV-group target.
Wires the new ``_gqa_decode`` and ``_gqa_prefill`` kernels through 4 Drives the LLaMA-3.1-70B single-KV-head-group prefill panel through
panels with real GQA (h_q = G·h_kv, G > 1 on the decode side), writing ``_gqa_attention_prefill_long`` (C=8 snake Ring KV + intra-CUBE PE-SP,
per-panel ``op_log_summary`` into ``sweep.json``. Independent from the all 64 ranks), writing ``op_log_summary`` into ``sweep.json``.
existing ``milestone-gqa-llama70b`` validation-scale bench (which stays Independent from the existing ``milestone-gqa-llama70b`` validation-scale
on the legacy baseline kernels). bench (which stays on the legacy baseline kernels).
Decode panels live in ``milestone_gqa_decode_long_ctx_4cases.py`` (the
4-cases long-context decode comparative study).
Restrictions: Restrictions:
- Decode side capped at C ≤ 4 (1D chain reduce; 2D mesh wiring on
the decode panels deferred to the 4-cases decode comparative study)
- Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP - Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP
headline deferred per ADR-0060 §B-item-1) headline deferred per ADR-0060 §B-item-1)
- T_q = S_kv = 1024 (scratch-limited; 32K headline awaits Q-axis
kernel tiling)
- No figure renderers (defer to a separate cycle) - No figure renderers (defer to a separate cycle)
Panels: Panels:
single_kv_group_prefill_gqa_c8_p8: prefill C=8 snake Ring KV + single_kv_group_prefill_gqa_c8_p8: prefill C=8 snake Ring KV +
intra-CUBE PE-SP (all 64 ranks), intra-CUBE PE-SP, T_q=S_kv=1K,
T_q=S_kv=1K, d_head=128 — the d_head=128 — the LLaMA-3.1-70B
LLaMA-3.1-70B single-KV-group target single-KV-group prefill target.
(scratch-limited; 32K headline awaits
Q-axis kernel tiling)
Decode panels live in ``milestone_gqa_decode_4cases.py`` (the 4-cases
comparative study). Legacy C=1 / C=4 panels were dropped — pytest
regression already covers those configurations.
Gated by ``GQA_HEADLINE_RUN=1``. Gated by ``GQA_HEADLINE_RUN=1``.
""" """
@@ -33,14 +30,10 @@ import json
import os import os
from pathlib import Path from pathlib import Path
from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel
from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel
from kernbench.benches.registry import bench from kernbench.benches.registry import bench
from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
from kernbench.ccl.sfr_config import ( from kernbench.ccl.sfr_config import configure_sfr_intercube_ring
configure_sfr_intercube_multisip,
configure_sfr_intercube_ring,
)
from kernbench.policy.placement.dp import DPPolicy from kernbench.policy.placement.dp import DPPolicy
_OUTPUT_DIR = Path(__file__).resolve().parent / "1H_milestone_output" / "gqa_headline" _OUTPUT_DIR = Path(__file__).resolve().parent / "1H_milestone_output" / "gqa_headline"
@@ -51,10 +44,7 @@ _SWEEP_JSON = _OUTPUT_DIR / "sweep.json"
_DTYPE = "f16" _DTYPE = "f16"
_D_HEAD = 64 _D_HEAD = 64
_T_Q_PREFILL = 4 _T_Q_PREFILL = 4
_T_Q_DECODE = 1
_S_KV_PREFILL = 16 _S_KV_PREFILL = 16
_H_Q_DECODE = 8 # real GQA: G = H_Q_DECODE / H_KV_DECODE = 8
_H_KV_DECODE = 1
_PANELS = ( _PANELS = (
"single_kv_group_prefill_gqa_c8_p8", "single_kv_group_prefill_gqa_c8_p8",
@@ -133,43 +123,14 @@ def _run_prefill_panel(
) )
def _run_decode_panel(
ctx, *, panel: str, C: int, P: int, S_kv: int,
sub_w: int = 0,
T_q: int = _T_Q_DECODE,
d_head: int = _D_HEAD,
h_q: int = _H_Q_DECODE,
h_kv: int = _H_KV_DECODE,
) -> None:
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_q * d_head),
dtype=_DTYPE, dp=dp_full, name=f"{panel}_q")
k = ctx.zeros((S_kv, h_kv * d_head),
dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k")
v = ctx.zeros((S_kv, h_kv * d_head),
dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v")
o = ctx.empty((T_q, h_q * d_head),
dtype=_DTYPE, dp=dp_full, name=f"{panel}_o")
ctx.launch(
panel, gqa_attention_decode_long_kernel,
q, k, v, o,
T_q, S_kv, h_q, h_kv, d_head, C, P, sub_w,
_auto_dim_remap=False,
)
def _make_bench_fn(panel: str): def _make_bench_fn(panel: str):
kind, params = _PANEL_DISPATCH[panel] kind, params = _PANEL_DISPATCH[panel]
assert kind == "prefill", (
f"milestone-gqa-headline only registers prefill panels; got {kind!r}"
)
def _bench_fn(ctx): def _bench_fn(ctx):
if kind == "prefill":
_run_prefill_panel(ctx, panel=panel, **params) _run_prefill_panel(ctx, panel=panel, **params)
else:
_run_decode_panel(ctx, panel=panel, **params)
return _bench_fn return _bench_fn
@@ -234,10 +195,10 @@ def _run_panel(panel: str, topology: str) -> dict:
@bench( @bench(
name="milestone-gqa-headline", name="milestone-gqa-headline",
description="Headline GQA milestone — real GQA h_q=8/h_kv=1 + 2-level SP (decode) + Ring KV (prefill).", description="Headline GQA prefill milestone — LLaMA-3.1-70B single-KV-group target (C=8, P=8).",
) )
def run(torch) -> None: def run(torch) -> None:
"""Drive 4 headline panels through the new GQA kernels; write sweep.json. """Drive the headline prefill panel; write sweep.json.
Gated by GQA_HEADLINE_RUN=1. Gated by GQA_HEADLINE_RUN=1.
""" """
@@ -254,10 +215,7 @@ def run(torch) -> None:
"panels": list(_PANELS), "panels": list(_PANELS),
"config": { "config": {
"T_q_prefill": _T_Q_PREFILL, "T_q_prefill": _T_Q_PREFILL,
"T_q_decode": _T_Q_DECODE,
"S_kv_prefill": _S_KV_PREFILL, "S_kv_prefill": _S_KV_PREFILL,
"h_q_decode": _H_Q_DECODE,
"h_kv_decode": _H_KV_DECODE,
"d_head": _D_HEAD, "d_head": _D_HEAD,
}, },
"rows": rows, "rows": rows,
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}" 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}"
)