chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[CPU] Use covt_e4m3_bf16 to optim BF16 to FP8 convert (#12191)

Browse Source
This commit is contained in:
wangyxbh
2025-11-14 01:36:51 -08:00
committed by GitHub
parent 3701f34dab
commit 5027739f2c
1 changed files with 4 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
sgl-kernel/csrc/cpu/vec.h
View File

@@ -52,14 +52,10 @@ convert_from_float_ext<at::BFloat16>(const Vectorized<float>& a, const Vectorize
// this doesn't handle NaN.
inline __m512bh cvt_e4m3_bf16_intrinsic_no_nan(__m256i fp8_vec) {
const __m512i x = _mm512_cvtepu8_epi16(fp8_vec);
const __m512i mant = _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(0x07)), 4);
const __m512i raw_exp = _mm512_srli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(0x78)), 3);
const __m512i exp = _mm512_slli_epi16(_mm512_add_epi16(raw_exp, _mm512_set1_epi16(120)), 7);
const __m512i nonsign = _mm512_or_si512(exp, mant);
const __m512i sign = _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(0x80)), 8);
const __m512i combined = _mm512_or_si512(nonsign, sign);
__m512i combined = _mm512_add_epi16(x, _mm512_set1_epi16(0x0780));
combined = _mm512_slli_epi16(combined, 4);
combined = _mm512_and_si512(combined, _mm512_set1_epi16(0x87f0));
combined = _mm512_add_epi16(combined, _mm512_set1_epi16(0x3c00));
const __mmask32 is_nonzero = _mm512_cmpneq_epi16_mask(x, _mm512_setzero_si512());
return (__m512bh)_mm512_maskz_mov_epi16(is_nonzero, combined);