From 5f99fb30d25d75661acd0900dc97f6392606f67b Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Thu, 11 Apr 2024 22:11:28 +0300 Subject: [PATCH 1/5] cooperative groups and fused scale kernel --- dev/cuda/attention_forward.cu | 116 ++++++++++++++++++++++++++++++++++ 1 file changed, 116 insertions(+) diff --git a/dev/cuda/attention_forward.cu b/dev/cuda/attention_forward.cu index 35e9d62..f313ea7 100644 --- a/dev/cuda/attention_forward.cu +++ b/dev/cuda/attention_forward.cu @@ -25,6 +25,8 @@ this turns out to be ~20X faster than (1) nice #include #include #include +#include +#include // ---------------------------------------------------------------------------- // CUDA utils @@ -302,6 +304,48 @@ __global__ void softmax_forward_kernel4(float* out, float* inp, int N, int C) { } } + +__global__ void softmax_forward_kernel5(float* out, float inv_temperature, const float* inp, int N, int T) { + // shape: (N, T, T) + + 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; + + // one row of inp, i.e. inp[idx, :] of shape (T,) + const float* x = inp + idx * T; + + // reduce to max + float maxval = -INFINITY; + for (int i = warp.thread_rank(); i <= own_pos; i += warp.size()) { + maxval = fmaxf(maxval, x[i]); + } + float offset = cg::reduce(warp, maxval, cg::greater{}); + + + // compute exp and sum + float sumval = 0.0f; + for (int i = warp.thread_rank(); i <= own_pos; i += warp.size()) { + // subtract max for numerical stability + float ev = expf(inv_temperature * (__ldcs(x + i) - offset)); + out[idx * T + i] = ev; + sumval += ev; + } + float sum = cg::reduce(warp, sumval, cg::plus{}); + float norm = 1.f / sum; + + // divide the whole row by the sum + for (int i = warp.thread_rank(); i <= own_pos; i += warp.size()) { + out[idx * T + i] *= norm; + } +} + + __global__ void attention_value_kernel1(float* out, float* att, float* inp, int B, int T, int C, int NH) { int idx = blockIdx.x * blockDim.x + threadIdx.x; @@ -679,6 +723,75 @@ void attention_forward3(float* out, float* vaccum, float* qkvr, float* preatt, f cublasDestroy(handle); } + +void attention_forward4(float* out, float* vaccum, float* qkvr, float* preatt, float* att, + 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<<>>(q, k, v, inp, B, T, NH, HS); + + // batched matrix multiply with cuBLAS + cublasHandle_t handle; + cublasStatus_t stat = cublasCreate(&handle); + const float alpha = 1.0f; + const float beta = 0.0f; + stat = cublasSgemmStridedBatched(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); + if (stat != CUBLAS_STATUS_SUCCESS) { + printf("cublasSgemm failed\n"); + exit(1); + } + + // multiply all elements of preatt elementwise by scale + float scale = 1.0 / sqrtf(HS); + int softmax_block_size = 256; + int grid_size = B * NH * T * 32 / softmax_block_size; + softmax_forward_kernel5<<>>(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) + stat = cublasSgemmStridedBatched(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); + if (stat != CUBLAS_STATUS_SUCCESS) { + printf("cublasSgemm failed\n"); + exit(1); + } + + // 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<<>>(vaccum, out, B, T, NH, HS); + + // cleanups + cublasDestroy(handle); +} + // kernel version dispatch void attention_forward(int kernel_num, float* out, float* vaccum, float* qkvr, float* preatt, float* att, @@ -695,6 +808,9 @@ void attention_forward(int kernel_num, 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; default: printf("Invalid kernel number\n"); exit(1); From baba3cfa01fa29e7146dcca3660bde41d4a6dbad Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Fri, 12 Apr 2024 02:20:20 +0300 Subject: [PATCH 2/5] recalculate instead of memory round-trip --- dev/cuda/attention_forward.cu | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/dev/cuda/attention_forward.cu b/dev/cuda/attention_forward.cu index f313ea7..50c2808 100644 --- a/dev/cuda/attention_forward.cu +++ b/dev/cuda/attention_forward.cu @@ -332,16 +332,16 @@ __global__ void softmax_forward_kernel5(float* out, float inv_temperature, const float sumval = 0.0f; for (int i = warp.thread_rank(); i <= own_pos; i += warp.size()) { // subtract max for numerical stability - float ev = expf(inv_temperature * (__ldcs(x + i) - offset)); - out[idx * T + i] = ev; - sumval += ev; + sumval += expf(inv_temperature * (x[i] - offset)); } float sum = cg::reduce(warp, sumval, cg::plus{}); float norm = 1.f / sum; // divide the whole row by the sum for (int i = warp.thread_rank(); i <= own_pos; i += warp.size()) { - out[idx * T + i] *= norm; + // recalculation is faster than doing the round-trip through memory. + float ev = expf(inv_temperature * (__ldcs(x + i) - offset)); + __stcs(out + idx * T + i, ev * norm); } } From de98a9a44a734a0e890b5ac414eaa7f18eec6a6b Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Fri, 12 Apr 2024 11:24:56 +0300 Subject: [PATCH 3/5] vector loads --- dev/cuda/attention_forward.cu | 42 +++++++++++++++++++++++++++++------ 1 file changed, 35 insertions(+), 7 deletions(-) diff --git a/dev/cuda/attention_forward.cu b/dev/cuda/attention_forward.cu index 50c2808..40bea59 100644 --- a/dev/cuda/attention_forward.cu +++ b/dev/cuda/attention_forward.cu @@ -23,6 +23,7 @@ this turns out to be ~20X faster than (1) nice #include #include +#include #include #include #include @@ -305,9 +306,17 @@ __global__ void softmax_forward_kernel4(float* out, float* inp, int N, int C) { } +__device__ float& vec_at(float4& vec, int index) { + return reinterpret_cast(&vec)[index]; +} + +__device__ float vec_at(const float4& vec, int index) { + return reinterpret_cast(&vec)[index]; +} + __global__ void softmax_forward_kernel5(float* out, float inv_temperature, const float* inp, int N, int T) { // shape: (N, T, T) - + 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); @@ -316,25 +325,44 @@ __global__ void softmax_forward_kernel5(float* out, float inv_temperature, const 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; // reduce to max float maxval = -INFINITY; - for (int i = warp.thread_rank(); i <= own_pos; i += warp.size()) { - maxval = fmaxf(maxval, x[i]); - } - float offset = cg::reduce(warp, maxval, cg::greater{}); + const float4* x_vec = reinterpret_cast(x); + for (int i = warp.thread_rank(); i < pos_by_4; i += warp.size()) { + float4 v = x_vec[i]; + for(int k = 0; k < 4; ++k) { + maxval = fmaxf(maxval, vec_at(v, k)); + } + } + + if(4*pos_by_4 + warp.thread_rank() <= own_pos) { + maxval = fmaxf(maxval, x[4*pos_by_4 + warp.thread_rank()]); + } + + float offset = cg::reduce(warp, maxval, cg::greater{}); // compute exp and sum float sumval = 0.0f; - for (int i = warp.thread_rank(); i <= own_pos; i += warp.size()) { + for (int i = warp.thread_rank(); i < pos_by_4; i += warp.size()) { // subtract max for numerical stability - sumval += expf(inv_temperature * (x[i] - offset)); + float4 v = x_vec[i]; + for(int k = 0; k < 4; ++k) { + sumval += expf(inv_temperature * (vec_at(v, k) - offset)); + } } + + if(4*pos_by_4 + warp.thread_rank() <= own_pos) { + sumval += expf(inv_temperature * (x[4*pos_by_4 + warp.thread_rank()] - offset)); + } + float sum = cg::reduce(warp, sumval, cg::plus{}); + float norm = 1.f / sum; // divide the whole row by the sum From 8e87aac34172fa604120d0617ec8e83ec2ea398a Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Fri, 12 Apr 2024 12:47:48 +0300 Subject: [PATCH 4/5] online softmax --- dev/cuda/attention_forward.cu | 35 ++++++++++++++++------------------- 1 file changed, 16 insertions(+), 19 deletions(-) diff --git a/dev/cuda/attention_forward.cu b/dev/cuda/attention_forward.cu index 40bea59..6219a23 100644 --- a/dev/cuda/attention_forward.cu +++ b/dev/cuda/attention_forward.cu @@ -24,6 +24,7 @@ this turns out to be ~20X faster than (1) nice #include #include #include +#include #include #include #include @@ -330,45 +331,41 @@ __global__ void softmax_forward_kernel5(float* out, float inv_temperature, const // one row of inp, i.e. inp[idx, :] of shape (T,) const float* x = inp + idx * T; - // reduce to max - float maxval = -INFINITY; + // 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(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)); } - } - - if(4*pos_by_4 + warp.thread_rank() <= own_pos) { - maxval = fmaxf(maxval, x[4*pos_by_4 + warp.thread_rank()]); - } - - float offset = cg::reduce(warp, maxval, cg::greater{}); - - // compute exp and sum - float sumval = 0.0f; - for (int i = warp.thread_rank(); i < pos_by_4; i += warp.size()) { - // subtract max for numerical stability - float4 v = x_vec[i]; + sumval *= expf(inv_temperature * (old_maxval - maxval)); for(int k = 0; k < 4; ++k) { - sumval += expf(inv_temperature * (vec_at(v, k) - offset)); + sumval += expf(inv_temperature * (vec_at(v, k) - maxval)); } } if(4*pos_by_4 + warp.thread_rank() <= own_pos) { - sumval += expf(inv_temperature * (x[4*pos_by_4 + warp.thread_rank()] - offset)); + 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 sum = cg::reduce(warp, sumval, cg::plus{}); + float global_maxval = cg::reduce(warp, maxval, cg::greater{}); + sumval *= expf(inv_temperature * (maxval - global_maxval)); + float sum = cg::reduce(warp, sumval, cg::plus{}); 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) - offset)); + float ev = expf(inv_temperature * (__ldcs(x + i) - global_maxval)); __stcs(out + idx * T + i, ev * norm); } } From e048825f5aeb6940d8b01adc1c3a533c61489743 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sat, 13 Apr 2024 01:20:09 +0300 Subject: [PATCH 5/5] div_ceil --- dev/cuda/attention_forward.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dev/cuda/attention_forward.cu b/dev/cuda/attention_forward.cu index 6219a23..c8ccf75 100644 --- a/dev/cuda/attention_forward.cu +++ b/dev/cuda/attention_forward.cu @@ -789,7 +789,7 @@ void attention_forward4(float* out, float* vaccum, float* qkvr, float* preatt, f // multiply all elements of preatt elementwise by scale float scale = 1.0 / sqrtf(HS); int softmax_block_size = 256; - int grid_size = B * NH * T * 32 / softmax_block_size; + int grid_size = CEIL_DIV(B * NH * T * 32, softmax_block_size); softmax_forward_kernel5<<>>(att, scale, preatt, B * NH, T); // new approach: first cuBLAS another batched matmul