// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Unfused-but-GPU-resident causal attention with GQA (M4 correctness path;
// the flash-style fused kernel is M5). Probabilities P [B,H,T,T] are
// materialized in f32 for the backward pass — exactly mirroring the CPU
// reference so parity is bit-for-bit meaningful.
//
// Parallelization: one THREAD per query row (b,h,i) in fwd/dq (grid
// B*H*T), and one thread per KV row (b,h,j) in dkv. dkv iterates the
// q-heads sharing its kv-head, so dk/dv accumulate privately — no atomics,
// deterministic.
//
// Layouts: q/o [B,T,H*hd], k/v [B,T,Hkv*hd], P [B,H,T,T] row-major.
#include <metal_stdlib>
using namespace metal;

struct AttnParams {
    uint B, T, H, HKV, HD;
    float scale;
    uint causal;
    uint window;   // sliding window (keys kept, self incl.); 0 = full
    float softcap; // cap*tanh(s/cap) pre-softmax; 0 = off
};

kernel void attention_fwd_f32(device const float* Q [[buffer(0)]],
                              device const float* K [[buffer(1)]],
                              device const float* V [[buffer(2)]],
                              device float*       O [[buffer(3)]],
                              device float*       P [[buffer(4)]],
                              constant AttnParams& p [[buffer(5)]],
                              uint gid [[thread_position_in_grid]]) {
    const uint i = gid % p.T;
    const uint h = (gid / p.T) % p.H;
    const uint b = gid / (p.T * p.H);
    if (b >= p.B) return;

    const uint rep = p.H / p.HKV;
    const uint hkv = h / rep;
    const uint Cq = p.H * p.HD;
    const uint Ckv = p.HKV * p.HD;

    device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * p.HD;
    device float* prow_base = P + ((ulong(b) * p.H + h) * p.T + i) * p.T;
    const uint jmax = p.causal ? i : p.T - 1;
    const uint jmin = (p.window > 0 && i + 1 > p.window) ? i + 1 - p.window : 0;

    float m = -FLT_MAX;
    for (uint j = jmin; j <= jmax; ++j) {
        device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;
        float s = 0.0f;
        for (uint d = 0; d < p.HD; ++d) s = fma(qi[d], kj[d], s);
        s *= p.scale;
        if (p.softcap > 0.0f) s = p.softcap * precise::tanh(s / p.softcap);
        prow_base[j] = s;
        m = max(m, s);
    }
    float sum = 0.0f;
    for (uint j = jmin; j <= jmax; ++j) {
        const float e = exp(prow_base[j] - m);
        prow_base[j] = e;
        sum += e;
    }
    const float inv = 1.0f / sum;
    for (uint j = 0; j < jmin; ++j) prow_base[j] = 0.0f;
    for (uint j = jmin; j <= jmax; ++j) prow_base[j] *= inv;
    for (uint j = jmax + 1; j < p.T; ++j) prow_base[j] = 0.0f;

    device float* oi = O + (ulong(b) * p.T + i) * Cq + h * p.HD;
    for (uint d = 0; d < p.HD; ++d) oi[d] = 0.0f;
    for (uint j = jmin; j <= jmax; ++j) {
        const float prob = prow_base[j];
        if (prob == 0.0f) continue;
        device const float* vj = V + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;
        for (uint d = 0; d < p.HD; ++d) oi[d] = fma(prob, vj[d], oi[d]);
    }
}

// D[b,h,i] = dO_i · O_i. This is the FlashAttention-2 preprocessing term and
// it equals rowsum(dP ∘ P) exactly:
//   Σ_j P_ij (dO_i · V_j) = dO_i · (Σ_j P_ij V_j) = dO_i · O_i
// Computing it once per query row instead of per (i,j) pair takes the dk/dv
// kernel from O(T³·hd) to O(T²·hd).
kernel void attention_bwd_d_f32(device const float* O  [[buffer(0)]],
                                device const float* dO [[buffer(1)]],
                                device float*       D  [[buffer(2)]],
                                constant AttnParams& p [[buffer(3)]],
                                uint gid [[thread_position_in_grid]]) {
    const uint i = gid % p.T;
    const uint h = (gid / p.T) % p.H;
    const uint b = gid / (p.T * p.H);
    if (b >= p.B) return;

    const uint Cq = p.H * p.HD;
    device const float* oi = O + (ulong(b) * p.T + i) * Cq + h * p.HD;
    device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * p.HD;
    float acc = 0.0f;
    for (uint d = 0; d < p.HD; ++d) acc = fma(doi[d], oi[d], acc);
    D[(ulong(b) * p.H + h) * p.T + i] = acc;
}

