diff --git a/ggml/src/ggml-cpu/arch-fallback.h b/ggml/src/ggml-cpu/arch-fallback.h index 4dbd1982b..91dcc0149 100644 --- a/ggml/src/ggml-cpu/arch-fallback.h +++ b/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 diff --git a/ggml/src/ggml-cpu/arch/x86/repack.cpp b/ggml/src/ggml-cpu/arch/x86/repack.cpp index af1cebad1..d3c08f468 100644 --- a/ggml/src/ggml-cpu/arch/x86/repack.cpp +++ b/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; diff --git a/ggml/src/ggml-cpu/repack.cpp b/ggml/src/ggml-cpu/repack.cpp index d56db9802..f8f3e7f53 100644 --- a/ggml/src/ggml-cpu/repack.cpp +++ b/ggml/src/ggml-cpu/repack.cpp @@ -15,6 +15,7 @@ #include #include #include // for GGML_ASSERT +#include #include "repack.h" @@ -4547,6 +4548,33 @@ template <> void gemm(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 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 && + 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 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 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 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( + 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( + 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( - 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; diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index a9038f1f4..18d02864d 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/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 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));