gemm_kernels_q5_0_sse_v2.c File Reference

SSE-optimized GEMM kernels for Q5_0 x Q8_K quantization. More...

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

Go to the source code of this file.

Functions
static float	dot_q5_0_q8_k_32_sse (const block_q5_0 *bw, const block_q8_K *ba, int q8_offset)
void	gemm_nt_q5_0_ref (const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
	GEMM with transposed Q5_0 weights: C = A @ B^T.
void	gemm_nt_q5_0_sse_v2 (const float *A, const void *B, const float *bias, float *C, int M, int N, int K)
void	quantize_row_q8_k (const float *x, void *vy, int k)

Detailed Description

SSE-optimized GEMM kernels for Q5_0 x Q8_K quantization.

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

Definition in file gemm_kernels_q5_0_sse_v2.c.

Function Documentation

dot_q5_0_q8_k_32_sse()

static float dot_q5_0_q8_k_32_sse ( const block_q5_0 *  bw, const block_q8_K *  ba, int  q8_offset  )

inlinestatic

Definition at line 25 of file gemm_kernels_q5_0_sse_v2.c.

                                                                                                    {
    const uint8_t *qs_w = bw->qs;
    const int8_t  *qs_a = ba->qs + q8_offset;
    
    uint32_t qh;
    memcpy(&qh, bw->qh, sizeof(qh));
 
    // Vectorize bit extraction
    // Load 16 bytes of low nibbles
    __m128i qs_vec = _mm_loadu_si128((const __m128i *)qs_w);
    __m128i mask_0f = _mm_set1_epi8(0x0F);
    
    __m128i w_lo = _mm_and_si128(qs_vec, mask_0f);
    __m128i w_hi = _mm_and_si128(_mm_srli_epi16(qs_vec, 4), mask_0f);
 
    // Load high bits from qh
    // This is still a bit scalar but we can use shuffle for some of it if needed
    // For now, let's just make sure we handle the -16 offset correctly in SIMD.
    
    uint8_t w[32];
    for (int j = 0; j < 16; j++) {
        w[j]    = (qs_w[j] & 0x0F) | (((qh >> (j + 0)) << 4) & 0x10);
        w[j+16] = (qs_w[j] >> 4)   | ((qh >> (j + 12)) & 0x10);
    }
 
    __m128i vw0 = _mm_loadu_si128((const __m128i *)&w[0]);
    __m128i vw1 = _mm_loadu_si128((const __m128i *)&w[16]);
    __m128i va0 = _mm_loadu_si128((const __m128i *)&qs_a[0]);
    __m128i va1 = _mm_loadu_si128((const __m128i *)&qs_a[16]);
 
    // Dot product: unsigned 8-bit * signed 8-bit -> signed 16-bit
    __m128i p0 = _mm_maddubs_epi16(vw0, va0);
    __m128i p1 = _mm_maddubs_epi16(vw1, va1);
 
    // Sum to i32
    __m128i one = _mm_set1_epi16(1);
    __m128i s0 = _mm_madd_epi16(p0, one);
    __m128i s1 = _mm_madd_epi16(p1, one);
    __m128i acc_i32 = _mm_add_epi32(s0, s1);
 
    // Horizontal sum of i32
    acc_i32 = _mm_add_epi32(acc_i32, _mm_shuffle_epi32(acc_i32, _MM_SHUFFLE(1, 0, 3, 2)));
    acc_i32 = _mm_add_epi32(acc_i32, _mm_shuffle_epi32(acc_i32, _MM_SHUFFLE(0, 1, 0, 1)));
    int32_t dot_wa = _mm_cvtsi128_si32(acc_i32);
 
    // sum((w - 16) * a) = sum(w*a) - 16 * sum(a)
    int32_t sum_a = (int32_t)ba->bsums[q8_offset/16] + (int32_t)ba->bsums[q8_offset/16 + 1];
    
    float result = ((float)dot_wa - 16.0f * (float)sum_a) * CK_FP16_TO_FP32(bw->d) * ba->d;
    return result;
}

References block_q8_K::bsums, CK_FP16_TO_FP32, block_q5_0::d, block_q8_K::d, block_q5_0::qh, block_q8_K::qs, and block_q5_0::qs.

Referenced by gemm_nt_q5_0_sse_v2().

gemm_nt_q5_0_ref()

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

GEMM with transposed Q5_0 weights: C = A @ B^T.

Parameters

A	Input activations [M x K], row-major FP32
B	Weight matrix in Q5_0 format [N x K], row-major quantized
bias	Optional bias [N], NULL if not used
C	Output [M x N], row-major FP32
M	Batch size (number of tokens)
N	Output dimension (number of rows in B)
K	Input dimension

Definition at line 790 of file gemm_kernels_q5_0.c.

{
    const block_q5_0 *blocks = (const block_q5_0 *)B;
    const int blocks_per_row = K / QK5_0;
 
    for (int m = 0; m < M; m++) {
        const float *a_row = &A[m * K];
 
        for (int n = 0; n < N; n++) {
            float sum = 0.0f;
 
            for (int b = 0; b < blocks_per_row; b++) {
                const block_q5_0 *block = &blocks[n * blocks_per_row + b];
                const float d = CK_FP16_TO_FP32(block->d);
                const float *ap = &a_row[b * QK5_0];
 
                uint32_t qh;
                memcpy(&qh, block->qh, sizeof(qh));
 
                /* llama.cpp Q5_0 layout - note j+12 for second weight high bit */
                for (int j = 0; j < QK5_0 / 2; j++) {
                    const uint8_t packed = block->qs[j];
                    const int lo = (packed & 0x0F);
                    const int hi = (packed >> 4);
                    const int xh_0 = ((qh >> (j + 0)) << 4) & 0x10;
                    const int xh_1 = ((qh >> (j + 12))) & 0x10;
                    const int q0 = (lo | xh_0) - 16;
                    const int q1 = (hi | xh_1) - 16;
 
                    sum += d * (float)q0 * ap[j];
                    sum += d * (float)q1 * ap[j + 16];
                }
            }
 
            C[m * N + n] = sum + (bias ? bias[n] : 0.0f);
        }
    }
}

References C, CK_FP16_TO_FP32, block_q5_0::d, block_q5_0::qh, QK5_0, and block_q5_0::qs.

Referenced by gemm_nt_q5_0_sse_v2().

gemm_nt_q5_0_sse_v2()

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

Definition at line 77 of file gemm_kernels_q5_0_sse_v2.c.

{
    if (K % QK_K != 0) {
        gemm_nt_q5_0_ref(A, B, bias, C, M, N, K);
        return;
    }
 
    size_t q8_size = (K / QK_K) * sizeof(block_q8_K);
    block_q8_K *A_q8 = (block_q8_K *)alloca(q8_size);
 
    const block_q5_0 *weights = (const block_q5_0 *)B;
    const int blocks_per_row = K / 32;
 
    for (int m = 0; m < M; m++) {
        quantize_row_q8_k(&A[m * K], A_q8, K);
 
        for (int n = 0; n < N; n++) {
            float sumf = 0.0f;
            const block_q5_0 *w_row = weights + n * blocks_per_row;
 
            for (int b = 0; b < blocks_per_row; b++) {
                int q8_block_idx = (b * 32) / QK_K;
                int q8_offset = (b * 32) % QK_K;
                sumf += dot_q5_0_q8_k_32_sse(&w_row[b], &A_q8[q8_block_idx], q8_offset);
            }
 
            C[m * N + n] = sumf + (bias ? bias[n] : 0.0f);
        }
    }
}

References C, dot_q5_0_q8_k_32_sse(), gemm_nt_q5_0_ref(), QK_K, and quantize_row_q8_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 gemm_nt_q5_0_sse_v2().