/*Kernels for attention backward pass. Compile example:nvcc -O3 --use_fast_math -lcublas -lcublasLt attention_backward.cu -o attention_backward version 1 is a naive first versionOMP_NUM_THREADS=32 ./attention_backward 1 version 2 much ensures better load-balancing by having independent threads for each batch and attention headOMP_NUM_THREADS=32 ./attention_backward 2 version 3 uses a full warp to calculate each result (instead of a thread), which enables coalesced memory accessOMP_NUM_THREADS=32 ./attention_backward 3 version 4 improves data reuse in registers by doing 8 values of t3 in one warp.OMP_NUM_THREADS=32 ./attention_backward 4 version 5 reduces the amount of non-fp32 instructions needed by avoiding ifsOMP_NUM_THREADS=32 ./attention_backward 5*/ #include <stdio.h>#include <stdlib.h>#include <assert.h>#include <float.h>#include <cublas_v2.h>#include <cuda_runtime.h>#include <cooperative_groups.h>#include <cooperative_groups/reduce.h>#include <cooperative_groups/scan.h>#include "common.h" // ----------------------------------------------------------------------------// CPU code reference /*NOTE:This version of attention_forward is modified to be consistent with theattention_forward GPU kernel in the following way small but important way:- preatt is only QUERY @ KEY, without the scale- the scale instead moved and fused into the softmax- the full preatt matrix is materialized, even the parts that get masked out    - this doesn't actually change anything due to masking, but it lets us      easily compare to the GPU version, which also does the full, dense sgemmIn this way we'll be able to make sure that preatt and att agree CPU vs GPU*/void attention_forward_cpu(float* out, float* preatt, float* att,                            float* inp,                            int B, int T, int C, int NH) {    // input is (B, T, 3C) holding the query, key, value (Q, K, V) vectors    // preatt, att are (B, NH, T, T). NH = number of heads, T = sequence length    // that holds the pre-attention and post-attention scores (used in backward)    // output is (B, T, C)    // attention is the only layer that mixes information across time    // every other operation is applied at every (b,t) position independently    // (and of course, no layer mixes information across batch)    int C3 = C*3;    int hs = C / NH; // head size    float scale = 1.0 / sqrtf(hs);     #pragma omp parallel for collapse(3)    for (int b = 0; b < B; b++) {        for (int t = 0; t < T; t++) {            for (int h = 0; h < NH; h++) {                float* query_t = inp + b * T * C3 + t * C3 + h * hs;                float* preatt_bth = preatt + b*NH*T*T + h*T*T + t*T;                float* att_bth = att + b*NH*T*T + h*T*T + t*T;                 // pass 1: calculate query dot key and maxval                float maxval = -10000.0f; // TODO something better                for (int t2 = 0; t2 < T; t2++) { // used to be t2 <= t                    float* key_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C; // +C because it's key                     // (query_t) dot (key_t2)                    float val = 0.0f;                    for (int i = 0; i < hs; i++) {                        val += query_t[i] * key_t2[i];                    }                    if (val > maxval) {                        maxval = val;                    }                     preatt_bth[t2] = val;                }                 // pass 2: calculate the exp and keep track of sum                // maxval is being calculated and subtracted only for numerical stability                float expsum = 0.0f;                for (int t2 = 0; t2 <= t; t2++) {                    float expv = expf(scale * (preatt_bth[t2] - maxval));                    expsum += expv;                    att_bth[t2] = expv;                }                float expsum_inv = expsum == 0.0f ? 0.0f : 1.0f / expsum;                 // pass 3: normalize to get the softmax                for (int t2 = 0; t2 < T; t2++) {                    if (t2 <= t) {                        att_bth[t2] *= expsum_inv;                    } else {                        // causal attention mask. not strictly necessary to set to zero here                        // only doing this explicitly for debugging and checking to PyTorch                        att_bth[t2] = 0.0f;                    }                }                 // pass 4: accumulate weighted values into the output of attention                float* out_bth = out + b * T * C + t * C + h * hs;                for (int i = 0; i < hs; i++) { out_bth[i] = 0.0f; }                for (int t2 = 0; t2 <= t; t2++) {                    float* value_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C*2; // +C*2 because it's value                    float att_btht2 = att_bth[t2];                    for (int i = 0; i < hs; i++) {                        out_bth[i] += att_btht2 * value_t2[i];                    }                }            }        }    }} // NOTE: Also contains the re-shuffling of the exact position of "scale"// and when it is applied (after preatt, not "during" preatt)// also, full matrices are materialized, even the parts that get masked outvoid attention_backward_cpu(float* dinp, float* dpreatt, float* datt,                            float* dout, float* inp, float* att,                            int B, int T, int C, int NH) {    // inp/dinp are (B, T, 3C) Q,K,V    // att/datt/dpreatt are (B, NH, T, T)    // dout is (B, T, C)    int C3 = C*3;    int hs = C / NH; // head size    float scale = 1.0 / sqrtf(hs);     for (int b = 0; b < B; b++) {        for (int t = 0; t < T; t++) {            for (int h = 0; h < NH; h++) {                float* att_bth = att + b*NH*T*T + h*T*T + t*T;                float* datt_bth = datt + b*NH*T*T + h*T*T + t*T;                float* dpreatt_bth = dpreatt + b*NH*T*T + h*T*T + t*T;                float* dquery_t = dinp + b * T * C3 + t * C3 + h * hs;                float* query_t = inp + b * T * C3 + t * C3 + h * hs;                 // backward pass 4, through the value accumulation                float* dout_bth = dout + b * T * C + t * C + h * hs;                for (int t2 = 0; t2 < T; t2++) { // ADJUSTED! this was t2 <= t (see note on function)                    float* value_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C*2; // +C*2 because it's value                    float* dvalue_t2 = dinp + b * T * C3 + t2 * C3 + h * hs + C*2;                    for (int i = 0; i < hs; i++) {                        // in the forward pass this was:                        // out_bth[i] += att_bth[t2] * value_t2[i];                        // so now we have:                        datt_bth[t2] += value_t2[i] * dout_bth[i];                        dvalue_t2[i] += att_bth[t2] * dout_bth[i];                    }                }                 // backward pass 2 & 3, the softmax                // note that softmax (like e.g. tanh) doesn't need the input (preatt) to backward                for (int t2 = 0; t2 <= t; t2++) {                    for (int t3 = 0; t3 <= t; t3++) {                        float indicator = t2 == t3 ? 1.0f : 0.0f;                        float local_derivative = att_bth[t2] * (indicator - att_bth[t3]);                        dpreatt_bth[t3] += scale * local_derivative * datt_bth[t2];                    }                }                 // backward pass 1, the query @ key matmul                for (int t2 = 0; t2 <= t; t2++) {                    float* key_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C; // +C because it's key                    float* dkey_t2 = dinp + b * T * C3 + t2 * C3 + h * hs + C; // +C because it's key                    for (int i = 0; i < hs; i++) {                        // in the forward pass this was:                        // preatt_bth[t2] += query_t[i] * key_t2[i]                        // so now we have:                        dquery_t[i] += key_t2[i] * dpreatt_bth[t2];                        dkey_t2[i] += query_t[i] * dpreatt_bth[t2];                    }                }            }        }    }} // ----------------------------------------------------------------------------// GPU kernels// the forward pass that is the sequence [permute, sgemm, softmax, sgemm, unpermute] __global__ void permute_kernel(float* q, float* k, float* v,                               const float* inp,                               int B, int N, int NH, int d) {    // okay so now, this kernel wants Q,K,V to all be of shape (B, NH, N, d)    // but instead, we have a single tensor QKV (inp) of shape (B, N, 3, NH, d)    int idx = blockIdx.x * blockDim.x + threadIdx.x;     // Q[b][nh_][n][d_] = inp[b][n][0][nh_][d_]    if (idx < B * NH * N * d) {        int b = idx / (NH * N * d);        int rest = idx % (NH * N * d);        int nh_ = rest / (N * d);        rest = rest % (N * d);        int n = rest / d;        int d_ = rest % d;         int inp_idx = (b * N * 3 * NH * d) + (n * 3 * NH * d) + (0 * NH * d) + (nh_ * d) + d_;        q[idx] = inp[inp_idx];        k[idx] = inp[inp_idx + NH * d];        v[idx] = inp[inp_idx + 2 * (NH * d)];    }} __global__ void permute_kernel_backward(float* dinp,                                        const float* dq, const float* dk, const float* dv,                                        int B, int N, int NH, int d) {    int idx = blockIdx.x * blockDim.x + threadIdx.x;    if (idx < B * NH * N * d) {        int b = idx / (NH * N * d);        int rest = idx % (NH * N * d);        int nh_ = rest / (N * d);        rest = rest % (N * d);        int n = rest / d;        int d_ = rest % d;         int inp_idx = (b * N * 3 * NH * d) + (n * 3 * NH * d) + (0 * NH * d) + (nh_ * d) + d_;        dinp[inp_idx] += dq[idx];        dinp[inp_idx + NH * d] += dk[idx];        dinp[inp_idx + 2 * (NH * d)] += dv[idx];    }} __global__ void unpermute_kernel(const float* inp, float *out, int B, int N, int NH, int d) {   // out has shape (B, nh, N, d) but we need to unpermute it to (B, N, nh, d)    int idx = blockIdx.x * blockDim.x + threadIdx.x;     // out[b][n][nh_][d_] <- inp[b][nh_][n][d_]    if (idx < B * NH * N * d) {        int b = idx / (NH * N * d);        int rest = idx % (NH * N * d);        int nh_ = rest / (N * d);        rest = rest % (N * d);        int n = rest / d;        int d_ = rest % d;         int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_;        out[other_idx] = inp[idx];    }} __global__ void unpermute_kernel_backward(float* dinp, const float *dout, int B, int N, int NH, int d) {    int idx = blockIdx.x * blockDim.x + threadIdx.x;    if (idx < B * NH * N * d) {        int b = idx / (NH * N * d);        int rest = idx % (NH * N * d);        int nh_ = rest / (N * d);        rest = rest % (N * d);        int n = rest / d;        int d_ = rest % d;         int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_;        dinp[idx] += dout[other_idx];    }} __device__ float& vec_at(float4& vec, int index) {    return reinterpret_cast<float*>(&vec)[index];} __device__ float vec_at(const float4& vec, int index) {    return reinterpret_cast<const float*>(&vec)[index];} __global__ void softmax_forward_kernel5(float* out, float inv_temperature, const float* inp, int N, int T) {    // inp, out shape: (N, T, T), where N = B * NH    // fuses the multiplication by scale inside attention    // directly autoregressive, so we only compute the lower triangular part    // uses the online softmax algorithm    assert(T % 4  == 0);    namespace cg = cooperative_groups;    cg::thread_block block = cg::this_thread_block();    cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);    int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();    if(idx >= N * T) {        return;    }    int own_pos = idx % T;    int pos_by_4 = own_pos / 4;     // one row of inp, i.e. inp[idx, :] of shape (T,)    const float* x = inp + idx * T;     // not INF, so we don't get NaNs accidentally when subtracting two values.    float maxval = -FLT_MAX;    float sumval = 0.0f;     const float4* x_vec = reinterpret_cast<const float4*>(x);    for (int i = warp.thread_rank(); i < pos_by_4; i += warp.size()) {        float4 v = x_vec[i];        float old_maxval = maxval;        for(int k = 0; k < 4; ++k) {            maxval = fmaxf(maxval, vec_at(v, k));        }        sumval *= expf(inv_temperature * (old_maxval - maxval));        for(int k = 0; k < 4; ++k) {            sumval += expf(inv_temperature * (vec_at(v, k) - maxval));        }    }     if(4*pos_by_4 + warp.thread_rank() <= own_pos) {        float old_maxval = maxval;        maxval = fmaxf(maxval, x[4*pos_by_4 + warp.thread_rank()]);        sumval *= expf(inv_temperature * (old_maxval - maxval));        sumval += expf(inv_temperature * (x[4*pos_by_4 + warp.thread_rank()] - maxval));    }     float global_maxval = cg::reduce(warp, maxval, cg::greater<float>{});    sumval *= expf(inv_temperature * (maxval - global_maxval));     float sum = cg::reduce(warp, sumval, cg::plus<float>{});    float norm = 1.f / sum;     // divide the whole row by the sum    for (int i = warp.thread_rank(); i <= own_pos; i += warp.size()) {        // recalculation is faster than doing the round-trip through memory.        float ev = expf(inv_temperature * (__ldcs(x + i) - global_maxval));        __stcs(out + idx * T + i, ev * norm);    }} // naive kernel to backward through an autoregressive softmax, just to get correctness__global__ void softmax_autoregressive_backward_kernel1(float* dpreatt, const float* datt, const float* att,                                                     int B, int T, int C, int NH) {    // dpreatt, datt, att are all (B, NH, T, T)    int t3 = blockIdx.x * blockDim.x + threadIdx.x;    if (t3 < T) {        int hs = C / NH; // head size        float scale = 1.0f / sqrtf(hs);        for (int b = 0; b < B; b++) {            for (int h = 0; h < NH; h++) {                for (int t = t3; t < T; t++) {                    const float* att_bth = att + b*NH*T*T + h*T*T + t*T;                    const float* datt_bth = datt + b*NH*T*T + h*T*T + t*T;                    float* dpreatt_bth = dpreatt + b*NH*T*T + h*T*T + t*T;                    float accum = 0.0f;                    for (int t2 = 0; t2 <= t; t2++) {                        float indicator = t2 == t3 ? 1.0f : 0.0f;                        float local_derivative = att_bth[t2] * (indicator - att_bth[t3]);                        accum +=  scale * local_derivative * datt_bth[t2];                    }                    dpreatt_bth[t3] = accum;                }            }        }    }} // parallelize across t,b,h__global__ void softmax_autoregressive_backward_kernel2(float* dpreatt, const float* datt, const float* att,                                                     int B, int T, int C, int NH) {    int t3 = blockIdx.x * blockDim.x + threadIdx.x;    int idx = blockIdx.y * T * T;    if (t3 >= T) { return; }     int hs = C / NH; // head size    float scale = 1.0f / sqrtf(hs);    for (int t = t3; t < T; t++) {        float result = 0.0;        const float* att_bth = att + idx + t*T;        const float* datt_bth = datt + idx + t*T;        float* dpreatt_bth = dpreatt + idx + t*T;         for (int t2 = 0; t2 <= t; t2++) {            float indicator = t2 == t3 ? 1.0f : 0.0f;            float local_derivative = att_bth[t2] * (indicator - att_bth[t3]);            result += scale * local_derivative * datt_bth[t2];        }         dpreatt_bth[t3] = result;    }} // parallelize across t,b,h__global__ void softmax_autoregressive_backward_kernel3(float* dpreatt, const float* datt, const float* att,                                                     int B, int T, int C, int NH) {    namespace cg = cooperative_groups;    cg::thread_block block = cg::this_thread_block();    cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);    int t3 = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();     int idx = blockIdx.y * T * T;    if (t3 >= T) { return; }     int hs = C / NH; // head size    float scale = 1.0f / sqrtf(hs);    for (int t = t3; t < T; t++) {        float result = 0.0;        const float* att_bth = att + idx + t*T;        const float* datt_bth = datt + idx + t*T;        float* dpreatt_bth = dpreatt + idx + t*T;        const float att_at_t3 = att_bth[t3];         for (int t2 = warp.thread_rank(); t2 <= t; t2 += warp.size()) {            float indicator = t2 == t3 ? 1.0f : 0.0f;            float local_derivative = att_bth[t2] * (indicator - att_at_t3);            result += local_derivative * datt_bth[t2];        }         result = cg::reduce(warp, result, cg::plus<float>());        if(warp.thread_rank() == 0) {            dpreatt_bth[t3] = scale * result;        }    }}__global__ void softmax_autoregressive_backward_kernel4(float* __restrict__ dpreatt, const float* __restrict__ datt,                                                        const float* __restrict__ att,                                                        int B, int T, int C, int NH) {    constexpr int UNROLL = 8;    namespace cg = cooperative_groups;    cg::thread_block block = cg::this_thread_block();    cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);    int t3 = UNROLL * (blockIdx.x * warp.meta_group_size() + warp.meta_group_rank());     int idx = blockIdx.y * T * T;    if (t3 >= T) { return; }     int hs = C / NH; // head size    float scale = 1.0f / sqrtf(hs);     // the innermost loop combines different values of t2 with different values of t.    // by handling [t3, t3 + UNROLL) in one thread, we get much better memory reuse:    // any t3/t-dependent value can be loaded once before the t2 loop.    // within the t2 loop, we can combine each loaded value with each of the UNROLL    // pre-loaded values, thus cutting memory ready by a factor of ~UNROLL.     // one iteration of this loop has to handle the cases    // this may lead to some invalid indices; therefore, we have several    // early-outs in the iteration over k below.    for (int t = t3; t < T; t++) {        float result[UNROLL] = {};        const float* att_bth = att + idx + t * T;        const float* datt_bth = datt + idx + t * T;        float* dpreatt_bth = dpreatt + idx + t * T;         float att_at_t3[UNROLL];        for(int k = 0; k < UNROLL; ++k) {            if (t < t3 + k) continue;            att_at_t3[k] = att_bth[t3 + k];        }         for (int t2 = warp.thread_rank(); t2 <= t; t2 += warp.size()) {            float att_t2 = att_bth[t2];            float datt_t2 = datt_bth[t2];            for(int k = 0; k < UNROLL; ++k) {                if (t < t3 + k) continue;                float indicator = t2 == (t3 + k) ? 1.0f : 0.0f;                float local_derivative = att_t2 * (indicator - att_at_t3[k]);                result[k] += local_derivative * datt_t2;            }        }         for(int k = 0; k < UNROLL; ++k) {            result[k] = cg::reduce(warp, result[k], cg::plus<float>());        }        if (warp.thread_rank() < UNROLL) {            dpreatt_bth[t3 + warp.thread_rank()] = scale * result[warp.thread_rank()];        }    }} __global__ void softmax_autoregressive_backward_kernel5(float* __restrict__ dpreatt, const float* __restrict__ datt,                                                        const float* __restrict__ att,                                                        int B, int T, int C, int NH) {    constexpr int UNROLL = 8;    namespace cg = cooperative_groups;    cg::thread_block block = cg::this_thread_block();    cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);    int t3 = UNROLL * (blockIdx.x * warp.meta_group_size() + warp.meta_group_rank());     int idx = blockIdx.y * T * T;    if (t3 >= T) { return; }     int hs = C / NH; // head size    float scale = 1.0f / sqrtf(hs);    for (int t = t3; t < T; t++) {        float result[UNROLL] = {};        const float* att_bth = att + idx + t * T;        const float* datt_bth = datt + idx + t * T;        float* dpreatt_bth = dpreatt + idx + t * T;         float att_at_t3[UNROLL];        for(int k = 0; k < UNROLL; ++k) {            // if t < t3+k, we're out of bounds.            // in that case, we don't care what we read, because later on,            // we won't write the corresponding result. So just clip to            // make sure this is a valid (in-bounds) memory access.            att_at_t3[k] = att_bth[min(t, t3 + k)];        }         // the code below is actually just a for loop; except,        // we have to do something special in one iteration in

Lines 1–500 of 1154. Load more