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 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).
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(
q_ptr: int,
k_ptr: int,
src/kernbench/benches/_gqa_attention_decode_long_ctx_cube_sp_pe_sp.py
+270 -21
View File
@@ -1,26 +1,60 @@
"""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:
- K, V split 64-way (cube=row_wise, pe=row_wise); each rank owns
S_local = S_kv / (C·P).
- Two-level reduce on the partial (m, ℓ, O): intra-CUBE 8-way
(row chain + col bridge over the 2×4 PE grid) then inter-CUBE
2-phase lrab over the 4×2 cube sub-mesh — converges at cube 6.
- PE 0 of CUBE 6 stores O.
``S_local = S_kv / (C·P)``.
- Local attention via per-rank S_kv-axis tile sweep (ADR-0063 §A.2)
keeps scratch bounded by ``TILE_S_KV``.
- Two-level reduce on the partial ``(m, ℓ, 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 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(
@@ -39,8 +73,223 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel(
tl,
) -> None:
"""Case-4 decode: Cube-SP × PE-SP, lrab center-root at cube 6."""
gqa_attention_decode_long_kernel(
q_ptr, k_ptr, v_ptr, o_ptr,
T_q, S_kv, h_q, h_kv, d_head, C, P, 4, # sub_w=4 baked in
tl=tl,
)
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
# 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 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).
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(
q_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:
"""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
as ``(d_head, S_local)``). The KV slice is split into two sub-tiles so
the second one drives the ``softmax_merge`` recipe composite.
Layout mirrors ``gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel``
(M-fold Q, K loaded as ``(d_head, S_local)``). The KV slice is split
into two sub-tiles so the second one drives the ``softmax_merge``
recipe composite.
"""
G = h_q // h_kv
S_local = S_kv // (C * P)
src/kernbench/benches/milestone_gqa_headline.py
+21 -63
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
panels with real GQA (h_q = G·h_kv, G > 1 on the decode side), writing
per-panel ``op_log_summary`` into ``sweep.json``. Independent from the
existing ``milestone-gqa-llama70b`` validation-scale bench (which stays
on the legacy baseline kernels).
Drives the LLaMA-3.1-70B single-KV-head-group prefill panel through
``_gqa_attention_prefill_long`` (C=8 snake Ring KV + intra-CUBE PE-SP,
all 64 ranks), writing ``op_log_summary`` into ``sweep.json``.
Independent from the existing ``milestone-gqa-llama70b`` validation-scale
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:
- 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
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)
Panels:
single_kv_group_prefill_gqa_c8_p8: prefill C=8 snake Ring KV +
intra-CUBE PE-SP (all 64 ranks),
T_q=S_kv=1K, d_head=128 — the
LLaMA-3.1-70B single-KV-group 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.
intra-CUBE PE-SP, T_q=S_kv=1K,
d_head=128 — the LLaMA-3.1-70B
single-KV-group prefill target.
Gated by ``GQA_HEADLINE_RUN=1``.
"""
@@ -33,14 +30,10 @@ import json
import os
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.registry import bench
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.ccl.sfr_config import configure_sfr_intercube_ring
from kernbench.policy.placement.dp import DPPolicy
_OUTPUT_DIR = Path(__file__).resolve().parent / "1H_milestone_output" / "gqa_headline"
@@ -51,10 +44,7 @@ _SWEEP_JSON = _OUTPUT_DIR / "sweep.json"
_DTYPE = "f16"
_D_HEAD = 64
_T_Q_PREFILL = 4
_T_Q_DECODE = 1
_S_KV_PREFILL = 16
_H_Q_DECODE = 8 # real GQA: G = H_Q_DECODE / H_KV_DECODE = 8
_H_KV_DECODE = 1
_PANELS = (
"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):
kind, params = _PANEL_DISPATCH[panel]
assert kind == "prefill", (
f"milestone-gqa-headline only registers prefill panels; got {kind!r}"
)
def _bench_fn(ctx):
if kind == "prefill":
_run_prefill_panel(ctx, panel=panel, **params)
else:
_run_decode_panel(ctx, panel=panel, **params)
_run_prefill_panel(ctx, panel=panel, **params)
return _bench_fn
@@ -234,10 +195,10 @@ def _run_panel(panel: str, topology: str) -> dict:
@bench(
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:
"""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.
"""
@@ -254,10 +215,7 @@ def run(torch) -> None:
"panels": list(_PANELS),
"config": {
"T_q_prefill": _T_Q_PREFILL,
"T_q_decode": _T_Q_DECODE,
"S_kv_prefill": _S_KV_PREFILL,
"h_q_decode": _H_Q_DECODE,
"h_kv_decode": _H_KV_DECODE,
"d_head": _D_HEAD,
},
"rows": rows,