gqa: tile-granular Ring KV (P3c) + rename to gqa_attention_* + ADR-0060/62/63/64 → Accepted
Three logically distinct changes, bundled for atomic test green:
1. **P3c — prefill_long tile-granular Ring KV** (ADR-0060 §5.5.1 amendment).
Convert the ring from slice-granular (one full ``(d_head, S_local)``
KV slice per step) to tile-granular (``n_tiles`` tiles of
``TILE_S_KV`` per step). Nested loop with outer tile, inner ring step:
each tile propagates through all C ring positions before the next
tile starts, so IPCQ in-flight depth stays at 1 per direction.
Bootstrap at ``(t=0, k=0)`` outside the scratch_scope establishes the
persistent ``(m, ℓ, O)``; every other iteration scope-wraps + persists
via ``copy_to``. Per-rank persistent scratch shrinks to ~1 KB; per-tile
scope bounded by TILE_S_KV regardless of S_local. Headline:
prefill_long now completes at S_kv=128K (previously overflowed).
New: ``tests/attention/test_gqa_prefill_long_tile_ring.py``
(3 tests — ceiling-lift + tile-granular ipcq_copy count +
per-CUBE distributed output regression guard).
2. **Rename ``gqa_*`` → ``gqa_attention_*``** across kernel files,
function names, and importers. The "attention" name makes the role
explicit (GQA is grouped-query attention) and matches upstream Triton
FlashAttention naming conventions. Renames:
_gqa_decode_long.py -> _gqa_attention_decode_long.py
_gqa_decode_short.py -> _gqa_attention_decode_short.py
_gqa_prefill_long.py -> _gqa_attention_prefill_long.py
_gqa_prefill_short.py -> _gqa_attention_prefill_short.py
And function names ``gqa_<phase>_<context>_kernel`` →
``gqa_attention_<phase>_<context>_kernel``. Updated 1 bench file
(milestone_gqa_headline.py) and 10 test files.
3. **ADR-0060 / 0062 / 0063 / 0064: Proposed → Accepted**.
All four are reflected in production code and covered by tests:
- ADR-0060 (GQA fused attention): 4 kernels deployed; §5.5.1
amendment added for the tile-granular Ring KV introduced by P3c
(EN + KO mirror).
- ADR-0062 (lazy tl.load): LoadFuture + _await_pending live in
tl_context.py.
- ADR-0063 (tl.scratch_scope + tl.copy_to): used in every chain
reduce + tile sweep + ring step. EN-only previously; KO
translation authored as part of this commit (CLAUDE.md
bidirectional rule).
- ADR-0064 (per-op-type CPU issue cost): cpu_issue_cost.py +
issue_cost_table wiring in tl_context.py (Phase E).
Files git mv'd from docs/adr-proposed/ to docs/adr/ (EN) and
docs/adr-ko/ (KO). ADR-0061 (tl.broadcast) stays Proposed — no
implementation; documented as optional convenience primitive in
the ADR itself.
Tests: 88/88 focused regression green
(tests/attention/ + Phase E + TL discipline).
ADR pair verification: ``python tools/verify_adr_lang_pairs.py`` OK.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,154 @@
|
||||
"""GQA fused-attention decode kernel — short context (ADR-0060 §B.split.2).
|
||||
|
||||
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.
|
||||
|
||||
The local attention uses an S_kv-axis tile sweep (ADR-0063 §A.2) so
|
||||
per-rank scratch is bounded by ``TILE_S_KV``.
|
||||
|
||||
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.
|
||||
|
||||
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``).
|
||||
"""
|
||||
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_short_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,
|
||||
kv_per_cube: int,
|
||||
*,
|
||||
tl,
|
||||
) -> None:
|
||||
"""Short-context GQA decode with PE-parallel heads + intra-group PE-SP."""
|
||||
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
|
||||
|
||||
# ── 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
|
||||
|
||||
# Tile 0: establishes persistent (m_local, l_local, O_local).
|
||||
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: 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)
|
||||
Reference in New Issue
Block a user