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

Merge branch 'lancerts-encoder_forward-float4'

Browse source
This commit is contained in:
Andrej Karpathy 2024-04-28 20:10:48 +00:00
commit 4b6f68a9a9
2 changed files with 59 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

40
dev/cuda/encoder_forward.cu
View file

@ -9,12 +9,16 @@ 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>
#include <stdlib.h>
#include <cuda_runtime.h>
#include "common.h"
#include <cassert>
// ----------------------------------------------------------------------------
// CPU code reference
@ -81,6 +85,28 @@ __global__ void encoder_forward_kernel2(float* out,
}
}
__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] = add_float4(wte[ix * C4 + c4], wpe[t * C4 + c4]);
}
}
// ----------------------------------------------------------------------------
// kernel launcher
@ -104,6 +130,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 / 4, block_size);
encoder_forward_kernel3<<<grid_size, block_size>>>((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 +154,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);

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());
}