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:
+1
-1
@@ -35,7 +35,7 @@ def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
|
||||
return m_new, l_new, O_new
|
||||
|
||||
|
||||
def gqa_decode_long_kernel(
|
||||
def gqa_attention_decode_long_kernel(
|
||||
q_ptr: int,
|
||||
k_ptr: int,
|
||||
v_ptr: int,
|
||||
+1
-1
@@ -44,7 +44,7 @@ def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
|
||||
return m_new, l_new, O_new
|
||||
|
||||
|
||||
def gqa_decode_short_kernel(
|
||||
def gqa_attention_decode_short_kernel(
|
||||
q_ptr: int,
|
||||
k_ptr: int,
|
||||
v_ptr: int,
|
||||
@@ -0,0 +1,151 @@
|
||||
"""GQA fused-attention prefill kernel — long context (ADR-0060 §5.5).
|
||||
|
||||
Head-parallel: each CUBE owns one Q head and one KV slice. Over C ring
|
||||
steps the KV blocks rotate around the C CUBEs (W-send / E-recv) so every
|
||||
CUBE sees every block; the online-softmax merge folds each step into the
|
||||
running ``(m, ℓ, O)``. No inter-CUBE reduce — each CUBE writes its own
|
||||
head's output.
|
||||
|
||||
The ring is **tile-granular** (ADR-0060 §5.5 + ADR-0063 §A.2): each
|
||||
ring step transmits ``n_tiles`` tiles of size ``TILE_S_KV`` rather than
|
||||
one full ``S_local`` slice. The kernel loop is nested over
|
||||
``(ring_step, tile_idx)``, with the bootstrap at ``(k=0, t=0)``
|
||||
establishing the persistent ``(m, ℓ, O)`` and every subsequent
|
||||
iteration folding its tile in inside a ``tl.scratch_scope``. Per-rank
|
||||
persistent scratch is ``(m, ℓ, O)`` only (~1 KB); per-tile in-scope
|
||||
scratch is bounded by ``TILE_S_KV`` regardless of ``S_local``.
|
||||
|
||||
Topology / SFR:
|
||||
- Requires ``configure_sfr_intercube_ring(ring_size=C)`` (1D ring of
|
||||
C CUBEs with wrap at the CUBE level).
|
||||
- Only PE 0 of each CUBE participates (head-parallel; intra-CUBE PE
|
||||
parallelism is a separate phase).
|
||||
|
||||
Layout caveats:
|
||||
- GEMMs use ``tl.dot`` (no composite epilogue / ``softmax_scale``).
|
||||
- K loaded as ``(d_head, tile_s)`` via byte-conserving reshape of
|
||||
the deployed ``(tile_s, d_head)`` shard (ADR-0060 §3
|
||||
reshape-not-transpose caveat).
|
||||
- No causal masking / step-skip; blocking ``tl.recv``.
|
||||
- ``TILE_S_KV`` is assumed to divide ``S_local`` evenly; last-tile
|
||||
padding is not modelled.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
|
||||
|
||||
|
||||
def _merge_running(m, l, O, m_step, l_step, O_step, *, tl):
|
||||
"""Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
|
||||
m_new = tl.maximum(m, m_step)
|
||||
scale_old = tl.exp(m - m_new)
|
||||
scale_step = tl.exp(m_step - m_new)
|
||||
l_new = l * scale_old + l_step * scale_step
|
||||
O_new = O * scale_old + O_step * scale_step
|
||||
return m_new, l_new, O_new
|
||||
|
||||
|
||||
def gqa_attention_prefill_long_kernel(
|
||||
q_ptr: int,
|
||||
k_ptr: int,
|
||||
v_ptr: int,
|
||||
o_ptr: int,
|
||||
T_q: int,
|
||||
S_kv: int,
|
||||
d_head: int,
|
||||
C: int,
|
||||
*,
|
||||
tl,
|
||||
) -> None:
|
||||
"""Head-parallel prefill attention with tile-granular Ring KV (ADR-0060 §5.5).
|
||||
|
||||
Tensor layout:
|
||||
Q : (T_q, d_head) one head per CUBE; replicated.
|
||||
K : (S_kv, d_head) sharded cube_row_wise → each CUBE owns
|
||||
(S_local, d_head). Tiles loaded as (d_head, tile_s) via
|
||||
byte-conserving reshape.
|
||||
V : (S_kv, d_head) sharded cube_row_wise → each CUBE owns
|
||||
(S_local, d_head). Tiles loaded as (tile_s, d_head).
|
||||
O : (T_q * C, d_head) sharded cube_row_wise → each CUBE
|
||||
writes its own (T_q, d_head) slice. NO reduce.
|
||||
|
||||
Algorithm: nested loop over (ring_step k, tile_idx t). At k=0 each
|
||||
CUBE loads its own block's tiles from HBM; at k > 0 it receives
|
||||
tiles from its E neighbour. Tiles are forwarded W to the next
|
||||
ring step. Each tile's partial is folded into the running
|
||||
(m, ℓ, O) via online-softmax merge.
|
||||
"""
|
||||
pe_id = tl.program_id(axis=0)
|
||||
# Head-parallel: only PE 0 of each CUBE participates.
|
||||
if pe_id != 0:
|
||||
return
|
||||
|
||||
S_local = S_kv // C
|
||||
n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
|
||||
KV_ROW_BYTES = d_head * 2 # f16
|
||||
Q = tl.load(q_ptr, shape=(T_q, d_head), dtype="f16")
|
||||
|
||||
# ── Bootstrap: (t=0, k=0) — load my own tile 0, establish (m, ℓ, O) ──
|
||||
# Outside ``scratch_scope`` so the persistent running state survives
|
||||
# subsequent ring iterations.
|
||||
tile_s = min(TILE_S_KV, S_local)
|
||||
K_t = tl.load(k_ptr, shape=(d_head, tile_s), dtype="f16")
|
||||
V_t = tl.load(v_ptr, shape=(tile_s, d_head), dtype="f16")
|
||||
if C > 1:
|
||||
tl.send(dir="W", src=K_t)
|
||||
tl.send(dir="W", src=V_t)
|
||||
scores = tl.dot(Q, K_t)
|
||||
m = tl.max(scores, axis=-1)
|
||||
exp_scores = tl.exp(scores - m)
|
||||
l = tl.sum(exp_scores, axis=-1)
|
||||
O = tl.dot(exp_scores, V_t)
|
||||
|
||||
# ── Nested loop: outer tile, inner ring step ──
|
||||
# Each tile propagates all the way through the ring before the next
|
||||
# tile starts; IPCQ in-flight depth stays at 1 per direction.
|
||||
#
|
||||
# Per outer ``t``:
|
||||
# k=0 : load my own tile ``t`` from HBM
|
||||
# k=1..C-1 : receive tile ``t`` of a peer's block from E
|
||||
# (sent by my E neighbour during their previous k step)
|
||||
# k<C-1 : forward the tile to W
|
||||
#
|
||||
# The ``(t=0, k=0)`` iteration is the bootstrap above; the loop skips
|
||||
# it via ``k_start``. Every other iteration uses ``scratch_scope`` +
|
||||
# ``copy_to`` to recycle per-tile intermediates.
|
||||
#
|
||||
# Triton port: drop ``with tl.scratch_scope():`` and replace each
|
||||
# ``copy_to`` with a Python rebind (e.g. ``m = m_new``).
|
||||
for t in range(n_tiles):
|
||||
tile_start = t * TILE_S_KV
|
||||
tile_s = min(TILE_S_KV, S_local - tile_start)
|
||||
k_start = 1 if t == 0 else 0
|
||||
for k in range(k_start, C):
|
||||
with tl.scratch_scope():
|
||||
if k == 0:
|
||||
K_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")
|
||||
else:
|
||||
K_t = tl.recv(dir="E", shape=(d_head, tile_s), dtype="f16")
|
||||
V_t = tl.recv(dir="E", shape=(tile_s, d_head), dtype="f16")
|
||||
if k < C - 1:
|
||||
tl.send(dir="W", src=K_t)
|
||||
tl.send(dir="W", src=V_t)
|
||||
scores = tl.dot(Q, K_t)
|
||||
m_step = tl.max(scores, axis=-1)
|
||||
exp_scores = tl.exp(scores - m_step)
|
||||
l_step = tl.sum(exp_scores, axis=-1)
|
||||
O_step = tl.dot(exp_scores, V_t)
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m, l, O, m_step, l_step, O_step, tl=tl,
|
||||
)
|
||||
tl.copy_to(m, m_new)
|
||||
tl.copy_to(l, l_new)
|
||||
tl.copy_to(O, O_new)
|
||||
|
||||
# ── Final normalise + store (each CUBE writes its own head) ──
|
||||
O_final = O / l
|
||||
tl.store(o_ptr, O_final)
|
||||
+1
-1
@@ -27,7 +27,7 @@ def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
|
||||
return m_new, l_new, O_new
|
||||
|
||||
|
||||
def gqa_prefill_short_kernel(
|
||||
def gqa_attention_prefill_short_kernel(
|
||||
q_ptr: int,
|
||||
k_ptr: int,
|
||||
v_ptr: int,
|
||||
@@ -1,104 +0,0 @@
|
||||
"""GQA fused-attention prefill kernel — long context (ADR-0060 §5.5).
|
||||
|
||||
Head-parallel: each CUBE owns one Q head and one KV slice. Over C ring
|
||||
steps the KV blocks rotate around the C CUBEs (W-send / E-recv) so every
|
||||
CUBE sees every block; the online-softmax merge folds each step into the
|
||||
running ``(m, ℓ, O)``. No inter-CUBE reduce — each CUBE writes its own
|
||||
head's output.
|
||||
|
||||
Topology / SFR:
|
||||
- Requires ``configure_sfr_intercube_ring(ring_size=C)`` (1D ring of
|
||||
C CUBEs with wrap at the CUBE level).
|
||||
- Only PE 0 of each CUBE participates (head-parallel; intra-CUBE PE
|
||||
parallelism is a separate phase).
|
||||
|
||||
Layout caveats:
|
||||
- GEMMs use ``tl.dot`` (no composite epilogue / ``softmax_scale``).
|
||||
- K loaded as ``(d_head, S_local)`` via byte-conserving reshape of
|
||||
the deployed ``(S_local, d_head)`` slice (ADR-0060 §3
|
||||
reshape-not-transpose caveat).
|
||||
- No causal masking / step-skip; blocking ``tl.recv``.
|
||||
- Step-0 local partial is NOT tile-swept — its scratch bound is set
|
||||
by the ring's full-slice ``Kc``/``Vc`` carry, which dominates the
|
||||
score-stack regardless. Tile-granular ring is a separate phase.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
def gqa_prefill_long_kernel(
|
||||
q_ptr: int,
|
||||
k_ptr: int,
|
||||
v_ptr: int,
|
||||
o_ptr: int,
|
||||
T_q: int,
|
||||
S_kv: int,
|
||||
d_head: int,
|
||||
C: int,
|
||||
*,
|
||||
tl,
|
||||
) -> None:
|
||||
"""Head-parallel prefill attention with Ring KV (C>1) — ADR-0060 §5.5.
|
||||
|
||||
Tensor layout consumed by this kernel:
|
||||
Q : (T_q, d_head) one head per CUBE; replicated.
|
||||
K : (S_kv, d_head) sharded cube_row_wise → each CUBE owns
|
||||
(S_kv/C, d_head); kernel loads as (d_head, S_local) via
|
||||
byte-conserving reshape (reshape-not-transpose caveat).
|
||||
V : (S_kv, d_head) sharded cube_row_wise → each CUBE owns
|
||||
(S_local, d_head).
|
||||
O : (T_q * C, d_head) sharded cube_row_wise → each CUBE
|
||||
writes its own (T_q, d_head) slice. NO reduce.
|
||||
|
||||
Algorithm: each CUBE computes a local partial against its current
|
||||
KV block, then over C-1 ring steps the K and V blocks rotate W
|
||||
while online-softmax merges each step into running (m, ℓ, O).
|
||||
"""
|
||||
pe_id = tl.program_id(axis=0)
|
||||
# Head-parallel: only PE 0 of each CUBE participates.
|
||||
if pe_id != 0:
|
||||
return
|
||||
|
||||
S_local = S_kv // C
|
||||
Q = tl.load(q_ptr, shape=(T_q, d_head), dtype="f16")
|
||||
Kc = tl.load(k_ptr, shape=(d_head, S_local), dtype="f16")
|
||||
Vc = tl.load(v_ptr, shape=(S_local, d_head), dtype="f16")
|
||||
|
||||
# ── Local attention: initial partial against own KV block ──
|
||||
# Establishes the persistent (m, ℓ, O) running state.
|
||||
scores = tl.dot(Q, Kc)
|
||||
m = tl.max(scores, axis=-1)
|
||||
exp_scores = tl.exp(scores - m)
|
||||
l = tl.sum(exp_scores, axis=-1)
|
||||
O = tl.dot(exp_scores, Vc)
|
||||
|
||||
# ── Communication: Ring KV rotation + online-softmax merge ──
|
||||
# Each step sends K, V to W and receives from E. Per-step
|
||||
# intermediates are scope-recycled; the merged (m, ℓ, O) is
|
||||
# persisted via tl.copy_to. Triton port: drop the scope and
|
||||
# replace each copy_to with a Python rebind.
|
||||
for _ in range(1, C):
|
||||
tl.send(dir="W", src=Kc)
|
||||
tl.send(dir="W", src=Vc)
|
||||
Kc = tl.recv(dir="E", shape=(d_head, S_local), dtype="f16")
|
||||
Vc = tl.recv(dir="E", shape=(S_local, d_head), dtype="f16")
|
||||
|
||||
with tl.scratch_scope():
|
||||
scores = tl.dot(Q, Kc)
|
||||
m_step = tl.max(scores, axis=-1)
|
||||
exp_scores = tl.exp(scores - m_step)
|
||||
l_step = tl.sum(exp_scores, axis=-1)
|
||||
O_step = tl.dot(exp_scores, Vc)
|
||||
|
||||
m_new = tl.maximum(m, m_step)
|
||||
scale_old = tl.exp(m - m_new)
|
||||
scale_step = tl.exp(m_step - m_new)
|
||||
l_new = l * scale_old + l_step * scale_step
|
||||
O_new = O * scale_old + O_step * scale_step
|
||||
|
||||
tl.copy_to(m, m_new)
|
||||
tl.copy_to(l, l_new)
|
||||
tl.copy_to(O, O_new)
|
||||
|
||||
# ── Final normalise + store (each CUBE writes its own head) ──
|
||||
O_final = O / l
|
||||
tl.store(o_ptr, O_final)
|
||||
@@ -28,8 +28,8 @@ import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from kernbench.benches._gqa_decode_long import gqa_decode_long_kernel
|
||||
from kernbench.benches._gqa_prefill_long import gqa_prefill_long_kernel
|
||||
from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel
|
||||
from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel
|
||||
from kernbench.benches.registry import bench
|
||||
from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
|
||||
from kernbench.ccl.sfr_config import (
|
||||
@@ -96,7 +96,7 @@ def _run_prefill_panel(ctx, *, panel: str, C: int, S_kv: int) -> None:
|
||||
o = ctx.empty((_T_Q_PREFILL * C, _D_HEAD),
|
||||
dtype=_DTYPE, dp=dp_o, name=f"{panel}_o")
|
||||
ctx.launch(
|
||||
panel, gqa_prefill_long_kernel,
|
||||
panel, gqa_attention_prefill_long_kernel,
|
||||
q, k, v, o,
|
||||
_T_Q_PREFILL, S_kv, _D_HEAD, C,
|
||||
_auto_dim_remap=False,
|
||||
@@ -118,7 +118,7 @@ def _run_decode_panel(ctx, *, panel: str, C: int, P: int, S_kv: int) -> None:
|
||||
o = ctx.empty((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD),
|
||||
dtype=_DTYPE, dp=dp_full, name=f"{panel}_o")
|
||||
ctx.launch(
|
||||
panel, gqa_decode_long_kernel,
|
||||
panel, gqa_attention_decode_long_kernel,
|
||||
q, k, v, o,
|
||||
_T_Q_DECODE, S_kv, _H_Q_DECODE, _H_KV_DECODE, _D_HEAD, C, P,
|
||||
_auto_dim_remap=False,
|
||||
|
||||
Reference in New Issue
Block a user