From 2186cd6fe7a71c2b0aca785769702e5114527a8a Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Mon, 3 Jun 2024 22:19:47 +0000 Subject: [PATCH] move attention --- llmc/attention.cuh | 294 +++++++++++++++++++++++++++++++++++++++++++++ train_gpt2.cu | 287 +------------------------------------------ 2 files changed, 298 insertions(+), 283 deletions(-) create mode 100644 llmc/attention.cuh diff --git a/llmc/attention.cuh b/llmc/attention.cuh new file mode 100644 index 0000000..746761f --- /dev/null +++ b/llmc/attention.cuh @@ -0,0 +1,294 @@ +/* +Attention, as a fallback when we do not use the Flash Attention from cuDNN +*/ +#include +// llmc internal imports +#include "cuda_common.h" +#include "cuda_utils.cuh" +#include "cublas_common.h" + +// ---------------------------------------------------------------------------- +// CUDA kernels + +// inputs floatX, outputs FP32 (for current FP32-only activation path for this WIP) +__global__ void permute_kernel(floatX* q, floatX* k, floatX* v, + const floatX* 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; + if (idx >= B * NH * N * d) { return; } + + // Q[b][nh_][n][d_] = inp[b][n][0][nh_][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] = __ldcs(&inp[inp_idx]); + k[idx] = __ldcs(&inp[inp_idx + NH * d]); + v[idx] = __ldcs(&inp[inp_idx + 2 * (NH * d)]); +} + +__global__ void permute_kernel_backward(floatX* dinp, + const floatX* dq, const floatX* dk, const floatX* dv, + int B, int N, int NH, int d) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= B * NH * N * d) { return; } + + 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(floatX* inp, floatX *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) { return; } + + 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] = __ldcs(&inp[idx]); +} + +__global__ void unpermute_kernel_backward(floatX* dinp, const floatX *dout, int B, int N, int NH, int d) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= B * NH * N * d) { return; } + + 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] = (floatX)dout[other_idx]; +} + +__global__ void softmax_forward_kernel5(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); + int lane_id = threadIdx.x % WARP_SIZE; + int warp_id = threadIdx.x / WARP_SIZE; + int num_warps = blockDim.x / WARP_SIZE; + + // micro-optimization: we iterate backwards so that + // after the softmax backward operation completes, the cache retains the + // part of the matrix close to the upper left corner, which benefits the + // matmul operation that immediately follows. + // int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); // forward order + int idx = (gridDim.x - blockIdx.x - 1) * num_warps + warp_id; // backward order + 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. + const float flt_max = 340282346638528859811704183484516925440.0f; // to avoid including float.h + float maxval = -flt_max; + float sumval = 0.0f; + + const floatX* x_aligned = reinterpret_cast(__builtin_assume_aligned(x, 16)); + for (int i = lane_id; i < pos_by_4; i += WARP_SIZE) { + float regarray[4]; + for (int k = 0; k < 4; ++k) { + regarray[k] = (float)x_aligned[4*i + k]; + } + float old_maxval = maxval; + for(int k = 0; k < 4; ++k) { + maxval = fmaxf(maxval, regarray[k]); + } + sumval *= expf(inv_temperature * (old_maxval - maxval)); + for(int k = 0; k < 4; ++k) { + sumval += expf(inv_temperature * (regarray[k] - maxval)); + } + } + + if(4*pos_by_4 + lane_id <= own_pos) { + float old_maxval = maxval; + maxval = fmaxf(maxval, (float)x[4*pos_by_4 + lane_id]); + sumval *= expf(inv_temperature * (old_maxval - maxval)); + sumval += expf(inv_temperature * ((float)x[4*pos_by_4 + lane_id] - maxval)); + } + + float global_maxval = warpReduceMax(maxval); + sumval *= expf(inv_temperature * (maxval - global_maxval)); + + float sum = warpReduceSum(sumval); + float norm = 1.f / sum; + + // divide the whole row by the sum + for (int i = lane_id; 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 softmax_autoregressive_backward_kernel(floatX* dpreatt, const floatX* datt, const floatX* att, + int B, int T, int C, float scale) { + constexpr const int BlockSize = 256; + constexpr int T_per_block = 4; + + // go through blocks in reverse order, so the slowest block starts first + int t0 = T - 1 - T_per_block*blockIdx.x; + int idx = blockIdx.y; + + att += idx * T * T; + datt += idx * T * T; + dpreatt += idx * T * T; + + for(int to = 0; to < T_per_block; ++to) { + int t = t0 - to; + if(t < 0) return; + const floatX* att_bth = att + t * T; + const floatX* datt_bth = datt + t * T; + floatX* dpreatt_bth = dpreatt + t * T; + + float local_sum = 0; + for (int t2 = threadIdx.x; t2 <= t; t2 += BlockSize) { + local_sum += (float)att_bth[t2] * (float)datt_bth[t2]; + } + + local_sum = blockReduce(local_sum); + + for (int t3 = threadIdx.x; t3 <= t; t3 += BlockSize) { + // don't touch the cache. Some parts will still be here from the previous loop, and + // we want to exploit those. + float acc = (float)__ldcs(att_bth + t3) * ((float)__ldcs(datt_bth + t3) - local_sum); + __stcs(dpreatt_bth + t3, (floatX)(scale * acc)); + } + } +} + +// ---------------------------------------------------------------------------- +// kernel launchers + +void attention_forward(floatX* out, floatX* qkvr, floatX* att, + floatX* inp, + int B, int T, int C, int NH) { + NVTX_RANGE_FN(); + // Note: `inp` is not needed for backward pass, so we re-use it as a scratch buffer. + // Its contents will be overwritten by this function. + const int block_size = 256; + const float alpha = 1.0f, beta = 0.0f; + + // 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) + floatX *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); + + floatX* preatt = inp; + cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, + CUBLAS_OP_T, CUBLAS_OP_N, + T, T, HS, &alpha, + k, CUBLAS_LOWP, HS, T * HS, + q, CUBLAS_LOWP, HS, T * HS, + &beta, preatt, CUBLAS_LOWP, T, T * T, + B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT)); + + // multiply all elements of preatt elementwise by scale + float scale = 1.0 / sqrtf(HS); + int grid_size = CEIL_DIV(B * NH * T * WARP_SIZE, block_size); + softmax_forward_kernel5<<>>(att, scale, preatt, B * NH, T); + + // new approach: first cuBLAS another batched matmul + floatX* vaccum = inp; + // 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, + v, CUBLAS_LOWP, HS, T * HS, + att, CUBLAS_LOWP, T, T * T, + &beta, vaccum, CUBLAS_LOWP, HS, T * HS, + B * NH, cublas_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); + unpermute_kernel<<>>(vaccum, out, B, T, NH, HS); + cudaCheck(cudaGetLastError()); +} + +// the sequence of transformations in this compound op is: +// inp (B,T,3C) -> qkvr (B,T,3C) -> preatt (B,NH,T,T) -> att (B,NH,T,T) -> vaccum (B,T,C) -> out (B,T,C) +void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* datt, floatX* scratch, + const floatX* dout, + const floatX* qkvr, const floatX* att, + int B, int T, int C, int NH) { + NVTX_RANGE_FN(); + const int block_size = 256; + int HS = C / NH; // head size + const float alpha = 1.0f, beta = 0.0f; + + // unpack convenience pointers into q, k, v + const floatX *q, *k, *v; + q = qkvr + 0 * B * T * C; + k = qkvr + 1 * B * T * C; + v = qkvr + 2 * B * T * C; + floatX *dq, *dk, *dv; + dq = dqkvr + 0 * B * T * C; + dk = dqkvr + 1 * B * T * C; + dv = dqkvr + 2 * B * T * C; + + // backward through the unpermute operation + int num_blocks = CEIL_DIV(B * T * C, block_size); + unpermute_kernel_backward<<>>(scratch, dout, B, T, NH, HS); + // backward into datt + cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, &alpha, + v, CUBLAS_LOWP, HS, T * HS, scratch, CUBLAS_LOWP, HS, T * HS, &beta, + datt, CUBLAS_LOWP, T, T * T, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT)); + // backward into dv + cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &alpha, + scratch, CUBLAS_LOWP, HS, T * HS, att, CUBLAS_LOWP, T, T * T, &beta, + dv, CUBLAS_LOWP, HS, T * HS, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT)); + // backward into preatt + int hs = C / NH; // head size + float scale = 1.0f / sqrtf(hs); + softmax_autoregressive_backward_kernel<<>>(dpreatt, datt, att, B, T, C, scale); + // backward into q + cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, &alpha, + k, CUBLAS_LOWP, HS, T * HS, dpreatt, CUBLAS_LOWP, T, T * T, &beta, + dq, CUBLAS_LOWP, HS, T * HS, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT)); + // backward into k + cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &alpha, + q, CUBLAS_LOWP, HS, T * HS, dpreatt, CUBLAS_LOWP, T, T * T, &beta, + dk, CUBLAS_LOWP, HS, T * HS, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT)); + // backward into inp + num_blocks = CEIL_DIV(B * NH * T * HS, block_size); + permute_kernel_backward<<>>(dinp, dq, dk, dv, B, T, NH, HS); + cudaCheck(cudaGetLastError()); +} diff --git a/train_gpt2.cu b/train_gpt2.cu index 44b5753..2b16abf 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -46,9 +46,12 @@ GPT-2 Transformer Neural Net training loop. See README.md for usage. #include "llmc/encoder.cuh" // defines: gelu_forward, gelu_backward_inplace #include "llmc/gelu.cuh" -// defines: create_cudnn, destroy_cudnn, attention_forward_cudnn, attention_backward_cudnn #ifdef ENABLE_CUDNN +// defines: create_cudnn, destroy_cudnn, attention_forward_cudnn, attention_backward_cudnn #include "llmc/cudnn_att.h" +#else +// defines: attention_forward, attention_backward +#include "llmc/attention.cuh" #endif // ---------------------------------------------------------------------------- @@ -371,146 +374,6 @@ __global__ void fused_residual_forward_kernel5(floatX* residual, floatX* normed, } } - -// inputs floatX, outputs FP32 (for current FP32-only activation path for this WIP) -__global__ void permute_kernel(floatX* q, floatX* k, floatX* v, - const floatX* 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; - if (idx >= B * NH * N * d) { return; } - - // Q[b][nh_][n][d_] = inp[b][n][0][nh_][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] = __ldcs(&inp[inp_idx]); - k[idx] = __ldcs(&inp[inp_idx + NH * d]); - v[idx] = __ldcs(&inp[inp_idx + 2 * (NH * d)]); -} - -__global__ void permute_kernel_backward(floatX* dinp, - const floatX* dq, const floatX* dk, const floatX* dv, - int B, int N, int NH, int d) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= B * NH * N * d) { return; } - - 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(floatX* inp, floatX *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) { return; } - - 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] = __ldcs(&inp[idx]); -} - -__global__ void unpermute_kernel_backward(floatX* dinp, const floatX *dout, int B, int N, int NH, int d) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= B * NH * N * d) { return; } - - 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] = (floatX)dout[other_idx]; -} - -__global__ void softmax_forward_kernel5(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); - int lane_id = threadIdx.x % WARP_SIZE; - int warp_id = threadIdx.x / WARP_SIZE; - int num_warps = blockDim.x / WARP_SIZE; - - // micro-optimization: we iterate backwards so that - // after the softmax backward operation completes, the cache retains the - // part of the matrix close to the upper left corner, which benefits the - // matmul operation that immediately follows. - // int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); // forward order - int idx = (gridDim.x - blockIdx.x - 1) * num_warps + warp_id; // backward order - 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. - const float flt_max = 340282346638528859811704183484516925440.0f; // to avoid including float.h - float maxval = -flt_max; - float sumval = 0.0f; - - const floatX* x_aligned = reinterpret_cast(__builtin_assume_aligned(x, 16)); - for (int i = lane_id; i < pos_by_4; i += WARP_SIZE) { - float regarray[4]; - for (int k = 0; k < 4; ++k) { - regarray[k] = (float)x_aligned[4*i + k]; - } - float old_maxval = maxval; - for(int k = 0; k < 4; ++k) { - maxval = fmaxf(maxval, regarray[k]); - } - sumval *= expf(inv_temperature * (old_maxval - maxval)); - for(int k = 0; k < 4; ++k) { - sumval += expf(inv_temperature * (regarray[k] - maxval)); - } - } - - if(4*pos_by_4 + lane_id <= own_pos) { - float old_maxval = maxval; - maxval = fmaxf(maxval, (float)x[4*pos_by_4 + lane_id]); - sumval *= expf(inv_temperature * (old_maxval - maxval)); - sumval += expf(inv_temperature * ((float)x[4*pos_by_4 + lane_id] - maxval)); - } - - float global_maxval = warpReduceMax(maxval); - sumval *= expf(inv_temperature * (maxval - global_maxval)); - - float sum = warpReduceSum(sumval); - float norm = 1.f / sum; - - // divide the whole row by the sum - for (int i = lane_id; 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 residual_forward_kernel(floatX* out, const floatX* inp1, const floatX* inp2) { int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; @@ -807,42 +670,6 @@ __global__ void __launch_bounds__(512, 2) // todo - any warnings on Turing with } } -__global__ void softmax_autoregressive_backward_kernel(floatX* dpreatt, const floatX* datt, const floatX* att, - int B, int T, int C, float scale) { - constexpr const int BlockSize = 256; - constexpr int T_per_block = 4; - - // go through blocks in reverse order, so the slowest block starts first - int t0 = T - 1 - T_per_block*blockIdx.x; - int idx = blockIdx.y; - - att += idx * T * T; - datt += idx * T * T; - dpreatt += idx * T * T; - - for(int to = 0; to < T_per_block; ++to) { - int t = t0 - to; - if(t < 0) return; - const floatX* att_bth = att + t * T; - const floatX* datt_bth = datt + t * T; - floatX* dpreatt_bth = dpreatt + t * T; - - float local_sum = 0; - for (int t2 = threadIdx.x; t2 <= t; t2 += BlockSize) { - local_sum += (float)att_bth[t2] * (float)datt_bth[t2]; - } - - local_sum = blockReduce(local_sum); - - for (int t3 = threadIdx.x; t3 <= t; t3 += BlockSize) { - // don't touch the cache. Some parts will still be here from the previous loop, and - // we want to exploit those. - float acc = (float)__ldcs(att_bth + t3) * ((float)__ldcs(datt_bth + t3) - local_sum); - __stcs(dpreatt_bth + t3, (floatX)(scale * acc)); - } - } -} - // Implements linear interpolation using only two floating-point operations (as opposed to three in a naive implementation). // Reference: https://developer.nvidia.com/blog/lerp-faster-cuda __device__ float lerp(float start, float end, float weight) { @@ -1131,62 +958,6 @@ void matmul_forward_cublaslt(floatX* out, cublasCheck(cublasLtMatrixLayoutDestroy(biasLayout)); } -void attention_forward(floatX* out, floatX* qkvr, floatX* att, - floatX* inp, - int B, int T, int C, int NH) { - NVTX_RANGE_FN(); - // Note: `inp` is not needed for backward pass, so we re-use it as a scratch buffer. - // Its contents will be overwritten by this function. - const int block_size = 256; - const float alpha = 1.0f, beta = 0.0f; - - // 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) - floatX *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); - - - floatX* preatt = inp; - cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, - CUBLAS_OP_T, CUBLAS_OP_N, - T, T, HS, &alpha, - k, CUBLAS_LOWP, HS, T * HS, - q, CUBLAS_LOWP, HS, T * HS, - &beta, preatt, CUBLAS_LOWP, T, T * T, - B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT)); - - // multiply all elements of preatt elementwise by scale - float scale = 1.0 / sqrtf(HS); - int grid_size = CEIL_DIV(B * NH * T * WARP_SIZE, block_size); - softmax_forward_kernel5<<>>(att, scale, preatt, B * NH, T); - - // new approach: first cuBLAS another batched matmul - floatX* vaccum = inp; - // 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, - v, CUBLAS_LOWP, HS, T * HS, - att, CUBLAS_LOWP, T, T * T, - &beta, vaccum, CUBLAS_LOWP, HS, T * HS, - B * NH, cublas_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); - unpermute_kernel<<>>(vaccum, out, B, T, NH, HS); - cudaCheck(cudaGetLastError()); -} - void residual_forward(floatX* out, const floatX* inp1, const floatX* inp2, int N) { NVTX_RANGE_FN(); const int block_size = 256; @@ -1279,56 +1050,6 @@ void layernorm_backward(floatX* dinp, floatX* dweight, floatX* dbias, float* scr cudaCheck(cudaGetLastError()); } -// the sequence of transformations in this compound op is: -// inp (B,T,3C) -> qkvr (B,T,3C) -> preatt (B,NH,T,T) -> att (B,NH,T,T) -> vaccum (B,T,C) -> out (B,T,C) -void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* datt, floatX* scratch, - const floatX* dout, - const floatX* qkvr, const floatX* att, - int B, int T, int C, int NH) { - NVTX_RANGE_FN(); - const int block_size = 256; - int HS = C / NH; // head size - const float alpha = 1.0f, beta = 0.0f; - - // unpack convenience pointers into q, k, v - const floatX *q, *k, *v; - q = qkvr + 0 * B * T * C; - k = qkvr + 1 * B * T * C; - v = qkvr + 2 * B * T * C; - floatX *dq, *dk, *dv; - dq = dqkvr + 0 * B * T * C; - dk = dqkvr + 1 * B * T * C; - dv = dqkvr + 2 * B * T * C; - - // backward through the unpermute operation - int num_blocks = CEIL_DIV(B * T * C, block_size); - unpermute_kernel_backward<<>>(scratch, dout, B, T, NH, HS); - // backward into datt - cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, &alpha, - v, CUBLAS_LOWP, HS, T * HS, scratch, CUBLAS_LOWP, HS, T * HS, &beta, - datt, CUBLAS_LOWP, T, T * T, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT)); - // backward into dv - cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &alpha, - scratch, CUBLAS_LOWP, HS, T * HS, att, CUBLAS_LOWP, T, T * T, &beta, - dv, CUBLAS_LOWP, HS, T * HS, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT)); - // backward into preatt - int hs = C / NH; // head size - float scale = 1.0f / sqrtf(hs); - softmax_autoregressive_backward_kernel<<>>(dpreatt, datt, att, B, T, C, scale); - // backward into q - cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, &alpha, - k, CUBLAS_LOWP, HS, T * HS, dpreatt, CUBLAS_LOWP, T, T * T, &beta, - dq, CUBLAS_LOWP, HS, T * HS, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT)); - // backward into k - cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &alpha, - q, CUBLAS_LOWP, HS, T * HS, dpreatt, CUBLAS_LOWP, T, T * T, &beta, - dk, CUBLAS_LOWP, HS, T * HS, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT)); - // backward into inp - num_blocks = CEIL_DIV(B * NH * T * HS, block_size); - permute_kernel_backward<<>>(dinp, dq, dk, dv, B, T, NH, HS); - cudaCheck(cudaGetLastError()); -} - // replaces logits with logit gradients template void fused_classifier(Type* logits, Type* losses,