qk_norm_kernels.c File Reference

Per-head RMSNorm on Q and K (Qwen3-style QK norm) More...

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

Go to the source code of this file.

Enumerations
enum	QKNormISA {   QK_NORM_ISA_SCALAR = 0 , QK_NORM_ISA_AVX = 1 , QK_NORM_ISA_AVX2 = 2 , QK_NORM_ISA_AVX_VNNI = 3 ,   QK_NORM_ISA_AUTO = -1 }

Functions
void	q_norm_forward (float *q, const float *q_gamma, int num_heads, int num_tokens, int head_dim, float eps)
void	qk_norm_backward (const float *d_q_out, const float *d_k_out, const float *q_in, const float *k_in, const float *q_gamma, const float *k_gamma, float *d_q_in, float *d_k_in, float *d_q_gamma, float *d_k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)
int	qk_norm_backward_last_isa (void)
static void	qk_norm_compute_rstd (const float *input, float *rstd_cache, int rows, int head_dim, float eps)
static void	qk_norm_compute_rstd_scalar (const float *input, float *rstd_cache, int rows, int head_dim, float eps)
void	qk_norm_forward (float *q, float *k, const float *q_gamma, const float *k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)
void	qk_norm_forward_fp64_sum (float *q, float *k, const float *q_gamma, const float *k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)
void	qk_norm_forward_llama_production (float *q, float *k, const float *q_gamma, const float *k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)
void	qk_norm_forward_muse_unweighted_scaled_pytorch_bf16_storage (float *q, float *k, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps, float q_scale)
void	qk_norm_forward_pytorch_bf16_storage (float *q, float *k, const float *q_gamma, const float *k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)
void	qk_norm_forward_qwen4_pytorch_bf16_storage (float *q, float *k, const float *q_gamma, const float *k_gamma, int num_heads, int num_kv_heads, int num_tokens, int head_dim, float eps)
static int	qk_norm_isa_compiled (QKNormISA isa)
static QKNormISA	qk_norm_parse_forced_isa (void)
static QKNormISA	qk_norm_select_isa (void)
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)
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_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_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_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)
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)

Variables
static int	g_qk_norm_last_isa = QK_NORM_ISA_SCALAR

Detailed Description

Per-head RMSNorm on Q and K (Qwen3-style QK norm)

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 v7-qk-norm-backward-parity-isa python unittest/test_qk_norm.py

QK Norm normalizes each head's query/key vectors independently before RoPE. This stabilizes Q*K^T dot products before softmax, preventing attention collapse from large magnitude vectors.

Why only Q and K, not V? V does not participate in the attention score computation (Q*K^T). The softmax saturation problem comes from large Q*K^T values, so only Q and K magnitudes matter. V is linearly combined after softmax weights are computed – normalizing it would change output scale but not fix attention stability.

Data layout after QKV projection (head-major): Q: [num_heads, num_tokens, head_dim] contiguous K: [num_kv_heads, num_tokens, head_dim] contiguous

We treat Q as [num_heads * num_tokens] rows of [head_dim] elements. rmsnorm_forward normalizes each row independently. The gamma weight [head_dim] is shared across all heads (Qwen3 design: one gamma per Q, one per K).

Definition in file qk_norm_kernels.c.

Enumeration Type Documentation

QKNormISA

enum QKNormISA

Enumerator
QK_NORM_ISA_SCALAR
QK_NORM_ISA_AVX
QK_NORM_ISA_AVX2
QK_NORM_ISA_AVX_VNNI
QK_NORM_ISA_AUTO

Definition at line 110 of file qk_norm_kernels.c.

             {
    QK_NORM_ISA_SCALAR = 0,
    QK_NORM_ISA_AVX = 1,
    QK_NORM_ISA_AVX2 = 2,
    QK_NORM_ISA_AVX_VNNI = 3,
    QK_NORM_ISA_AUTO = -1
} QKNormISA;

Function Documentation

q_norm_forward()

void q_norm_forward	(	float *	q,
		const float *	q_gamma,
		int	num_heads,
		int	num_tokens,
		int	head_dim,
		float	eps
	)

Forward pass for Gemma4-assistant q-only per-head RMSNorm.

Some Gemma4 assistant/drafter checkpoints project only Q and then reuse Q as the shared K/V stream. This wrapper keeps that public kernel contract explicit while reusing the same row-wise RMSNorm implementation as qk_norm_forward.

Definition at line 434 of file qk_norm_kernels.c.

