From 104bf3bafb1b2f89a108b8b5a7e8c8d0f33fc87a Mon Sep 17 00:00:00 2001 From: AlexGabbia Date: Fri, 18 Sep 2026 02:14:22 +0200 Subject: [PATCH] ggml-cpu: add x86 AVX-VNNI dot product for PTQ1_0 PTQ1_0 had only the generic scalar vec_dot on x86. Add an AVX2 + AVX-VNNI / AVX-512-VNNI implementation of ggml_vec_dot_ptq1_0_q8_0, following the same shape as the existing PQ2_0 kernel: - decode the base-3 packed trits ((b * 3^p) & 0xFF) * 3 >> 8 to {0,1,2} codes with 16-bit lane arithmetic, no lookup tables - dot(code - 1, qy) = dpbusd(code, qy) - dpbusd(ones, qy) - per 32-wide q8_0 sub-block the integer sum and float accumulation order match the generic kernel exactly test-quantize-fns passes. Ternary-Bonsai-2-27B PTQ1_0, CPU only, 16 threads, Core Ultra 9 275HX (AVX-VNNI, no AVX-512): generic this kernel MSVC 19.44 0.79 tg64 / 0.85 pp256 2.87 tg64 / 3.78 pp256 (3.6x / 4.4x) clang 22 4.58 tg128 / 6.03 pp512 4.65 tg128 / 6.10 pp512 (parity) MSVC does not auto-vectorize the generic decode, so the explicit kernel gives shipped Windows binaries a 3.6-4.4x CPU speedup. Under clang the generic auto-vectorizes to near-parity; the kernel pins the fast path independently of the compiler. --- ggml/src/ggml-cpu/arch-fallback.h | 2 - ggml/src/ggml-cpu/arch/x86/quants.c | 88 +++++++++++++++++++++++++++++ 2 files changed, 88 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cpu/arch-fallback.h b/ggml/src/ggml-cpu/arch-fallback.h index d283c2ed51be..e4e220920be3 100644 --- a/ggml/src/ggml-cpu/arch-fallback.h +++ b/ggml/src/ggml-cpu/arch-fallback.h @@ -94,8 +94,6 @@ #define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K #define ggml_gemm_pq2_0_4x8_q8_0_generic ggml_gemm_pq2_0_4x8_q8_0 #elif defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) || defined(_M_X64) -// 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 // quants.c #define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0 // repack.cpp diff --git a/ggml/src/ggml-cpu/arch/x86/quants.c b/ggml/src/ggml-cpu/arch/x86/quants.c index 7fc8e7713b45..6f4b0922a931 100644 --- a/ggml/src/ggml-cpu/arch/x86/quants.c +++ b/ggml/src/ggml-cpu/arch/x86/quants.c @@ -628,6 +628,94 @@ void ggml_vec_dot_pq2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo *s = sumf; } +// PTQ1_0 x Q8_0 with AVX-VNNI / AVX-512-VNNI. +// Decode trits to codes in {0,1,2}, then dot(code - 1, qy) = dpbusd(code, qy) - dpbusd(ones, qy). +void ggml_vec_dot_ptq1_0_q8_0(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) { +#if (defined(__AVX512VNNI__) && defined(__AVX512VL__)) || defined(__AVXVNNI__) + const int qk = QK_PTQ1_0; + const int nb = n / qk; + + assert(n % qk == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_ptq1_0 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + static const uint8_t pow3[5] = {1, 3, 9, 27, 81}; + const __m256i ones = _mm256_set1_epi8(1); + const __m256i mask = _mm256_set1_epi16(0x00FF); + const __m256i three = _mm256_set1_epi16(3); + + float sumf = 0.0f; + + for (int i = 0; i < nb; i++) { + // codes_a[p]: qs[0..15] at power p -> elements 16p..16p+15 + // codes_b[p]: qs[16..23] at power p -> elements 80+8p..87+8p + // codes_h: qh[0..1] at powers 0..3 -> elements 120..127 + __m128i codes_a[5], codes_b[5]; + + const __m128i qa = _mm_loadu_si128((const __m128i *) &x[i].qs[0]); + const __m128i qb = _mm_loadl_epi64((const __m128i *) &x[i].qs[16]); + + for (int p = 0; p < 5; ++p) { + const __m256i p3 = _mm256_set1_epi16(pow3[p]); + for (int src = 0; src < 2; ++src) { + const __m128i v = src ? qb : qa; + const __m256i w = _mm256_cvtepu8_epi16(v); + const __m256i t = _mm256_and_si256(_mm256_mullo_epi16(w, p3), mask); + const __m256i c = _mm256_srli_epi16(_mm256_mullo_epi16(t, three), 8); + const __m256i pk = _mm256_packus_epi16(c, _mm256_setzero_si256()); + const __m128i d = _mm_unpacklo_epi64(_mm256_castsi256_si128(pk), + _mm256_extracti128_si256(pk, 1)); + if (src) { + codes_b[p] = d; + } else { + codes_a[p] = d; + } + } + } + + const uint8_t qh0 = x[i].qh[0]; + const uint8_t qh1 = x[i].qh[1]; + const __m128i codes_h = _mm_setr_epi8( + ((qh0 * 1) & 0xFF) * 3 >> 8, ((qh1 * 1) & 0xFF) * 3 >> 8, + ((qh0 * 3) & 0xFF) * 3 >> 8, ((qh1 * 3) & 0xFF) * 3 >> 8, + ((qh0 * 9) & 0xFF) * 3 >> 8, ((qh1 * 9) & 0xFF) * 3 >> 8, + ((qh0 * 27) & 0xFF) * 3 >> 8, ((qh1 * 27) & 0xFF) * 3 >> 8, + 0, 0, 0, 0, 0, 0, 0, 0); + + // per 32-wide q8_0 sub-block: a0 a1 | a2 a3 | a4 b0 b1 | b2 b3 b4 h + const __m256i cc[4] = { + _mm256_set_m128i(codes_a[1], codes_a[0]), + _mm256_set_m128i(codes_a[3], codes_a[2]), + _mm256_set_m128i(_mm_unpacklo_epi64(codes_b[0], codes_b[1]), codes_a[4]), + _mm256_set_m128i(_mm_unpacklo_epi64(codes_b[4], codes_h), + _mm_unpacklo_epi64(codes_b[2], codes_b[3])), + }; + + const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d); + float sumi = 0.0f; + 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 int dp = hsum_i32_8(GGML_DPBUSD_256(_mm256_setzero_si256(), cc[k], qy)); + const int sy = hsum_i32_8(GGML_DPBUSD_256(_mm256_setzero_si256(), ones, qy)); + sumi += d1 * (float) (dp - sy); + } + sumf += d0 * sumi; + } + + *s = sumf; +#else + ggml_vec_dot_ptq1_0_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + void ggml_vec_dot_q1_0_q8_0(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) { const int qk = QK1_0; const int nb = n / qk;