layernorm_kernels.c File Reference

LayerNorm forward/backward kernels with SIMD (SSE/AVX/AVX512) More...

#include <math.h>

Go to the source code of this file.

Functions
void	layernorm_backward_kernel (const float *d_output, const float *input, const float *gamma, const float *mean, const float *rstd, float *d_input, float *d_gamma, float *d_beta, int tokens, int d_model, int aligned_embed_dim)
void	layernorm_forward_rolled_slice (const float *__restrict input_slice_base, const float *__restrict gamma, const float *__restrict beta, float *__restrict output_slice_base, float *__restrict mean_cache_slice, float *__restrict rstd_cache_slice, int num_tokens_in_slice, int d_model, int aligned_embed_dim, float eps)
void	layernorm_forward_unrolled_slice (const float *__restrict input_slice_base, const float *__restrict gamma, const float *__restrict beta, float *__restrict output_slice_base, float *__restrict mean_cache_slice, float *__restrict rstd_cache_slice, int num_tokens_in_slice, int d_model, float eps)
static void	layernorm_forward_unrolled_slice_scalar (const float *__restrict input_slice_base, const float *__restrict gamma, const float *__restrict beta, float *__restrict output_slice_base, float *__restrict mean_cache_slice, float *__restrict rstd_cache_slice, int num_tokens_in_slice, int d_model, float eps)
void	layernorm_naive_serial (const float *input, const float *gamma, const float *beta, float *output, float *mean_cache, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
void	layernorm_naive_serial_matched_precision (const float *input, const float *gamma, const float *beta, float *output, float *mean_cache, float *rstd_cache, int tokens, int d_model, float eps)
static void	zero_layernorm_padding (float *out_ptr, int d_model, int aligned_embed_dim)

Detailed Description

LayerNorm forward/backward kernels with SIMD (SSE/AVX/AVX512)

CK-ENGINE KERNEL RULES:

1. NO malloc/free - memory via bump allocator, pointers passed in
2. NO OpenMP - parallelization at orchestrator/codegen layer
3. API must define: inputs, outputs, workspace, and memory layouts
4. Pure computation - deterministic, no side effects

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

LayerNorm: y = gamma * (x - mean) / sqrt(var + eps) + beta

Definition in file layernorm_kernels.c.

Function Documentation

layernorm_backward_kernel()

void layernorm_backward_kernel	(	const float *	d_output,
		const float *	input,
		const float *	gamma,
		const float *	mean,
		const float *	rstd,
		float *	d_input,
		float *	d_gamma,
		float *	d_beta,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim
	)

Definition at line 668 of file layernorm_kernels.c.

{
    int T = tokens;
    int D = d_model;
    int aligned_D = aligned_embed_dim;
 
    // Per-token input gradients
    for (int t = 0; t < T; ++t) {
        float mean_t = mean[t];
        float rstd_t = rstd[t];
 
        float d_y_gamma_sum = 0.0f;
        float d_y_gamma_xhat_sum = 0.0f;
 
        // First pass: compute sums
        for (int d = 0; d < D; ++d) {
            float x = input[t * aligned_D + d];
            float x_hat = (x - mean_t) * rstd_t;
            float d_y = d_output[t * aligned_D + d];
            float d_y_gamma = d_y * gamma[d];
 
            d_y_gamma_sum += d_y_gamma;
            d_y_gamma_xhat_sum += d_y_gamma * x_hat;
        }
 
        // Second pass: compute input gradients
        float scale = rstd_t / (float)D;
        for (int d = 0; d < D; ++d) {
            float x = input[t * aligned_D + d];
            float x_hat = (x - mean_t) * rstd_t;
            float d_y = d_output[t * aligned_D + d];
 
            d_input[t * aligned_D + d] =
                scale * ((float)D * d_y * gamma[d] - d_y_gamma_sum - x_hat * d_y_gamma_xhat_sum);
        }
 
        // Zero padding for aligned dimension beyond D
        for (int d = D; d < aligned_D; ++d) {
            d_input[t * aligned_D + d] = 0.0f;
        }
    }
 
    // Parameter gradients (gamma, beta)
    for (int d = 0; d < D; ++d) {
        float gamma_grad = 0.0f;
        float beta_grad = 0.0f;
 
        for (int t = 0; t < T; ++t) {
            float x = input[t * aligned_D + d];
            float x_hat = (x - mean[t]) * rstd[t];
            float d_y = d_output[t * aligned_D + d];
 
            gamma_grad += d_y * x_hat;
            beta_grad += d_y;
        }
 
        d_gamma[d] += gamma_grad;
        d_beta[d] += beta_grad;
    }
}

Referenced by layernorm_backward_kernel_bf16().

layernorm_forward_rolled_slice()

void layernorm_forward_rolled_slice	(	const float *__restrict	input_slice_base,
		const float *__restrict	gamma,
		const float *__restrict	beta,
		float *__restrict	output_slice_base,
		float *__restrict	mean_cache_slice,
		float *__restrict	rstd_cache_slice,
		int	num_tokens_in_slice,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 274 of file layernorm_kernels.c.

{
#if defined(__AVX512F__)
    layernorm_forward_rolled_slice_avx512(input_slice_base, gamma, beta,
                                           output_slice_base, mean_cache_slice, rstd_cache_slice,
                                           num_tokens_in_slice, d_model, aligned_embed_dim, eps);
#elif defined(__AVX2__) || defined(__AVX__)
    layernorm_forward_rolled_slice_avx256(input_slice_base, gamma, beta,
                                           output_slice_base, mean_cache_slice, rstd_cache_slice,
                                           num_tokens_in_slice, d_model, aligned_embed_dim, eps);
#else
    layernorm_naive_serial(input_slice_base, gamma, beta,
                           output_slice_base, mean_cache_slice, rstd_cache_slice,
                           num_tokens_in_slice, d_model, aligned_embed_dim, eps);
#endif
}

References layernorm_naive_serial().

Referenced by layernorm_forward_rolled_slice_bf16().

layernorm_forward_unrolled_slice()

void layernorm_forward_unrolled_slice	(	const float *__restrict	input_slice_base,
		const float *__restrict	gamma,
		const float *__restrict	beta,
		float *__restrict	output_slice_base,
		float *__restrict	mean_cache_slice,
		float *__restrict	rstd_cache_slice,
		int	num_tokens_in_slice,
		int	d_model,
		float	eps
	)

Definition at line 598 of file layernorm_kernels.c.

{
#if defined(__AVX512F__)
    layernorm_forward_unrolled_slice_avx512(input_slice_base, gamma, beta,
                                             output_slice_base, mean_cache_slice, rstd_cache_slice,
                                             num_tokens_in_slice, d_model, eps);
#elif defined(__AVX2__) || defined(__AVX__)
    layernorm_forward_unrolled_slice_avx256(input_slice_base, gamma, beta,
                                             output_slice_base, mean_cache_slice, rstd_cache_slice,
                                             num_tokens_in_slice, d_model, eps);
#else
    layernorm_forward_unrolled_slice_scalar(input_slice_base, gamma, beta,
                                            output_slice_base, mean_cache_slice, rstd_cache_slice,
                                            num_tokens_in_slice, d_model, eps);
#endif
}

References layernorm_forward_unrolled_slice_scalar().

Referenced by layernorm_forward_unrolled_slice_bf16().

layernorm_forward_unrolled_slice_scalar()

static void layernorm_forward_unrolled_slice_scalar ( const float *__restrict  input_slice_base, const float *__restrict  gamma, const float *__restrict  beta, float *__restrict  output_slice_base, float *__restrict  mean_cache_slice, float *__restrict  rstd_cache_slice, int  num_tokens_in_slice, int  d_model, float  eps  )

static

Definition at line 582 of file layernorm_kernels.c.

{
    layernorm_naive_serial_matched_precision(input_slice_base, gamma, beta,
                                             output_slice_base, mean_cache_slice, rstd_cache_slice,
                                             num_tokens_in_slice, d_model, eps);
}

References layernorm_naive_serial_matched_precision().

Referenced by layernorm_forward_unrolled_slice().

layernorm_naive_serial()

void layernorm_naive_serial	(	const float *	input,
		const float *	gamma,
		const float *	beta,
		float *	output,
		float *	mean_cache,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 51 of file layernorm_kernels.c.

{
    for (int t = 0; t < tokens; ++t) {
        const float *in_ptr = input + t * aligned_embed_dim;
        float *out_ptr = output + t * aligned_embed_dim;
 
        float sum_val = 0.0f;
        for (int i = 0; i < d_model; ++i) {
            sum_val += in_ptr[i];
        }
        float mean = sum_val / (float)d_model;
 
        float sum_sq_diff = 0.0f;
        for (int i = 0; i < d_model; ++i) {
            float diff = in_ptr[i] - mean;
            sum_sq_diff += diff * diff;
        }
        float variance = sum_sq_diff / (float)d_model + eps;
 
        double var_double = (double)variance;
        float inv_std = (float)(1.0 / sqrt(var_double));
 
        for (int i = 0; i < d_model; ++i) {
            float normalized_val = (in_ptr[i] - mean) * inv_std;
            out_ptr[i] = normalized_val * gamma[i] + beta[i];
        }
 
        if (mean_cache) {
            mean_cache[t] = mean;
        }
        if (rstd_cache) {
            rstd_cache[t] = inv_std;
        }
        /* Keep aligned padding quiet so future GEMMs see deterministic memory. */
        if (aligned_embed_dim > d_model) {
            /* Keep padded lanes zeroed so subsequent GEMMs never read stale data. */
            for (int i = d_model; i < aligned_embed_dim; ++i) {
                out_ptr[i] = 0.0f;
            }
        }
    }
}

Referenced by layernorm_forward_rolled_slice().

layernorm_naive_serial_matched_precision()

void layernorm_naive_serial_matched_precision	(	const float *	input,
		const float *	gamma,
		const float *	beta,
		float *	output,
		float *	mean_cache,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		float	eps
	)

Definition at line 624 of file layernorm_kernels.c.

{
    for (int t = 0; t < tokens; ++t) {
        const float *in_ptr = input + t * d_model;
        float *out_ptr = output + t * d_model;
 
        float sum_val = 0.0f;
        for (int i = 0; i < d_model; ++i) {
            sum_val += in_ptr[i];
        }
        float mean = sum_val / (float)d_model;
 
        float sum_sq_diff = 0.0f;
        for (int i = 0; i < d_model; ++i) {
            float diff = in_ptr[i] - mean;
            sum_sq_diff += diff * diff;
        }
        float variance = sum_sq_diff / (float)d_model + eps;
 
        double var_double = (double)variance;
        float inv_std = (float)(1.0 / sqrt(var_double));
 
        for (int i = 0; i < d_model; ++i) {
            float normalized_val = (in_ptr[i] - mean) * inv_std;
            out_ptr[i] = normalized_val * gamma[i] + beta[i];
        }
 
        if (mean_cache) {
            mean_cache[t] = mean;
        }
        if (rstd_cache) {
            rstd_cache[t] = inv_std;
        }
    }
}

Referenced by layernorm_forward_unrolled_slice_scalar().

zero_layernorm_padding()

static void zero_layernorm_padding ( float *  out_ptr, int  d_model, int  aligned_embed_dim  )

inlinestatic

Definition at line 22 of file layernorm_kernels.c.

{
    for (int idx = d_model; idx < aligned_embed_dim; ++idx) {
        out_ptr[idx] = 0.0f;
    }
}