"""GQA decode opt2 variant — two-composite inner attention (ADR-0065).

opt2 of decode (ADR-0060 §8 item 4 / ADR-0065): the per-tile attention is
expressed as **two composites** instead of opt3's long primitive sequence —

  #1  Q·Kᵀ      → a GEMM composite writing the score tile ``scores``.
  #2  softmax   → a ``softmax_merge`` recipe composite (online-softmax merge
      + P·V        of ``(m, l, O)``) whose head GEMM is P·V, with an ``add``
                   epilogue folding the P·V contribution into ``O``.

Fewer PE_CPU-issued commands ⇒ lower dispatch cost under the ADR-0064 Rev2
structural model (the CPU-offload win). Tile 0 establishes the running
state with primitives (kernbench has no scratch-backed ``-inf`` initializer,
same limitation as the opt3 kernel); the next tile exercises the recipe.

**Scope (P5 B):** runnable in op_log mode (latency only). Full data-mode
numeric parity (computing the recipe's 8 MATH ops in the DataExecutor) is a
separate follow-up — see DDD-0065 / the P5-numerics note.
"""
from __future__ import annotations


def gqa_attention_decode_opt2_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:
    """Single-rank (C=P=1) GQA decode using the opt2 two-composite form.

    Layout mirrors ``gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel``
    (M-fold Q, K loaded as ``(d_head, S_local)``). The KV slice is split
    into two sub-tiles so the second one drives the ``softmax_merge``
    recipe composite.
    """
    G = h_q // h_kv
    S_local = S_kv // (C * P)
    KV_ROW_BYTES = d_head * 2  # f16

    Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")

    # Split the local KV slice in two: tile 0 establishes the running state,
    # tile 1 merges via the opt2 two-composite path.
    half = max(1, S_local // 2)
    rest = S_local - half

    # ── Tile 0: establish running (m_local, l_local, O_local) ──
    K_T0 = tl.load(k_ptr, shape=(d_head, half), dtype="f16")
    V0 = tl.load(v_ptr, shape=(half, d_head), dtype="f16")
    scores0 = tl.dot(Q, K_T0)
    m_local = tl.max(scores0, axis=-1)
    exp0 = tl.exp(scores0 - m_local)
    l_local = tl.sum(exp0, axis=-1)
    O_local = tl.dot(exp0, V0)

    # ── Tile 1: opt2 two-composite merge (skipped when S_local == 1) ──
    if rest > 0:
        K_T1 = tl.load(k_ptr + half * KV_ROW_BYTES,
                       shape=(d_head, rest), dtype="f16")
        V1 = tl.ref(v_ptr + half * KV_ROW_BYTES, shape=(rest, d_head),
                    dtype="f16")
        # #1: Q·Kᵀ composite → score tile. No `out` → auto-allocated TCM
        # scratch returned as a handle (ADR-0065 D8), consumed by #2.
        scores1 = tl.composite(op="gemm", a=Q, b=K_T1)
        # #2: softmax_merge recipe (online merge of (m,l,O)) + P·V + add.
        tl.composite(
            prologue=[{"op": "softmax_merge", "s": scores1,
                       "m": m_local, "l": l_local, "O": O_local}],
            op="gemm", b=V1, out=O_local,
            epilogue=[{"op": "add", "other": O_local}],
        )

    # ── Final normalise + store (root only) ──
    if tl.program_id(axis=0) == 0 and tl.program_id(axis=1) == 0:
        O_final = O_local / l_local
        tl.store(o_ptr, O_final)
