GEMM/GEMV kernels with Q4_K quantized weights. More...

#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include <stdlib.h>
#include "ckernel_quant.h"

Macros
#define	CK_Q4K_STACK_Q8_BLOCKS 128

Functions
static int	ck_q4k_debug_q8_contract (void)

float	dot_q4_k (const void w_q4k, const float x, int K)
	Compute dot product of Q4_K row with FP32 vector.

void	gemm_nt_q4_k (const float A, const void B, const float bias, float C, int M, int N, int K)

void	gemm_q4_k (float Y, const void W, const float *X, int M, int N, int K)
	Auto-dispatch GEMM based on available SIMD.

void	gemm_q4_k_backward (float dX, const void W, const float *dY, int M, int N, int K)
	Batched backward pass.

void	gemm_q4_k_ref (float Y, const void W, const float *X, int M, int N, int K)
	Matrix-matrix multiply with Q4_K weights (scalar reference)

void	gemv_q4_k (float y, const void W, const float *x, int M, int K)
	Auto-dispatch GEMV based on available SIMD.

void	gemv_q4_k_backward (float dX, const void W, const float *dY, int M, int K)
	Auto-dispatch backward.

void	gemv_q4_k_backward_ref (float dX, const void W, const float *dY, int M, int K)
	Backward pass: compute input gradient (scalar reference)

void	gemv_q4_k_q8_k (float y, const void W, const void *x_q8, int M, int K)

void	gemv_q4_k_ref (float y, const void W, const float *x, int M, int K)
	Matrix-vector multiply with Q4_K weights (scalar reference)

void	quantize_row_q8_k (const float x, void vy, int k)

Detailed Description

GEMM/GEMV kernels with Q4_K quantized weights.

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

Implements matrix multiplication where:

Activations (input): FP32
Weights: Q4_K (4.5 bits/weight, nested scales)
Output: FP32

Key optimization: Fused dequantization - weights are dequantized in registers and immediately used in FMA, never written to memory.

Operations:

gemv_q4_k: Matrix-vector multiply (batch=1, token generation)
gemm_q4_k: Matrix-matrix multiply (batch>1, prefill)

Definition in file gemm_kernels_q4k.c.

Macro Definition Documentation

◆ CK_Q4K_STACK_Q8_BLOCKS

#define CK_Q4K_STACK_Q8_BLOCKS 128

Definition at line 38 of file gemm_kernels_q4k.c.

Function Documentation

◆ ck_q4k_debug_q8_contract()

static int ck_q4k_debug_q8_contract ( void )

static

Definition at line 43 of file gemm_kernels_q4k.c.

{
    static int cached = -1;
    if (cached < 0) {
        const char *env = getenv("CK_DEBUG_Q4K_Q8_CONTRACT");
        cached = (env && env[0] && env[0] != '0') ? 1 : 0;
    }
    return cached;
}

Referenced by gemv_q4_k().

◆ dot_q4_k()

float dot_q4_k	(	const void *	w_q4k,
		const float *	x,
		int	K
	)

Compute dot product of Q4_K row with FP32 vector.

Parameters

w_q4k	Q4_K blocks for one row
x	FP32 input vector
K	Vector length (must be multiple of 256)

Returns: Dot product result

Definition at line 509 of file gemm_kernels_q4k.c.

{
    float result;
    gemv_q4_k(&result, w_q4k, x, 1, K);
    return result;
}

References gemv_q4_k().

◆ gemm_nt_q4_k()

void gemm_nt_q4_k	(	const float *	A,
		const void *	B,
		const float *	bias,
		float *	C,
		int	M,
		int	N,
		int	K
	)

Definition at line 708 of file gemm_kernels_q4k.c.

{
    if (!A || !B || !C) {
        return;
    }
    if (M <= 0 || N <= 0 || K <= 0) {
        return;
    }
 
    /* gemm_q4_k produces Y as [batch x M_out]. Here:
     *   batch = M (tokens)
     *   M_out = N (output channels) */
    gemm_q4_k(C, B, A, /*M_out=*/N, /*N_batch=*/M, K);
 
    if (!bias) {
        return;
    }
 
    for (int i = 0; i < M; ++i) {
        float *row = C + (size_t)i * (size_t)N;
        for (int j = 0; j < N; ++j) {
            row[j] += bias[j];
        }
    }
}

