// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Fused (flash-style) causal attention with GQA — forward + backward, f32.
// Nothing of size T² is ever written: the forward keeps the online-softmax
// state in registers and stores only L = m + log(l), one float per query
// row, which the backward uses to recompute P tile-by-tile. That is what
// lets context length scale — the unfused path in attention.metal needs
// B·H·T² floats of probabilities per layer (2.1 GB for batch 64 × T 1024 ×
// 8 heads), this needs B·H·T.
//
// Parallelization: one THREAD per output row —
//   forward / dQ : one thread per (b, h, i)      grid B·H·T
//   dK,dV        : one thread per (b, hkv, j)    grid B·HKV·T
// so dK/dV accumulate privately per KV row (no atomics, deterministic) and
// each simdgroup's 32 threads are consecutive query rows of the same head,
// which makes their K/V reads a broadcast that the cache serves once.
//
// head_dim is a TEMPLATE parameter, not a runtime value: the per-thread
// q/o/accumulator arrays must be compile-time sized and fully unrolled or
// they land in thread-local (device-backed) memory instead of registers.
// Instantiated per supported head_dim below; the host falls back to the
// unfused kernel for other sizes.
//
// Online softmax (FlashAttention-2, RESEARCH.md §6):
//   m_new = max(m, s) ; corr = exp(m − m_new) ; e = exp(s − m_new)
//   l = l·corr + e    ; o = o·corr + e·v      ; divide by l once at the end
// m starts at −FLT_MAX, never −INFINITY: exp(−inf − (−inf)) is NaN.
#include <metal_stdlib>
using namespace metal;

struct FlashParams {
    uint B, T, H, HKV;
    float scale;
    uint causal;
    uint window; // sliding window (keys kept, self incl.); 0 = full attention
};

// ---------------------------------------------------------------- forward

// Threadgroup = TGQ consecutive query rows of one (b, h). K/V arrive through
// threadgroup memory in blocks of BKV rows, staged cooperatively, so each
// K/V element is fetched from device once per threadgroup instead of once
// per thread. Causal blocks entirely above the diagonal are skipped outright
// (kb_lim), and only the diagonal block pays the per-element mask test.
enum : uint { TGQ = 64, BKV = 16 };

template <uint HD>
kernel void flash_attn_fwd(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*       L  [[buffer(4)]],
                           constant FlashParams& p [[buffer(5)]],
                           uint tgid [[threadgroup_position_in_grid]],
                           uint tid  [[thread_index_in_threadgroup]]) {
    threadgroup float Ks[BKV * HD];
    threadgroup float Vs[BKV * HD];

    const uint q_blocks = (p.T + TGQ - 1) / TGQ;
    const uint qb = tgid % q_blocks;
    const uint h = (tgid / q_blocks) % p.H;
    const uint b = tgid / (q_blocks * p.H);

    const uint i = qb * TGQ + tid;
    const bool active = (i < p.T) && (b < p.B);

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

    float q[HD], o[HD];
    if (active) {
        device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * HD;
#pragma clang loop unroll(full)
        for (uint d = 0; d < HD; ++d) {
            q[d] = qi[d];
            o[d] = 0.0f;
        }
    } else {
#pragma clang loop unroll(full)
        for (uint d = 0; d < HD; ++d) {
            q[d] = 0.0f;
            o[d] = 0.0f;
        }
    }

    float m = -FLT_MAX;
    float l = 0.0f;

    // Causal: this block's largest query index bounds the KV blocks we touch.
    const uint row_max = min(qb * TGQ + TGQ - 1, p.T - 1);
    const uint j_end = p.causal ? (row_max + 1) : p.T;
    const uint diag_start = p.causal ? (qb * TGQ) : p.T; // blocks below need no mask
    // Sliding window: this row's first visible key, and (uniform across the
    // threadgroup) the first KV block any row here can see — earlier blocks
    // are skipped outright, so compute scales with the window, not with T.
    const uint jmin = (p.window > 0 && i + 1 > p.window) ? i + 1 - p.window : 0;
    const uint tg_jmin =
        (p.window > 0 && qb * TGQ + 1 > p.window) ? qb * TGQ + 1 - p.window : 0;
    const uint jb_begin = (tg_jmin / BKV) * BKV;

    for (uint jb = jb_begin; jb < j_end; jb += BKV) {
        const uint block_len = min(BKV, j_end - jb);

        threadgroup_barrier(mem_flags::mem_threadgroup);
        for (uint e = tid; e < block_len * HD; e += TGQ) {
            const uint jr = e / HD;
            const uint d = e % HD;
            const ulong src = (ulong(b) * p.T + jb + jr) * Ckv + hkv * HD + d;
            Ks[jr * HD + d] = K[src];
            Vs[jr * HD + d] = V[src];
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);

        if (!active) continue;
        const bool needs_mask = p.causal && (jb + BKV > diag_start);
        for (uint jr = 0; jr < block_len; ++jr) {
            const uint j = jb + jr;
            if (needs_mask && j > i) break; // rest of the block is masked too
            if (j < jmin) continue;         // below this row's window

            threadgroup const float* kj = Ks + jr * HD;
            float s = 0.0f;
#pragma clang loop unroll(full)
            for (uint d = 0; d < HD; ++d) s = fma(q[d], kj[d], s);
            s *= p.scale;

            const float m_new = max(m, s);
            const float corr = exp(m - m_new); // 0 on the first iteration
            const float e = exp(s - m_new);
            l = l * corr + e;

            threadgroup const float* vj = Vs + jr * HD;
#pragma clang loop unroll(full)
            for (uint d = 0; d < HD; ++d) o[d] = fma(o[d], corr, e * vj[d]);
            m = m_new;
        }
    }

    if (!active) return;
    const float inv = (l > 0.0f) ? (1.0f / l) : 0.0f;
    device float* oi = O + (ulong(b) * p.T + i) * Cq + h * HD;
#pragma clang loop unroll(full)
    for (uint d = 0; d < HD; ++d) oi[d] = o[d] * inv;
    // logsumexp; a fully-masked row would give -inf, guarded like the divide
    L[(ulong(b) * p.H + h) * p.T + i] = (l > 0.0f) ? (m + log(l)) : 0.0f;
}

