From 7d90d53d4d2700985e25cefac8c9e99093d2ce10 Mon Sep 17 00:00:00 2001 From: eusonice <98355549+eusonice@users.noreply.github.com> Date: Tue, 23 Jun 2026 16:19:39 -0700 Subject: [PATCH] short context length attention kernel --- .../short_ctx/_gqa_attention_decode_short.py | 170 +++++++----- .../short_ctx/_gqa_attention_prefill_short.py | 257 ++++++++++++------ 2 files changed, 279 insertions(+), 148 deletions(-) diff --git a/src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_decode_short.py b/src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_decode_short.py index 48919f6..8dd5526 100644 --- a/src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_decode_short.py +++ b/src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_decode_short.py @@ -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) diff --git a/src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_prefill_short.py b/src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_prefill_short.py index c9a9c81..59fdbcc 100644 --- a/src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_prefill_short.py +++ b/src/kernbench/benches/gqa_helpers/short_ctx/_gqa_attention_prefill_short.py @@ -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: - O_final = O_local / l_local - tl.store(o_ptr, O_final) + # 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 + 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)