// dq for one query row (b,h,i): dq_i = Σ_j dS_ij * K_j * scale, where
// dS = P ∘ (dP − D_i) and dP_ij = dO_i · V_j.
kernel void attention_bwd_dq_f32(device const float* Q  [[buffer(0)]],
                                 device const float* K  [[buffer(1)]],
                                 device const float* V  [[buffer(2)]],
                                 device const float* P  [[buffer(3)]],
                                 device const float* dO [[buffer(4)]],
                                 device const float* D  [[buffer(5)]],
                                 device float*       dQ [[buffer(6)]],
                                 constant AttnParams& p [[buffer(7)]],
                                 uint gid [[thread_position_in_grid]]) {
    const uint i = gid % p.T;
    const uint h = (gid / p.T) % p.H;
    const uint b = gid / (p.T * p.H);
    if (b >= p.B) return;

    const uint rep = p.H / p.HKV;
    const uint hkv = h / rep;
    const uint Cq = p.H * p.HD;
    const uint Ckv = p.HKV * p.HD;

    device const float* prow = P + ((ulong(b) * p.H + h) * p.T + i) * p.T;
    device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * p.HD;
    device float* dqi = dQ + (ulong(b) * p.T + i) * Cq + h * p.HD;

    const float row_dot = D[(ulong(b) * p.H + h) * p.T + i];
    for (uint j = 0; j < p.T; ++j) {
        const float prob = prow[j];
        if (prob == 0.0f) continue;
        device const float* vj = V + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;
        device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;
        float dp = 0.0f;
        for (uint d = 0; d < p.HD; ++d) dp = fma(doi[d], vj[d], dp);
        float ds = prob * (dp - row_dot) * p.scale;
        if (p.softcap > 0.0f) {
            device const float* qq = Q + (ulong(b) * p.T + i) * Cq + h * p.HD;
            float s = 0.0f;
            for (uint d = 0; d < p.HD; ++d) s = fma(qq[d], kj[d], s);
            const float t = precise::tanh(s * p.scale / p.softcap);
            ds *= 1.0f - t * t;
        }
        for (uint d = 0; d < p.HD; ++d) dqi[d] = fma(ds, kj[d], dqi[d]);
    }
}

// dk/dv for one KV row (b,hkv,j): sums over the q-heads sharing this
// kv-head and all query rows i (P_ij = 0 above the diagonal already).
kernel void attention_bwd_dkv_f32(device const float* Q  [[buffer(0)]],
                                  device const float* K  [[buffer(1)]],
                                  device const float* V  [[buffer(2)]],
                                  device const float* P  [[buffer(3)]],
                                  device const float* dO [[buffer(4)]],
                                  device const float* D  [[buffer(5)]],
                                  device float*       dK [[buffer(6)]],
                                  device float*       dV [[buffer(7)]],
                                  constant AttnParams& p [[buffer(8)]],
                                  uint gid [[thread_position_in_grid]]) {
    const uint j = gid % p.T;
    const uint hkv = (gid / p.T) % p.HKV;
    const uint b = gid / (p.T * p.HKV);
    if (b >= p.B) return;

    const uint rep = p.H / p.HKV;
    const uint Cq = p.H * p.HD;
    const uint Ckv = p.HKV * p.HD;

    device const float* vj = V + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;
    device float* dkj = dK + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;
    device float* dvj = dV + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;

    for (uint r = 0; r < rep; ++r) {
        const uint h = hkv * rep + r;
        for (uint i = 0; i < p.T; ++i) {
            const float prob = P[((ulong(b) * p.H + h) * p.T + i) * p.T + j];
            if (prob == 0.0f) continue;
            device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * p.HD;
            device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * p.HD;

            const float row_dot = D[(ulong(b) * p.H + h) * p.T + i];
            float dp = 0.0f;
            for (uint d = 0; d < p.HD; ++d) dp = fma(doi[d], vj[d], dp);
            float ds = prob * (dp - row_dot) * p.scale;
            if (p.softcap > 0.0f) {
                device const float* kk = K + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;
                float s = 0.0f;
                for (uint d = 0; d < p.HD; ++d) s = fma(qi[d], kk[d], s);
                const float t = precise::tanh(s * p.scale / p.softcap);
                ds *= 1.0f - t * t;
            }
            for (uint d = 0; d < p.HD; ++d) {
                dvj[d] = fma(prob, doi[d], dvj[d]);
                dkj[d] = fma(ds, qi[d], dkj[d]);
            }
        }
    }
}