{
    rmsnorm_forward(q, q_gamma, q, NULL,
                    num_heads * num_tokens, head_dim, head_dim, eps);
}

References rmsnorm_forward().

qk_norm_backward()

void qk_norm_backward	(	const float *	d_q_out,
		const float *	d_k_out,
		const float *	q_in,
		const float *	k_in,
		const float *	q_gamma,
		const float *	k_gamma,
		float *	d_q_in,
		float *	d_k_in,
		float *	d_q_gamma,
		float *	d_k_gamma,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		float	eps
	)

Backward pass for per-head QK RMSNorm.

This computes:

• d_q / d_k for the Q and K activations
• d_q_gamma / d_k_gamma for shared per-head gamma vectors

Implementation is reference-first and deterministic: 1) recompute row rstd values from saved q/k inputs 2) call rmsnorm_backward on flattened [rows, head_dim] views

Definition at line 456 of file qk_norm_kernels.c.

{
    int q_rows = num_heads * num_tokens;
    int k_rows = num_kv_heads * num_tokens;
 
    if (q_rows > 0) {
        float q_rstd_cache[q_rows];
        qk_norm_compute_rstd(q_in, q_rstd_cache, q_rows, head_dim, eps);
        rmsnorm_backward(d_q_out, q_in, q_gamma, q_rstd_cache,
                         d_q_in, d_q_gamma, q_rows, head_dim, head_dim);
    }
 
    if (k_rows > 0) {
        float k_rstd_cache[k_rows];
        qk_norm_compute_rstd(k_in, k_rstd_cache, k_rows, head_dim, eps);
        rmsnorm_backward(d_k_out, k_in, k_gamma, k_rstd_cache,
                         d_k_in, d_k_gamma, k_rows, head_dim, head_dim);
    }
}

References qk_norm_compute_rstd(), and rmsnorm_backward().

qk_norm_backward_last_isa()

int qk_norm_backward_last_isa

(

void

)

Definition at line 120 of file qk_norm_kernels.c.

{
    return g_qk_norm_last_isa;
}

References g_qk_norm_last_isa.

qk_norm_compute_rstd()

static void qk_norm_compute_rstd ( const float *  input, float *  rstd_cache, int  rows, int  head_dim, float  eps  )

static

Definition at line 288 of file qk_norm_kernels.c.

{
    QKNormISA selected = qk_norm_select_isa();
    g_qk_norm_last_isa = (int)selected;
    switch (selected) {
#if defined(__AVXVNNI__)
    case QK_NORM_ISA_AVX_VNNI:
        qk_norm_compute_rstd_avx_vnni(input, rstd_cache, rows, head_dim, eps);
        return;
#endif
#if defined(__AVX2__)
    case QK_NORM_ISA_AVX2:
        qk_norm_compute_rstd_avx2(input, rstd_cache, rows, head_dim, eps);
        return;
#endif
#if defined(__AVX__)
    case QK_NORM_ISA_AVX:
        qk_norm_compute_rstd_avx(input, rstd_cache, rows, head_dim, eps);
        return;
#endif
    case QK_NORM_ISA_SCALAR:
    case QK_NORM_ISA_AUTO:
    default:
        qk_norm_compute_rstd_scalar(input, rstd_cache, rows, head_dim, eps);
        return;
    }
}

References g_qk_norm_last_isa, qk_norm_compute_rstd_scalar(), QK_NORM_ISA_AUTO, QK_NORM_ISA_AVX, QK_NORM_ISA_AVX2, QK_NORM_ISA_AVX_VNNI, QK_NORM_ISA_SCALAR, and qk_norm_select_isa().

Referenced by qk_norm_backward().

qk_norm_compute_rstd_scalar()

static void qk_norm_compute_rstd_scalar ( const float *  input, float *  rstd_cache, int  rows, int  head_dim, float  eps  )

static

Definition at line 187 of file qk_norm_kernels.c.

{
    for (int r = 0; r < rows; ++r) {
        const float *x = input + (size_t)r * (size_t)head_dim;
        double sum_sq = 0.0;
        for (int d = 0; d < head_dim; ++d) {
            double v = (double)x[d];
            sum_sq += v * v;
        }
        float mean_sq = (float)(sum_sq / (double)head_dim);
        rstd_cache[r] = 1.0f / sqrtf(mean_sq + eps);
    }
}

Referenced by qk_norm_compute_rstd().

qk_norm_forward()

void qk_norm_forward	(	float *	q,
		float *	k,
		const float *	q_gamma,
		const float *	k_gamma,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		float	eps
	)

