Batch GEMM kernels for quantized weights with INT8 activations. More...

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

Macros
#define	AMX_TILE_K 64

#define	AMX_TILE_M 16

#define	AMX_TILE_N 16

#define	HAS_AMX 0

#define	QK5_0 32 /* Q5_0: 32 weights per block */

#define	QK8_0 32 /* Q8_0: 32 weights per block */

Functions
const char *	gemm_batch_int8_impl_name (void)
	Get the best implementation name for logging/debugging. More...

void	gemm_nt_q5_0_q8_0_ref (const void A, const void B, float *C, int M, int N, int K)
	Dispatcher for gemm_nt_q8_0_q8_0. More...

void	gemm_nt_q8_0_q8_0 (const void A, const void B, const float bias, float C, int M, int N, int K)
	gemm_nt_q8_0_q8_0 with optional bias (matches header signature) More...

void	gemm_nt_q8_0_q8_0_ref (const void A, const void B, float *C, int M, int N, int K)
	Scalar reference: gemm_nt_q8_0_q8_0. More...

Detailed Description

Batch GEMM kernels for quantized weights with INT8 activations.

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 batch matrix multiplication where:

Activations (A): Q8_0 quantized (INT8 + scale)
Weights (B): Q5_0 or Q8_0 quantized
Output (C): FP32

Operation: C[M,N] = A[M,K] @ B[N,K]^T (B is transposed/row-major weights)

Instruction Set Implementations:

Scalar: Reference implementation for correctness verification
AVX: 256-bit SIMD (8 floats, or 32 int8s)
AVX-512: 512-bit SIMD (16 floats, or 64 int8s)
AMX: Intel Advanced Matrix Extensions (tile-based, requires Sapphire Rapids+)

Design Philosophy:

Every kernel MUST produce bit-identical results to scalar reference
Comprehensive testing against llama.cpp ensures correctness
Performance optimizations never compromise accuracy

Author: C-Kernel-Engine Team

Date: 2024

Definition in file gemm_batch_int8.c.

Macro Definition Documentation

◆ AMX_TILE_K

#define AMX_TILE_K 64

Definition at line 65 of file gemm_batch_int8.c.

◆ AMX_TILE_M

#define AMX_TILE_M 16

Definition at line 63 of file gemm_batch_int8.c.

◆ AMX_TILE_N

#define AMX_TILE_N 16

Definition at line 64 of file gemm_batch_int8.c.

◆ HAS_AMX

#define HAS_AMX 0

Definition at line 52 of file gemm_batch_int8.c.

◆ QK5_0

#define QK5_0 32 /* Q5_0: 32 weights per block */

Definition at line 60 of file gemm_batch_int8.c.

◆ QK8_0

#define QK8_0 32 /* Q8_0: 32 weights per block */

Definition at line 59 of file gemm_batch_int8.c.

Function Documentation

◆ gemm_batch_int8_impl_name()

const char* gemm_batch_int8_impl_name ( void )

Get the best implementation name for logging/debugging.

Definition at line 553 of file gemm_batch_int8.c.

 {
 #if HAS_AMX
     return "AMX";
 #elif defined(__AVX512VNNI__)
     return "AVX-512 VNNI";
 #elif defined(__AVX512F__)
     return "AVX-512";
 #elif defined(__AVX2__)
     return "AVX2";
 #elif defined(__AVX__)
     return "AVX";
 #else
     return "Scalar";
 #endif
 }

◆ gemm_nt_q5_0_q8_0_ref()

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

Dispatcher for gemm_nt_q8_0_q8_0.

Selects the best available implementation at runtime.

Scalar reference: gemm_nt_q5_0_q8_0

Q5_0 weight reconstruction: weight[j] = d * ((qs_nibble | (qh_bit << 4)) - 16)

For j in 0..15: use low nibble + qh bit j For j in 16..31: use high nibble + qh bit (j+16) -> actually bit (j) for j=16..31

Parameters

