"""GQA decode kernel — Case 2 (Cube-Repl × PE-TP).

Per GQA_full_deck.pptx slide 11:
  - K, V replicated across all 8 cubes × 8 PEs (the 8 KB/tok/PE
    memory waste — this is the inherent cost Case 2 demonstrates).
  - PEs nominally split on the batch dim (PE-TP). For B=1 (single-
    user decode, the slide-11 default), only PE 0 of CUBE 0 has
    work; the other 63 ranks idle.
  - NO inter-rank communication (each rank has full KV — slide 11
    lists comm cost as "none").

This kernel is the simplest of the 4 cases by design: one active
rank does the full attention locally; everyone else early-returns.
The DPPolicy at the call site models the cluster-wide 8× memory
waste even though only one rank reads from HBM.

Tensor layout (B=1):
  Q : (T_q, h_q · d_head) replicated on every rank; loaded as
      (G · T_q, d_head) on the active rank.
  K : (S_kv, h_kv · d_head) replicated on every rank.
  V : (S_kv, h_kv · d_head) replicated on every rank.
  O : (T_q, h_q · d_head) — only PE 0 of CUBE 0 stores.

Topology / SFR:
  - configure_sfr_intercube_multisip is fine but not strictly required
    (no inter-rank sends/recvs happen).
"""
from __future__ import annotations


TILE_S_KV = 1024  # match decode_long — per-tile S_kv width (ADR-0063 §A.2).


def gqa_attention_decode_long_ctx_cube_repl_pe_tp_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:
    """Case-2 decode: single-rank attention; full KV per rank; no comm."""
    pe_id = tl.program_id(axis=0)
    cube_id = tl.program_id(axis=1)
    # B=1 single-user decode + PE-TP: only one rank has work.
    # Slide-11 acknowledges this PE-TP waste at B=1.
    if pe_id != 0 or cube_id != 0:
        return

    G = h_q // h_kv
    n_tiles = (S_kv + TILE_S_KV - 1) // TILE_S_KV
    KV_ROW_BYTES = d_head * 2  # f16

    # ── Load Q (full; replicated; M-folded for GQA reuse) ──
    Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")

    # ── Tile 0: bootstrap persistent (m, ℓ, O) ──
    tile_s0 = min(TILE_S_KV, S_kv)
    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: fold via online-softmax merge in scratch_scope ──
    for tile_idx in range(1, n_tiles):
        tile_start = tile_idx * TILE_S_KV
        tile_s = min(TILE_S_KV, S_kv - 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 = tl.maximum(m_local, m_tile)
            scale_old = tl.exp(m_local - m_new)
            scale_new = tl.exp(m_tile - m_new)
            l_new = l_local * scale_old + l_tile * scale_new
            O_new = O_local * scale_old + O_tile * scale_new
            tl.copy_to(m_local, m_new)
            tl.copy_to(l_local, l_new)
            tl.copy_to(O_local, O_new)

    # ── Final normalise + store ──
    O_final = O_local / l_local
    tl.store(o_ptr, O_final)