// --------------------------------------------------------------- backward
// D[b,h,i] = dO_i · O_i == rowsum(dP ∘ P) (FA2 identity), computed by
// attention_bwd_d_f32 in attention.metal and passed in here.

template <uint HD>
kernel void flash_attn_bwd_dq(device const float* Q  [[buffer(0)]],
                              device const float* K  [[buffer(1)]],
                              device const float* V  [[buffer(2)]],
                              device const float* dO [[buffer(3)]],
                              device const float* L  [[buffer(4)]],
                              device const float* D  [[buffer(5)]],
                              device float*       dQ [[buffer(6)]],
                              constant FlashParams& 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 hkv = h / (p.H / p.HKV);
    const uint Cq = p.H * HD;
    const uint Ckv = p.HKV * HD;

    device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * HD;
    device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * HD;

    float q[HD], dq[HD], go[HD];
#pragma clang loop unroll(full)
    for (uint d = 0; d < HD; ++d) {
        q[d] = qi[d];
        go[d] = doi[d];
        dq[d] = 0.0f;
    }

    const float li = L[(ulong(b) * p.H + h) * p.T + i];
    const float di = D[(ulong(b) * p.H + h) * p.T + i];
    const uint jmax = p.causal ? i : (p.T - 1);
    const uint jmin = (p.window > 0 && i + 1 > p.window) ? i + 1 - p.window : 0;

    for (uint j = jmin; j <= jmax; ++j) {
        device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * HD;
        device const float* vj = V + (ulong(b) * p.T + j) * Ckv + hkv * HD;

        float s = 0.0f, dp = 0.0f;
#pragma clang loop unroll(full)
        for (uint d = 0; d < HD; ++d) {
            s = fma(q[d], kj[d], s);
            dp = fma(go[d], vj[d], dp);
        }
        const float prob = exp(s * p.scale - li); // recomputed, never stored
        const float ds = prob * (dp - di) * p.scale;
#pragma clang loop unroll(full)
        for (uint d = 0; d < HD; ++d) dq[d] = fma(ds, kj[d], dq[d]);
    }

    device float* dqi = dQ + (ulong(b) * p.T + i) * Cq + h * HD;
#pragma clang loop unroll(full)
    for (uint d = 0; d < HD; ++d) dqi[d] += dq[d];
}

// dK and dV are separate kernels on purpose. Combined, one thread holds
// dk[HD] + dv[HD] and spills: measured 150 ms -> 118 ms for gpt-10m just by
// moving the read-only k/v out of registers, so the accumulators matter too.
// Split, each thread carries a single HD-sized accumulator; the price is
// recomputing the q·k dot in both kernels, which is cheaper than the spill.