Per-head RMSNorm on Q and K.

Parameters

q	Q scratch buffer [num_heads * num_tokens * head_dim], in-place
k	K scratch buffer [num_kv_heads * num_tokens * head_dim], in-place
q_gamma	Q norm gamma weights [head_dim]
k_gamma	K norm gamma weights [head_dim]
num_heads	Number of query heads (e.g. 32 for Qwen3-8B)
num_kv_heads	Number of KV heads (e.g. 8 for Qwen3-8B with GQA)
num_tokens	Number of tokens (1 for decode, T for prefill)
head_dim	Dimension per head (e.g. 128)
eps	RMSNorm epsilon (e.g. 1e-6)

Test:

unittest/test_qk_norm.py

Definition at line 335 of file qk_norm_kernels.c.

{
    /* Q norm: [num_heads * num_tokens] rows of [head_dim]
     * Each row is one head's vector for one token. */
    rmsnorm_forward(q, q_gamma, q, NULL,
                    num_heads * num_tokens, head_dim, head_dim, eps);
 
    /* K norm: [num_kv_heads * num_tokens] rows of [head_dim]
     * Same logic, fewer rows when using GQA. */
    rmsnorm_forward(k, k_gamma, k, NULL,
                    num_kv_heads * num_tokens, head_dim, head_dim, eps);
}

References rmsnorm_forward().

qk_norm_forward_fp64_sum()

void qk_norm_forward_fp64_sum	(	float *	q,
		float *	k,
		const float *	q_gamma,
		const float *	k_gamma,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		float	eps
	)

Definition at line 351 of file qk_norm_kernels.c.

{
    rmsnorm_forward_fp64_sum(q, q_gamma, q, NULL,
                             num_heads * num_tokens, head_dim, head_dim, eps);
    rmsnorm_forward_fp64_sum(k, k_gamma, k, NULL,
                             num_kv_heads * num_tokens, head_dim, head_dim, eps);
}

References rmsnorm_forward_fp64_sum().

qk_norm_forward_llama_production()

void qk_norm_forward_llama_production	(	float *	q,
		float *	k,
		const float *	q_gamma,
		const float *	k_gamma,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		float	eps
	)

Definition at line 362 of file qk_norm_kernels.c.

{
    rmsnorm_forward_llama_production(
        q, q_gamma, q, NULL,
        num_heads * num_tokens, head_dim, head_dim, eps);
    rmsnorm_forward_llama_production(
        k, k_gamma, k, NULL,
        num_kv_heads * num_tokens, head_dim, head_dim, eps);
}

References rmsnorm_forward_llama_production().

qk_norm_forward_muse_unweighted_scaled_pytorch_bf16_storage()

void qk_norm_forward_muse_unweighted_scaled_pytorch_bf16_storage	(	float *	q,
		float *	k,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		float	eps,
		float	q_scale
	)

Definition at line 407 of file qk_norm_kernels.c.

{
    const int q_elements = num_heads * num_tokens * head_dim;
    rmsnorm_forward_muse_unweighted_pytorch_bf16_storage(
        q, q, NULL, num_heads * num_tokens, head_dim, head_dim, eps);
    rmsnorm_forward_muse_unweighted_pytorch_bf16_storage(
        k, k, NULL, num_kv_heads * num_tokens, head_dim, head_dim, eps);
    for (int index = 0; index < q_elements; ++index) {
        q[index] = bf16_to_float(float_to_bf16(q[index] * q_scale));
    }
}

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

qk_norm_forward_pytorch_bf16_storage()

void qk_norm_forward_pytorch_bf16_storage	(	float *	q,
		float *	k,
		const float *	q_gamma,
		const float *	k_gamma,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		float	eps
	)

Definition at line 375 of file qk_norm_kernels.c.

{
    rmsnorm_forward_pytorch_bf16_storage(
        q, q_gamma, q, NULL,
        num_heads * num_tokens, head_dim, head_dim, eps);
    rmsnorm_forward_pytorch_bf16_storage(
        k, k_gamma, k, NULL,
        num_kv_heads * num_tokens, head_dim, head_dim, eps);
}

References rmsnorm_forward_pytorch_bf16_storage().

qk_norm_forward_qwen4_pytorch_bf16_storage()

void qk_norm_forward_qwen4_pytorch_bf16_storage	(	float *	q,
		float *	k,
		const float *	q_gamma,
		const float *	k_gamma,
		int	num_heads,
		int	num_kv_heads,
		int	num_tokens,
		int	head_dim,
		float	eps
	)

