geglu_kernels.c File Reference

GeGLU kernels split from gelu_kernels.c. More...

#include <math.h>
#include <stddef.h>
#include <stdint.h>
#include "ckernel_engine.h"
#include "bf16_utils.h"

Go to the source code of this file.

Functions
static float	ck_gelu_tanh_parity_f32 (float x)
void	geglu_backward_bf16_mixed (const uint16_t *x, const uint16_t *d_out, float *d_x, int tokens, int dim)
void	geglu_backward_fp32 (const float *x, const float *d_out, float *d_x, int tokens, int dim)
void	geglu_forward_bf16 (const uint16_t *x, uint16_t *out, int tokens, int dim, float *scratch)
void	geglu_forward_exact (const float *x, float *out, int tokens, int dim)
void	geglu_forward_fp32 (const float *x, float *out, int tokens, int dim)
void	geglu_forward_ggml_native (const float *x, float *out, int tokens, int dim)
void	gelu_fast_inplace (float *data, size_t n)
void	gelu_ggml_native_inplace (float *data, size_t n)

Detailed Description

GeGLU kernels split from gelu_kernels.c.

Definition in file geglu_kernels.c.

Function Documentation

ck_gelu_tanh_parity_f32()

static float ck_gelu_tanh_parity_f32 ( float  x )

inlinestatic

Definition at line 33 of file geglu_kernels.c.

{
    const float sqrt_2_over_pi = 0.7978845608f;
    const float coeff = 0.044715f;
    const float x3 = x * x * x;
    const float inner = sqrt_2_over_pi * (x + coeff * x3);
    return 0.5f * x * (1.0f + tanhf(inner));
}

Referenced by geglu_forward_exact().

geglu_backward_bf16_mixed()

void geglu_backward_bf16_mixed	(	const uint16_t *	x,
		const uint16_t *	d_out,
		float *	d_x,
		int	tokens,
		int	dim
	)

Definition at line 139 of file geglu_kernels.c.

{
    if (!x || !d_out || !d_x || tokens <= 0 || dim <= 0) {
        return;
    }
 
    const float sqrt_2_over_pi = 0.7978845608f;
    const float coeff = 0.044715f;
    const int inner_dim = dim * 2;
 
    for (int t = 0; t < tokens; ++t) {
        const uint16_t *x_ptr = x + (size_t)t * (size_t)inner_dim;
        const uint16_t *d_out_ptr = d_out + (size_t)t * (size_t)dim;
        float *d_x_ptr = d_x + (size_t)t * (size_t)inner_dim;
 
        for (int d = 0; d < dim; ++d) {
            const float a = bf16_to_float(x_ptr[d]);
            const float b = bf16_to_float(x_ptr[dim + d]);
            const float dout = bf16_to_float(d_out_ptr[d]);
 
            const float a2 = a * a;
            const float a3 = a2 * a;
            const float g = sqrt_2_over_pi * (a + coeff * a3);
            const float tanh_g = tanhf(g);
            const float sech2_g = 1.0f - tanh_g * tanh_g;
            const float g_prime = sqrt_2_over_pi * (1.0f + 3.0f * coeff * a2);
 
            const float d_gelu = 0.5f * (1.0f + tanh_g) + 0.5f * a * sech2_g * g_prime;
            d_x_ptr[d] = dout * d_gelu * b;
 
            const float gelu_a = 0.5f * a * (1.0f + tanh_g);
            d_x_ptr[dim + d] = dout * gelu_a;
        }
    }
}

References bf16_to_float().

geglu_backward_fp32()

void geglu_backward_fp32	(	const float *	x,
		const float *	d_out,
		float *	d_x,
		int	tokens,
		int	dim
	)

Definition at line 99 of file geglu_kernels.c.

{
    if (!x || !d_out || !d_x || tokens <= 0 || dim <= 0) {
        return;
    }
 
    const float sqrt_2_over_pi = 0.7978845608f;
    const float coeff = 0.044715f;
    const int inner_dim = dim * 2;
 
    for (int t = 0; t < tokens; ++t) {
        const float *x_ptr = x + (size_t)t * (size_t)inner_dim;
        const float *d_out_ptr = d_out + (size_t)t * (size_t)dim;
        float *d_x_ptr = d_x + (size_t)t * (size_t)inner_dim;
 
        for (int d = 0; d < dim; ++d) {
            float a = x_ptr[d];
            float b = x_ptr[dim + d];
            float dout = d_out_ptr[d];
 
            float a2 = a * a;
            float a3 = a2 * a;
            float g = sqrt_2_over_pi * (a + coeff * a3);
            float tanh_g = tanhf(g);
            float sech2_g = 1.0f - tanh_g * tanh_g;
            float g_prime = sqrt_2_over_pi * (1.0f + 3.0f * coeff * a2);
 
            float d_gelu = 0.5f * (1.0f + tanh_g) + 0.5f * a * sech2_g * g_prime;
            d_x_ptr[d] = dout * d_gelu * b;
 
            float gelu_a = 0.5f * a * (1.0f + tanh_g);
            d_x_ptr[dim + d] = dout * gelu_a;
        }
    }
}

