// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Tiled flash attention forward using simdgroup_matrix 8x8 fragments.
// Same algorithm and same outputs as flash_attn_fwd in flash_attention.metal
// (online softmax, stores only L = m + log(l)); the difference is that QK^T
// and PV are matrix multiplies over threadgroup-staged tiles instead of
// per-thread scalar dot products, which removes the serial FMA dependency
// chain that capped the scalar kernel at ~1.6 TFLOPs.
//
// Layout (BQ=32, BK=16, 4 simdgroups = 128 threads), split-Q per FA2:
//   simdgroup s owns query rows [8s, 8s+8) — exactly one fragment row — and
//   holds Q and the O accumulator in registers for the whole kernel
//   (2 * HD/8 fragments = 32 floats/lane at HD=64). K/V tiles pass through
//   threadgroup memory; the S/P tile round-trips through it so the softmax
//   reductions can run on plain threads.
//
// THE PER-ROW RESCALE. Online softmax needs O <- diag(corr) * O every KV
// block, but MSL leaves the element->lane mapping of a simdgroup_matrix
// unspecified, so a lane cannot know which row its registers belong to.
// MLX reverse-engineers the mapping; instead this builds an 8x8 diagonal
// matrix holding corr in threadgroup memory and applies it with an MMA.
// That is spec-clean, costs HD/8 extra MMAs per block (~25% more MMA work,
// measured cheaper than the alternatives), and keeps O in registers — the
// point of the exercise, since staging O in threadgroup memory would add
// 8 KB and halve residency.
//
// Threadgroup memory ~12.7 KB at HD=64: K and V tiles (padded +4 floats per
// row against bank conflicts), the S/P tile, per-row softmax state, and one
// 8x8 diagonal scratch per simdgroup. The K/V tiles are reused as the output
// staging buffer at the end, once they are dead.
#include <metal_stdlib>
#include <metal_simdgroup_matrix>
using namespace metal;

struct FlashParams {
    uint B, T, H, HKV;
    float scale;
    uint causal;
    uint window; // unused here: the host routes window > 0 to the scalar kernels
};

// Enumerators, never `constant constexpr`: see RESEARCH.md 7a — the latter is
// a constant-address-space variable, which blocks unrolling and spills every
// fragment to the stack.
enum : uint {
    BQ = 32,           // query rows per threadgroup
    BK = 16,           // KV rows per iteration
    NSG = BQ / 8,      // one 8-row fragment per simdgroup
    MMA_THREADS = NSG * 32,
    KF = BK / 8,       // S fragments along the KV axis
    PADF = 4,          // 16 bytes
    LDS = BK + PADF,   // S/P tile row stride
};