Definition at line 390 of file qk_norm_kernels.c.

{
    rmsnorm_forward_qwen3next_pytorch_bf16_storage(
        q, q_gamma, q, NULL,
        num_heads * num_tokens, head_dim, head_dim, eps);
    rmsnorm_forward_qwen3next_pytorch_bf16_storage(
        k, k_gamma, k, NULL,
        num_kv_heads * num_tokens, head_dim, head_dim, eps);
}

References rmsnorm_forward_qwen3next_pytorch_bf16_storage().

qk_norm_isa_compiled()

static int qk_norm_isa_compiled ( QKNormISA  isa )

static

Definition at line 147 of file qk_norm_kernels.c.

{
    switch (isa) {
    case QK_NORM_ISA_SCALAR:
        return 1;
#if defined(__AVX__)
    case QK_NORM_ISA_AVX:
        return 1;
#endif
#if defined(__AVX2__)
    case QK_NORM_ISA_AVX2:
        return 1;
#endif
#if defined(__AVXVNNI__)
    case QK_NORM_ISA_AVX_VNNI:
        return 1;
#endif
    default:
        return 0;
    }
}

References QK_NORM_ISA_AVX, QK_NORM_ISA_AVX2, QK_NORM_ISA_AVX_VNNI, and QK_NORM_ISA_SCALAR.

Referenced by qk_norm_select_isa().

qk_norm_parse_forced_isa()

static QKNormISA qk_norm_parse_forced_isa ( void  )

static

Definition at line 125 of file qk_norm_kernels.c.

{
    const char *forced = getenv("CK_QK_NORM_BACKWARD_ISA");
    if (!forced || forced[0] == '\0' || strcmp(forced, "auto") == 0) {
        return QK_NORM_ISA_AUTO;
    }
    if (strcmp(forced, "scalar") == 0) {
        return QK_NORM_ISA_SCALAR;
    }
    if (strcmp(forced, "avx") == 0) {
        return QK_NORM_ISA_AVX;
    }
    if (strcmp(forced, "avx2") == 0) {
        return QK_NORM_ISA_AVX2;
    }
    if (strcmp(forced, "avx_vnni") == 0) {
        return QK_NORM_ISA_AVX_VNNI;
    }
    /* Unknown value -> keep behavior deterministic by falling back. */
    return QK_NORM_ISA_SCALAR;
}

References QK_NORM_ISA_AUTO, QK_NORM_ISA_AVX, QK_NORM_ISA_AVX2, QK_NORM_ISA_AVX_VNNI, and QK_NORM_ISA_SCALAR.

Referenced by qk_norm_select_isa().

qk_norm_select_isa()

static QKNormISA qk_norm_select_isa ( void  )

static

Definition at line 169 of file qk_norm_kernels.c.

{
    QKNormISA forced = qk_norm_parse_forced_isa();
    if (forced != QK_NORM_ISA_AUTO) {
        return qk_norm_isa_compiled(forced) ? forced : QK_NORM_ISA_SCALAR;
    }
 
#if defined(__AVXVNNI__)
    return QK_NORM_ISA_AVX_VNNI;
#elif defined(__AVX2__)
    return QK_NORM_ISA_AVX2;
#elif defined(__AVX__)
    return QK_NORM_ISA_AVX;
#else
    return QK_NORM_ISA_SCALAR;
#endif
}

References QK_NORM_ISA_AUTO, QK_NORM_ISA_AVX, QK_NORM_ISA_AVX2, QK_NORM_ISA_AVX_VNNI, qk_norm_isa_compiled(), QK_NORM_ISA_SCALAR, and qk_norm_parse_forced_isa().

Referenced by qk_norm_compute_rstd().

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;
        }
    }
}

Referenced by qk_norm_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
    );
}

Referenced by q_norm_forward(), qk_norm_forward(), and rmsnorm_forward_kv_lora().

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_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;
        }
    }
}

Referenced by qk_norm_forward_llama_production().

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);
}

Referenced by qk_norm_forward_muse_unweighted_scaled_pytorch_bf16_storage().

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);
}

Referenced by qk_norm_forward_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);
}

Referenced by qk_norm_forward_qwen4_pytorch_bf16_storage().

Variable Documentation

g_qk_norm_last_isa

int g_qk_norm_last_isa = QK_NORM_ISA_SCALAR

static

Definition at line 118 of file qk_norm_kernels.c.

Referenced by qk_norm_backward_last_isa(), and qk_norm_compute_rstd().