cpu_extension.cmake 8.93 KB
Newer Older
1
2
3
include(FetchContent)

set(CMAKE_CXX_STANDARD_REQUIRED ON)
4
set(CMAKE_CXX_STANDARD 17)
5
set(CMAKE_CXX_EXTENSIONS ON)
6
7
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)

8
9
10
11
12
if (${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
    set(MACOSX_FOUND TRUE)
endif()


13
14
15
#
# Define environment variables for special configurations
#
16
17
set(ENABLE_AVX512BF16 $ENV{VLLM_CPU_AVX512BF16})
set(ENABLE_AVX512VNNI $ENV{VLLM_CPU_AVX512VNNI})
18
19
20

include_directories("${CMAKE_SOURCE_DIR}/csrc")

21
22
23

set (ENABLE_NUMA TRUE)

24
25
26
#
# Check the compile flags
#
27
28

if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64")
29
30
    list(APPEND CXX_COMPILE_FLAGS
        "-mf16c"
31
    )
32
endif()
33

34
35
36
37
38
39
40
41
if(MACOSX_FOUND)
    list(APPEND CXX_COMPILE_FLAGS
        "-DVLLM_CPU_EXTENSION")
else()
    list(APPEND CXX_COMPILE_FLAGS
        "-fopenmp"
        "-DVLLM_CPU_EXTENSION")
endif()
42

43
44
45
46
47
48
49
if (NOT MACOSX_FOUND)
    execute_process(COMMAND cat /proc/cpuinfo
                    RESULT_VARIABLE CPUINFO_RET
                    OUTPUT_VARIABLE CPUINFO)
    if (NOT CPUINFO_RET EQUAL 0)
        message(FATAL_ERROR "Failed to check CPU features via /proc/cpuinfo")
    endif()
50
51
endif()

52

53
54
55
56
57
58
59
60
61
function (find_isa CPUINFO TARGET OUT)
    string(FIND ${CPUINFO} ${TARGET} ISA_FOUND)
    if(NOT ISA_FOUND EQUAL -1)
        set(${OUT} ON PARENT_SCOPE)
    else()
        set(${OUT} OFF PARENT_SCOPE)
    endif()
endfunction()

62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

function(check_sysctl TARGET OUT)
    execute_process(COMMAND sysctl -n "${TARGET}"
                    RESULT_VARIABLE SYSCTL_RET
                    OUTPUT_VARIABLE SYSCTL_INFO
                    ERROR_QUIET
                    OUTPUT_STRIP_TRAILING_WHITESPACE)
    if(SYSCTL_RET EQUAL 0 AND
      (SYSCTL_INFO STREQUAL "1" OR SYSCTL_INFO GREATER 0))
        set(${OUT} ON PARENT_SCOPE)
    else()
        set(${OUT} OFF PARENT_SCOPE)
    endif()
endfunction()


78
79
80
81
82
83
84
85
86
87
88
function (is_avx512_disabled OUT)
    set(DISABLE_AVX512 $ENV{VLLM_CPU_DISABLE_AVX512})
    if(DISABLE_AVX512 AND DISABLE_AVX512 STREQUAL "true")
        set(${OUT} ON PARENT_SCOPE)
    else()
        set(${OUT} OFF PARENT_SCOPE)
    endif()
endfunction()

is_avx512_disabled(AVX512_DISABLED)

89
if (MACOSX_FOUND AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
90
    message(STATUS "Apple Silicon Detected")
91
    set(APPLE_SILICON_FOUND TRUE)
92
93
94
    set(ENABLE_NUMA OFF)
    check_sysctl(hw.optional.neon ASIMD_FOUND)
    check_sysctl(hw.optional.arm.FEAT_BF16 ARM_BF16_FOUND)
95
96
97
else()
    find_isa(${CPUINFO} "avx2" AVX2_FOUND)
    find_isa(${CPUINFO} "avx512f" AVX512_FOUND)
98
    find_isa(${CPUINFO} "Power11" POWER11_FOUND)
99
100
101
102
    find_isa(${CPUINFO} "POWER10" POWER10_FOUND)
    find_isa(${CPUINFO} "POWER9" POWER9_FOUND)
    find_isa(${CPUINFO} "asimd" ASIMD_FOUND) # Check for ARM NEON support
    find_isa(${CPUINFO} "bf16" ARM_BF16_FOUND) # Check for ARM BF16 support
103
    find_isa(${CPUINFO} "S390" S390_FOUND)
104
105
endif()

106
if (AVX512_FOUND AND NOT AVX512_DISABLED)
107
108
109
110
111
112
113
114
    list(APPEND CXX_COMPILE_FLAGS
        "-mavx512f"
        "-mavx512vl"
        "-mavx512bw"
        "-mavx512dq")

    find_isa(${CPUINFO} "avx512_bf16" AVX512BF16_FOUND)
    if (AVX512BF16_FOUND OR ENABLE_AVX512BF16)
115
116
        if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
            CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3)
117
            list(APPEND CXX_COMPILE_FLAGS "-mavx512bf16")
118
            set(ENABLE_AVX512BF16 ON)
119
        else()
120
            set(ENABLE_AVX512BF16 OFF)
121
122
123
            message(WARNING "Disable AVX512-BF16 ISA support, requires gcc/g++ >= 12.3")
        endif()
    else()
124
        set(ENABLE_AVX512BF16 OFF)
125
126
        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.")
    endif()
127
128

    find_isa(${CPUINFO} "avx512_vnni" AVX512VNNI_FOUND)
129
130
131
132
133
134
135
136
137
138
139
140
141
    if (AVX512VNNI_FOUND OR ENABLE_AVX512VNNI)
        if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
            CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3)
            list(APPEND CXX_COMPILE_FLAGS "-mavx512vnni")
            set(ENABLE_AVX512VNNI ON)
        else()
            set(ENABLE_AVX512VNNI OFF)
            message(WARNING "Disable AVX512-VNNI ISA support, requires gcc/g++ >= 12.3")
        endif()
    else()
        set(ENABLE_AVX512VNNI OFF)
        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.")
    endif()
142
    
143
144
145
elseif (AVX2_FOUND)
    list(APPEND CXX_COMPILE_FLAGS "-mavx2")
    message(WARNING "vLLM CPU backend using AVX2 ISA")
146
    
147
elseif (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
148
    message(STATUS "PowerPC detected")
149
150
151
152
153
154
155
156
157
158
159
    if (POWER9_FOUND)
        list(APPEND CXX_COMPILE_FLAGS
            "-mvsx"
            "-mcpu=power9"
            "-mtune=power9")
    elseif (POWER10_FOUND OR POWER11_FOUND)
        list(APPEND CXX_COMPILE_FLAGS
            "-mvsx"
            "-mcpu=power10"
            "-mtune=power10")
    endif()
160
161
162
163
164
165
166
167
168
169
170
171

elseif (ASIMD_FOUND)
    message(STATUS "ARMv8 or later architecture detected")
    if(ARM_BF16_FOUND)
        message(STATUS "BF16 extension detected")
        set(MARCH_FLAGS "-march=armv8.2-a+bf16+dotprod+fp16")
        add_compile_definitions(ARM_BF16_SUPPORT)
    else()
        message(WARNING "BF16 functionality is not available")
        set(MARCH_FLAGS "-march=armv8.2-a+dotprod+fp16")  
    endif()
    list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS})     
