repack + gemm

ggml/src/ggml-cpu/arch-fallback.h

@@ -96,7 +96,8 @@
 #define ggml_gemv_iq4_nl_4x4_q8_0_generic ggml_gemv_iq4_nl_4x4_q8_0
 #define ggml_gemv_mxfp4_4x4_q8_0_generic ggml_gemv_mxfp4_4x4_q8_0
 #define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0
-#define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0
+// ggml_gemv/gemm_q8_0_4x8_q8_0: x86 supplies AVX-512 VNNI implementations in
+// arch/x86/repack.cpp (transcribe.cpp) — NOT remapped to generic here.
 #define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0
 #define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0
 #define ggml_gemm_q4_K_8x4_q8_K_generic ggml_gemm_q4_K_8x4_q8_K
@@ -107,7 +108,6 @@
 #define ggml_gemm_iq4_nl_4x4_q8_0_generic ggml_gemm_iq4_nl_4x4_q8_0
 #define ggml_gemm_mxfp4_4x4_q8_0_generic ggml_gemm_mxfp4_4x4_q8_0
 #define ggml_gemm_q8_0_4x4_q8_0_generic ggml_gemm_q8_0_4x4_q8_0
-#define ggml_gemm_q8_0_4x8_q8_0_generic ggml_gemm_q8_0_4x8_q8_0
 #elif defined(__POWERPC__) || defined(__powerpc__)
 // ref: https://github.com/ggml-org/llama.cpp/pull/14146#issuecomment-2972561679
 // quants.c

ggml/src/ggml-cpu/arch/x86/repack.cpp

