Skip to content

ggml-cpu: AVX2/AVX-VNNI kernels for PQ2_0 + Q8_K activations + tiled GEMM (x86 without AVX-512) #206

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bri-prism merged 3 commits into from
Sep 21, 2026
Merged

ggml-cpu: AVX2/AVX-VNNI kernels for PQ2_0 + Q8_K activations + tiled GEMM (x86 without AVX-512) #206

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
164c337
ggml-cpu: AVX2/AVX-VNNI kernels for PQ2_0 + Q8_K activations + tiled …
alenizy361 Sep 19, 2026
10df298
ggml-cpu, llama-quant: fix PQ2_0 non-x86 link break and 128-column sh…
alenizy361 Sep 20, 2026
d2c9ddd
ggml-cpu: keep PQ2_0 rows that are not multiples of 256 on the Q8_0 a…
alenizy361 Sep 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions ggml/src/ggml-cpu/arch-fallback.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
#define ggml_vec_dot_pq2_0_q8_0_generic ggml_vec_dot_pq2_0_q8_0
// PQ2_0 x Q8_K has only the generic vec_dot outside x86; alias it until a SIMD version lands
#define ggml_vec_dot_pq2_0_q8_K_generic ggml_vec_dot_pq2_0_q8_K
#define ggml_vec_dot_ptq1_0_q8_0_generic ggml_vec_dot_ptq1_0_q8_0
#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K
Expand Down Expand Up @@ -82,6 +84,8 @@
#elif defined(__aarch64__) || defined(__arm__) || defined(_M_ARM) || defined(_M_ARM64)
// PTQ1_0 currently has only the generic vec_dot; alias it here until a SIMD version lands
#define ggml_vec_dot_ptq1_0_q8_0_generic ggml_vec_dot_ptq1_0_q8_0
// PQ2_0 x Q8_K has only the generic vec_dot outside x86; alias it until a SIMD version lands
#define ggml_vec_dot_pq2_0_q8_K_generic ggml_vec_dot_pq2_0_q8_K
// repack.cpp
#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
#define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8
Expand Down Expand Up @@ -134,6 +138,8 @@
#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
#define ggml_vec_dot_pq2_0_q8_0_generic ggml_vec_dot_pq2_0_q8_0
// PQ2_0 x Q8_K has only the generic vec_dot outside x86; alias it until a SIMD version lands
#define ggml_vec_dot_pq2_0_q8_K_generic ggml_vec_dot_pq2_0_q8_K
#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K
#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K
Expand Down Expand Up @@ -193,6 +199,8 @@
#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
#define ggml_vec_dot_pq2_0_q8_0_generic ggml_vec_dot_pq2_0_q8_0
// PQ2_0 x Q8_K has only the generic vec_dot outside x86; alias it until a SIMD version lands
#define ggml_vec_dot_pq2_0_q8_K_generic ggml_vec_dot_pq2_0_q8_K
#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
// repack.cpp
#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
Expand Down Expand Up @@ -244,6 +252,8 @@
#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
#define ggml_vec_dot_pq2_0_q8_0_generic ggml_vec_dot_pq2_0_q8_0
// PQ2_0 x Q8_K has only the generic vec_dot outside x86; alias it until a SIMD version lands
#define ggml_vec_dot_pq2_0_q8_K_generic ggml_vec_dot_pq2_0_q8_K
// repack.cpp
#define ggml_quantize_mat_q8_0_4x1_generic ggml_quantize_mat_q8_0_4x1
#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
Expand Down Expand Up @@ -294,6 +304,8 @@
#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
#define ggml_vec_dot_pq2_0_q8_0_generic ggml_vec_dot_pq2_0_q8_0
// PQ2_0 x Q8_K has only the generic vec_dot outside x86; alias it until a SIMD version lands
#define ggml_vec_dot_pq2_0_q8_K_generic ggml_vec_dot_pq2_0_q8_K
#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
#define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K
#define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K
Expand Down Expand Up @@ -367,6 +379,8 @@
#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
#define ggml_vec_dot_pq2_0_q8_0_generic ggml_vec_dot_pq2_0_q8_0
// PQ2_0 x Q8_K has only the generic vec_dot outside x86; alias it until a SIMD version lands
#define ggml_vec_dot_pq2_0_q8_K_generic ggml_vec_dot_pq2_0_q8_K
#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
// repack.cpp
#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
Expand Down
118 changes: 96 additions & 22 deletions ggml/src/ggml-cpu/arch/x86/quants.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -574,34 +574,51 @@ void ggml_vec_dot_pq2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo

float sumf = 0.0f;

#if (defined(__AVX512VNNI__) && defined(__AVX512VL__)) || defined(__AVXVNNI__)
// AVX-VNNI / AVX-512-VNNI: unpack 2-bit codes c in {0,1,2,3} (value = c-1), then
// dot((c-1), qy) = dpbusd(c, qy) - dpbusd(1, qy). Group 128 = four q8_0 sub-blocks.
const __m256i ones = _mm256_set1_epi8(1);
const __m128i idxlo = _mm_setr_epi8(0,0,0,0,1,1,1,1,2,2,2,2,3,3,3,3);
const __m128i idxhi = _mm_setr_epi8(4,4,4,4,5,5,5,5,6,6,6,6,7,7,7,7);
const __m256i mul = _mm256_setr_epi16(64,16,4,1, 64,16,4,1, 64,16,4,1, 64,16,4,1); // <<(6-2c)
const __m256i three = _mm256_set1_epi16(3);
#if defined(__AVX2__)
// AVX2 (+AVX-VNNI / AVX-512-VNNI dot when available). 2-bit codes c in {0,1,2,3}
// (value = c-1); byte b of a 32-weight sub-block holds weights 4b..4b+3, LSB pair first.
// Group 128 = four q8_0 sub-blocks.
//
// Unpack: broadcast the 8 code bytes into all four 64-bit lanes, shift lane j right by
// 2j and mask with 3 -> byte [8j+b] = code of weight 4b+j (3 instructions, no
// per-sub-block horizontal sums, which dominated the previous kernel: 2 hsums per 32
// weights and a scalar float chain). The q8 activations are permuted once per sub-block
// to the same [8j+b] order, and the int32 partial sums stay vectorized; the block scale
// is applied with one FMA and there is a single horizontal sum at the very end.
const __m256i shifts = _mm256_setr_epi64x(0, 2, 4, 6);
const __m256i mask3 = _mm256_set1_epi8(3);
const __m256i ones_8 = _mm256_set1_epi8(1);
#if !defined(GGML_DPBUSD_256)
const __m256i ones_16 = _mm256_set1_epi16(1);
#endif
// weight order [4b+j] -> code order [8j+b]: in-lane 4x4 byte transpose, then interleave
// the two 128-bit lanes dword-wise (lo j0, hi j0, lo j1, hi j1, ...).
const __m256i qy_shuf = _mm256_setr_epi8(0,4,8,12, 1,5,9,13, 2,6,10,14, 3,7,11,15,
0,4,8,12, 1,5,9,13, 2,6,10,14, 3,7,11,15);
const __m256i qy_perm = _mm256_setr_epi32(0,4,1,5,2,6,3,7);
__m256 acc = _mm256_setzero_ps();
for (int i = 0; i < nb; i++) {
const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d);
float sumi = 0.0f;
__m256 acc_block = _mm256_setzero_ps();
for (int k = 0; k < 4; k++) {
const block_q8_0 * GGML_RESTRICT yb = &y[i * 4 + k];
const float d1 = GGML_CPU_FP16_TO_FP32(yb->d);
const __m256i qy = _mm256_loadu_si256((const __m256i *) yb->qs);
const __m128i src = _mm_loadl_epi64((const __m128i *) &x[i].qs[k * 8]); // 8 bytes
const __m256i rep = _mm256_set_m128i(_mm_shuffle_epi8(src, idxhi), _mm_shuffle_epi8(src, idxlo));
__m256i r0 = _mm256_cvtepu8_epi16(_mm256_castsi256_si128(rep));
__m256i r1 = _mm256_cvtepu8_epi16(_mm256_extracti128_si256(rep, 1));
r0 = _mm256_and_si256(_mm256_srli_epi16(_mm256_mullo_epi16(r0, mul), 6), three);
r1 = _mm256_and_si256(_mm256_srli_epi16(_mm256_mullo_epi16(r1, mul), 6), three);
__m256i codes = _mm256_permute4x64_epi64(_mm256_packus_epi16(r0, r1), 0xD8);
const int dp = hsum_i32_8(GGML_DPBUSD_256(_mm256_setzero_si256(), codes, qy));
const int sy = hsum_i32_8(GGML_DPBUSD_256(_mm256_setzero_si256(), ones, qy));
sumi += d1 * (float)(dp - sy);
const __m256i qy = _mm256_loadu_si256((const __m256i *) yb->qs);
const __m256i qyp = _mm256_permutevar8x32_epi32(_mm256_shuffle_epi8(qy, qy_shuf), qy_perm);
int64_t xq; memcpy(&xq, &x[i].qs[k * 8], sizeof(xq));
const __m256i codes = _mm256_and_si256(_mm256_srlv_epi64(_mm256_set1_epi64x(xq), shifts), mask3);
// sum((c-1)*qy) = sum(c*qy) - sum(qy), per 32-bit lane (4 products each)
#if defined(GGML_DPBUSD_256)
const __m256i s32 = _mm256_sub_epi32(GGML_DPBUSD_256(_mm256_setzero_si256(), codes, qyp),
GGML_DPBUSD_256(_mm256_setzero_si256(), ones_8, qyp));
#else
const __m256i s32 = _mm256_sub_epi32(_mm256_madd_epi16(_mm256_maddubs_epi16(codes, qyp), ones_16),
_mm256_madd_epi16(_mm256_maddubs_epi16(ones_8, qyp), ones_16));
#endif
acc_block = _mm256_fmadd_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(yb->d)), _mm256_cvtepi32_ps(s32), acc_block);
}
sumf += d0 * sumi;
acc = _mm256_fmadd_ps(_mm256_set1_ps(d0), acc_block, acc);
}
sumf = hsum_float_8(acc);
#else
for (int i = 0; i < nb; i++) {
const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d);
Expand Down Expand Up @@ -4182,3 +4199,60 @@ void ggml_vec_dot_iq4_xs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const v
ggml_vec_dot_iq4_xs_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
#endif
}


