"""GQA decode kernel — Case 6, **composite_extended** command form.

Same placement and (m, ℓ, O) reduce as the primitive / composite-GEMM
Case-6 kernels; the local attention is the opt2 **two-composite** form
(ADR-0060 §8 item 4 / ADR-0065), issued coarsely:

  establish  a small first KV slice computes the running (m, ℓ, O) with
             primitives (kernbench has no scratch-backed ``-inf``
             initializer, so the recipe — which *merges* into an existing
             accumulator — needs a seed).
  #1  Q·Kᵀ   one coarse composite GEMM over the remaining S_local
             (K passed as an HBM ref → PE_SCHEDULER streams + tiles it).
  #2  softmax → one ``softmax_merge`` recipe composite (online-softmax
      + P·V        merge of (m, ℓ, O)) whose head GEMM is P·V over the same
                   remaining slice (V as an HBM ref), with an ``add``
                   epilogue folding the P·V contribution into ``O``.

So the bulk of the KV slice is two coarse PE_CPU commands, and the recipe
additionally folds the per-tile online merge + P·V into one descriptor —
the fewest / cheapest commands of the three forms (the headline
CPU-offload win, ADR-0064 Rev2).

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

from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
    reduce_mlo,
    root_cube_for,
)

# Small primitive seed slice that establishes the running (m, ℓ, O) the
# recipe then merges into. One scheduler K-tile wide (ADR-0064 TILE_K).
_SEED_S_KV = 64


def gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_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-6 decode (Cube-SP × PE-SP) — softmax_merge recipe command form."""
    G = h_q // h_kv
    n_ranks = C * P
    S_local = S_kv // n_ranks
    pe_id = tl.program_id(axis=0)
    cube_id = tl.program_id(axis=1)
    KV_ROW_BYTES = d_head * 2  # f16

    # ── Establish running (m, ℓ, O) on a small seed slice (primitive) ──
    Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
    seed = min(_SEED_S_KV, S_local)
    K_T0 = tl.load(k_ptr, shape=(d_head, seed), dtype="f16")
    V0 = tl.load(v_ptr, shape=(seed, 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)

    # ── Remaining slice: one Q·Kᵀ composite + one softmax_merge recipe ──
    rest = S_local - seed
    if rest > 0:
        K_T1 = tl.ref(k_ptr + seed * KV_ROW_BYTES,
                      shape=(d_head, rest), dtype="f16")
        V1 = tl.ref(v_ptr + seed * KV_ROW_BYTES,
                    shape=(rest, d_head), dtype="f16")
        scores1 = tl.composite(op="gemm", a=Q, b=K_T1)
        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}],
        )

    # ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ──
    reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl)

    # ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
    if pe_id == 0 and cube_id == root_cube_for(C):
        O_final = O_local / l_local
        tl.store(o_ptr, O_final)
