102 lines
2.7 KiB
Plaintext
102 lines
2.7 KiB
Plaintext
#ifndef __LAYERNORM_KERNEL_CUH__
|
|
#define __LAYERNORM_KERNEL_CUH__
|
|
|
|
#include <cuda_runtime.h>
|
|
#include <math.h>
|
|
#include <float.h>
|
|
#include "../utils/cuda_utils.cuh"
|
|
|
|
#define BLOCK_DIM 256
|
|
#define WARP_SIZE 32
|
|
|
|
//lecture 15 slide ~32
|
|
__device__ __inline__ float warp_reduce(float val) {
|
|
float partialSum = val;
|
|
for(unsigned int stride = WARP_SIZE/2; stride > 0; stride /= 2) {
|
|
partialSum += __shfl_down_sync(0xffffffff, partialSum, stride);
|
|
}
|
|
return partialSum;
|
|
}
|
|
|
|
__global__ void layernorm_forward_kernel(float* out, float* mean, float* rstd, const float* inp, const float* weight,
|
|
const float* bias, int B, int T, int C) {
|
|
// Implement this
|
|
int bt = blockIdx.x;
|
|
if(bt >= B*T){
|
|
return;
|
|
}
|
|
|
|
const float* x = inp + bt * C;
|
|
float* o = out + bt * C;
|
|
|
|
// computer mean
|
|
float partial_sum = 0.0f;
|
|
for(int c = threadIdx.x ; c < C ; c += BLOCK_DIM){
|
|
partial_sum += x[c];
|
|
}
|
|
|
|
partial_sum = warp_reduce(partial_sum);
|
|
|
|
__shared__ float partial_sum_s[BLOCK_DIM / WARP_SIZE];
|
|
if(threadIdx.x % WARP_SIZE == 0){
|
|
partial_sum_s[threadIdx.x / WARP_SIZE] = partial_sum;
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
if(threadIdx.x < WARP_SIZE){
|
|
partial_sum = (threadIdx.x < BLOCK_DIM / WARP_SIZE) ? partial_sum_s[threadIdx.x] : 0.0f;
|
|
partial_sum = warp_reduce(partial_sum);
|
|
}
|
|
|
|
//thread 0 will have the final sum
|
|
__shared__ float s_mean;
|
|
if (threadIdx.x == 0) {
|
|
s_mean = partial_sum / C;
|
|
}
|
|
|
|
__syncthreads();
|
|
float m = s_mean;
|
|
|
|
//compute variance
|
|
partial_sum = 0.0f;
|
|
for (int c = threadIdx.x; c < C; c += BLOCK_DIM) {
|
|
float diff = x[c] - m;
|
|
partial_sum += diff * diff;
|
|
}
|
|
|
|
partial_sum = warp_reduce(partial_sum);
|
|
|
|
if (threadIdx.x % WARP_SIZE == 0) {
|
|
partial_sum_s[threadIdx.x / WARP_SIZE] = partial_sum;
|
|
}
|
|
__syncthreads();
|
|
|
|
if (threadIdx.x < WARP_SIZE) {
|
|
partial_sum = (threadIdx.x < BLOCK_DIM / WARP_SIZE) ? partial_sum_s[threadIdx.x] : 0.0f;
|
|
partial_sum = warp_reduce(partial_sum);
|
|
}
|
|
|
|
__shared__ float s_rstd;
|
|
if (threadIdx.x == 0) {
|
|
s_rstd = rsqrtf(partial_sum / C + 1e-5f);
|
|
}
|
|
__syncthreads();
|
|
float s = s_rstd;
|
|
|
|
//normalize
|
|
for (int c = threadIdx.x; c < C; c += BLOCK_DIM) {
|
|
float n = s * (x[c] - m);
|
|
o[c] = n * weight[c] + bias[c];
|
|
}
|
|
}
|
|
|
|
// Launch kernel here
|
|
void layernorm_forward(float* out, float* mean, float* rstd, float* inp, float* weight, float* bias,
|
|
int B, int T, int C) {
|
|
// Implement this
|
|
int num_blocks = B*T;
|
|
layernorm_forward_kernel<<<num_blocks, BLOCK_DIM>>>(out,mean, rstd, inp, weight, bias,B,T, C);
|
|
}
|
|
|
|
#endif // __LAYERNORM_KERNEL_CUH__ |