References C, and gemm_q4_k().

◆ gemm_q4_k()

void gemm_q4_k	(	float *	Y,
		const void *	W,
		const float *	X,
		int	M,
		int	N,
		int	K
	)

Auto-dispatch GEMM based on available SIMD.

Definition at line 486 of file gemm_kernels_q4k.c.

{
    /* Use reference implementation for correctness
     * TODO: Fix AVX-512 version to match llama.cpp layout */
    gemm_q4_k_ref(Y, W, X, M, N, K);
}

References gemm_q4_k_ref().

Referenced by gemm_nt_q4_k().

◆ gemm_q4_k_backward()

void gemm_q4_k_backward	(	float *	dX,
		const void *	W,
		const float *	dY,
		int	M,
		int	N,
		int	K
	)

Batched backward pass.

Definition at line 681 of file gemm_kernels_q4k.c.

{
    for (int n = 0; n < N; n++) {
        gemv_q4_k_backward(&dX[n * K], W, &dY[n * M], M, K);
    }
}

References gemv_q4_k_backward().

◆ gemm_q4_k_ref()

void gemm_q4_k_ref	(	float *	Y,
		const void *	W,
		const float *	X,
		int	M,
		int	N,
		int	K
	)

Matrix-matrix multiply with Q4_K weights (scalar reference)

Parameters

Y	Output matrix [M x N]
W	Weight matrix in Q4_K format [M x K]
X	Input matrix [K x N] (column-major for cache efficiency)
M	Number of output rows
N	Batch size (number of columns)
K	Hidden dimension

Definition at line 341 of file gemm_kernels_q4k.c.

{
    /* For each column in batch, use the dispatching gemv_q4_k
     * which automatically selects AVX/AVX-512/scalar based on CPU */
    for (int n = 0; n < N; n++) {
        gemv_q4_k(&Y[n * M], W, &X[n * K], M, K);
    }
}

References gemv_q4_k().

Referenced by gemm_q4_k().

◆ gemv_q4_k()

void gemv_q4_k	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Auto-dispatch GEMV based on available SIMD.

Definition at line 301 of file gemm_kernels_q4k.c.

{
    if (ck_q4k_debug_q8_contract() && K > 0 && (K % QK_K) == 0) {
        const int nb = K / QK_K;
        if (nb <= CK_Q4K_STACK_Q8_BLOCKS) {
            block_q8_K x_q8[CK_Q4K_STACK_Q8_BLOCKS];
            quantize_row_q8_k(x, x_q8, K);
            gemv_q4_k_q8_k(y, W, x_q8, M, K);
            return;
        }
    }
#ifdef __AVX512F__
    gemv_q4_k_avx512(y, W, x, M, K);
#elif defined(__AVX__)
    gemv_q4_k_avx(y, W, x, M, K);
#else
    gemv_q4_k_ref(y, W, x, M, K);
#endif
}

References ck_q4k_debug_q8_contract(), CK_Q4K_STACK_Q8_BLOCKS, gemv_q4_k_q8_k(), gemv_q4_k_ref(), QK_K, and quantize_row_q8_k().

Referenced by attention_mlp_fused_q4k(), dot_q4_k(), gemm_q4_k_ref(), layer_fused_attn_mlp_qkv_q4k(), and rmsnorm_qkv_q4k_fused().

◆ gemv_q4_k_backward()

void gemv_q4_k_backward	(	float *	dX,
		const void *	W,
		const float *	dY,
		int	M,
		int	K
	)

Auto-dispatch backward.

Definition at line 666 of file gemm_kernels_q4k.c.

{
#ifdef __AVX512F__
    gemv_q4_k_backward_avx512(dX, W, dY, M, K);
#else
    gemv_q4_k_backward_ref(dX, W, dY, M, K);
#endif
}

References gemv_q4_k_backward_ref().

Referenced by gemm_q4_k_backward().

