wehub-resource-sync
1391 lines
53 KiB
Plaintext
1391 lines
53 KiB
Plaintext
/*
|
|
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 from
|
|
https://github.com/tspeterkim/flash-attention-minimal
|
|
and with help from
|
|
https://github.com/leloykun/flash-hyperbolic-attention-minimal
|
|
sadly, 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 implementation
|
|
cuBLAS is used both to calculate the QK^T and the final weighted sum
|
|
the softmax is calculated using a custom, efficient kernel as well
|
|
this 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 setup
|
|
static 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 = -FLT_MAX;
|
|
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 = -FLT_MAX;
|
|
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;
|
|
}
|
|
|
|
// implement softmax with causal masking
|
|
// P = exp(S - row_m), row_l = rowsum(P)
|
|
// P[tx][y] = exp(S[tx][y] - row_m)
|
|
float row_l = 0;
|
|
for (int y = 0; y < Bc; y++) {
|
|
if (j * Bc + y >= N) {
|
|
break;
|
|
}
|
|
if (i * Br + tx < j * Bc + y)
|
|
S[(Bc * tx) + y] = 0;
|
|
else
|
|
S[(Bc * tx) + y] = __expf(S[(Bc * tx) + y] - row_m);
|
|
row_l += S[(Bc * tx) + y];
|
|
}
|
|
|
|
// Compute new m and l
|
|
float row_m_new = max(row_m_prev, row_m);
|
|
float row_l_new = (__expf(row_m_prev - row_m_new) * row_l_prev) + (__expf(row_m - row_m_new) * row_l);
|
|
|
|
// Write O, l, m to HBM
|
|
for (int x = 0; x < d; x++) {
|
|
float pv = 0; // Pij * Vj
|
|
for (int y = 0; y < Bc; y++) {
|
|
if (j * Bc + y >= N) {
|
|
break;
|
|
}
|
|
pv += S[(Bc * tx) + y] * Vj[(y * d) + x];
|
|
}
|
|
O[qkv_offset + (tile_size * i) + (tx * d) + x] = (1 / row_l_new) \
|
|
* ((row_l_prev * __expf(row_m_prev - row_m_new) * O[qkv_offset + (tile_size * i) + (tx * d) + x]) \
|
|
+ (__expf(row_m - row_m_new) * pv));
|
|
}
|
|
m[lm_offset + (Br * i) + tx] = row_m_new;
|
|
l[lm_offset + (Br * i) + tx] = row_l_new;
|
|
}
|
|
__syncthreads(); // otherwise, thread can use the wrong Kj, Vj in inner loop
|
|
}
|
|
}
|
|
|
|
__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 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 scale_kernel(float* inp, float scale, int B, int NH, int T) {
|
|
// scales the pre-softmax attention scores by scale
|
|
// and sets the autoregressive locations to -INFINITY
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (idx < B * NH * T * T) {
|
|
int rest = idx % (NH * T * T);
|
|
rest = rest % (T * T);
|
|
int t2 = rest / T;
|
|
int t = rest % T;
|
|
if (t > t2) {
|
|
inp[idx] = -INFINITY;
|
|
} else {
|
|
inp[idx] *= scale;
|
|
}
|
|
}
|
|
}
|
|
|
|
// direct translation of the CPU kernel. Each warp handles ont (b, h, t) combination.
|
|
// The important changes compared to the CPU version:
|
|
// - each inner loop is handled by a warp
|
|
// - don't write non-autoregressive parts
|
|
// - reordered the last loops so that we can do all writing in the outer loop.
|
|
__global__ void attention_forward_fused1(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);
|
|
|
|
namespace cg = cooperative_groups;
|
|
cg::thread_block block = cg::this_thread_block();
|
|
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
|
|
int t = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();
|
|
int h = blockIdx.y;
|
|
int b = blockIdx.z;
|
|
|
|
if(t >= T) return;
|
|
|
|
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 = -INFINITY;
|
|
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 = warp.thread_rank(); i < hs; i += warp.size()) {
|
|
val += query_t[i] * key_t2[i];
|
|
}
|
|
val = cg::reduce(warp, val, cg::plus<float>{});
|
|
val *= scale;
|
|
maxval = max(maxval, val);
|
|
if(warp.thread_rank() == 0) {
|
|
preatt_bth[t2] = val;
|
|
}
|
|
}
|
|
|
|
// pass 2: calculate the exp and keep track of sum
|
|
float expsum = 0.0f;
|
|
for (int t2 = warp.thread_rank(); t2 <= t; t2 += warp.size()) {
|
|
float expv = expf(preatt_bth[t2] - maxval);
|
|
expsum += expv;
|
|
}
|
|
|
|
expsum = cg::reduce(warp, expsum, cg::plus<float>{});
|
|
|
|
float expsum_inv = expsum == 0.0f ? 0.0f : 1.0f / expsum;
|
|
|
|
// pass 3: normalize to get the softmax is combined with the next loop to reduce memory round-trips
|
|
for (int t2 = warp.thread_rank(); t2 <= t; t2 += warp.size()) {
|
|
att_bth[t2] = expf(preatt_bth[t2] - maxval) * expsum_inv;
|
|
}
|
|
|
|
// pass 4: accumulate weighted values into the output of attention
|
|
float* out_bth = out + b * T * C + t * C + h * hs;
|
|
for (int i = warp.thread_rank(); i < hs; i += warp.size()) {
|
|
float o = 0.f;
|
|
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];
|
|
o += att_btht2 * value_t2[i];
|
|
}
|
|
out_bth[i] = o;
|
|
}
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// kernel launcher
|
|
|
|
void attention_forward1(float* out, float* preatt, float* att,
|
|
const float* inp,
|
|
int B, int T, int C, int NH,
|
|
const int block_size) {
|
|
// attention calculation
|
|
int total_threads = B * NH * T * T;
|
|
int num_blocks = ceil_div(total_threads, block_size);
|
|
attention_query_key_kernel1<<<num_blocks, block_size>>>(preatt, inp, B, T, C, NH);
|
|
// softmax and value accumulation
|
|
total_threads = B * T * NH;
|
|
num_blocks = ceil_div(total_threads, block_size);
|
|
attention_softmax_kernel1<<<num_blocks, block_size>>>(att, preatt, B, T, NH);
|
|
attention_value_kernel1<<<num_blocks, block_size>>>(out, att, inp, B, T, C, NH);
|
|
}
|
|
|
|
|
|
void attention_forward2(float* out,
|
|
const float* inp,
|
|
int B, int T, int C, int NH,
|
|
const int block_size) {
|
|
// TODO there should be no mallocs inside any of these functions!
|
|
// not fixing this because we don't intend to use attention_forward2,
|
|
// it seems to be way too slow as is
|
|
|
|
// these are hardcoded to 32 for now
|
|
const int Bc = 32;
|
|
const int Br = 32;
|
|
// renaming these to be consistent with the kernel
|
|
// const int B = B;
|
|
const int nh = NH;
|
|
const int N = T;
|
|
const int d = C / NH;
|
|
// more
|
|
const int Tc = ceil((float) N / Bc);
|
|
const int Tr = ceil((float) N / Br);
|
|
const float softmax_scale = 1.0 / sqrt(d);
|
|
// create some temporary memory
|
|
float* l;
|
|
float* m;
|
|
cudaCheck(cudaMalloc(&l, B * nh * N * sizeof(float)));
|
|
cudaCheck(cudaMalloc(&m, B * nh * N * sizeof(float)));
|
|
cudaCheck(cudaMemset(l, 0, B * nh * N * sizeof(float)));
|
|
cudaCheck(cudaMemset(m, -10000.0f, B * nh * N * sizeof(float)));
|
|
|
|
// calculate SRAM size needed per block, ensure we have enough shared memory
|
|
int col_tile_size = Bc * d; // size of Kj, Vj
|
|
int row_tile_size = Br * d; // size of Qi
|
|
const int sram_size =
|
|
(2 * col_tile_size * sizeof(float)) // SRAM size for Kj, Vj
|
|
+ (row_tile_size * sizeof(float)) // SRAM size for Qi
|
|
+ (Bc * Br * sizeof(float)); // SRAM size for S
|
|
int max_sram_size;
|
|
cudaDeviceGetAttribute(&max_sram_size, cudaDevAttrMaxSharedMemoryPerBlock, 0);
|
|
if (sram_size > max_sram_size) {
|
|
printf("Max shared memory: %d, requested shared memory: %d \n", max_sram_size, sram_size);
|
|
printf("SRAM size exceeds maximum shared memory per block\n");
|
|
printf("Try decreasing col_tile_size or row_tile_size further\n");
|
|
exit(1);
|
|
}
|
|
|
|
// grid and block dims
|
|
dim3 grid_dim(B, nh); // batch_size x num_heads
|
|
dim3 block_dim(Br); // Br threads per block
|
|
|
|
// 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)
|
|
// so we have to permute the tensor using a kernel with block_size
|
|
float *q, *k, *v;
|
|
cudaCheck(cudaMalloc(&q, B * T * C * sizeof(float)));
|
|
cudaCheck(cudaMalloc(&k, B * T * C * sizeof(float)));
|
|
cudaCheck(cudaMalloc(&v, B * T * C * sizeof(float)));
|
|
int total_threads = B * N * nh * d;
|
|
int num_blocks = ceil_div(total_threads, block_size);
|
|
permute_kernel<<<num_blocks, block_size>>>(q, k, v, inp, B, N, nh, d);
|
|
|
|
// now actually call the flash attention kernel
|
|
attention_forward_kernel2<<<grid_dim, block_dim, sram_size>>>(
|
|
q, k, v,
|
|
N, d, Tc, Tr, Bc, Br, softmax_scale,
|
|
l, m, out
|
|
);
|
|
|
|
// out has shape (B, nh, N, d) but we need to unpermute it to (B, N, nh, d)
|
|
unpermute_kernel<<<num_blocks, block_size>>>(out, q, B, N, nh, d);
|
|
cudaCheck(cudaMemcpy(out, q, B * T * C * sizeof(float), cudaMemcpyDeviceToDevice));
|
|
|
|
// free memory
|
|
cudaCheck(cudaFree(l));
|
|
cudaCheck(cudaFree(m));
|
|
cudaCheck(cudaFree(q));
|
|
cudaCheck(cudaFree(k));
|
|
cudaCheck(cudaFree(v));
|
|
}
|
|
|
|
void attention_forward3(float* out, float* vaccum, float* qkvr, float* preatt, float* att,
|
|
const float* inp,
|
|
int B, int T, int C, int NH,
|
|
const int block_size) {
|
|
// inp is (B, T, 3C) QKV
|
|
// preatt, att are (B, NH, T, T)
|
|
// output is (B, T, C)
|
|
int HS = C / NH; // head size
|
|
|
|
// permute and separate inp from (B, T, 3, NH, HS) to 3X (B, NH, T, HS)
|
|
float *q, *k, *v;
|
|
q = qkvr + 0 * B * T * C;
|
|
k = qkvr + 1 * B * T * C;
|
|
v = qkvr + 2 * B * T * C;
|
|
int total_threads = B * NH * T * HS;
|
|
int num_blocks = ceil_div(total_threads, block_size);
|
|
permute_kernel<<<num_blocks, block_size>>>(q, k, v, inp, B, T, NH, HS);
|
|
|
|
// batched matrix multiply with cuBLAS
|
|
const float alpha = 1.0f;
|
|
const float beta = 0.0f;
|
|
cublasCheck(cublasSgemmStridedBatched(cublas_handle,
|
|
CUBLAS_OP_T, CUBLAS_OP_N,
|
|
T, T, HS,
|
|
&alpha,
|
|
k, HS, T * HS,
|
|
q, HS, T * HS,
|
|
&beta,
|
|
preatt, T, T * T,
|
|
B * NH));
|
|
|
|
// multiply all elements of preatt elementwise by scale
|
|
float scale = 1.0f / sqrtf(HS);
|
|
total_threads = B * NH * T * T;
|
|
num_blocks = ceil_div(total_threads, block_size);
|
|
scale_kernel<<<num_blocks, block_size>>>(preatt, scale, B, NH, T);
|
|
|
|
// softmax. preatt is (B, NH, T, T) but we view it as (B * NH * T, T) and use the softmax kernel
|
|
int softmax_block_size = 256;
|
|
int grid_size = B * NH * T;
|
|
size_t shared_mem_size = 2 * softmax_block_size / 32 * sizeof(float);
|
|
softmax_forward_kernel4<<<grid_size, softmax_block_size, shared_mem_size>>>(att, preatt, B * NH * T, T);
|
|
|
|
// new approach: first cuBLAS another batched matmul
|
|
// y = att @ v # (B, nh, T, T) @ (B, nh, T, hs) -> (B, nh, T, hs)
|
|
cublasCheck(cublasSgemmStridedBatched(cublas_handle,
|
|
CUBLAS_OP_N, CUBLAS_OP_N,
|
|
HS, T, T,
|
|
&alpha,
|
|
v, HS, T * HS,
|
|
att, T, T * T,
|
|
&beta,
|
|
vaccum, HS, T * HS,
|
|
B * NH));
|
|
|
|
// now unpermute
|
|
// y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
|
|
num_blocks = ceil_div(B * T * C, block_size);
|
|
unpermute_kernel<<<num_blocks, block_size>>>(vaccum, out, B, T, NH, HS);
|
|
}
|
|
|
|
void attention_forward4(float* out, float* vaccum, float* qkvr, float* preatt, float* att,
|
|
const float* inp,
|
|
int B, int T, int C, int NH,
|
|
const int block_size) {
|
|
// inp is (B, T, 3C) QKV
|
|
// preatt, att are (B, NH, T, T)
|
|
// output is (B, T, C)
|
|
int HS = C / NH; // head size
|
|
|
|
// permute and separate inp from (B, T, 3, NH, HS) to 3X (B, NH, T, HS)
|
|
float *q, *k, *v;
|
|
q = qkvr + 0 * B * T * C;
|
|
k = qkvr + 1 * B * T * C;
|
|
v = qkvr + 2 * B * T * C;
|
|
int total_threads = B * NH * T * HS;
|
|
int num_blocks = ceil_div(total_threads, block_size);
|
|
permute_kernel<<<num_blocks, block_size>>>(q, k, v, inp, B, T, NH, HS);
|
|
|
|
// batched matrix multiply with cuBLAS
|
|
const float alpha = 1.0f;
|
|
const float beta = 0.0f;
|
|
|
|
cublasCheck(cublasSgemmStridedBatched(cublas_handle,
|
|
CUBLAS_OP_T, CUBLAS_OP_N,
|
|
T, T, HS,
|
|
&alpha,
|
|
k, HS, T * HS,
|
|
q, HS, T * HS,
|
|
&beta,
|
|
preatt, T, T * T,
|
|
B * NH));
|
|
|
|
// multiply all elements of preatt elementwise by scale
|
|
float scale = 1.0 / sqrtf(HS);
|
|
int softmax_block_size = 256;
|
|
int grid_size = ceil_div(B * NH * T * 32, softmax_block_size);
|
|
softmax_forward_kernel5<<<grid_size, softmax_block_size>>>(att, scale, preatt, B * NH, T);
|
|
|
|
// new approach: first cuBLAS another batched matmul
|
|
// y = att @ v # (B, nh, T, T) @ (B, nh, T, hs) -> (B, nh, T, hs)
|
|
cublasCheck(cublasSgemmStridedBatched(cublas_handle,
|
|
CUBLAS_OP_N, CUBLAS_OP_N,
|
|
HS, T, T,
|
|
&alpha,
|
|
v, HS, T * HS,
|
|
att, T, T * T,
|
|
&beta,
|
|
vaccum, HS, T * HS,
|
|
B * NH));
|
|
|
|
// now unpermute
|
|
// y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
|
|
num_blocks = ceil_div(B * T * C, block_size);
|
|
unpermute_kernel<<<num_blocks, block_size>>>(vaccum, out, B, T, NH, HS);
|
|
}
|
|
|
|
|
|
__global__ void softmax_forward_kernel5_lowp(floatX* out, float inv_temperature,
|
|
const floatX* 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 floatX* 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;
|
|
|
|
// Same thing but without float4, one at a time
|
|
for (int i = warp.thread_rank(); i < pos_by_4; i += warp.size()) {
|
|
float old_maxval = maxval;
|
|
for(int k = 0; k < 4; ++k) {
|
|
maxval = fmaxf(maxval, (float)x[4*i + k]);
|
|
}
|
|
sumval *= expf(inv_temperature * (old_maxval - maxval));
|
|
for(int k = 0; k < 4; ++k) {
|
|
sumval += expf(inv_temperature * ((float)x[4*i + k] - maxval));
|
|
}
|
|
}
|
|
|
|
if(4*pos_by_4 + warp.thread_rank() <= own_pos) {
|
|
float old_maxval = maxval;
|
|
maxval = fmaxf(maxval, (float)x[4*pos_by_4 + warp.thread_rank()]);
|
|
sumval *= expf(inv_temperature * (old_maxval - maxval));
|
|
sumval += expf(inv_temperature * ((float)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 * ((float)__ldcs(x + i) - global_maxval));
|
|
__stcs(out + idx * T + i, (floatX)(ev * norm));
|
|
}
|
|
}
|
|
|
|
__global__ void permute_kernel_lowp(floatX* q, floatX* k, floatX* 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] = (floatX)inp[inp_idx];
|
|
k[idx] = (floatX)inp[inp_idx + NH * d];
|
|
v[idx] = (floatX)inp[inp_idx + 2 * (NH * d)];
|
|
}
|
|
}
|
|
|
|
__global__ void unpermute_kernel_lowp(const floatX* 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] = (float)inp[idx];
|
|
}
|
|
}
|
|
|
|
void attention_forward5(float* out, floatX* vaccum, floatX* qkvr, floatX* preatt, floatX* att,
|
|
const float* inp,
|
|
int B, int T, int C, int NH,
|
|
const int block_size, bool skip_permute=false) {
|
|
// FP16 version of kernel 4 (with permute/unpermute doing FP32<->FP16)
|
|
// That permute can be skipped on perf runs to analyse its performance impact
|
|
// inp is (B, T, 3C) QKV
|
|
// preatt, att are (B, NH, T, T)
|
|
// output is (B, T, C)
|
|
|
|
// permute and separate inp from (B, T, 3, NH, HS) to 3X (B, NH, T, HS)
|
|
int HS = C / NH; // head size
|
|
floatX *q = qkvr + 0 * B * T * C;
|
|
floatX *k = qkvr + 1 * B * T * C;
|
|
floatX* v = qkvr + 2 * B * T * C;
|
|
|
|
int total_threads = B * NH * T * HS;
|
|
int num_blocks = ceil_div(total_threads, block_size);
|
|
if (!skip_permute || first_run_validation) {
|
|
permute_kernel_lowp<<<num_blocks, block_size>>>(q, k, v, inp, B, T, NH, HS);
|
|
}
|
|
|
|
// IMPORTANT: alpha/beta are FP32 for CUBLAS_COMPUTE_32F even if FP16 inputs/outputs
|
|
// But need FP16 scale for CUBLAS_COMPUTE_16F (no errors otherwise, just garbage results *sigh*)
|
|
const float alpha = 1.0f;
|
|
const float beta = 0.0f;
|
|
const floatX alpha_lowp = (floatX)alpha;
|
|
const floatX beta_lowp = (floatX)beta;
|
|
void* alpha_ptr = CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F ? (void*)&alpha_lowp : (void*)α
|
|
void* beta_ptr = CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F ? (void*)&beta_lowp : (void*)β
|
|
|
|
// batched matrix multiply with cuBLAS
|
|
cublasCheck(cublasGemmStridedBatchedEx(cublas_handle,
|
|
CUBLAS_OP_T, CUBLAS_OP_N,
|
|
T, T, HS,
|
|
alpha_ptr,
|
|
k, CUBLAS_LOWP, HS, T * HS,
|
|
q, CUBLAS_LOWP, HS, T * HS,
|
|
beta_ptr,
|
|
preatt, CUBLAS_LOWP, T, T * T,
|
|
B * NH,
|
|
CUBLAS_LOWP_COMPUTE,
|
|
CUBLAS_GEMM_DEFAULT));
|
|
|
|
// multiply all elements of preatt elementwise by scale
|
|
float scale = 1.0f / sqrtf(HS);
|
|
int softmax_block_size = 256;
|
|
int grid_size = ceil_div(B * NH * T * 32, softmax_block_size);
|
|
softmax_forward_kernel5_lowp<<<grid_size, softmax_block_size>>>(att, scale, preatt, B * NH, T);
|
|
|
|
// new approach: first cuBLAS another batched matmul
|
|
// y = att @ v # (B, nh, T, T) @ (B, nh, T, hs) -> (B, nh, T, hs)
|
|
cublasCheck(cublasGemmStridedBatchedEx(cublas_handle,
|
|
CUBLAS_OP_N, CUBLAS_OP_N,
|
|
HS, T, T,
|
|
alpha_ptr,
|
|
v, CUBLAS_LOWP, HS, T * HS,
|
|
att, CUBLAS_LOWP, T, T * T,
|
|
beta_ptr,
|
|
vaccum, CUBLAS_LOWP, HS, T * HS,
|
|
B * NH,
|
|
CUBLAS_LOWP_COMPUTE,
|
|
CUBLAS_GEMM_DEFAULT));
|
|
|
|
// now unpermute
|
|
// y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
|
|
num_blocks = ceil_div(B * T * C, block_size);
|
|
if(!skip_permute || first_run_validation) {
|
|
unpermute_kernel_lowp<<<num_blocks, block_size>>>(vaccum, out, B, T, NH, HS);
|
|
}
|
|
}
|
|
|
|
#ifdef ENABLE_CUDNN
|
|
using graph_tensors_fwd = std::tuple<std::shared_ptr<fe::graph::Graph>,
|
|
std::shared_ptr<fe::graph::Tensor_attributes>, // Q,
|
|
std::shared_ptr<fe::graph::Tensor_attributes>, // K,
|
|
std::shared_ptr<fe::graph::Tensor_attributes>, // V,
|
|
std::shared_ptr<fe::graph::Tensor_attributes>, // Attn_scale,
|
|
std::shared_ptr<fe::graph::Tensor_attributes>, // O
|
|
std::shared_ptr<fe::graph::Tensor_attributes>>; // Stats
|
|
|
|
// Need a cache because graph->build_operation_graph() is slow but everything else seems fast
|
|
using cache_type_fwd = std::unordered_map<std::size_t, graph_tensors_fwd>;
|
|
|
|
// Loosely based on cuDNN frontend samples functions and massively simplified
|
|
template <typename... Args>
|
|
auto lookup_cache_or_build_graph_fwd(Args... args) {
|
|
static cache_type_fwd user_maintained_cache_fwd;
|
|
auto [B, H, T, HS, is_inference_only] = std::make_tuple(args...);
|
|
|
|
auto graph = std::make_shared<fe::graph::Graph>();
|
|
graph->set_io_data_type(CUDNN_16BIT)
|
|
.set_intermediate_data_type(fe::DataType_t::FLOAT)
|
|
.set_compute_data_type(fe::DataType_t::FLOAT);
|
|
|
|
// QKV is (B, T, 3, NH, HS) which cuDNN can handle directly without an external permute
|
|
auto Q = graph->tensor(fe::graph::Tensor_attributes()
|
|
.set_name("Q")
|
|
.set_dim({B, H, T, HS})
|
|
.set_stride({3 * H * HS * T, HS, 3 * H * HS, 1}));
|
|
auto K = graph->tensor(fe::graph::Tensor_attributes()
|
|
.set_name("K")
|
|
.set_dim({B, H, T, HS})
|
|
.set_stride({3 * H * HS * T, HS, 3 * H * HS, 1}));
|
|
auto V = graph->tensor(fe::graph::Tensor_attributes()
|
|
.set_name("V")
|
|
.set_dim({B, H, T, HS})
|
|
.set_stride({3 * H * HS * T, HS, 3 * H * HS, 1}));
|
|
auto attn_scale = graph->tensor(fe::graph::Tensor_attributes()
|
|
.set_name("attn_scale")
|
|
.set_dim({1, 1, 1, 1})
|
|
.set_stride({1, 1, 1, 1})
|
|
.set_is_pass_by_value(true)
|
|
.set_data_type(fe::DataType_t::FLOAT));
|
|
|
|
auto sdpa_options = fe::graph::SDPA_attributes().set_name("flash_attention");
|
|
sdpa_options.set_is_inference(is_inference_only);
|
|
sdpa_options.set_attn_scale(attn_scale);
|
|
sdpa_options.set_causal_mask(true);
|
|
|
|
// Create the graph operation and get the output tensors back
|
|
auto [O, stats] = graph->sdpa(Q, K, V, sdpa_options);
|
|
|
|
// Output is (B, T, NH, HS) BF16/FP16 and stats for backward pass is (B, NH, T) FP32
|
|
O->set_output(true).set_dim({B, H, T, HS}).set_stride({H * HS * T, HS, H * HS, 1});
|
|
|
|
assert(stats == nullptr || is_inference_only == false);
|
|
if (is_inference_only == false) {
|
|
stats->set_output(true).set_data_type(fe::DataType_t::FLOAT)
|
|
.set_dim({B, H, T, 1})
|
|
.set_stride({H * T, T, 1, 1});
|
|
}
|
|
|
|
assert(graph->validate().is_good());
|
|
auto key = graph->key();
|
|
auto it = user_maintained_cache_fwd.find(key);
|
|
if (it != user_maintained_cache_fwd.end()) {
|
|
return it->second;
|
|
}
|
|
|
|
// Build the operation graph and execution part (this is the VERY SLOW PART)
|
|
assert(graph->build_operation_graph(cudnn_handle).is_good());
|
|
auto plans = graph->create_execution_plans({fe::HeurMode_t::A});
|
|
assert(graph->check_support(cudnn_handle).is_good());
|
|
assert(graph->build_plans(cudnn_handle).is_good());
|
|
|
|
auto tuple = std::make_tuple(graph, Q, K, V, attn_scale, O, stats);
|
|
user_maintained_cache_fwd.insert({key, tuple});
|
|
return tuple;
|
|
}
|
|
|
|
// Used on first run only so we can validate against the CPU results
|
|
__global__ void fp32_to_lowp_kernel(floatX* out, const float* inp) {
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
out[idx] = (floatX)inp[idx];
|
|
}
|
|
|
|
__global__ void lowp_to_fp32_kernel(const floatX* inp, float *out) {
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
out[idx] = (float)inp[idx];
|
|
}
|
|
|
|
void attention_forward_cudnn(floatX* out, // output: (B, T, NH, HS)
|
|
float* stats, // output for backward pass: (B, NH, T)
|
|
floatX* inp, // input: (B, T, 3, NH, HS) QKV
|
|
float* in_fp32, // fp32 input
|
|
float* out_fp32, // fp32 output for validation
|
|
int B, int T, int C, int NH) {
|
|
static bool first_run_validation = true;
|
|
int HS = C / NH; // number of features per head
|
|
bool is_inference_only = (stats == nullptr);
|
|
|
|
// Convert from FP32 to FP16/BF16 on 1st run to get correct results
|
|
const int block_size = 64; // smallest full occupancy block size on modern GPUs
|
|
if (first_run_validation) {
|
|
int total_threads = B * T * C * 3;
|
|
assert(total_threads % block_size == 0);
|
|
int num_blocks = total_threads / block_size;
|
|
fp32_to_lowp_kernel<<<num_blocks, block_size>>>(inp, in_fp32);
|
|
}
|
|
|
|
// Get graph and tensors from cache (or generate it on first use)
|
|
auto [graph, Q, K, V, attn_scale, O, softmax_stats] =
|
|
lookup_cache_or_build_graph_fwd(B, NH, T, HS, is_inference_only);
|
|
|
|
// Prepare all the tensor pointers for executing the graph
|
|
void* devPtrQ = inp;
|
|
void* devPtrK = (inp + C);
|
|
void* devPtrV = (inp + 2 * C);
|
|
float attn_scale_cpu = 1.0 / sqrtf(HS);
|
|
void* devPtrO = out;
|
|
|
|
// Build variant pack
|
|
std::unordered_map<std::shared_ptr<fe::graph::Tensor_attributes>, void*> variant_pack = {
|
|
{Q, devPtrQ}, {K, devPtrK}, {V, devPtrV}, {attn_scale, &attn_scale_cpu}, {O, devPtrO}};
|
|
|
|
// Add the stats tensor unless we are only doing inference (only needed for backward pass)
|
|
if (is_inference_only == false) {
|
|
variant_pack[softmax_stats] = stats;
|
|
}
|
|
|
|
// Reallocate the workspace if the required size is greater than the current workspace
|
|
// By default, cuDNN uses up to 256MiB of workspace, so we don't want to just allocate the maximum
|
|
if (graph->get_workspace_size() > cudnn_workspace_size) {
|
|
if (cudnn_workspace_size > 0) {
|
|
cudaCheck(cudaFree(cudnn_workspace));
|
|
}
|
|
cudnn_workspace_size = graph->get_workspace_size();
|
|
cudaCheck(cudaMalloc(&cudnn_workspace, cudnn_workspace_size));
|
|
}
|
|
|
|
// Execute graph
|
|
assert(graph->execute(cudnn_handle, variant_pack, cudnn_workspace).is_good());
|
|
cudaCheck(cudaGetLastError());
|
|
|
|
// Optionally convert back from FP16/BF16 to FP32
|
|
if (first_run_validation) {
|
|
int total_threads = B * T * C;
|
|
assert(total_threads % block_size == 0);
|
|
int num_blocks = total_threads / block_size;
|
|
lowp_to_fp32_kernel<<<num_blocks, block_size>>>(out, out_fp32);
|
|
}
|
|
cudaCheck(cudaGetLastError());
|
|
first_run_validation = false;
|
|
}
|
|
|
|
#endif // ENABLE_CUDNN
|
|
|
|
// kernel version dispatch
|
|
void attention_forward(int kernel_num,
|
|
float* out, float* stats, float* vaccum,
|
|
float* qkvr, float* preatt, float* att,
|
|
float* inp,
|
|
int B, int T, int C, int NH,
|
|
const int block_size) {
|
|
switch (kernel_num) {
|
|
case 1:
|
|
attention_forward1(out, preatt, att, inp, B, T, C, NH, block_size);
|
|
break;
|
|
case 2:
|
|
attention_forward2(out, inp, B, T, C, NH, block_size);
|
|
break;
|
|
case 3:
|
|
attention_forward3(out, vaccum, qkvr, preatt, att, inp, B, T, C, NH, block_size);
|
|
break;
|
|
case 4:
|
|
attention_forward4(out, vaccum, qkvr, preatt, att, inp, B, T, C, NH, block_size);
|
|
break;
|
|
case 5:
|
|
attention_forward5(out, (floatX*)vaccum, (floatX*)qkvr,
|
|
(floatX*)preatt, (floatX*)att,
|
|
inp, B, T, C, NH, block_size, false);
|
|
break;
|
|
case 6: // skip permutes for perf passes (to analyse perf as if in/out were truly 16-bit)
|
|
attention_forward5(out, (floatX*)vaccum, (floatX*)qkvr,
|
|
(floatX*)preatt, (floatX*)att,
|
|
inp, B, T, C, NH, block_size, true);
|
|
break;
|
|
#ifdef ENABLE_CUDNN
|
|
case 10:
|
|
// note: validation only cares about out, which is out_fp32 of the function
|
|
// inp is hackily converted to FP16 into qkvr only on the first run
|
|
// similarly, vaccum is converted to FP32 into out only on the first run
|
|
attention_forward_cudnn((floatX*)vaccum, stats, (floatX*)qkvr, inp, out, B, T, C, NH);
|
|
break;
|
|
#endif
|
|
default:
|
|
printf("Invalid kernel number\n");
|
|
exit(1);
|
|
}
|
|
}
|
|
// ----------------------------------------------------------------------------
|
|
|
|
int main(int argc, char **argv) {
|
|
setup_main();
|
|
|
|
int B = 8;
|
|
int T = 1024;
|
|
int C = 768;
|
|
int NH = 12;
|
|
|
|
int deviceIdx = 0;
|
|
cudaCheck(cudaSetDevice(deviceIdx));
|
|
cudaDeviceProp deviceProp;
|
|
cudaGetDeviceProperties(&deviceProp, deviceIdx);
|
|
|
|
// setup cuBLAS (and cuDNN if needed)
|
|
cublasCreate(&cublas_handle);
|
|
int enable_tf32 = deviceProp.major >= 8 ? 1 : 0;
|
|
printf("enable_tf32: %d\n", enable_tf32);
|
|
cublasMath_t cublas_math_mode = enable_tf32 ? CUBLAS_TF32_TENSOR_OP_MATH : CUBLAS_DEFAULT_MATH;
|
|
cublasCheck(cublasSetMathMode(cublas_handle, cublas_math_mode));
|
|
|
|
#ifdef ENABLE_CUDNN
|
|
checkCudnnErr(cudnnCreate(&cudnn_handle));
|
|
#endif
|
|
|
|
// create host memory of random numbers
|
|
float* out = (float*)malloc(B * T * C * sizeof(float));
|
|
float* preatt = (float*)malloc(B * NH * T * T * sizeof(float));
|
|
float* att = (float*)malloc(B * NH * T * T * sizeof(float));
|
|
//float* inp = make_random_float(B * T * 3 * C, 10.0f);
|
|
float* inp = make_random_float(B * T * 3 * C);
|
|
|
|
// move to GPU
|
|
float* d_out;
|
|
float* d_stats; // for cuDNN
|
|
float* d_vaccum;
|
|
float* d_qkvr;
|
|
float* d_preatt;
|
|
float* d_att;
|
|
float* d_inp;
|
|
cudaCheck(cudaMalloc(&d_out, B * T * C * sizeof(float)));
|
|
cudaCheck(cudaMalloc(&d_stats, B * NH * T * sizeof(float)));
|
|
cudaCheck(cudaMalloc(&d_vaccum, B * T * C * sizeof(float)));
|
|
cudaCheck(cudaMalloc(&d_qkvr, B * T * 3 * C * sizeof(float)));
|
|
cudaCheck(cudaMalloc(&d_preatt, B * NH * T * T * sizeof(float)));
|
|
cudaCheck(cudaMalloc(&d_att, B * NH * T * T * sizeof(float)));
|
|
cudaCheck(cudaMalloc(&d_inp, B * T * 3 * C * sizeof(float)));
|
|
cudaCheck(cudaMemcpy(d_inp, inp, B * T * 3 * C * sizeof(float), cudaMemcpyHostToDevice));
|
|
|
|
// read kernel_num from command line
|
|
int kernel_num = 1;
|
|
if (argc > 1) {
|
|
kernel_num = atoi(argv[1]);
|
|
}
|
|
printf("Using kernel %d\n", kernel_num);
|
|
int block_sizes[] = {32, 64, 128, 256, 512};
|
|
|
|
// Lower accuracy requirements for FP16 (1e-4f also too much for TF32 on kernels 3 & 4)
|
|
float accuracy_threshold = (kernel_num <= 4) ? 1e-3f : 1e-2f;
|
|
|
|
// first check the correctness of the kernel
|
|
attention_forward_cpu(out, preatt, att, inp, B, T, C, NH);
|
|
for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {
|
|
int block_size = block_sizes[j];
|
|
printf("Checking block size %d.\n", block_size);
|
|
attention_forward(kernel_num, d_out, d_stats, d_vaccum, d_qkvr, d_preatt, d_att, d_inp, B, T, C, NH, block_size);
|
|
// all kernels should produce the correct output out
|
|
// todo - make accuracy threshold dynamic and depend on FP16 vs FP32?
|
|
validate_result(d_out, out, "out", B * T * C, accuracy_threshold);
|
|
// but as for preatt and att, things get a bit more complicated:
|
|
if (kernel_num != 2 && kernel_num < 5) {
|
|
// kernel 2 (knowingly) fails att/preatt because it uses a different algorithm
|
|
// that estimates the softmax online and never materializes preatt/att
|
|
validate_result(d_att, att, "att", B * NH * T * T, accuracy_threshold);
|
|
}
|
|
if (kernel_num != 2 && kernel_num < 4) {
|
|
// kernel 4 (knowingly) fails preatt because it fuses the scale normalization
|
|
// into the softmax, so preatt is off by 1.0f / sqrt(HS)
|
|
// but att and out (checked below) should match.
|
|
validate_result(d_preatt, preatt, "preatt", B * NH * T * T, accuracy_threshold);
|
|
}
|
|
}
|
|
printf("All results match. Starting benchmarks.\n\n");
|
|
first_run_validation = false;
|
|
|
|
// benchmark speed of the kernel
|
|
for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {
|
|
int block_size = block_sizes[j];
|
|
int repeat_times = 100;
|
|
|
|
float elapsed_time = benchmark_kernel(repeat_times, attention_forward,
|
|
kernel_num, d_out, d_stats, d_vaccum, d_qkvr, d_preatt, d_att,
|
|
d_inp, B, T, C, NH, block_size);
|
|
|
|
printf("block_size %4d | time %f ms\n", block_size, elapsed_time);
|
|
}
|
|
|
|
// free memory
|
|
free(out);
|
|
free(preatt);
|
|
free(att);
|
|
free(inp);
|
|
cudaCheck(cudaFree(d_out));
|
|
cudaCheck(cudaFree(d_vaccum));
|
|
cudaCheck(cudaFree(d_qkvr));
|
|
cudaCheck(cudaFree(d_preatt));
|
|
cudaCheck(cudaFree(d_att));
|
|
cudaCheck(cudaFree(d_inp));
|
|
cudaCheck(cudaFree(d_stats));
|
|
cublasDestroy(cublas_handle);
|
|
|
|
#ifdef ENABLE_CUDNN
|
|
cudnnDestroy(cudnn_handle);
|
|
if (cudnn_workspace_size > 0) {
|
|
cudaCheck(cudaFree(cudnn_workspace));
|
|
}
|
|
#endif
|
|
|
|
return 0;
|
|
} |