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

short context length attention kernel

Browse Source
This commit is contained in:
eusonice
2026-06-23 16:19:39 -07:00
parent faf011dae5
commit 7d90d53d4d
2 changed files with 279 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_decode_short.py
+107 -63
View File
@@ -1,37 +1,74 @@
"""GQA fused-attention decode kernel — short context (ADR-0060 §B.split.2). """GQA decode kernel: short context, attention only, multi-tile per PE.
Short context (``S_kv < 256K``): each CUBE owns ``kv_per_cube`` whole Unified A1/A2/A4/B decode mapping per ADR-0070 (phase mirror of
KV heads, with no S_kv sharding across CUBEs and no inter-CUBE reduce. ``_gqa_attention_prefill_short.py``). Mode selected at launch via
PE-SP within each CUBE: the ``P`` PEs split into ``kv_per_cube`` groups ``kv_per_cube ∈ {1, 2, 4, 8}``:
of ``P/kv_per_cube`` PEs each; each group does PE-SP across the group
for one owned head, then the group's root PE stores its head's output.
The local attention uses an S_kv-axis tile sweep (ADR-0063 §A.2) so Mode kv_per_cube C group_size Reduce topology
per-rank scratch is bounded by ``TILE_S_KV``. ---- ----------- ------- ---------- -----------------------------
A1 1 h_kv P (=8) row 0 + col bridge + row 1
A2 2 h_kv/2 P/2 (=4) row chain only
A4 4 h_kv/4 P/4 (=2) single intra_W hop
B 8 1 1 (no reduce, single PE)
Group layout on the 2×4 PE grid: Cube ``h`` owns ``kv_per_cube`` whole KV heads. Within each cube the
kv_per_cube=1, group=8 PEs (full 2×4): row chain + col bridge. 8 PEs split into ``kv_per_cube`` groups of ``group_size = P/kv_per_cube``
kv_per_cube=2, group=4 PEs (one row): row chain only. PEs; each PE owns ``S_local = S_kv/group_size`` KV tokens (sequence
kv_per_cube=4, group=2 PEs (adj cols): 1-step chain. shard). FA2 fuses ``G = h_q/h_kv`` Q heads into the M dim of one GEMM
kv_per_cube=8, group=1 PE: no chain — direct write. per tile. After ``n_tiles_per_pe = S_local/TILE_S_KV`` tiles, partials
``(m, ℓ, O)`` chain-reduce up to group root (PE 0), which normalizes
and stores the cube's O slab.
Layout caveats: Per ADR-0011 D-VA1 the kernel computes its own shard base offset from
- K, V: ``(h_kv·S_kv, d_head)`` head-stacked, deployed with ``program_id(axis=0)`` (PE id) and ``program_id(axis=1)`` (cube id);
``dp = (cube=row_wise, pe=row_wise)`` so each PE's chunk is deploy places shards in HBM but the kernel must address them.
contiguously ``(S_local, d_head)`` at its own shard. K loaded as
``(d_head, S_local)`` via byte-conserving reshape (ADR-0060 §3). Tensor layouts (host-side, mode-invariant byte totals):
- Q: replicated ``(T_q, h_q·d_head)``, reshaped byte-conservingly to - Q: ``(kv_per_cube·T_q, h_q·d_head/kv_per_cube)``, T_q=1.
``(h_q·T_q, d_head)``. Kernel computes attention for ALL Q rows dp=(cube=column_wise, pe=replicate). Caller pre-scales by
against the group's owned K head; only the group's owned head rows ``1/sqrt(d_head)``.
are semantically meaningful (correct for zero / symmetric inputs). - K: ``(h_kv·S_kv·d_head/TILE_S_KV, TILE_S_KV)`` tile-major.
- O: replicated; each group root writes its head's dp=(cube=row_wise, pe=row_wise).
``(h_q·T_q, d_head)`` result at disjoint PE-local addresses. - V: ``(h_kv·S_kv, d_head)`` native. Same dp as K.
- GEMMs use ``tl.dot`` (no composite epilogue / ``softmax_scale``). - O: same dp as Q; only group root (pe_in_group==0) stores.
Configuration constraints:
kv_per_cube ∈ {1, 2, 4, 8}, T_q == 1, P == 8,
C == h_kv/kv_per_cube, h_q % h_kv == 0,
S_kv % (group_size·TILE_S_KV) == 0.
Out of scope: causal mask (decode is auto-causal), f32 accumulator.
""" """
from __future__ import annotations from __future__ import annotations
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width). TILE_S_KV = 1024
def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
h_q: int, h_kv: int, S_kv: int) -> None:
"""Validate decode kernel configuration before the run.
"""
if kv_per_cube not in (1, 2, 4, 8):
raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
if T_q != 1:
raise ValueError(f"decode requires T_q == 1; got {T_q}")
if P != 8:
raise ValueError(f"2x4 mesh chain-reduce geometry requires P == 8; got {P}")
if h_q % h_kv != 0:
raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
if C != h_kv // kv_per_cube:
raise ValueError(
f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
)
group_size = P // kv_per_cube
if S_kv % (group_size * TILE_S_KV) != 0:
raise ValueError(
f"S_kv ({S_kv}) must be a multiple of group_size·TILE_S_KV "
f"({group_size * TILE_S_KV}); each PE's sequence shard must "
f"be a whole number of tiles"
)
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl): def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
@@ -60,32 +97,40 @@ def gqa_attention_decode_short_kernel(
*, *,
tl, tl,
) -> None: ) -> None:
"""Short-context GQA decode with PE-parallel heads + intra-group PE-SP.""" """Unified decode: sequence-shard + chain reduce + FA2 (ADR-0070).
"""
group_size = P // kv_per_cube group_size = P // kv_per_cube
pe_id = tl.program_id(axis=0)
pe_in_group = pe_id % group_size
S_local = S_kv // group_size S_local = S_kv // group_size
n_tiles_per_pe = S_local // TILE_S_KV
# ── Local attention (S_kv-axis tile sweep, ADR-0063 §A.2) ── G = h_q // h_kv
Q = tl.load(q_ptr, shape=(h_q * T_q, d_head), dtype="f16") pe_id = tl.program_id(axis=0)
n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV cube_id = tl.program_id(axis=1)
KV_ROW_BYTES = d_head * 2 # f16 pe_in_group = pe_id % group_size
group_id_in_cube = pe_id // group_size
# Tile 0: establishes persistent (m_local, l_local, O_local). KV_ROW_BYTES = d_head * 2
# K_TILE_BYTES = d_head * TILE_S_KV * 2
# Cannot be folded into the Tiles 1..N loop (kernbench-only limitation): Q_ROW_BYTES = G * d_head * 2
# - persistent (m, ℓ, O) must live OUTSIDE ``tl.scratch_scope``, K_HEAD_BYTES = S_kv * d_head * 2
# otherwise scope teardown discards them before the next tile's V_HEAD_BYTES = S_kv * KV_ROW_BYTES
# merge can read them;
# - kernbench has no scratch-backed initializer — ``tl.zeros`` / # Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
# ``tl.full`` return addr=0 handles with no backing storage, so q_base = q_ptr + cube_id * kv_per_cube * Q_ROW_BYTES + group_id_in_cube * Q_ROW_BYTES
# they cannot be overwritten via ``tl.copy_to`` to seed (-inf, 0, 0). k_shard_base = (k_ptr
# So Tile 0 computes the initial running state directly; Tiles 1..N + cube_id * kv_per_cube * K_HEAD_BYTES
# fold into it. Triton port: limitation does not apply (SSA tensors + group_id_in_cube * K_HEAD_BYTES
# stay live across iterations) — a single unified loop suffices. + pe_in_group * n_tiles_per_pe * K_TILE_BYTES)
tile_s0 = min(TILE_S_KV, S_local) v_shard_base = (v_ptr
K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16") + cube_id * kv_per_cube * V_HEAD_BYTES
V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16") + group_id_in_cube * V_HEAD_BYTES
+ pe_in_group * n_tiles_per_pe * TILE_S_KV * KV_ROW_BYTES)
Q = tl.load(q_base, shape=(G, d_head), dtype="f16")
# ── Tile 0: establish (m, ℓ, O) ──────────────────────────────────
K_T = tl.load(k_shard_base, shape=(d_head, TILE_S_KV), dtype="f16")
V = tl.load(v_shard_base, shape=(TILE_S_KV, d_head), dtype="f16")
scores = tl.dot(Q, K_T) scores = tl.dot(Q, K_T)
m_local = tl.max(scores, axis=-1) m_local = tl.max(scores, axis=-1)
centered = scores - m_local centered = scores - m_local
@@ -93,17 +138,13 @@ def gqa_attention_decode_short_kernel(
l_local = tl.sum(exp_scores, axis=-1) l_local = tl.sum(exp_scores, axis=-1)
O_local = tl.dot(exp_scores, V) O_local = tl.dot(exp_scores, V)
# Tiles 1..n_tiles-1: fold into running state via online-softmax merge. # ── Tiles 1..n_tiles_per_pe-1: sweep this PE's sequence shard ──
# Triton port: drop the ``with tl.scratch_scope():`` line and replace for tile_idx in range(1, n_tiles_per_pe):
# 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(): with tl.scratch_scope():
K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES, K_T_t = tl.load(k_shard_base + tile_idx * K_TILE_BYTES,
shape=(d_head, tile_s), dtype="f16") shape=(d_head, TILE_S_KV), dtype="f16")
V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES, V_t = tl.load(v_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
shape=(tile_s, d_head), dtype="f16") shape=(TILE_S_KV, d_head), dtype="f16")
scores_t = tl.dot(Q, K_T_t) scores_t = tl.dot(Q, K_T_t)
m_tile = tl.max(scores_t, axis=-1) m_tile = tl.max(scores_t, axis=-1)
centered_t = scores_t - m_tile centered_t = scores_t - m_tile
@@ -117,13 +158,13 @@ def gqa_attention_decode_short_kernel(
tl.copy_to(l_local, l_new) tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new) tl.copy_to(O_local, O_new)
# ── Communication: within-group chain reduce-to-root (Level-2 only) ── # ── 2x4 mesh chain reduce geometry within the group ──
group_cols = min(4, group_size) group_cols = min(4, group_size)
group_rows = (group_size + group_cols - 1) // group_cols group_rows = (group_size + group_cols - 1) // group_cols
pe_col_in_group = pe_in_group % group_cols pe_col_in_group = pe_in_group % group_cols
pe_row_in_group = pe_in_group // group_cols pe_row_in_group = pe_in_group // group_cols
# Row chain (within group's row, along intra_W, leftward). # Row chain (intra_W, leftward) — every group with group_cols > 1.
if group_cols > 1: if group_cols > 1:
if pe_col_in_group < group_cols - 1: if pe_col_in_group < group_cols - 1:
with tl.scratch_scope(): with tl.scratch_scope():
@@ -141,7 +182,7 @@ def gqa_attention_decode_short_kernel(
tl.send(dir="intra_W", src=l_local) tl.send(dir="intra_W", src=l_local)
tl.send(dir="intra_W", src=O_local) tl.send(dir="intra_W", src=O_local)
# Col bridge (within group, along intra_N, row-1 → row-0). # Col bridge (intra_N, row-1 col-0 → row-0 col-0) — only A1 (group_rows > 1).
if pe_col_in_group == 0 and group_rows > 1: if pe_col_in_group == 0 and group_rows > 1:
if pe_row_in_group < group_rows - 1: if pe_row_in_group < group_rows - 1:
with tl.scratch_scope(): with tl.scratch_scope():
@@ -159,7 +200,10 @@ def gqa_attention_decode_short_kernel(
tl.send(dir="intra_N", src=l_local) tl.send(dir="intra_N", src=l_local)
tl.send(dir="intra_N", src=O_local) tl.send(dir="intra_N", src=O_local)
# ── Final normalise + store (group root only) ── # ── Final normalize + store (group root, pe_in_group == 0) ──
if pe_in_group == 0: if pe_in_group == 0:
O_final = O_local / l_local O_final = O_local / l_local
tl.store(o_ptr, O_final) o_base = (o_ptr
+ cube_id * kv_per_cube * Q_ROW_BYTES
+ group_id_in_cube * Q_ROW_BYTES)
tl.store(o_base, O_final)
src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_prefill_short.py
+172 -85
View File
@@ -1,15 +1,45 @@
"""GQA fused-attention prefill kernel — short context (ADR-0060 §B.split.2). """GQA prefill kernel: short context, attention only, multi-tile.
Prefill analogue of ``_gqa_decode_short.py`` — same CUBE/PE layout Unified A1/A2/A4/B prefill mapping per ADR-0070 (supersedes
(``kv_per_cube`` heads per CUBE, group-PE-SP, within-group chain reduce, ADR-0060 §B.split.2 prefill-short clause). Mode selected at launch
no inter-CUBE reduce). The only structural difference from short decode: via ``kv_per_cube ∈ {1, 2, 4, 8}``:
``T_q`` may be > 1 (prefill processes multiple query tokens) and Q is
shaped ``(T_q, h_kv·d_head)`` — one Q head per KV head, no GQA M-fold.
The local attention uses an S_kv-axis tile sweep (ADR-0063 §A.2) so Mode kv_per_cube C group_size Broadcast topology
per-rank scratch is bounded by ``TILE_S_KV``. ---- ----------- ------- ---------- ---------------------------
A1 1 h_kv P (=8) row 0 + col bridge + row 1
A2 2 h_kv/2 P/2 (=4) row chain only
A4 4 h_kv/4 P/4 (=2) single intra_E hop
B 8 1 1 (no broadcast, single PE)
No Ring KV here — each owned head is fully resident at its CUBE. Cube ``h`` owns ``kv_per_cube`` whole KV heads. Within each cube the
8 PEs split into ``kv_per_cube`` groups of ``group_size = P/kv_per_cube``
PEs; the group's T_q rows are floor-balanced across its PEs (Q-tile
split). FA2 fuses ``G = h_q/h_kv`` Q heads into one batched GEMM per
PE per tile. Group root (pe_in_group==0) loads K/V from HBM and IPCQ-
broadcasts each tile to the rest of its group; q-tiles are
independent so no intra-group reduce — every PE writes its own
``(T_q_pe·G, d_head)`` slab to disjoint rows of O.
Per ADR-0011 D-VA1 the kernel computes its own shard base offset from
``program_id(axis=0)`` (PE id) and ``program_id(axis=1)`` (cube id);
deploy places shards in HBM but the kernel must address them.
Tensor layouts (host-side, mode-invariant byte totals):
- Q: ``(kv_per_cube·T_q, h_q·d_head/kv_per_cube)``.
dp=(cube=column_wise, pe=replicate). Caller pre-scales by
``1/sqrt(d_head)``.
- K: ``(h_kv·n_tiles·d_head, TILE_S_KV)`` tile-major.
dp=(cube=row_wise, pe=replicate).
- V: ``(h_kv·S_kv, d_head)`` native. Same dp as K.
- O: same dp as Q; every PE stores.
Configuration constraints:
kv_per_cube ∈ {1, 2, 4, 8}, P == 8,
C == h_kv/kv_per_cube, h_q % h_kv == 0,
T_q >= group_size, S_kv % TILE_S_KV == 0.
Out of scope: causal mask, f32 accumulator. See ADR-0070 §Known
Limitations.
""" """
from __future__ import annotations from __future__ import annotations
@@ -17,6 +47,33 @@ from __future__ import annotations
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 _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
h_q: int, h_kv: int, S_kv: int) -> None:
"""Validate prefill kernel configuration before the run.
"""
if kv_per_cube not in (1, 2, 4, 8):
raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
if P != 8:
raise ValueError(f"2x4 mesh broadcast geometry requires P == 8; got {P}")
if h_q % h_kv != 0:
raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
if C != h_kv // kv_per_cube:
raise ValueError(
f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
)
group_size = P // kv_per_cube
if T_q < group_size:
raise ValueError(
f"T_q ({T_q}) must be >= group_size ({group_size}) so every PE "
f"has q-tile work for the broadcast chain"
)
if S_kv % TILE_S_KV != 0:
raise ValueError(
f"S_kv ({S_kv}) must be a multiple of TILE_S_KV ({TILE_S_KV})"
)
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl): 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.""" """Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
m_new = tl.maximum(m_local, m_other) m_new = tl.maximum(m_local, m_other)
@@ -34,6 +91,7 @@ def gqa_attention_prefill_short_kernel(
o_ptr: int, o_ptr: int,
T_q: int, T_q: int,
S_kv: int, S_kv: int,
h_q: int,
h_kv: int, h_kv: int,
d_head: int, d_head: int,
C: int, C: int,
@@ -42,32 +100,79 @@ def gqa_attention_prefill_short_kernel(
*, *,
tl, tl,
) -> None: ) -> None:
"""Short-context prefill with PE-parallel heads + intra-group PE-SP.""" """Unified prefill: Q-tile split + FA2 + IPCQ KV broadcast (ADR-0070)."""
group_size = P // kv_per_cube group_size = P // kv_per_cube
G = h_q // h_kv
pe_id = tl.program_id(axis=0) pe_id = tl.program_id(axis=0)
cube_id = tl.program_id(axis=1)
pe_in_group = pe_id % group_size pe_in_group = pe_id % group_size
S_local = S_kv // group_size group_id_in_cube = pe_id // group_size
# ── Local attention (S_kv-axis tile sweep, ADR-0063 §A.2) ── # Floor-balanced Q-tile partition.
Q = tl.load(q_ptr, shape=(h_kv * T_q, d_head), dtype="f16") q_start = (T_q * pe_in_group) // group_size
n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV q_end = (T_q * (pe_in_group + 1)) // group_size
KV_ROW_BYTES = d_head * 2 # f16 T_q_pe = q_end - q_start
n_tiles = S_kv // TILE_S_KV
Q_ROW_BYTES = G * d_head * 2
KV_ROW_BYTES = d_head * 2
K_TILE_BYTES = d_head * TILE_S_KV * 2
K_HEAD_BYTES = n_tiles * K_TILE_BYTES
V_HEAD_BYTES = S_kv * KV_ROW_BYTES
# Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
q_base = (q_ptr
+ cube_id * kv_per_cube * T_q * Q_ROW_BYTES
+ (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
k_head_shard_base = (k_ptr
+ cube_id * kv_per_cube * K_HEAD_BYTES
+ group_id_in_cube * K_HEAD_BYTES)
v_head_shard_base = (v_ptr
+ cube_id * kv_per_cube * V_HEAD_BYTES
+ group_id_in_cube * V_HEAD_BYTES)
Q = tl.load(q_base, shape=(T_q_pe * G, d_head), dtype="f16")
# 2x4 mesh broadcast geometry within the group.
group_cols = min(4, group_size)
group_rows = (group_size + group_cols - 1) // group_cols
pe_col_in_group = pe_in_group % group_cols
pe_row_in_group = pe_in_group // group_cols
# ──────────────────────────────────────────────────────────
# Tile 0: KV broadcast + establish persistent (m, ℓ, O).
# Persistent state lives OUTSIDE scratch_scope (ADR-0063 §A.3).
# ──────────────────────────────────────────────────────────
if pe_in_group == 0:
K_T = tl.load(k_head_shard_base,
shape=(d_head, TILE_S_KV), dtype="f16")
V = tl.load(v_head_shard_base,
shape=(TILE_S_KV, d_head), dtype="f16")
if group_cols > 1:
tl.send(dir="intra_E", src=K_T)
tl.send(dir="intra_E", src=V)
if group_rows > 1:
tl.send(dir="intra_S", src=K_T)
tl.send(dir="intra_S", src=V)
elif pe_col_in_group == 0 and pe_row_in_group > 0:
K_T = tl.recv(dir="intra_N",
shape=(d_head, TILE_S_KV), dtype="f16")
V = tl.recv(dir="intra_N",
shape=(TILE_S_KV, d_head), dtype="f16")
if group_cols > 1:
tl.send(dir="intra_E", src=K_T)
tl.send(dir="intra_E", src=V)
else:
K_T = tl.recv(dir="intra_W",
shape=(d_head, TILE_S_KV), dtype="f16")
V = tl.recv(dir="intra_W",
shape=(TILE_S_KV, d_head), dtype="f16")
if pe_col_in_group < group_cols - 1:
tl.send(dir="intra_E", src=K_T)
tl.send(dir="intra_E", src=V)
# 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) scores = tl.dot(Q, K_T)
m_local = tl.max(scores, axis=-1) m_local = tl.max(scores, axis=-1)
centered = scores - m_local centered = scores - m_local
@@ -75,17 +180,38 @@ def gqa_attention_prefill_short_kernel(
l_local = tl.sum(exp_scores, axis=-1) l_local = tl.sum(exp_scores, axis=-1)
O_local = tl.dot(exp_scores, V) 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 # Tiles 1..n_tiles-1: broadcast + fold into running state.
# each ``copy_to`` with a Python rebind. # ──────────────────────────────────────────────────────────
for tile_idx in range(1, n_tiles): 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(): with tl.scratch_scope():
K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES, if pe_in_group == 0:
shape=(d_head, tile_s), dtype="f16") K_T_t = tl.load(k_head_shard_base + tile_idx * K_TILE_BYTES,
V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES, shape=(d_head, TILE_S_KV), dtype="f16")
shape=(tile_s, d_head), dtype="f16") V_t = tl.load(v_head_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
shape=(TILE_S_KV, d_head), dtype="f16")
if group_cols > 1:
tl.send(dir="intra_E", src=K_T_t)
tl.send(dir="intra_E", src=V_t)
if group_rows > 1:
tl.send(dir="intra_S", src=K_T_t)
tl.send(dir="intra_S", src=V_t)
elif pe_col_in_group == 0 and pe_row_in_group > 0:
K_T_t = tl.recv(dir="intra_N",
shape=(d_head, TILE_S_KV), dtype="f16")
V_t = tl.recv(dir="intra_N",
shape=(TILE_S_KV, d_head), dtype="f16")
if group_cols > 1:
tl.send(dir="intra_E", src=K_T_t)
tl.send(dir="intra_E", src=V_t)
else:
K_T_t = tl.recv(dir="intra_W",
shape=(d_head, TILE_S_KV), dtype="f16")
V_t = tl.recv(dir="intra_W",
shape=(TILE_S_KV, d_head), dtype="f16")
if pe_col_in_group < group_cols - 1:
tl.send(dir="intra_E", src=K_T_t)
tl.send(dir="intra_E", src=V_t)
scores_t = tl.dot(Q, K_T_t) scores_t = tl.dot(Q, K_T_t)
m_tile = tl.max(scores_t, axis=-1) m_tile = tl.max(scores_t, axis=-1)
centered_t = scores_t - m_tile centered_t = scores_t - m_tile
@@ -99,49 +225,10 @@ def gqa_attention_prefill_short_kernel(
tl.copy_to(l_local, l_new) tl.copy_to(l_local, l_new)
tl.copy_to(O_local, O_new) tl.copy_to(O_local, O_new)
# ── Communication: within-group chain reduce-to-root (Level-2 only) ── # No intra-group reduce — q-tiles independent. Each PE writes its own
group_cols = min(4, group_size) # slab to disjoint rows of the cube's column-wise O slab.
group_rows = (group_size + group_cols - 1) // group_cols O_final = O_local / l_local
pe_col_in_group = pe_in_group % group_cols o_base = (o_ptr
pe_row_in_group = pe_in_group // group_cols + cube_id * kv_per_cube * T_q * Q_ROW_BYTES
+ (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
# Row chain (within group's row, along intra_W, leftward). tl.store(o_base, O_final)
if group_cols > 1:
if pe_col_in_group < group_cols - 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_in_group > 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)
# Col bridge (within group, along intra_N, row-1 → row-0).
if pe_col_in_group == 0 and group_rows > 1:
if pe_row_in_group < group_rows - 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_in_group > 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)
# ── Final normalise + store (group root only) ──
if pe_in_group == 0:
O_final = O_local / l_local
tl.store(o_ptr, O_final)