◆ gemv_q4_k_backward_ref()

void gemv_q4_k_backward_ref	(	float *	dX,
		const void *	W,
		const float *	dY,
		int	M,
		int	K
	)

Backward pass: compute input gradient (scalar reference)

Parameters

dX	Output gradient w.r.t. input [K]
W	Weight matrix in Q4_K format [M x K]
dY	Gradient w.r.t. output [M]
M	Number of output rows
K	Number of columns (input dimension)

Definition at line 536 of file gemm_kernels_q4k.c.

{
    const block_q4_K *blocks = (const block_q4_K *)W;
    const int blocks_per_row = K / QK_K;
 
    /* Zero output gradient */
    memset(dX, 0, K * sizeof(float));
 
    /* Accumulate: dX += W^T @ dY
     * Uses llama.cpp layout: 4 iterations of 64 weights each */
    for (int row = 0; row < M; row++) {
        const float dy = dY[row];
 
        for (int b = 0; b < blocks_per_row; b++) {
            const block_q4_K *block = &blocks[row * blocks_per_row + b];
            const float d = CK_FP16_TO_FP32(block->d);
            const float dmin = CK_FP16_TO_FP32(block->dmin);
 
            uint8_t sc[8], m[8];
            unpack_q4_k_scales(block->scales, sc, m);
 
            /* llama.cpp layout: 4 iterations of 64 weights each */
            for (int iter = 0; iter < 4; iter++) {
                const float d1 = d * (float)sc[2 * iter];
                const float m1 = dmin * (float)m[2 * iter];
                const float d2 = d * (float)sc[2 * iter + 1];
                const float m2 = dmin * (float)m[2 * iter + 1];
 
                const uint8_t *qs = &block->qs[iter * 32];
                float *dxp = &dX[b * QK_K + iter * 64];
 
                /* First 32 weights: low nibbles */
                for (int l = 0; l < 32; l++) {
                    const int q = (qs[l] & 0x0F);
                    const float w = d1 * (float)q - m1;
                    dxp[l] += w * dy;
                }
 
                /* Next 32 weights: high nibbles */
                for (int l = 0; l < 32; l++) {
                    const int q = (qs[l] >> 4);
                    const float w = d2 * (float)q - m2;
                    dxp[32 + l] += w * dy;
                }
            }
        }
    }
}

References CK_FP16_TO_FP32, block_q4_K::d, block_q4_K::dmin, QK_K, block_q4_K::qs, block_q4_K::scales, and unpack_q4_k_scales().

Referenced by gemv_q4_k_backward().

◆ gemv_q4_k_q8_k()

void gemv_q4_k_q8_k	(	float *	y,
		const void *	W,
		const void *	x_q8,
		int	M,
		int	K
	)

Definition at line 273 of file gemm_kernels_q4k_q8k.c.

{
    if (ck_q4k_q8k_force_ref()) {
        gemv_q4_k_q8_k_ref(y, W, x_q8, M, K);
        return;
    }
#if defined(__AVX512VNNI__) && defined(__AVX512VL__) && !defined(CK_NO_AVX512_VNNI)
    /* VNNI: Best for decode (single token) - INT8 dot product acceleration */
    gemv_q4_k_q8_k_vnni(y, W, x_q8, M, K);
#elif defined(__AVX2__)
    gemv_q4_k_q8_k_avx2(y, W, x_q8, M, K);
#elif defined(__AVX__)
    /* AVX version uses maddubs_epi16 (more efficient than SSE) */
    gemv_q4_k_q8_k_avx(y, W, x_q8, M, K);
#elif defined(__SSE4_1__)
    gemv_q4_k_q8_k_sse(y, W, x_q8, M, K);
#else
    gemv_q4_k_q8_k_ref(y, W, x_q8, M, K);
#endif
}

Referenced by gemm_q4_k_q8_k(), gemm_q4_k_q8_k_thread_fn(), and gemv_q4_k().

◆ gemv_q4_k_ref()

void gemv_q4_k_ref	(	float *	y,
		const void *	W,
		const float *	x,
		int	M,
		int	K
	)

