rmsnorm_q8_k_fused.c File Reference

Fused RMSNorm + Q8_K Quantization kernel. More...

#include <immintrin.h>
#include <math.h>
#include <stdint.h>
#include <string.h>
#include "ckernel_quant.h"

Go to the source code of this file.

Functions
void	rmsnorm_q8_k_fused (const float *input, const float *gamma, void *vy, int tokens, int d_model, int aligned_embed_dim, float eps)

Detailed Description

Fused RMSNorm + Q8_K Quantization kernel.

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

FUSION BENEFIT: Eliminates intermediate FP32 buffer between RMSNorm and quantization, keeping normalized values in registers/L1.

Definition in file rmsnorm_q8_k_fused.c.

Function Documentation

rmsnorm_q8_k_fused()

void rmsnorm_q8_k_fused	(	const float *	input,
		const float *	gamma,
		void *	vy,
		int	tokens,
		int	d_model,
		int	aligned_embed_dim,
		float	eps
	)

Definition at line 144 of file rmsnorm_q8_k_fused.c.

{
    const int T = tokens;
    const int D = d_model;
    block_q8_K *y = (block_q8_K *)vy;
 
    for (int t = 0; t < T; ++t) {
        const float *x = input + (size_t)t * (size_t)aligned_embed_dim;
 
        float sum_sq = 0.0f;
        for (int d = 0; d < D; ++d) {
            sum_sq += x[d] * x[d];
        }
        const float rstd = 1.0f / sqrtf(sum_sq / (float)D + eps);
 
        for (int b = 0; b < D / QK_K; ++b) {
            const float *xb = x + b * QK_K;
            const float *gb = gamma + b * QK_K;
            block_q8_K *out_block = &y[t * (D / QK_K) + b];
 
            float norm_buf[QK_K];
            float max_val = 0.0f;
            for (int d = 0; d < QK_K; ++d) {
                const float normalized = xb[d] * rstd * gb[d];
                norm_buf[d] = normalized;
                const float abs_val = fabsf(normalized);
                if (abs_val > max_val) {
                    max_val = abs_val;
                }
            }
 
            if (max_val == 0.0f) {
                out_block->d = 0.0f;
                memset(out_block->qs, 0, QK_K);
                memset(out_block->bsums, 0, sizeof(out_block->bsums));
                continue;
            }
 
            const float iscale = -127.0f / max_val;
            out_block->d = 1.0f / iscale;
            for (int j = 0; j < QK_K; j += 16) {
                int bsum = 0;
                for (int k = 0; k < 16; ++k) {
                    int q = (int)lrintf(norm_buf[j + k] * iscale);
                    if (q < -128) q = -128;
                    if (q > 127) q = 127;
                    out_block->qs[j + k] = (int8_t)q;
                    bsum += q;
                }
                out_block->bsums[j / 16] = (int16_t)bsum;
            }
        }
    }
}

References block_q8_K::bsums, block_q8_K::d, QK_K, and block_q8_K::qs.