/*Kernels for attention forward pass. If you do not have CUDNN, you can remove ENABLE_CUDNN to run the other kernels See the README for cuDNN install instructions Compile example with cuDNN:nvcc -I/PATH/TO/cudnn-frontend/include -DENABLE_CUDNN -O3 --use_fast_math --lcublas -lcublasLt -lcudnn attention_forward.cu -o attention_forward Compile example without cuDNN:nvcc -O3 --use_fast_math -lcublas -lcublasLt attention_forward.cu -o attention_forward version 1 is naive port from CPU code to kernel, parallelize over batch, time, heads only./attention_forward 1 version 2 is a naive implementation of flash attention, taken, adapted fromhttps://github.com/tspeterkim/flash-attention-minimaland with help fromhttps://github.com/leloykun/flash-hyperbolic-attention-minimalsadly, this flash attention version seems about 3X slower than the naive version./attention_forward 2 version 3 is a cuBLAS + softmax version, similar to the PyTorch implementationcuBLAS is used both to calculate the QK^T and the final weighted sumthe softmax is calculated using a custom, efficient kernel as wellthis turns out to be ~20X faster than (1) nice./attention_forward 3 version 4 is a further optimized kernel that fuses the scale operation,uses a directly autoregressive softmax, and uses the online softmax algorithm../attention_forward 4 version 5 is a FP16 version of kernel 4./attention_forward 5 version 6 is kernel 5 skipping (un)permute (unrealistic but useful comparison point) version 10 is using cuDNN Flash Attention using FP16 or BF16, see:https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Attention.md./attention_forward 10 version 11 is kernel 10 skipping FP16/FP32 conversions (full FP16/BF16 network)./attention_forward 11*///#define ENABLE_CUDNN // can be enabled via nvcc "-DENABLE_CUDNN"#include <stdio.h>#include <stdlib.h>#include <assert.h>#include <float.h>#include <cublas_v2.h>#include <cuda_runtime.h>#include <cuda_bf16.h>#include <cooperative_groups.h>#include <cooperative_groups/reduce.h> #define ENABLE_BF16#include "common.h" // ----------------------------------------------------------------------------// CUDA & cuDNN setupstatic bool first_run_validation = true; // always run e.g. permute on 1st run #ifdef ENABLE_CUDNN#include <cudnn_frontend.h>namespace fe = cudnn_frontend;#if CUBLAS_LOWP == CUDA_R_16BF#define CUDNN_16BIT fe::DataType_t::BFLOAT16#else#define CUDNN_16BIT fe::DataType_t::HALF#endif static cudnnHandle_t cudnn_handle;static size_t cudnn_workspace_size = 0; // dynamically allocated as needed (up to 256MiB!)static void* cudnn_workspace = NULL; #define checkCudaErr(err) assert((int)err == 0);#define checkCudnnErr(err) assert((int)err == 0);#endif // ENABLE_CUDNN// ----------------------------------------------------------------------------// CPU code reference void attention_forward_cpu(float* out, float* preatt, float* att,                       const float* inp,                       int B, int T, int C, int NH) {    // input is (B, T, 3C) Q,K,V    // preatt, att are (B, NH, T, T)    // output 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++) {                const 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++) {                    const 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];                    }                    val *= scale;                    if (val > maxval) {                        maxval = val;                    }                     preatt_bth[t2] = val;                }                // pad with -INFINITY outside of autoregressive region for debugging comparisons                for (int t2 = t+1; t2 < T; t2++) {                    preatt_bth[t2] = -INFINITY;                }                 // pass 2: calculate the exp and keep track of sum                float expsum = 0.0f;                for (int t2 = 0; t2 <= t; t2++) {                    float expv = expf(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++) {                    const 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];                    }                }            }        }    }} // ----------------------------------------------------------------------------// GPU kernels __global__ void attention_query_key_kernel1(float* preatt, const float* inp,                                           int B, int T, int C, int NH) {    int idx = blockIdx.x * blockDim.x + threadIdx.x;    int total_threads = B * NH * T * T;     if (idx < total_threads) {        int t2 = idx % T;        int t = (idx / T) % T;        if (t2 > t) {            // autoregressive mask            preatt[idx] = -INFINITY;            return;        }        int h = (idx / (T * T)) % NH;        int b = idx / (NH * T * T);         int C3 = C*3;        int hs = C / NH; // head size        const float* query_t = inp + b * T * C3 + t * C3 + h * hs;        const 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];        }        val *= 1.0 / sqrtf(hs);         preatt[idx] = val;    }} __global__ void attention_softmax_kernel1(float* att, const float* preatt,                                         int B, int T, int NH) {    int idx = blockIdx.x * blockDim.x + threadIdx.x;    int total_threads = B * T * NH;     if (idx < total_threads) {        int h = idx % NH;        int t = (idx / NH) % T;        int b = idx / (NH * T);         const 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;         // find maxval        float maxval = -10000.0f; // TODO something better        for (int t2 = 0; t2 <= t; t2++) {            if (preatt_bth[t2] > maxval) {                maxval = preatt_bth[t2];            }        }         // calculate the exp and keep track of sum        float expsum = 0.0f;        for (int t2 = 0; t2 <= t; t2++) {            float expv = expf(preatt_bth[t2] - maxval);            expsum += expv;            att_bth[t2] = expv;        }        float expsum_inv = expsum == 0.0f ? 0.0f : 1.0f / expsum;         // 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;            }        }    }} // warp-level reduction for finding the maximum value__device__ float warpReduceMax(float val) {    for (int offset = 16; offset > 0; offset /= 2) {        val = fmaxf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));    }    return val;} __global__ void softmax_forward_kernel4(float* out, const float* inp, int N, int C) {    // out is (N, C) just like inp. Each row of inp will get softmaxed.    // same as kernel3, but can handle any block size (multiple of 32)    // each row of C elements is handled by block_size threads    // furthermore, each block_size threads get executed in warps of 32 threads     // special reduction operations warpReduceMax/warpReduceSum are used for intra-warp reductions    // shared memory is used for inter-warp reduction    extern __shared__ float shared[];    int idx = blockIdx.x;    int tid = threadIdx.x;    int warpId = threadIdx.x / 32; // warp index within a block    int laneId = threadIdx.x % 32; // thread index within a warp     // the number of warps per block. recall that blockDim.x is block_size    int warpsPerBlock = blockDim.x / 32;     // shared[] must be allocated to have 2 * warpsPerBlock elements    // first half for max values, the second half for sum values    float* maxvals = shared;    float* sumvals = &shared[warpsPerBlock];     // one row of inp, i.e. inp[idx, :] of shape (C,)    const float* x = inp + idx * C;     // first, thread coarsening by directly accessing global memory in series    float maxval = -INFINITY;    for (int i = tid; i < C; i += blockDim.x) {        maxval = fmaxf(maxval, x[i]);    }    // now within-warp reductions for maxval    maxval = warpReduceMax(maxval);     // the 0th thread of each warp writes the maxval of that warp to shared memory    if (laneId == 0) maxvals[warpId] = maxval;    __syncthreads();     // now the 0th thread reduces the maxvals in shared memory, i.e. across warps    if (tid == 0) {        float val = maxvals[tid];        for (int i = 1; i < warpsPerBlock; i++) {            val = fmaxf(val, maxvals[i]);        }        // store the final max in the first position        maxvals[0] = val;    }    __syncthreads();    // broadcast the max to all threads    float offset = maxvals[0];     // compute expf and write the result to global memory    for (int i = tid; i < C; i += blockDim.x) {        // subtract max for numerical stability        out[idx * C + i] = expf(x[i] - offset);    }     // okay now we calculated exp(x - max(x))    // step 2: sum all the values and divide by the sum     // thread coarsening for sum    x = out + idx * C;    float sumval = 0.0f;    for (int i = tid; i < C; i += blockDim.x) {        sumval += x[i];    }    // within-warp reduction for sumval    sumval = warpReduceSum(sumval);     // write sumval to shared memory    if (laneId == 0) sumvals[warpId] = sumval;    __syncthreads();     // inter-thread reduction of sum    if (tid == 0) {        float val = sumvals[tid];        for (int i = 1; i < warpsPerBlock; ++i) {            val += sumvals[i];        }        sumvals[0] = val;    }    __syncthreads();    // broadcast the sum to all threads    float sum = sumvals[0];     // divide the whole row by the sum    for (int i = tid; i < C; i += blockDim.x) {        out[idx * C + i] = x[i] / sum;    }}  __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);    }}  __global__ void attention_value_kernel1(float* out, const float* att, const float* inp,                                       int B, int T, int C, int NH) {    int idx = blockIdx.x * blockDim.x + threadIdx.x;    int total_threads = B * T * NH;     if (idx < total_threads) {        int h = idx % NH;        int t = (idx / NH) % T;        int b = idx / (NH * T);         int C3 = C*3;        int hs = C / NH; // head size         float* out_bth = out + b * T * C + t * C + h * hs;        const float* att_bth = att + b*NH*T*T + h*T*T + t*T;         for (int i = 0; i < hs; i++) { out_bth[i] = 0.0f; }        for (int t2 = 0; t2 <= t; t2++) {           const  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];            }        }    }} __global__void attention_forward_kernel2(    const float* Q,    const float* K,    const float* V,    const int N,    const int d,    const int Tc,    const int Tr,    const int Bc,    const int Br,    const float softmax_scale,    float* l,    float* m,    float* O) {    int tx = threadIdx.x;    int bx = blockIdx.x; int by = blockIdx.y;  // batch and head index     // Offset into Q,K,V,O,l,m - different for each batch and head    int qkv_offset = (bx * gridDim.y * N * d) + (by * N * d);  // gridDim.y = nh    int lm_offset = (bx * gridDim.y * N) + (by * N);  // offset for l and m     // Define SRAM for Q,K,V,S    extern __shared__ float sram[];    int tile_size = Bc * d;  // size of Qi, Kj, Vj    float* Qi = sram;    float* Kj = &sram[tile_size];    float* Vj = &sram[tile_size * 2];    float* S = &sram[tile_size * 3];     for (int j = 0; j < Tc; j++) {         // Load Kj, Vj to SRAM        for (int x = 0; x < d; x++) {            Kj[(tx * d) + x] = K[qkv_offset + (tile_size * j) + (tx * d) + x];            Vj[(tx * d) + x] = V[qkv_offset + (tile_size * j) + (tx * d) + x];        }        __syncthreads();  // such that the inner loop can use the correct Kj, Vj         for (int i = 0; i < Tr; i++)  {            // if past the end of the sequence, break            if (i * Br + tx >= N) {                break;            }             // Load Qi to SRAM, l and m to registers            for (int x = 0; x < d; x++) {                Qi[(tx * d) + x] = Q[qkv_offset + (tile_size * i) + (tx * d) + x];            }            float row_m_prev = m[lm_offset + (Br * i) + tx];            float row_l_prev = l[lm_offset + (Br * i) + tx];             // S = QK^T, row_m = rowmax(S)            // S[tx][y] = Sum_{x = 0}^{d-1} {Qi[tx][x] * Kj[y][x]}            // row_m = Max_{y = 0}^{Bc-1} S[tx][y]            // with causal masking            float row_m = -INFINITY;            for (int y = 0; y < Bc; y++) {                if (j * Bc + y >= N) {                    break;                }                float sum = 0;                for (int x = 0; x < d; x++) {                    sum += Qi[(tx * d) + x] * Kj[(y * d) + x];                }                sum *= softmax_scale;                if (i * Br + tx < j * Bc + y)                    sum = -INFINITY;                S[(Bc * tx) + y] = sum;                 if (sum > row_m)                    row_m = sum;            }

Lines 1–500 of 1390. Load more