geglu_forward_bf16()

void geglu_forward_bf16	(	const uint16_t *	x,
		uint16_t *	out,
		int	tokens,
		int	dim,
		float *	scratch
	)

Definition at line 83 of file geglu_kernels.c.

{
    if (!x || !out || !scratch || tokens <= 0 || dim <= 0) {
        return;
    }
 
    const size_t fp32_size = (size_t)tokens * (size_t)dim;
    const size_t input_size = fp32_size * 2;
    float *fp32_input = scratch;
    float *fp32_output = scratch + input_size;
 
    bf16_tensor_to_float(x, fp32_input, input_size);
    geglu_forward_fp32(fp32_input, fp32_output, tokens, dim);
    float_tensor_to_bf16(fp32_output, out, fp32_size);
}

References bf16_tensor_to_float(), float_tensor_to_bf16(), and geglu_forward_fp32().

geglu_forward_exact()

void geglu_forward_exact	(	const float *	x,
		float *	out,
		int	tokens,
		int	dim
	)

Definition at line 42 of file geglu_kernels.c.

{
    const int inner_dim = dim * 2;
    for (int t = 0; t < tokens; ++t) {
        const float *x_ptr = x + (size_t)t * (size_t)inner_dim;
        float *out_ptr = out + (size_t)t * (size_t)dim;
 
        for (int d = 0; d < dim; ++d) {
            out_ptr[d] = ck_gelu_tanh_parity_f32(x_ptr[d]) * x_ptr[dim + d];
        }
    }
}

References ck_gelu_tanh_parity_f32().

Referenced by geglu_forward_fp32().

geglu_forward_fp32()

void geglu_forward_fp32	(	const float *	x,
		float *	out,
		int	tokens,
		int	dim
	)

Definition at line 55 of file geglu_kernels.c.

{
    if (!x || !out || tokens <= 0 || dim <= 0) {
        return;
    }
 
    if (ck_strict_parity_enabled()) {
        geglu_forward_exact(x, out, tokens, dim);
        return;
    }
 
    const int inner_dim = dim * 2;
    for (int t = 0; t < tokens; ++t) {
        const float *x_ptr = x + (size_t)t * (size_t)inner_dim;
        float *out_ptr = out + (size_t)t * (size_t)dim;
 
        for (int d = 0; d < dim; ++d) {
            out_ptr[d] = x_ptr[d];
        }
 
        gelu_fast_inplace(out_ptr, (size_t)dim);
 
        for (int d = 0; d < dim; ++d) {
            out_ptr[d] *= x_ptr[dim + d];
        }
    }
}

References ck_strict_parity_enabled(), geglu_forward_exact(), and gelu_fast_inplace().

Referenced by geglu_forward_bf16().

geglu_forward_ggml_native()

void geglu_forward_ggml_native	(	const float *	x,
		float *	out,
		int	tokens,
		int	dim
	)

Definition at line 17 of file geglu_kernels.c.

{
    if (!x || !out || tokens <= 0 || dim <= 0) return;
    /* Preserve ascending compaction when output aliases the gate/up input. */
    for (int t = 0; t < tokens; ++t) {
        const float *row = x + (size_t)t * (size_t)dim * 2;
        float *dst = out + (size_t)t * dim;
        for (int d = 0; d < dim; ++d) {
            const float gate = row[d], up = row[dim + d];
            float value = gate;
            gelu_ggml_native_inplace(&value, 1);
            dst[d] = gate <= -10.0f ? 0.0f : value * up;
        }
    }
}

References gelu_ggml_native_inplace().

gelu_fast_inplace()

void gelu_fast_inplace ( float *  data, size_t  n  )

extern

GELU activation forward (fast approximation, in-place)

Test:

test_gelu.py::TestGELUForward::test_gelu_fast_inplace

test_gelu.py::TestGELUForward::test_gelu_vs_exact

test_parity.py::test_gelu_parity

Fast GELU approximation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) In-place on contiguous buffer.

After changes: make test && make llamacpp-parity-full

Definition at line 243 of file gelu_kernels.c.

