"""GQA prefill kernel — Case 3 (Cube-Repl × PE-SP) at single-KV-head group.

Per GQA_full_deck.pptx slide 11 (prefill counterpart, T_q ≫ 1):
  - K, V replicated across all 8 cubes (the 8× memory waste).
  - PEs SP on S_kv inside each cube: each PE attends to its
    ``S_local = S_kv / P`` slice.
  - Q replicated on every rank (every cube does full T_q × S_local).
  - Intra-CUBE 8-way reduce on the partial ``(m, ℓ, O)`` triple
    per Q-tile (row chain along intra_W + col bridge along intra_N
    over the 2×4 PE grid).
  - NO inter-CUBE communication — every cube ends with the full
    answer (8× redundant compute is the inherent cost of Case 3).
  - Designated writer: only PE 0 of CUBE 0 stores ``O`` to avoid
    8 redundant DMA writes.

Q-axis tiling (vs decode Case 3): an outer Q-tile loop wraps the
entire per-tile pipeline (Q load → bootstrap → S_kv-tile fold →
intra-CUBE AR → store) in ``tl.scratch_scope``. Peak scratch is
bounded by ``TILE_Q × TILE_S_KV`` instead of ``T_q × S_kv``, so the
kernel runs at large T_q without exceeding the ~1 MB per-PE budget.
Each Q-tile's (m, ℓ, O) is independent of other Q-tiles, so no
state leaks across iterations.

Tensor layout:
  Q : (T_q, h_q · d_head) replicated on every rank; loaded per
      Q-tile as ``(TILE_Q, d_head)`` slices (G is folded into the
      T_q row axis, total ``G·T_q`` rows).
  K : (S_kv, h_kv · d_head) with cube=replicate, pe=row_wise — each
      PE owns its ``(S_local, h_kv·d_head)`` shard within its cube.
  V : same as K.
  O : (T_q, h_q · d_head) — only PE 0 of CUBE 0 stores, one
      ``(TILE_Q, d_head)`` slice per outer Q-tile iteration.

Topology / SFR:
  - Requires intra_* lanes; ``configure_sfr_intercube_ring`` (with the
    snake submesh) or ``configure_sfr_intercube_multisip`` both install
    them.
"""
from __future__ import annotations


TILE_S_KV = 64  # per-tile S_kv width (small to keep ``scores`` bounded).
TILE_Q = 64     # per-tile Q row count (outer-loop tile size).


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_prefill_long_ctx_cube_repl_pe_sp_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-3 prefill: PE-SP S_kv split + Q-tiled intra-CUBE AR."""
    G = h_q // h_kv
    S_local = S_kv // P  # each PE owns S_kv/P (cube=replicate, pe=row_wise)
    pe_id = tl.program_id(axis=0)
    cube_id = tl.program_id(axis=1)

    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

    KV_ROW_BYTES = d_head * 2  # f16
    Q_ROW_BYTES = d_head * 2   # G is folded into the row axis, not d_head
    total_q_rows = G * T_q
    n_q_tiles = (total_q_rows + TILE_Q - 1) // TILE_Q
    n_kv_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV

    # ── Outer Q-tile loop ──
    # Per Q-tile: load Q-slice, run full S_kv attention, intra-CUBE AR,
    # store. All scratch is recycled at outer-scope exit.
    for q_idx in range(n_q_tiles):
        q_start = q_idx * TILE_Q
        q_rows = min(TILE_Q, total_q_rows - q_start)
        with tl.scratch_scope():
            Q = tl.load(q_ptr + q_start * Q_ROW_BYTES,
                        shape=(q_rows, d_head), dtype="f16")

            # ── Bootstrap (S_kv tile 0) ──
            # K_T, V, scores, exp_scores live in the OUTER (Q-tile)
            # scope. They survive the inner S_kv-tile loop (each
            # iteration's allocations are recycled) and are freed
            # together when the outer scope exits.
            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)
            exp_scores = tl.exp(scores - m_local)
            l_local = tl.sum(exp_scores, axis=-1)
            O_local = tl.dot(exp_scores, V)

            # ── S_kv tiles 1..N: fold via online-softmax merge ──
            for tile_idx in range(1, n_kv_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)
                    exp_scores_t = tl.exp(scores_t - m_tile)
                    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)

            # ── Intra-CUBE 8-way reduce-to-PE0 (row chain + col bridge) ──
            # (m_local, l_local, O_local) are still alive in the outer
            # scope; the send/recv operate on them directly.
            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)

            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)

            # ── Store this Q-tile (designated writer: cube 0, PE 0) ──
            if pe_id == 0 and cube_id == 0:
                O_final = O_local / l_local
                tl.store(o_ptr + q_start * Q_ROW_BYTES, O_final)
