mirror of
https://github.com/karpathy/llm.c.git
synced 2026-07-27 20:25:09 -04:00
Added packing for gelu forwards kernel
This commit is contained in:
Christopher Dryden
parent
2490f7896d
commit
568615fef1
2 changed files with 67 additions and 26 deletions
|
|
@ -21,6 +21,22 @@ version 2 uses the Packed128 data structure
|
|||
#include <cuda_runtime.h>
|
||||
#include "common.h"
|
||||
|
||||
// turn on bf16 as default, done up here for now
|
||||
//#define ENABLE_BF16
|
||||
|
||||
#if defined(ENABLE_BF16)
|
||||
typedef __nv_bfloat16 floatX;
|
||||
typedef __nv_bfloat16 floatN;
|
||||
#elif defined(ENABLE_FP16)
|
||||
typedef half floatX;
|
||||
typedef half floatN;
|
||||
#else
|
||||
typedef float floatX;
|
||||
typedef float floatN;
|
||||
#endif
|
||||
|
||||
typedef Packed128<floatX> x128;
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// CPU code reference
|
||||
|
||||
|
|
@ -48,19 +64,19 @@ __global__ void gelu_kernel(float* out, const float* inp, int N) {
|
|||
}
|
||||
|
||||
// elementwise ops are nice and ez
|
||||
__global__ void gelu_kernel2(float* out, const float* inp, int N) {
|
||||
int i = (blockIdx.x * blockDim.x + threadIdx.x) * f128::size;
|
||||
__global__ void gelu_kernel2(floatX* out, const floatX* inp, int N) {
|
||||
int i = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size;
|
||||
if (i < N) {
|
||||
f128 packet_out;
|
||||
f128 packet_in = load128cs(inp + i); // load and do not keep in cache
|
||||
for(int k = 0; k < packet_in.size; ++k) {
|
||||
float xi = packet_in[k];
|
||||
x128 packed_out;
|
||||
x128 packed_inp = load128cs(inp + i); // load and do not keep in cache
|
||||
for(int k = 0; k < packed_inp.size; ++k) {
|
||||
float xi = (float)packed_inp[k];
|
||||
float cube = 0.044715f * xi * xi * xi;
|
||||
packet_out[k] = 0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube)));
|
||||
packed_out[k] = (float)(0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube))));
|
||||
}
|
||||
// store instead of storecs (without cache streaming) in case it is useful for the
|
||||
// data to be in the cache for the next operation after this GeLU
|
||||
store128(out + i, packet_out);
|
||||
store128(out + i, packed_out);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -73,25 +89,29 @@ void gelu_forward1(float* out, const float* inp, int N, const int block_size) {
|
|||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
void gelu_forward2(float* out, const float* inp, int N, const int block_size) {
|
||||
const int grid_size = ceil_div(N, block_size) / 4;
|
||||
void gelu_forward2(floatX* out, const floatX* inp, int N, const int block_size) {
|
||||
const int grid_size = ceil_div(N, block_size)/x128::size;
|
||||
gelu_kernel2<<<grid_size, block_size>>>(out, inp, N);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
// kernel version dispatch
|
||||
void gelu_forward(int kernel_num,
|
||||
float* out,
|
||||
const float* inp,
|
||||
floatX* out,
|
||||
const floatX* inp,
|
||||
int B, int T, int C,
|
||||
int block_size) {
|
||||
switch (kernel_num) {
|
||||
#if !defined(ENABLE_BF16) && !defined(ENABLE_FP16)
|
||||
case 1:
|
||||
gelu_forward1(out, inp, B * T * C, block_size);
|
||||
break;
|
||||
#endif
|
||||
#if defined(ENABLE_BF16)
|
||||
case 2:
|
||||
gelu_forward2(out, inp, B * T * C, block_size);
|
||||
break;
|
||||
#endif
|
||||
default:
|
||||
printf("Invalid kernel number\n");
|
||||
exit(1);
|
||||
|
|
@ -114,13 +134,6 @@ int main(int argc, char **argv) {
|
|||
float* out = (float*)malloc(B * T * C * sizeof(float));
|
||||
float* inp = make_random_float(B * T * C);
|
||||
|
||||
// move to GPU
|
||||
float* d_out;
|
||||
float* d_inp;
|
||||
cudaCheck(cudaMalloc(&d_out, B * T * C * sizeof(float)));
|
||||
cudaCheck(cudaMalloc(&d_inp, B * T * C * sizeof(float)));
|
||||
cudaCheck(cudaMemcpy(d_inp, inp, B * T * C * sizeof(float), cudaMemcpyHostToDevice));
|
||||
|
||||
// read kernel_num from command line
|
||||
int kernel_num = 1;
|
||||
if (argc > 1) {
|
||||
|
|
@ -131,6 +144,19 @@ int main(int argc, char **argv) {
|
|||
// first check the correctness of the kernel
|
||||
gelu_forward_cpu(out, inp, B * T * C);
|
||||
|
||||
// move to GPU
|
||||
floatX* d_out;
|
||||
floatX* d_inp;
|
||||
cudaCheck(cudaMalloc(&d_out, B * T * C * sizeof(floatX)));
|
||||
cudaCheck(cudaMalloc(&d_inp, B * T * C * sizeof(floatX)));
|
||||
|
||||
floatX* inpX = (floatX*)malloc(B * T * C * sizeof(floatX));
|
||||
|
||||
for (int i = 0; i < B * T * C; i++) {
|
||||
inpX[i] = (floatX)inp[i];
|
||||
}
|
||||
|
||||
cudaCheck(cudaMemcpy(d_inp, inpX, B * T * C * sizeof(floatX), cudaMemcpyHostToDevice));
|
||||
|
||||
// time the kernel at different block sizes
|
||||
int block_sizes[] = {32, 64, 128, 256, 512, 1024};
|
||||
|
|
@ -138,7 +164,12 @@ int main(int argc, char **argv) {
|
|||
int block_size = block_sizes[j];
|
||||
printf("Checking block size %d.\n", block_size);
|
||||
gelu_forward(kernel_num, d_out, d_inp, B, T, C, block_size);
|
||||
#if !defined(ENABLE_BF16) && !defined(ENABLE_FP16)
|
||||
validate_result(d_out, out, "out", B * T * C, 1e-5f);
|
||||
#endif
|
||||
#if defined(ENABLE_BF16)
|
||||
#endif
|
||||
validate_result(d_out, out, "out", B * T * C, 1e-2f);
|
||||
}
|
||||
|
||||
printf("All results match. Starting benchmarks.\n\n");
|
||||
|
|
@ -164,8 +195,9 @@ int main(int argc, char **argv) {
|
|||
// free memory
|
||||
free(out);
|
||||
free(inp);
|
||||
free(inpX);
|
||||
|
||||
cudaCheck(cudaFree(d_out));
|
||||
cudaCheck(cudaFree(d_inp));
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
|
@ -227,6 +227,7 @@ struct alignas(16) Packed128 {
|
|||
|
||||
// short-form typedef
|
||||
typedef Packed128<float> f128;
|
||||
typedef Packed128<floatX> x128;
|
||||
|
||||
// load a Packed128 from an aligned memory address
|
||||
template<class ElementType>
|
||||
|
|
@ -566,7 +567,8 @@ __global__ void permute_kernel_backward(floatX* dinp,
|
|||
|
||||
__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;
|
||||
|
||||
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);
|
||||
|
|
@ -668,11 +670,18 @@ __global__ void residual_forward_kernel(floatX* out, floatX* inp1, floatX* inp2,
|
|||
|
||||
#define GELU_SCALING_FACTOR sqrtf(2.0f / M_PI)
|
||||
__global__ void gelu_forward_kernel(floatX* out, const floatX* inp, int N) {
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int i = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size;
|
||||
if (i < N) {
|
||||
float xi = (float)inp[i];
|
||||
float cube = 0.044715f * xi * xi * xi;
|
||||
out[i] = (floatX)(0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube))));
|
||||
x128 packed_out;
|
||||
x128 packed_inp = load128cs(inp + i); // load and do not keep in cache
|
||||
for(int k = 0; k < packed_inp.size; ++k) {
|
||||
float xi =(float)packed_inp[k];
|
||||
float cube = 0.044715f * xi * xi * xi;
|
||||
packed_out[k] = (floatX)(0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube))));
|
||||
}
|
||||
// store instead of storecs (without cache streaming) in case it is useful for the
|
||||
// data to be in the cache for the next operation after this GeLU
|
||||
store128(out + i, packed_out);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -1192,7 +1201,7 @@ void residual_forward(floatX* out, floatX* inp1, floatX* inp2, int N) {
|
|||
|
||||
void gelu_forward(floatX* out, const floatX* inp, int N) {
|
||||
const int block_size = 128;
|
||||
const int grid_size = CEIL_DIV(N, block_size);
|
||||
const int grid_size = CEIL_DIV(N/x128::size, block_size);
|
||||
gelu_forward_kernel<<<grid_size, block_size>>>(out, inp, N);
|
||||
cudaCheck(cudaGetLastError());
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue