From 89e99202f285d1dc1168f87f6cc3e26e930970ff Mon Sep 17 00:00:00 2001 From: almayne Date: Wed, 1 Jul 2026 07:24:40 +0100 Subject: [PATCH] [CPU][Perf]Added tanh AOR for faster gelu activations. (#44639) Signed-off-by: Anna Mayne Signed-off-by: almayne Co-authored-by: Li, Jiang --- cmake/cpu_extension.cmake | 1 + csrc/cpu/activation.cpp | 12 ++ csrc/cpu/cpu_tanhf_neon.hpp | 128 ++++++++++++++++++++++ csrc/cpu/cpu_types_arm.hpp | 22 ++++ csrc/cpu/torch_bindings.cpp | 4 + csrc/ops.h | 2 + tests/kernels/core/test_cpu_activation.py | 101 +++++++++++++++++ vllm/model_executor/layers/activation.py | 49 ++++++++- vllm/platforms/cpu.py | 12 ++ 9 files changed, 328 insertions(+), 3 deletions(-) create mode 100644 csrc/cpu/cpu_tanhf_neon.hpp diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 9d8796c0d7a..ddec286a0ca 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -427,6 +427,7 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) set(VLLM_EXT_SRC "csrc/cpu/shm.cpp" "csrc/cpu/activation_lut_bf16.cpp" + "csrc/cpu/cpu_tanhf_neon.hpp" "csrc/cpu/cpu_fused_moe.cpp" ${VLLM_EXT_SRC}) endif() diff --git a/csrc/cpu/activation.cpp b/csrc/cpu/activation.cpp index 039b8d5c30d..2f06813a194 100644 --- a/csrc/cpu/activation.cpp +++ b/csrc/cpu/activation.cpp @@ -126,6 +126,18 @@ void gelu_tanh_and_mul(torch::Tensor& out, // [..., d] }); } +void gelu_tanh(torch::Tensor& out, torch::Tensor& input) { + int num_tokens = input.numel() / input.size(-1); + int d = input.size(-1); + + VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_tanh_impl", [&] { + CPU_KERNEL_GUARD_IN(gelu_tanh_impl) + activation_kernel( + num_tokens, d, input.data_ptr(), out.data_ptr()); + CPU_KERNEL_GUARD_OUT(gelu_tanh_impl) + }); +} + void gelu_new(torch::Tensor& out, torch::Tensor& input) { int num_tokens = input.numel() / input.size(-1); int d = input.size(-1); diff --git a/csrc/cpu/cpu_tanhf_neon.hpp b/csrc/cpu/cpu_tanhf_neon.hpp new file mode 100644 index 00000000000..2ea7f336513 --- /dev/null +++ b/csrc/cpu/cpu_tanhf_neon.hpp @@ -0,0 +1,128 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#ifndef CPU_TANHF_NEON_HPP +#define CPU_TANHF_NEON_HPP + +#include +#include + +namespace vec_op { + +namespace { + +struct TanhfConstants { + float32x4_t special_bound; + float32x4_t two; + float32x4_t c0; + float32x4_t c2; + int32x4_t exponent_bias; + float c1; + float c3; + float two_over_ln2; + float c4; + float ln2_hi; + float ln2_lo; +}; + +const TanhfConstants kTanhfConstants = { + // 9.01, above which tanhf rounds to 1 (or -1 for negative). + .special_bound = vdupq_n_f32(0x1.205966p+3f), + .two = vdupq_n_f32(0x1.0p+1f), + .c0 = vdupq_n_f32(0x1.fffffep-2f), + .c2 = vdupq_n_f32(0x1.555736p-5f), + .exponent_bias = vdupq_n_s32(0x3f800000), + .c1 = 0x1.5554aep-3f, + .c3 = 0x1.12287cp-7f, + .two_over_ln2 = 0x1.715476p+1f, + .c4 = 0x1.6b55a2p-10f, + .ln2_hi = 0x1.62e4p-1f, + .ln2_lo = 0x1.7f7d1cp-20f, +}; + +// Return the ptr but hide it's value from the compiler so accesses +// through it can't be optimised based on contents. +template +inline const T* ptr_barrier(const T* ptr) { + const T* opaque_ptr = ptr; + __asm__("" : "+r"(opaque_ptr)); + return opaque_ptr; +} + +// Check whether any lanes in the mask are set +inline bool any_u32(uint32x4_t x) { return vmaxvq_u32(x) != 0; } + +// e^2x - 1 inline helper +inline float32x4_t e2xm1f_inline(float32x4_t x, const TanhfConstants* d) { + float32x2_t ln2 = vld1_f32(&d->ln2_hi); + float32x4_t lane_consts = vld1q_f32(&d->c1); + + // Reduce argument: f in [-ln2/2, ln2/2], i is exact. + float32x4_t j = vrndaq_f32(vmulq_laneq_f32(x, lane_consts, 2)); + int32x4_t i = vcvtq_s32_f32(j); + float32x4_t f = vaddq_f32(x, x); + f = vfmsq_lane_f32(f, j, ln2, 0); + f = vfmsq_lane_f32(f, j, ln2, 1); + + // Approximate expm1(f) with polynomial P, expm1(f) ~= f + f^2 * P(f) + float32x4_t f2 = vmulq_f32(f, f); + float32x4_t f4 = vmulq_f32(f2, f2); + float32x4_t p01 = vfmaq_laneq_f32(d->c0, f, lane_consts, 0); + float32x4_t p23 = vfmaq_laneq_f32(d->c2, f, lane_consts, 1); + float32x4_t poly = vfmaq_f32(p01, f2, p23); + poly = vfmaq_laneq_f32(poly, f4, lane_consts, 3); + poly = vfmaq_f32(f, f2, poly); + + // scale = 2^i + int32x4_t u = vaddq_s32(vshlq_n_s32(i, 23), d->exponent_bias); + float32x4_t scale = vreinterpretq_f32_s32(u); + return vfmaq_f32(vsubq_f32(scale, vdupq_n_f32(1.0f)), poly, scale); +} + +// Calculate the result tanh(x) = q / (q+2) and set special lanes to ±1 +inline float32x4_t special_case(float32x4_t x, float32x4_t q, + uint32x4_t special) { + const TanhfConstants* d = ptr_barrier(&kTanhfConstants); + + float32x4_t y = vdivq_f32(q, vaddq_f32(q, d->two)); + uint32x4_t ix = vreinterpretq_u32_f32(x); + uint32x4_t one_bits = vreinterpretq_u32_s32(d->exponent_bias); + uint32x4_t sign_mask = vdupq_n_u32(0x80000000u); + uint32x4_t special_bits = vbslq_u32(sign_mask, ix, one_bits); + float32x4_t special_y = vreinterpretq_f32_u32(special_bits); + return vbslq_f32(special, special_y, y); +} + +} // namespace + +// Implementation of tanhf adapted from Arm Optimized Routines (tanhf +// AdvSIMD) +// https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/tanhf.c +// +// Approximation for single-precision vector tanh(x), using a simplified +// version of expm1f. The maximum error is 2.08 + 0.5 ULP: +// _ZGVnN4v_tanhf (0x1.fa5eep-5) got 0x1.f9ba02p-5 want 0x1.f9ba08p-5. +inline float32x4_t fast_tanhf_f32x4(float32x4_t x) { + const TanhfConstants* d = ptr_barrier(&kTanhfConstants); + + // tanh(x) = (e^2x - 1) / (e^2x + 1) + // q = e^2x -1 + float32x4_t q = e2xm1f_inline(x, d); + + // Check for special cases + uint32x4_t special = vcagtq_f32(x, d->special_bound); + + // Fall back to vectorised special case for any lanes which would cause + // expm1 to overflow + if (any_u32(special)) { + return special_case(x, q, special); + } + + // Complete fast path if no special lanes + // tanh(x) = q / (q+2) + return vdivq_f32(q, vaddq_f32(q, d->two)); +} + +} // namespace vec_op + +#endif // CPU_TANHF_NEON_HPP \ No newline at end of file diff --git a/csrc/cpu/cpu_types_arm.hpp b/csrc/cpu/cpu_types_arm.hpp index fc987f706a5..294dee90bd8 100644 --- a/csrc/cpu/cpu_types_arm.hpp +++ b/csrc/cpu/cpu_types_arm.hpp @@ -3,6 +3,8 @@ #include +#include "cpu/cpu_tanhf_neon.hpp" + #include #include #include @@ -345,6 +347,10 @@ struct FP32Vec4 : public VectorizedRegWrapper { explicit FP32Vec4(float32x4_t data) : Base(VectorizedT(data)) {}; explicit FP32Vec4(const FP32Vec4& data) : Base(data) {}; + + FORCE_INLINE FP32Vec4 tanh() const { + return FP32Vec4(fast_tanhf_f32x4(reg.val[0])); + } }; struct FP32Vec8 : public VectorizedRegWrapper { @@ -391,6 +397,13 @@ struct FP32Vec8 : public VectorizedRegWrapper { reg.val[1] = Vectorized(data.val[1]); } + FORCE_INLINE FP32Vec8 tanh() const { + FP32Vec8 r(uninit); + r.reg.val[0] = Vectorized(fast_tanhf_f32x4(reg.val[0])); + r.reg.val[1] = Vectorized(fast_tanhf_f32x4(reg.val[1])); + return r; + } + FORCE_INLINE float reduce_sum() const noexcept { float answer = 0; std::plus add; @@ -497,6 +510,15 @@ struct FP32Vec16 : public VectorizedRegWrapper { reg.val[3] = Vectorized(vcvt_f32_f16(vget_high_f16(v.reg.val[1]))); }; + FORCE_INLINE FP32Vec16 tanh() const { + FP32Vec16 r(uninit); + r.reg.val[0] = Vectorized(fast_tanhf_f32x4(reg.val[0])); + r.reg.val[1] = Vectorized(fast_tanhf_f32x4(reg.val[1])); + r.reg.val[2] = Vectorized(fast_tanhf_f32x4(reg.val[2])); + r.reg.val[3] = Vectorized(fast_tanhf_f32x4(reg.val[3])); + return r; + } + static FORCE_INLINE void load_even_odd(const float* ptr, FP32Vec16& even, FP32Vec16& odd) noexcept { const float32x4x2_t x01 = vuzpq_f32(vld1q_f32(ptr), vld1q_f32(ptr + 4)); diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index bc02511eb80..e17c9ab3a7e 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -298,6 +298,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()"); ops.impl("gelu_tanh_and_mul", torch::kCPU, &gelu_tanh_and_mul); + // GELU tanh implementation. + ops.def("gelu_tanh(Tensor! out, Tensor input) -> ()"); + ops.impl("gelu_tanh", torch::kCPU, &gelu_tanh); + // GELU implementation used in GPT-2. ops.def("gelu_new(Tensor! out, Tensor input) -> ()"); ops.impl("gelu_new", torch::kCPU, &gelu_new); diff --git a/csrc/ops.h b/csrc/ops.h index c310bd59ff5..0cf73f6bfb3 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -35,6 +35,8 @@ void gelu_and_mul(torch::Tensor& out, torch::Tensor& input); void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input); +void gelu_tanh(torch::Tensor& out, torch::Tensor& input); + void gelu_new(torch::Tensor& out, torch::Tensor& input); void gelu_fast(torch::Tensor& out, torch::Tensor& input); diff --git a/tests/kernels/core/test_cpu_activation.py b/tests/kernels/core/test_cpu_activation.py index 40b5f045468..110c92042e0 100644 --- a/tests/kernels/core/test_cpu_activation.py +++ b/tests/kernels/core/test_cpu_activation.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + import pytest import torch @@ -109,3 +110,103 @@ def test_cpu_unary_activation( if not (activation_cls is GELU and dtype != torch.bfloat16): raw_out = torch.empty_like(x) opcheck(fn, (raw_out, x, *op_args)) + + +@pytest.mark.parametrize("dtype", DTYPES) +@torch.inference_mode() +def test_cpu_gelu_tanh_and_mul( + default_vllm_config, + dtype: torch.dtype, +) -> None: + gate = torch.tensor( + [ + [ + -12.0, + -10.0, + -9.01, + -5.0, + -2.0, + -1.0, + -0.0, + 0.0, + 0.5, + 1.0, + 2.0, + 5.0, + 9.01, + 10.0, + 12.0, + 11.0, + ], + [ + -7.5, + -4.5, + -3.0, + -1.5, + -0.75, + -0.25, + 0.25, + 0.75, + 1.5, + 3.0, + 4.5, + 7.5, + -11.0, + 11.0, + 8.75, + -8.75, + ], + ], + dtype=dtype, + ) + val = torch.tensor( + [ + [ + 0.25, + -0.5, + 0.75, + -1.0, + 1.25, + -1.5, + 1.75, + -2.0, + 2.25, + -2.5, + 2.75, + -3.0, + 3.25, + -3.5, + 3.75, + -4.0, + ], + [ + -0.4, + 0.6, + -0.8, + 1.0, + -1.2, + 1.4, + -1.6, + 1.8, + -2.0, + 2.2, + -2.4, + 2.6, + -2.8, + 3.0, + -3.2, + 3.4, + ], + ], + dtype=dtype, + ) + + x = torch.cat((val, gate), dim=-1).contiguous() + kernel_out = torch.empty_like(val) + torch.ops._C.gelu_tanh_and_mul(kernel_out, x) + + torch_ref = torch.nn.functional.gelu(val, approximate="tanh") * gate + + atol = get_default_atol(kernel_out) + rtol = get_default_rtol(kernel_out) + torch.testing.assert_close(kernel_out, torch_ref, atol=atol, rtol=rtol) diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index 80bf251b2d8..0115912ce4c 100644 --- a/vllm/model_executor/layers/activation.py +++ b/vllm/model_executor/layers/activation.py @@ -313,8 +313,10 @@ class GELU(CustomOp): def __init__(self): super().__init__() - if current_platform.get_cpu_architecture() == CpuArchEnum.ARM and hasattr( - torch.ops._C, "activation_lut_bf16" + if ( + current_platform.is_cpu() + and current_platform.get_cpu_architecture() == CpuArchEnum.ARM + and hasattr(torch.ops._C, "activation_lut_bf16") ): self.op = torch.ops._C.activation_lut_bf16 else: @@ -334,6 +336,36 @@ class GELU(CustomOp): return self.forward_native(x) +# --8<-- [start:gelu_tanh] +@CustomOp.register("gelu_tanh") +class GELUTanh(CustomOp): + # --8<-- [end:gelu_tanh] + + def __init__(self): + super().__init__() + if ( + current_platform.is_cpu() + and current_platform.get_cpu_architecture() == CpuArchEnum.ARM + and hasattr(torch.ops._C, "gelu_tanh") + ): + self.op = torch.ops._C.gelu_tanh + else: + self.op = None + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + return F.gelu(x, approximate="tanh") + + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: + if self.op: + out = torch.empty_like(x) + self.op(out, x) + return out + return self.forward_native(x) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_native(x) + + # --8<-- [start:gelu_and_mul] @CustomOp.register("gelu_and_mul") class GeluAndMul(CustomOp): @@ -385,6 +417,11 @@ class GeluAndMul(CustomOp): self.op(out, x) return out + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: + if self.op: + return self.forward_cuda(x) + return self.native(x) + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: return self.forward_cuda(x) @@ -739,7 +776,8 @@ _ACTIVATION_REGISTRY = LazyDict( def _get_gelu_pytorch_tanh() -> nn.Module: - """Get PyTorch GELU with tanh approximation, with ROCm fallback.""" + """Get PyTorch GELU with tanh approximation, with ROCm fallback + and fast GELU for ARM.""" if current_platform.is_rocm(): # TODO:[ROCm] PyTorch native GELU with tanh is unstable with torch.compile logger.warning_once( @@ -747,6 +785,11 @@ def _get_gelu_pytorch_tanh() -> nn.Module: "Falling back to GELU(approximate='none')." ) return nn.GELU(approximate="none") + if ( + current_platform.is_cpu() + and current_platform.get_cpu_architecture() == CpuArchEnum.ARM + ): + return GELUTanh() return nn.GELU(approximate="tanh") diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 369e07dd256..571a8c9c2cc 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -193,6 +193,18 @@ class CpuPlatform(Platform): and "-gelu" not in compilation_config.custom_ops ): compilation_config.custom_ops.append("+gelu") + if ( + cls.get_cpu_architecture() == CpuArchEnum.ARM + and "+gelu_tanh" not in compilation_config.custom_ops + and "-gelu_tanh" not in compilation_config.custom_ops + ): + compilation_config.custom_ops.append("+gelu_tanh") + if ( + cls.get_cpu_architecture() == CpuArchEnum.ARM + and "+gelu_and_mul" not in compilation_config.custom_ops + and "-gelu_and_mul" not in compilation_config.custom_ops + ): + compilation_config.custom_ops.append("+gelu_and_mul") vllm_config.profiler_config.torch_profiler_dump_cuda_time_total = False