@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 90805ff464
Fork 0

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 7, 14:00 UTC

[CI/Build] CPU release supports both of AVX2 and AVX512 (#35466)

Signed-off-by: jiang1.li <jiang1.li@intel.com> Co-authored-by: jiang1.li <jiang1.li@intel.com>

Ma Jian committed Feb 28, 2026 at 12:35 UTC 90805ff46434eb81f3c609498146e2b1a2477c2d
6 files changed +161 -147
cmake/cpu_extension.cmake
+118 -132
@@ -13,28 +13,16 @@ endif()
13 #
14 # Define environment variables for special configurations
15 #
16 -set(ENABLE_AVX2 $ENV{VLLM_CPU_AVX2})
17 -set(ENABLE_AVX512 $ENV{VLLM_CPU_AVX512})
18 -set(ENABLE_AVX512BF16 $ENV{VLLM_CPU_AVX512BF16})
19 -set(ENABLE_AVX512VNNI $ENV{VLLM_CPU_AVX512VNNI})
20 -set(ENABLE_AMXBF16 $ENV{VLLM_CPU_AMXBF16})
16 +set(ENABLE_X86_ISA $ENV{VLLM_CPU_X86})
17 set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16})
18
19 include_directories("${CMAKE_SOURCE_DIR}/csrc")
20
25 -
21 set (ENABLE_NUMA TRUE)
22
23 #
24 # Check the compile flags
25 #
31 -
32 -if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64")
33 - list(APPEND CXX_COMPILE_FLAGS
34 - "-mf16c"
35 - )
36 -endif()
37 -
26 if(MACOSX_FOUND)
27 list(APPEND CXX_COMPILE_FLAGS
28 "-DVLLM_CPU_EXTENSION")
@@ -78,18 +66,6 @@ function(check_sysctl TARGET OUT)
66 endif()
67 endfunction()
68
81 -
82 -function (is_avx512_disabled OUT)
83 - set(DISABLE_AVX512 $ENV{VLLM_CPU_DISABLE_AVX512})
84 - if(DISABLE_AVX512 AND DISABLE_AVX512 STREQUAL "true")
85 - set(${OUT} ON PARENT_SCOPE)
86 - else()
87 - set(${OUT} OFF PARENT_SCOPE)
88 - endif()
89 -endfunction()
90 -
91 -is_avx512_disabled(AVX512_DISABLED)
92 -
69 if (MACOSX_FOUND AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
70 message(STATUS "Apple Silicon Detected")
71 set(APPLE_SILICON_FOUND TRUE)
@@ -97,8 +73,6 @@ if (MACOSX_FOUND AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
73 check_sysctl(hw.optional.neon ASIMD_FOUND)
74 check_sysctl(hw.optional.arm.FEAT_BF16 ARM_BF16_FOUND)
75 else()
100 - find_isa(${CPUINFO} "avx2" AVX2_FOUND)
101 - find_isa(${CPUINFO} "avx512f" AVX512_FOUND)
76 find_isa(${CPUINFO} "Power11" POWER11_FOUND)
77 find_isa(${CPUINFO} "POWER10" POWER10_FOUND)
78 find_isa(${CPUINFO} "POWER9" POWER9_FOUND)
@@ -108,77 +82,32 @@ else()
82 find_isa(${CPUINFO} "v" RVV_FOUND) # Check for RISC-V RVV support
83
84 # Support cross-compilation by allowing override via environment variables
111 - if (ENABLE_AVX2)
112 - set(AVX2_FOUND ON)
113 - message(STATUS "AVX2 support enabled via VLLM_CPU_AVX2 environment variable")
114 - endif()
115 - if (ENABLE_AVX512)
116 - set(AVX512_FOUND ON)
117 - message(STATUS "AVX512 support enabled via VLLM_CPU_AVX512 environment variable")
118 - endif()
85 if (ENABLE_ARM_BF16)
86 set(ARM_BF16_FOUND ON)
87 message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable")
88 endif()
89 endif()
90
125 -if (AVX512_FOUND AND NOT AVX512_DISABLED)
126 - list(APPEND CXX_COMPILE_FLAGS
91 +if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA)
92 + set(ENABLE_X86_ISA ON)
93 + if (NOT (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
94 + CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3))
95 + message(FATAL_ERROR "X86 backend requires gcc/g++ >= 12.3")
96 + endif()
97 + list(APPEND CXX_COMPILE_FLAGS "-mf16c")
98 + list(APPEND CXX_COMPILE_FLAGS_AVX512 ${CXX_COMPILE_FLAGS})
99 + list(APPEND CXX_COMPILE_FLAGS_AVX2 ${CXX_COMPILE_FLAGS})
100 + list(APPEND CXX_COMPILE_FLAGS_AVX512
101 "-mavx512f"
102 "-mavx512vl"
103 "-mavx512bw"
130 - "-mavx512dq")
131 -
132 - find_isa(${CPUINFO} "avx512_bf16" AVX512BF16_FOUND)
133 - if (AVX512BF16_FOUND OR ENABLE_AVX512BF16)
134 - if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
135 - CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3)
136 - list(APPEND CXX_COMPILE_FLAGS "-mavx512bf16")
137 - set(ENABLE_AVX512BF16 ON)
138 - else()
139 - set(ENABLE_AVX512BF16 OFF)
140 - message(WARNING "Disable AVX512-BF16 ISA support, requires gcc/g++ >= 12.3")
141 - endif()
142 - else()
143 - set(ENABLE_AVX512BF16 OFF)
144 - message(WARNING "Disable AVX512-BF16 ISA support, no avx512_bf16 found in local CPU flags." " If cross-compilation is required, please set env VLLM_CPU_AVX512BF16=1.")
145 - endif()
146 -
147 - find_isa(${CPUINFO} "avx512_vnni" AVX512VNNI_FOUND)
148 - if (AVX512VNNI_FOUND OR ENABLE_AVX512VNNI)
149 - if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
150 - CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3)
151 - list(APPEND CXX_COMPILE_FLAGS "-mavx512vnni")
152 - set(ENABLE_AVX512VNNI ON)
153 - else()
154 - set(ENABLE_AVX512VNNI OFF)
155 - message(WARNING "Disable AVX512-VNNI ISA support, requires gcc/g++ >= 12.3")
156 - endif()
157 - else()
158 - set(ENABLE_AVX512VNNI OFF)
159 - message(WARNING "Disable AVX512-VNNI ISA support, no avx512_vnni found in local CPU flags." " If cross-compilation is required, please set env VLLM_CPU_AVX512VNNI=1.")
160 - endif()
161 -
162 - find_isa(${CPUINFO} "amx_bf16" AMXBF16_FOUND)
163 - if (AMXBF16_FOUND OR ENABLE_AMXBF16)
164 - if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
165 - CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3)
166 - list(APPEND CXX_COMPILE_FLAGS "-mamx-bf16" "-mamx-tile")
167 - set(ENABLE_AMXBF16 ON)
168 - add_compile_definitions(-DCPU_CAPABILITY_AMXBF16)
169 - else()
170 - set(ENABLE_AMXBF16 OFF)
171 - message(WARNING "Disable AMX_BF16 ISA support, requires gcc/g++ >= 12.3")
172 - endif()
173 - else()
174 - set(ENABLE_AMXBF16 OFF)
175 - message(WARNING "Disable AMX_BF16 ISA support, no amx_bf16 found in local CPU flags." " If cross-compilation is required, please set env VLLM_CPU_AMXBF16=1.")
176 - endif()
177 -
178 -elseif (AVX2_FOUND)
179 - list(APPEND CXX_COMPILE_FLAGS "-mavx2")
180 - message(WARNING "vLLM CPU backend using AVX2 ISA")
181 -
104 + "-mavx512dq"
105 + "-mavx512bf16"
106 + "-mavx512vnni"
107 + "-mamx-bf16"
108 + "-mamx-tile")
109 + list(APPEND CXX_COMPILE_FLAGS_AVX2
110 + "-mavx2")
111 elseif (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
112 message(STATUS "PowerPC detected")
113 if (POWER9_FOUND)
@@ -219,12 +148,12 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
148 list(APPEND CXX_COMPILE_FLAGS "-march=rv64gc")
149 endif()
150 else()
222 - message(FATAL_ERROR "vLLM CPU backend requires AVX512, AVX2, Power9+ ISA, S390X ISA, ARMv8 or RISC-V support.")
151 + message(FATAL_ERROR "vLLM CPU backend requires X86, Power9+ ISA, S390X ISA, ARMv8 or RISC-V support.")
152 endif()
153
154
226 -# Build oneDNN for GEMM kernels (only for x86-AVX512 /ARM platforms)
227 -if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
155 +# Build oneDNN for GEMM kernels
156 +if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
157 # Fetch and build Arm Compute Library (ACL) as oneDNN's backend for AArch64
158 # TODO [fadara01]: remove this once ACL can be fetched and built automatically as a dependency of oneDNN
159 set(ONEDNN_AARCH64_USE_ACL OFF CACHE BOOL "")
@@ -329,13 +258,21 @@ if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR (ASIMD_FOUND AND NOT APPLE_SILICON
258 set(ONEDNN_ENABLE_WORKLOAD "INFERENCE")
259 set(ONEDNN_ENABLE_PRIMITIVE "MATMUL;REORDER")
260 set(ONEDNN_BUILD_GRAPH "OFF")
332 - set(ONEDNN_ENABLE_JIT_PROFILING "OFF")
261 + set(ONEDNN_ENABLE_JIT_PROFILING "ON")
262 set(ONEDNN_ENABLE_ITT_TASKS "OFF")
334 - set(ONEDNN_ENABLE_MAX_CPU_ISA "OFF")
335 - set(ONEDNN_ENABLE_CPU_ISA_HINTS "OFF")
336 - set(ONEDNN_VERBOSE "OFF")
263 + set(ONEDNN_ENABLE_MAX_CPU_ISA "ON")
264 + set(ONEDNN_ENABLE_CPU_ISA_HINTS "ON")
265 + set(ONEDNN_VERBOSE "ON")
266 set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
267
268 + # TODO: Refactor this
269 + if (ENABLE_X86_ISA)
270 + # Note: only enable oneDNN for AVX512
271 + list(APPEND DNNL_COMPILE_FLAGS ${CXX_COMPILE_FLAGS_AVX512})
272 + else()
273 + list(APPEND DNNL_COMPILE_FLAGS ${CXX_COMPILE_FLAGS})
274 + endif()
275 +
276 set(VLLM_BUILD_TYPE ${CMAKE_BUILD_TYPE})
277 set(CMAKE_BUILD_TYPE "Release") # remove oneDNN debug symbols to reduce size
278 FetchContent_MakeAvailable(oneDNN)
@@ -348,14 +285,20 @@ if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR (ASIMD_FOUND AND NOT APPLE_SILICON
285 PRIVATE ${oneDNN_SOURCE_DIR}/src
286 )
287 target_link_libraries(dnnl_ext dnnl torch)
351 - target_compile_options(dnnl_ext PRIVATE ${CXX_COMPILE_FLAGS} -fPIC)
288 + target_compile_options(dnnl_ext PRIVATE ${DNNL_COMPILE_FLAGS} -fPIC)
289 list(APPEND LIBS dnnl_ext)
290 set(USE_ONEDNN ON)
291 else()
292 set(USE_ONEDNN OFF)
293 endif()
294
358 -message(STATUS "CPU extension compile flags: ${CXX_COMPILE_FLAGS}")
295 +# TODO: Refactor this
296 +if (ENABLE_X86_ISA)
297 + message(STATUS "CPU extension (AVX512) compile flags: ${CXX_COMPILE_FLAGS_AVX512}")
298 + message(STATUS "CPU extension (AVX2) compile flags: ${CXX_COMPILE_FLAGS_AVX2}")
299 +else()
300 + message(STATUS "CPU extension compile flags: ${CXX_COMPILE_FLAGS}")
301 +endif()
302
303 if(ENABLE_NUMA)
304 list(APPEND LIBS numa)
@@ -390,25 +333,6 @@ set(VLLM_EXT_SRC
333 "csrc/cpu/cpu_attn.cpp"
334 "csrc/cpu/torch_bindings.cpp")
335
393 -if (AVX512_FOUND AND NOT AVX512_DISABLED)
394 - set(VLLM_EXT_SRC
395 - "csrc/cpu/shm.cpp"
396 - "csrc/cpu/cpu_wna16.cpp"
397 - "csrc/cpu/cpu_fused_moe.cpp"
398 - ${VLLM_EXT_SRC})
399 - if (ENABLE_AVX512BF16 AND ENABLE_AVX512VNNI)
400 - set(VLLM_EXT_SRC
401 - "csrc/cpu/sgl-kernels/gemm.cpp"
402 - "csrc/cpu/sgl-kernels/gemm_int8.cpp"
403 - "csrc/cpu/sgl-kernels/gemm_fp8.cpp"
404 - "csrc/cpu/sgl-kernels/moe.cpp"
405 - "csrc/cpu/sgl-kernels/moe_int8.cpp"
406 - "csrc/cpu/sgl-kernels/moe_fp8.cpp"
407 - ${VLLM_EXT_SRC})
408 - add_compile_definitions(-DCPU_CAPABILITY_AVX512)
409 - endif()
410 -endif()
411 -
336 if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
337 set(VLLM_EXT_SRC
338 "csrc/cpu/shm.cpp"
@@ -421,21 +345,83 @@ if(USE_ONEDNN)
345 ${VLLM_EXT_SRC})
346 endif()
347
424 -message(STATUS "CPU extension source files: ${VLLM_EXT_SRC}")
425 -
426 -#
427 -# Define extension targets
428 -#
348 +if (ENABLE_X86_ISA)
349 + set(VLLM_EXT_SRC_AVX512
350 + "csrc/cpu/sgl-kernels/gemm.cpp"
351 + "csrc/cpu/sgl-kernels/gemm_int8.cpp"
352 + "csrc/cpu/sgl-kernels/gemm_fp8.cpp"
353 + "csrc/cpu/sgl-kernels/moe.cpp"
354 + "csrc/cpu/sgl-kernels/moe_int8.cpp"
355 + "csrc/cpu/sgl-kernels/moe_fp8.cpp"
356 + "csrc/cpu/shm.cpp"
357 + "csrc/cpu/cpu_wna16.cpp"
358 + "csrc/cpu/cpu_fused_moe.cpp"
359 + "csrc/cpu/utils.cpp"
360 + "csrc/cpu/cpu_attn.cpp"
361 + "csrc/cpu/dnnl_kernels.cpp"
362 + "csrc/cpu/torch_bindings.cpp"
363 + # TODO: Remove these files
364 + "csrc/cpu/activation.cpp"
365 + "csrc/cpu/layernorm.cpp"
366 + "csrc/cpu/mla_decode.cpp"
367 + "csrc/cpu/pos_encoding.cpp"
368 + "csrc/moe/dynamic_4bit_int_moe_cpu.cpp")
369 +
370 + set(VLLM_EXT_SRC_AVX2
371 + "csrc/cpu/utils.cpp"
372 + "csrc/cpu/cpu_attn.cpp"
373 + "csrc/cpu/torch_bindings.cpp"
374 + # TODO: Remove these files
375 + "csrc/cpu/activation.cpp"
376 + "csrc/cpu/layernorm.cpp"
377 + "csrc/cpu/mla_decode.cpp"
378 + "csrc/cpu/pos_encoding.cpp"
379 + "csrc/moe/dynamic_4bit_int_moe_cpu.cpp")
380 +
381 + message(STATUS "CPU extension (AVX512) source files: ${VLLM_EXT_SRC_AVX512}")
382 + message(STATUS "CPU extension (AVX2) source files: ${VLLM_EXT_SRC_AVX2}")
383 +
384 + define_extension_target(
385 + _C
386 + DESTINATION vllm
387 + LANGUAGE CXX
388 + SOURCES ${VLLM_EXT_SRC_AVX512}
389 + LIBRARIES ${LIBS}
390 + COMPILE_FLAGS ${CXX_COMPILE_FLAGS_AVX512}
391 + USE_SABI 3
392 + WITH_SOABI
393 + )
394
430 -define_extension_target(
431 - _C
432 - DESTINATION vllm
433 - LANGUAGE CXX
434 - SOURCES ${VLLM_EXT_SRC}
435 - LIBRARIES ${LIBS}
436 - COMPILE_FLAGS ${CXX_COMPILE_FLAGS}
437 - USE_SABI 3
438 - WITH_SOABI
439 -)
395 + # For SGL kernels
396 + target_compile_definitions(_C PRIVATE "-DCPU_CAPABILITY_AVX512")
397 + # For AMX kernels
398 + target_compile_definitions(_C PRIVATE "-DCPU_CAPABILITY_AMXBF16")
399 +
400 + define_extension_target(
401 + _C_AVX2
402 + DESTINATION vllm
403 + LANGUAGE CXX
404 + SOURCES ${VLLM_EXT_SRC_AVX2}
405 + LIBRARIES ${LIBS}
406 + COMPILE_FLAGS ${CXX_COMPILE_FLAGS_AVX2}
407 + USE_SABI 3
408 + WITH_SOABI
409 + )
410 +else()
411 + message(STATUS "CPU extension source files: ${VLLM_EXT_SRC}")
412 + #
413 + # Define extension targets
414 + #
415 + define_extension_target(
416 + _C
417 + DESTINATION vllm
418 + LANGUAGE CXX
419 + SOURCES ${VLLM_EXT_SRC}
420 + LIBRARIES ${LIBS}
421 + COMPILE_FLAGS ${CXX_COMPILE_FLAGS}
422 + USE_SABI 3
423 + WITH_SOABI
424 + )
425 +endif()
426
427 message(STATUS "Enabling C extension.")
csrc/cpu/torch_bindings.cpp
+7 -10
@@ -4,6 +4,10 @@
4
5 #include <torch/library.h>
6
7 +// Note: overwrite the external defination for sharing same name between
8 +// libraries use different ISAs.
9 +#define TORCH_EXTENSION_NAME _C
10 +
11 std::string init_cpu_threads_env(const std::string& cpu_ids);
12
13 void release_dnnl_matmul_handler(int64_t handler);
@@ -324,19 +328,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
328 "str act, str isa) -> ()");
329 ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
330 #endif
327 -}
328 -
329 -TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _utils), utils) {
330 - // CPU utils
331 - utils.def("init_cpu_threads_env(str cpu_ids) -> str", &init_cpu_threads_env);
332 -}
333 -
334 -TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cpu), cpu_ops) {
335 - cpu_ops.def(
331 + ops.def("init_cpu_threads_env(str cpu_ids) -> str", &init_cpu_threads_env);
332 + ops.def(
333 "mla_decode_kvcache("
334 " Tensor! out, Tensor query, Tensor kv_cache,"
335 " float scale, Tensor block_tables, Tensor seq_lens) -> ()");
339 - cpu_ops.impl("mla_decode_kvcache", torch::kCPU, &mla_decode_kvcache);
336 + ops.impl("mla_decode_kvcache", torch::kCPU, &mla_decode_kvcache);
337 }
338
339 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
setup.py
+10 -1
@@ -818,7 +818,7 @@ def _is_xpu() -> bool:
818
819
820 def _build_custom_ops() -> bool:
821 - return _is_cuda() or _is_hip() or _is_cpu()
821 + return _is_cuda() or _is_hip()
822
823
824 def get_rocm_version():
@@ -987,6 +987,15 @@ if _is_cuda():
987 CMakeExtension(name="vllm._flashmla_extension_C", optional=True)
988 )
989
990 +if _is_cpu():
991 + import platform
992 +
993 + if platform.machine() in ("x86_64", "AMD64"):
994 + ext_modules.append(CMakeExtension(name="vllm._C"))
995 + ext_modules.append(CMakeExtension(name="vllm._C_AVX2"))
996 + else:
997 + ext_modules.append(CMakeExtension(name="vllm._C"))
998 +
999 if _build_custom_ops():
1000 ext_modules.append(CMakeExtension(name="vllm._C"))
1001
vllm/_custom_ops.py
+1 -3
@@ -178,9 +178,7 @@ def mla_decode_kvcache_cpu(
178 block_tables: torch.Tensor,
179 seq_lens: torch.Tensor,
180 ) -> None:
181 - torch.ops._C_cpu.mla_decode_kvcache(
182 - out, query, kv_cache, scale, block_tables, seq_lens
183 - )
181 + torch.ops._C.mla_decode_kvcache(out, query, kv_cache, scale, block_tables, seq_lens)
182
183
184 # merge attn states ops
vllm/platforms/cpu.py
+24
@@ -483,3 +483,27 @@ class CpuPlatform(Platform):
483 @classmethod
484 def support_hybrid_kv_cache(cls) -> bool:
485 return True
486 +
487 + @classmethod
488 + def import_kernels(cls) -> None:
489 + if Platform.get_cpu_architecture() in (CpuArchEnum.X86,):
490 + if torch._C._cpu._is_avx512_supported():
491 + try:
492 + import vllm._C # noqa: F401
493 + except ImportError as e:
494 + logger.warning("Failed to import from vllm._C: %r", e)
495 + else:
496 + # Note: The lib name is _C_AVX2, but the module name is _C.
497 + # This will cause a exception "dynamic module does define
498 + # module export function". But the library is imported
499 + # successfully. So ignore the exception for now, until we find
500 + # a solution.
501 + try:
502 + import vllm._C_AVX2 # noqa: F401
503 + except ImportError as e:
504 + logger.warning("Failed to import from vllm._C_AVX2: %r", e)
505 + else:
506 + try:
507 + import vllm._C # noqa: F401
508 + except ImportError as e:
509 + logger.warning("Failed to import from vllm._C: %r", e)
vllm/v1/worker/cpu_worker.py
+1 -1
@@ -85,7 +85,7 @@ class CPUWorker(Worker):
85 self.local_omp_cpuid = omp_cpuids_list[self.rank]
86
87 if self.local_omp_cpuid != "nobind":
88 - ret = torch.ops._C_utils.init_cpu_threads_env(self.local_omp_cpuid)
88 + ret = torch.ops._C.init_cpu_threads_env(self.local_omp_cpuid)
89 if ret:
90 logger.info(ret)
91