ggml : add AVX1 Q4_K kernels and prefill

Assisted-by: OpenAI Codex

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

@@ -94,8 +94,6 @@
 #define ggml_gemv_q1_0_4x4_q8_0_generic ggml_gemv_q1_0_4x4_q8_0
 #define ggml_gemv_q1_0_4x8_q8_0_generic ggml_gemv_q1_0_4x8_q8_0
 #define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0
-#define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0
-#define ggml_gemv_q4_K_8x4_q8_K_generic ggml_gemv_q4_K_8x4_q8_K
 #define ggml_gemv_q5_K_8x4_q8_K_generic ggml_gemv_q5_K_8x4_q8_K
 #define ggml_gemv_q5_K_8x8_q8_K_generic ggml_gemv_q5_K_8x8_q8_K
 #define ggml_gemv_q6_K_8x4_q8_K_generic ggml_gemv_q6_K_8x4_q8_K
@@ -108,7 +106,6 @@
 #define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0
 #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
 #define ggml_gemm_q5_K_8x4_q8_K_generic ggml_gemm_q5_K_8x4_q8_K
 #define ggml_gemm_q5_K_8x8_q8_K_generic ggml_gemm_q5_K_8x8_q8_K
 #define ggml_gemm_q6_K_8x4_q8_K_generic ggml_gemm_q6_K_8x4_q8_K

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

@@ -175,6 +175,171 @@ static inline __m256i mul_sum_i8_pairs_acc_int32x8(const __m256i acc, const __m2
 }
 #endif
 
