recurrent_norm_kernels.c File Reference

#include "bf16_utils.h"
#include "ckernel_engine.h"
#include <math.h>
#include <string.h>

Go to the source code of this file.

Functions
void	recurrent_norm_gate_backward (const float *d_out, const float *x, const float *gate, const float *weight, float *d_x, float *d_gate, float *d_weight, int rows, int num_heads, int head_dim, float eps)
void	recurrent_norm_gate_forward (const float *x, const float *gate, const float *weight, float *out, int rows, int num_heads, int head_dim, float eps)
void	recurrent_norm_gate_llama_avx2_forward (const float *x, const float *gate, const float *weight, float *out, int rows, int num_heads, int head_dim, float eps)
void	recurrent_norm_gate_pytorch_bf16_storage (const float *x, const float *gate, const float *weight, float *out, int rows, int num_heads, int head_dim, float eps)
void	recurrent_norm_sigmoid_gate_llama_avx2_forward (const float *x, const float *gate, const float *weight, float *out, int rows, int num_heads, int head_dim, float eps)
void	recurrent_norm_sigmoid_gate_pytorch_bf16_storage (const float *x, const float *gate, const float *weight, float *out, int rows, int num_heads, int head_dim, float eps)
static float	recurrent_sigmoid_local (float x)

Function Documentation

recurrent_norm_gate_backward()

void recurrent_norm_gate_backward	(	const float *	d_out,
		const float *	x,
		const float *	gate,
		const float *	weight,
		float *	d_x,
		float *	d_gate,
		float *	d_weight,
		int	rows,
		int	num_heads,
		int	head_dim,
		float	eps
	)

Definition at line 182 of file recurrent_norm_kernels.c.

                                             {
    const int inner_dim = num_heads * head_dim;
    memset(d_weight, 0, (size_t) head_dim * sizeof(float));
 
    for (int row = 0; row < rows; ++row) {
        const float *d_out_row = d_out + (size_t) row * (size_t) inner_dim;
        const float *x_row = x + (size_t) row * (size_t) inner_dim;
        const float *gate_row = gate + (size_t) row * (size_t) inner_dim;
        float *d_x_row = d_x + (size_t) row * (size_t) inner_dim;
        float *d_gate_row = d_gate + (size_t) row * (size_t) inner_dim;
 
        for (int head = 0; head < num_heads; ++head) {
            const float *x_head = x_row + (size_t) head * (size_t) head_dim;
            const float *gate_head = gate_row + (size_t) head * (size_t) head_dim;
            const float *d_out_head = d_out_row + (size_t) head * (size_t) head_dim;
            float *d_x_head = d_x_row + (size_t) head * (size_t) head_dim;
            float *d_gate_head = d_gate_row + (size_t) head * (size_t) head_dim;
 
            float ms = 0.0f;
            for (int col = 0; col < head_dim; ++col) {
                ms += x_head[col] * x_head[col];
            }
            ms /= (float) head_dim;
            const float inv_rms = 1.0f / sqrtf(ms + eps);
            const float inv_rms3_over_dim = (inv_rms * inv_rms * inv_rms) / (float) head_dim;
 
            float dot = 0.0f;
            for (int col = 0; col < head_dim; ++col) {
                const float g = gate_head[col];
                const float sig = recurrent_sigmoid_local(g);
                const float silu = g * sig;
                dot += d_out_head[col] * weight[col] * silu * x_head[col];
            }
 
            for (int col = 0; col < head_dim; ++col) {
                const float g = gate_head[col];
                const float sig = recurrent_sigmoid_local(g);
                const float silu = g * sig;
                const float scaled = inv_rms * weight[col] * silu;
                d_x_head[col] = d_out_head[col] * scaled - x_head[col] * inv_rms3_over_dim * dot;
                d_weight[col] += d_out_head[col] * (x_head[col] * inv_rms) * silu;
                d_gate_head[col] = d_out_head[col] * (x_head[col] * inv_rms * weight[col]) *
                                   (sig + g * sig * (1.0f - sig));
            }
        }
    }
}

