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>
2026-06-10 16:17:32 -07:00
parent a8c50238c6
commit 7fad0371c5
22 changed files with 667 additions and 152 deletions
@@ -0,0 +1,174 @@
+"""GQA fused-attention decode kernel — long context (ADR-0060).
+
+Each rank holds an ``S_local = S_kv / (C·P)`` slice of K, V and the full
+Q (replicated). The local attention is computed via an S_kv-axis tile
+sweep (ADR-0063 §A.2) so per-rank scratch is bounded by ``TILE_S_KV``
+regardless of ``S_local``. The partial ``(m, ℓ, O)`` is then reduced
+to PE 0 of CUBE 0 via a 2-level chain (intra-CUBE row+col, then
+inter-CUBE), and the root writes the final output.
+
+Topology / SFR:
+  - Requires ``configure_sfr_intercube_multisip`` when ``P > 1`` or
+    ``C > 1`` (provides disjoint ``intra_*`` and ``E/W/N/S`` namespaces).
+  - Intra-CUBE PEs are arranged as a 2×4 grid (no wrap).
+  - Inter-CUBE CUBEs are arranged as a 1D row (no wrap).
+
+Layout caveats:
+  - GEMMs use ``tl.dot`` (no composite epilogue / ``softmax_scale``).
+  - K is loaded as ``(d_head, S_local)`` via byte-conserving reshape of
+    the deployed ``(S_local, h_kv·d_head)`` slice — correct for zero /
+    symmetric inputs (ADR-0060 §3 reshape-not-transpose caveat).
+"""
+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_long_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,
+    *,
+    tl,
+) -> None:
+    """GQA decode with M-fold + S_kv tile sweep + 2-level chain reduce-to-root.
+
+    Tensor layout:
+      Q : (T_q, h_q · d_head) replicated on every rank; loaded as
+          (G·T_q, d_head) — byte-conserving and math-correct for T_q=1.
+      K : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank
+          loads its (d_head, S_local) slice via byte-conserving reshape.
+      V : (S_kv, h_kv · d_head) sharded row_wise by (cube, pe); each rank
+          loads its (S_local, d_head) slice.
+      O : (T_q, h_q · d_head) — only PE 0 of CUBE 0 stores.
+    """
+    G = h_q // h_kv
+    n_ranks = C * P
+    S_local = S_kv // n_ranks
+    pe_id = tl.program_id(axis=0)
+    cube_id = tl.program_id(axis=1)
+
+    # ── Local attention (S_kv-axis tile sweep, ADR-0063 §A.2) ──
+    Q = tl.load(q_ptr, shape=(G * 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: chain reduce-to-root at PE 0 of CUBE 0 ──
+    PE_GRID_COLS = 4
+    pe_col = pe_id % PE_GRID_COLS
+    pe_row = pe_id // PE_GRID_COLS
+    pe_cols_used = min(PE_GRID_COLS, P)
+    pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS
+
+    # Level-2 row chain (intra-CUBE, along intra_W, leftward).
+    if pe_cols_used > 1:
+        if pe_col < pe_cols_used - 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 > 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)
+
+    # Level-2 col bridge (intra-CUBE, along intra_N, row-1 → row-0).
+    if pe_col == 0 and pe_rows_used > 1:
+        if pe_row < pe_rows_used - 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 > 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)
+
+    # Level-1 inter-CUBE chain (along W, leftward; only PE 0 of each CUBE).
+    if pe_id == 0 and C > 1:
+        if cube_id < C - 1:
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="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 cube_id > 0:
+            tl.send(dir="W", src=m_local)
+            tl.send(dir="W", src=l_local)
+            tl.send(dir="W", src=O_local)
+
+    # ── Final normalise + store (root only) ──
+    if pe_id == 0 and cube_id == 0:
+        O_final = O_local / l_local
+        tl.store(o_ptr, O_final)