172
173
174
175
176
177
178
179
elseif (S390_FOUND)
    message(STATUS "S390 detected")
    # Check for S390 VXE support
    list(APPEND CXX_COMPILE_FLAGS
        "-mvx"
        "-mzvector"
        "-march=native"
        "-mtune=native")
180
else()
181
    message(FATAL_ERROR "vLLM CPU backend requires AVX512, AVX2, Power9+ ISA, S390X ISA or ARMv8 support.")
182
183
endif()

184
#
185
186
# Build oneDNN for W8A8 GEMM kernels (only for x86-AVX512 /ARM platforms)
# Flag to enable ACL kernels for AARCH64 platforms
187
if (VLLM_BUILD_ACL STREQUAL "ON")
188
189
190
191
192
    set(USE_ACL ON)
else()
    set(USE_ACL OFF)
endif()

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
195
196
    FetchContent_Declare(
        oneDNN
        GIT_REPOSITORY https://github.com/oneapi-src/oneDNN.git
197
        GIT_TAG v3.9
198
199
200
201
        GIT_PROGRESS TRUE
        GIT_SHALLOW TRUE
    )

202
203
204
205
206
207
208
    if(USE_ACL)
        find_library(ARM_COMPUTE_LIBRARY NAMES arm_compute PATHS $ENV{ACL_ROOT_DIR}/build/)
        if(NOT ARM_COMPUTE_LIBRARY)
            message(FATAL_ERROR "Could not find ARM Compute Library: please set ACL_ROOT_DIR")
        endif()
        set(ONEDNN_AARCH64_USE_ACL "ON")
        set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wl,-rpath,$ENV{ACL_ROOT_DIR}/build/")
