"""GQA decode kernel — Case 6, **composite-GEMM** command form.

Identical placement and (m, ℓ, O) reduce as the primitive Case-6 kernel
(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp``); the difference is the
command *granularity* of the local attention. The primitive kernel walks
its KV slice in ``TILE_S_KV``-wide tiles, issuing per-tile loads + GEMMs
+ a manual online-softmax merge — O(n_tiles) PE_CPU commands. This kernel
instead issues **one coarse composite GEMM over the whole ``S_local``**
for each matrix product, passing K and V as HBM refs so PE_SCHEDULER
streams and tiles them (the DMA fan-out moves off PE_CPU). The
online-softmax stays primitive but now runs once over the full score row.

So the kernel issues O(1) coarse commands; the scheduler expands each
composite into the same per-tile DMA + MAC work the primitive kernel
issued by hand. Fewer, coarser PE_CPU commands ⇒ lower dispatch cost
under the ADR-0064 Rev2 structural model (the CPU-offload win).
"""
from __future__ import annotations

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


def gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_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) — coarse composite-GEMM 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)

    # ── Local attention as two coarse composite GEMMs ──
    # K, V are HBM refs: PE_SCHEDULER streams + tiles them over S_local
    # (the per-tile DMA fan-out the primitive kernel issued by hand).
    Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
    K_T = tl.ref(k_ptr, shape=(d_head, S_local), dtype="f16")
    V = tl.ref(v_ptr, shape=(S_local, d_head), dtype="f16")

    scores = tl.composite(op="gemm", a=Q, b=K_T)          # Q·Kᵀ, one command
    m_local = tl.max(scores, axis=-1)
    centered = scores - m_local
    exp_scores = tl.exp(centered)
    l_local = tl.sum(exp_scores, axis=-1)
    # P·V into a pre-allocated TCM output (explicit out — the composite
    # output must be a real TCM handle, not auto-allocated scratch).
    O_local = tl.zeros((G * T_q, d_head), dtype="f16")
    tl.composite(op="gemm", a=exp_scores, b=V, out=O_local)  # P·V, one command

    # ── 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)
