"""GQA prefill kernel — Case 2 (Cube-Repl × PE-TP) 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 — this is
    the inherent cost Case 2 demonstrates). PE-replicate within each
    cube too: every PE on every cube holds an HBM copy.
  - PE-TP on T_q: only CUBE 0 has work; within CUBE 0 the P PEs split
    the T_q axis disjointly. CUBEs 1..7 idle (slide-11 "PE-TP only
    uses one cube; replicate is pure waste here").
  - NO inter-rank communication (each active rank has full KV and
    disjoint T_q rows — slide 11 lists comm cost as "none").

Distinction from decode Case 2: decode T_q=1 makes PE-TP wasteful
(only PE 0 of CUBE 0 works → 1 of 64 ranks active). Prefill T_q≫1
makes PE-TP useful — all P PEs of CUBE 0 work on disjoint T_q rows
(P of 64 ranks active; CUBEs 1..7 are pure memory waste).

Q-axis tiling: an outer Q-tile loop wraps the entire per-tile
pipeline (Q load → bootstrap → S_kv-tile fold → store) in
``tl.scratch_scope``. Peak scratch is bounded by
``TILE_Q × TILE_S_KV`` instead of ``T_q_local × S_kv``, so the
kernel runs at large T_q without exceeding the ~1 MB per-PE budget.
Each Q-tile is independent (disjoint output rows; no inter-tile state).

Tensor layout:
  Q : (T_q, h_q · d_head) cube=replicate, pe=row_wise on T_q. Only
      CUBE 0's PEs load — each owns ``G · T_q_local`` total rows with
      ``T_q_local = T_q / P``; loaded per Q-tile.
  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) cube=replicate, pe=row_wise on T_q —
      only CUBE 0's PEs write their disjoint ``T_q_local`` rows,
      one ``(TILE_Q, d_head)`` slice per outer iteration.

Topology / SFR:
  - SFR setup is benign (no inter-rank sends/recvs); any of the
    ``configure_sfr_intercube_*`` helpers works.

Requires ``T_q % P == 0``.
"""
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 gqa_attention_prefill_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 prefill: CUBE 0 only; Q-tiled PE-TP; full KV per rank; no comm."""
    if T_q % P != 0:
        raise ValueError(
            f"Case 2 prefill requires T_q={T_q} divisible by P={P}"
        )

    cube_id = tl.program_id(axis=1)
    # CUBE 0 only — CUBEs 1..7 are pure memory-replication waste, no work.
    if cube_id != 0:
        return

    G = h_q // h_kv
    T_q_local = T_q // P
    n_kv_tiles = (S_kv + TILE_S_KV - 1) // TILE_S_KV
    KV_ROW_BYTES = d_head * 2  # f16
    Q_ROW_BYTES = d_head * 2
    total_q_rows = G * T_q_local
    n_q_tiles = (total_q_rows + TILE_Q - 1) // TILE_Q

    # ── Outer Q-tile loop ──
    # Per Q-tile: load Q-slice, run full S_kv attention, store. All
    # scratch is recycled at outer-scope exit, so peak usage is bounded
    # by TILE_Q × TILE_S_KV.
    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) ──
            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)
            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_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)
                    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 = 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)

            # ── Store this Q-tile (each active PE writes its slice) ──
            O_final = O_local / l_local
            tl.store(o_ptr + q_start * Q_ROW_BYTES, O_final)