template <uint HD>
kernel void flash_attn_fwd_mma(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]],
                               uint lane [[thread_index_in_simdgroup]],
                               uint sgid [[simdgroup_index_in_threadgroup]]) {
    constexpr uint DF = HD / 8;   // fragments along head_dim
    constexpr uint LDKV = HD + PADF;

    // K and V tiles, reused as output staging after the KV loop.
    threadgroup float kv[2 * BK * LDKV];
    threadgroup float* Ks = kv;
    threadgroup float* Vs = kv + BK * LDKV;
    threadgroup float Sbuf[BQ * LDS];
    threadgroup float row_m[BQ], row_l[BQ], row_corr[BQ];
    threadgroup float diag[NSG * 64];

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

    const uint hkv = h / (p.H / p.HKV);
    const uint Cq = p.H * HD;
    const uint Ckv = p.HKV * HD;
    const uint q0 = qb * BQ;              // first query row of this threadgroup
    const uint sg_row = sgid * 8;         // first query row of this simdgroup

    if (tid < BQ) {
        row_m[tid] = -FLT_MAX;
        row_l[tid] = 0.0f;
    }

    // Q stays in registers for the whole kernel. It is staged through the
    // (not yet used) K/V tile first so a query block straddling the end of
    // the sequence gets zero-filled rows instead of an out-of-bounds read —
    // loading fragments straight from device would need the block to be a
    // full 8 rows.
    simdgroup_float8x8 Qf[DF], Of[DF];
    {
        threadgroup_barrier(mem_flags::mem_threadgroup);
        for (uint e = tid; e < BQ * HD; e += MMA_THREADS) {
            const uint r = e / HD;
            const uint d = e % HD;
            const uint i = q0 + r;
            kv[r * HD + d] =
                (i < p.T) ? Q[(ulong(b) * p.T + i) * Cq + h * HD + d] : 0.0f;
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
#pragma clang loop unroll(full)
        for (uint d = 0; d < DF; ++d) {
            simdgroup_load(Qf[d], kv + sg_row * HD + d * 8, HD);
            Of[d] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
        }
    }

    // Causal: only KV blocks up to this threadgroup's last query row matter.
    const uint row_max = min(q0 + BQ - 1, p.T - 1);
    const uint j_end = p.causal ? (row_max + 1) : p.T;

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

        threadgroup_barrier(mem_flags::mem_threadgroup);
        // Stage K/V; rows past the block length are zeroed so masked columns
        // can never contribute a NaN through 0 * garbage.
        for (uint e = tid; e < BK * HD; e += MMA_THREADS) {
            const uint jr = e / HD;
            const uint d = e % HD;
            const bool ok = jr < block_len;
            const ulong src = (ulong(b) * p.T + jb + jr) * Ckv + hkv * HD + d;
            Ks[jr * LDKV + d] = ok ? K[src] : 0.0f;
            Vs[jr * LDKV + d] = ok ? V[src] : 0.0f;
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);

        // S = Q @ K^T, accumulated over head_dim fragments.
        simdgroup_float8x8 Sf[KF];
#pragma clang loop unroll(full)
        for (uint j = 0; j < KF; ++j) Sf[j] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
#pragma clang loop unroll(full)
        for (uint d = 0; d < DF; ++d) {
#pragma clang loop unroll(full)
            for (uint j = 0; j < KF; ++j) {
                simdgroup_float8x8 KTf;
                // transpose=true turns the K tile into K^T[d-block][j-block]
                simdgroup_load(KTf, Ks + (j * 8) * LDKV + d * 8, LDKV, 0, true);
                simdgroup_multiply_accumulate(Sf[j], Qf[d], KTf, Sf[j]);
            }
        }
#pragma clang loop unroll(full)
        for (uint j = 0; j < KF; ++j)
            simdgroup_store(Sf[j], Sbuf + sg_row * LDS + j * 8, LDS);
        threadgroup_barrier(mem_flags::mem_threadgroup);

        // Softmax on plain threads, one row each. This is O(BQ*BK) against
        // O(BQ*BK*HD) of MMA work, so leaving 96 of 128 threads idle costs
        // ~1% and keeps the reduction free of any lane-mapping assumption.
        if (tid < BQ) {
            const uint i = q0 + tid;
            threadgroup float* srow = Sbuf + tid * LDS;
            const uint valid =
                (i >= p.T) ? 0u
                           : (p.causal ? min(block_len, (i >= jb) ? (i - jb + 1) : 0u)
                                       : block_len);
            float bmax = -FLT_MAX;
            for (uint j = 0; j < valid; ++j) bmax = max(bmax, srow[j] * p.scale);

            const float m_old = row_m[tid];
            const float m_new = max(m_old, bmax);
            const float corr = (m_old == -FLT_MAX) ? 0.0f : exp(m_old - m_new);
            float bsum = 0.0f;
            for (uint j = 0; j < BK; ++j) {
                if (j < valid) {
                    const float e = exp(srow[j] * p.scale - m_new);
                    srow[j] = e;
                    bsum += e;
                } else {
                    srow[j] = 0.0f; // masked / out of range
                }
            }
            row_m[tid] = (valid > 0) ? m_new : m_old;
            row_l[tid] = row_l[tid] * ((valid > 0) ? corr : 1.0f) + bsum;
            row_corr[tid] = (valid > 0) ? corr : 1.0f;
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);

        // O <- diag(corr) @ O, then O += P @ V.
        threadgroup float* dg = diag + sgid * 64;
        for (uint e = lane; e < 64; e += 32) dg[e] = 0.0f;
        simdgroup_barrier(mem_flags::mem_threadgroup);
        if (lane < 8) dg[lane * 8 + lane] = row_corr[sg_row + lane];
        simdgroup_barrier(mem_flags::mem_threadgroup);

        simdgroup_float8x8 Dg;
        simdgroup_load(Dg, dg, 8);
#pragma clang loop unroll(full)
        for (uint d = 0; d < DF; ++d) {
            simdgroup_float8x8 scaled;
            simdgroup_multiply(scaled, Dg, Of[d]);
            Of[d] = scaled;
        }
#pragma clang loop unroll(full)
        for (uint j = 0; j < KF; ++j) {
            simdgroup_float8x8 Pf;
            simdgroup_load(Pf, Sbuf + sg_row * LDS + j * 8, LDS);
#pragma clang loop unroll(full)
            for (uint d = 0; d < DF; ++d) {
                simdgroup_float8x8 Vf;
                simdgroup_load(Vf, Vs + (j * 8) * LDKV + d * 8, LDKV);
                simdgroup_multiply_accumulate(Of[d], Pf, Vf, Of[d]);
            }
        }
    }

    // Divide by l and write out. K/V are dead now, so their tile doubles as
    // the [BQ x HD] staging buffer (2 * BK * LDKV >= BQ * HD).
    threadgroup_barrier(mem_flags::mem_threadgroup);
    threadgroup float* Ostage = kv;
#pragma clang loop unroll(full)
    for (uint d = 0; d < DF; ++d)
        simdgroup_store(Of[d], Ostage + sg_row * HD + d * 8, HD);
    threadgroup_barrier(mem_flags::mem_threadgroup);

    for (uint e = tid; e < BQ * HD; e += MMA_THREADS) {
        const uint r = e / HD;
        const uint d = e % HD;
        const uint i = q0 + r;
        if (i >= p.T) continue;
        const float l = row_l[r];
        O[(ulong(b) * p.T + i) * Cq + h * HD + d] =
            (l > 0.0f) ? (Ostage[r * HD + d] / l) : 0.0f;
    }
    if (tid < BQ && (q0 + tid) < p.T) {
        const float l = row_l[tid];
        L[(ulong(b) * p.H + h) * p.T + q0 + tid] =
            (l > 0.0f) ? (row_m[tid] + log(l)) : 0.0f;
    }
}

