vrr/kvcache-ai-ktransformers
2
0
Fork 0
mirror of https://github.com/kvcache-ai/ktransformers.git synced 2026-10-19 10:13:32 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

[fix]: fix exp_avx512 for act_fn (#1797)
Some checks failed
Book-CI / test (push) Has been cancelled
Details
Book-CI / test-1 (push) Has been cancelled
Details
Book-CI / test-2 (push) Has been cancelled
Details
Deploy / deploy (macos-latest) (push) Has been cancelled
Details
Deploy / deploy (ubuntu-latest) (push) Has been cancelled
Details
Deploy / deploy (windows-latest) (push) Has been cancelled
Details

Browse source
This commit is contained in:
mrhaoxx 2026-01-20 11:07:22 +08:00 • committed by GitHub
parent 027832c590
commit b27de4068b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 8 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

8
kt-kernel/operators/amx/la/amx.hpp
View file

@ -35,10 +35,10 @@ static inline __m512 exp_avx512(__m512 x) {
const __m512 poly_6 = _mm512_set1_ps(0.0013333558f);
__m512 frac_exp = _mm512_fmadd_ps(
frac_part, poly_6,
_mm512_fmadd_ps(frac_part, poly_5,
_mm512_fmadd_ps(frac_part, poly_4,
_mm512_fmadd_ps(frac_part, poly_3, _mm512_fmadd_ps(frac_part, poly_2, poly_1)))));
_mm512_fmadd_ps(_mm512_fmadd_ps(_mm512_fmadd_ps(_mm512_fmadd_ps(poly_6, frac_part, poly_5), frac_part, poly_4),
frac_part, poly_3),
frac_part, poly_2),
frac_part, poly_1);
__m512 two_pow_i = _mm512_scalef_ps(_mm512_set1_ps(1.0f), _mm512_cvtepi32_ps(int_part));
return _mm512_mul_ps(two_pow_i, frac_exp);

8
kt-sft/csrc/ktransformers_ext/operators/amx/moe.hpp
View file

@ -51,10 +51,10 @@ static inline __m512 exp_avx512(__m512 x) {
const __m512 poly_6 = _mm512_set1_ps(0.0013333558f);
__m512 frac_exp = _mm512_fmadd_ps(
frac_part, poly_6,
_mm512_fmadd_ps(frac_part, poly_5,
_mm512_fmadd_ps(frac_part, poly_4,
_mm512_fmadd_ps(frac_part, poly_3, _mm512_fmadd_ps(frac_part, poly_2, poly_1)))));
_mm512_fmadd_ps(_mm512_fmadd_ps(_mm512_fmadd_ps(_mm512_fmadd_ps(poly_6, frac_part, poly_5), frac_part, poly_4),
frac_part, poly_3),
frac_part, poly_2),
frac_part, poly_1);
__m512 two_pow_i = _mm512_scalef_ps(_mm512_set1_ps(1.0f), _mm512_cvtepi32_ps(int_part));
return _mm512_mul_ps(two_pow_i, frac_exp);