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 4, 17:30 UTC
[CPU] Added faster exp routine for lower precision data types. (#38112)
Signed-off-by: Anna Mayne <anna.mayne@arm.com> Co-authored-by: Fadi Arafeh <fadi.arafeh@arm.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
almayne committed
Apr 23, 2026 at 14:14 UTC
2f314bc5e6706bffe77933ae5b13756d04641ed8
2 files changed
+76
-6
csrc/cpu/cpu_arch_macros.h
+48
-2
@@ -61,8 +61,23 @@
61
#endif
62
63
#ifdef __aarch64__
64
- // Implementation copied from Arm Optimized Routines (expf AdvSIMD)
64
+ // Implementation of neon_expf copied from Arm Optimized Routines (expf
65
+ // AdvSIMD)
66
// https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/expf.c
67
+ //
68
+ // Additional fast exponential intended for cases where outputs will be
69
+ // downcasted to FP16 / BF16 (e.g. attention softmax). Accurate within 1 ULP
70
+ // for FP16 Accurate within 1 ULP for BF16 for inputs in [-87.683, 88.376] &
71
+ // clamps inputs outside this range to 0 / inf. Implementation is similar to
72
+ // exp_u20, but:
73
+ // - uses a third degree polynomial approximation for exp(r) instead of a
74
+ // fifth degree one, with coefficients re-tuned.
75
+ // - does not split natural log (ln) into high / low parts
76
+ // - clamps exp(x) to 0 for x < -87.683113f and inf for x > 88.3762589f
77
+ // exp(x) = 2^n (exp(r))
78
+ // r = x - n*ln2, with n = round(x/ln2)
79
+ // exp(r) ~ poly(r) = 1 + r + r^2 * (c3 + c2 * r)
80
+ // n = round(x / ln2), r = x - n*ln2
81
#include <limits>
82
#define DEFINE_FAST_EXP \
83
const float32x4_t inv_ln2 = vdupq_n_f32(0x1.715476p+0f); \
@@ -106,7 +121,38 @@
121
result.val[2] = neon_expf(vec.reg.val[2]); \
122
result.val[3] = neon_expf(vec.reg.val[3]); \
123
return vec_op::FP32Vec16(result); \
109
- };
124
+ }; \
125
+ const float32x4_t lower_bound = vdupq_n_f32(-0x1.5ebb82p+6f); \
126
+ const float32x4_t upper_bound = vdupq_n_f32(0x1.61814ap+6f); \
127
+ constexpr float ln2 = 0x1.62e43p-1f; \
128
+ constexpr float f_c2 = 0x1.5592ecp-3f; \
129
+ const float32x4_t f_c3 = vdupq_n_f32(0x1.017d34p-1f); \
130
+ auto neon_expf_f16 = [&](float32x4_t values) __attribute__(( \
131
+ always_inline)) { \
132
+ const uint32x4_t lt_lower = vcltq_f32(values, lower_bound); \
133
+ const uint32x4_t gt_upper = vcgtq_f32(values, upper_bound); \
134
+ float32x4_t n = vrndaq_f32(vmulq_f32(values, inv_ln2)); \
135
+ float32x4_t r = vfmsq_n_f32(values, n, ln2); \
136
+ uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_s32(vcvtq_s32_f32(n)), 23); \
137
+ float32x4_t r2 = vmulq_f32(r, r); \
138
+ float32x4_t q = vfmaq_n_f32(f_c3, r, f_c2); \
139
+ float32x4_t s = vaddq_f32(vdupq_n_f32(1.0f), r); \
140
+ float32x4_t p = vfmaq_f32(s, q, r2); \
141
+ float32x4_t y = \
142
+ vreinterpretq_f32_u32(vaddq_u32(vreinterpretq_u32_f32(p), e)); \
143
+ y = vbslq_f32(lt_lower, vdupq_n_f32(0.0f), y); \
144
+ y = vbslq_f32(gt_upper, vdupq_n_f32(INFINITY), y); \
145
+ return y; \
146
+ }; \
147
+ auto fast_exp_f16 = [&](const vec_op::FP32Vec16& vec) \
148
+ __attribute__((always_inline)) { \
149
+ float32x4x4_t result; \
150
+ result.val[0] = neon_expf_f16(vec.reg.val[0]); \
151
+ result.val[1] = neon_expf_f16(vec.reg.val[1]); \
152
+ result.val[2] = neon_expf_f16(vec.reg.val[2]); \
153
+ result.val[3] = neon_expf_f16(vec.reg.val[3]); \
154
+ return vec_op::FP32Vec16(result); \
155
+ };
156
157
#endif // __aarch64__
158
csrc/cpu/cpu_attn_impl.hpp
+28
-4
@@ -1152,7 +1152,11 @@ class AttentionMainLoop {
1152
bool use_sink) {
1153
#ifdef DEFINE_FAST_EXP
1154
DEFINE_FAST_EXP
1155
+ bool constexpr IsReducedPrecision =
1156
+ std::is_same_v<query_t, c10::BFloat16> ||
1157
+ std::is_same_v<query_t, c10::Half>;
1158
#endif
1159
+
1160
using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;
1161
static_assert(sizeof(prob_buffer_t) <= sizeof(logits_buffer_t));
1162
@@ -1201,8 +1205,17 @@ class AttentionMainLoop {
1205
vec = vec - max_vec;
1206
1207
// compute exp
1204
-#ifdef DEFINE_FAST_EXP
1205
- vec = fast_exp(vec);
1208
+
1209
+#if defined(DEFINE_FAST_EXP)
1210
+ #ifdef __aarch64__
1211
+ if constexpr (IsReducedPrecision) {
1212
+ vec = fast_exp_f16(vec);
1213
+ } else
1214
+ #endif
1215
+ {
1216
+ vec = fast_exp(vec);
1217
+ }
1218
+
1219
prob_buffer_vec_t output_vec(vec);
1220
output_vec.save(curr_prob_buffer_iter);
1221
#else
@@ -1258,7 +1271,11 @@ class AttentionMainLoop {
1271
int32_t kv_tile_token_num, float softcap_scale) {
1272
#ifdef DEFINE_FAST_EXP
1273
DEFINE_FAST_EXP
1274
+ bool constexpr IsReducedPrecision =
1275
+ std::is_same_v<query_t, c10::BFloat16> ||
1276
+ std::is_same_v<query_t, c10::Half>;
1277
#endif
1278
+
1279
float inv_softcap_scale = 1.0 / softcap_scale;
1280
vec_op::FP32Vec16 softcap_scale_vec(softcap_scale);
1281
vec_op::FP32Vec16 inv_softcap_scale_vec(inv_softcap_scale);
@@ -1272,8 +1289,15 @@ class AttentionMainLoop {
1289
vec_op::FP32Vec16 vec(curr_logits_buffer_iter);
1290
vec = vec * inv_softcap_scale_vec;
1291
1275
-#ifdef DEFINE_FAST_EXP
1276
- vec = fast_exp(vec);
1292
+#if defined(DEFINE_FAST_EXP)
1293
+ #ifdef __aarch64__
1294
+ if constexpr (IsReducedPrecision) {
1295
+ vec = fast_exp_f16(vec);
1296
+ } else
1297
+ #endif
1298
+ {
1299
+ vec = fast_exp(vec);
1300
+ }
1301
vec_op::FP32Vec16 inv_vec = ones_vec / vec;
1302
vec = (vec - inv_vec) / (vec + inv_vec);
1303
#else