[CPU][Perf]Added tanh AOR for faster gelu activations. (#44639)
Signed-off-by: Anna Mayne <anna.mayne@arm.com> Signed-off-by: almayne <anna.mayne@arm.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
almayne committed
Jul 1, 2026 at 07:24 UTC
89e99202f285d1dc1168f87f6cc3e26e930970ff
9 files changed
+328
-3
cmake/cpu_extension.cmake
+1
@@ -427,6 +427,7 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
427
set(VLLM_EXT_SRC
428
"csrc/cpu/shm.cpp"
429
"csrc/cpu/activation_lut_bf16.cpp"
430
+ "csrc/cpu/cpu_tanhf_neon.hpp"
431
"csrc/cpu/cpu_fused_moe.cpp"
432
${VLLM_EXT_SRC})
433
endif()
csrc/cpu/activation.cpp
+12
@@ -126,6 +126,18 @@ void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
126
});
127
}
128
129
+void gelu_tanh(torch::Tensor& out, torch::Tensor& input) {
130
+ int num_tokens = input.numel() / input.size(-1);
131
+ int d = input.size(-1);
132
+
133
+ VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_tanh_impl", [&] {
134
+ CPU_KERNEL_GUARD_IN(gelu_tanh_impl)
135
+ activation_kernel<scalar_t, gelu_tanh_act, false>(
136
+ num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
137
+ CPU_KERNEL_GUARD_OUT(gelu_tanh_impl)
138
+ });
139
+}
140
+
141
void gelu_new(torch::Tensor& out, torch::Tensor& input) {
142
int num_tokens = input.numel() / input.size(-1);
143
int d = input.size(-1);
csrc/cpu/cpu_tanhf_neon.hpp
new
+128
@@ -0,0 +1,128 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+#ifndef CPU_TANHF_NEON_HPP
5
+#define CPU_TANHF_NEON_HPP
6
+
7
+#include <cstdint>
8
+#include <arm_neon.h>
9
+
10
+namespace vec_op {
11
+
12
+namespace {
13
+
14
+struct TanhfConstants {
15
+ float32x4_t special_bound;
16
+ float32x4_t two;
17
+ float32x4_t c0;
18
+ float32x4_t c2;
19
+ int32x4_t exponent_bias;
20
+ float c1;
21
+ float c3;
22
+ float two_over_ln2;
23
+ float c4;
24
+ float ln2_hi;
25
+ float ln2_lo;
26
+};
27
+
28
+const TanhfConstants kTanhfConstants = {
29
+ // 9.01, above which tanhf rounds to 1 (or -1 for negative).
30
+ .special_bound = vdupq_n_f32(0x1.205966p+3f),
31
+ .two = vdupq_n_f32(0x1.0p+1f),
32
+ .c0 = vdupq_n_f32(0x1.fffffep-2f),
33
+ .c2 = vdupq_n_f32(0x1.555736p-5f),
34
+ .exponent_bias = vdupq_n_s32(0x3f800000),
35
+ .c1 = 0x1.5554aep-3f,
36
+ .c3 = 0x1.12287cp-7f,
37
+ .two_over_ln2 = 0x1.715476p+1f,
38
+ .c4 = 0x1.6b55a2p-10f,
39
+ .ln2_hi = 0x1.62e4p-1f,
40
+ .ln2_lo = 0x1.7f7d1cp-20f,
41
+};
42
+
43
+// Return the ptr but hide it's value from the compiler so accesses
44
+// through it can't be optimised based on contents.
45
+template <typename T>
46
+inline const T* ptr_barrier(const T* ptr) {
47
+ const T* opaque_ptr = ptr;
48
+ __asm__("" : "+r"(opaque_ptr));
49
+ return opaque_ptr;
50
+}
51
+
52
+// Check whether any lanes in the mask are set
53
+inline bool any_u32(uint32x4_t x) { return vmaxvq_u32(x) != 0; }
54
+
55
+// e^2x - 1 inline helper
56
+inline float32x4_t e2xm1f_inline(float32x4_t x, const TanhfConstants* d) {
57
+ float32x2_t ln2 = vld1_f32(&d->ln2_hi);
58
+ float32x4_t lane_consts = vld1q_f32(&d->c1);
59
+
60
+ // Reduce argument: f in [-ln2/2, ln2/2], i is exact.
61
+ float32x4_t j = vrndaq_f32(vmulq_laneq_f32(x, lane_consts, 2));
62
+ int32x4_t i = vcvtq_s32_f32(j);
63
+ float32x4_t f = vaddq_f32(x, x);
64
+ f = vfmsq_lane_f32(f, j, ln2, 0);
65
+ f = vfmsq_lane_f32(f, j, ln2, 1);
66
+
67
+ // Approximate expm1(f) with polynomial P, expm1(f) ~= f + f^2 * P(f)
68
+ float32x4_t f2 = vmulq_f32(f, f);
69
+ float32x4_t f4 = vmulq_f32(f2, f2);
70
+ float32x4_t p01 = vfmaq_laneq_f32(d->c0, f, lane_consts, 0);
71
+ float32x4_t p23 = vfmaq_laneq_f32(d->c2, f, lane_consts, 1);
72
+ float32x4_t poly = vfmaq_f32(p01, f2, p23);
73
+ poly = vfmaq_laneq_f32(poly, f4, lane_consts, 3);
74
+ poly = vfmaq_f32(f, f2, poly);
75
+
76
+ // scale = 2^i
77
+ int32x4_t u = vaddq_s32(vshlq_n_s32(i, 23), d->exponent_bias);
78
+ float32x4_t scale = vreinterpretq_f32_s32(u);
79
+ return vfmaq_f32(vsubq_f32(scale, vdupq_n_f32(1.0f)), poly, scale);
80
+}
81
+
82
+// Calculate the result tanh(x) = q / (q+2) and set special lanes to ±1
83
+inline float32x4_t special_case(float32x4_t x, float32x4_t q,
84
+ uint32x4_t special) {
85
+ const TanhfConstants* d = ptr_barrier(&kTanhfConstants);
86
+
87
+ float32x4_t y = vdivq_f32(q, vaddq_f32(q, d->two));
88
+ uint32x4_t ix = vreinterpretq_u32_f32(x);
89
+ uint32x4_t one_bits = vreinterpretq_u32_s32(d->exponent_bias);
90
+ uint32x4_t sign_mask = vdupq_n_u32(0x80000000u);
91
+ uint32x4_t special_bits = vbslq_u32(sign_mask, ix, one_bits);
92
+ float32x4_t special_y = vreinterpretq_f32_u32(special_bits);
93
+ return vbslq_f32(special, special_y, y);
94
+}
95
+
96
+} // namespace
97
+
98
+// Implementation of tanhf adapted from Arm Optimized Routines (tanhf
99
+// AdvSIMD)
100
+// https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/tanhf.c
101
+//
102
+// Approximation for single-precision vector tanh(x), using a simplified
103
+// version of expm1f. The maximum error is 2.08 + 0.5 ULP:
104
+// _ZGVnN4v_tanhf (0x1.fa5eep-5) got 0x1.f9ba02p-5 want 0x1.f9ba08p-5.
105
+inline float32x4_t fast_tanhf_f32x4(float32x4_t x) {
106
+ const TanhfConstants* d = ptr_barrier(&kTanhfConstants);
107
+
108
+ // tanh(x) = (e^2x - 1) / (e^2x + 1)
109
+ // q = e^2x -1
110
+ float32x4_t q = e2xm1f_inline(x, d);
111
+
112
+ // Check for special cases
113
+ uint32x4_t special = vcagtq_f32(x, d->special_bound);
114
+
115
+ // Fall back to vectorised special case for any lanes which would cause
116
+ // expm1 to overflow
117
+ if (any_u32(special)) {
118
+ return special_case(x, q, special);
119
+ }
120
+
121
+ // Complete fast path if no special lanes
122
+ // tanh(x) = q / (q+2)
123
+ return vdivq_f32(q, vaddq_f32(q, d->two));
124
+}
125
+
126
+} // namespace vec_op
127
+
128
+#endif // CPU_TANHF_NEON_HPP
\ No newline at end of file
csrc/cpu/cpu_types_arm.hpp
+22
@@ -3,6 +3,8 @@
3
4
#include <arm_neon.h>
5
6
+#include "cpu/cpu_tanhf_neon.hpp"
7
+
8
#include <torch/all.h>
9
#include <ATen/cpu/vec/functional.h>
10
#include <ATen/cpu/vec/vec.h>
@@ -345,6 +347,10 @@ struct FP32Vec4 : public VectorizedRegWrapper<FP32Vec4, 1, float> {
347
explicit FP32Vec4(float32x4_t data) : Base(VectorizedT(data)) {};
348
349
explicit FP32Vec4(const FP32Vec4& data) : Base(data) {};
350
+
351
+ FORCE_INLINE FP32Vec4 tanh() const {
352
+ return FP32Vec4(fast_tanhf_f32x4(reg.val[0]));
353
+ }
354
};
355
356
struct FP32Vec8 : public VectorizedRegWrapper<FP32Vec8, 2, float> {
@@ -391,6 +397,13 @@ struct FP32Vec8 : public VectorizedRegWrapper<FP32Vec8, 2, float> {
397
reg.val[1] = Vectorized<float>(data.val[1]);
398
}
399
400
+ FORCE_INLINE FP32Vec8 tanh() const {
401
+ FP32Vec8 r(uninit);
402
+ r.reg.val[0] = Vectorized<float>(fast_tanhf_f32x4(reg.val[0]));
403
+ r.reg.val[1] = Vectorized<float>(fast_tanhf_f32x4(reg.val[1]));
404
+ return r;
405
+ }
406
+
407
FORCE_INLINE float reduce_sum() const noexcept {
408
float answer = 0;
409
std::plus<VectorizedT> add;
@@ -497,6 +510,15 @@ struct FP32Vec16 : public VectorizedRegWrapper<FP32Vec16, 4, float> {
510
reg.val[3] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v.reg.val[1])));
511
};
512
513
+ FORCE_INLINE FP32Vec16 tanh() const {
514
+ FP32Vec16 r(uninit);
515
+ r.reg.val[0] = Vectorized<float>(fast_tanhf_f32x4(reg.val[0]));
516
+ r.reg.val[1] = Vectorized<float>(fast_tanhf_f32x4(reg.val[1]));
517
+ r.reg.val[2] = Vectorized<float>(fast_tanhf_f32x4(reg.val[2]));
518
+ r.reg.val[3] = Vectorized<float>(fast_tanhf_f32x4(reg.val[3]));
519
+ return r;
520
+ }
521
+
522
static FORCE_INLINE void load_even_odd(const float* ptr, FP32Vec16& even,
523
FP32Vec16& odd) noexcept {
524
const float32x4x2_t x01 = vuzpq_f32(vld1q_f32(ptr), vld1q_f32(ptr + 4));
csrc/cpu/torch_bindings.cpp
+4
@@ -298,6 +298,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
298
ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()");
299
ops.impl("gelu_tanh_and_mul", torch::kCPU, &gelu_tanh_and_mul);
300
301
+ // GELU tanh implementation.
302
+ ops.def("gelu_tanh(Tensor! out, Tensor input) -> ()");
303
+ ops.impl("gelu_tanh", torch::kCPU, &gelu_tanh);
304
+
305
// GELU implementation used in GPT-2.
306
ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
307
ops.impl("gelu_new", torch::kCPU, &gelu_new);
csrc/ops.h
+2
@@ -35,6 +35,8 @@ void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);
35
36
void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input);
37
38
+void gelu_tanh(torch::Tensor& out, torch::Tensor& input);
39
+
40
void gelu_new(torch::Tensor& out, torch::Tensor& input);
41
42
void gelu_fast(torch::Tensor& out, torch::Tensor& input);
tests/kernels/core/test_cpu_activation.py
+101
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+
5
import pytest
6
import torch
7
@@ -109,3 +110,103 @@ def test_cpu_unary_activation(
110
if not (activation_cls is GELU and dtype != torch.bfloat16):
111
raw_out = torch.empty_like(x)
112
opcheck(fn, (raw_out, x, *op_args))
113
+
114
+
115
+@pytest.mark.parametrize("dtype", DTYPES)
116
+@torch.inference_mode()
117
+def test_cpu_gelu_tanh_and_mul(
118
+ default_vllm_config,
119
+ dtype: torch.dtype,
120
+) -> None:
121
+ gate = torch.tensor(
122
+ [
123
+ [
124
+ -12.0,
125
+ -10.0,
126
+ -9.01,
127
+ -5.0,
128
+ -2.0,
129
+ -1.0,
130
+ -0.0,
131
+ 0.0,
132
+ 0.5,
133
+ 1.0,
134
+ 2.0,
135
+ 5.0,
136
+ 9.01,
137
+ 10.0,
138
+ 12.0,
139
+ 11.0,
140
+ ],
141
+ [
142
+ -7.5,
143
+ -4.5,
144
+ -3.0,
145
+ -1.5,
146
+ -0.75,
147
+ -0.25,
148
+ 0.25,
149
+ 0.75,
150
+ 1.5,
151
+ 3.0,
152
+ 4.5,
153
+ 7.5,
154
+ -11.0,
155
+ 11.0,
156
+ 8.75,
157
+ -8.75,
158
+ ],
159
+ ],
160
+ dtype=dtype,
161
+ )
162
+ val = torch.tensor(
163
+ [
164
+ [
165
+ 0.25,
166
+ -0.5,
167
+ 0.75,
168
+ -1.0,
169
+ 1.25,
170
+ -1.5,
171
+ 1.75,
172
+ -2.0,
173
+ 2.25,
174
+ -2.5,
175
+ 2.75,
176
+ -3.0,
177
+ 3.25,
178
+ -3.5,
179
+ 3.75,
180
+ -4.0,
181
+ ],
182
+ [
183
+ -0.4,
184
+ 0.6,
185
+ -0.8,
186
+ 1.0,
187
+ -1.2,
188
+ 1.4,
189
+ -1.6,
190
+ 1.8,
191
+ -2.0,
192
+ 2.2,
193
+ -2.4,
194
+ 2.6,
195
+ -2.8,
196
+ 3.0,
197
+ -3.2,
198
+ 3.4,
199
+ ],
200
+ ],
201
+ dtype=dtype,
202
+ )
203
+
204
+ x = torch.cat((val, gate), dim=-1).contiguous()
205
+ kernel_out = torch.empty_like(val)
206
+ torch.ops._C.gelu_tanh_and_mul(kernel_out, x)
207
+
208
+ torch_ref = torch.nn.functional.gelu(val, approximate="tanh") * gate
209
+
210
+ atol = get_default_atol(kernel_out)
211
+ rtol = get_default_rtol(kernel_out)
212
+ torch.testing.assert_close(kernel_out, torch_ref, atol=atol, rtol=rtol)
vllm/model_executor/layers/activation.py
+46
-3
@@ -313,8 +313,10 @@ class GELU(CustomOp):
313
314
def __init__(self):
315
super().__init__()
316
- if current_platform.get_cpu_architecture() == CpuArchEnum.ARM and hasattr(
317
- torch.ops._C, "activation_lut_bf16"
316
+ if (
317
+ current_platform.is_cpu()
318
+ and current_platform.get_cpu_architecture() == CpuArchEnum.ARM
319
+ and hasattr(torch.ops._C, "activation_lut_bf16")
320
):
321
self.op = torch.ops._C.activation_lut_bf16
322
else:
@@ -334,6 +336,36 @@ class GELU(CustomOp):
336
return self.forward_native(x)
337
338
339
+# --8<-- [start:gelu_tanh]
340
+@CustomOp.register("gelu_tanh")
341
+class GELUTanh(CustomOp):
342
+ # --8<-- [end:gelu_tanh]
343
+
344
+ def __init__(self):
345
+ super().__init__()
346
+ if (
347
+ current_platform.is_cpu()
348
+ and current_platform.get_cpu_architecture() == CpuArchEnum.ARM
349
+ and hasattr(torch.ops._C, "gelu_tanh")
350
+ ):
351
+ self.op = torch.ops._C.gelu_tanh
352
+ else:
353
+ self.op = None
354
+
355
+ def forward_native(self, x: torch.Tensor) -> torch.Tensor:
356
+ return F.gelu(x, approximate="tanh")
357
+
358
+ def forward_cpu(self, x: torch.Tensor) -> torch.Tensor:
359
+ if self.op:
360
+ out = torch.empty_like(x)
361
+ self.op(out, x)
362
+ return out
363
+ return self.forward_native(x)
364
+
365
+ def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
366
+ return self.forward_native(x)
367
+
368
+
369
# --8<-- [start:gelu_and_mul]
370
@CustomOp.register("gelu_and_mul")
371
class GeluAndMul(CustomOp):
@@ -385,6 +417,11 @@ class GeluAndMul(CustomOp):
417
self.op(out, x)
418
return out
419
420
+ def forward_cpu(self, x: torch.Tensor) -> torch.Tensor:
421
+ if self.op:
422
+ return self.forward_cuda(x)
423
+ return self.native(x)
424
+
425
def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
426
return self.forward_cuda(x)
427
@@ -739,7 +776,8 @@ _ACTIVATION_REGISTRY = LazyDict(
776
777
778
def _get_gelu_pytorch_tanh() -> nn.Module:
742
- """Get PyTorch GELU with tanh approximation, with ROCm fallback."""
779
+ """Get PyTorch GELU with tanh approximation, with ROCm fallback
780
+ and fast GELU for ARM."""
781
if current_platform.is_rocm():
782
# TODO:[ROCm] PyTorch native GELU with tanh is unstable with torch.compile
783
logger.warning_once(
@@ -747,6 +785,11 @@ def _get_gelu_pytorch_tanh() -> nn.Module:
785
"Falling back to GELU(approximate='none')."
786
)
787
return nn.GELU(approximate="none")
788
+ if (
789
+ current_platform.is_cpu()
790
+ and current_platform.get_cpu_architecture() == CpuArchEnum.ARM
791
+ ):
792
+ return GELUTanh()
793
return nn.GELU(approximate="tanh")
794
795
vllm/platforms/cpu.py
+12
@@ -193,6 +193,18 @@ class CpuPlatform(Platform):
193
and "-gelu" not in compilation_config.custom_ops
194
):
195
compilation_config.custom_ops.append("+gelu")
196
+ if (
197
+ cls.get_cpu_architecture() == CpuArchEnum.ARM
198
+ and "+gelu_tanh" not in compilation_config.custom_ops
199
+ and "-gelu_tanh" not in compilation_config.custom_ops
200
+ ):
201
+ compilation_config.custom_ops.append("+gelu_tanh")
202
+ if (
203
+ cls.get_cpu_architecture() == CpuArchEnum.ARM
204
+ and "+gelu_and_mul" not in compilation_config.custom_ops
205
+ and "-gelu_and_mul" not in compilation_config.custom_ops
206
+ ):
207
+ compilation_config.custom_ops.append("+gelu_and_mul")
208
209
vllm_config.profiler_config.torch_profiler_dump_cuda_time_total = False
210