209
    endif()
210

211
212
213
214
215
216
217
218
219
220
221
    set(ONEDNN_LIBRARY_TYPE "STATIC")
    set(ONEDNN_BUILD_DOC "OFF")
    set(ONEDNN_BUILD_EXAMPLES "OFF")
    set(ONEDNN_BUILD_TESTS "OFF")
    set(ONEDNN_ENABLE_WORKLOAD "INFERENCE")
    set(ONEDNN_ENABLE_PRIMITIVE "MATMUL;REORDER")
    set(ONEDNN_BUILD_GRAPH "OFF")
    set(ONEDNN_ENABLE_JIT_PROFILING "OFF")
    set(ONEDNN_ENABLE_ITT_TASKS "OFF")
    set(ONEDNN_ENABLE_MAX_CPU_ISA "OFF")
    set(ONEDNN_ENABLE_CPU_ISA_HINTS "OFF")
222
    set(ONEDNN_VERBOSE "OFF")
223
224
225
    set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)

    FetchContent_MakeAvailable(oneDNN)
226
227
228
229
230
231
    add_library(dnnl_ext OBJECT "csrc/cpu/dnnl_helper.cpp")
    target_include_directories(
        dnnl_ext
        PUBLIC ${oneDNN_SOURCE_DIR}/include
        PUBLIC ${oneDNN_BINARY_DIR}/include
        PRIVATE ${oneDNN_SOURCE_DIR}/src
232
    )
233
234
235
236
237
238
    target_link_libraries(dnnl_ext dnnl)
    target_compile_options(dnnl_ext PRIVATE ${CXX_COMPILE_FLAGS} -fPIC)
    list(APPEND LIBS dnnl_ext)
    set(USE_ONEDNN ON)
else()
    set(USE_ONEDNN OFF)
239
240
endif()

241
242
message(STATUS "CPU extension compile flags: ${CXX_COMPILE_FLAGS}")

243
244
245
246
247
248
if(ENABLE_NUMA)
    list(APPEND LIBS numa)
else()
    message(STATUS "NUMA is disabled")
    add_compile_definitions(-DVLLM_NUMA_DISABLED)
endif()
249

250
251
252
253
254
255
256
#
# _C extension
#
set(VLLM_EXT_SRC
    "csrc/cpu/activation.cpp"
    "csrc/cpu/attention.cpp"
    "csrc/cpu/cache.cpp"
257
    "csrc/cpu/utils.cpp"
258
    "csrc/cpu/layernorm.cpp"
Thien Tran's avatar
Thien Tran committed
259
    "csrc/cpu/mla_decode.cpp"
260
    "csrc/cpu/pos_encoding.cpp"
261
    "csrc/cpu/torch_bindings.cpp")
262

263
264
if (AVX512_FOUND AND NOT AVX512_DISABLED)
    set(VLLM_EXT_SRC
265
        "csrc/cpu/shm.cpp"
266
        ${VLLM_EXT_SRC})
267
268
269
270
271
272
273
274
275
276
277
    if (ENABLE_AVX512BF16 AND ENABLE_AVX512VNNI)
        set(VLLM_EXT_SRC
            "csrc/cpu/sgl-kernels/gemm.cpp"
            "csrc/cpu/sgl-kernels/gemm_int8.cpp"
            "csrc/cpu/sgl-kernels/gemm_fp8.cpp"
            "csrc/cpu/sgl-kernels/moe.cpp"
            "csrc/cpu/sgl-kernels/moe_int8.cpp"
            "csrc/cpu/sgl-kernels/moe_fp8.cpp"
            ${VLLM_EXT_SRC})
        add_compile_definitions(-DCPU_CAPABILITY_AVX512)
    endif()
278
endif()
279
280

if(USE_ONEDNN)
281
    set(VLLM_EXT_SRC
282
        "csrc/cpu/dnnl_kernels.cpp"
283
284
        ${VLLM_EXT_SRC})
endif()
285

286
287
message(STATUS "CPU extension source files: ${VLLM_EXT_SRC}")

288
289
290
291
#
# Define extension targets
#

292
293
294
295
296
define_gpu_extension_target(
    _C
    DESTINATION vllm
    LANGUAGE CXX
    SOURCES ${VLLM_EXT_SRC}
297
    LIBRARIES ${LIBS}
298
    COMPILE_FLAGS ${CXX_COMPILE_FLAGS}
299
300
    USE_SABI 3
    WITH_SOABI
301
302
)

303
message(STATUS "Enabling C extension.")