aparler/c-llm
1
0
Fork 0
mirror of https://github.com/karpathy/llm.c.git synced 2026-07-28 20:35:09 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

incorporate faster encoder_forward kernel to fp32 CUDA version

Browse source
This commit is contained in:
Andrej Karpathy 2024-04-28 20:10:33 +00:00
parent 4b6a532b9b
commit 327eef3f22
2 changed files with 26 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

13
dev/cuda/encoder_forward.cu
View file

@ -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 <stdio.h>
@ -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]);
}
}

36
train_gpt2_fp32.cu
View file

@ -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<<<grid_size, block_size>>>(out, inp, wte, wpe, B, T, C);
const int grid_size = CEIL_DIV(N / 4, block_size);
encoder_forward_kernel3<<<grid_size, block_size>>>((float4*) out, inp, (float4*) wte, (float4*) wpe, B, T, C);
cudaCheck(cudaGetLastError());
}