use faster covnersion from float8_e4m3fn to bfloat16 (#12316)

This commit is contained in:
Ma Mingfei
2025-12-04 12:34:05 +08:00
committed by GitHub
parent b2b09f5f24
commit 46f21a5956
2 changed files with 33 additions and 5 deletions
+24
View File
@@ -121,6 +121,30 @@ inline __m512bh CVT_FP8_TO_BF16(__m256i a) {
#endif
}
// faster version of float8_e4m3fn conversion to bfloat16
//
// we mapped cuda implementation from below link and vectorized with avx512:
// https://github.com/thu-pacman/chitu/blob/1ed2078ec26581ebdca05b7306d4385f86edaa7c/csrc/cuda/marlin/marlin_gemm/dequant.h#L387
//
inline __attribute__((always_inline)) __m512bh CVT_FP8_TO_BF16_EXT(__m256i a) {
const __m512i mask0 = _mm512_set1_epi16(0x80); // sign bit
const __m512i mask1 = _mm512_set1_epi16(0x7F); // exponent and mantissa
const __m512i mask2 = _mm512_set1_epi16(0x4000);
__m512i x = _mm512_cvtepu8_epi16(a);
__m512i vsign = _mm512_and_si512(x, mask0);
vsign = _mm512_slli_epi16(vsign, 8);
__m512i vexp_and_mant = _mm512_and_si512(x, mask1);
vexp_and_mant = _mm512_slli_epi16(vexp_and_mant, 4);
// _MM_TERNLOG_A | _MM_TERNLOG_B | _MM_TERNLOG_C: 0b11111110
return (__m512bh)(_mm512_ternarylogic_epi32(vsign, mask2, vexp_and_mant, 0b11111110));
}
// bias for conversion of fp8 to bf16 1/256 in float32
#define kFP8_BIAS 0x3b800000
#endif
// vector to scalar reduction