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 5, 02:15 UTC
[perf][cpu] Accelerate BF16 GELU with LUT impl on Arm CPUs (#37469)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Fadi Arafeh committed
Apr 16, 2026 at 07:26 UTC
445b7093fdbbc6c52875d2deef28df1d5da56089
9 files changed
+240
-2
.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
+1
@@ -51,6 +51,7 @@ function cpu_tests() {
51
set -e
52
pytest -x -v -s tests/kernels/test_onednn.py
53
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
54
+ pytest -x -v -s tests/kernels/core/test_cpu_activation.py
55
pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic"
56
57
# basic online serving
cmake/cpu_extension.cmake
+1
@@ -360,6 +360,7 @@ set(VLLM_EXT_SRC
360
if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
361
set(VLLM_EXT_SRC
362
"csrc/cpu/shm.cpp"
363
+ "csrc/cpu/activation_lut_bf16.cpp"
364
${VLLM_EXT_SRC})
365
endif()
366
csrc/cpu/activation_lut_bf16.cpp
new
+71
@@ -0,0 +1,71 @@
1
+#include "cpu_types.hpp"
2
+
3
+#include <array>
4
+#include <cstdint>
5
+#include <mutex>
6
+#include <string>
7
+
8
+#include <ATen/ops/empty.h>
9
+#include <ATen/ops/gelu.h>
10
+#include <c10/util/BFloat16.h>
11
+
12
+constexpr uint32_t ActivationLutSize = 1u << 16;
13
+
14
+at::Tensor gelu_reference(const at::Tensor& x) { return at::gelu(x, "none"); }
15
+
16
+void maybe_init_activation_lut_bf16(
17
+ uint16_t* lut, std::once_flag& once,
18
+ at::Tensor (*activation)(const at::Tensor&)) {
19
+ std::call_once(once, [&]() {
20
+ auto lut_input =
21
+ at::empty({static_cast<int64_t>(ActivationLutSize)},
22
+ at::TensorOptions().device(at::kCPU).dtype(at::kFloat));
23
+ auto* lut_input_ptr = lut_input.data_ptr<float>();
24
+#pragma omp parallel for
25
+ for (uint32_t i = 0; i < ActivationLutSize; ++i) {
26
+ lut_input_ptr[i] = c10::detail::f32_from_bits(static_cast<uint16_t>(i));
27
+ }
28
+
29
+ auto lut_output = activation(lut_input);
30
+ const auto* lut_output_ptr = lut_output.data_ptr<float>();
31
+#pragma omp parallel for
32
+ for (uint32_t i = 0; i < ActivationLutSize; ++i) {
33
+ lut[i] = c10::detail::round_to_nearest_even(lut_output_ptr[i]);
34
+ }
35
+ });
36
+}
37
+
38
+void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
39
+ const uint16_t* lut, const char* op_name) {
40
+ TORCH_CHECK(input.scalar_type() == at::kBFloat16, op_name,
41
+ ": input must be bfloat16");
42
+ TORCH_CHECK(out.scalar_type() == at::kBFloat16, op_name,
43
+ ": out must be bfloat16");
44
+ TORCH_CHECK(input.is_contiguous(), op_name, ": input must be contiguous");
45
+ TORCH_CHECK(out.is_contiguous(), op_name, ": out must be contiguous");
46
+
47
+ const auto* src =
48
+ reinterpret_cast<const uint16_t*>(input.data_ptr<at::BFloat16>());
49
+ auto* dst = reinterpret_cast<uint16_t*>(out.data_ptr<at::BFloat16>());
50
+ const int64_t n = input.numel();
51
+
52
+ CPU_KERNEL_GUARD_IN(activation_lut_bf16_impl)
53
+#pragma omp parallel for
54
+ for (int64_t i = 0; i < n; ++i) {
55
+ dst[i] = lut[src[i]];
56
+ }
57
+ CPU_KERNEL_GUARD_OUT(activation_lut_bf16_impl)
58
+}
59
+
60
+void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
61
+ const std::string& activation) {
62
+ if (activation == "gelu") {
63
+ static std::array<uint16_t, ActivationLutSize> lut{};
64
+ static std::once_flag once;
65
+ maybe_init_activation_lut_bf16(lut.data(), once, gelu_reference);
66
+ activation_lut_bf16(out, input, lut.data(), "gelu_lut");
67
+ return;
68
+ }
69
+
70
+ TORCH_CHECK(false, "Unsupported activation: ", activation);
71
+}
csrc/cpu/torch_bindings.cpp
+12
@@ -85,6 +85,9 @@ at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros,
85
at::Tensor& w_scales,
86
std::optional<at::Tensor> bias);
87
88
+void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
89
+ const std::string& activation);
90
+
91
torch::Tensor get_scheduler_metadata(
92
const int64_t num_req, const int64_t num_heads_q,
93
const int64_t num_heads_kv, const int64_t head_dim,
@@ -231,6 +234,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
234
ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
235
ops.impl("gelu_quick", torch::kCPU, &gelu_quick);
236
237
+#if (defined(__aarch64__) && !defined(__APPLE__))
238
+
239
+ ops.def(
240
+ "activation_lut_bf16(Tensor! out, Tensor input, str activation)"
241
+ " -> ()");
242
+ ops.impl("activation_lut_bf16", torch::kCPU, &activation_lut_bf16);
243
+
244
+#endif // (defined(__aarch64__) && !defined(__APPLE__))
245
+
246
// Layernorm
247
// Apply Root Mean Square (RMS) Normalization to the input tensor.
248
ops.def(
csrc/ops.h
+1
@@ -1,6 +1,7 @@
1
#pragma once
2
3
#include <optional>
4
+#include <string>
5
#include <torch/library.h>
6
#include <tuple>
7
tests/kernels/core/test_cpu_activation.py
new
+111
@@ -0,0 +1,111 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from tests.kernels.allclose_default import get_default_atol, get_default_rtol
8
+from tests.kernels.utils import opcheck
9
+from vllm.platforms import CpuArchEnum, current_platform
10
+from vllm.utils.torch_utils import set_random_seed
11
+
12
+if not current_platform.is_cpu():
13
+ pytest.skip("skipping CPU-only tests", allow_module_level=True)
14
+
15
+from vllm.model_executor.layers.activation import (
16
+ GELU,
17
+ FastGELU,
18
+ GeluAndMul,
19
+ NewGELU,
20
+ QuickGELU,
21
+ SiluAndMul,
22
+)
23
+
24
+DTYPES = [torch.bfloat16, torch.float32]
25
+NUM_TOKENS = [7, 83]
26
+D = [512, 2048]
27
+SEEDS = [0]
28
+
29
+
30
+@pytest.mark.parametrize(
31
+ ("activation_cls", "fn"),
32
+ [
33
+ (SiluAndMul, torch.ops._C.silu_and_mul),
34
+ (GeluAndMul, torch.ops._C.gelu_and_mul),
35
+ (GeluAndMul, torch.ops._C.gelu_tanh_and_mul),
36
+ ],
37
+)
38
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
39
+@pytest.mark.parametrize("d", D)
40
+@pytest.mark.parametrize("dtype", DTYPES)
41
+@pytest.mark.parametrize("seed", SEEDS)
42
+@torch.inference_mode()
43
+def test_cpu_act_and_mul(
44
+ default_vllm_config,
45
+ activation_cls: type[torch.nn.Module],
46
+ fn: object,
47
+ num_tokens: int,
48
+ d: int,
49
+ dtype: torch.dtype,
50
+ seed: int,
51
+) -> None:
52
+ set_random_seed(seed)
53
+ x = torch.randn(num_tokens, 2 * d, dtype=dtype)
54
+
55
+ layer = activation_cls()
56
+ out = layer(x)
57
+ ref_out = layer.forward_native(x)
58
+
59
+ torch.testing.assert_close(
60
+ out, ref_out, atol=get_default_atol(out), rtol=get_default_rtol(out)
61
+ )
62
+
63
+ output_shape = x.shape[:-1] + (x.shape[-1] // 2,)
64
+ raw_out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
65
+ opcheck(fn, (raw_out, x))
66
+
67
+
68
+@pytest.mark.parametrize(
69
+ ("activation_cls", "fn", "op_args"),
70
+ [
71
+ (NewGELU, torch.ops._C.gelu_new, ()),
72
+ (FastGELU, torch.ops._C.gelu_fast, ()),
73
+ (QuickGELU, torch.ops._C.gelu_quick, ()),
74
+ pytest.param(
75
+ GELU,
76
+ getattr(torch.ops._C, "activation_lut_bf16", None),
77
+ ("gelu",),
78
+ marks=pytest.mark.skipif(
79
+ current_platform.get_cpu_architecture() != CpuArchEnum.ARM,
80
+ reason="activation_lut_bf16 is only built on Arm CPU",
81
+ ),
82
+ ),
83
+ ],
84
+)
85
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
86
+@pytest.mark.parametrize("d", D)
87
+@pytest.mark.parametrize("dtype", DTYPES)
88
+@pytest.mark.parametrize("seed", SEEDS)
89
+@torch.inference_mode()
90
+def test_cpu_unary_activation(
91
+ default_vllm_config,
92
+ activation_cls: type[torch.nn.Module],
93
+ fn: object,
94
+ op_args: tuple[str, ...],
95
+ num_tokens: int,
96
+ d: int,
97
+ dtype: torch.dtype,
98
+ seed: int,
99
+) -> None:
100
+ set_random_seed(seed)
101
+ x = torch.randn(num_tokens, d, dtype=dtype)
102
+ layer = activation_cls()
103
+ out = layer(x)
104
+ ref_out = layer.forward_native(x)
105
+ torch.testing.assert_close(
106
+ out, ref_out, atol=get_default_atol(out), rtol=get_default_rtol(out)
107
+ )
108
+ # gelu with activation_lut_bf16 only makes sense for BF16
109
+ if not (activation_cls is GELU and dtype != torch.bfloat16):
110
+ raw_out = torch.empty_like(x)
111
+ opcheck(fn, (raw_out, x, *op_args))
vllm/_custom_ops.py
+6
@@ -3300,6 +3300,12 @@ def cpu_gemm_wna16(
3300
return output
3301
3302
3303
+def cpu_activation_lut_bf16(input: torch.Tensor, activation: str) -> torch.Tensor:
3304
+ out = torch.empty_like(input)
3305
+ torch.ops._C.activation_lut_bf16(out, input, activation)
3306
+ return out
3307
+
3308
+
3309
def cpu_prepack_moe_weight(
3310
weight: torch.Tensor,
3311
isa: str,
vllm/model_executor/layers/activation.py
+30
-2
@@ -16,7 +16,7 @@ from vllm.distributed import (
16
from vllm.logger import init_logger
17
from vllm.model_executor.custom_op import CustomOp
18
from vllm.model_executor.utils import set_weight_attrs
19
-from vllm.platforms import current_platform
19
+from vllm.platforms import CpuArchEnum, current_platform
20
from vllm.triton_utils import tl, triton
21
from vllm.utils.collection_utils import LazyDict
22
@@ -247,6 +247,34 @@ class GeluAndMulSparse(CustomOp):
247
return self.forward_native(x)
248
249
250
+# --8<-- [start:gelu]
251
+@CustomOp.register("gelu")
252
+class GELU(CustomOp):
253
+ # --8<-- [end:gelu]
254
+
255
+ def __init__(self):
256
+ super().__init__()
257
+ if current_platform.get_cpu_architecture() == CpuArchEnum.ARM and hasattr(
258
+ torch.ops._C, "activation_lut_bf16"
259
+ ):
260
+ self.op = torch.ops._C.activation_lut_bf16
261
+ else:
262
+ self.op = None
263
+
264
+ def forward_native(self, x: torch.Tensor) -> torch.Tensor:
265
+ return F.gelu(x, approximate="none")
266
+
267
+ def forward_cpu(self, x: torch.Tensor) -> torch.Tensor:
268
+ if self.op and x.dtype == torch.bfloat16 and x.is_contiguous():
269
+ out = torch.empty_like(x)
270
+ self.op(out, x, "gelu")
271
+ return out
272
+ return self.forward_native(x)
273
+
274
+ def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
275
+ return self.forward_native(x)
276
+
277
+
278
# --8<-- [start:gelu_and_mul]
279
@CustomOp.register("gelu_and_mul")
280
class GeluAndMul(CustomOp):
@@ -635,7 +663,7 @@ class ScaledActivation(nn.Module):
663
664
_ACTIVATION_REGISTRY = LazyDict(
665
{
638
- "gelu": lambda: nn.GELU(),
666
+ "gelu": lambda: GELU(),
667
"gelu_fast": lambda: FastGELU(),
668
"gelu_new": lambda: NewGELU(),
669
"gelu_pytorch_tanh": lambda: (
vllm/platforms/cpu.py
+7
@@ -246,6 +246,13 @@ class CpuPlatform(Platform):
246
if vllm_config.lora_config is not None:
247
compilation_config.mode = CompilationMode.NONE
248
249
+ if (
250
+ cls.get_cpu_architecture() == CpuArchEnum.ARM
251
+ and "+gelu" not in compilation_config.custom_ops
252
+ and "-gelu" not in compilation_config.custom_ops
253
+ ):
254
+ compilation_config.custom_ops.append("+gelu")
255
+
256
vllm_config.profiler_config.torch_profiler_dump_cuda_time_total = False
257
258
assert vllm_config.device_config.device_type == "cpu"