short context length attention kernel
This commit is contained in:
eusonice
@@ -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
|
||||
KV heads, with no S_kv sharding across CUBEs and no inter-CUBE reduce.
|
||||
PE-SP within each CUBE: the ``P`` PEs split into ``kv_per_cube`` groups
|
||||
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.
|
||||
Unified A1/A2/A4/B decode mapping per ADR-0070 (phase mirror of
|
||||
``_gqa_attention_prefill_short.py``). Mode selected at launch via
|
||||
``kv_per_cube ∈ {1, 2, 4, 8}``:
|
||||
|
||||
The local attention uses an S_kv-axis tile sweep (ADR-0063 §A.2) so
|
||||
per-rank scratch is bounded by ``TILE_S_KV``.
|
||||
Mode kv_per_cube C group_size Reduce topology
|
||||
---- ----------- ------- ---------- -----------------------------
|
||||
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:
|
||||
kv_per_cube=1, group=8 PEs (full 2×4): row chain + col bridge.
|
||||
kv_per_cube=2, group=4 PEs (one row): row chain only.
|
||||
kv_per_cube=4, group=2 PEs (adj cols): 1-step chain.
|
||||
kv_per_cube=8, group=1 PE: no chain — direct write.
|
||||
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; each PE owns ``S_local = S_kv/group_size`` KV tokens (sequence
|
||||
shard). FA2 fuses ``G = h_q/h_kv`` Q heads into the M dim of one GEMM
|
||||
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:
|
||||
- K, V: ``(h_kv·S_kv, d_head)`` head-stacked, deployed with
|
||||
``dp = (cube=row_wise, pe=row_wise)`` so each PE's chunk is
|
||||
contiguously ``(S_local, d_head)`` at its own shard. K loaded as
|
||||
``(d_head, S_local)`` via byte-conserving reshape (ADR-0060 §3).
|
||||
- Q: replicated ``(T_q, h_q·d_head)``, reshaped byte-conservingly to
|
||||
``(h_q·T_q, d_head)``. Kernel computes attention for ALL Q rows
|
||||
against the group's owned K head; only the group's owned head rows
|
||||
are semantically meaningful (correct for zero / symmetric inputs).
|
||||
- O: replicated; each group root writes its head's
|
||||
``(h_q·T_q, d_head)`` result at disjoint PE-local addresses.
|
||||
- GEMMs use ``tl.dot`` (no composite epilogue / ``softmax_scale``).
|
||||
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)``, T_q=1.
|
||||
dp=(cube=column_wise, pe=replicate). Caller pre-scales by
|
||||
``1/sqrt(d_head)``.
|
||||
- K: ``(h_kv·S_kv·d_head/TILE_S_KV, TILE_S_KV)`` tile-major.
|
||||
dp=(cube=row_wise, pe=row_wise).
|
||||
- V: ``(h_kv·S_kv, d_head)`` native. Same dp as K.
|
||||
- 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
|
||||
|
||||
|
||||
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):
|
||||
@@ -60,32 +97,40 @@ def gqa_attention_decode_short_kernel(
|
||||
*,
|
||||
tl,
|
||||
) -> 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
|
||||
pe_id = tl.program_id(axis=0)
|
||||
pe_in_group = pe_id % 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) ──
|
||||
Q = tl.load(q_ptr, shape=(h_q * T_q, d_head), dtype="f16")
|
||||
n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
|
||||
KV_ROW_BYTES = d_head * 2 # f16
|
||||
G = h_q // h_kv
|
||||
pe_id = tl.program_id(axis=0)
|
||||
cube_id = tl.program_id(axis=1)
|
||||
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).
|
||||
#
|
||||
# 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")
|
||||
KV_ROW_BYTES = d_head * 2
|
||||
K_TILE_BYTES = d_head * TILE_S_KV * 2
|
||||
Q_ROW_BYTES = G * d_head * 2
|
||||
K_HEAD_BYTES = S_kv * d_head * 2
|
||||
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 * Q_ROW_BYTES + group_id_in_cube * Q_ROW_BYTES
|
||||
k_shard_base = (k_ptr
|
||||
+ cube_id * kv_per_cube * K_HEAD_BYTES
|
||||
+ group_id_in_cube * K_HEAD_BYTES
|
||||
+ pe_in_group * n_tiles_per_pe * K_TILE_BYTES)
|
||||
v_shard_base = (v_ptr
|
||||
+ cube_id * kv_per_cube * V_HEAD_BYTES
|
||||
+ 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)
|
||||
m_local = tl.max(scores, axis=-1)
|
||||
centered = scores - m_local
|
||||
@@ -93,17 +138,13 @@ def gqa_attention_decode_short_kernel(
|
||||
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)
|
||||
# ── Tiles 1..n_tiles_per_pe-1: sweep this PE's sequence shard ──
|
||||
for tile_idx in range(1, n_tiles_per_pe):
|
||||
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")
|
||||
K_T_t = tl.load(k_shard_base + tile_idx * K_TILE_BYTES,
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V_t = tl.load(v_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
|
||||
shape=(TILE_S_KV, 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
|
||||
@@ -117,13 +158,13 @@ def gqa_attention_decode_short_kernel(
|
||||
tl.copy_to(l_local, l_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_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
|
||||
|
||||
# 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 pe_col_in_group < group_cols - 1:
|
||||
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=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_row_in_group < group_rows - 1:
|
||||
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=O_local)
|
||||
|
||||
# ── Final normalise + store (group root only) ──
|
||||
# ── Final normalize + store (group root, pe_in_group == 0) ──
|
||||
if pe_in_group == 0:
|
||||
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)
|
||||
|
||||
@@ -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
|
||||
(``kv_per_cube`` heads per CUBE, group-PE-SP, within-group chain reduce,
|
||||
no inter-CUBE reduce). The only structural difference from short decode:
|
||||
``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.
|
||||
Unified A1/A2/A4/B prefill mapping per ADR-0070 (supersedes
|
||||
ADR-0060 §B.split.2 prefill-short clause). Mode selected at launch
|
||||
via ``kv_per_cube ∈ {1, 2, 4, 8}``:
|
||||
|
||||
The local attention uses an S_kv-axis tile sweep (ADR-0063 §A.2) so
|
||||
per-rank scratch is bounded by ``TILE_S_KV``.
|
||||
Mode kv_per_cube C group_size Broadcast topology
|
||||
---- ----------- ------- ---------- ---------------------------
|
||||
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
|
||||
|
||||
@@ -17,6 +47,33 @@ from __future__ import annotations
|
||||
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):
|
||||
"""Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
|
||||
m_new = tl.maximum(m_local, m_other)
|
||||
@@ -34,6 +91,7 @@ def gqa_attention_prefill_short_kernel(
|
||||
o_ptr: int,
|
||||
T_q: int,
|
||||
S_kv: int,
|
||||
h_q: int,
|
||||
h_kv: int,
|
||||
d_head: int,
|
||||
C: int,
|
||||
@@ -42,32 +100,79 @@ def gqa_attention_prefill_short_kernel(
|
||||
*,
|
||||
tl,
|
||||
) -> 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
|
||||
|
||||
G = h_q // h_kv
|
||||
pe_id = tl.program_id(axis=0)
|
||||
cube_id = tl.program_id(axis=1)
|
||||
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) ──
|
||||
Q = tl.load(q_ptr, shape=(h_kv * T_q, d_head), dtype="f16")
|
||||
n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
|
||||
KV_ROW_BYTES = d_head * 2 # f16
|
||||
# Floor-balanced Q-tile partition.
|
||||
q_start = (T_q * pe_in_group) // group_size
|
||||
q_end = (T_q * (pe_in_group + 1)) // group_size
|
||||
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)
|
||||
m_local = tl.max(scores, axis=-1)
|
||||
centered = scores - m_local
|
||||
@@ -75,17 +180,38 @@ def gqa_attention_prefill_short_kernel(
|
||||
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.
|
||||
# ──────────────────────────────────────────────────────────
|
||||
# Tiles 1..n_tiles-1: broadcast + fold into running state.
|
||||
# ──────────────────────────────────────────────────────────
|
||||
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")
|
||||
if pe_in_group == 0:
|
||||
K_T_t = tl.load(k_head_shard_base + tile_idx * K_TILE_BYTES,
|
||||
shape=(d_head, TILE_S_KV), 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)
|
||||
m_tile = tl.max(scores_t, axis=-1)
|
||||
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(O_local, O_new)
|
||||
|
||||
# ── Communication: within-group chain reduce-to-root (Level-2 only) ──
|
||||
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
|
||||
|
||||
# Row chain (within group's row, along intra_W, leftward).
|
||||
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:
|
||||
# No intra-group reduce — q-tiles independent. Each PE writes its own
|
||||
# slab to disjoint rows of the cube's column-wise O slab.
|
||||
O_final = O_local / l_local
|
||||
tl.store(o_ptr, O_final)
|
||||
o_base = (o_ptr
|
||||
+ cube_id * kv_per_cube * T_q * Q_ROW_BYTES
|
||||
+ (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
|
||||
tl.store(o_base, O_final)
|
||||
|
||||
Reference in New Issue
Block a user