Files
gpt-2-inference-engine/kernels_req_5/local_softmax.cuh
T
2026-06-07 13:43:20 -05:00

37 lines
1007 B
Plaintext

#ifndef __SOFTMAX_KERNEL_CUH__
#define __SOFTMAX_KERNEL_CUH__
#include <cuda_runtime.h>
#include <math.h>
#include <float.h>
#include "../utils/cuda_utils.cuh"
__global__ void local_softmax_forward_kernel(float* out, float inv_temperature, const float* inp, int N, int T) {
// Implement this
int x = blockDim.x * blockIdx.x + threadIdx.x;
if(x < N*T){
//row offset based on WIndow size
int rowoff = x * WINDOW;
float maxval = -FLT_MAX;
for (int i = 0; i < WINDOW; i++) {
maxval = fmaxf(maxval, inp[i + rowoff]);
}
float sumval = 0.0f;
//Traverse Window Length Again
for (int i = 0; i < WINDOW; i++) {
float ev = expf(inv_temperature * (inp[i + rowoff] - maxval));
sumval += ev;
out[i + rowoff] = ev;
}
float norm = 1.0f / sumval;
for (int i = 0; i < WINDOW; i++) {
out[i + rowoff] *= norm;
}
}
}
#endif // __SOFTMAX_KERNEL_CUH__