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

Per GQA_full_deck.pptx slide 11 (prefill counterpart, T_q ≫ 1):
  - K, V split 64-way (cube=row_wise, pe=row_wise on S_kv); each rank
    owns ``S_local_pe = S_kv / (C·P)``.
  - Q, O split T_q-wise across cubes (cube=row_wise on T_q),
    pe=replicate within each cube — each cube owns
    ``T_q_cube = T_q / C`` rows, all P PEs of the cube hold a copy.
  - Inter-CUBE Ring KV (snake over 4×2 sub-mesh, ADR-0060 §5.5):
    P parallel rings — each PE drives its own same-lane ring via
    independent IPCQ channels. Over C ring steps each PE has
    processed its strided lane of S_kv (= ``S_kv / P`` tokens) for
    the cube's T_q_cube query rows.
  - Intra-CUBE 8-way reduce on the partial ``(m, ℓ, O)`` triple
    (row chain along intra_W + col bridge along intra_N over the
    2×4 PE grid) — combines the P PE-lanes per cube so cube's PE 0
    holds the full S_kv attention for its T_q_cube rows.
  - NO inter-CUBE reduce — each cube writes its own disjoint
    T_q_cube rows of O (PE 0 of each cube is the writer).

Distinction from decode Case 4: decode T_q=1 forces an inter-CUBE
lrab reduce-to-root over (m,ℓ,O). Prefill T_q≫1 allows Q to be
T_q-sharded across cubes, so disjoint output rows avoid the
inter-cube reduce — the prefill-native pattern (per user's
"Ring KV" choice).

Tensor layout:
  Q : (T_q, h_q · d_head) sharded cube=row_wise on T_q, pe=replicate
      within cube; each rank loads ``(G · T_q_cube, d_head)``.
  K : (S_kv, h_kv · d_head) cube=row_wise + pe=row_wise on S_kv;
      each rank owns ``(S_local_pe, h_kv·d_head)``.
  V : same as K.
  O : (T_q, h_q · d_head) sharded same as Q — PE 0 of each cube
      writes its ``(G · T_q_cube, d_head)`` slice; PEs 1..P-1 idle
      the write.

Topology / SFR:
  - Requires ``configure_sfr_intercube_ring(submesh_shape=(2, 4))`` —
    that helper installs both the snake E/W ring across 8 cubes and
    the intra_* lanes needed for the intra-CUBE reduce.

Requires ``T_q % C == 0`` and ``S_kv % (C · P) == 0``.
"""
from __future__ import annotations


TILE_S_KV = 1024  # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).


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_sp_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-4 prefill: Ring KV across cubes + intra-CUBE AR; disjoint T_q/C rows per cube."""
    if T_q % C != 0:
        raise ValueError(
            f"Case 4 prefill requires T_q={T_q} divisible by C={C}"
        )
    if S_kv % (C * P) != 0:
        raise ValueError(
            f"Case 4 prefill requires S_kv={S_kv} divisible by C·P={C * P}"
        )

    G = h_q // h_kv
    T_q_cube = T_q // C
    S_local_pe = S_kv // (C * P)
    KV_ROW_BYTES = d_head * 2  # f16
    n_tiles = (S_local_pe + TILE_S_KV - 1) // TILE_S_KV

    pe_id = tl.program_id(axis=0)

    # Q is this cube's T_q slice (pe=replicate within cube ⇒ every PE
    # loads the same G·T_q_cube rows).
    Q = tl.load(q_ptr, shape=(G * T_q_cube, d_head), dtype="f16")

    # ── Bootstrap (t=0, k=0): load own rank's tile 0; send W for ring ──
    # Persistent (m, ℓ, O) must live OUTSIDE tl.scratch_scope (kernbench
    # scope teardown discards in-scope tensors).
    tile_s = min(TILE_S_KV, S_local_pe)
    K_T = tl.load(k_ptr, shape=(d_head, tile_s), dtype="f16")
    V = tl.load(v_ptr, shape=(tile_s, d_head), dtype="f16")
    if C > 1:
        tl.send(dir="W", src=K_T)
        tl.send(dir="W", src=V)
    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)

    # ── Nested loop (outer tile, inner ring step) — P parallel rings ──
    for t in range(n_tiles):
        tile_start = t * TILE_S_KV
        tile_s = min(TILE_S_KV, S_local_pe - tile_start)
        k_start = 1 if t == 0 else 0
        for k in range(k_start, C):
            with tl.scratch_scope():
                if k == 0:
                    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")
                else:
                    K_T_t = tl.recv(dir="E", shape=(d_head, tile_s), dtype="f16")
                    V_t = tl.recv(dir="E", shape=(tile_s, d_head), dtype="f16")
                if k < C - 1:
                    tl.send(dir="W", src=K_T_t)
                    tl.send(dir="W", src=V_t)
                scores_t = tl.dot(Q, K_T_t)
                m_step = tl.max(scores_t, axis=-1)
                exp_scores_t = tl.exp(scores_t - m_step)
                l_step = tl.sum(exp_scores_t, axis=-1)
                O_step = tl.dot(exp_scores_t, V_t)
                m_new, l_new, O_new = _merge_running(
                    m_local, l_local, O_local, m_step, l_step, O_step, 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) ──
    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

    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)

    # ── Final normalise + store (PE 0 of each cube writes its T_q_cube rows) ──
    if pe_id == 0:
        O_final = O_local / l_local
        tl.store(o_ptr, O_final)