{
    const float sqrt_2_over_pi = 0.7978845608f;
    const float coeff = 0.044715f;
 
#if defined(__AVX512F__)
    const __m512 sqrt_2_pi_vec = _mm512_set1_ps(sqrt_2_over_pi);
    const __m512 coeff_vec = _mm512_set1_ps(coeff);
    const __m512 half_vec = _mm512_set1_ps(0.5f);
    const __m512 one_vec = _mm512_set1_ps(1.0f);
 
    size_t i = 0;
    for (; i + 16 <= n; i += 16) {
        __m512 x = _mm512_loadu_ps(&data[i]);
        __m512 x2 = _mm512_mul_ps(x, x);
        __m512 x3 = _mm512_mul_ps(x2, x);
 
        // inner = sqrt(2/pi) * (x + 0.044715 * x^3)
        __m512 inner = _mm512_fmadd_ps(coeff_vec, x3, x);
        inner = _mm512_mul_ps(sqrt_2_pi_vec, inner);
 
        // result = 0.5 * x * (1 + tanh(inner))
        __m512 tanh_val = tanh512_fast(inner);
        __m512 one_plus_tanh = _mm512_add_ps(one_vec, tanh_val);
        __m512 result = _mm512_mul_ps(half_vec, _mm512_mul_ps(x, one_plus_tanh));
 
        _mm512_storeu_ps(&data[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = data[i];
        float x3 = x * x * x;
        float inner = sqrt_2_over_pi * (x + coeff * x3);
        data[i] = 0.5f * x * (1.0f + tanhf(inner));
    }
 
#elif defined(__AVX2__)
    const __m256 sqrt_2_pi_vec = _mm256_set1_ps(sqrt_2_over_pi);
    const __m256 coeff_vec = _mm256_set1_ps(coeff);
    const __m256 half_vec = _mm256_set1_ps(0.5f);
    const __m256 one_vec = _mm256_set1_ps(1.0f);
 
    size_t i = 0;
    for (; i + 8 <= n; i += 8) {
        __m256 x = _mm256_loadu_ps(&data[i]);
        __m256 x2 = _mm256_mul_ps(x, x);
        __m256 x3 = _mm256_mul_ps(x2, x);
 
        // inner = sqrt(2/pi) * (x + 0.044715 * x^3)
        __m256 inner = _mm256_fmadd_ps(coeff_vec, x3, x);
        inner = _mm256_mul_ps(sqrt_2_pi_vec, inner);
 
        // result = 0.5 * x * (1 + tanh(inner))
        __m256 tanh_val = tanh256_fast(inner);
        __m256 one_plus_tanh = _mm256_add_ps(one_vec, tanh_val);
        __m256 result = _mm256_mul_ps(half_vec, _mm256_mul_ps(x, one_plus_tanh));
 
        _mm256_storeu_ps(&data[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = data[i];
        float x3 = x * x * x;
        float inner = sqrt_2_over_pi * (x + coeff * x3);
        data[i] = 0.5f * x * (1.0f + tanhf(inner));
    }
 
#elif defined(__AVX__)
    // AVX1: Vectorize arithmetic, use scalar tanh
    const __m256 sqrt_2_pi_vec = _mm256_set1_ps(sqrt_2_over_pi);
    const __m256 coeff_vec = _mm256_set1_ps(coeff);
    const __m256 half_vec = _mm256_set1_ps(0.5f);
    const __m256 one_vec = _mm256_set1_ps(1.0f);
 
    size_t i = 0;
    float inner_arr[8] __attribute__((aligned(32)));
    float tanh_arr[8] __attribute__((aligned(32)));
 
    for (; i + 8 <= n; i += 8) {
        __m256 x = _mm256_loadu_ps(&data[i]);
        __m256 x2 = _mm256_mul_ps(x, x);
        __m256 x3 = _mm256_mul_ps(x2, x);
 
        // inner = sqrt(2/pi) * (x + 0.044715 * x^3)
        __m256 coeff_x3 = _mm256_mul_ps(coeff_vec, x3);
        __m256 inner = _mm256_mul_ps(sqrt_2_pi_vec, _mm256_add_ps(x, coeff_x3));
 
        // Compute tanh scalarly
        _mm256_store_ps(inner_arr, inner);
        for (int j = 0; j < 8; ++j) {
            tanh_arr[j] = tanhf(inner_arr[j]);
        }
        __m256 tanh_val = _mm256_load_ps(tanh_arr);
 
        // result = 0.5 * x * (1 + tanh(inner))
        __m256 one_plus_tanh = _mm256_add_ps(one_vec, tanh_val);
        __m256 result = _mm256_mul_ps(half_vec, _mm256_mul_ps(x, one_plus_tanh));
 
        _mm256_storeu_ps(&data[i], result);
    }
    // Handle remaining elements
    for (; i < n; ++i) {
        float x = data[i];
        float x3 = x * x * x;
        float inner = sqrt_2_over_pi * (x + coeff * x3);
        data[i] = 0.5f * x * (1.0f + tanhf(inner));
    }
 
#else
    // Scalar fallback
    for (size_t i = 0; i < n; ++i) {
        float x = data[i];
        float x3 = x * x * x;
        float inner = sqrt_2_over_pi * (x + coeff * x3);
        data[i] = 0.5f * x * (1.0f + tanhf(inner));
    }
#endif
}

Referenced by geglu_forward_fp32().

gelu_ggml_native_inplace()

void gelu_ggml_native_inplace ( float *  data, size_t  n  )

extern

Definition at line 741 of file gelu_kernels.c.

{
    pthread_once(&ck_gelu_ggml_table_once, ck_gelu_ggml_table_init);
    for (size_t i = 0; i < n; ++i) {
        const float x = data[i];
        if (x <= -10.0f) data[i] = 0.0f;
        else if (x >= 10.0f) data[i] = x;
        else data[i] = ggml_fp16_to_fp32(
            ck_gelu_ggml_table_f16[(uint16_t)ggml_fp32_to_fp16(x)]);
    }
}

Referenced by geglu_forward_ggml_native().