From 164c33700e48de0189a83f2c7b8fed549c8d5dfe Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Sat, 19 Sep 2026 10:28:24 +0300 Subject: [PATCH 1/3] ggml-cpu: AVX2/AVX-VNNI kernels for PQ2_0 + Q8_K activations + tiled GEMM (x86 without AVX-512) The fork's fast PQ2_0 paths (gemv/gemm_pq2_0_4x8) are gated on AVX-512, so 12th-14th gen Intel and other AVX2/AVX-VNNI CPUs fell back to the generic scalar dot. This adds: - ggml_vec_dot_pq2_0_q8_0 / ggml_vec_dot_pq2_0_q8_K (AVX2, AVX-VNNI, AVX-512VNNI-VL): 2-bit codes unpacked with srlv+and, dpbusd against the permuted activation bytes, unsigned-code correction via a lane-wise sum(y) so per-lane partials stay exact. - PQ2_0 vec_dot_type switched to Q8_K (one activation scale per 256): 4 dpbusd + 1 sub + 1 cvt + 1 fmadd per 128 weights instead of 4x(2 dpbusd + sub + cvt + fmadd). Perplexity unchanged within 0.004 (Bonsai-2 27B: 7.0542 -> 7.0583 +/- 0.58). - llamafile/sgemm.cpp: tinyBLAS_PQ2K_AVX / tinyBLAS_PQ2_AVX tiled GEMMs. A 128-weight block is unpacked once per row tile and reused across RN token columns; the activation permute and lane sums are staged once per matmul call per thread; jobs are row tiles (each weight row streams from DRAM exactly once) claimed dynamically from the threadpool counter in runs of 8 so P- and E-cores share the compute-bound loop. Float accumulation order and reduction are identical to the vec_dot path, so batched and single-token results are bit-identical (PPL identical with PQ2_SGEMM=0/1). Measured on an i7-13620H (6P+4E, AVX2+AVX-VNNI, single-channel DDR5-5200 ~28 GB/s effective), Ternary-Bonsai-2-27B PQ2_0: decode (N=1) 1.3 -> 3.9-4.2 tok/s (generic -> vec_dot) batch cost, ms/token N=4: 123 -> 78, N=8: 99 -> 68, N=16: 89 -> 54 (t=12, 65 W), N=64: 89 -> 69 prompt processing 6.7 -> 16 tok/s MTP speculative verify (N=3) becomes ~1 DRAM sweep, giving 6.8-8.3 tok/s greedy decode with the grafted nextn head. Env: PQ2_SGEMM=0 disables the tiled path (A/B). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- ggml/src/ggml-cpu/arch/x86/quants.c | 118 +++++++++-- ggml/src/ggml-cpu/ggml-cpu.c | 4 +- ggml/src/ggml-cpu/llamafile/sgemm.cpp | 279 ++++++++++++++++++++++++++ ggml/src/ggml-cpu/quants.c | 33 +++ ggml/src/ggml-cpu/quants.h | 2 + 5 files changed, 412 insertions(+), 24 deletions(-) diff --git a/ggml/src/ggml-cpu/arch/x86/quants.c b/ggml/src/ggml-cpu/arch/x86/quants.c index 7fc8e7713b45..bd572432f48f 100644 --- a/ggml/src/ggml-cpu/arch/x86/quants.c +++ b/ggml/src/ggml-cpu/arch/x86/quants.c @@ -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); @@ -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; +} diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 9cd2c2d39a5a..48c3f5f0890f 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -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, .nrows = 1, }, [GGML_TYPE_PTQ1_0] = { diff --git a/ggml/src/ggml-cpu/llamafile/sgemm.cpp b/ggml/src/ggml-cpu/llamafile/sgemm.cpp index 99b7d5afa2f9..db7a26ed19a1 100644 --- a/ggml/src/ggml-cpu/llamafile/sgemm.cpp +++ b/ggml/src/ggml-cpu/llamafile/sgemm.cpp @@ -56,6 +56,8 @@ #include #include +#include +#include #ifdef _MSC_VER #define NOINLINE __declspec(noinline) @@ -3772,6 +3774,265 @@ class tinyBLAS_PPC { #endif } // namespace +// ---- PQ2_0 x Q8_K tiled GEMM (tier 2, 2026-09-18). Same tiling as tinyBLAS_PQ2_AVX; with a single activation +// scale per 256 the per-(row,col) work per 128 weights collapses to 4 dpbusd + sub + cvt + fmadd, the activation +// permute and sum(y) are shared across RM rows, the unpacked codes across RN columns. Per-block float op and the +// final reduction match ggml_vec_dot_pq2_0_q8_K exactly, so the batched and single-token paths agree bit-for-bit. +#if defined(__AVX2__) +static inline float pq2_hsum(__m256 x); // defined with the tier-1 class below +static inline __m256i pq2k_sg_dpbusd(__m256i acc, __m256i u, __m256i s) { +#if defined(__AVX512VNNI__) && defined(__AVX512VL__) + return _mm256_dpbusd_epi32(acc, u, s); +#elif defined(__AVXVNNI__) + return _mm256_dpbusd_avx_epi32(acc, u, s); +#else + return _mm256_add_epi32(acc, _mm256_madd_epi16(_mm256_maddubs_epi16(u, s), _mm256_set1_epi16(1))); +#endif +} +#include +class tinyBLAS_PQ2K_AVX { + public: + tinyBLAS_PQ2K_AVX(int64_t k, const block_pq2_0 *A, int64_t lda, const block_q8_K *B, int64_t ldb, float *C, int64_t ldc, int ith, int nth, + const struct ggml_compute_params *params = nullptr) + : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth), params(params) {} + void matmul(int64_t m, int64_t n) { prep(n); mnpack(0, m, 0, n); } + private: + // Staging (2026-09-19): the permute of every 32-byte Q8_K group into code order and the lane-wise sum(y) of + // every 128-weight block depend only on (column, block), but the first tile recomputed both inside every row + // tile (m/RM times per matmul; perf: 91% of a 64-token batch was this kernel, ~1/3 of it that redundancy). + // Each thread now stages all n columns once per call into a thread-local buffer: per (column, PQ2 block) + // 4 x 32 permuted bytes followed by the 8 x int32 lane sums. Same integer sums in the same lanes, so the + // float accumulation order is unchanged and results stay bit-identical to ggml_vec_dot_pq2_0_q8_K. + static constexpr int64_t BP_BLOCK = 160; + const uint8_t *Bp = nullptr; + void prep(int64_t n) { + static thread_local std::vector buf; + const size_t need = (size_t) n * (size_t) k * BP_BLOCK + 64; + if (buf.size() < need) buf.resize(need); + uint8_t *base = (uint8_t *) (((uintptr_t) buf.data() + 31) & ~(uintptr_t) 31); + 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); + for (int64_t j = 0; j < n; ++j) { + for (int64_t l = 0; l < k; ++l) { + const int8_t *q8 = B[ldb * j + (l >> 1)].qs + 128 * (int) (l & 1); + uint8_t *dst = base + ((size_t) j * (size_t) k + (size_t) l) * BP_BLOCK; + __m256i sumq = _mm256_setzero_si256(); + for (int s = 0; s < 4; ++s) { + const __m256i qy = _mm256_loadu_si256((const __m256i *) (q8 + 32 * s)); + const __m256i qyp = _mm256_permutevar8x32_epi32(_mm256_shuffle_epi8(qy, qy_shuf), qy_perm); + _mm256_store_si256((__m256i *) (dst + 32 * s), qyp); + sumq = pq2k_sg_dpbusd(sumq, ones_8, qyp); + } + _mm256_store_si256((__m256i *) (dst + 128), sumq); + } + } + Bp = base; + } + void mnpack(int64_t m0, int64_t m, int64_t n0, int64_t n) { + int64_t mc, nc, mp, np; + // (RM=1,RN=8 measured L2-bound: 128 B of staged activations per (row,col) block vs 64 B for RM=2 — N=64 + // went 88 -> 103 ms/token. Wide-and-two-rows stays the tile for every n.) + // Wide tiles first: 3-4 token columns share one weight pass (RM=2 keeps it at 15 ymm). n=2 keeps RM=4. + switch ((MIN(m - m0, 4) << 4) | MIN(n - n0, 4)) { + case 0x44: case 0x34: case 0x24: mc = 2; nc = 4; gemm<2, 4>(m0, m, n0, n); break; + case 0x43: case 0x33: case 0x23: mc = 2; nc = 3; gemm<2, 3>(m0, m, n0, n); break; + case 0x14: mc = 1; nc = 4; gemm<1, 4>(m0, m, n0, n); break; + case 0x13: mc = 1; nc = 3; gemm<1, 3>(m0, m, n0, n); break; + case 0x42: mc = 4; nc = 2; gemm<4, 2>(m0, m, n0, n); break; + case 0x41: mc = 4; nc = 1; gemm<4, 1>(m0, m, n0, n); break; + case 0x32: mc = 3; nc = 2; gemm<3, 2>(m0, m, n0, n); break; + case 0x31: mc = 3; nc = 1; gemm<3, 1>(m0, m, n0, n); break; + case 0x22: mc = 2; nc = 2; gemm<2, 2>(m0, m, n0, n); break; + case 0x21: mc = 2; nc = 1; gemm<2, 1>(m0, m, n0, n); break; + case 0x12: mc = 1; nc = 2; gemm<1, 2>(m0, m, n0, n); break; + case 0x11: mc = 1; nc = 1; gemm<1, 1>(m0, m, n0, n); break; + default: return; + } + mp = m0 + (m - m0) / mc * mc; + np = n0 + (n - n0) / nc * nc; + mnpack(mp, m, n0, np); + mnpack(m0, m, np, n); + } + template + NOINLINE void gemm(int64_t m0, int64_t m, int64_t n0, int64_t n) { + const int64_t ytiles = (m - m0) / RM; + const int64_t xtiles = (n - n0) / RN; + // A job is one row tile with ALL its column tiles, so every weight row streams from DRAM exactly once + // whichever thread takes it. With several threads and enough work the row tiles are claimed from the + // threadpool's shared counter (same pattern as the fp32 path above): this loop is compute-bound for + // n >= 4 and P-/E-cores run it at different speeds, so a static split left the P-cores idle at the end + // (-t 10 static: 101 ms/token vs 70 with -t 6). + const bool dynamic = nth > 1 && params != nullptr && xtiles * ytiles >= 64; + int64_t start, end; + if (dynamic) { + if (ith == 0) ggml_threadpool_chunk_set(params->threadpool, nth); + ggml_barrier(params->threadpool); + start = ith; end = ytiles; + } else { + const int64_t duty = (ytiles + nth - 1) / nth; + start = duty * ith; end = start + duty; if (end > ytiles) end = ytiles; + } + const __m256i shifts = _mm256_setr_epi64x(0, 2, 4, 6); + const __m256i mask3 = _mm256_set1_epi8(3); + // Claims come in runs of CH row tiles to keep the shared counter off the hot path (ytiles is 2.5k-8.7k here). + constexpr int64_t CH = 8; + if (dynamic) { start = (int64_t) ith * CH; } + for (int64_t y0 = start; y0 < end; y0 = dynamic ? ggml_threadpool_chunk_add(params->threadpool, 1) * CH : end) { + const int64_t y1 = dynamic ? MIN(y0 + CH, end) : end; + for (int64_t yt = y0; yt < y1; ++yt) { + const int64_t ii = m0 + yt * RM; + for (int64_t xt = 0; xt < xtiles; ++xt) { + const int64_t jj = n0 + xt * RN; + const uint8_t *bp[RN]; + for (int j = 0; j < RN; ++j) bp[j] = Bp + (size_t) (jj + j) * (size_t) k * BP_BLOCK; + __m256 Cv[RN][RM] = {}; + for (int64_t l = 0; l < k; ++l) { // k = number of PQ2_0 blocks per row + const int64_t lk = l >> 1; // Q8_K block index + const size_t lo = (size_t) l * BP_BLOCK; + __m256i acc[RN][RM]; + for (int j = 0; j < RN; ++j) for (int i = 0; i < RM; ++i) acc[j][i] = _mm256_setzero_si256(); + for (int s = 0; s < 4; ++s) { + __m256i codes[RM]; + for (int i = 0; i < RM; ++i) { + int64_t xq; memcpy(&xq, &A[lda * (ii + i) + l].qs[s * 8], sizeof(xq)); + codes[i] = _mm256_and_si256(_mm256_srlv_epi64(_mm256_set1_epi64x(xq), shifts), mask3); + } + for (int j = 0; j < RN; ++j) { + const __m256i qyp = _mm256_load_si256((const __m256i *) (bp[j] + lo + 32 * s)); + for (int i = 0; i < RM; ++i) acc[j][i] = pq2k_sg_dpbusd(acc[j][i], codes[i], qyp); + } + } + for (int i = 0; i < RM; ++i) { + const float da = unhalf(A[lda * (ii + i) + l].d); + for (int j = 0; j < RN; ++j) { + const __m256i sumq = _mm256_load_si256((const __m256i *) (bp[j] + lo + 128)); + const __m256 sc = _mm256_set1_ps(da * B[ldb * (jj + j) + lk].d); + Cv[j][i] = _mm256_fmadd_ps(sc, _mm256_cvtepi32_ps(_mm256_sub_epi32(acc[j][i], sumq)), Cv[j][i]); + } + } + } + for (int j = 0; j < RN; ++j) + for (int i = 0; i < RM; ++i) + C[ldc * (jj + j) + (ii + i)] = pq2_hsum(Cv[j][i]); + } + } + } + if (dynamic) ggml_barrier(params->threadpool); + } + const block_pq2_0 *const A; + const block_q8_K *const B; + float *const C; + const int64_t k, lda, ldb, ldc; + const int ith, nth; + const struct ggml_compute_params *const params; +}; +#endif // __AVX2__ + +// ---- PQ2_0 (PrismML 2-bit ternary, group 128) x Q8_0 tiled GEMM for AVX2 / AVX-VNNI. Added 2026-09-18. ---- +// The vec_dot path re-unpacks every 128-weight block once PER TOKEN, so for a batch of n tokens (speculative +// verify, prompt eval) unpack dominates: measured batch cost was a flat 168 ms/token regardless of n. This tile +// unpacks a block once per row and reuses it across RN token columns, and computes the activation permute and +// the sum(y) term once per column and reuses them across RM rows. Per (row, col) the float accumulation order +// and the horizontal reduction are the SAME as ggml_vec_dot_pq2_0_q8_0, so results are bit-identical to the +// vec_dot path. PQ2_SGEMM=0 in the environment disables this case (A/B and exactness checks). +#if defined(__AVX2__) +static inline __m256i pq2_dpbusd(__m256i u, __m256i s) { +#if defined(__AVX512VNNI__) && defined(__AVX512VL__) + return _mm256_dpbusd_epi32(_mm256_setzero_si256(), u, s); +#elif defined(__AVXVNNI__) + return _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), u, s); +#else + return _mm256_madd_epi16(_mm256_maddubs_epi16(u, s), _mm256_set1_epi16(1)); +#endif +} +static inline float pq2_hsum(__m256 x) { // identical reduction order to hsum_float_8 in arch/x86/quants.c + __m128 res = _mm256_extractf128_ps(x, 1); + res = _mm_add_ps(res, _mm256_castps256_ps128(x)); + res = _mm_add_ps(res, _mm_movehl_ps(res, res)); + res = _mm_add_ss(res, _mm_movehdup_ps(res)); + return _mm_cvtss_f32(res); +} +class tinyBLAS_PQ2_AVX { + public: + tinyBLAS_PQ2_AVX(int64_t k, const block_pq2_0 *A, int64_t lda, const block_q8_0 *B, int64_t ldb, float *C, int64_t ldc, int ith, int nth) + : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) {} + void matmul(int64_t m, int64_t n) { mnpack(0, m, 0, n); } + private: + void mnpack(int64_t m0, int64_t m, int64_t n0, int64_t n) { + int64_t mc, nc, mp, np; + switch ((MIN(m - m0, 4) << 4) | MIN(n - n0, 2)) { + case 0x42: mc = 4; nc = 2; gemm<4, 2>(m0, m, n0, n); break; + case 0x41: mc = 4; nc = 1; gemm<4, 1>(m0, m, n0, n); break; + case 0x32: mc = 3; nc = 2; gemm<3, 2>(m0, m, n0, n); break; + case 0x31: mc = 3; nc = 1; gemm<3, 1>(m0, m, n0, n); break; + case 0x22: mc = 2; nc = 2; gemm<2, 2>(m0, m, n0, n); break; + case 0x21: mc = 2; nc = 1; gemm<2, 1>(m0, m, n0, n); break; + case 0x12: mc = 1; nc = 2; gemm<1, 2>(m0, m, n0, n); break; + case 0x11: mc = 1; nc = 1; gemm<1, 1>(m0, m, n0, n); break; + default: return; + } + mp = m0 + (m - m0) / mc * mc; + np = n0 + (n - n0) / nc * nc; + mnpack(mp, m, n0, np); + mnpack(m0, m, np, n); + } + template + NOINLINE void gemm(int64_t m0, int64_t m, int64_t n0, int64_t n) { + const int64_t ytiles = (m - m0) / RM; + const int64_t xtiles = (n - n0) / RN; + const int64_t tiles = xtiles * ytiles; + const int64_t duty = (tiles + nth - 1) / nth; + const int64_t start = duty * ith; + int64_t end = start + duty; + if (end > tiles) end = tiles; + 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); + for (int64_t job = start; job < end; ++job) { + const int64_t ii = m0 + job / xtiles * RM; + const int64_t jj = n0 + job % xtiles * RN; + __m256 Cv[RN][RM] = {}; + for (int64_t l = 0; l < k; ++l) { + __m256 blk[RN][RM] = {}; + for (int s = 0; s < 4; ++s) { + __m256i codes[RM]; + for (int i = 0; i < RM; ++i) { + int64_t xq; memcpy(&xq, &A[lda * (ii + i) + l].qs[s * 8], sizeof(xq)); + codes[i] = _mm256_and_si256(_mm256_srlv_epi64(_mm256_set1_epi64x(xq), shifts), mask3); + } + for (int j = 0; j < RN; ++j) { + const block_q8_0 *yb = &B[ldb * (jj + j) + 4 * l + s]; + const __m256i qy = _mm256_loadu_si256((const __m256i *) yb->qs); + const __m256i qyp = _mm256_permutevar8x32_epi32(_mm256_shuffle_epi8(qy, qy_shuf), qy_perm); + const __m256i sy = pq2_dpbusd(ones_8, qyp); + const __m256 db = _mm256_set1_ps(unhalf(yb->d)); + for (int i = 0; i < RM; ++i) { + const __m256i s32 = _mm256_sub_epi32(pq2_dpbusd(codes[i], qyp), sy); + blk[j][i] = _mm256_fmadd_ps(db, _mm256_cvtepi32_ps(s32), blk[j][i]); + } + } + } + for (int i = 0; i < RM; ++i) { + const __m256 da = _mm256_set1_ps(unhalf(A[lda * (ii + i) + l].d)); + for (int j = 0; j < RN; ++j) Cv[j][i] = _mm256_fmadd_ps(da, blk[j][i], Cv[j][i]); + } + } + for (int j = 0; j < RN; ++j) + for (int i = 0; i < RM; ++i) + C[ldc * (jj + j) + (ii + i)] = pq2_hsum(Cv[j][i]); + } + } + const block_pq2_0 *const A; + const block_q8_0 *const B; + float *const C; + const int64_t k, lda, ldb, ldc; + const int ith, nth; +}; +#endif // __AVX2__ + /** * Performs optimized matrix multiplication on CPU. * @@ -4038,6 +4299,24 @@ bool llamafile_sgemm(const struct ggml_compute_params * params, int64_t m, int64 return false; } + case GGML_TYPE_PQ2_0: { +#if defined(__AVX2__) + { static int off = -1; if (off < 0) { const char * e = getenv("PQ2_SGEMM"); off = (e && e[0] == '0') ? 1 : 0; } if (off) return false; } + if (Btype == GGML_TYPE_Q8_K) { + tinyBLAS_PQ2K_AVX tb{ k, (const block_pq2_0 *)A, lda, (const block_q8_K *)B, ldb, (float *)C, ldc, params->ith, params->nth, params }; + tb.matmul(m, n); + return true; + } + if (Btype == GGML_TYPE_Q8_0) { + tinyBLAS_PQ2_AVX tb{ k, (const block_pq2_0 *)A, lda, (const block_q8_0 *)B, ldb, (float *)C, ldc, params->ith, params->nth }; + tb.matmul(m, n); + return true; + } + return false; +#else + return false; +#endif + } case GGML_TYPE_Q8_0: { if (Btype != GGML_TYPE_Q8_0) return false; diff --git a/ggml/src/ggml-cpu/quants.c b/ggml/src/ggml-cpu/quants.c index 3e275dc70bf3..e8a3067cfa6f 100644 --- a/ggml/src/ggml-cpu/quants.c +++ b/ggml/src/ggml-cpu/quants.c @@ -1461,3 +1461,36 @@ void quantize_row_iq4_xs(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, assert(k % QK_K == 0); quantize_iq4_xs(x, y, 1, k, NULL); } + +// PQ2_0 x Q8_K reference dot (tier 2, 2026-09-18): one activation scale per 256 elements instead of per 32, so a +// 128-weight block's integer dot can be accumulated in one int32 and scaled once. Scalar reference for the x86 +// kernel's tolerance test; also the non-x86 fallback. Two PQ2_0 blocks map onto one Q8_K block (halves). +void ggml_vec_dot_pq2_0_q8_K_generic(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; + 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); + int sumi = 0; + for (int k = 0; k < 4; k++) { + for (int b = 0; b < 8; b++) { + const uint8_t byte = x[i].qs[8 * k + b]; + for (int j = 0; j < 4; j++) { + sumi += (((byte >> (2 * j)) & 3) - 1) * q8[32 * k + 4 * b + j]; + } + } + } + sumf += (GGML_CPU_FP16_TO_FP32(x[i].d) * yb->d) * (float) sumi; + } + *s = sumf; +} diff --git a/ggml/src/ggml-cpu/quants.h b/ggml/src/ggml-cpu/quants.h index f00f35c3fe85..2981ec951d78 100644 --- a/ggml/src/ggml-cpu/quants.h +++ b/ggml/src/ggml-cpu/quants.h @@ -43,6 +43,7 @@ void quantize_row_iq4_xs (const float * GGML_RESTRICT x, void * GGML_RESTRICT y, 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); void ggml_vec_dot_q2_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); void ggml_vec_dot_pq2_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); +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); 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); void ggml_vec_dot_q4_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); void ggml_vec_dot_q4_1_q8_1(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); @@ -78,6 +79,7 @@ void quantize_row_q8_1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRI void quantize_row_q8_K_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); void ggml_vec_dot_q1_0_q8_0_generic(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); void ggml_vec_dot_pq2_0_q8_0_generic(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); +void ggml_vec_dot_pq2_0_q8_K_generic(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); void ggml_vec_dot_ptq1_0_q8_0_generic(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); void ggml_vec_dot_q2_0_q8_0_generic(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); void ggml_vec_dot_q4_0_q8_0_generic(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); From 10df298811c35948ab31f3fa0e2d9032809fe339 Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Sun, 20 Sep 2026 11:36:43 +0300 Subject: [PATCH 2/3] ggml-cpu, llama-quant: fix PQ2_0 non-x86 link break and 128-column shape contract Addresses both findings from the review of 164c3370. P1 - non-x86 dispatch. Switching PQ2_0's vec_dot to Q8_K left ggml_vec_dot_pq2_0_q8_K referenced unconditionally from type_traits_cpu while the symbol is only defined in arch/x86/quants.c. The portable implementation exists as ggml_vec_dot_pq2_0_q8_K_generic in quants.c but had no alias in arch-fallback.h, so every non-x86 target linked against a missing symbol. Adds the alias to all seven non-x86 sections (generic, arm, powerpc, loongarch, riscv, s390x, wasm), matching the existing _q8_0 pattern. ARM is aliased next to its ptq1_0 entry since it has a native pq2_0_q8_0 and so lacks the usual anchor line. Verified by cross-compiling for aarch64 (gcc-aarch64-linux-gnu) before and after. GNU ld tolerates undefined symbols in a shared object, so the failure reproduces in the symbol table rather than at link time: before: U ggml_vec_dot_pq2_0_q8_K after: T ggml_vec_dot_pq2_0_q8_K (0x4d504) That undefined symbol is what the macOS linker reports as "Undefined symbols for architecture arm64". x86-64 is unaffected and still resolves to the AVX2/AVX-VNNI implementation. P2 - 128-column rows. PQ2_0 blocks are 128 wide, so a row that is a multiple of 128 is valid for the weight type, but the Q8_K activations are allocated and quantized in 256-element blocks. ggml_row_size rounds down, so such a row sizes its activation buffer to zero bytes: ncols row_size(PQ2_0) row_size(Q8_K) 128 34 0 <- corrupt 256 68 292 384 102 0 <- corrupt 512 136 584 tensor_type_fallback compared ncols against ggml_blck_size(target_type), which is 128 for PQ2_0, so these tensors passed the shape check and were quantized to PQ2_0 anyway. It now compares against the activation granularity for PQ2_0, so they fall back to q4_0 like the other 256-block types already do. Verified on a synthetic model whose rows are all 128 wide: before: "converting to pq2_0" for all seven 128-column tensors, no warning after: "ncols 128 not divisible by 256 (required for type pq2_0) -> falling back to q4_0", 9 of 12 tensors reported as fallbacks Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- ggml/src/ggml-cpu/arch-fallback.h | 14 ++++++++++++++ src/llama-quant.cpp | 10 +++++++++- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-cpu/arch-fallback.h b/ggml/src/ggml-cpu/arch-fallback.h index d283c2ed51be..e116558224b4 100644 --- a/ggml/src/ggml-cpu/arch-fallback.h +++ b/ggml/src/ggml-cpu/arch-fallback.h @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 121e610abf5a..4c5313a2f966 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -375,7 +375,15 @@ static ggml_type tensor_type_fallback(quantize_state_impl & qs, const ggml_tenso ggml_type return_type = target_type; const int64_t ncols = t->ne[0]; - const int64_t qk_k = ggml_blck_size(target_type); + int64_t qk_k = ggml_blck_size(target_type); + + // PQ2_0 weights are stored in 128-wide blocks, but the CPU vec_dot consumes Q8_K + // activations, which are allocated and quantized in 256-element blocks. A row that is + // only a multiple of 128 would size the activation row to zero bytes + // (ggml_row_size rounds down), so the shape must satisfy the activation granularity. + if (target_type == GGML_TYPE_PQ2_0) { + qk_k = ggml_blck_size(GGML_TYPE_Q8_K); + } if (ncols % qk_k != 0) { // this tensor's shape is incompatible with this quant LLAMA_LOG_WARN("warning: %-36s - ncols %6" PRId64 " not divisible by %3" PRId64 " (required for type %7s) ", From d2c9ddd053aa97ed19f9705cfccd487378d11950 Mon Sep 17 00:00:00 2001 From: abdulaziz Date: Mon, 21 Sep 2026 20:39:12 +0300 Subject: [PATCH 3/3] ggml-cpu: keep PQ2_0 rows that are not multiples of 256 on the Q8_0 activation path P2 from the review of 10df29881. The Q8_K activation switch was applied to every PQ2_0 tensor via the static trait table, but PQ2_0 stores 128 weights per block, so 128, 384, 640, ... are all legal row widths and appear in existing models. Q8_K only exists in 256-element blocks, so those rows cannot be expressed in it. Guarding the quantizer, as the previous commit did, does not help: tensors inside an already published model never pass through tensor_type_fallback. That guard is reverted here - it only cost PQ2_0 coverage on shapes that are now handled. Selection is now shape-aware. ggml_cpu_vec_dot_type / ggml_cpu_vec_dot return Q8_0 + ggml_vec_dot_pq2_0_q8_0 when ne[0] % QK_K != 0 and the trait-table entry otherwise, and are used at all six sites that previously read the table directly, including both ggml_graph_plan work-size computations - a mismatch there would size the activation buffer for a different type than the compute path consumes. llamafile_sgemm already dispatches on Btype, so those rows also pick up tinyBLAS_PQ2_AVX rather than falling back to the scalar dot. Rows that are multiples of 256 - which is every PQ2_0 tensor in a typical model - take an unchanged code path and keep the faster Q8_K kernels. New test tests/test-pq2-row-shapes.cpp exercises pre-existing PQ2_0 tensors rather than the quantizer: it reference-encodes a PQ2_0 weight tensor the way a .gguf holds one, runs ggml_mul_mat through ggml_graph_compute_with_ctx for K in {128, 256, 384, 512, 640} x N in {1, 4, 16}, and checks the result against dequantized weights times exact float activations. Before (Release, x86-64 AVX2 + AVX-VNNI) - the assert the report hit is not enabled in this configuration, so the failure is silently wrong output and reads past the activation buffer: K= 128 N= 1 FAIL rel_rms=1.0000 K= 384 N= 1 FAIL rel_rms=0.4793 K= 128 N= 4 FAIL rel_rms=1.0000 K= 384 N= 4 FAIL rel_rms=0.9874 K= 128 N=16 FAIL rel_rms=1.0000 K= 384 N=16 FAIL rel_rms=4.3e36 K= 256 all OK rel_rms<=0.0053 K= 640 all FAIL up to rel_rms=3.3e36 K= 512 all OK rel_rms<=0.0068 After: 15/15 OK, rel_rms <= 0.0068. The K=256 and K=512 figures are unchanged to all digits, so the fast path is untouched. Regression checks: test-backend-ops -o MUL_MAT -b CPU 1294/1294 test-backend-ops -o MUL_MAT_ID -b CPU 1022/1022 test-quantize-fns pass llama-bench Bonsai-2-27B-PQ2_0 tg32 3.97 +/- 0.09 t/s (PR range 3.9-4.2) That model's 402 PQ2_0 tensors are all 256-multiples, so its graph provably does not change; the branch is evaluated once per matmul node, not per row. aarch64 cross-build (gcc-aarch64-linux-gnu), all defined, none undefined: T ggml_vec_dot_pq2_0_q8_0 T ggml_vec_dot_pq2_0_q8_0_generic T ggml_vec_dot_pq2_0_q8_K Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R --- ggml/src/ggml-cpu/ggml-cpu.c | 43 ++++++++++--- src/llama-quant.cpp | 10 +-- tests/CMakeLists.txt | 1 + tests/test-pq2-row-shapes.cpp | 115 ++++++++++++++++++++++++++++++++++ 4 files changed, 152 insertions(+), 17 deletions(-) create mode 100644 tests/test-pq2-row-shapes.cpp diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 48c3f5f0890f..f63c87b03ca0 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -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, @@ -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; @@ -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; @@ -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; @@ -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 @@ -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])); @@ -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) { diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 4c5313a2f966..121e610abf5a 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -375,15 +375,7 @@ static ggml_type tensor_type_fallback(quantize_state_impl & qs, const ggml_tenso ggml_type return_type = target_type; const int64_t ncols = t->ne[0]; - int64_t qk_k = ggml_blck_size(target_type); - - // PQ2_0 weights are stored in 128-wide blocks, but the CPU vec_dot consumes Q8_K - // activations, which are allocated and quantized in 256-element blocks. A row that is - // only a multiple of 128 would size the activation row to zero bytes - // (ggml_row_size rounds down), so the shape must satisfy the activation granularity. - if (target_type == GGML_TYPE_PQ2_0) { - qk_k = ggml_blck_size(GGML_TYPE_Q8_K); - } + const int64_t qk_k = ggml_blck_size(target_type); if (ncols % qk_k != 0) { // this tensor's shape is incompatible with this quant LLAMA_LOG_WARN("warning: %-36s - ncols %6" PRId64 " not divisible by %3" PRId64 " (required for type %7s) ", diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 67ff970be0b3..3e5874ce61c4 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -334,6 +334,7 @@ if (NOT GGML_BACKEND_DL) llama_build_and_test(test-quantize-fns.cpp) llama_build_and_test(test-ptq1_0-element-map.cpp) llama_build_and_test(test-ptq1_0-cuda-dot.cpp) +llama_build_and_test(test-pq2-row-shapes.cpp) llama_build_and_test(test-quantize-perf.cpp) llama_build_and_test(test-rope.cpp) llama_build_and_test(test-col2im-1d.cpp) diff --git a/tests/test-pq2-row-shapes.cpp b/tests/test-pq2-row-shapes.cpp new file mode 100644 index 000000000000..b2fc23de4ea9 --- /dev/null +++ b/tests/test-pq2-row-shapes.cpp @@ -0,0 +1,115 @@ +// PQ2_0 rows are 128 wide, so a model may legitimately hold tensors whose row length is any +// multiple of 128 -- 128 and 384 as much as 256. The AVX2 kernels added for PQ2_0 consume Q8_K +// activations, which only exist in 256-element blocks, so routing every PQ2_0 matmul through +// Q8_K breaks exactly those rows. Guarding the quantizer is not enough: tensors already inside +// published models never pass through it. +// +// This exercises pre-existing PQ2_0 tensors directly (reference-encoded, as a model file holds +// them) for K in {128, 256, 384}, single-token and batched, and checks the numbers rather than +// just survival. + +#include "ggml.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include +#include +#include + +static std::vector make_weights(int64_t n, uint32_t seed) { + std::mt19937 rng(seed); + std::normal_distribution dist(0.0f, 0.5f); + std::vector v(n); + for (int64_t i = 0; i < n; i++) { + v[i] = dist(rng); + } + return v; +} + +// rms(computed - reference) / rms(reference), with the reference taken from the dequantized +// weights and exact float activations. Activation quantization (Q8_K or Q8_0) is the only +// remaining error term, so this stays far below the threshold when the path is correct and +// explodes when it is not. +static double rel_rms(const std::vector & got, const std::vector & ref) { + double se = 0.0; + double sr = 0.0; + for (size_t i = 0; i < ref.size(); i++) { + const double d = (double) got[i] - (double) ref[i]; + se += d * d; + sr += (double) ref[i] * (double) ref[i]; + } + if (sr == 0.0) { + return se == 0.0 ? 0.0 : INFINITY; + } + return std::sqrt(se / sr); +} + +static bool run_shape(int64_t K, int64_t M, int64_t N) { + const size_t mem_size = 64u * 1024 * 1024; + struct ggml_init_params ip = { mem_size, nullptr, false }; + struct ggml_context * ctx = ggml_init(ip); + if (!ctx) { + printf(" K=%4lld M=%lld N=%lld : ggml_init failed\n", (long long) K, (long long) M, (long long) N); + return false; + } + + // A pre-existing PQ2_0 weight tensor: encoded by the reference row quantizer, which is what + // a .gguf holds. It never sees llama-quant.cpp's type fallback. + struct ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_PQ2_0, K, M); + const std::vector w = make_weights(K * M, 1234 + (uint32_t) K); + ggml_quantize_chunk(GGML_TYPE_PQ2_0, w.data(), a->data, 0, M, K, nullptr); + + struct ggml_tensor * b = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, K, N); + const std::vector x = make_weights(K * N, 99); + memcpy(b->data, x.data(), x.size() * sizeof(float)); + + struct ggml_tensor * c = ggml_mul_mat(ctx, a, b); + struct ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, c); + ggml_graph_compute_with_ctx(ctx, gf, 4); + + // reference: dequantized weights x exact float activations + std::vector a_deq(K * M); + const ggml_type_traits * tt = ggml_get_type_traits(GGML_TYPE_PQ2_0); + tt->to_float(a->data, a_deq.data(), K * M); + + std::vector ref(M * N); + for (int64_t j = 0; j < N; j++) { + for (int64_t i = 0; i < M; i++) { + double s = 0.0; + for (int64_t k = 0; k < K; k++) { + s += (double) a_deq[i * K + k] * (double) x[j * K + k]; + } + ref[j * M + i] = (float) s; + } + } + + std::vector got(M * N); + memcpy(got.data(), c->data, got.size() * sizeof(float)); + + const double err = rel_rms(got, ref); + const bool ok = err < 0.05; + + printf(" K=%4lld M=%lld N=%2lld : %-4s rel_rms=%.4f\n", + (long long) K, (long long) M, (long long) N, ok ? "OK" : "FAIL", err); + + ggml_free(ctx); + return ok; +} + +int main() { + printf("test-pq2-row-shapes: PQ2_0 matmul over every legal row width\n"); + + bool all_ok = true; + // 128 and 384 are legal PQ2_0 row widths that are NOT multiples of the Q8_K block (256). + for (int64_t K : { (int64_t) 128, (int64_t) 256, (int64_t) 384, (int64_t) 512, (int64_t) 640 }) { + for (int64_t N : { (int64_t) 1, (int64_t) 4, (int64_t) 16 }) { + all_ok &= run_shape(K, 4, N); + } + } + + printf("%s\n", all_ok ? "PASS" : "FAIL"); + return all_ok ? 0 : 1; +}