mukesh
9e1242039b
End-to-end wires C=8 P=8 d_head=128 prefill with snake-ring inter-CUBE SFR, intra-CUBE PE-SP (all 64 ranks active), and the milestone bench panel. Decode kernel gains lrab-adapted center-root reduce for the 2×4 sub-mesh per ADR-0060 §4.2. Increment 1 — SFR multi-row snake src/kernbench/ccl/sfr_config.py: configure_sfr_intercube_ring gains submesh_shape / submesh_origin kwargs; installs a Hamiltonian snake ring through a rectangular sub-mesh (every hop is 1-hop physical neighbour). Backward-compat: 1D-row behaviour preserved when submesh_shape is None. tests/test_intercube_snake_ring.py (12 tests) Increment 2 — Decode lrab-adapted center-root reduce src/kernbench/benches/_gqa_attention_decode_long.py: new sub_w param (default 0 = existing 1D-chain). sub_w >= 2 selects the ADR-0060 §4.2 prescribed lrab-adapted Phase 1+2 reduce (bidirectional row + bidirectional col converge to the center cube), with log-sum-exp _merge_running replacing the plain + of lrab. tests/attention/test_gqa_decode_long_2d_reduce.py (4 tests) Increment 3 — Prefill kernel at C=8 (no production change) Verified by inspection that the existing prefill_long kernel + Increment 1's snake SFR already work at C=8 without any kernel edit. The kernel speaks logical W/E; the snake routes it. tests/attention/test_gqa_prefill_long_c8_snake.py (3 tests) Increment 4 — Intra-CUBE PE-SP in prefill (all 64 ranks) src/kernbench/benches/_gqa_attention_prefill_long.py: new P param (default 1 = existing PE-0-only). P > 1 splits T_q query-axis-wise across the P PEs of each CUBE; output rows are disjoint per PE so no intra-CUBE reduce is needed; each PE drives its own same-lane ring (P parallel rings). tests/attention/test_gqa_prefill_long_pe_sp.py (5 tests) Increment 5 — LLaMA-scale milestone bench panel src/kernbench/benches/milestone_gqa_headline.py: new panel single_kv_group_prefill_gqa_c8_p8 (C=8, P=8, T_q=S_kv=32K, d_head=128). _run_prefill_panel extended with P/T_q/d_head defaults; routes snake SFR when C > mesh_w. tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py (3 tests) Total: 4 production files modified, 5 new test files, 27 new tests. Followed the Phase 1/2 protocol per CLAUDE.md throughout. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
354 lines
16 KiB
Python
354 lines
16 KiB
Python
"""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)
|