48 lines
1.4 KiB
Plaintext
48 lines
1.4 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"
|
|
|
|
__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 b = blockDim.x * blockIdx.x + threadIdx.x;
|
|
int t = blockDim.y * blockIdx.y + threadIdx.y;
|
|
|
|
if (b >= B || t >= T) {
|
|
return;
|
|
}
|
|
|
|
float sum = 0.0f;
|
|
for (int c = 0; c < C; ++c) {
|
|
sum += inp[b * T * C + t * C + c];
|
|
}
|
|
float m = sum / C;
|
|
|
|
sum = 0.0f;
|
|
for (int c = 0; c < C; ++c) {
|
|
float diff = inp[b * T * C + t * C + c] - m;
|
|
sum += diff * diff;
|
|
}
|
|
float s = rsqrtf(sum / C + 1e-5f);
|
|
|
|
for (int c = 0; c < C; ++c) {
|
|
float n = s * (inp[b * T * C + t * C + c] - m);
|
|
out[b * T * C + t * C + 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
|
|
dim3 block(16,16);
|
|
dim3 grid((B+block.x-1)/block.x, (T+block.y-1)/block.y);
|
|
|
|
layernorm_forward_kernel<<<grid, block>>>(out,mean, rstd, inp, weight, bias,B,T, C);
|
|
}
|
|
|
|
#endif // __LAYERNORM_KERNEL_CUH__ |