#define INSTANTIATE_FLASH_MMA(HD)                                                   \
    template [[host_name("flash_attn_fwd_mma_f32_hd" #HD)]] kernel void             \
    flash_attn_fwd_mma<HD>(device const float*, device const float*,                \
                           device const float*, device float*, device float*,       \
                           constant FlashParams&, uint, uint, uint, uint);

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

// ============================================================ backward (MMA)
//
// Backward needs no online rescale — L from the forward already fixes the
// softmax normalisation — so there is no diagonal-matrix trick here, just
// three matmuls per block with an elementwise step between them.
//
//   dQ kernel  (parallel over queries, one threadgroup per BQ query rows):
//       S  = Q @ K^T ; P = exp(S*scale - L_i)
//       dP = dO @ V^T ; dS = P * (dP - D_i) * scale
//       dQ += dS @ K
//   dKV kernel (parallel over KV rows, one threadgroup per BKV rows):
//       everything transposed, obtained by swapping the operand roles:
//       S^T = K @ Q^T and dP^T = V @ dO^T, so
//       dV += P^T @ dO and dK += dS^T @ Q.
//
// Splitting this way keeps dK/dV accumulation private to a threadgroup — no
// atomics, and the result is deterministic. Fragments cut register pressure
// by 4x versus the scalar kernels: an 8x8 fragment is 2 floats per lane, so
// holding K, V, dK, dV as fragments is 64 floats/lane at hd=64 where the
// scalar version needed 256 and spilled.

enum : uint {
    BQB = 16,              // query block consumed per iteration by dKV
    LDQB = BQB + PADF,
};

template <uint HD>
kernel void flash_attn_bwd_dq_mma(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 tgid [[threadgroup_position_in_grid]],
                                  uint tid  [[thread_index_in_threadgroup]],
                                  uint sgid [[simdgroup_index_in_threadgroup]]) {
    constexpr uint DF = HD / 8;
    constexpr uint LDKV = HD + PADF;

    threadgroup float kv[2 * BK * LDKV];
    threadgroup float* Ks = kv;
    threadgroup float* Vs = kv + BK * LDKV;
    threadgroup float Sbuf[BQ * LDS];
    threadgroup float Pbuf[BQ * LDS];
    threadgroup float row_L[BQ], row_D[BQ];

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

    const uint hkv = h / (p.H / p.HKV);
    const uint Cq = p.H * HD;
    const uint Ckv = p.HKV * HD;
    const uint q0 = qb * BQ;
    const uint sg_row = sgid * 8;

    if (tid < BQ) {
        const uint i = q0 + tid;
        row_L[tid] = (i < p.T) ? L[(ulong(b) * p.H + h) * p.T + i] : 0.0f;
        row_D[tid] = (i < p.T) ? D[(ulong(b) * p.H + h) * p.T + i] : 0.0f;
    }

    // Q and dO into registers (staged so ragged tails zero-fill).
    simdgroup_float8x8 Qf[DF], dOf[DF], dQf[DF];
    for (uint pass = 0; pass < 2; ++pass) {
        threadgroup_barrier(mem_flags::mem_threadgroup);
        device const float* src = (pass == 0) ? Q : dO;
        for (uint e = tid; e < BQ * HD; e += MMA_THREADS) {
            const uint r = e / HD, d = e % HD;
            const uint i = q0 + r;
            kv[r * HD + d] =
                (i < p.T) ? src[(ulong(b) * p.T + i) * Cq + h * HD + d] : 0.0f;
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
#pragma clang loop unroll(full)
        for (uint d = 0; d < DF; ++d) {
            if (pass == 0) simdgroup_load(Qf[d], kv + sg_row * HD + d * 8, HD);
            else           simdgroup_load(dOf[d], kv + sg_row * HD + d * 8, HD);
        }
    }
#pragma clang loop unroll(full)
    for (uint d = 0; d < DF; ++d) dQf[d] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);

    const uint row_max = min(q0 + BQ - 1, p.T - 1);
    const uint j_end = p.causal ? (row_max + 1) : p.T;

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

        threadgroup_barrier(mem_flags::mem_threadgroup);
        for (uint e = tid; e < BK * HD; e += MMA_THREADS) {
            const uint jr = e / HD, d = e % HD;
            const bool ok = jr < block_len;
            const ulong src = (ulong(b) * p.T + jb + jr) * Ckv + hkv * HD + d;
            Ks[jr * LDKV + d] = ok ? K[src] : 0.0f;
            Vs[jr * LDKV + d] = ok ? V[src] : 0.0f;
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);

        // S = Q @ K^T and dP = dO @ V^T
        simdgroup_float8x8 Sf[KF], Pf[KF];
#pragma clang loop unroll(full)
        for (uint j = 0; j < KF; ++j) {
            Sf[j] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
            Pf[j] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
        }
#pragma clang loop unroll(full)
        for (uint d = 0; d < DF; ++d) {
#pragma clang loop unroll(full)
            for (uint j = 0; j < KF; ++j) {
                simdgroup_float8x8 KTf, VTf;
                simdgroup_load(KTf, Ks + (j * 8) * LDKV + d * 8, LDKV, 0, true);
                simdgroup_load(VTf, Vs + (j * 8) * LDKV + d * 8, LDKV, 0, true);
                simdgroup_multiply_accumulate(Sf[j], Qf[d], KTf, Sf[j]);
                simdgroup_multiply_accumulate(Pf[j], dOf[d], VTf, Pf[j]);
            }
        }
#pragma clang loop unroll(full)
        for (uint j = 0; j < KF; ++j) {
            simdgroup_store(Sf[j], Sbuf + sg_row * LDS + j * 8, LDS);
            simdgroup_store(Pf[j], Pbuf + sg_row * LDS + j * 8, LDS);
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);

        // dS = P * (dP - D_i) * scale, written back over Sbuf
        if (tid < BQ) {
            const uint i = q0 + tid;
            threadgroup float* srow = Sbuf + tid * LDS;
            threadgroup const float* prow = Pbuf + tid * LDS;
            const uint valid =
                (i >= p.T) ? 0u
                           : (p.causal ? min(block_len, (i >= jb) ? (i - jb + 1) : 0u)
                                       : block_len);
            const float li = row_L[tid], di = row_D[tid];
            for (uint j = 0; j < BK; ++j) {
                if (j < valid) {
                    const float prob = exp(srow[j] * p.scale - li);
                    srow[j] = prob * (prow[j] - di) * p.scale;
                } else {
                    srow[j] = 0.0f;
                }
            }
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);

        // dQ += dS @ K
#pragma clang loop unroll(full)
        for (uint j = 0; j < KF; ++j) {
            simdgroup_float8x8 dSf;
            simdgroup_load(dSf, Sbuf + sg_row * LDS + j * 8, LDS);
#pragma clang loop unroll(full)
            for (uint d = 0; d < DF; ++d) {
                simdgroup_float8x8 Kf;
                simdgroup_load(Kf, Ks + (j * 8) * LDKV + d * 8, LDKV);
                simdgroup_multiply_accumulate(dQf[d], dSf, Kf, dQf[d]);
            }
        }
    }

    threadgroup_barrier(mem_flags::mem_threadgroup);
#pragma clang loop unroll(full)
    for (uint d = 0; d < DF; ++d)
        simdgroup_store(dQf[d], kv + sg_row * HD + d * 8, HD);
    threadgroup_barrier(mem_flags::mem_threadgroup);
    for (uint e = tid; e < BQ * HD; e += MMA_THREADS) {
        const uint r = e / HD, d = e % HD;
        const uint i = q0 + r;
        if (i >= p.T) continue;
        dQ[(ulong(b) * p.T + i) * Cq + h * HD + d] += kv[r * HD + d];
    }
}

// dK and dV are separate kernels. Combined, one thread holds K, V, dK and dV
// as fragments and the compiler spills: gpudebug reported 4352 spilled bytes
// and 111 temp registers for the fused version, which made it the single most
// expensive backward kernel. Split, dV carries K+dV and dK carries K+V+dK, and
// both fit. The cost is recomputing S^T = K @ Q^T in each.

template <uint HD>
kernel void flash_attn_bwd_dv_mma(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 tgid [[threadgroup_position_in_grid]],
                                  uint tid  [[thread_index_in_threadgroup]],
                                  uint sgid [[simdgroup_index_in_threadgroup]]) {
    constexpr uint DF = HD / 8;
    constexpr uint LDQ = HD + PADF;
    constexpr uint IF = BQB / 8;

    threadgroup float qo[2 * BQB * LDQ];
    threadgroup float* Qs = qo;
    threadgroup float* dOs = qo + BQB * LDQ;
    threadgroup float PTbuf[BQ * LDQB];
    threadgroup float col_L[BQB];

    const uint kv_blocks = (p.T + BQ - 1) / BQ;
    const uint jbk = tgid % kv_blocks;
    const uint hkv = (tgid / kv_blocks) % p.HKV;
    const uint b = tgid / (kv_blocks * p.HKV);

    const uint rep = p.H / p.HKV;
    const uint Cq = p.H * HD;
    const uint Ckv = p.HKV * HD;
    const uint j0 = jbk * BQ;
    const uint sg_row = sgid * 8;

    simdgroup_float8x8 Kf[DF], dVf[DF];
    {
        threadgroup_barrier(mem_flags::mem_threadgroup);
        for (uint e = tid; e < BQ * HD; e += MMA_THREADS) {
            const uint r = e / HD, d = e % HD;
            const uint j = j0 + r;
            qo[r * HD + d] =
                (j < p.T) ? K[(ulong(b) * p.T + j) * Ckv + hkv * HD + d] : 0.0f;
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
#pragma clang loop unroll(full)
        for (uint d = 0; d < DF; ++d) {
            simdgroup_load(Kf[d], qo + sg_row * HD + d * 8, HD);
            dVf[d] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
        }
    }

    const uint i_start = p.causal ? (j0 / BQB) * BQB : 0;
    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 ib = i_start; ib < p.T; ib += BQB) {
            const uint blk = min(BQB, p.T - ib);
            threadgroup_barrier(mem_flags::mem_threadgroup);
            for (uint e = tid; e < BQB * HD; e += MMA_THREADS) {
                const uint ir = e / HD, d = e % HD;
                const bool ok = ir < blk;
                const ulong src = (ulong(b) * p.T + ib + ir) * Cq + h * HD + d;
                Qs[ir * LDQ + d] = ok ? Q[src] : 0.0f;
                dOs[ir * LDQ + d] = ok ? dO[src] : 0.0f;
            }
            if (tid < BQB) col_L[tid] = (tid < blk) ? Lh[ib + tid] : 0.0f;
            threadgroup_barrier(mem_flags::mem_threadgroup);

            simdgroup_float8x8 STf[IF];
#pragma clang loop unroll(full)
            for (uint i = 0; i < IF; ++i)
                STf[i] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
#pragma clang loop unroll(full)
            for (uint d = 0; d < DF; ++d) {
#pragma clang loop unroll(full)
                for (uint i = 0; i < IF; ++i) {
                    simdgroup_float8x8 QTf;
                    simdgroup_load(QTf, Qs + (i * 8) * LDQ + d * 8, LDQ, 0, true);
                    simdgroup_multiply_accumulate(STf[i], Kf[d], QTf, STf[i]);
                }
            }
#pragma clang loop unroll(full)
            for (uint i = 0; i < IF; ++i)
                simdgroup_store(STf[i], PTbuf + sg_row * LDQB + i * 8, LDQB);
            threadgroup_barrier(mem_flags::mem_threadgroup);

            if (tid < BQ) {
                const uint j = j0 + tid;
                threadgroup float* prow = PTbuf + tid * LDQB;
                for (uint i = 0; i < BQB; ++i) {
                    const uint iq = ib + i;
                    const bool ok = (i < blk) && (j < p.T) && (!p.causal || iq >= j);
                    prow[i] = ok ? exp(prow[i] * p.scale - col_L[i]) : 0.0f;
                }
            }
            threadgroup_barrier(mem_flags::mem_threadgroup);

#pragma clang loop unroll(full)
            for (uint i = 0; i < IF; ++i) {
                simdgroup_float8x8 PTx;
                simdgroup_load(PTx, PTbuf + sg_row * LDQB + i * 8, LDQB);
#pragma clang loop unroll(full)
                for (uint d = 0; d < DF; ++d) {
                    simdgroup_float8x8 dOf;
                    simdgroup_load(dOf, dOs + (i * 8) * LDQ + d * 8, LDQ);
                    simdgroup_multiply_accumulate(dVf[d], PTx, dOf, dVf[d]);
                }
            }
        }
    }

    threadgroup_barrier(mem_flags::mem_threadgroup);
#pragma clang loop unroll(full)
    for (uint d = 0; d < DF; ++d)
        simdgroup_store(dVf[d], qo + sg_row * HD + d * 8, HD);
    threadgroup_barrier(mem_flags::mem_threadgroup);
    for (uint e = tid; e < BQ * HD; e += MMA_THREADS) {
        const uint r = e / HD, d = e % HD;
        const uint j = j0 + r;
        if (j >= p.T) continue;
        dV[(ulong(b) * p.T + j) * Ckv + hkv * HD + d] += qo[r * HD + d];
    }
}

template <uint HD>
kernel void flash_attn_bwd_dk_mma(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 tgid [[threadgroup_position_in_grid]],
                                  uint tid  [[thread_index_in_threadgroup]],
                                  uint sgid [[simdgroup_index_in_threadgroup]]) {
    constexpr uint DF = HD / 8;
    constexpr uint LDQ = HD + PADF;
    constexpr uint IF = BQB / 8;

    threadgroup float qo[2 * BQB * LDQ];
    threadgroup float* Qs = qo;
    threadgroup float* dOs = qo + BQB * LDQ;
    threadgroup float STbuf[BQ * LDQB];
    threadgroup float PTbuf[BQ * LDQB];
    threadgroup float col_L[BQB], col_D[BQB];

    const uint kv_blocks = (p.T + BQ - 1) / BQ;
    const uint jbk = tgid % kv_blocks;
    const uint hkv = (tgid / kv_blocks) % p.HKV;
    const uint b = tgid / (kv_blocks * p.HKV);

    const uint rep = p.H / p.HKV;
    const uint Cq = p.H * HD;
    const uint Ckv = p.HKV * HD;
    const uint j0 = jbk * BQ;
    const uint sg_row = sgid * 8;

    simdgroup_float8x8 Kf[DF], Vf[DF], dKf[DF];
    for (uint pass = 0; pass < 2; ++pass) {
        threadgroup_barrier(mem_flags::mem_threadgroup);
        device const float* src = (pass == 0) ? K : V;
        for (uint e = tid; e < BQ * HD; e += MMA_THREADS) {
            const uint r = e / HD, d = e % HD;
            const uint j = j0 + r;
            qo[r * HD + d] =
                (j < p.T) ? src[(ulong(b) * p.T + j) * Ckv + hkv * HD + d] : 0.0f;
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
#pragma clang loop unroll(full)
        for (uint d = 0; d < DF; ++d) {
            if (pass == 0) simdgroup_load(Kf[d], qo + sg_row * HD + d * 8, HD);
            else           simdgroup_load(Vf[d], qo + sg_row * HD + d * 8, HD);
        }
    }
#pragma clang loop unroll(full)
    for (uint d = 0; d < DF; ++d) dKf[d] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);

    const uint i_start = p.causal ? (j0 / BQB) * BQB : 0;
    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 ib = i_start; ib < p.T; ib += BQB) {
            const uint blk = min(BQB, p.T - ib);
            threadgroup_barrier(mem_flags::mem_threadgroup);
            for (uint e = tid; e < BQB * HD; e += MMA_THREADS) {
                const uint ir = e / HD, d = e % HD;
                const bool ok = ir < blk;
                const ulong src = (ulong(b) * p.T + ib + ir) * Cq + h * HD + d;
                Qs[ir * LDQ + d] = ok ? Q[src] : 0.0f;
                dOs[ir * LDQ + d] = ok ? dO[src] : 0.0f;
            }
            if (tid < BQB) {
                col_L[tid] = (tid < blk) ? Lh[ib + tid] : 0.0f;
                col_D[tid] = (tid < blk) ? Dh[ib + tid] : 0.0f;
            }
            threadgroup_barrier(mem_flags::mem_threadgroup);

            simdgroup_float8x8 STf[IF], PTf[IF];
#pragma clang loop unroll(full)
            for (uint i = 0; i < IF; ++i) {
                STf[i] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
                PTf[i] = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);
            }
#pragma clang loop unroll(full)
            for (uint d = 0; d < DF; ++d) {
#pragma clang loop unroll(full)
                for (uint i = 0; i < IF; ++i) {
                    simdgroup_float8x8 QTf, dOTf;
                    simdgroup_load(QTf, Qs + (i * 8) * LDQ + d * 8, LDQ, 0, true);
                    simdgroup_load(dOTf, dOs + (i * 8) * LDQ + d * 8, LDQ, 0, true);
                    simdgroup_multiply_accumulate(STf[i], Kf[d], QTf, STf[i]);
                    simdgroup_multiply_accumulate(PTf[i], Vf[d], dOTf, PTf[i]);
                }
            }
#pragma clang loop unroll(full)
            for (uint i = 0; i < IF; ++i) {
                simdgroup_store(STf[i], STbuf + sg_row * LDQB + i * 8, LDQB);
                simdgroup_store(PTf[i], PTbuf + sg_row * LDQB + i * 8, LDQB);
            }
            threadgroup_barrier(mem_flags::mem_threadgroup);

            if (tid < BQ) {
                const uint j = j0 + tid;
                threadgroup float* srow = STbuf + tid * LDQB;
                threadgroup const float* prow = PTbuf + tid * LDQB;
                for (uint i = 0; i < BQB; ++i) {
                    const uint iq = ib + i;
                    const bool ok = (i < blk) && (j < p.T) && (!p.causal || iq >= j);
                    srow[i] = ok ? exp(srow[i] * p.scale - col_L[i]) *
                                       (prow[i] - col_D[i]) * p.scale
                                 : 0.0f;
                }
            }
            threadgroup_barrier(mem_flags::mem_threadgroup);

#pragma clang loop unroll(full)
            for (uint i = 0; i < IF; ++i) {
                simdgroup_float8x8 STx;
                simdgroup_load(STx, STbuf + sg_row * LDQB + i * 8, LDQB);
#pragma clang loop unroll(full)
                for (uint d = 0; d < DF; ++d) {
                    simdgroup_float8x8 Qf2;
                    simdgroup_load(Qf2, Qs + (i * 8) * LDQ + d * 8, LDQ);
                    simdgroup_multiply_accumulate(dKf[d], STx, Qf2, dKf[d]);
                }
            }
        }
    }

    threadgroup_barrier(mem_flags::mem_threadgroup);