// PQ2_0 x Q8_K AVX2 / AVX-VNNI dot (tier 2, 2026-09-18). With one activation scale per 256 the four sub-block
// dots of a 128-weight block are accumulated in int32 and scaled ONCE (4 dpbusd + 1 sub + 1 cvt + 1 fmadd per
// 128 weights, against 4x(2 dpbusd + sub + cvt + fmadd) on the Q8_0 path). The sum(y) correction for the
// unsigned codes is kept per lane (dpbusd(ones, y)) so the per-lane partials stay exact; the tiled GEMM in
// llamafile/sgemm.cpp uses the same per-block float op and the same reduction, so both paths agree bit-for-bit.
static inline __m256i pq2k_dpbusd_acc(__m256i acc, __m256i u, __m256i s) {
#if defined(GGML_DPBUSD_256)
return GGML_DPBUSD_256(acc, u, s);
#else
return _mm256_add_epi32(acc, _mm256_madd_epi16(_mm256_maddubs_epi16(u, s), _mm256_set1_epi16(1)));
#endif
}
void ggml_vec_dot_pq2_0_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
assert(n % QK_K == 0);
assert(nrc == 1);
UNUSED(nrc);
UNUSED(bx);
UNUSED(by);
UNUSED(bs);

const block_pq2_0 * GGML_RESTRICT x = vx;
const block_q8_K * GGML_RESTRICT y = vy;
const int nb = n / QK_PQ2_0;

