A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 23:45 UTC
[CPU][Perf] Add fast vectorized exp impl from Arm Optimized Routines (#30068)
Signed-off-by: Ubuntu <ubuntu@ip-10-252-30-150.eu-west-1.compute.internal> Signed-off-by: Elham Harirpoush <elham.harirpoush@arm.com> Co-authored-by: Ubuntu <ubuntu@ip-10-252-30-150.eu-west-1.compute.internal>
Elham committed
Dec 5, 2025 at 08:09 UTC
9843e332da56307597deb2739bb83b85c18c5dde
2 files changed
+50
-13
csrc/cpu/cpu_attn_impl.hpp
-13
@@ -1246,14 +1246,8 @@ class AttentionMainLoop {
1246
// rescale sum and partial outputs
1247
if (need_rescale) {
1248
// compute rescale factor
1249
-#ifdef DEFINE_FAST_EXP
1250
- vec_op::FP32Vec16 rescale_factor_vec(rescale_factor);
1251
- rescale_factor_vec = fast_exp(rescale_factor_vec);
1252
- rescale_factor = rescale_factor_vec.get_last_elem();
1253
-#else
1249
rescale_factor = std::exp(rescale_factor);
1250
vec_op::FP32Vec16 rescale_factor_vec(rescale_factor);
1256
-#endif
1251
1252
// rescale sum
1253
new_sum_val += rescale_factor * init_sum_val;
@@ -1889,15 +1883,8 @@ class AttentionMainLoop {
1883
: curr_output_buffer;
1884
float rescale_factor = final_max > curr_max ? curr_max - final_max
1885
: final_max - curr_max;
1892
-
1893
-#ifdef DEFINE_FAST_EXP
1894
- vec_op::FP32Vec16 rescale_factor_vec(rescale_factor);
1895
- rescale_factor_vec = fast_exp(rescale_factor_vec);
1896
- rescale_factor = rescale_factor_vec.get_last_elem();
1897
-#else
1886
rescale_factor = std::exp(rescale_factor);
1887
vec_op::FP32Vec16 rescale_factor_vec(rescale_factor);
1900
-#endif
1888
1889
local_sum[head_idx] = final_max > curr_max
1890
? final_sum + rescale_factor * curr_sum
csrc/cpu/cpu_attn_macros.h
+50
@@ -60,4 +60,54 @@
60
61
#endif
62
63
+#ifdef __aarch64__
64
+ // Implementation copied from Arm Optimized Routines (expf AdvSIMD)
65
+ // https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/expf.c
66
+ #include <limits>
67
+ #define DEFINE_FAST_EXP \
68
+ const float32x4_t inv_ln2 = vdupq_n_f32(0x1.715476p+0f); \
69
+ const float ln2_hi = 0x1.62e4p-1f; \
70
+ const float ln2_lo = 0x1.7f7d1cp-20f; \
71
+ const float c0 = 0x1.0e4020p-7f; \
72
+ const float c2 = 0x1.555e66p-3f; \
73
+ const float32x4_t ln2_c02 = {ln2_hi, ln2_lo, c0, c2}; \
74
+ const uint32x4_t exponent_bias = vdupq_n_u32(0x3f800000); \
75
+ const float32x4_t c1 = vdupq_n_f32(0x1.573e2ep-5f); \
76
+ const float32x4_t c3 = vdupq_n_f32(0x1.fffdb6p-2f); \
77
+ const float32x4_t c4 = vdupq_n_f32(0x1.ffffecp-1f); \
78
+ const float32x4_t pos_special_bound = vdupq_n_f32(0x1.5d5e2ap+6f); \
79
+ const float32x4_t neg_special_bound = vnegq_f32(pos_special_bound); \
80
+ const float32x4_t inf = \
81
+ vdupq_n_f32(std::numeric_limits<float>::infinity()); \
82
+ const float32x4_t zero = vdupq_n_f32(0.0f); \
83
+ auto neon_expf = [&](float32x4_t values) __attribute__((always_inline)) { \
84
+ float32x4_t n = vrndaq_f32(vmulq_f32(values, inv_ln2)); \
85
+ float32x4_t r = vfmsq_laneq_f32(values, n, ln2_c02, 0); \
86
+ r = vfmsq_laneq_f32(r, n, ln2_c02, 1); \
87
+ uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_s32(vcvtq_s32_f32(n)), 23); \
88
+ float32x4_t scale = vreinterpretq_f32_u32(vaddq_u32(e, exponent_bias)); \
89
+ float32x4_t r2 = vmulq_f32(r, r); \
90
+ float32x4_t p = vfmaq_laneq_f32(c1, r, ln2_c02, 2); \
91
+ float32x4_t q = vfmaq_laneq_f32(c3, r, ln2_c02, 3); \
92
+ q = vfmaq_f32(q, p, r2); \
93
+ p = vmulq_f32(c4, r); \
94
+ float32x4_t poly = vfmaq_f32(p, q, r2); \
95
+ poly = vfmaq_f32(scale, poly, scale); \
96
+ const uint32x4_t hi_mask = vcgeq_f32(values, pos_special_bound); \
97
+ const uint32x4_t lo_mask = vcleq_f32(values, neg_special_bound); \
98
+ poly = vbslq_f32(hi_mask, inf, poly); \
99
+ return vbslq_f32(lo_mask, zero, poly); \
100
+ }; \
101
+ auto fast_exp = [&](vec_op::FP32Vec16& vec) \
102
+ __attribute__((always_inline)) { \
103
+ float32x4x4_t result; \
104
+ result.val[0] = neon_expf(vec.reg.val[0]); \
105
+ result.val[1] = neon_expf(vec.reg.val[1]); \
106
+ result.val[2] = neon_expf(vec.reg.val[2]); \
107
+ result.val[3] = neon_expf(vec.reg.val[3]); \
108
+ return vec_op::FP32Vec16(result); \
109
+ };
110
+
111
+#endif // __aarch64__
112
+
113
#endif
\ No newline at end of file