References recurrent_sigmoid_local(), and silu().

recurrent_norm_gate_forward()

void recurrent_norm_gate_forward	(	const float *	x,
		const float *	gate,
		const float *	weight,
		float *	out,
		int	rows,
		int	num_heads,
		int	head_dim,
		float	eps
	)

Definition at line 18 of file recurrent_norm_kernels.c.

                                            {
    const int inner_dim = num_heads * head_dim;
    for (int row = 0; row < rows; ++row) {
        const float *x_row = x + (size_t) row * (size_t) inner_dim;
        const float *gate_row = gate + (size_t) row * (size_t) inner_dim;
        float *out_row = out + (size_t) row * (size_t) inner_dim;
 
        for (int head = 0; head < num_heads; ++head) {
            const float *x_head = x_row + (size_t) head * (size_t) head_dim;
            const float *gate_head = gate_row + (size_t) head * (size_t) head_dim;
            float *out_head = out_row + (size_t) head * (size_t) head_dim;
 
            float ms = 0.0f;
            for (int col = 0; col < head_dim; ++col) {
                ms += x_head[col] * x_head[col];
            }
            ms /= (float) head_dim;
            const float inv_rms = 1.0f / sqrtf(ms + eps);
 
            for (int col = 0; col < head_dim; ++col) {
                const float g = gate_head[col];
                const float silu = g * recurrent_sigmoid_local(g);
                out_head[col] = x_head[col] * inv_rms * weight[col] * silu;
            }
        }
    }
}

References recurrent_sigmoid_local(), and silu().

Referenced by ck_test_recurrent_norm_gate().

recurrent_norm_gate_llama_avx2_forward()

void recurrent_norm_gate_llama_avx2_forward	(	const float *	x,
		const float *	gate,
		const float *	weight,
		float *	out,
		int	rows,
		int	num_heads,
		int	head_dim,
		float	eps
	)

Definition at line 53 of file recurrent_norm_kernels.c.

                                                        {
    if (!x || !gate || !weight || !out || rows <= 0 || num_heads <= 0 ||
        head_dim <= 0 || head_dim > 4096) {
        return;
    }
    const int inner_dim = num_heads * head_dim;
    float normalized[4096];
    float silu[4096];
    for (int row = 0; row < rows; ++row) {
        for (int head = 0; head < num_heads; ++head) {
            const size_t offset = (size_t) row * (size_t) inner_dim
                                + (size_t) head * (size_t) head_dim;
            rmsnorm_forward_llama_production(
                x + offset, weight, normalized, NULL, 1, head_dim, head_dim, eps);
            recurrent_silu_forward_ggml(gate + offset, silu, 1, head_dim);
            for (int col = 0; col < head_dim; ++col) {
                out[offset + (size_t) col] = normalized[col] * silu[col];
            }
        }
    }
}

References recurrent_silu_forward_ggml(), rmsnorm_forward_llama_production(), and silu().

recurrent_norm_gate_pytorch_bf16_storage()

void recurrent_norm_gate_pytorch_bf16_storage	(	const float *	x,
		const float *	gate,
		const float *	weight,
		float *	out,
		int	rows,
		int	num_heads,
		int	head_dim,
		float	eps
	)

Definition at line 111 of file recurrent_norm_kernels.c.

{
    if (!x || !gate || !weight || !out || rows <= 0 || num_heads <= 0 ||
        head_dim <= 0 || head_dim > 4096) {
        return;
    }
    const int inner_dim = num_heads * head_dim;
    float normalized[4096];
    float silu[4096];
    for (int row = 0; row < rows; ++row) {
        for (int head = 0; head < num_heads; ++head) {
            const size_t offset = (size_t)row * (size_t)inner_dim +
                (size_t)head * (size_t)head_dim;
            /*
             * Qwen3Next order:
             * FP32 RMS statistics -> BF16 normalized value -> BF16 weight
             * multiply -> FP32 SiLU(gate) -> BF16 final output.
             */
            rmsnorm_forward_pytorch_bf16_storage(
                x + offset, weight, normalized, NULL,
                1, head_dim, head_dim, eps);
            recurrent_silu_forward_pytorch_bf16_input_fp32_output(
                gate + offset, silu, 1, head_dim);
            for (int col = 0; col < head_dim; ++col) {
                out[offset + (size_t)col] = bf16_to_float(float_to_bf16(
                    normalized[col] * silu[col]));
            }
        }
    }
}

