Gaussian Error Gated Linear Unit

1
#include <cuda_runtime.h>
2

3
__global__ void geglu_kernel(const float* input, float* output, int halfN) {
4
    int idx = blockDim.x * blockIdx.x + threadIdx.x;
5
    if (idx >= halfN) return;
6

7
    float x1 = input[idx];
8
    float x2 = input[idx + halfN];
9

10

11
    const float inv_sqrt2 = 0.70710678118f;
12

13
    float gelu = 0.5f * x2 * (1.0f + erff(x2 * inv_sqrt2));
14

15
    output[idx] = x1 * gelu;
16
}
17

18
// input, output are device pointers
19
extern "C" void solve(const float* input, float* output, int N) {
20
    int halfN = N / 2;
21
    int threadsPerBlock = 256;
22
    int blocksPerGrid = (halfN + threadsPerBlock - 1) / threadsPerBlock;
23

24
    geglu_kernel<<<blocksPerGrid, threadsPerBlock>>>(input, output, halfN);
25
    cudaDeviceSynchronize();
26
}