@@ -1445,6 +1445,112 @@ static void gemm_q4_b32_8x8_q8_0_lut_avx(int n, float * GGML_RESTRICT s, size_t
 
 #endif // defined(__AVX2__) || defined(__AVX512F__)
 
+// ===========================================================================
+// transcribe.cpp: AVX-512 VNNI (AVX2 fallback) Q8_0 x4 blocked GEMM/GEMV.
+//
+// Upstream ggml has no x86 blocked GEMM for Q8_0 weights (only Q4_0/Q4_K/IQ4_NL),
+// so Q8_0 matmuls fall back to the per-row vec_dot — ~2x slower than a tuned int8
+// GEMM (the ONNX/MLAS gap). These mirror the q8_0_4x8 generic layout exactly
+// (weight block_q8_0x4 qs[k*32 + j*8 + i], 4 cols, blocklen 8; GEMM activation
+// block_q8_0x4 qs[k*32 + m*8 + i]; GEMV activation plain block_q8_0 qs[k*8 + i])
+// and are simpler than the Q4_0 kernels — no nibble unpack/LUT.
+//
+// s8*s8 via VNNI dpbusd (which is u8*s8): dpbusd(|a|, b*sign(a)) == sum(a*b),
+// using AVX2 _mm256_sign_epi8 for |a| and the signed b; falls back to
+// maddubs+madd on AVX2-without-VNNI (no saturation: |a|<=127,|b|<=127 ->
+// pair sum <= 32258 < 32767). The 8 int32 lanes from one 32-byte (4col x 8)
+// dpbusd are pair-reduced to 4 per-column sums.
+#if defined(__AVX2__) || defined(__AVX512F__)
+static inline __m256i q8_0_x4_acc_dot(__m256i acc, __m256i a /*s8 act*/, __m256i b /*s8 weight*/) {
+    const __m256i ax = _mm256_sign_epi8(a, a); // |a|  -> u8
+    const __m256i sy = _mm256_sign_epi8(b, a); // b * sign(a)  -> s8
+#if defined(__AVX512VNNI__) && defined(__AVX512VL__)
+    return _mm256_dpbusd_epi32(acc, ax, sy);
+#elif defined(__AVXVNNI__)
+    return _mm256_dpbusd_avx_epi32(acc, ax, sy);
+#else
+    const __m256i p16 = _mm256_maddubs_epi16(ax, sy);
+    return _mm256_add_epi32(acc, _mm256_madd_epi16(_mm256_set1_epi16(1), p16));
+#endif
+}
+// acc lanes [c0a,c0b,c1a,c1b,c2a,c2b,c3a,c3b] -> [c0,c1,c2,c3]
+static inline __m128i q8_0_x4_reduce(__m256i acc) {
+    const __m256i h = _mm256_hadd_epi32(acc, acc); // lo:[c0,c1,c0,c1] hi:[c2,c3,c2,c3]
+    return _mm_unpacklo_epi64(_mm256_castsi256_si128(h), _mm256_extracti128_si256(h, 1));
+}
+#endif
+
+void ggml_gemv_q8_0_4x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) {
+#if defined(__AVX2__) || defined(__AVX512F__)
+    {
+        const int nb = n / QK8_0;
+        const block_q8_0 * GGML_RESTRICT a_ptr = (const block_q8_0 *) vy;
+        for (int x = 0; x < nc / 4; x++) {
+            const block_q8_0x4 * GGML_RESTRICT b_ptr = (const block_q8_0x4 *) vx + x * nb;
+            __m128 sumf = _mm_setzero_ps();
+            for (int l = 0; l < nb; l++) {
+                __m256i acc = _mm256_setzero_si256();
+                for (int k = 0; k < 4; k++) {
+                    const __m256i W = _mm256_loadu_si256((const __m256i *)(b_ptr[l].qs + k * 32));
+                    int64_t am; memcpy(&am, a_ptr[l].qs + k * 8, 8);
+                    acc = q8_0_x4_acc_dot(acc, _mm256_set1_epi64x(am), W);
+                }
+                const __m128 db = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *) b_ptr[l].d));
+                const __m128 cf = _mm_cvtepi32_ps(q8_0_x4_reduce(acc));
+                sumf = _mm_fmadd_ps(cf, _mm_mul_ps(db, _mm_set1_ps(GGML_CPU_FP16_TO_FP32(a_ptr[l].d))), sumf);
+            }
+            _mm_storeu_ps(s + x * 4, sumf);
+        }
+        return;
+    }
+#endif
+    ggml_gemv_q8_0_4x8_q8_0_generic(n, s, bs, vx, vy, nr, nc);
+}
+
+void ggml_gemm_q8_0_4x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) {
+#if defined(__AVX2__) || defined(__AVX512F__)
+    {
+        const int nb = n / QK8_0;
+        for (int y = 0; y < nr / 4; y++) {
+            const block_q8_0x4 * GGML_RESTRICT a_ptr = (const block_q8_0x4 *) vy + y * nb;
+            for (int x = 0; x < nc / 4; x++) {
+                const block_q8_0x4 * GGML_RESTRICT b_ptr = (const block_q8_0x4 *) vx + x * nb;
+                __m128 sumf0 = _mm_setzero_ps(), sumf1 = _mm_setzero_ps();
+                __m128 sumf2 = _mm_setzero_ps(), sumf3 = _mm_setzero_ps();
+                for (int l = 0; l < nb; l++) {
+                    __m256i acc0 = _mm256_setzero_si256(), acc1 = _mm256_setzero_si256();
+                    __m256i acc2 = _mm256_setzero_si256(), acc3 = _mm256_setzero_si256();
+                    for (int k = 0; k < 4; k++) {
+                        const __m256i W = _mm256_loadu_si256((const __m256i *)(b_ptr[l].qs + k * 32));
+                        int64_t a0, a1, a2, a3;
+                        memcpy(&a0, a_ptr[l].qs + k * 32 + 0 * 8, 8);
+                        memcpy(&a1, a_ptr[l].qs + k * 32 + 1 * 8, 8);
+                        memcpy(&a2, a_ptr[l].qs + k * 32 + 2 * 8, 8);
+                        memcpy(&a3, a_ptr[l].qs + k * 32 + 3 * 8, 8);
+                        acc0 = q8_0_x4_acc_dot(acc0, _mm256_set1_epi64x(a0), W);
+                        acc1 = q8_0_x4_acc_dot(acc1, _mm256_set1_epi64x(a1), W);
+                        acc2 = q8_0_x4_acc_dot(acc2, _mm256_set1_epi64x(a2), W);
+                        acc3 = q8_0_x4_acc_dot(acc3, _mm256_set1_epi64x(a3), W);
+                    }
+                    const __m128 db = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *) b_ptr[l].d));
+                    float da[4]; _mm_storeu_ps(da, _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *) a_ptr[l].d)));
+                    sumf0 = _mm_fmadd_ps(_mm_cvtepi32_ps(q8_0_x4_reduce(acc0)), _mm_mul_ps(db, _mm_set1_ps(da[0])), sumf0);
+                    sumf1 = _mm_fmadd_ps(_mm_cvtepi32_ps(q8_0_x4_reduce(acc1)), _mm_mul_ps(db, _mm_set1_ps(da[1])), sumf1);
+                    sumf2 = _mm_fmadd_ps(_mm_cvtepi32_ps(q8_0_x4_reduce(acc2)), _mm_mul_ps(db, _mm_set1_ps(da[2])), sumf2);
+                    sumf3 = _mm_fmadd_ps(_mm_cvtepi32_ps(q8_0_x4_reduce(acc3)), _mm_mul_ps(db, _mm_set1_ps(da[3])), sumf3);
+                }
+                _mm_storeu_ps(s + (y * 4 + 0) * bs + x * 4, sumf0);
+                _mm_storeu_ps(s + (y * 4 + 1) * bs + x * 4, sumf1);
+                _mm_storeu_ps(s + (y * 4 + 2) * bs + x * 4, sumf2);
+                _mm_storeu_ps(s + (y * 4 + 3) * bs + x * 4, sumf3);
+            }
+        }
+        return;
+    }
+#endif
+    ggml_gemm_q8_0_4x8_q8_0_generic(n, s, bs, vx, vy, nr, nc);
+}
+
 void ggml_gemv_q4_0_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) {
 #if defined(__AVX2__) || defined(__AVX512F__)
     {

ggml/src/ggml-cpu/cmake/FindSIMD.cmake

@@ -40,6 +40,23 @@ set(AVX2_CODE "
     }
 ")
 
+# AVX-512 VNNI: MSVC has no flag to enable only this subset and does not define
+# __AVX512VNNI__, so it must be probed and the macro set manually (mirrors how
+# ggml-cpu/CMakeLists.txt handles GGML_AVX512_VNNI). check_c_source_runs both
+# COMPILES the dpbusd intrinsic (needs /arch:AVX512) and RUNS it (confirms the
+# build host actually has VNNI — SIGILL-safe, stays off when absent / cross-building).
+set(AVX512_VNNI_CODE "
+    #include <immintrin.h>
+    int main()
+    {
+        __m512i acc = _mm512_setzero_si512();
+        __m512i a   = _mm512_setzero_si512();
+        __m512i b   = _mm512_setzero_si512();
+        acc = _mm512_dpbusd_epi32(acc, a, b);
+        return _mm_cvtsi128_si32(_mm512_castsi512_si128(acc));
+    }
+")
+
 set(FMA_CODE "
     #include <immintrin.h>
     int main()
@@ -97,4 +114,11 @@ if (NOT ${AVX512_FOUND})
     set(GGML_AVX512 OFF)
 else()
     set(GGML_AVX512 ON)
+    # Probe the VNNI subset (dpbusd) only when the AVX-512 base is present.
+    # Enabling it lets ggml's int8 dot + the transcribe.cpp Q8_0 blocked GEMM
+    # use _mm*_dpbusd_epi32 instead of the slower maddubs+madd path.
+    check_sse("AVX512_VNNI" " ;/arch:AVX512")
+    if (${AVX512_VNNI_FOUND})
+        set(GGML_AVX512_VNNI ON)
+    endif()
 endif()

ggml/src/ggml-cpu/ggml-cpu.cpp

@@ -432,7 +432,14 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st
     // note: only the first sources are checked for extra buffer types to reduce overhead, increase if necessary
     for (int i = 0; i < 4; i++) {
         if (op->src[i] && op->src[i]->buffer &&
-            ggml_backend_cpu_is_extra_buffer_type(op->src[i]->buffer->buft)) {
+            ggml_backend_cpu_is_extra_buffer_type(op->src[i]->buffer->buft) &&
+            // Only delegate to the extra buffer's op support when THIS source was
+            // actually converted by it (tensor->extra holds the repack traits). A
+            // pass-through tensor that merely happens to live in the extra buffer
+            // (e.g. an F32 conv weight in a mixed-dtype weight set placed wholesale
+            // in CPU_REPACK) must fall through to the normal CPU op support, or ops
+            // like IM2COL on that weight would be wrongly rejected.
+            op->src[i]->extra != nullptr) {
             auto * buf_extra = (ggml::cpu::extra_buffer_type *) op->src[i]->buffer->buft->context;
             return buf_extra->supports_op(dev, op);
         }

ggml/src/ggml-cpu/repack.cpp

@@ -4697,6 +4697,11 @@ static const ggml::cpu::tensor_traits * ggml_repack_get_optimal_repack_type(cons
             }
         }
     } else if (cur->type == GGML_TYPE_Q8_0) {
+        if (ggml_cpu_has_avx2()) {
+            if (cur->ne[1] % 4 == 0) {
+                return &q8_0_4x8_q8_0;
+            }
+        }
         if (ggml_cpu_has_neon() && ggml_cpu_has_matmul_int8()) {
             if (cur->ne[1] % 4 == 0) {
                 return &q8_0_4x8_q8_0;
@@ -4736,12 +4741,32 @@ static void ggml_backend_cpu_repack_buffer_set_tensor(ggml_backend_buffer_t buff
     GGML_ASSERT(size == ggml_nbytes(tensor));
 
     auto tensor_traits = (ggml::cpu::repack::tensor_traits_base *) tensor->extra;
+    if (tensor_traits == nullptr) {
+        // Not a repackable tensor (no optimal repack type for it): store the
+        // bytes verbatim, exactly like the plain CPU buffer. This lets a
+        // consumer place a whole mixed-dtype weight set in the repack buffer —
+        // only the matmul weights with a repack type get interleaved; norms,
+        // biases, F32/F16 tensors pass through untouched.
+        memcpy((char *) tensor->data + offset, data, size);
+        return;
+    }
     auto OK            = tensor_traits->repack(tensor, data, size);
 
     GGML_ASSERT(OK == 0);
     GGML_UNUSED(buffer);
 }
 
+static void ggml_backend_cpu_repack_buffer_get_tensor(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor,
+                                                       void * data, size_t offset, size_t size) {
+    // Readback is only meaningful for pass-through (non-repacked) tensors; a
+    // repacked tensor's bytes are in an interleaved layout the caller can't
+    // interpret (and during the CPU_REPACK measurement hack the decoder reads
+    // some weights back — we let it read raw interleaved bytes rather than
+    // abort, since the transcript is already declared invalid in that mode).
+    memcpy(data, (const char *) tensor->data + offset, size);
+    GGML_UNUSED(buffer);
+}
+
 static const char * ggml_backend_cpu_repack_buffer_type_get_name(ggml_backend_buffer_type_t buft) {
     return "CPU_REPACK";
 
@@ -4758,7 +4783,7 @@ static ggml_backend_buffer_t ggml_backend_cpu_repack_buffer_type_alloc_buffer(gg
     buffer->buft              = buft;
     buffer->iface.init_tensor = ggml_backend_cpu_repack_buffer_init_tensor;
     buffer->iface.set_tensor  = ggml_backend_cpu_repack_buffer_set_tensor;
-    buffer->iface.get_tensor  = nullptr;
+    buffer->iface.get_tensor  = ggml_backend_cpu_repack_buffer_get_tensor;
     buffer->iface.cpy_tensor  = nullptr;
     return buffer;
 }