spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// Unfused-but-GPU-resident causal attention with GQA (M4 correctness path;4// the flash-style fused kernel is M5). Probabilities P [B,H,T,T] are5// materialized in f32 for the backward pass — exactly mirroring the CPU6// reference so parity is bit-for-bit meaningful.7//8// Parallelization: one THREAD per query row (b,h,i) in fwd/dq (grid9// B*H*T), and one thread per KV row (b,h,j) in dkv. dkv iterates the10// q-heads sharing its kv-head, so dk/dv accumulate privately — no atomics,11// deterministic.12//13// Layouts: q/o [B,T,H*hd], k/v [B,T,Hkv*hd], P [B,H,T,T] row-major.14#include <metal_stdlib>15using namespace metal;1617struct AttnParams {18    uint B, T, H, HKV, HD;19    float scale;20    uint causal;21    uint window;   // sliding window (keys kept, self incl.); 0 = full22    float softcap; // cap*tanh(s/cap) pre-softmax; 0 = off23};2425kernel void attention_fwd_f32(device const float* Q [[buffer(0)]],26                              device const float* K [[buffer(1)]],27                              device const float* V [[buffer(2)]],28                              device float*       O [[buffer(3)]],29                              device float*       P [[buffer(4)]],30                              constant AttnParams& p [[buffer(5)]],31                              uint gid [[thread_position_in_grid]]) {32    const uint i = gid % p.T;33    const uint h = (gid / p.T) % p.H;34    const uint b = gid / (p.T * p.H);35    if (b >= p.B) return;3637    const uint rep = p.H / p.HKV;38    const uint hkv = h / rep;39    const uint Cq = p.H * p.HD;40    const uint Ckv = p.HKV * p.HD;4142    device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * p.HD;43    device float* prow_base = P + ((ulong(b) * p.H + h) * p.T + i) * p.T;44    const uint jmax = p.causal ? i : p.T - 1;45    const uint jmin = (p.window > 0 && i + 1 > p.window) ? i + 1 - p.window : 0;4647    float m = -FLT_MAX;48    for (uint j = jmin; j <= jmax; ++j) {49        device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;50        float s = 0.0f;51        for (uint d = 0; d < p.HD; ++d) s = fma(qi[d], kj[d], s);52        s *= p.scale;53        if (p.softcap > 0.0f) s = p.softcap * precise::tanh(s / p.softcap);54        prow_base[j] = s;55        m = max(m, s);56    }57    float sum = 0.0f;58    for (uint j = jmin; j <= jmax; ++j) {59        const float e = exp(prow_base[j] - m);60        prow_base[j] = e;61        sum += e;62    }63    const float inv = 1.0f / sum;64    for (uint j = 0; j < jmin; ++j) prow_base[j] = 0.0f;65    for (uint j = jmin; j <= jmax; ++j) prow_base[j] *= inv;66    for (uint j = jmax + 1; j < p.T; ++j) prow_base[j] = 0.0f;6768    device float* oi = O + (ulong(b) * p.T + i) * Cq + h * p.HD;69    for (uint d = 0; d < p.HD; ++d) oi[d] = 0.0f;70    for (uint j = jmin; j <= jmax; ++j) {71        const float prob = prow_base[j];72        if (prob == 0.0f) continue;73        device const float* vj = V + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;74        for (uint d = 0; d < p.HD; ++d) oi[d] = fma(prob, vj[d], oi[d]);75    }76}7778// D[b,h,i] = dO_i · O_i. This is the FlashAttention-2 preprocessing term and79// it equals rowsum(dP ∘ P) exactly:80//   Σ_j P_ij (dO_i · V_j) = dO_i · (Σ_j P_ij V_j) = dO_i · O_i81// Computing it once per query row instead of per (i,j) pair takes the dk/dv82// kernel from O(T³·hd) to O(T²·hd).83kernel void attention_bwd_d_f32(device const float* O  [[buffer(0)]],84                                device const float* dO [[buffer(1)]],85                                device float*       D  [[buffer(2)]],86                                constant AttnParams& p [[buffer(3)]],87                                uint gid [[thread_position_in_grid]]) {88    const uint i = gid % p.T;89    const uint h = (gid / p.T) % p.H;90    const uint b = gid / (p.T * p.H);91    if (b >= p.B) return;9293    const uint Cq = p.H * p.HD;94    device const float* oi = O + (ulong(b) * p.T + i) * Cq + h * p.HD;95    device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * p.HD;96    float acc = 0.0f;97    for (uint d = 0; d < p.HD; ++d) acc = fma(doi[d], oi[d], acc);98    D[(ulong(b) * p.H + h) * p.T + i] = acc;99}100101// dq for one query row (b,h,i): dq_i = Σ_j dS_ij * K_j * scale, where102// dS = P ∘ (dP − D_i) and dP_ij = dO_i · V_j.103kernel void attention_bwd_dq_f32(device const float* Q  [[buffer(0)]],104                                 device const float* K  [[buffer(1)]],105                                 device const float* V  [[buffer(2)]],106                                 device const float* P  [[buffer(3)]],107                                 device const float* dO [[buffer(4)]],108                                 device const float* D  [[buffer(5)]],109                                 device float*       dQ [[buffer(6)]],110                                 constant AttnParams& p [[buffer(7)]],111                                 uint gid [[thread_position_in_grid]]) {112    const uint i = gid % p.T;113    const uint h = (gid / p.T) % p.H;114    const uint b = gid / (p.T * p.H);115    if (b >= p.B) return;116117    const uint rep = p.H / p.HKV;118    const uint hkv = h / rep;119    const uint Cq = p.H * p.HD;120    const uint Ckv = p.HKV * p.HD;121122    device const float* prow = P + ((ulong(b) * p.H + h) * p.T + i) * p.T;123    device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * p.HD;124    device float* dqi = dQ + (ulong(b) * p.T + i) * Cq + h * p.HD;125126    const float row_dot = D[(ulong(b) * p.H + h) * p.T + i];127    for (uint j = 0; j < p.T; ++j) {128        const float prob = prow[j];129        if (prob == 0.0f) continue;130        device const float* vj = V + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;131        device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;132        float dp = 0.0f;133        for (uint d = 0; d < p.HD; ++d) dp = fma(doi[d], vj[d], dp);134        float ds = prob * (dp - row_dot) * p.scale;135        if (p.softcap > 0.0f) {136            device const float* qq = Q + (ulong(b) * p.T + i) * Cq + h * p.HD;137            float s = 0.0f;138            for (uint d = 0; d < p.HD; ++d) s = fma(qq[d], kj[d], s);139            const float t = precise::tanh(s * p.scale / p.softcap);140            ds *= 1.0f - t * t;141        }142        for (uint d = 0; d < p.HD; ++d) dqi[d] = fma(ds, kj[d], dqi[d]);143    }144}145146// dk/dv for one KV row (b,hkv,j): sums over the q-heads sharing this147// kv-head and all query rows i (P_ij = 0 above the diagonal already).148kernel void attention_bwd_dkv_f32(device const float* Q  [[buffer(0)]],149                                  device const float* K  [[buffer(1)]],150                                  device const float* V  [[buffer(2)]],151                                  device const float* P  [[buffer(3)]],152                                  device const float* dO [[buffer(4)]],153                                  device const float* D  [[buffer(5)]],154                                  device float*       dK [[buffer(6)]],155                                  device float*       dV [[buffer(7)]],156                                  constant AttnParams& p [[buffer(8)]],157                                  uint gid [[thread_position_in_grid]]) {158    const uint j = gid % p.T;159    const uint hkv = (gid / p.T) % p.HKV;160    const uint b = gid / (p.T * p.HKV);161    if (b >= p.B) return;162163    const uint rep = p.H / p.HKV;164    const uint Cq = p.H * p.HD;165    const uint Ckv = p.HKV * p.HD;166167    device const float* vj = V + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;168    device float* dkj = dK + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;169    device float* dvj = dV + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;170171    for (uint r = 0; r < rep; ++r) {172        const uint h = hkv * rep + r;173        for (uint i = 0; i < p.T; ++i) {174            const float prob = P[((ulong(b) * p.H + h) * p.T + i) * p.T + j];175            if (prob == 0.0f) continue;176            device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * p.HD;177            device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * p.HD;178179            const float row_dot = D[(ulong(b) * p.H + h) * p.T + i];180            float dp = 0.0f;181            for (uint d = 0; d < p.HD; ++d) dp = fma(doi[d], vj[d], dp);182            float ds = prob * (dp - row_dot) * p.scale;183            if (p.softcap > 0.0f) {184                device const float* kk = K + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;185                float s = 0.0f;186                for (uint d = 0; d < p.HD; ++d) s = fma(qi[d], kk[d], s);187                const float t = precise::tanh(s * p.scale / p.softcap);188                ds *= 1.0f - t * t;189            }190            for (uint d = 0; d < p.HD; ++d) {191                dvj[d] = fma(prob, doi[d], dvj[d]);192                dkj[d] = fma(ds, qi[d], dkj[d]);193            }194        }195    }196}197