gemm_kernels_amx.c File Reference

AMX (Advanced Matrix Extensions) GEMM kernels. More...

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

Go to the source code of this file.

Functions
bool	amx_available (void)
void	gemv_q4_k_q8_k_amx (float *y, const void *W, const void *x_q8, int M, int K)
void	gemv_q4_k_q8_k_avx (float *y, const void *W, const void *x_q8, int M, int K)
void	gemv_q4_k_q8_k_avx2 (float *y, const void *W, const void *x_q8, int M, int K)
void	gemv_q4_k_q8_k_ref (float *y, const void *W, const void *x_q8, int M, int K)
void	gemv_q4_k_q8_k_vnni (float *y, const void *W, const void *x_q8, int M, int K)

Detailed Description

AMX (Advanced Matrix Extensions) GEMM kernels.

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

Intel AMX provides dedicated matrix multiply hardware:

• 8 tile registers (TMM0-TMM7), each up to 1KB
• TDPBSSD: INT8 signed dot product (A signed, B signed)
• TDPBSUD: INT8 mixed sign (A signed, B unsigned)
• TDPBUSD: INT8 mixed sign (A unsigned, B signed)
• TDPBUUD: INT8 unsigned dot product
• TDPBF16PS: BF16 dot product to FP32

Tile dimensions:

• Max: 16 rows x 64 bytes (1024 bytes per tile)
• For INT8: 16x64 elements
• For BF16: 16x32 elements

Performance:

• AMX INT8: ~2000 INT8 ops/cycle (vs ~256 for AVX-512 VNNI)
• AMX BF16: ~1000 BF16 ops/cycle
• Expected 8-16x speedup over AVX-512 for large GEMM

Requirements:

• Sapphire Rapids or newer (4th Gen Xeon)
• Linux kernel 5.16+ with AMX support
• Compiler: GCC 11+, Clang 12+, ICX 2022+

Definition in file gemm_kernels_amx.c.

Function Documentation

amx_available()

bool amx_available

(

void

)

Definition at line 271 of file gemm_kernels_amx.c.

                         {
    return false;
}

gemv_q4_k_q8_k_amx()

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

Definition at line 256 of file gemm_kernels_amx.c.

                                                                                 {
    /* No AMX support - cascade through fallbacks: AVX-512 VNNI → AVX2 → AVX → ref */
#if defined(__AVX512VNNI__) && defined(__AVX512VL__)
    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__)
    gemv_q4_k_q8_k_avx(y, W, x_q8, M, K);
#else
    gemv_q4_k_q8_k_ref(y, W, x_q8, M, K);
#endif
}

References gemv_q4_k_q8_k_avx(), gemv_q4_k_q8_k_avx2(), gemv_q4_k_q8_k_ref(), and gemv_q4_k_q8_k_vnni().

gemv_q4_k_q8_k_avx()

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

Definition at line 251 of file gemm_kernels_q4k_avx.c.

{
    gemv_q4_k_q8_k_ref(y, W, x_q8, M, K);
}

References gemv_q4_k_q8_k_ref().

Referenced by gemv_q4_k_q8_k_amx().

gemv_q4_k_q8_k_avx2()

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

Definition at line 118 of file gemm_kernels_q4k_q8k_avx2.c.

{
#if defined(__AVX2__)
    if (!y || !W || !x_q8 || M <= 0 || K <= 0) {
        return;
    }
 
    const block_q4_K *blocks = (const block_q4_K *)W;
    const block_q8_K *x = (const block_q8_K *)x_q8;
    const int blocks_per_row = K / QK_K;
 
    for (int row = 0; row < M; ++row) {
        const block_q4_K *w_row = blocks + (size_t)row * (size_t)blocks_per_row;
        y[row] = dot_q4_k_q8_k_avx2(w_row, x, K);
    }
#else
    gemv_q4_k_q8_k_ref(y, W, x_q8, M, K);
#endif
}

References gemv_q4_k_q8_k_ref(), and QK_K.

Referenced by gemv_q4_k_q8_k_amx().

gemv_q4_k_q8_k_ref()

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

Definition at line 201 of file gemm_kernels_q4k_q8k.c.

{
    if (!y || !W || !x_q8 || M <= 0 || K <= 0) {
        return;
    }
 
    const block_q4_K *blocks = (const block_q4_K *)W;
    const block_q8_K *x = (const block_q8_K *)x_q8;
    const int blocks_per_row = K / QK_K;
 
    for (int row = 0; row < M; ++row) {
        const block_q4_K *w_row = blocks + (size_t)row * (size_t)blocks_per_row;
        y[row] = dot_q4_k_q8_k_ref(w_row, x, K);
    }
}

Referenced by gemv_q4_k_q8_k_amx().

gemv_q4_k_q8_k_vnni()

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

Definition at line 3534 of file gemm_kernels_q4k_q8k_vnni.c.

{
#if defined(__AVX512VNNI__) && defined(__AVX512VL__)
    const char *fast_env = getenv("CK_ENABLE_Q4K_Q8K_VNNI_FAST");
    const int fast_disabled = fast_env && fast_env[0] && fast_env[0] == '0';
    if (!fast_disabled && !ck_strict_parity_enabled()) {
        if (!y || !W || !x_q8 || M <= 0 || K <= 0) {
            return;
        }
 
        const block_q4_K *blocks = (const block_q4_K *)W;
        const block_q8_K *x = (const block_q8_K *)x_q8;
        const int blocks_per_row = K / QK_K;
 
        for (int row = 0; row < M; ++row) {
            const block_q4_K *w_row = blocks + (size_t)row * (size_t)blocks_per_row;
            float sum = 0.0f;
            for (int b = 0; b < blocks_per_row; ++b) {
                sum += dot_q4_k_q8_k_vnni_block(&w_row[b], &x[b]);
            }
            y[row] = sum;
        }
        return;
    }
#endif
 
    /* Strict/debug parity keeps the llama-style scalar accumulation path.
     * Production AVX-512 hosts use VNNI by default; set
     * CK_ENABLE_Q4K_Q8K_VNNI_FAST=0 or CK_DEBUG_Q4K_Q8_REF=1 when attributing
     * borderline logit movement against scalar/reference behavior.
     */
    gemv_q4_k_q8_k_ref(y, W, x_q8, M, K);
}

References ck_strict_parity_enabled(), gemv_q4_k_q8_k_ref(), and QK_K.

Referenced by gemv_q4_k_q8_k_amx().