use faster covnersion from float8_e4m3fn to bfloat16 (#12316)
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user