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:
@@ -24,8 +24,8 @@ After chain reduce-to-group-root, the group's root PE writes its
|
||||
owned head's output. No inter-CUBE reduce.
|
||||
|
||||
Phase 1 (this commit): tests only — production code lands in Phase 2.
|
||||
All tests fail because the short kernels (``_gqa_decode_short.py`` and
|
||||
``_gqa_prefill_short.py``) do not exist yet.
|
||||
All tests fail because the short kernels (``_gqa_attention_decode_short.py`` and
|
||||
``_gqa_attention_prefill_short.py``) do not exist yet.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -76,7 +76,7 @@ def _run_decode_short(*, h_q: int, h_kv: int, kv_per_cube: int,
|
||||
O: replicated; each group root writes the full byte-conserving
|
||||
(h_q·T_q, D_HEAD) result.
|
||||
"""
|
||||
from kernbench.benches._gqa_decode_short import gqa_decode_short_kernel # Phase 2
|
||||
from kernbench.benches._gqa_attention_decode_short import gqa_attention_decode_short_kernel # Phase 2
|
||||
|
||||
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
|
||||
|
||||
@@ -98,7 +98,7 @@ def _run_decode_short(*, h_q: int, h_kv: int, kv_per_cube: int,
|
||||
dtype=DTYPE, dp=dp_full, name=f"o_short_{kv_per_cube}_{C}")
|
||||
ctx.launch(
|
||||
f"gqa_decode_short_{kv_per_cube}_{C}",
|
||||
gqa_decode_short_kernel,
|
||||
gqa_attention_decode_short_kernel,
|
||||
q, k, v, o,
|
||||
1, S_kv, h_q, h_kv, D_HEAD, C, P, kv_per_cube,
|
||||
_auto_dim_remap=False,
|
||||
@@ -196,7 +196,7 @@ def _run_prefill_short(*, h_kv: int, kv_per_cube: int,
|
||||
Layout: same head-stacked K/V scheme as decode short, with Q/O
|
||||
replicated and byte-conserving reshape inside the kernel.
|
||||
"""
|
||||
from kernbench.benches._gqa_prefill_short import gqa_prefill_short_kernel # Phase 2
|
||||
from kernbench.benches._gqa_attention_prefill_short import gqa_attention_prefill_short_kernel # Phase 2
|
||||
|
||||
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
|
||||
|
||||
@@ -218,7 +218,7 @@ def _run_prefill_short(*, h_kv: int, kv_per_cube: int,
|
||||
dtype=DTYPE, dp=dp_o, name=f"o_pre_short_{kv_per_cube}_{C}")
|
||||
ctx.launch(
|
||||
f"gqa_prefill_short_{kv_per_cube}_{C}",
|
||||
gqa_prefill_short_kernel,
|
||||
gqa_attention_prefill_short_kernel,
|
||||
q, k, v, o,
|
||||
T_q, S_kv, h_kv, D_HEAD, C, P, kv_per_cube,
|
||||
_auto_dim_remap=False,
|
||||
|
||||
Reference in New Issue
Block a user