layernorm_kernels.c File Reference

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

#include "ckernel_engine.h"
#include "bf16_utils.h"
#include <math.h>
#include <stdlib.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)
static void	layernorm_forward_ggml_exact (const float *input, const float *gamma, const float *beta, float *output, float *mean_cache, float *rstd_cache, int tokens, int d_model, int input_stride, int output_stride, int aligned_embed_dim, float eps)
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_bf16_storage (const float *input, const float *gamma, const float *beta, float *output, float *mean_cache, float *rstd_cache, int tokens, int d_model, 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)
void	layernorm_pytorch_welford_bf16_storage (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 1007 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_ggml_exact()

static void layernorm_forward_ggml_exact ( const float *  input, const float *  gamma, const float *  beta, float *  output, float *  mean_cache, float *  rstd_cache, int  tokens, int  d_model, int  input_stride, int  output_stride, int  aligned_embed_dim, float  eps  )

static

Definition at line 46 of file layernorm_kernels.c.

{
    for (int t = 0; t < tokens; ++t) {
        const float *x = input + (size_t)t * (size_t)input_stride;
        float *y = output + (size_t)t * (size_t)output_stride;
 
        double sum_acc = 0.0;
#if defined(__clang__)
#pragma clang loop vectorize(disable)
#pragma clang loop interleave(disable)
#endif
        for (int i = 0; i < d_model; ++i) {
            sum_acc += (double)x[i];
        }
        const float sum = (float)sum_acc;
        const float mean = sum / (float)d_model;
 
        double var_acc = 0.0;
        int i = 0;
#if defined(__AVX2__) && defined(__FMA__)
        /* This provider promises ggml's ordered AVX2 reduction. Keep that
         * arithmetic identity on AVX-512 hosts instead of widening it based
         * on the compiler target. */
        for (; i + 7 < d_model; i += 8) {
            __m256 val = _mm256_sub_ps(_mm256_loadu_ps(x + i),
                                       _mm256_set1_ps(mean));
            _mm256_storeu_ps(y + i, val);
            val = _mm256_mul_ps(val, val);
            __m128 val2 = _mm_add_ps(_mm256_extractf128_ps(val, 1),
                                     _mm256_castps256_ps128(val));
            val2 = _mm_add_ps(val2, _mm_movehl_ps(val2, val2));
            val2 = _mm_add_ss(val2, _mm_movehdup_ps(val2));
            var_acc += (double)_mm_cvtss_f32(val2);
        }
#elif defined(__SSE2__)
        for (; i + 3 < d_model; i += 4) {
            __m128 val = _mm_sub_ps(_mm_loadu_ps(x + i),
                                    _mm_set1_ps(mean));
            _mm_storeu_ps(y + i, val);
            val = _mm_mul_ps(val, val);
#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
            val = _mm_add_ps(val, _mm_movehl_ps(val, val));
            val = _mm_add_ss(val, _mm_movehdup_ps(val));
#else
            __m128 tmp = _mm_shuffle_ps(val, val, _MM_SHUFFLE(2, 3, 0, 1));
            val = _mm_add_ps(val, tmp);
            tmp = _mm_movehl_ps(tmp, val);
            val = _mm_add_ss(val, tmp);
#endif
            var_acc += (double)_mm_cvtss_f32(val);
        }
#endif
#if defined(__clang__)
#pragma clang loop vectorize(disable)
#pragma clang loop interleave(disable)
#endif
        for (; i < d_model; ++i) {
            const float centered = x[i] - mean;
            y[i] = centered;
            var_acc += (double)(centered * centered);
        }
        const float variance = (float)(var_acc / (double)d_model);
        const float scale = 1.0f / sqrtf(variance + eps);
 
        if (mean_cache) {
            mean_cache[t] = mean;
        }
        if (rstd_cache) {
            rstd_cache[t] = scale;
        }
 
#if defined(__clang__)
#pragma clang loop vectorize(disable)
#pragma clang loop interleave(disable)
#endif
        for (int i = 0; i < d_model; ++i) {
            y[i] *= scale;
        }
        if (gamma) {
#if defined(__clang__)
#pragma clang loop vectorize(disable)
#pragma clang loop interleave(disable)
#endif
            for (int i = 0; i < d_model; ++i) {
                y[i] *= gamma[i];
            }
        }
        if (beta) {
#if defined(__clang__)
#pragma clang loop vectorize(disable)
#pragma clang loop interleave(disable)
#endif
            for (int i = 0; i < d_model; ++i) {
                y[i] += beta[i];
            }
        }
 
        if (aligned_embed_dim > d_model) {
            zero_layernorm_padding(y, d_model, aligned_embed_dim);
        }
    }
}

References zero_layernorm_padding().

Referenced by layernorm_forward_rolled_slice(), layernorm_forward_unrolled_slice(), layernorm_naive_serial_bf16_storage(), and layernorm_naive_serial_matched_precision().

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 398 of file layernorm_kernels.c.

{
    if (ck_strict_parity_enabled()) {
        layernorm_forward_ggml_exact(input_slice_base, gamma, beta,
                                     output_slice_base, mean_cache_slice, rstd_cache_slice,
                                     num_tokens_in_slice, d_model,
                                     aligned_embed_dim, aligned_embed_dim, aligned_embed_dim, eps);
        return;
    }
 
#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 ck_strict_parity_enabled(), layernorm_forward_ggml_exact(), and 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 730 of file layernorm_kernels.c.

{
    if (ck_strict_parity_enabled()) {
        layernorm_forward_ggml_exact(input_slice_base, gamma, beta,
                                     output_slice_base, mean_cache_slice, rstd_cache_slice,
                                     num_tokens_in_slice, d_model,
                                     d_model, d_model, d_model, eps);
        return;
    }
 
#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 ck_strict_parity_enabled(), layernorm_forward_ggml_exact(), and 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 714 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 175 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_bf16_storage()

void layernorm_naive_serial_bf16_storage	(	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 786 of file layernorm_kernels.c.

{
    const size_t count = (size_t)tokens * (size_t)d_model;
    for (size_t i = 0; i < count; ++i) {
        output[i] = bf16_to_float(float_to_bf16(input[i]));
    }
    layernorm_forward_ggml_exact(output, gamma, beta,
                                 output, mean_cache, rstd_cache,
                                 tokens, d_model,
                                 d_model, d_model, d_model, eps);
    for (size_t i = 0; i < count; ++i) {
        output[i] = bf16_to_float(float_to_bf16(output[i]));
    }
}

References bf16_to_float(), float_to_bf16(), and layernorm_forward_ggml_exact().

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 764 of file layernorm_kernels.c.

{
    layernorm_forward_ggml_exact(input, gamma, beta,
                                 output, mean_cache, rstd_cache,
                                 tokens, d_model,
                                 d_model, d_model, d_model, eps);
}

References layernorm_forward_ggml_exact().

Referenced by layernorm_forward_unrolled_slice_scalar().

layernorm_pytorch_welford_bf16_storage()

void layernorm_pytorch_welford_bf16_storage	(	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 952 of file layernorm_kernels.c.

{
#if !defined(__AVX2__) || !defined(__FMA__)
    (void)input; (void)gamma; (void)beta; (void)output;
    (void)mean_cache; (void)rstd_cache; (void)tokens; (void)d_model; (void)eps;
    abort();
#else
    for (int t = 0; t < tokens; ++t) {
        const float *x = input + (size_t)t * (size_t)d_model;
        float *y = output + (size_t)t * (size_t)d_model;
        float mean;
        float variance;
        layernorm_pytorch_bf16_rowwise_moments_avx2(x, d_model, &mean, &variance);
        const float rstd = 1.0f / sqrtf(variance + eps);
        const float bias = -rstd * mean;
        int i = 0;
        for (; i + 7 < d_model; i += 8) {
            const __m256 x_vec = _mm256_loadu_ps(x + i);
            const __m256 gamma_vec = gamma ? _mm256_loadu_ps(gamma + i) : _mm256_set1_ps(1.0f);
            const __m256 beta_vec = beta ? _mm256_loadu_ps(beta + i) : _mm256_setzero_ps();
            const __m256 normalized = _mm256_fmadd_ps(
                x_vec, _mm256_set1_ps(rstd), _mm256_set1_ps(bias));
            const __m256 transformed = _mm256_fmadd_ps(normalized, gamma_vec, beta_vec);
            float lanes[8];
            _mm256_storeu_ps(lanes, transformed);
            for (int lane = 0; lane < 8; ++lane) {
                y[i + lane] = bf16_to_float(float_to_bf16(lanes[lane]));
            }
        }
        for (; i < d_model; ++i) {
            const float gamma_v = gamma ? gamma[i] : 1.0f;
            const float beta_v = beta ? beta[i] : 0.0f;
            const float value = fmaf(fmaf(x[i], rstd, bias), gamma_v, beta_v);
            y[i] = bf16_to_float(float_to_bf16(value));
        }
        if (mean_cache) {
            mean_cache[t] = mean;
        }
        if (rstd_cache) {
            rstd_cache[t] = rstd;
        }
    }
#endif
}

References bf16_to_float(), and float_to_bf16().

zero_layernorm_padding()

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

inlinestatic

Definition at line 26 of file layernorm_kernels.c.

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

Referenced by layernorm_forward_ggml_exact().