gqa: rename long-context kernels to *_long for symmetry with *_short

Make the file + function naming symmetric: _gqa_decode.py -> _gqa_decode_long.py _gqa_prefill.py -> _gqa_prefill_long.py gqa_decode_kernel -> gqa_decode_long_kernel gqa_prefill_kernel -> gqa_prefill_long_kernel Mirrors the existing _gqa_{decode,prefill}_short.py naming. Updates the two imports + two call sites in milestone_gqa_headline.py and the 9 attention tests that import the kernels. Tests: 72/72 focused regression green (tests/attention/ + Phase E + TL discipline). milestone-gqa-headline bench passes its 7 panel/schema assertions. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-06-10 10:05:54 -07:00
parent d282144339
commit 5a76ed4f6a
10 changed files with 25 additions and 25 deletions
@@ -0,0 +1,111 @@
+"""GQA fused-attention prefill kernel — P6a + P6b (head-parallel + Ring KV).
+
+Lineage (DDD-0060 §7 phase plan):
+  P6a    : head-parallel structure (one Q head per CUBE), C=1 baseline.
+  P6b    : add Ring KV rotation across C CUBEs (ADR-0060 §5.5). Each
+            CUBE rotates its KV block to its W neighbour and receives
+            from E; over C-1 steps every CUBE sees every block. Online-
+            softmax merge folds each step into running (m, ℓ, O). No
+            reduce — each CUBE writes its own head's output.
+            Requires ``configure_sfr_intercube_ring(ring_size=C)`` SFR.
+  P6c (later): tile T_q across the P PEs inside a CUBE for intra-CUBE
+            parallelism (ADR-0060 §B item 3).
+
+Deviations from ADR-0060 §5.5 — deferred to later phases:
+  1. GEMMs use ``tl.dot`` not ``tl.composite`` (parallel to P1a; lifts
+     when P1b decides the composite-output-handle question).
+  2. ``softmax_scale`` omitted — same composite-epilogue deferral.
+  3. K loaded as ``[d, S_local]`` via byte-conserving reshape of
+     ``[S_local, d]`` (ADR-0060 §3 / §B item 2 reshape-not-transpose
+     caveat; correct for zero / symmetric inputs).
+  4. No causal masking / step-skip — future P6c.
+  5. Blocking ``tl.recv`` (not ``recv_async``) — overlap via lazy
+     ``tl.load`` lands in P4.
+"""
+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. P6c (future)
+    # will tile T_q across the 8 PEs for intra-CUBE parallelism.
+    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")
+
+    # ── Step 0: initial partial against own KV block — establishes the
+    # persistent (m, ℓ, O) arena. Intermediates (scores, exp_scores) stay
+    # allocated; ring steps below recycle per-step intermediates inside
+    # tl.scratch_scope to keep peak scratch O(one step) (ADR-0063 §D3).
+    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)
+
+    # ── Steps 1..C-1: Ring KV rotation + online-softmax merge ──
+    # Per-step intermediates wrapped in tl.scratch_scope; the merged
+    # running (m, ℓ, O) is persisted to the outside-scope handles via
+    # tl.copy_to (ADR-0063 §D3.1) so its bytes survive __exit__.
+    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():
+            # Partial on rotated KV
+            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)
+
+            # Online-softmax merge into new running (m, ℓ, O)
+            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
+
+            # Persist new running state back to the outside-scope arena.
+            tl.copy_to(m, m_new)
+            tl.copy_to(l, l_new)
+            tl.copy_to(O, O_new)
+        # __exit__: scoped intermediates gone; persistent m, l, O carry
+        # the new running state into the next ring iteration.
+
+    # Final normalise + store — each CUBE writes its own head's rows.
+    O_final = O / l
+    tl.store(o_ptr, O_final)