template <uint HD>
kernel void flash_attn_bwd_dv(device const float* Q  [[buffer(0)]],
                              device const float* K  [[buffer(1)]],
                              device const float* dO [[buffer(2)]],
                              device const float* L  [[buffer(3)]],
                              device float*       dV [[buffer(4)]],
                              constant FlashParams& p [[buffer(5)]],
                              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 * HD;
    const uint Ckv = p.HKV * HD;
    device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * HD;

    float dv[HD];
#pragma clang loop unroll(full)
    for (uint d = 0; d < HD; ++d) dv[d] = 0.0f;

    const uint imin = p.causal ? j : 0;
    // Sliding window: only queries within `window` of j ever saw it.
    const uint imax = (p.window > 0) ? min(p.T, j + p.window) : p.T;
    for (uint r = 0; r < rep; ++r) {
        const uint h = hkv * rep + r;
        device const float* Lh = L + (ulong(b) * p.H + h) * p.T;
        for (uint i = imin; i < imax; ++i) {
            device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * HD;
            device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * HD;
            float s = 0.0f;
#pragma clang loop unroll(full)
            for (uint d = 0; d < HD; ++d) s = fma(qi[d], kj[d], s);
            const float prob = exp(s * p.scale - Lh[i]);
#pragma clang loop unroll(full)
            for (uint d = 0; d < HD; ++d) dv[d] = fma(prob, doi[d], dv[d]);
        }
    }

    device float* dvj = dV + (ulong(b) * p.T + j) * Ckv + hkv * HD;
#pragma clang loop unroll(full)
    for (uint d = 0; d < HD; ++d) dvj[d] += dv[d];
}

template <uint HD>
kernel void flash_attn_bwd_dk(device const float* Q  [[buffer(0)]],
                              device const float* K  [[buffer(1)]],
                              device const float* V  [[buffer(2)]],
                              device const float* dO [[buffer(3)]],
                              device const float* L  [[buffer(4)]],
                              device const float* D  [[buffer(5)]],
                              device float*       dK [[buffer(6)]],
                              constant FlashParams& p [[buffer(7)]],
                              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 * HD;
    const uint Ckv = p.HKV * HD;
    device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * HD;
    device const float* vj = V + (ulong(b) * p.T + j) * Ckv + hkv * HD;

    float dk[HD];
#pragma clang loop unroll(full)
    for (uint d = 0; d < HD; ++d) dk[d] = 0.0f;

    const uint imin = p.causal ? j : 0;
    const uint imax = (p.window > 0) ? min(p.T, j + p.window) : p.T;
    for (uint r = 0; r < rep; ++r) {
        const uint h = hkv * rep + r;
        device const float* Lh = L + (ulong(b) * p.H + h) * p.T;
        device const float* Dh = D + (ulong(b) * p.H + h) * p.T;
        for (uint i = imin; i < imax; ++i) {
            device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * HD;
            device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * HD;
            float s = 0.0f, dp = 0.0f;
#pragma clang loop unroll(full)
            for (uint d = 0; d < HD; ++d) {
                s = fma(qi[d], kj[d], s);
                dp = fma(doi[d], vj[d], dp);
            }
            const float prob = exp(s * p.scale - Lh[i]);
            const float ds = prob * (dp - Dh[i]) * p.scale;
#pragma clang loop unroll(full)
            for (uint d = 0; d < HD; ++d) dk[d] = fma(ds, qi[d], dk[d]);
        }
    }

    device float* dkj = dK + (ulong(b) * p.T + j) * Ckv + hkv * HD;
#pragma clang loop unroll(full)
    for (uint d = 0; d < HD; ++d) dkj[d] += dk[d];
}

// ---------------------------------------------------------- instantiations
// Every config in configs/ uses head_dim 64 (d_model / n_heads); the others
// are here so the fused path covers common variants. Unsupported sizes fall
// back to the unfused kernels on the host side.
#define INSTANTIATE_FLASH(HD)                                                       \
    template [[host_name("flash_attn_fwd_f32_hd" #HD)]] kernel void                 \
    flash_attn_fwd<HD>(device const float*, device const float*, device const float*, \
                       device float*, device float*, constant FlashParams&, uint,   \
                       uint);                                                        \
    template [[host_name("flash_attn_bwd_dq_f32_hd" #HD)]] kernel void              \
    flash_attn_bwd_dq<HD>(device const float*, device const float*,                 \
                          device const float*, device const float*,                 \
                          device const float*, device const float*, device float*,  \
                          constant FlashParams&, uint);                             \
    template [[host_name("flash_attn_bwd_dv_f32_hd" #HD)]] kernel void              \
    flash_attn_bwd_dv<HD>(device const float*, device const float*,                 \
                          device const float*, device const float*, device float*,  \
                          constant FlashParams&, uint);                             \
    template [[host_name("flash_attn_bwd_dk_f32_hd" #HD)]] kernel void              \
    flash_attn_bwd_dk<HD>(device const float*, device const float*,                 \
                          device const float*, device const float*,                 \
                          device const float*, device const float*, device float*,  \
                          constant FlashParams&, uint);

INSTANTIATE_FLASH(16)
INSTANTIATE_FLASH(32)
INSTANTIATE_FLASH(48)
INSTANTIATE_FLASH(64)
INSTANTIATE_FLASH(80)
INSTANTIATE_FLASH(96)
INSTANTIATE_FLASH(128)