From 6dbc1174b8f46a8e064259985d84c81a472beea4 Mon Sep 17 00:00:00 2001 From: ragz4125 <65285549+ragz4125@users.noreply.github.com> Date: Wed, 1 Jul 2026 13:01:20 +0530 Subject: [PATCH] ggml-cpu: add AVX2 optimization for nvfp4 dot product and use UE4M3 LUT (#23961) --- ggml/src/ggml-common.h | 5 +- ggml/src/ggml-cpu/arch-fallback.h | 1 - ggml/src/ggml-cpu/arch/x86/quants.c | 146 +++++++++++++++++++++++++++- ggml/src/ggml-cpu/ggml-cpu.c | 8 ++ ggml/src/ggml-cpu/simd-mappings.h | 11 +++ 5 files changed, 164 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-common.h b/ggml/src/ggml-common.h index f05683b44..a51b37dcd 100644 --- a/ggml/src/ggml-common.h +++ b/ggml/src/ggml-common.h @@ -1111,11 +1111,12 @@ GGML_TABLE_BEGIN(int8_t, kvalues_iq4nl, 16) -127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113, GGML_TABLE_END() -// e2m1 values (doubled) +// e2m1 values (doubled), shared by MXFP4 and NVFP4 // ref: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf -GGML_TABLE_BEGIN(int8_t, kvalues_mxfp4, 16) +GGML_TABLE_BEGIN(int8_t, kvalues_fp4, 16) 0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12, GGML_TABLE_END() +#define kvalues_mxfp4 kvalues_fp4 #define NGRID_IQ1S 2048 #define IQ1S_DELTA 0.125f diff --git a/ggml/src/ggml-cpu/arch-fallback.h b/ggml/src/ggml-cpu/arch-fallback.h index 1fc2b4b71..da2ac98a4 100644 --- a/ggml/src/ggml-cpu/arch-fallback.h +++ b/ggml/src/ggml-cpu/arch-fallback.h @@ -82,7 +82,6 @@ #define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K #elif defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) || defined(_M_X64) // quants.c -#define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 // repack.cpp #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 #define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 diff --git a/ggml/src/ggml-cpu/arch/x86/quants.c b/ggml/src/ggml-cpu/arch/x86/quants.c index 94b19b82b..ea54cfe44 100644 --- a/ggml/src/ggml-cpu/arch/x86/quants.c +++ b/ggml/src/ggml-cpu/arch/x86/quants.c @@ -934,7 +934,7 @@ void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo #if defined __AVX2__ - const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_mxfp4); + const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_fp4); const __m128i m4b = _mm_set1_epi8(0x0f); const __m256i mone = _mm256_set1_epi16(1); @@ -963,7 +963,7 @@ void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo sumf = hsum_float_8(_mm256_add_ps(accum1, accum2)); #elif defined __AVX__ - const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_mxfp4); + const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_fp4); const __m128i m4b = _mm_set1_epi8(0x0f); __m256 accum = _mm256_setzero_ps(); @@ -993,14 +993,152 @@ void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo int sumi1 = 0; int sumi2 = 0; for (int j = 0; j < QK_MXFP4/2; ++j) { - sumi1 += y[ib].qs[j + 0] * kvalues_mxfp4[x[ib].qs[j] & 0xf]; - sumi2 += y[ib].qs[j + QK_MXFP4/2] * kvalues_mxfp4[x[ib].qs[j] >> 4]; + sumi1 += y[ib].qs[j + 0] * kvalues_fp4[x[ib].qs[j] & 0xf]; + sumi2 += y[ib].qs[j + QK_MXFP4/2] * kvalues_fp4[x[ib].qs[j] >> 4]; } sumf += d * (sumi1 + sumi2); } *s = sumf; } +void ggml_vec_dot_nvfp4_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) { + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + assert(n % QK_NVFP4 == 0); + + const block_nvfp4 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + const int nb = n / QK_NVFP4; + int ib = 0; + float sumf = 0; + +#if defined(__AVX2__) + + const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_fp4); + const __m128i m4b = _mm_set1_epi8(0x0f); + const __m256i mone = _mm256_set1_epi16(1); + + __m256 accum = _mm256_setzero_ps(); + for(; ib < nb; ib++){ + + const __m128i q4bits_01 = _mm_loadu_si128((const __m128i *)(x[ib].qs + 0)); + const __m128i q4bits_23 = _mm_loadu_si128((const __m128i *)(x[ib].qs + 16)); + + const __m256i q8_01 = _mm256_loadu_si256((const __m256i *)y[2*ib + 0].qs); + const __m256i q8_23 = _mm256_loadu_si256((const __m256i *)y[2*ib + 1].qs); + + const __m128i q4_01_lo = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_01, m4b)); + const __m128i q4_01_hi = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_01, 4), m4b)); + const __m128i q4_23_lo = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_23, m4b)); + const __m128i q4_23_hi = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_23, 4), m4b)); + + //reordering + const __m256i q4_01 = MM256_SET_M128I(_mm_unpackhi_epi64(q4_01_lo,q4_01_hi), _mm_unpacklo_epi64(q4_01_lo,q4_01_hi)); + const __m256i q4_23 = MM256_SET_M128I(_mm_unpackhi_epi64(q4_23_lo,q4_23_hi),_mm_unpacklo_epi64(q4_23_lo,q4_23_hi)); + + const __m256i p01 = mul_add_epi8(q4_01,q8_01); + const __m256i p_1 = _mm256_madd_epi16(p01, mone); + + const __m256i p23 = mul_add_epi8(q4_23,q8_23); + const __m256i p_2 = _mm256_madd_epi16(p23, mone); + + const float dy0 = GGML_CPU_FP16_TO_FP32(y[2*ib].d); + const float dy1 = GGML_CPU_FP16_TO_FP32(y[2*ib+1].d); + + const float s0 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[0]) * dy0; + const float s1 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[1]) * dy0; + const float s2 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[2]) * dy1; + const float s3 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[3]) * dy1; + + const __m256 scales01 = _mm256_set_m128(_mm_set1_ps(s1), _mm_set1_ps(s0)); + const __m256 scales23 = _mm256_set_m128(_mm_set1_ps(s3), _mm_set1_ps(s2)); + + accum = _mm256_fmadd_ps(scales01, _mm256_cvtepi32_ps(p_1), accum); + accum = _mm256_fmadd_ps(scales23, _mm256_cvtepi32_ps(p_2), accum); + } + sumf = hsum_float_8(accum); + +#elif defined(__AVX__) + + const __m128i values128 = _mm_loadu_si128((const __m128i*)kvalues_fp4); + const __m128i m4b = _mm_set1_epi8(0x0f); + + __m256 accum = _mm256_setzero_ps(); + for(; ib < nb; ib++){ + + const __m128i q4bits_01 = _mm_loadu_si128((const __m128i *)(x[ib].qs + 0)); + const __m128i q4bits_23 = _mm_loadu_si128((const __m128i *)(x[ib].qs + 16)); + + const __m128i q8_0 = _mm_loadu_si128((const __m128i *)(y[2*ib + 0].qs + 0)); + const __m128i q8_1 = _mm_loadu_si128((const __m128i *)(y[2*ib + 0].qs + 16)); + const __m128i q8_2 = _mm_loadu_si128((const __m128i *)(y[2*ib + 1].qs + 0)); + const __m128i q8_3 = _mm_loadu_si128((const __m128i *)(y[2*ib + 1].qs + 16)); + + const __m128i q4_01_lo = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_01, m4b)); + const __m128i q4_01_hi = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_01, 4), m4b)); + const __m128i q4_23_lo = _mm_shuffle_epi8(values128, _mm_and_si128(q4bits_23, m4b)); + const __m128i q4_23_hi = _mm_shuffle_epi8(values128, _mm_and_si128(_mm_srli_epi16(q4bits_23, 4), m4b)); + + const __m128i q4_0 = _mm_unpacklo_epi64(q4_01_lo, q4_01_hi); + const __m128i q4_1 = _mm_unpackhi_epi64(q4_01_lo, q4_01_hi); + const __m128i q4_2 = _mm_unpacklo_epi64(q4_23_lo, q4_23_hi); + const __m128i q4_3 = _mm_unpackhi_epi64(q4_23_lo, q4_23_hi); + + const __m128i p0_i32 = mul_sum_i8_pairs(q4_0, q8_0); + const __m128i p1_i32 = mul_sum_i8_pairs(q4_1, q8_1); + const __m128i p2_i32 = mul_sum_i8_pairs(q4_2, q8_2); + const __m128i p3_i32 = mul_sum_i8_pairs(q4_3, q8_3); + + const __m128 p0 = _mm_cvtepi32_ps(p0_i32); + const __m128 p1 = _mm_cvtepi32_ps(p1_i32); + const __m128 p2 = _mm_cvtepi32_ps(p2_i32); + const __m128 p3 = _mm_cvtepi32_ps(p3_i32); + + const __m256 p01 = _mm256_set_m128(p1, p0); + const __m256 p23 = _mm256_set_m128(p3, p2); + + const float dy0 = GGML_CPU_FP16_TO_FP32(y[2*ib].d); + const float dy1 = GGML_CPU_FP16_TO_FP32(y[2*ib+1].d); + + const float s0 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[0]) * dy0; + const float s1 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[1]) * dy0; + const float s2 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[2]) * dy1; + const float s3 = GGML_CPU_UE4M3_TO_FP32(x[ib].d[3]) * dy1; + + const __m256 scales01 = _mm256_set_m128(_mm_set1_ps(s1), _mm_set1_ps(s0)); + const __m256 scales23 = _mm256_set_m128(_mm_set1_ps(s3), _mm_set1_ps(s2)); + + accum = _mm256_add_ps(accum, _mm256_mul_ps(p01, scales01)); + accum = _mm256_add_ps(accum, _mm256_mul_ps(p23, scales23)); + } + sumf = hsum_float_8(accum); + +#endif + + for (;ib < nb; ++ib) { + for (int s_idx = 0; s_idx < 4; ++s_idx) { + const float d = GGML_CPU_UE4M3_TO_FP32(x[ib].d[s_idx]); + const int q8_block = s_idx / 2; + const int q8_off = (s_idx % 2) * QK_NVFP4_SUB; + const float dy = GGML_CPU_FP16_TO_FP32(y[2*ib + q8_block].d); + + int sumi_lo = 0, sumi_hi = 0; + for (int j = 0; j < QK_NVFP4_SUB/2; ++j) { + const uint8_t qv = x[ib].qs[s_idx*(QK_NVFP4_SUB/2) + j]; + sumi_lo += y[2*ib + q8_block].qs[q8_off + j + 0] * kvalues_fp4[qv & 0xf]; + sumi_hi += y[2*ib + q8_block].qs[q8_off + j + QK_NVFP4_SUB/2] * kvalues_fp4[qv >> 4]; + } + + sumf += dy * d * (sumi_lo + sumi_hi); + } + } + *s = sumf; +} + void ggml_vec_dot_q5_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { const int qk = QK8_0; const int nb = n / qk; diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index eb8341c9a..8a0e5c255 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -82,6 +82,9 @@ float ggml_table_f32_f16[1 << 16]; // precomputed f32 table for e8m0 half (1 KB) (simd-mappings.h) float ggml_table_f32_e8m0_half[1 << 8]; +// precomputed f32 table for ue4m3 (1 KB) (simd-mappings.h) +float ggml_table_f32_ue4m3[1 << 8]; + #if defined(__ARM_ARCH) struct ggml_arm_arch_features_type { int sve_cnt; @@ -3798,6 +3801,11 @@ void ggml_cpu_init(void) { ggml_table_f32_e8m0_half[i] = GGML_E8M0_TO_FP32_HALF(i); } + // initialize UE4M3 table (256 entries) + for (int i = 0; i < (1 << 8); ++i) { + ggml_table_f32_ue4m3[i] = ggml_ue4m3_to_fp32(i); + } + const uint64_t t_end = ggml_time_us(); UNUSED(t_end); GGML_PRINT_DEBUG("%s: GELU, Quick GELU, SILU and EXP tables initialized in %f ms\n", __func__, (t_end - t_start)/1000.0); diff --git a/ggml/src/ggml-cpu/simd-mappings.h b/ggml/src/ggml-cpu/simd-mappings.h index 62e687201..be50c25c0 100644 --- a/ggml/src/ggml-cpu/simd-mappings.h +++ b/ggml/src/ggml-cpu/simd-mappings.h @@ -120,6 +120,10 @@ extern float ggml_table_f32_f16[1 << 16]; // defined in ggml-cpu.c, initialized in ggml_cpu_init() extern float ggml_table_f32_e8m0_half[1 << 8]; +// precomputed f32 table for ue4m3 (1 KB) +// defined in ggml-cpu.c, initialized in ggml_cpu_init() +extern float ggml_table_f32_ue4m3[1 << 8]; + // Use lookup table for E8M0 on x86 (faster than bit manipulation) #if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) #define GGML_CPU_E8M0_TO_FP32_HALF(x) ggml_table_f32_e8m0_half[(uint8_t)(x)] @@ -127,6 +131,13 @@ extern float ggml_table_f32_e8m0_half[1 << 8]; #define GGML_CPU_E8M0_TO_FP32_HALF(x) GGML_E8M0_TO_FP32_HALF(x) #endif +// Use lookup table for UE4M3 on x86 (faster than bit manipulation) +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) +#define GGML_CPU_UE4M3_TO_FP32(x) ggml_table_f32_ue4m3[(uint8_t)(x)] +#else +#define GGML_CPU_UE4M3_TO_FP32(x) ggml_ue4m3_to_fp32(x) +#endif + // On ARM NEON, it's quicker to directly convert x -> x instead of calling into ggml_lookup_fp16_to_fp32, // so we define GGML_CPU_FP16_TO_FP32 and GGML_CPU_FP32_TO_FP16 elsewhere for NEON. // This is also true for POWER9.