karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[CPU][Perf]Added tanh AOR for faster gelu activations. (#44639)

Browse Source 
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>
This commit is contained in:
almayne
2026-06-30 23:24:40 -07:00
committed by GitHub
co-authored by Li, Jiang <jiang1.li@intel.com>
parent b446792306
commit 89e99202f2
9 changed files with 328 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmake/cpu_extension.cmake
+1
View File
@@ -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()
csrc/cpu/activation.cpp
+12
View File
@@ -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<scalar_t, gelu_tanh_act, false>(
num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
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);
csrc/cpu/cpu_tanhf_neon.hpp
+128
View File
@@ -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 <cstdint>
#include <arm_neon.h>
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 <typename T>
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
csrc/cpu/cpu_types_arm.hpp
+22
View File
@@ -3,6 +3,8 @@
#include <arm_neon.h>
#include "cpu/cpu_tanhf_neon.hpp"
#include <torch/all.h>
#include <ATen/cpu/vec/functional.h>
#include <ATen/cpu/vec/vec.h>
@@ -345,6 +347,10 @@ struct FP32Vec4 : public VectorizedRegWrapper<FP32Vec4, 1, float> {
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<FP32Vec8, 2, float> {
@@ -391,6 +397,13 @@ struct FP32Vec8 : public VectorizedRegWrapper<FP32Vec8, 2, float> {
reg.val[1] = Vectorized<float>(data.val[1]);
}
FORCE_INLINE FP32Vec8 tanh() const {
FP32Vec8 r(uninit);
r.reg.val[0] = Vectorized<float>(fast_tanhf_f32x4(reg.val[0]));
r.reg.val[1] = Vectorized<float>(fast_tanhf_f32x4(reg.val[1]));
return r;
}
FORCE_INLINE float reduce_sum() const noexcept {
float answer = 0;
std::plus<VectorizedT> add;
@@ -497,6 +510,15 @@ struct FP32Vec16 : public VectorizedRegWrapper<FP32Vec16, 4, float> {
reg.val[3] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v.reg.val[1])));
};
FORCE_INLINE FP32Vec16 tanh() const {
FP32Vec16 r(uninit);
r.reg.val[0] = Vectorized<float>(fast_tanhf_f32x4(reg.val[0]));
r.reg.val[1] = Vectorized<float>(fast_tanhf_f32x4(reg.val[1]));
r.reg.val[2] = Vectorized<float>(fast_tanhf_f32x4(reg.val[2]));
r.reg.val[3] = Vectorized<float>(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));
csrc/cpu/torch_bindings.cpp
+4
View File
@@ -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);
csrc/ops.h
+2
View File
@@ -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);
tests/kernels/core/test_cpu_activation.py
+101
View File
@@ -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)
vllm/model_executor/layers/activation.py
+46 -3
View File
@@ -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")
vllm/platforms/cpu.py
+12
View File
@@ -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