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

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

Functions
const char *	ck_rmsnorm_provider_name (void)

void	gemma4_v_norm_forward (const float input, float output, float *rstd_cache, int tokens, int num_kv_heads, int head_dim, float eps)

void	rmsnorm_backward (const float d_output, const float input, const float gamma, const float rstd_cache, float d_input, float d_gamma, int tokens, int d_model, int aligned_embed_dim)

static void	rmsnorm_backward_strict_scalar (const float d_output, const float input, const float gamma, const float rstd_cache, float d_input, float d_gamma, int tokens, int d_model, int aligned_embed_dim)

void	rmsnorm_forward (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_fp64_sum (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_kv_lora (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_llama_production (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_muse_centered_pytorch_bf16_storage (const float input, const float zero_centered_weight, float output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_muse_unweighted_pytorch_bf16_storage (const float input, float output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_muse_weighted_pytorch_bf16_storage (const float input, const float weight, float output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_no_weight (const float input, float output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_no_weight_llama_production (const float input, float output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

void	rmsnorm_forward_pytorch_bf16_storage (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

static void	rmsnorm_forward_pytorch_bf16_storage_impl (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int input_stride, int output_stride, float eps, int weight_mode, int use_pow_inverse)

void	rmsnorm_forward_qwen3next_pytorch_bf16_storage (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)

static void	rmsnorm_forward_strict_scalar (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int input_stride, int output_stride, float eps)

void	rmsnorm_forward_strided_f32 (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int input_stride, int output_stride, float eps)

void	rmsnorm_forward_strided_pytorch_bf16_storage (const float input, const float gamma, float output, float rstd_cache, int tokens, int d_model, int input_stride, int output_stride, float eps)

static float	rmsnorm_llama_production_rstd (float mean_eps)

Detailed Description

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

CK-ENGINE KERNEL RULES:

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

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

RMSNorm: y[i] = gamma[i] * x[i] / sqrt(mean(x^2) + eps)

Definition in file rmsnorm_kernels.c.

Function Documentation

◆ ck_rmsnorm_provider_name()

const char * ck_rmsnorm_provider_name ( void )

Definition at line 859 of file rmsnorm_kernels.c.

{
    const char *exact_env = getenv("CK_RMSNORM_EXACT");
    if (ck_strict_parity_enabled() || (exact_env && atoi(exact_env) != 0)) {
        return "rmsnorm_strict_scalar";
    }
#if defined(__AVX512F__)
    return "rmsnorm_avx512";
#elif defined(__AVX__)
    return "rmsnorm_avx";
#elif defined(__ARM_NEON) || defined(__aarch64__)
    return "rmsnorm_neon";
#else
    return "rmsnorm_scalar";
#endif
}

References ck_strict_parity_enabled().

◆ gemma4_v_norm_forward()

void gemma4_v_norm_forward	(	const float *	input,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	num_kv_heads,
		int	head_dim,
		float	eps
	)

Definition at line 839 of file rmsnorm_kernels.c.

{
    if (!input || !output || tokens <= 0 || num_kv_heads <= 0 || head_dim <= 0) {
        return;
    }
    rmsnorm_forward_no_weight_llama_production(
        input, output, rstd_cache,
        tokens * num_kv_heads, head_dim, head_dim, eps);
}

References rmsnorm_forward_no_weight_llama_production().

◆ rmsnorm_backward()

void rmsnorm_backward	(	const float *	d_output,
		const float *	input,
		const float *	gamma,
		const float *	rstd_cache,
		float *	d_input,
		float *	d_gamma,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim
	)

RMSNorm backward pass

Test:

test_rmsnorm.py::TestRMSNormBackward::test_backward_tokens

test_rmsnorm.py::TestRMSNormBackward::test_backward_single

test_parity.py::test_rmsnorm_backward_parity

Computes dX and dGamma given dY, X, gamma, and cached rstd. dX_i = rstd * (dY_i * gamma_i - x_hat_i * m) dGamma_i = sum_t (dY_i * x_hat_i)

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

Definition at line 888 of file rmsnorm_kernels.c.

{
    int T = tokens;
    int D = d_model;
    int aligned = aligned_embed_dim;
 
    if (ck_strict_parity_enabled()) {
        rmsnorm_backward_strict_scalar(d_output, input, gamma, rstd_cache, d_input, d_gamma, T, D, aligned);
        return;
    }
 
    // Zero parameter gradients
#if defined(__AVX512F__)
    {
        int d = 0;
        for (; d + 16 <= D; d += 16) {
            _mm512_storeu_ps(&d_gamma[d], _mm512_setzero_ps());
        }
        for (; d < D; ++d) {
            d_gamma[d] = 0.0f;
        }
    }
#elif defined(__AVX__)
    {
        int d = 0;
        for (; d + 8 <= D; d += 8) {
            _mm256_storeu_ps(&d_gamma[d], _mm256_setzero_ps());
        }
        for (; d < D; ++d) {
            d_gamma[d] = 0.0f;
        }
    }
#else
    for (int d = 0; d < D; ++d) {
        d_gamma[d] = 0.0f;
    }
#endif
 
    for (int t = 0; t < T; ++t) {
        const float *x = input + (size_t)t * aligned;
        const float *dY = d_output + (size_t)t * aligned;
        float *dX = d_input + (size_t)t * aligned;
 
        float rstd = rstd_cache[t];
 
#if defined(__AVX512F__)
        // Compute m = (1/D) * sum_j (dY_j * gamma_j * x_hat_j)
        __m512 rstd_vec = _mm512_set1_ps(rstd);
        __m512 sum_vec = _mm512_setzero_ps();
        int d = 0;
 
        for (; d + 16 <= D; d += 16) {
            __m512 xv = _mm512_loadu_ps(&x[d]);
            __m512 dyv = _mm512_loadu_ps(&dY[d]);
            __m512 gv = _mm512_loadu_ps(&gamma[d]);
            __m512 x_hat = _mm512_mul_ps(xv, rstd_vec);
            // sum += dY * gamma * x_hat
            __m512 prod = _mm512_mul_ps(dyv, gv);
            sum_vec = _mm512_fmadd_ps(prod, x_hat, sum_vec);
        }
        float sum_dY_g_xhat = _mm512_reduce_add_ps(sum_vec);
 
        // Handle remaining elements
        for (; d < D; ++d) {
            float x_hat = x[d] * rstd;
            sum_dY_g_xhat += dY[d] * gamma[d] * x_hat;
        }
        float m = sum_dY_g_xhat / (float)D;
 
        // Compute dX and accumulate dGamma (vectorized)
        __m512 m_vec = _mm512_set1_ps(m);
        d = 0;
        for (; d + 16 <= D; d += 16) {
            __m512 xv = _mm512_loadu_ps(&x[d]);
            __m512 dyv = _mm512_loadu_ps(&dY[d]);
            __m512 gv = _mm512_loadu_ps(&gamma[d]);
            __m512 dgv = _mm512_loadu_ps(&d_gamma[d]);
 
            __m512 x_hat = _mm512_mul_ps(xv, rstd_vec);
 
            // dX = rstd * (dY * gamma - x_hat * m)
            __m512 dy_g = _mm512_mul_ps(dyv, gv);
            __m512 xhat_m = _mm512_mul_ps(x_hat, m_vec);
            __m512 diff = _mm512_sub_ps(dy_g, xhat_m);
            __m512 dxv = _mm512_mul_ps(rstd_vec, diff);
            _mm512_storeu_ps(&dX[d], dxv);
 
            // d_gamma += dY * x_hat
            dgv = _mm512_fmadd_ps(dyv, x_hat, dgv);
            _mm512_storeu_ps(&d_gamma[d], dgv);
        }
        // Handle remaining elements
        for (; d < D; ++d) {
            float x_hat = x[d] * rstd;
            float dy = dY[d];
            dX[d] = rstd * (dy * gamma[d] - x_hat * m);
            d_gamma[d] += dy * x_hat;
        }
 
#elif defined(__AVX__)
        // Compute m = (1/D) * sum_j (dY_j * gamma_j * x_hat_j)
        __m256 rstd_vec = _mm256_set1_ps(rstd);
        __m256 sum_vec = _mm256_setzero_ps();
        int d = 0;
 
        for (; d + 8 <= D; d += 8) {
            __m256 xv = _mm256_loadu_ps(&x[d]);
            __m256 dyv = _mm256_loadu_ps(&dY[d]);
            __m256 gv = _mm256_loadu_ps(&gamma[d]);
            __m256 x_hat = _mm256_mul_ps(xv, rstd_vec);
            // sum += dY * gamma * x_hat (no FMA, use mul + mul + add)
            __m256 prod = _mm256_mul_ps(dyv, gv);
            __m256 prod2 = _mm256_mul_ps(prod, x_hat);
            sum_vec = _mm256_add_ps(sum_vec, prod2);
        }
        float sum_dY_g_xhat = hsum256_ps_rmsnorm(sum_vec);
 
        // Handle remaining elements
        for (; d < D; ++d) {
            float x_hat = x[d] * rstd;
            sum_dY_g_xhat += dY[d] * gamma[d] * x_hat;
        }
        float m = sum_dY_g_xhat / (float)D;
 
        // Compute dX and accumulate dGamma (vectorized)
        __m256 m_vec = _mm256_set1_ps(m);
        d = 0;
        for (; d + 8 <= D; d += 8) {
            __m256 xv = _mm256_loadu_ps(&x[d]);
            __m256 dyv = _mm256_loadu_ps(&dY[d]);
            __m256 gv = _mm256_loadu_ps(&gamma[d]);
            __m256 dgv = _mm256_loadu_ps(&d_gamma[d]);
 
            __m256 x_hat = _mm256_mul_ps(xv, rstd_vec);
 
            // dX = rstd * (dY * gamma - x_hat * m)
            __m256 dy_g = _mm256_mul_ps(dyv, gv);
            __m256 xhat_m = _mm256_mul_ps(x_hat, m_vec);
            __m256 diff = _mm256_sub_ps(dy_g, xhat_m);
            __m256 dxv = _mm256_mul_ps(rstd_vec, diff);
            _mm256_storeu_ps(&dX[d], dxv);
 
            // d_gamma += dY * x_hat
            __m256 dy_xhat = _mm256_mul_ps(dyv, x_hat);
            dgv = _mm256_add_ps(dgv, dy_xhat);
            _mm256_storeu_ps(&d_gamma[d], dgv);
        }
        // Handle remaining elements
        for (; d < D; ++d) {
            float x_hat = x[d] * rstd;
            float dy = dY[d];
            dX[d] = rstd * (dy * gamma[d] - x_hat * m);
            d_gamma[d] += dy * x_hat;
        }
 
#else
        // Scalar fallback
        // Compute m = (1/D) * sum_j (dY_j * gamma_j * x_hat_j)
        float sum_dY_g_xhat = 0.0f;
        for (int d = 0; d < D; ++d) {
            float x_hat = x[d] * rstd;
            sum_dY_g_xhat += dY[d] * gamma[d] * x_hat;
        }
        float m = sum_dY_g_xhat / (float)D;
 
        // Compute dX and accumulate dGamma
        for (int d = 0; d < D; ++d) {
            float x_hat = x[d] * rstd;
            float dy = dY[d];
            dX[d] = rstd * (dy * gamma[d] - x_hat * m);
            d_gamma[d] += dy * x_hat;
        }
#endif
 
        // Zero padding gradients (if any)
        for (int d = D; d < aligned; ++d) {
            dX[d] = 0.0f;
        }
    }
}

References ck_strict_parity_enabled(), and rmsnorm_backward_strict_scalar().

Referenced by ck_layer_backward_rmsnorm_swiglu(), qk_norm_backward(), rmsnorm_backward_int4(), and rmsnorm_backward_int8().

◆ rmsnorm_backward_strict_scalar()

static void rmsnorm_backward_strict_scalar	(	const float *	d_output,
		const float *	input,
		const float *	gamma,
		const float *	rstd_cache,
		float *	d_input,
		float *	d_gamma,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim
	)

static

Definition at line 514 of file rmsnorm_kernels.c.

{
    const float inv_d = 1.0f / (float)d_model;
    for (int d = 0; d < d_model; ++d) {
        d_gamma[d] = 0.0f;
    }
 
    for (int t = 0; t < tokens; ++t) {
        const float *x = input + (size_t)t * (size_t)aligned_embed_dim;
        const float *dY = d_output + (size_t)t * (size_t)aligned_embed_dim;
        float *dX = d_input + (size_t)t * (size_t)aligned_embed_dim;
        const float rstd = rstd_cache[t];
 
        float sum_dY_g_xhat = 0.0f;
        for (int d = 0; d < d_model; ++d) {
            const float x_hat = x[d] * rstd;
            const float grad_x_hat = dY[d] * gamma[d];
            sum_dY_g_xhat += x_hat * grad_x_hat;
        }
 
        for (int d = 0; d < d_model; ++d) {
            const float x_hat = x[d] * rstd;
            const float grad_x_hat = dY[d] * gamma[d];
            dX[d] = (grad_x_hat - (x_hat * inv_d) * sum_dY_g_xhat) * rstd;
            d_gamma[d] += dY[d] * x_hat;
        }
        for (int d = d_model; d < aligned_embed_dim; ++d) {
            dX[d] = 0.0f;
        }
    }
}

Referenced by rmsnorm_backward().

◆ rmsnorm_forward()

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

Definition at line 735 of file rmsnorm_kernels.c.

{
    rmsnorm_forward_strided_f32(
        input,
        gamma,
        output,
        rstd_cache,
        tokens,
        d_model,
        aligned_embed_dim,
        aligned_embed_dim,
        eps
    );
}

References rmsnorm_forward_strided_f32().

◆ rmsnorm_forward_fp64_sum()

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

Definition at line 154 of file rmsnorm_kernels.c.

{
    for (int t = 0; t < tokens; ++t) {
        const float *x = input + (size_t)t * (size_t)aligned_embed_dim;
        float *y = output + (size_t)t * (size_t)aligned_embed_dim;
        /* This provider's contract requires an ascending scalar reduction.
         * Keep the accumulator volatile so whole-program optimization cannot
         * reassociate the sum or replace it with SIMD partial reductions. */
        volatile double sum_sq = 0.0;
        for (int d = 0; d < d_model; ++d) {
            const float square = x[d] * x[d];
            sum_sq = sum_sq + (double)square;
        }
        const float mean_sq = (float)(sum_sq / (double)d_model);
        const float rstd = 1.0f / sqrtf(mean_sq + eps);
        if (rstd_cache) {
            rstd_cache[t] = rstd;
        }
        for (int d = 0; d < d_model; ++d) {
            const float normalized = x[d] * rstd;
            y[d] = normalized * gamma[d];
        }
        for (int d = d_model; d < aligned_embed_dim; ++d) {
            y[d] = 0.0f;
        }
    }
}

Referenced by qk_norm_forward_fp64_sum().

◆ rmsnorm_forward_kv_lora()

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

Definition at line 757 of file rmsnorm_kernels.c.

{
    rmsnorm_forward(input, gamma, output, rstd_cache, tokens, d_model, aligned_embed_dim, eps);
}

References rmsnorm_forward().

◆ rmsnorm_forward_llama_production()

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

Definition at line 213 of file rmsnorm_kernels.c.

{
    for (int t = 0; t < tokens; ++t) {
        const float *x = input + (size_t)t * (size_t)aligned_embed_dim;
        float *y = output + (size_t)t * (size_t)aligned_embed_dim;
        volatile double sum_sq = 0.0;
        for (int d = 0; d < d_model; ++d) {
            const float square = x[d] * x[d];
            sum_sq = sum_sq + (double)square;
        }
        const float mean_sq = (float)(sum_sq / (double)d_model);
        const float rstd = rmsnorm_llama_production_rstd(mean_sq + eps);
        if (rstd_cache) {
            rstd_cache[t] = rstd;
        }
        for (int d = 0; d < d_model; ++d) {
            /*
             * Keep the RMSNorm + scale expression fused at the source level.
             * llama.cpp's CPU graph fuses GGML_OP_RMS_NORM followed by
             * GGML_OP_MUL and evaluates this left-associative expression in
             * one kernel.  Materializing the normalized value as a named
             * float introduces a store/load rounding boundary under ICX and
             * differs by one ULP for otherwise identical inputs.
             */
            y[d] = x[d] * rstd * gamma[d];
        }
        for (int d = d_model; d < aligned_embed_dim; ++d) {
            y[d] = 0.0f;
        }
    }
}

References rmsnorm_llama_production_rstd().

Referenced by ck_gemma4_embed_range(), qk_norm_forward_llama_production(), qwen4_group_rmsnorm_llama(), recurrent_norm_gate_llama_avx2_forward(), and recurrent_norm_sigmoid_gate_llama_avx2_forward().

◆ rmsnorm_forward_muse_centered_pytorch_bf16_storage()

void rmsnorm_forward_muse_centered_pytorch_bf16_storage	(	const float *	input,
		const float *	zero_centered_weight,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 470 of file rmsnorm_kernels.c.

{
    rmsnorm_forward_pytorch_bf16_storage_impl(
        input, zero_centered_weight, output, rstd_cache, tokens, d_model,
        aligned_embed_dim, aligned_embed_dim, eps, 2, 0);
}

References rmsnorm_forward_pytorch_bf16_storage_impl().

◆ rmsnorm_forward_muse_unweighted_pytorch_bf16_storage()

void rmsnorm_forward_muse_unweighted_pytorch_bf16_storage	(	const float *	input,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 500 of file rmsnorm_kernels.c.

{
    rmsnorm_forward_pytorch_bf16_storage_impl(
        input, NULL, output, rstd_cache, tokens, d_model,
        aligned_embed_dim, aligned_embed_dim, eps, 3, 0);
}

References rmsnorm_forward_pytorch_bf16_storage_impl().

Referenced by qk_norm_forward_muse_unweighted_scaled_pytorch_bf16_storage().

◆ rmsnorm_forward_muse_weighted_pytorch_bf16_storage()

void rmsnorm_forward_muse_weighted_pytorch_bf16_storage	(	const float *	input,
		const float *	weight,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 485 of file rmsnorm_kernels.c.

{
    rmsnorm_forward_pytorch_bf16_storage_impl(
        input, weight, output, rstd_cache, tokens, d_model,
        aligned_embed_dim, aligned_embed_dim, eps, 1, 0);
}

References rmsnorm_forward_pytorch_bf16_storage_impl().

◆ rmsnorm_forward_no_weight()

void rmsnorm_forward_no_weight	(	const float *	input,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 769 of file rmsnorm_kernels.c.

{
    if (!input || !output || tokens <= 0 || d_model <= 0 || aligned_embed_dim <= 0) {
        return;
    }
    const float inv_d = 1.0f / (float)d_model;
    for (int t = 0; t < tokens; ++t) {
        const float *x = input + (size_t)t * (size_t)aligned_embed_dim;
        float *y = output + (size_t)t * (size_t)aligned_embed_dim;
        double sum_sq = 0.0;
        for (int d = 0; d < d_model; ++d) {
            sum_sq += (double)x[d] * (double)x[d];
        }
        const float rstd = 1.0f / sqrtf((float)(sum_sq * (double)inv_d) + eps);
        if (rstd_cache) {
            rstd_cache[t] = rstd;
        }
        for (int d = 0; d < d_model; ++d) {
            y[d] = x[d] * rstd;
        }
        for (int d = d_model; d < aligned_embed_dim; ++d) {
            y[d] = 0.0f;
        }
    }
}

◆ rmsnorm_forward_no_weight_llama_production()

void rmsnorm_forward_no_weight_llama_production	(	const float *	input,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 806 of file rmsnorm_kernels.c.

{
    if (!input || !output || tokens <= 0 || d_model <= 0 || aligned_embed_dim < d_model) {
        return;
    }
    for (int t = 0; t < tokens; ++t) {
        const float *x = input + (size_t)t * (size_t)aligned_embed_dim;
        float *y = output + (size_t)t * (size_t)aligned_embed_dim;
        volatile double sum_sq = 0.0;
        for (int d = 0; d < d_model; ++d) {
            const float square = x[d] * x[d];
            sum_sq = sum_sq + (double)square;
        }
        const float mean_sq = (float)(sum_sq / (double)d_model);
        const float rstd = rmsnorm_llama_production_rstd(mean_sq + eps);
        if (rstd_cache) {
            rstd_cache[t] = rstd;
        }
        for (int d = 0; d < d_model; ++d) {
            y[d] = x[d] * rstd;
        }
        for (int d = d_model; d < aligned_embed_dim; ++d) {
            y[d] = 0.0f;
        }
    }
}

References rmsnorm_llama_production_rstd().

Referenced by gemma4_v_norm_forward().

◆ rmsnorm_forward_pytorch_bf16_storage()

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

Definition at line 426 of file rmsnorm_kernels.c.

{
    rmsnorm_forward_pytorch_bf16_storage_impl(
        input, gamma, output, rstd_cache, tokens, d_model,
        aligned_embed_dim, aligned_embed_dim, eps, 0, 0);
}

References rmsnorm_forward_pytorch_bf16_storage_impl().

Referenced by qk_norm_forward_pytorch_bf16_storage(), recurrent_norm_gate_pytorch_bf16_storage(), and recurrent_norm_sigmoid_gate_pytorch_bf16_storage().

◆ rmsnorm_forward_pytorch_bf16_storage_impl()

static void rmsnorm_forward_pytorch_bf16_storage_impl	(	const float *	input,
		const float *	gamma,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	input_stride,
		int	output_stride,
		float	eps,
		int	weight_mode,
		int	use_pow_inverse
	)

static

Definition at line 271 of file rmsnorm_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;
        float sum_sq = 0.0f;
 
#if defined(__AVX2__)
        /* Match the installed ATen build's AVX2 floating-point cascade_sum
         * contract. It treats the
         * contiguous row as four interleaved vector streams, accumulates
         * 16-item blocks through four hierarchy levels, then folds the four
         * streams and vector lanes left-to-right. PyTorch materializes pow(2)
         * before mean(), so keep the multiply separate from accumulation.
         * The host supports AVX-512, but this PyTorch build has AVX2 reduction
         * providers only; provider ISA is part of the numerical contract. */
        __m256 level[4][4];
        for (int hierarchy = 0; hierarchy < 4; ++hierarchy) {
            for (int stream = 0; stream < 4; ++stream) {
                level[hierarchy][stream] = _mm256_setzero_ps();
            }
        }
        int d = 0;
        const int vector_count = d_model / 8;
        const int cascade_items = vector_count / 4;
        int level_power = 4;
        if (cascade_items > 1) {
            int ceil_log2 = 0;
            unsigned int value = (unsigned int)(cascade_items - 1);
            while (value != 0) {
                value >>= 1;
                ++ceil_log2;
            }
            const int candidate = ceil_log2 / 4;
            if (candidate > level_power) level_power = candidate;
        }
        const int level_step = 1 << level_power;
        const int level_mask = level_step - 1;
        int item = 0;
        for (; item + level_step <= cascade_items;) {
            for (int block = 0; block < level_step; ++block, ++item) {
                for (int stream = 0; stream < 4; ++stream) {
                    const int offset = (item * 4 + stream) * 8;
                    const __m256 values = rmsnorm_load_bf16_values_avx2(x + offset);
                    const __m256 squared = rmsnorm_square_avx2_no_contract(values);
                    level[0][stream] = rmsnorm_add_avx2_ordered(
                        level[0][stream], squared
                    );
                }
            }
            for (int hierarchy = 1; hierarchy < 4; ++hierarchy) {
                for (int stream = 0; stream < 4; ++stream) {
                    level[hierarchy][stream] = rmsnorm_add_avx2_ordered(
                        level[hierarchy][stream], level[hierarchy - 1][stream]
                    );
                    level[hierarchy - 1][stream] = _mm256_setzero_ps();
                }
                const int mask = level_mask << (hierarchy * level_power);
                if ((item & mask) != 0) break;
            }
        }
        for (; item < cascade_items; ++item) {
            for (int stream = 0; stream < 4; ++stream) {
                const int offset = (item * 4 + stream) * 8;
                const __m256 values = rmsnorm_load_bf16_values_avx2(x + offset);
                const __m256 squared = rmsnorm_square_avx2_no_contract(values);
                level[0][stream] = rmsnorm_add_avx2_ordered(
                    level[0][stream], squared
                );
            }
        }
        for (int hierarchy = 1; hierarchy < 4; ++hierarchy) {
            for (int stream = 0; stream < 4; ++stream) {
                level[0][stream] = rmsnorm_add_avx2_ordered(
                    level[0][stream], level[hierarchy][stream]
                );
            }
        }
        __m256 reduced = level[0][0];
        for (int stream = 1; stream < 4; ++stream) {
            reduced = rmsnorm_add_avx2_ordered(reduced, level[0][stream]);
        }
        d = cascade_items * 4 * 8;
        for (; d + 8 <= d_model; d += 8) {
            const __m256 values = rmsnorm_load_bf16_values_avx2(x + d);
            const __m256 squared = rmsnorm_square_avx2_no_contract(values);
            reduced = rmsnorm_add_avx2_ordered(reduced, squared);
        }
        _Alignas(32) float lanes[8];
        _mm256_store_ps(lanes, reduced);
        volatile float ordered_sum = 0.0f;
        for (int lane = 0; lane < 8; ++lane) {
            ordered_sum = ordered_sum + lanes[lane];
        }
        sum_sq = ordered_sum;
        for (; d < d_model; ++d) {
            const float value = bf16_to_float(float_to_bf16(x[d]));
            sum_sq += value * value;
        }
#else
        for (int d = 0; d < d_model; ++d) {
            const float value = bf16_to_float(float_to_bf16(x[d]));
            sum_sq += value * value;
        }
#endif
 
#if defined(__i386__) || defined(__x86_64__)
        const float variance = rmsnorm_div_f32_ordered(sum_sq, (float)d_model);
        const float rstd = use_pow_inverse
            ? powf(variance + eps, -0.5f)
            : rmsnorm_div_f32_ordered(1.0f, sqrtf(variance + eps));
#else
        const float variance = sum_sq / (float)d_model;
        const float rstd = use_pow_inverse
            ? powf(variance + eps, -0.5f)
            : 1.0f / sqrtf(variance + eps);
#endif
        if (rstd_cache) rstd_cache[t] = rstd;
        for (int d = 0; d < d_model; ++d) {
            const float value = bf16_to_float(float_to_bf16(x[d]));
            if (weight_mode == 3) {
                /* Muse unweighted normalization rounds only when the FP32
                 * normalized value returns to the BF16 hidden-state dtype. */
                y[d] = bf16_to_float(float_to_bf16(value * rstd));
            } else if (weight_mode == 2) {
                /* Keep the checkpoint weight zero-centered. Form 1 + weight
                 * in FP32 at execution instead of pre-folding into BF16. */
                y[d] = bf16_to_float(float_to_bf16(
                    (value * rstd) * (1.0f + gamma[d])));
            } else if (weight_mode == 1) {
                /* Qwen3Next: (FP32 normalized * FP32 weight).to(BF16). */
                y[d] = bf16_to_float(float_to_bf16(
                    (value * rstd) * gamma[d]));
            } else {
                const float weight =
                    bf16_to_float(float_to_bf16(gamma[d]));
                const float normalized =
                    bf16_to_float(float_to_bf16(value * rstd));
                y[d] = bf16_to_float(float_to_bf16(normalized * weight));
            }
        }
        for (int d = d_model; d < output_stride; ++d) y[d] = 0.0f;
    }
}

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

Referenced by rmsnorm_forward_muse_centered_pytorch_bf16_storage(), rmsnorm_forward_muse_unweighted_pytorch_bf16_storage(), rmsnorm_forward_muse_weighted_pytorch_bf16_storage(), rmsnorm_forward_pytorch_bf16_storage(), rmsnorm_forward_qwen3next_pytorch_bf16_storage(), and rmsnorm_forward_strided_pytorch_bf16_storage().

◆ rmsnorm_forward_qwen3next_pytorch_bf16_storage()

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

Definition at line 455 of file rmsnorm_kernels.c.

{
    rmsnorm_forward_pytorch_bf16_storage_impl(
        input, gamma, output, rstd_cache, tokens, d_model,
        aligned_embed_dim, aligned_embed_dim, eps, 1, 0);
}

References rmsnorm_forward_pytorch_bf16_storage_impl().

Referenced by hyper_connection_mix_bf16(), qk_norm_forward_qwen4_pytorch_bf16_storage(), and qwen4_group_rmsnorm_pytorch_bf16().

◆ rmsnorm_forward_strict_scalar()

static void rmsnorm_forward_strict_scalar	(	const float *	input,
		const float *	gamma,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	input_stride,
		int	output_stride,
		float	eps
	)

static

Definition at line 113 of file rmsnorm_kernels.c.

{
    const float inv_d = 1.0f / (float)d_model;
    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;
 
        float sum_sq = 0.0f;
        for (int d = 0; d < d_model; ++d) {
            const float v = x[d];
            sum_sq += v * v;
        }
        const float mean_sq = sum_sq * inv_d;
        const float rstd = 1.0f / sqrtf(mean_sq + eps);
        if (rstd_cache) {
            rstd_cache[t] = rstd;
        }
 
        for (int d = 0; d < d_model; ++d) {
            const float x_hat = x[d] * rstd;
            y[d] = x_hat * gamma[d];
        }
        for (int d = d_model; d < output_stride; ++d) {
            y[d] = 0.0f;
        }
    }
}

Referenced by rmsnorm_forward_strided_f32().

◆ rmsnorm_forward_strided_f32()

void rmsnorm_forward_strided_f32	(	const float *	input,
		const float *	gamma,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	input_stride,
		int	output_stride,
		float	eps
	)

RMSNorm forward pass

Test:

test_rmsnorm.py::TestRMSNormForward::test_fp32_tokens

test_rmsnorm.py::TestRMSNormForward::test_fp32_single

test_rmsnorm.py::TestRMSNormForward::test_perf_rolled

test_layernorm.py::TestLayerNormForward::test_rmsnorm_compat

test_parity.py::test_rmsnorm_parity

RMSNorm: y[i] = gamma[i] * x[i] / sqrt(mean(x^2) + eps)

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

Definition at line 567 of file rmsnorm_kernels.c.

{
    int T = tokens;
    int D = d_model;
 
    const char *exact_env = getenv("CK_RMSNORM_EXACT");
    if (ck_strict_parity_enabled() || (exact_env && atoi(exact_env) != 0)) {
        rmsnorm_forward_strict_scalar(
            input, gamma, output, rstd_cache, T, D, input_stride, output_stride, eps
        );
        return;
    }
 
    for (int t = 0; t < T; ++t) {
        const float *x = input + (size_t)t * (size_t)input_stride;
        float *y = output + (size_t)t * (size_t)output_stride;
 
#if defined(__AVX512F__)
        // AVX-512: Process 16 floats at a time
        __m512 sum_sq_vec = _mm512_setzero_ps();
        int d = 0;
 
        // Vectorized sum of squares
        for (; d + 16 <= D; d += 16) {
            __m512 xv = _mm512_loadu_ps(&x[d]);
            sum_sq_vec = _mm512_fmadd_ps(xv, xv, sum_sq_vec);
        }
        float sum_sq = _mm512_reduce_add_ps(sum_sq_vec);
 
        // Handle remaining elements
        for (; d < D; ++d) {
            sum_sq += x[d] * x[d];
        }
 
        float mean_sq = sum_sq / (float)D;
        float rstd = 1.0f / sqrtf(mean_sq + eps);
        if (rstd_cache) {
            rstd_cache[t] = rstd;
        }
 
        // Apply normalization and scale (vectorized)
        __m512 rstd_vec = _mm512_set1_ps(rstd);
        d = 0;
        for (; d + 16 <= D; d += 16) {
            __m512 xv = _mm512_loadu_ps(&x[d]);
            __m512 gv = _mm512_loadu_ps(&gamma[d]);
            __m512 x_hat = _mm512_mul_ps(xv, rstd_vec);
            __m512 yv = _mm512_mul_ps(x_hat, gv);
            _mm512_storeu_ps(&y[d], yv);
        }
        // Handle remaining elements
        for (; d < D; ++d) {
            y[d] = x[d] * rstd * gamma[d];
        }
 
#elif defined(__AVX__)
        // AVX: Process 8 floats at a time
        __m256 sum_sq_vec = _mm256_setzero_ps();
        int d = 0;
 
        // Vectorized sum of squares (no FMA in AVX1, use mul + add)
        for (; d + 8 <= D; d += 8) {
            __m256 xv = _mm256_loadu_ps(&x[d]);
            __m256 xv_sq = _mm256_mul_ps(xv, xv);
            sum_sq_vec = _mm256_add_ps(sum_sq_vec, xv_sq);
        }
        float sum_sq = hsum256_ps_rmsnorm(sum_sq_vec);
 
        // Handle remaining elements
        for (; d < D; ++d) {
            sum_sq += x[d] * x[d];
        }
 
        float mean_sq = sum_sq / (float)D;
        float rstd = 1.0f / sqrtf(mean_sq + eps);
        if (rstd_cache) {
            rstd_cache[t] = rstd;
        }
 
        // Apply normalization and scale (vectorized)
        __m256 rstd_vec = _mm256_set1_ps(rstd);
        d = 0;
        for (; d + 8 <= D; d += 8) {
            __m256 xv = _mm256_loadu_ps(&x[d]);
            __m256 gv = _mm256_loadu_ps(&gamma[d]);
            __m256 x_hat = _mm256_mul_ps(xv, rstd_vec);
            __m256 yv = _mm256_mul_ps(x_hat, gv);
            _mm256_storeu_ps(&y[d], yv);
        }
        // Handle remaining elements
        for (; d < D; ++d) {
            y[d] = x[d] * rstd * gamma[d];
        }
 
#elif defined(__ARM_NEON) || defined(__aarch64__)
        /* NEON: 4 floats at a time. The sum-of-squares is reassociated
         * through four lanes (FMA accumulation, like the AVX-512 branch),
         * so results differ from the ascending scalar order by FP32
         * reassociation noise only; CK_RMSNORM_EXACT / strict parity mode
         * above still selects the scalar reference. */
        float32x4_t sum_sq_vec = vdupq_n_f32(0.0f);
        int d = 0;
 
        for (; d + 4 <= D; d += 4) {
            const float32x4_t xv = vld1q_f32(&x[d]);
            sum_sq_vec = vfmaq_f32(sum_sq_vec, xv, xv);
        }
        float sum_sq = hsum128_ps_rmsnorm(sum_sq_vec);
 
        for (; d < D; ++d) {
            sum_sq += x[d] * x[d];
        }
 
        float mean_sq = sum_sq / (float)D;
        float rstd = 1.0f / sqrtf(mean_sq + eps);
        if (rstd_cache) {
            rstd_cache[t] = rstd;
        }
 
        const float32x4_t rstd_vec = vdupq_n_f32(rstd);
        d = 0;
        for (; d + 4 <= D; d += 4) {
            const float32x4_t xv = vld1q_f32(&x[d]);
            const float32x4_t gv = vld1q_f32(&gamma[d]);
            /* Keep the two separate rounding boundaries of the scalar tail:
             * x_hat = x * rstd, then y = x_hat * gamma (no FMA fusion). */
            const float32x4_t x_hat = vmulq_f32(xv, rstd_vec);
            vst1q_f32(&y[d], vmulq_f32(x_hat, gv));
        }
        for (; d < D; ++d) {
            y[d] = x[d] * rstd * gamma[d];
        }
 
#else
        // Scalar fallback
        float sum_sq = 0.0f;
        for (int d = 0; d < D; ++d) {
            float v = x[d];
            sum_sq += v * v;
        }
        float mean_sq = sum_sq / (float)D;
        float rstd = 1.0f / sqrtf(mean_sq + eps);
        if (rstd_cache) {
            rstd_cache[t] = rstd;
        }
 
        // Apply normalization and scale
        for (int d = 0; d < D; ++d) {
            float x_hat = x[d] * rstd;
            y[d] = x_hat * gamma[d];
        }
#endif
 
        // Zero padding (if any)
        for (int d = D; d < output_stride; ++d) {
            y[d] = 0.0f;
        }
    }
}

References ck_strict_parity_enabled(), and rmsnorm_forward_strict_scalar().

Referenced by rmsnorm_forward().

◆ rmsnorm_forward_strided_pytorch_bf16_storage()

void rmsnorm_forward_strided_pytorch_bf16_storage	(	const float *	input,
		const float *	gamma,
		float *	output,
		float *	rstd_cache,
		int	tokens,
		int	d_model,
		int	input_stride,
		int	output_stride,
		float	eps
	)

Definition at line 440 of file rmsnorm_kernels.c.

{
    rmsnorm_forward_pytorch_bf16_storage_impl(
        input, gamma, output, rstd_cache, tokens, d_model,
        input_stride, output_stride, eps, 0, 0);
}

References rmsnorm_forward_pytorch_bf16_storage_impl().

◆ rmsnorm_llama_production_rstd()

static float rmsnorm_llama_production_rstd ( float mean_eps )

inlinestatic

Definition at line 189 of file rmsnorm_kernels.c.

{
    /*
     * ggml's production CPU RMSNorm emits scalar sqrt followed by scalar
     * division.  With -fno-math-errno ICX otherwise strength-reduces the C
     * expression to vrsqrt14ss plus one Newton step.  That estimate differs by
     * one ULP for some rows and the error is amplified by quantized
     * projections in deep recurrent models.
     */
#if defined(CK_TARGET_X86)
    const __m128 value = _mm_set_ss(mean_eps);
    const __m128 root = _mm_sqrt_ss(value);
    return _mm_cvtss_f32(_mm_div_ss(_mm_set_ss(1.0f), root));
#else
    const volatile float root = sqrtf(mean_eps);
    return 1.0f / root;
#endif
}

Referenced by rmsnorm_forward_llama_production(), and rmsnorm_forward_no_weight_llama_production().

Functions

Detailed Description

CK-ENGINE KERNEL RULES:

Function Documentation

◆ ck_rmsnorm_provider_name()

◆ gemma4_v_norm_forward()

◆ rmsnorm_backward()

◆ rmsnorm_backward_strict_scalar()

◆ rmsnorm_forward()

◆ rmsnorm_forward_fp64_sum()

◆ rmsnorm_forward_kv_lora()

◆ rmsnorm_forward_llama_production()

◆ rmsnorm_forward_muse_centered_pytorch_bf16_storage()

◆ rmsnorm_forward_muse_unweighted_pytorch_bf16_storage()

◆ rmsnorm_forward_muse_weighted_pytorch_bf16_storage()

◆ rmsnorm_forward_no_weight()

◆ rmsnorm_forward_no_weight_llama_production()

◆ rmsnorm_forward_pytorch_bf16_storage()

◆ rmsnorm_forward_pytorch_bf16_storage_impl()

◆ rmsnorm_forward_qwen3next_pytorch_bf16_storage()

◆ rmsnorm_forward_strict_scalar()

◆ rmsnorm_forward_strided_f32()

◆ rmsnorm_forward_strided_pytorch_bf16_storage()

◆ rmsnorm_llama_production_rstd()