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 14, 05:30 UTC
[cpu] Dispatch un-quantized linear to oneDNN/ACL by default for AArch64 (#27183)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com> Co-authored-by: Michael Yang <Michael.Yang@arm.com>
Fadi Arafeh committed
Oct 21, 2025 at 03:02 UTC
163965d183a980f28909003fac32100d53e1d2e9
3 files changed
+42
-23
cmake/cpu_extension.cmake
+39
-18
@@ -188,16 +188,47 @@ else()
188
message(FATAL_ERROR "vLLM CPU backend requires AVX512, AVX2, Power9+ ISA, S390X ISA, ARMv8 or RISC-V support.")
189
endif()
190
191
-#
192
-# Build oneDNN for W8A8 GEMM kernels (only for x86-AVX512 /ARM platforms)
193
-# Flag to enable ACL kernels for AARCH64 platforms
194
-if (VLLM_BUILD_ACL STREQUAL "ON")
195
- set(USE_ACL ON)
196
-else()
197
- set(USE_ACL OFF)
198
-endif()
191
192
+# Build oneDNN for GEMM kernels (only for x86-AVX512 /ARM platforms)
193
if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
194
+ # Fetch and build Arm Compute Library (ACL) as oneDNN's backend for AArch64
195
+ # TODO [fadara01]: remove this once ACL can be fetched and built automatically as a dependency of oneDNN
196
+ if(ASIMD_FOUND)
197
+ if(DEFINED ENV{ACL_ROOT_DIR} AND IS_DIRECTORY "$ENV{ACL_ROOT_DIR}")
198
+ message(STATUS "Using ACL from specified source directory: $ENV{ACL_ROOT_DIR}")
199
+ else()
200
+ message(STATUS "Downloading Arm Compute Library (ACL) from GitHub")
201
+ FetchContent_Populate(arm_compute
202
+ SUBBUILD_DIR "${FETCHCONTENT_BASE_DIR}/arm_compute-subbuild"
203
+ SOURCE_DIR "${FETCHCONTENT_BASE_DIR}/arm_compute-src"
204
+ GIT_REPOSITORY https://github.com/ARM-software/ComputeLibrary.git
205
+ GIT_TAG v52.2.0
206
+ GIT_SHALLOW TRUE
207
+ GIT_PROGRESS TRUE
208
+ )
209
+ set(ENV{ACL_ROOT_DIR} "${arm_compute_SOURCE_DIR}")
210
+ endif()
211
+
212
+ # Build ACL with scons
213
+ include(ProcessorCount)
214
+ ProcessorCount(_NPROC)
215
+ execute_process(
216
+ COMMAND scons -j${_NPROC}
217
+ Werror=0 debug=0 neon=1 examples=0 embed_kernels=0 os=linux
218
+ arch=armv8.2-a build=native benchmark_examples=0 fixed_format_kernels=1
219
+ multi_isa=1 openmp=1 cppthreads=0
220
+ WORKING_DIRECTORY "$ENV{ACL_ROOT_DIR}"
221
+ RESULT_VARIABLE _acl_rc
222
+ )
223
+ if(NOT _acl_rc EQUAL 0)
224
+ message(FATAL_ERROR "ACL SCons build failed (exit ${_acl_rc}).")
225
+ endif()
226
+
227
+ set(ONEDNN_AARCH64_USE_ACL "ON")
228
+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wl,-rpath,$ENV{ACL_ROOT_DIR}/build/")
229
+ add_compile_definitions(VLLM_USE_ACL)
230
+ endif()
231
+
232
set(FETCHCONTENT_SOURCE_DIR_ONEDNN "$ENV{FETCHCONTENT_SOURCE_DIR_ONEDNN}" CACHE PATH "Path to a local oneDNN source directory.")
233
234
if(FETCHCONTENT_SOURCE_DIR_ONEDNN)
@@ -217,16 +248,6 @@ if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR (ASIMD_FOUND AND NOT APPLE_SILICON
248
)
249
endif()
250
220
- if(USE_ACL)
221
- find_library(ARM_COMPUTE_LIBRARY NAMES arm_compute PATHS $ENV{ACL_ROOT_DIR}/build/)
222
- if(NOT ARM_COMPUTE_LIBRARY)
223
- message(FATAL_ERROR "Could not find ARM Compute Library: please set ACL_ROOT_DIR")
224
- endif()
225
- set(ONEDNN_AARCH64_USE_ACL "ON")
226
- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wl,-rpath,$ENV{ACL_ROOT_DIR}/build/")
227
- add_compile_definitions(VLLM_USE_ACL)
228
- endif()
229
-
251
set(ONEDNN_LIBRARY_TYPE "STATIC")
252
set(ONEDNN_BUILD_DOC "OFF")
253
set(ONEDNN_BUILD_EXAMPLES "OFF")
requirements/cpu-build.txt
+1
@@ -6,6 +6,7 @@ setuptools-scm>=8
6
--extra-index-url https://download.pytorch.org/whl/cpu
7
torch==2.8.0+cpu; platform_machine == "x86_64"
8
torch==2.8.0; platform_machine == "ppc64le" or platform_machine == "aarch64" or platform_system == "Darwin"
9
+scons; platform_machine == "aarch64" # needed to build Arm Compute Library (ACL)
10
wheel
11
jinja2>=3.1.6
12
regex
vllm/model_executor/layers/utils.py
+2
-5
@@ -8,7 +8,7 @@ import torch
8
9
from vllm import _custom_ops as ops
10
from vllm import envs
11
-from vllm.platforms import CpuArchEnum, current_platform
11
+from vllm.platforms import current_platform
12
from vllm.utils.torch_utils import direct_register_custom_op
13
14
@@ -178,10 +178,7 @@ def dispatch_cpu_unquantized_gemm(
178
)
179
if remove_weight:
180
layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
181
- elif ops._supports_onednn and (
182
- current_platform.get_cpu_architecture() == CpuArchEnum.X86
183
- or ops.is_onednn_acl_supported()
184
- ):
181
+ elif ops._supports_onednn:
182
origin_weight = layer.weight
183
if remove_weight:
184
layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)