float sumf = 0.0f;
#if defined(__AVX2__)
const __m256i shifts = _mm256_setr_epi64x(0, 2, 4, 6);
const __m256i mask3 = _mm256_set1_epi8(3);
const __m256i ones_8 = _mm256_set1_epi8(1);
const __m256i qy_shuf = _mm256_setr_epi8(0,4,8,12, 1,5,9,13, 2,6,10,14, 3,7,11,15,
0,4,8,12, 1,5,9,13, 2,6,10,14, 3,7,11,15);
const __m256i qy_perm = _mm256_setr_epi32(0,4,1,5,2,6,3,7);
__m256 acc = _mm256_setzero_ps();
for (int i = 0; i < nb; i++) {
const block_q8_K * GGML_RESTRICT yb = &y[i >> 1];
const int8_t * GGML_RESTRICT q8 = yb->qs + 128 * (i & 1);
__m256i acc_c = _mm256_setzero_si256();
__m256i acc_1 = _mm256_setzero_si256();
for (int k = 0; k < 4; k++) {
const __m256i qy = _mm256_loadu_si256((const __m256i *) (q8 + 32 * k));
const __m256i qyp = _mm256_permutevar8x32_epi32(_mm256_shuffle_epi8(qy, qy_shuf), qy_perm);
int64_t xq; memcpy(&xq, &x[i].qs[8 * k], sizeof(xq));
const __m256i codes = _mm256_and_si256(_mm256_srlv_epi64(_mm256_set1_epi64x(xq), shifts), mask3);
acc_c = pq2k_dpbusd_acc(acc_c, codes, qyp);
acc_1 = pq2k_dpbusd_acc(acc_1, ones_8, qyp);
}
const __m256i s32 = _mm256_sub_epi32(acc_c, acc_1);
acc = _mm256_fmadd_ps(_mm256_set1_ps(GGML_CPU_FP16_TO_FP32(x[i].d) * yb->d), _mm256_cvtepi32_ps(s32), acc);
}
sumf = hsum_float_8(acc);
#else
ggml_vec_dot_pq2_0_q8_K_generic(n, &sumf, bs, vx, bx, vy, by, nrc);
#endif
*s = sumf;
}
47 changes: 37 additions & 10 deletions ggml/src/ggml-cpu/ggml-cpu.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -239,8 +239,8 @@ static const struct ggml_type_traits_cpu type_traits_cpu[GGML_TYPE_COUNT] = {
},
[GGML_TYPE_PQ2_0] = {
.from_float = quantize_row_pq2_0,
.vec_dot = ggml_vec_dot_pq2_0_q8_0,
.vec_dot_type = GGML_TYPE_Q8_0,
.vec_dot = ggml_vec_dot_pq2_0_q8_K,
.vec_dot_type = GGML_TYPE_Q8_K,
Comment on lines +242 to +243
.nrows = 1,
},
[GGML_TYPE_PTQ1_0] = {
Expand Down Expand Up @@ -1174,6 +1174,33 @@ void ggml_set_f32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2,

// ggml_compute_forward_mul_mat

// PQ2_0 stores 128 weights per block, so any multiple of 128 is a legal row width and models in
// the wild contain rows of 128, 384, 640, ... The fast PQ2_0 kernels consume Q8_K activations,
// which only exist in 256-element blocks; a row that is not a multiple of 256 cannot be expressed
// in Q8_K at all, and forcing it through that path reads past the activation buffer and returns
// garbage (or trips ggml_row_size's block-size assert on builds where it is enabled).
//
// Those rows keep the Q8_0 activation path they used before the Q8_K kernels landed: correct, and
// still the AVX2 kernel rather than the generic scalar dot. Every other PQ2_0 row, which is the
// overwhelming majority in practice, is untouched and keeps the faster Q8_K path.
static inline bool ggml_cpu_pq2_needs_q8_0(const struct ggml_tensor * src0) {
return src0->type == GGML_TYPE_PQ2_0 && (src0->ne[0] % QK_K) != 0;
}

static inline enum ggml_type ggml_cpu_vec_dot_type(const struct ggml_tensor * src0) {
if (ggml_cpu_pq2_needs_q8_0(src0)) {
return GGML_TYPE_Q8_0;
}
return type_traits_cpu[src0->type].vec_dot_type;
}

static inline ggml_vec_dot_t ggml_cpu_vec_dot(const struct ggml_tensor * src0) {
if (ggml_cpu_pq2_needs_q8_0(src0)) {
return ggml_vec_dot_pq2_0_q8_0;
}
return type_traits_cpu[src0->type].vec_dot;
}

static void ggml_compute_forward_mul_mat_one_chunk(
const struct ggml_compute_params * params,
struct ggml_tensor * dst,
Expand All @@ -1191,8 +1218,8 @@ static void ggml_compute_forward_mul_mat_one_chunk(

const bool src1_cont = ggml_is_contiguous(src1);

ggml_vec_dot_t const vec_dot = type_traits_cpu[type].vec_dot;
enum ggml_type const vec_dot_type = type_traits_cpu[type].vec_dot_type;
ggml_vec_dot_t const vec_dot = ggml_cpu_vec_dot(src0);
enum ggml_type const vec_dot_type = ggml_cpu_vec_dot_type(src0);

// broadcast factors
const int64_t r2 = ne12 / ne02;
Expand Down Expand Up @@ -1282,7 +1309,7 @@ void ggml_compute_forward_mul_mat(
const int ith = params->ith;
const int nth = params->nth;

enum ggml_type const vec_dot_type = type_traits_cpu[src0->type].vec_dot_type;
enum ggml_type const vec_dot_type = ggml_cpu_vec_dot_type(src0);
ggml_from_float_t const from_float = type_traits_cpu[vec_dot_type].from_float;
int64_t const vec_dot_num_rows = type_traits_cpu[src0->type].nrows;

Expand Down Expand Up @@ -1501,8 +1528,8 @@ static void ggml_compute_forward_mul_mat_id_one_chunk(

const enum ggml_type type = src0->type;

ggml_vec_dot_t const vec_dot = type_traits_cpu[type].vec_dot;
enum ggml_type const vec_dot_type = type_traits_cpu[type].vec_dot_type;
ggml_vec_dot_t const vec_dot = ggml_cpu_vec_dot(src0);
enum ggml_type const vec_dot_type = ggml_cpu_vec_dot_type(src0);

const int64_t blck_0 = 16;
const int64_t blck_1 = 16;
Expand Down Expand Up @@ -1569,7 +1596,7 @@ static void ggml_compute_forward_mul_mat_id(

const bool src1_cont = ggml_is_contiguous(src1);

enum ggml_type const vec_dot_type = type_traits_cpu[type].vec_dot_type;
enum ggml_type const vec_dot_type = ggml_cpu_vec_dot_type(src0);
ggml_from_float_t const from_float = type_traits_cpu[vec_dot_type].from_float;

// we don't support permuted src0 or src1
Expand Down Expand Up @@ -2873,7 +2900,7 @@ struct ggml_cplan ggml_graph_plan(
} break;
case GGML_OP_MUL_MAT:
{
const enum ggml_type vec_dot_type = type_traits_cpu[node->src[0]->type].vec_dot_type;
const enum ggml_type vec_dot_type = ggml_cpu_vec_dot_type(node->src[0]);

if (node->src[1]->type != vec_dot_type) {
cur = ggml_row_size(vec_dot_type, ggml_nelements(node->src[1]));
Expand All @@ -2885,7 +2912,7 @@ struct ggml_cplan ggml_graph_plan(
const struct ggml_tensor * src0 = node->src[0];
const struct ggml_tensor * src1 = node->src[1];
const struct ggml_tensor * ids = node->src[2];
const enum ggml_type vec_dot_type = type_traits_cpu[src0->type].vec_dot_type;
const enum ggml_type vec_dot_type = ggml_cpu_vec_dot_type(src0);
const int n_as = src0->ne[2];
// src1
if (src1->type != vec_dot_type) {
Expand Down
Loading