Matrix-vector multiply with Q4_K weights (scalar reference)

Parameters

y	Output vector [M]
W	Weight matrix in Q4_K format [M x K], stored row-major
x	Input vector [K]
M	Number of output rows
K	Number of columns (must be multiple of 256)

Definition at line 69 of file gemm_kernels_q4k.c.

{
    const block_q4_K *blocks = (const block_q4_K *)W;
    const int blocks_per_row = K / QK_K;  /* QK_K = 256 */
 
    for (int row = 0; row < M; row++) {
        float sum = 0.0f;
 
        for (int b = 0; b < blocks_per_row; b++) {
            const block_q4_K *block = &blocks[row * blocks_per_row + b];
            const float d = GGML_FP16_TO_FP32(block->d);
            const float dmin = GGML_FP16_TO_FP32(block->dmin);
 
            /* Unpack sub-block scales */
            uint8_t sc[8], m[8];
            unpack_q4_k_scales(block->scales, sc, m);
 
            /* llama.cpp Q4_K layout: 4 iterations of 64 weights each
             * Each iteration uses 32 bytes of qs and 2 scales:
             *   - First 32 weights (indices 0-31): low nibbles with scale[2*iter]
             *   - Next 32 weights (indices 32-63): high nibbles with scale[2*iter+1]
             */
            for (int iter = 0; iter < 4; iter++) {
                const float d1 = d * (float)sc[2*iter];
                const float m1 = dmin * (float)m[2*iter];
                const float d2 = d * (float)sc[2*iter + 1];
                const float m2 = dmin * (float)m[2*iter + 1];
                const uint8_t *qs = &block->qs[iter * 32];
                const float *xp = &x[b * QK_K + iter * 64];
 
                /* First 32 weights: low nibbles of qs[0..31] */
                for (int l = 0; l < 32; l++) {
                    const int8_t q = (qs[l] & 0x0F);
                    sum += (d1 * (float)q - m1) * xp[l];
                }
                /* Next 32 weights: high nibbles of qs[0..31] */
                for (int l = 0; l < 32; l++) {
                    const int8_t q = (qs[l] >> 4);
                    sum += (d2 * (float)q - m2) * xp[l + 32];
                }
            }
        }
 
        y[row] = sum;
    }
}

References block_q4_K::d, block_q4_K::dmin, GGML_FP16_TO_FP32, QK_K, block_q4_K::qs, block_q4_K::scales, and unpack_q4_k_scales().

Referenced by gemv_q4_k().

◆ quantize_row_q8_k()

void quantize_row_q8_k	(	const float *	x,
		void *	vy,
		int	k
	)

Definition at line 121 of file gemm_kernels_q4k_q8k.c.

                                                        {
    const char *ref_env = getenv("CK_DEBUG_Q8K_REF");
    if (ref_env && atoi(ref_env) != 0) {
        quantize_row_q8_k_ref(x, vy, k);
        return;
    }
#if defined(__AVX512F__) && defined(__AVX512BW__)
    quantize_row_q8_k_avx512(x, vy, k);
#elif defined(__AVX2__)
    quantize_row_q8_k_avx2(x, vy, k);
#elif defined(__AVX__)
    quantize_row_q8_k_avx(x, vy, k);
#elif defined(__SSE4_1__)
    quantize_row_q8_k_sse(x, vy, k);
#else
    quantize_row_q8_k_ref(x, vy, k);
#endif
}

Referenced by gemv_q4_k().

Macros

Functions

Detailed Description

CK-ENGINE KERNEL RULES:

Macro Definition Documentation

◆ CK_Q4K_STACK_Q8_BLOCKS

Function Documentation

◆ ck_q4k_debug_q8_contract()

◆ dot_q4_k()

◆ gemm_nt_q4_k()

◆ gemm_q4_k()

◆ gemm_q4_k_backward()

◆ gemm_q4_k_ref()

◆ gemv_q4_k()

◆ gemv_q4_k_backward()

◆ gemv_q4_k_backward_ref()

◆ gemv_q4_k_q8_k()

◆ gemv_q4_k_ref()

◆ quantize_row_q8_k()