References bf16_to_float(), float_to_bf16(), recurrent_silu_forward_pytorch_bf16_input_fp32_output(), rmsnorm_forward_pytorch_bf16_storage(), and silu().

recurrent_norm_sigmoid_gate_llama_avx2_forward()

void recurrent_norm_sigmoid_gate_llama_avx2_forward	(	const float *	x,
		const float *	gate,
		const float *	weight,
		float *	out,
		int	rows,
		int	num_heads,
		int	head_dim,
		float	eps
	)

Definition at line 82 of file recurrent_norm_kernels.c.

                                                                {
    if (!x || !gate || !weight || !out || rows <= 0 || num_heads <= 0 ||
        head_dim <= 0 || head_dim > 4096) {
        return;
    }
    const int inner_dim = num_heads * head_dim;
    float normalized[4096];
    float sigmoid[4096];
    for (int row = 0; row < rows; ++row) {
        for (int head = 0; head < num_heads; ++head) {
            const size_t offset = (size_t) row * (size_t) inner_dim
                                + (size_t) head * (size_t) head_dim;
            rmsnorm_forward_llama_production(
                x + offset, weight, normalized, NULL, 1, head_dim, head_dim, eps);
            recurrent_sigmoid_forward_ggml(gate + offset, sigmoid, 1, head_dim);
            for (int col = 0; col < head_dim; ++col) {
                out[offset + (size_t) col] = normalized[col] * sigmoid[col];
            }
        }
    }
}

References recurrent_sigmoid_forward_ggml(), and rmsnorm_forward_llama_production().

recurrent_norm_sigmoid_gate_pytorch_bf16_storage()

void recurrent_norm_sigmoid_gate_pytorch_bf16_storage	(	const float *	x,
		const float *	gate,
		const float *	weight,
		float *	out,
		int	rows,
		int	num_heads,
		int	head_dim,
		float	eps
	)

Definition at line 149 of file recurrent_norm_kernels.c.

{
    if (!x || !gate || !weight || !out || rows <= 0 || num_heads <= 0 ||
        head_dim <= 0 || head_dim > 4096) {
        return;
    }
    const int inner_dim = num_heads * head_dim;
    float normalized[4096];
    float sigmoid[4096];
    for (int row = 0; row < rows; ++row) {
        for (int head = 0; head < num_heads; ++head) {
            const size_t offset = (size_t)row * (size_t)inner_dim +
                (size_t)head * (size_t)head_dim;
            rmsnorm_forward_pytorch_bf16_storage(
                x + offset, weight, normalized, NULL,
                1, head_dim, head_dim, eps);
            recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output(
                gate + offset, sigmoid, 1, head_dim);
            for (int col = 0; col < head_dim; ++col) {
                out[offset + (size_t)col] = bf16_to_float(float_to_bf16(
                    normalized[col] * sigmoid[col]));
            }
        }
    }
}

References bf16_to_float(), float_to_bf16(), recurrent_sigmoid_forward_pytorch_bf16_input_fp32_output(), and rmsnorm_forward_pytorch_bf16_storage().

recurrent_sigmoid_local()

static float recurrent_sigmoid_local ( float  x )

inlinestatic

Definition at line 7 of file recurrent_norm_kernels.c.

                                                     {
    if (x >= 0.0f) {
        const float z = expf(-x);
        return 1.0f / (1.0f + z);
    }
    {
        const float z = expf(x);
        return z / (1.0f + z);
    }
}

Referenced by recurrent_norm_gate_backward(), and recurrent_norm_gate_forward().