kernbench2/src/kernbench/benches/_gqa_prefill_long.py

"""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)