#pragma clang loop unroll(full)
    for (uint d = 0; d < DF; ++d)
        simdgroup_store(dKf[d], qo + sg_row * HD + d * 8, HD);
    threadgroup_barrier(mem_flags::mem_threadgroup);
    for (uint e = tid; e < BQ * HD; e += MMA_THREADS) {
        const uint r = e / HD, d = e % HD;
        const uint j = j0 + r;
        if (j >= p.T) continue;
        dK[(ulong(b) * p.T + j) * Ckv + hkv * HD + d] += qo[r * HD + d];
    }
}

#define INSTANTIATE_FLASH_MMA_BWD(HD)                                               \
    template [[host_name("flash_attn_bwd_dq_mma_f32_hd" #HD)]] kernel void          \
    flash_attn_bwd_dq_mma<HD>(device const float*, device const float*,             \
                              device const float*, device const float*,             \
                              device const float*, device const float*,             \
                              device float*, constant FlashParams&, uint, uint,     \
                              uint);                                                \
    template [[host_name("flash_attn_bwd_dv_mma_f32_hd" #HD)]] kernel void          \
    flash_attn_bwd_dv_mma<HD>(device const float*, device const float*,             \
                              device const float*, device const float*,             \
                              device float*, constant FlashParams&, uint, uint,     \
                              uint);                                                \
    template [[host_name("flash_attn_bwd_dk_mma_f32_hd" #HD)]] kernel void          \
    flash_attn_bwd_dk_mma<HD>(device const float*, device const float*,             \
                              device const float*, device const float*,             \
                              device const float*, device const float*,             \
                              device float*, constant FlashParams&, uint, uint,     \
                              uint);

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