From cfccd821e6c2e0c9496c3560cbbe1eb21288ada3 Mon Sep 17 00:00:00 2001 From: lancer Date: Sat, 27 Apr 2024 19:58:18 -0700 Subject: [PATCH 1/4] Include the float4 kernel --- dev/cuda/encoder_forward.cu | 58 +++++++++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/dev/cuda/encoder_forward.cu b/dev/cuda/encoder_forward.cu index c06a523..67d0c8f 100644 --- a/dev/cuda/encoder_forward.cu +++ b/dev/cuda/encoder_forward.cu @@ -15,6 +15,7 @@ version 2 is more optimized, parallelizes over all of B,T,C #include #include #include "common.h" +#include // ---------------------------------------------------------------------------- // CPU code reference @@ -81,6 +82,49 @@ __global__ void encoder_forward_kernel2(float* out, } } +// use float4 for memory coalescing +// optimized implementation: parallelize over all of B,T,C + +__device__ inline float4 operator+(const float4& a, const float4& b) { + return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); +} + +__device__ inline float& vec_at(float4& vec, int index) { + return reinterpret_cast(&vec)[index]; +} + +__device__ inline float vec_at(const float4& vec, int index) { + return reinterpret_cast(&vec)[index]; +} + +__global__ void encoder_forward_kernel3(float4* out, + const int* inp, const float4* wte, const float4* wpe, + int B, int T, int C) { + int C4 = C / 4; + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int N = B * T * C4; + + if (idx < N) { + int bt = idx / C4; + int b = bt / T; + int t = bt % T; + int c4 = idx % C4; + + int ix = inp[b * T + t]; + + // float4* out_btc4 = out + b * T * C4 + t * C4 + c4; + // const float4* wte_ix4 = wte + ix * C4 + c4; + // const float4* wpe_tc4 = wpe + t * C4 + c4; + // for (int i = 0; i < 4; i++) { + // vec_at(*out_btc4, i) = vec_at(*wte_ix4, i) + vec_at(*wpe_tc4, i); + // } + //*out_btc4 = *wte_ix4 + *wpe_tc4; + out[b * T * C4 + t * C4 + c4] = wte[ix * C4 + c4] + wpe[t * C4 + c4]; + } +} + + + // ---------------------------------------------------------------------------- // kernel launcher @@ -104,6 +148,17 @@ void encoder_forward2(float* out, cudaCheck(cudaGetLastError()); } +void encoder_forward3(float* out, + const int* inp, const float* wte, const float* wpe, + int B, int T, int C, + const int block_size) { + assert(C % 4 == 0); + const int N = B * T * C; + const int grid_size = ceil_div(N, block_size); + encoder_forward_kernel3<<>>((float4*) out, inp, (float4*) wte, (float4*) wpe, B, T, C); + cudaCheck(cudaGetLastError()); +} + // kernel version dispatch void encoder_forward(int kernel_num, float* out, @@ -117,6 +172,9 @@ void encoder_forward(int kernel_num, case 2: encoder_forward2(out, inp, wte, wpe, B, T, C, block_size); break; + case 3: + encoder_forward3(out, inp, wte, wpe, B, T, C, block_size); + break; default: printf("Invalid kernel number\n"); exit(1); From 61c5c05ea8115119d55f34eb9eb99012c26f6927 Mon Sep 17 00:00:00 2001 From: lancer Date: Sat, 27 Apr 2024 20:23:39 -0700 Subject: [PATCH 2/4] amend the float4 kernel --- dev/cuda/encoder_forward.cu | 21 ++------------------- 1 file changed, 2 insertions(+), 19 deletions(-) diff --git a/dev/cuda/encoder_forward.cu b/dev/cuda/encoder_forward.cu index 67d0c8f..1466132 100644 --- a/dev/cuda/encoder_forward.cu +++ b/dev/cuda/encoder_forward.cu @@ -89,14 +89,6 @@ __device__ inline float4 operator+(const float4& a, const float4& b) { return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); } -__device__ inline float& vec_at(float4& vec, int index) { - return reinterpret_cast(&vec)[index]; -} - -__device__ inline float vec_at(const float4& vec, int index) { - return reinterpret_cast(&vec)[index]; -} - __global__ void encoder_forward_kernel3(float4* out, const int* inp, const float4* wte, const float4* wpe, int B, int T, int C) { @@ -111,20 +103,11 @@ __global__ void encoder_forward_kernel3(float4* out, int c4 = idx % C4; int ix = inp[b * T + t]; - - // float4* out_btc4 = out + b * T * C4 + t * C4 + c4; - // const float4* wte_ix4 = wte + ix * C4 + c4; - // const float4* wpe_tc4 = wpe + t * C4 + c4; - // for (int i = 0; i < 4; i++) { - // vec_at(*out_btc4, i) = vec_at(*wte_ix4, i) + vec_at(*wpe_tc4, i); - // } - //*out_btc4 = *wte_ix4 + *wpe_tc4; + out[b * T * C4 + t * C4 + c4] = wte[ix * C4 + c4] + wpe[t * C4 + c4]; } } - - // ---------------------------------------------------------------------------- // kernel launcher @@ -154,7 +137,7 @@ void encoder_forward3(float* out, const int block_size) { assert(C % 4 == 0); const int N = B * T * C; - const int grid_size = ceil_div(N, block_size); + const int grid_size = ceil_div(N / 4, block_size); encoder_forward_kernel3<<>>((float4*) out, inp, (float4*) wte, (float4*) wpe, B, T, C); cudaCheck(cudaGetLastError()); } From 1c7d23ab3c015674728b38d114d92418c5217e80 Mon Sep 17 00:00:00 2001 From: lancer Date: Sat, 27 Apr 2024 20:26:19 -0700 Subject: [PATCH 3/4] amend the float4 kernel --- dev/cuda/encoder_forward.cu | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/dev/cuda/encoder_forward.cu b/dev/cuda/encoder_forward.cu index 1466132..5ab32bc 100644 --- a/dev/cuda/encoder_forward.cu +++ b/dev/cuda/encoder_forward.cu @@ -83,8 +83,6 @@ __global__ void encoder_forward_kernel2(float* out, } // use float4 for memory coalescing -// optimized implementation: parallelize over all of B,T,C - __device__ inline float4 operator+(const float4& a, const float4& b) { return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); } @@ -103,7 +101,7 @@ __global__ void encoder_forward_kernel3(float4* out, int c4 = idx % C4; int ix = inp[b * T + t]; - + out[b * T * C4 + t * C4 + c4] = wte[ix * C4 + c4] + wpe[t * C4 + c4]; } } From 327eef3f22d9bb0954a0091e49d46cb2b5a1fa95 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sun, 28 Apr 2024 20:10:33 +0000 Subject: [PATCH 4/4] incorporate faster encoder_forward kernel to fp32 CUDA version --- dev/cuda/encoder_forward.cu | 13 +++++++------ train_gpt2_fp32.cu | 36 +++++++++++++++++++----------------- 2 files changed, 26 insertions(+), 23 deletions(-) diff --git a/dev/cuda/encoder_forward.cu b/dev/cuda/encoder_forward.cu index 5ab32bc..f56a9c8 100644 --- a/dev/cuda/encoder_forward.cu +++ b/dev/cuda/encoder_forward.cu @@ -9,6 +9,9 @@ version 1 is naive port from CPU code to kernel: parallelizes over B,T, loops ov version 2 is more optimized, parallelizes over all of B,T,C ./encoder_forward 2 + +version 3 is like version 2 but uses float4 reads/writes +./encoder_forward 3 */ #include @@ -82,27 +85,25 @@ __global__ void encoder_forward_kernel2(float* out, } } -// use float4 for memory coalescing -__device__ inline float4 operator+(const float4& a, const float4& b) { +__device__ inline float4 add_float4(const float4& a, const float4& b) { return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); } +// use of float4 leads to using 128-bit LDG / STG instructions in SASS, +// very helpful in memory-bound kernels like encoder_forward __global__ void encoder_forward_kernel3(float4* out, const int* inp, const float4* wte, const float4* wpe, int B, int T, int C) { int C4 = C / 4; int idx = blockIdx.x * blockDim.x + threadIdx.x; int N = B * T * C4; - if (idx < N) { int bt = idx / C4; int b = bt / T; int t = bt % T; int c4 = idx % C4; - int ix = inp[b * T + t]; - - out[b * T * C4 + t * C4 + c4] = wte[ix * C4 + c4] + wpe[t * C4 + c4]; + out[b * T * C4 + t * C4 + c4] = add_float4(wte[ix * C4 + c4], wpe[t * C4 + c4]); } } diff --git a/train_gpt2_fp32.cu b/train_gpt2_fp32.cu index 19bc574..8b90e3c 100644 --- a/train_gpt2_fp32.cu +++ b/train_gpt2_fp32.cu @@ -153,24 +153,25 @@ __device__ float warpReduceSum(float val) { return val; } -__global__ void encoder_forward_kernel2(float* out, - int* inp, float* wte, float* wpe, - int B, int T, int C) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int N = B * T * C; +__device__ inline float4 add_float4(const float4& a, const float4& b) { + return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); +} +// use of float4 leads to using 128-bit LDG / STG instructions in SASS, +// very helpful in memory-bound kernels like encoder_forward +__global__ void encoder_forward_kernel3(float4* out, + const int* inp, const float4* wte, const float4* wpe, + int B, int T, int C) { + int C4 = C / 4; + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int N = B * T * C4; if (idx < N) { - int bt = idx / C; + int bt = idx / C4; int b = bt / T; int t = bt % T; - int c = idx % C; - + int c4 = idx % C4; int ix = inp[b * T + t]; - - float* out_btc = out + b * T * C + t * C + c; - float* wte_ix = wte + ix * C + c; - float* wpe_tc = wpe + t * C + c; - *out_btc = *wte_ix + *wpe_tc; + out[b * T * C4 + t * C4 + c4] = add_float4(wte[ix * C4 + c4], wpe[t * C4 + c4]); } } @@ -812,12 +813,13 @@ __global__ void fused_classifier_kernel3(float* logits, float* losses, float* pr // kernel launchers void encoder_forward(float* out, - int* inp, float* wte, float* wpe, + const int* inp, const float* wte, const float* wpe, int B, int T, int C) { + assert(C % 4 == 0); + const int block_size = 512; const int N = B * T * C; - const int block_size = 256; - const int grid_size = CEIL_DIV(N, block_size); - encoder_forward_kernel2<<>>(out, inp, wte, wpe, B, T, C); + const int grid_size = CEIL_DIV(N / 4, block_size); + encoder_forward_kernel3<<>>((float4*) out, inp, (float4*) wte, (float4*) wpe, B, T, C); cudaCheck(cudaGetLastError()); }