A	Input activations [M, K] in Q8_0 format
B	Weight matrix [N, K] in Q5_0 format
C	Output matrix [M, N] in FP32
M	Number of tokens (batch size)
N	Number of output features
K	Number of input features (must be multiple of 32)

Definition at line 391 of file gemm_batch_int8.c.

 {
     const int nb = K / QK5_0;
     const block_q8_0 *a_blocks = (const block_q8_0 *)A;
     const block_q5_0 *b_blocks = (const block_q5_0 *)B;
  
     for (int m = 0; m < M; m++) {
         const block_q8_0 *a_row = a_blocks + (size_t)m * nb;
  
         for (int n = 0; n < N; n++) {
             const block_q5_0 *b_row = b_blocks + (size_t)n * nb;
             float sum = 0.0f;
  
             for (int ib = 0; ib < nb; ib++) {
                 const float d_a = CK_FP16_TO_FP32(a_row[ib].d);
                 const float d_b = CK_FP16_TO_FP32(b_row[ib].d);
                 const float d = d_a * d_b;
  
                 /* Load high bits as 32-bit value */
                 uint32_t qh;
                 memcpy(&qh, b_row[ib].qh, sizeof(qh));
  
                 int32_t sumi = 0;
  
                 /* Process 32 weights: j=0..15 uses low nibble, j=16..31 uses high nibble */
                 for (int j = 0; j < 16; j++) {
                     /* First 16 weights: low nibble + qh bit j */
                     const uint8_t xh_0 = ((qh >> j) & 1) << 4;
                     const int8_t w0 = (int8_t)(((b_row[ib].qs[j] & 0x0F) | xh_0) - 16);
  
                     /* Second 16 weights: high nibble + qh bit (j+16) */
                     const uint8_t xh_1 = ((qh >> (j + 16)) & 1) << 4;
                     const int8_t w1 = (int8_t)(((b_row[ib].qs[j] >> 4) | xh_1) - 16);
  
                     /* Accumulate with activation values */
                     sumi += (int32_t)w0 * (int32_t)a_row[ib].qs[j];
                     sumi += (int32_t)w1 * (int32_t)a_row[ib].qs[j + 16];
                 }
  
                 sum += d * (float)sumi;
             }
  
             C[(size_t)m * N + n] = sum;
         }
     }
 }

References C, CK_FP16_TO_FP32, QK5_0, and block_q8_0::qs.

◆ gemm_nt_q8_0_q8_0()

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

gemm_nt_q8_0_q8_0 with optional bias (matches header signature)

C[m,n] = A[m,K] @ B[n,K]^T + bias[n]

Definition at line 582 of file gemm_batch_int8.c.

 {
     /* First compute GEMM */
 #if defined(__AVX512VNNI__)
     gemm_nt_q8_0_q8_0_vnni(A, B, C, M, N, K);
 #elif defined(__AVX512F__)
     gemm_nt_q8_0_q8_0_avx512(A, B, C, M, N, K);
 #elif defined(__AVX2__)
     gemm_nt_q8_0_q8_0_avx2(A, B, C, M, N, K);
 #elif defined(__AVX__)
     gemm_nt_q8_0_q8_0_avx(A, B, C, M, N, K);
 #else
     gemm_nt_q8_0_q8_0_ref(A, B, C, M, N, K);
 #endif
  
     /* Add bias if provided */
     if (bias != NULL) {
         for (int m = 0; m < M; m++) {
             for (int n = 0; n < N; n++) {
                 C[(size_t)m * N + n] += bias[n];
             }
         }
     }
 }

References C, and gemm_nt_q8_0_q8_0_ref().

Referenced by ck_test_gemm_q8_0(), gemm_nt_q8_0_dispatch(), and gemm_nt_q8_0_mlp_dispatch().

◆ gemm_nt_q8_0_q8_0_ref()

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

Scalar reference: gemm_nt_q8_0_q8_0.

C[m,n] = sum_k( dequant(A[m,k]) * dequant(B[n,k]) ) = sum_blocks( d_a * d_b * sum_j(a_qs[j] * b_qs[j]) )