+void ggml_gemv_q4_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) {
+    const int qk = QK8_0;
+    const int nb = n / qk;
+    const int ncols_interleaved = 4;
+
+    assert(n % qk == 0);
+    assert(nc % ncols_interleaved == 0);
+    GGML_ASSERT(nr == 1);
+
+    UNUSED(bs);
+    UNUSED(nr);
+
+#if defined(__AVX__) && defined(__SSSE3__) && defined(__SSE4_1__)
+    const __m128i m4b = _mm_set1_epi8(0x0F);
+    const __m128i sign = _mm_set1_epi8(0x08);
+    const __m128i ones = _mm_set1_epi16(1);
+
+    const block_q4_0x4 * b_ptr = (const block_q4_0x4 *) vx;
+
+    for (int c = 0; c < nc; c += ncols_interleaved) {
+        const block_q8_0 * a_ptr = (const block_q8_0 *) vy;
+        __m128 acc = _mm_setzero_ps();
+
+        for (int b = 0; b < nb; ++b) {
+            const __m128i b0 = _mm_loadu_si128((const __m128i *) (b_ptr->qs));
+            const __m128i b1 = _mm_loadu_si128((const __m128i *) (b_ptr->qs + 16));
+            const __m128i b2 = _mm_loadu_si128((const __m128i *) (b_ptr->qs + 32));
+            const __m128i b3 = _mm_loadu_si128((const __m128i *) (b_ptr->qs + 48));
+
+            const __m128i a0 = _mm_unpacklo_epi64(_mm_loadl_epi64((const __m128i *) (a_ptr->qs)), _mm_loadl_epi64((const __m128i *) (a_ptr->qs)));
+            const __m128i a1 = _mm_unpacklo_epi64(_mm_loadl_epi64((const __m128i *) (a_ptr->qs + 8)), _mm_loadl_epi64((const __m128i *) (a_ptr->qs + 8)));
+            const __m128i a2 = _mm_unpacklo_epi64(_mm_loadl_epi64((const __m128i *) (a_ptr->qs + 16)), _mm_loadl_epi64((const __m128i *) (a_ptr->qs + 16)));
+            const __m128i a3 = _mm_unpacklo_epi64(_mm_loadl_epi64((const __m128i *) (a_ptr->qs + 24)), _mm_loadl_epi64((const __m128i *) (a_ptr->qs + 24)));
+
+            const auto unpack = [m4b, sign](__m128i x) {
+                x = _mm_and_si128(x, m4b);
+                return _mm_sub_epi8(_mm_xor_si128(x, sign), sign);
+            };
+            const auto dot = [ones](__m128i x, __m128i y) {
+                const __m128i ax = _mm_sign_epi8(x, x);
+                const __m128i sy = _mm_sign_epi8(y, x);
+                return _mm_madd_epi16(_mm_maddubs_epi16(ax, sy), ones);
+            };
+
+            __m128i ret00 = dot(unpack(b0), a0);
+            ret00 = _mm_add_epi32(ret00, dot(unpack(_mm_srli_epi16(b0, 4)), a2));
+            __m128i ret01 = dot(unpack(b2), a1);
+            ret01 = _mm_add_epi32(ret01, dot(unpack(_mm_srli_epi16(b2, 4)), a3));
+
+            __m128i ret10 = dot(unpack(b1), a0);
+            ret10 = _mm_add_epi32(ret10, dot(unpack(_mm_srli_epi16(b1, 4)), a2));
+            __m128i ret11 = dot(unpack(b3), a1);
+            ret11 = _mm_add_epi32(ret11, dot(unpack(_mm_srli_epi16(b3, 4)), a3));
+
+            ret00 = _mm_hadd_epi32(ret00, ret00);
+            ret01 = _mm_hadd_epi32(ret01, ret01);
+            ret10 = _mm_hadd_epi32(ret10, ret10);
+            ret11 = _mm_hadd_epi32(ret11, ret11);
+            const __m128 sums0 = _mm_cvtepi32_ps(_mm_unpacklo_epi64(ret00, ret10));
+            const __m128 sums1 = _mm_cvtepi32_ps(_mm_unpacklo_epi64(ret01, ret11));
+
+#if defined(__F16C__)
+            const __m128 bd = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *) b_ptr->d));
+#else
+            const __m128 bd = _mm_set_ps(
+                    GGML_CPU_FP16_TO_FP32(b_ptr->d[3]),
+                    GGML_CPU_FP16_TO_FP32(b_ptr->d[2]),
+                    GGML_CPU_FP16_TO_FP32(b_ptr->d[1]),
+                    GGML_CPU_FP16_TO_FP32(b_ptr->d[0]));
+#endif
+            const __m128 ad = _mm_set1_ps(GGML_CPU_FP16_TO_FP32(a_ptr->d));
+            acc = _mm_add_ps(acc, _mm_mul_ps(_mm_mul_ps(sums0, bd), ad));
+            acc = _mm_add_ps(acc, _mm_mul_ps(_mm_mul_ps(sums1, bd), ad));
+
+            ++a_ptr;
+            ++b_ptr;
+        }
+
+        _mm_storeu_ps(s, acc);
+        s += ncols_interleaved;
+    }
+    return;
+#endif
+
+    ggml_gemv_q4_0_4x8_q8_0_generic(n, s, bs, vx, vy, nr, nc);
+}
+
+#if defined(__AVX__)
+static inline void quantize_q8_K_row_avx1(
+        const float * GGML_RESTRICT x,
+        float * GGML_RESTRICT d,
+        int8_t * GGML_RESTRICT qs) {
+    float amax = 0.0f;
+    float max = 0.0f;
+
+    // Keep the scalar maximum scan to preserve the reference implementation's
+    // signed first-maximum tie behavior exactly.
+    for (int i = 0; i < QK_K; ++i) {
+        const float value = x[i];
+        const float abs_value = fabsf(value);
+        if (amax < abs_value) {
+            amax = abs_value;
+            max = value;
+        }
+    }
+
+    const float iscale = amax ? -127.0f / max : 0.0f;
+    *d = amax ? 1.0f / iscale : 0.0f;
+    const __m256 scale = _mm256_set1_ps(iscale);
+
+    for (int i = 0; i < QK_K; i += 8) {
+        const __m256 values = _mm256_loadu_ps(x + i);
+        const __m256 rounded = _mm256_round_ps(
+                _mm256_mul_ps(values, scale),
+                _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
+        const __m256i q32 = _mm256_cvtps_epi32(rounded);
+        const __m128i q16 = _mm_packs_epi32(
+                _mm256_castsi256_si128(q32),
+                _mm256_extractf128_si256(q32, 1));
+        const __m128i q8 = _mm_packs_epi16(q16, _mm_setzero_si128());
+        _mm_storel_epi64((__m128i *) (qs + i), q8);
+    }
+}
+#endif
+
+void ggml_quantize_mat_q8_K_4x4(
+        const float * GGML_RESTRICT x,
+        void * GGML_RESTRICT vy,
+        int64_t k) {
+    GGML_ASSERT(QK_K == 256);
+    GGML_ASSERT(k % QK_K == 0);
+
+#if defined(__AVX__)
+    const int nb = k / QK_K;
+    block_q8_Kx4 * GGML_RESTRICT y = (block_q8_Kx4 *) vy;
+
+    for (int ib = 0; ib < nb; ++ib) {
+        alignas(32) int8_t quantized[4][QK_K];
+
+        for (int row = 0; row < 4; ++row) {
+            const float * row_src = x + row * k + ib * QK_K;
+            quantize_q8_K_row_avx1(row_src, &y[ib].d[row], quantized[row]);
+        }
+
+        memset(y[ib].bsums, 0, sizeof(y[ib].bsums));
+
+        // Preserve the generic Q8_Kx4 layout exactly. Keeping this layout loop
+        // explicit also makes bit-for-bit validation against the reference simple.
+        for (int j = 0; j < QK_K * 4; ++j) {
+            int src_offset = (j / 16) * 4;
+            const int src_id = (j % 16) / 4;
+            src_offset += j % 4;
+            const int index = (((j & 15) >> 2) << 2) + ((j >> 8) << 4) + ((j >> 6) & 3);
+
+            const int8_t value = quantized[src_id][src_offset];
+            y[ib].qs[j] = value;
+            y[ib].bsums[index] += value;
+        }
+    }
+    return;
+#endif
+
+    ggml_quantize_mat_q8_K_4x4_generic(x, vy, k);
+}
+
 void ggml_quantize_mat_q8_0_4x8(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) {
     assert(QK8_0 == 32);
     assert(k % QK8_0 == 0);
@@ -1461,6 +1626,258 @@ void ggml_gemv_q4_0_8x8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo
     ggml_gemv_q4_0_8x8_q8_0_generic(n, s, bs, vx, vy, nr, nc);
 }
 
+#if defined(__AVX__) && defined(__SSSE3__) && defined(__SSE4_1__)
+static inline int32_t q4_k_load_i32(const void * src) {
+    int32_t value;
+    memcpy(&value, src, sizeof(value));
+    return value;
+}
+
+static inline __m128i q4_k_dot(__m128i q4, __m128i q8) {
+    // q4 contains unsigned nibbles [0, 15], while q8 is signed. PMADDUBSW
+    // already has precisely that operand convention, so PSIGNB normalization
+    // is redundant here.
+    return _mm_madd_epi16(_mm_maddubs_epi16(q4, q8), _mm_set1_epi16(1));
+}
+
+static inline __m128i q4_k_dot_scaled(__m128i q4, __m128i q8, __m128i scale) {
+    return _mm_madd_epi16(_mm_maddubs_epi16(q4, q8), scale);
+}
+
+static inline __m128i q4_k_load_u8x4(const uint8_t * src) {
+    return _mm_cvtepu8_epi32(_mm_cvtsi32_si128(q4_k_load_i32(src)));
+}
+
+static inline __m128i q4_k_load_scale_pairs(const uint8_t * src) {
+    const __m128i scales = _mm_cvtepu8_epi16(_mm_cvtsi32_si128(q4_k_load_i32(src)));
+    return _mm_unpacklo_epi16(scales, scales);
+}
+
+static inline __m128i q4_k_repeat_i8x4(const int8_t * src) {
+    return _mm_set1_epi32(q4_k_load_i32(src));
+}
+
+static inline __m256 q4_k_load_f16x8(const ggml_fp16_t * src) {
+#if defined(__F16C__)
+    return _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *) src));
+#else
+    return _mm256_set_ps(
+            GGML_CPU_FP16_TO_FP32(src[7]),
+            GGML_CPU_FP16_TO_FP32(src[6]),
+            GGML_CPU_FP16_TO_FP32(src[5]),
+            GGML_CPU_FP16_TO_FP32(src[4]),
+            GGML_CPU_FP16_TO_FP32(src[3]),
+            GGML_CPU_FP16_TO_FP32(src[2]),
+            GGML_CPU_FP16_TO_FP32(src[1]),
+            GGML_CPU_FP16_TO_FP32(src[0]));
+#endif
+}
+#endif
+
+void ggml_gemv_q4_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) {
+    const int qk = QK_K;
+    const int nb = n / qk;
+    const int ncols_interleaved = 8;
+
+    assert(n % qk == 0);
+    assert(nc % ncols_interleaved == 0);
+    GGML_ASSERT(nr == 1);
+
+    UNUSED(bs);
+    UNUSED(nr);
+
+#if defined(__AVX__) && defined(__SSSE3__) && defined(__SSE4_1__)
+    const __m128i m4b = _mm_set1_epi8(0x0F);
+    const block_q8_K * a_ptr = (const block_q8_K *) vy;
+
+    for (int x = 0; x < nc / ncols_interleaved; ++x) {
+        const block_q4_Kx8 * b_ptr = (const block_q4_Kx8 *) vx + x * nb;
+        __m256 acc = _mm256_setzero_ps();
+
+        for (int b = 0; b < nb; ++b) {
+            __m128i sums_0123 = _mm_setzero_si128();
+            __m128i sums_4567 = _mm_setzero_si128();
+            __m128i mins_0123 = _mm_setzero_si128();
+            __m128i mins_4567 = _mm_setzero_si128();
+            uint32_t decoded_scales[32];
+            for (int sb = 0; sb < 8; ++sb) {
+                memcpy(decoded_scales + sb * 4, b_ptr[b].scales + sb * 12, 12);
+                decoded_scales[sb * 4 + 3] = ((decoded_scales[sb * 4 + 2] >> 4) & 0x0f0f0f0f) | (((decoded_scales[sb * 4 + 1] >> 6) & 0x03030303) << 4);
+                const uint32_t mins = decoded_scales[sb * 4 + 1] & 0x3f3f3f3f;
+                decoded_scales[sb * 4 + 1] = (decoded_scales[sb * 4 + 2] & 0x0f0f0f0f) | (((decoded_scales[sb * 4] >> 6) & 0x03030303) << 4);
+                decoded_scales[sb * 4 + 2] = mins;
+                decoded_scales[sb * 4] &= 0x3f3f3f3f;
+            }
+
+            const __m256 q4_d = q4_k_load_f16x8(b_ptr[b].d);
+            const __m256 q4_dmin = q4_k_load_f16x8(b_ptr[b].dmin);
+            const __m256 q8_d = _mm256_set1_ps(a_ptr[b].d);
+
+            for (int k = 0; k < 32; ++k) {
+                const uint8_t * scales_0 = (const uint8_t *) decoded_scales + (k / 8) * 32;
+                const uint8_t * scales_1 = scales_0 + 16;
+                const int8_t * q8 = a_ptr[b].qs + (k / 8) * 64 + (k % 8) * 4;
+                const __m128i q8_lo = q4_k_repeat_i8x4(q8);
+                const __m128i q8_hi = q4_k_repeat_i8x4(q8 + 32);
+                const __m128i q4_0123 = _mm_loadu_si128((const __m128i *) (b_ptr[b].qs + k * 32));
+                const __m128i q4_4567 = _mm_loadu_si128((const __m128i *) (b_ptr[b].qs + k * 32 + 16));
+
+                const __m128i dots_0123 = _mm_add_epi32(
+                        _mm_mullo_epi32(q4_k_dot(_mm_and_si128(q4_0123, m4b), q8_lo), q4_k_load_u8x4(scales_0)),
+                        _mm_mullo_epi32(q4_k_dot(_mm_and_si128(_mm_srli_epi16(q4_0123, 4), m4b), q8_hi), q4_k_load_u8x4(scales_1)));
+                const __m128i dots_4567 = _mm_add_epi32(
+                        _mm_mullo_epi32(q4_k_dot(_mm_and_si128(q4_4567, m4b), q8_lo), q4_k_load_u8x4(scales_0 + 4)),
+                        _mm_mullo_epi32(q4_k_dot(_mm_and_si128(_mm_srli_epi16(q4_4567, 4), m4b), q8_hi), q4_k_load_u8x4(scales_1 + 4)));
+
+                sums_0123 = _mm_add_epi32(sums_0123, dots_0123);
+                sums_4567 = _mm_add_epi32(sums_4567, dots_4567);
+            }
+
+            for (int sb = 0; sb < 8; ++sb) {
+                const uint8_t * mins = (const uint8_t *) decoded_scales + 8 + sb * 16;
+                const int bsum = a_ptr[b].bsums[sb * 2] + a_ptr[b].bsums[sb * 2 + 1];
+                mins_0123 = _mm_add_epi32(mins_0123, _mm_mullo_epi32(q4_k_load_u8x4(mins), _mm_set1_epi32(bsum)));
+                mins_4567 = _mm_add_epi32(mins_4567, _mm_mullo_epi32(q4_k_load_u8x4(mins + 4), _mm_set1_epi32(bsum)));
+            }
+
+            const __m256 sums = _mm256_insertf128_ps(_mm256_castps128_ps256(_mm_cvtepi32_ps(sums_0123)), _mm_cvtepi32_ps(sums_4567), 1);
+            const __m256 min_sums = _mm256_insertf128_ps(_mm256_castps128_ps256(_mm_cvtepi32_ps(mins_0123)), _mm_cvtepi32_ps(mins_4567), 1);
+            acc = _mm256_add_ps(acc, _mm256_mul_ps(
+                    _mm256_sub_ps(_mm256_mul_ps(sums, q4_d), _mm256_mul_ps(min_sums, q4_dmin)), q8_d));
+        }
+
+        _mm256_storeu_ps(s + x * ncols_interleaved, acc);
+    }
+    return;
+#endif
+
+    ggml_gemv_q4_K_8x4_q8_K_generic(n, s, bs, vx, vy, nr, nc);
+}
+
+void ggml_gemm_q4_K_8x4_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) {
+    const int qk = QK_K;
+    const int nb = n / qk;
+    const int ncols_interleaved = 8;
+
+    assert(n % qk == 0);
+    assert(nr % 4 == 0);
+    assert(nc % ncols_interleaved == 0);
+
+    UNUSED(s);
+    UNUSED(bs);
+    UNUSED(vx);
+    UNUSED(vy);
+    UNUSED(nr);
+    UNUSED(nc);
+    UNUSED(nb);
+
+#if defined(__AVX__) && defined(__SSSE3__) && defined(__SSE4_1__)
+    const __m128i m4b = _mm_set1_epi8(0x0F);
+    const block_q4_Kx8 * b_ptr_start = (const block_q4_Kx8 *) vx;
+    const block_q8_Kx4 * a_ptr_start = (const block_q8_Kx4 *) vy;
+
+    for (int y = 0; y < nr / 4; ++y) {
+        const block_q8_Kx4 * a_ptr = a_ptr_start + y * nb;
+
+        for (int x = 0; x < nc / ncols_interleaved; ++x) {
+            const block_q4_Kx8 * b_ptr = b_ptr_start + x * nb;
+            __m256 acc[4] = {
+                _mm256_setzero_ps(), _mm256_setzero_ps(),
+                _mm256_setzero_ps(), _mm256_setzero_ps(),
+            };
+
+            for (int b = 0; b < nb; ++b) {
+                __m128i sums_0123[4] = {
+                    _mm_setzero_si128(), _mm_setzero_si128(),
+                    _mm_setzero_si128(), _mm_setzero_si128(),
+                };
+                __m128i sums_4567[4] = {
+                    _mm_setzero_si128(), _mm_setzero_si128(),
+                    _mm_setzero_si128(), _mm_setzero_si128(),
+                };
+                uint32_t decoded_scales[32];
+
+                for (int sb = 0; sb < 8; ++sb) {
+                    memcpy(decoded_scales + sb * 4, b_ptr[b].scales + sb * 12, 12);
+                    decoded_scales[sb * 4 + 3] = ((decoded_scales[sb * 4 + 2] >> 4) & 0x0f0f0f0f) |
+                                                   (((decoded_scales[sb * 4 + 1] >> 6) & 0x03030303) << 4);
+                    const uint32_t mins = decoded_scales[sb * 4 + 1] & 0x3f3f3f3f;
+                    decoded_scales[sb * 4 + 1] = (decoded_scales[sb * 4 + 2] & 0x0f0f0f0f) |
+                                                   (((decoded_scales[sb * 4] >> 6) & 0x03030303) << 4);
+                    decoded_scales[sb * 4 + 2] = mins;
+                    decoded_scales[sb * 4] &= 0x3f3f3f3f;
+                }
+
+                const __m256 q4_d = q4_k_load_f16x8(b_ptr[b].d);
+                const __m256 q4_dmin = q4_k_load_f16x8(b_ptr[b].dmin);
+
+                for (int k = 0; k < 32; ++k) {
+                    const uint8_t * scales_0 = (const uint8_t *) decoded_scales + (k / 8) * 32;
+                    const uint8_t * scales_1 = scales_0 + 16;
+                    const __m128i scale_0123 = q4_k_load_scale_pairs(scales_0);
+                    const __m128i scale_4567 = q4_k_load_scale_pairs(scales_0 + 4);
+                    const __m128i scale_hi_0123 = q4_k_load_scale_pairs(scales_1);
+                    const __m128i scale_hi_4567 = q4_k_load_scale_pairs(scales_1 + 4);
+                    const __m128i q4_0123 = _mm_loadu_si128((const __m128i *) (b_ptr[b].qs + k * 32));
+                    const __m128i q4_4567 = _mm_loadu_si128((const __m128i *) (b_ptr[b].qs + k * 32 + 16));
+
+                    for (int m = 0; m < 4; ++m) {
+                        const int8_t * q8 = a_ptr[b].qs + (k / 8) * 256 + (k % 8) * 16 + m * 4;
+                        const __m128i q8_lo = q4_k_repeat_i8x4(q8);
+                        const __m128i q8_hi = q4_k_repeat_i8x4(q8 + 128);
+
+                        const __m128i dots_0123 = _mm_add_epi32(
+                                q4_k_dot_scaled(_mm_and_si128(q4_0123, m4b), q8_lo, scale_0123),
+                                q4_k_dot_scaled(_mm_and_si128(_mm_srli_epi16(q4_0123, 4), m4b), q8_hi, scale_hi_0123));
+                        const __m128i dots_4567 = _mm_add_epi32(
+                                q4_k_dot_scaled(_mm_and_si128(q4_4567, m4b), q8_lo, scale_4567),
+                                q4_k_dot_scaled(_mm_and_si128(_mm_srli_epi16(q4_4567, 4), m4b), q8_hi, scale_hi_4567));
+
+                        sums_0123[m] = _mm_add_epi32(sums_0123[m], dots_0123);
+                        sums_4567[m] = _mm_add_epi32(sums_4567[m], dots_4567);
+                    }
+                }
+
+                for (int m = 0; m < 4; ++m) {
+                    __m128i min_sums_0123 = _mm_setzero_si128();
+                    __m128i min_sums_4567 = _mm_setzero_si128();
+
+                    for (int sb = 0; sb < 8; ++sb) {
+                        const uint8_t * mins = (const uint8_t *) decoded_scales + 8 + sb * 16;
+                        const int16_t * bsums = a_ptr[b].bsums + (sb * 8) + (m * 4) - ((sb % 2) * 6);
+                        const __m128i bsum = _mm_set1_epi32(bsums[0] + bsums[1]);
+                        min_sums_0123 = _mm_add_epi32(
+                                min_sums_0123,
+                                _mm_mullo_epi32(q4_k_load_u8x4(mins), bsum));
+                        min_sums_4567 = _mm_add_epi32(
+                                min_sums_4567,
+                                _mm_mullo_epi32(q4_k_load_u8x4(mins + 4), bsum));
+                    }
+
+                    const __m128 sums_0 = _mm_cvtepi32_ps(sums_0123[m]);
+                    const __m128 sums_1 = _mm_cvtepi32_ps(sums_4567[m]);
+                    const __m128 mins_0 = _mm_cvtepi32_ps(min_sums_0123);
+                    const __m128 mins_1 = _mm_cvtepi32_ps(min_sums_4567);
+                    const __m256 sums = _mm256_insertf128_ps(_mm256_castps128_ps256(sums_0), sums_1, 1);
+                    const __m256 mins = _mm256_insertf128_ps(_mm256_castps128_ps256(mins_0), mins_1, 1);
+                    const __m256 q8_d = _mm256_set1_ps(a_ptr[b].d[m]);
+
+                    acc[m] = _mm256_add_ps(acc[m], _mm256_mul_ps(
+                            _mm256_sub_ps(_mm256_mul_ps(sums, q4_d), _mm256_mul_ps(mins, q4_dmin)), q8_d));
+                }
+            }
+
+            for (int m = 0; m < 4; ++m) {
+                _mm256_storeu_ps(s + (y * 4 + m) * bs + x * 8, acc[m]);
+            }
+        }
+    }
+    return;
+#endif
+
+    ggml_gemm_q4_K_8x4_q8_K_generic(n, s, bs, vx, vy, nr, nc);
+}
+
 void ggml_gemv_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) {
     const int qk = QK_K;
     const int nb = n / qk;

ggml/src/ggml-cpu/repack.cpp

@@ -15,6 +15,7 @@
 #include <cstring>
 #include <cassert>
 #include <cstdio>  // for GGML_ASSERT
+#include <type_traits>
 
 #include "repack.h"
 
@@ -4547,6 +4548,33 @@ template <> void gemm<block_q2_K, 1, 16, GGML_TYPE_Q8_K>(int n, float * s, size_
 }
 #endif
 
+static void pack_q8_K_rows_4x4(
+        const block_q8_K * const rows[4],
+        block_q8_Kx4 * GGML_RESTRICT dst,
+        int64_t n_per_row) {
+    GGML_ASSERT(n_per_row % QK_K == 0);
+    const int64_t nb = n_per_row / QK_K;
+
+    for (int64_t ib = 0; ib < nb; ++ib) {
+        for (int row = 0; row < 4; ++row) {
+            dst[ib].d[row] = rows[row][ib].d;
+        }
+        memset(dst[ib].bsums, 0, sizeof(dst[ib].bsums));
+
+        // This is the exact layout used by ggml_quantize_mat_q8_K_4x4_generic.
+        for (int j = 0; j < QK_K * 4; ++j) {
+            int src_offset = (j / 16) * 4;
+            const int src_id = (j % 16) / 4;
+            src_offset += j % 4;
+            const int index = (((j & 15) >> 2) << 2) + ((j >> 8) << 4) + ((j >> 6) & 3);
+
+            const int8_t value = rows[src_id][ib].qs[src_offset];
+            dst[ib].qs[j] = value;
+            dst[ib].bsums[index] += value;
+        }
+    }
+}
+
 class tensor_traits_base : public ggml::cpu::tensor_traits {
   public:
     virtual int repack(struct ggml_tensor * t, const void * data, size_t data_size) = 0;
@@ -4554,8 +4582,11 @@ class tensor_traits_base : public ggml::cpu::tensor_traits {
 
 template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS, ggml_type PARAM_TYPE> class tensor_traits : public tensor_traits_base {
 
-    bool work_size(int /* n_threads */, const struct ggml_tensor * op, size_t & size) override {
-        // not realy a GGML_TYPE_Q8_0 but same size.
+    static constexpr bool supports_grouped_mmid_gemm =
+            std::is_same_v<BLOC_TYPE, block_q4_K> &&
+            INTER_SIZE == 4 && NB_COLS == 8 && PARAM_TYPE == GGML_TYPE_Q8_K;
+
+    bool work_size(int n_threads, const struct ggml_tensor * op, size_t & size) override {
         switch (op->op) {
             case GGML_OP_MUL_MAT:
                 {
@@ -4564,20 +4595,25 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS, ggml_type PAR
                 }
             case GGML_OP_MUL_MAT_ID:
                 {
-                    size = ggml_row_size(PARAM_TYPE, ggml_nelements(op->src[1]));
-                    size = GGML_PAD(size, sizeof(int64_t)); // + padding for next block.
+                    const int64_t ne10 = op->src[1]->ne[0];
+                    const int64_t ne11 = op->src[1]->ne[1];
+                    const int64_t ne12 = op->src[1]->ne[2];
+                    const int64_t ne02 = op->src[0]->ne[2];
+                    const size_t nbw1 = ggml_row_size(PARAM_TYPE, ne10);
+                    const size_t quantized_size = GGML_PAD(nbw1 * ne11 * ne12, 64);
 
-                    const int64_t ne02 = op->src[0]->ne[2]; // n_as, n_expert
-                    const int64_t ne12 = op->src[1]->ne[2]; // n_tokens
+                    size = quantized_size;
+                    if constexpr (supports_grouped_mmid_gemm) {
+                        const size_t packed_stride = GGML_PAD(4 * nbw1, 64);
+                        size += (size_t) (n_threads > 0 ? n_threads : 1) * packed_stride;
+                    }
+                    size = GGML_PAD(size, sizeof(int64_t));
 
                     const size_t sizeof_mmid_row_mapping = sizeof(int64_t);
-
-                    size += sizeof_mmid_row_mapping*ne02*(ne12 + 1);
-
+                    size += sizeof_mmid_row_mapping * ne02 * (ne12 + 1);
                     return true;
                 }
             default:
-                // GGML_ABORT("fatal error");
                 break;
         }
         return false;
@@ -4790,73 +4826,72 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS, ggml_type PAR
 
         const int ith = params->ith;
         const int nth = params->nth;
-
         const ggml_from_float_t from_float = ggml_get_type_traits_cpu(PARAM_TYPE)->from_float;
 
-        // we don't support permuted src0 or src1
         GGML_ASSERT(nb00 == ggml_type_size(src0->type));
         GGML_ASSERT(nb10 == ggml_type_size(src1->type));
-
-        // dst cannot be transposed or permuted
         GGML_ASSERT(nb0 == sizeof(float));
-        GGML_ASSERT(nb0 <= nb1);
-        GGML_ASSERT(nb1 <= nb2);
-        GGML_ASSERT(nb2 <= nb3);
-
-        GGML_ASSERT(ne03 == 1);
-        GGML_ASSERT(ne13 == 1);
-        GGML_ASSERT(ne3  == 1);
-
+        GGML_ASSERT(nb0 <= nb1 && nb1 <= nb2 && nb2 <= nb3);
+        GGML_ASSERT(ne03 == 1 && ne13 == 1 && ne3 == 1);
         GGML_ASSERT(src1->type == GGML_TYPE_F32);
 
-        // row groups
-        const int n_ids = ids->ne[0]; // n_expert_used
-        const int n_as  = ne02;       // n_expert
+        const int n_ids = ids->ne[0];
+        const int n_as = ne02;
+        GGML_ASSERT(ids->ne[1] == ne12);
 
         const size_t nbw1 = ggml_row_size(PARAM_TYPE, ne10);
-        const size_t nbw2 = nbw1*ne11;
-        const size_t nbw3 = nbw2*ne12;
+        const size_t nbw2 = nbw1 * ne11;
+        const size_t nbw3 = nbw2 * ne12;
+        const size_t quantized_size = GGML_PAD(nbw3, 64);
+        const size_t packed_stride = supports_grouped_mmid_gemm ? GGML_PAD(4 * nbw1, 64) : 0;
+        const size_t packed_size = supports_grouped_mmid_gemm ? packed_stride * nth : 0;
+        const size_t metadata_offset = GGML_PAD(quantized_size + packed_size, sizeof(int64_t));
 
         struct mmid_row_mapping {
             int32_t i1;
             int32_t i2;
         };
+        static_assert(sizeof(mmid_row_mapping) == sizeof(int64_t));
 
-        GGML_ASSERT(params->wsize >=
-                (GGML_PAD(nbw3, sizeof(int64_t)) +
-                 n_as*(ne12 + 1)*sizeof(mmid_row_mapping))
-                );
+        const size_t required_size = metadata_offset +
+                (size_t) n_as * (ne12 + 1) * sizeof(mmid_row_mapping);
+        GGML_ASSERT(params->wsize >= required_size);
 
-        auto * wdata          = (char *)params->wdata;
-        auto * wdata_src1_end = (char *)wdata + GGML_PAD(nbw3, sizeof(int64_t));
+        auto * wdata = (char *) params->wdata;
+        auto * packed_workspace = wdata + quantized_size;
+        GGML_UNUSED(packed_workspace);
+        auto * matrix_row_counts = (int64_t *) (wdata + metadata_offset);
+        auto * matrix_rows = (mmid_row_mapping *) (matrix_row_counts + n_as);
 
-        // total of [n_as][ne12 + 1] elements of type mmid_row_mapping (2*int32_t = int64_t)
-        auto * matrix_row_counts = (int64_t *) (wdata_src1_end);                                        // [n_as]
-        struct mmid_row_mapping * matrix_rows = (struct mmid_row_mapping *) (matrix_row_counts + n_as); // [n_as][ne12]
-
-        // src1: float32 => param type
-        for (int64_t i12 = 0; i12 < ne12; ++i12) {
-            for (int64_t i11 = ith; i11 < ne11; i11 += nth) {
-                from_float((float *)((char *) src1->data + i12 * nb12 + i11 * nb11),
-                           (void *)               (wdata + i12 * nbw2 + i11 * nbw1),
-                           ne10);
-            }
+        // Flatten token and selected-expert dimensions. With top-k smaller than
+        // the thread count, the previous nested loop left most workers idle.
+        const int64_t src1_rows = ne11 * ne12;
+        for (int64_t row = ith; row < src1_rows; row += nth) {
+            const int64_t i12 = row / ne11;
+            const int64_t i11 = row - i12 * ne11;
+            from_float(
+                    (float *) ((char *) src1->data + i12 * nb12 + i11 * nb11),
+                    (void *) (wdata + i12 * nbw2 + i11 * nbw1),
+                    ne10);
         }
 
 #define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) * ne12 + (i1)]
 
         if (ith == 0) {
-            // initialize matrix_row_counts
             memset(matrix_row_counts, 0, n_as * sizeof(int64_t));
 
-            // group rows by src0 matrix
             for (int32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
                 for (int32_t id = 0; id < n_ids; ++id) {
-                    const int32_t i02 =
-                        *(const int32_t *) ((const char *) ids->data + iid1 * ids->nb[1] + id * ids->nb[0]);
-
+                    int32_t i02;
+                    memcpy(
+                            &i02,
+                            (const char *) ids->data + iid1 * ids->nb[1] + id * ids->nb[0],
+                            sizeof(i02));
                     GGML_ASSERT(i02 >= 0 && i02 < n_as);
-
+                    // Top-k routing should not select one expert twice for a token.
+                    // Guard the fixed per-expert mapping capacity instead of corrupting
+                    // the following expert's row list if malformed IDs are supplied.
+                    GGML_ASSERT(matrix_row_counts[i02] < ne12);
                     MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = { id, iid1 };
                     matrix_row_counts[i02] += 1;
                 }
@@ -4865,49 +4900,119 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS, ggml_type PAR
 
         ggml_barrier(params->threadpool);
 
-        // compute each matrix multiplication in sequence
-        for (int cur_a = 0; cur_a < n_as; ++cur_a) {
-            const int64_t cne1 = matrix_row_counts[cur_a];
+        if constexpr (supports_grouped_mmid_gemm) {
+            int64_t full_groups = 0;
+            for (int cur_a = 0; cur_a < n_as; ++cur_a) {
+                full_groups += matrix_row_counts[cur_a] / 4;
+            }
 
-            if (cne1 == 0) {
+            // For a substantial prompt, assign complete routed-row groups to workers
+            // and use the 4-row GEMM kernel. Small/decode batches retain output-column
+            // parallel GEMV so one routed row can still occupy all cores.
+            if (full_groups >= nth) {
+                GGML_ASSERT(ne01 % NB_COLS == 0);
+                auto * packed_src1 = (block_q8_Kx4 *) (packed_workspace + ith * packed_stride);
+                int64_t task_index = 0;
+
+                for (int cur_a = 0; cur_a < n_as; ++cur_a) {
+                    const int64_t count = matrix_row_counts[cur_a];
+                    if (count == 0) {
+                        continue;
+                    }
+
+                    const auto * src0_cur = (const char *) src0->data + cur_a * nb02;
+                    const int64_t n_groups = count / 4;
+
+                    for (int64_t group = 0; group < n_groups; ++group, ++task_index) {
+                        if (task_index % nth != ith) {
+                            continue;
+                        }
+
+                        mmid_row_mapping mappings[4];
+                        const block_q8_K * rows[4];
+                        for (int m = 0; m < 4; ++m) {
+                            mappings[m] = MMID_MATRIX_ROW(cur_a, group * 4 + m);
+                            const int64_t i11 = mappings[m].i1 % ne11;
+                            const int64_t i12 = mappings[m].i2;
+                            rows[m] = (const block_q8_K *) (wdata + i12 * nbw2 + i11 * nbw1);
+                        }
+                        pack_q8_K_rows_4x4(rows, packed_src1, ne10);
+
+                        for (int64_t col = 0; col < ne01; col += NB_COLS) {
+                            alignas(32) float tile[4 * NB_COLS];
+                            gemm<BLOC_TYPE, INTER_SIZE, NB_COLS, PARAM_TYPE>(
+                                    ne00,
+                                    tile,
+                                    NB_COLS,
+                                    src0_cur + col * nb01,
+                                    packed_src1,
+                                    4,
+                                    NB_COLS);
+
+                            for (int m = 0; m < 4; ++m) {
+                                float * dst_row = (float *) ((char *) dst->data +
+                                        mappings[m].i1 * nb1 + mappings[m].i2 * nb2);
+                                memcpy(dst_row + col, tile + m * NB_COLS, NB_COLS * sizeof(float));
+                            }
+                        }
+                    }
+
+                    for (int64_t ir1 = n_groups * 4; ir1 < count; ++ir1, ++task_index) {
+                        if (task_index % nth != ith) {
+                            continue;
+                        }
+
+                        const mmid_row_mapping mapping = MMID_MATRIX_ROW(cur_a, ir1);
+                        const int64_t i11 = mapping.i1 % ne11;
+                        const int64_t i12 = mapping.i2;
+                        const auto * src1_col = wdata + i12 * nbw2 + i11 * nbw1;
+                        float * dst_row = (float *) ((char *) dst->data +
+                                mapping.i1 * nb1 + mapping.i2 * nb2);
+
+                        gemv<BLOC_TYPE, INTER_SIZE, NB_COLS, PARAM_TYPE>(
+                                ne00, dst_row, ne01, src0_cur, src1_col, 1, ne01);
+                    }
+                }
+                return;
+            }
+        }
+
+        // Decode and small prompts: parallelize each routed row over output columns.
+        for (int cur_a = 0; cur_a < n_as; ++cur_a) {
+            const int64_t count = matrix_row_counts[cur_a];
+            if (count == 0) {
                 continue;
             }
 
-            const auto * src0_cur = (const char *) src0->data + cur_a*nb02;
-
-            //const int64_t nr0 = ne01; // src0 rows
-            const int64_t nr1 = cne1; // src1 rows
-
+            const auto * src0_cur = (const char *) src0->data + cur_a * nb02;
             int64_t src0_cur_start = (ith * ne01) / nth;
-            int64_t src0_cur_end   = ((ith + 1) * ne01) / nth;
+            int64_t src0_cur_end = ((ith + 1) * ne01) / nth;
 
-            // Align boundaries to NB_COLS - round up to ensure all data is included
-            src0_cur_start = (src0_cur_start % NB_COLS) ? src0_cur_start + NB_COLS - (src0_cur_start % NB_COLS) : src0_cur_start;
-            src0_cur_end   = (src0_cur_end   % NB_COLS) ? src0_cur_end   + NB_COLS - (src0_cur_end   % NB_COLS) : src0_cur_end;
-            if (src0_cur_end > ne01) {
-                src0_cur_end = ne01;
-            }
+            src0_cur_start = (src0_cur_start % NB_COLS) ?
+                    src0_cur_start + NB_COLS - (src0_cur_start % NB_COLS) : src0_cur_start;
+            src0_cur_end = (src0_cur_end % NB_COLS) ?
+                    src0_cur_end + NB_COLS - (src0_cur_end % NB_COLS) : src0_cur_end;
+            src0_cur_end = MIN(src0_cur_end, ne01);
 
             if (src0_cur_start >= src0_cur_end) {
-                return;
+                continue;
             }
 
-            for (int ir1 = 0; ir1 < nr1; ir1++) {
-                struct mmid_row_mapping row_mapping = MMID_MATRIX_ROW(cur_a, ir1);
-
-                const int id = row_mapping.i1;  // selected expert index
-
-                const int64_t i11 = id % ne11;
-                const int64_t i12 = row_mapping.i2;  // row index in src1
-
-                const int64_t i1 = id;               // selected expert index
-                const int64_t i2 = i12;              // row
-
-                const auto * src1_col = (const char *) wdata + (i11 * nbw1 + i12 * nbw2);
+            for (int64_t ir1 = 0; ir1 < count; ++ir1) {
+                const mmid_row_mapping mapping = MMID_MATRIX_ROW(cur_a, ir1);
+                const int64_t i11 = mapping.i1 % ne11;
+                const int64_t i12 = mapping.i2;
+                const auto * src1_col = wdata + i12 * nbw2 + i11 * nbw1;
 
                 gemv<BLOC_TYPE, INTER_SIZE, NB_COLS, PARAM_TYPE>(
-                    ne00, (float *) ((char *) dst->data + (i1 * nb1 + i2 * nb2)) + src0_cur_start, ne01,
-                    src0_cur + src0_cur_start * nb01, src1_col, 1, src0_cur_end - src0_cur_start);
+                        ne00,
+                        (float *) ((char *) dst->data + mapping.i1 * nb1 + mapping.i2 * nb2) +
+                                src0_cur_start,
+                        ne01,
+                        src0_cur + src0_cur_start * nb01,
+                        src1_col,
+                        1,
+                        src0_cur_end - src0_cur_start);
             }
         }
 #undef MMID_MATRIX_ROW
@@ -5009,6 +5114,11 @@ static const ggml::cpu::tensor_traits * ggml_repack_get_optimal_repack_type(cons
                 return &q4_K_8x8_q8_K;
             }
         }
+        if (ggml_cpu_has_avx()) {
+            if (cur->ne[1] % 8 == 0) {
+                return &q4_K_8x4_q8_K;
+            }
+        }
         if (ggml_cpu_has_neon() && ggml_cpu_has_matmul_int8()) {
             if (cur->ne[1] % 8 == 0) {
                 return &q4_K_8x8_q8_K;

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -802,6 +802,27 @@ static void ggml_backend_cuda_buffer_set_tensor_2d(ggml_backend_buffer_t buffer,
     ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *) buffer->context;
 
     ggml_cuda_set_device(ctx->device);
+    if (n_copies == 1 || (stride_tensor == size && stride_data == size)) {
+        CUDA_CHECK(cudaMemcpyAsync((char *) tensor->data + offset, data, size*n_copies, cudaMemcpyHostToDevice, cudaStreamPerThread));
+        CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+        return;
+    }
+    if (stride_tensor == size) {
+        // Strided host reads into a contiguous device range, e.g. uploading one device's slice
+        // of a split model tensor. A strided copy from pageable (often mmap-backed) host memory
+        // degrades to per-row processing in the driver; gathering the rows into a contiguous
+        // host buffer and issuing a single 1D copy is an order of magnitude faster.
+        static thread_local std::vector<char> staging;
+        if (staging.size() < size*n_copies) {
+            staging.resize(size*n_copies);
+        }
+        for (size_t i = 0; i < n_copies; i++) {
+            memcpy(staging.data() + i*size, (const char *) data + i*stride_data, size);
+        }
+        CUDA_CHECK(cudaMemcpyAsync((char *) tensor->data + offset, staging.data(), size*n_copies, cudaMemcpyHostToDevice, cudaStreamPerThread));
+        CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+        return;
+    }
     CUDA_CHECK(cudaMemcpy2DAsync(
         (char *) tensor->data + offset, stride_tensor, data, stride_data, size, n_copies, cudaMemcpyHostToDevice, cudaStreamPerThread));
     CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));