CMakeLists.txt 1.61 KB
Newer Older
1
2
3
cmake_minimum_required(VERSION 3.18)
project(sgl-kernel LANGUAGES CXX CUDA)

4
# Basic settings
5
6
7
8
9
10
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CUDA_STANDARD 17)
set(CMAKE_CUDA_STANDARD_REQUIRED ON)

11
12
set(CUTLASS_DIR "3rdparty/cutlass")

13
14
15
16
17
# Set CUDA architectures
set(CMAKE_CUDA_ARCHITECTURES "75;80;86;89;90")
message(STATUS "Building for CUDA architectures: ${CMAKE_CUDA_ARCHITECTURES}")

find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
18

19
# Find PyTorch
20
execute_process(
21
    COMMAND ${Python3_EXECUTABLE} -c "import torch; print(torch.utils.cmake_prefix_path)"
22
23
24
25
26
27
28
    OUTPUT_VARIABLE TORCH_CMAKE_PATH
    OUTPUT_STRIP_TRAILING_WHITESPACE
)
list(APPEND CMAKE_PREFIX_PATH "${TORCH_CMAKE_PATH}")

find_package(Torch REQUIRED)

29
# Warp Reduce library
Ke Bao's avatar
Ke Bao committed
30
add_library(_kernels SHARED
31
32
    src/sgl-kernel/csrc/trt_reduce_internal.cu
    src/sgl-kernel/csrc/trt_reduce_kernel.cu
Ke Bao's avatar
Ke Bao committed
33
34
    src/sgl-kernel/csrc/moe_align_kernel.cu
    src/sgl-kernel/csrc/sgl_kernel_ops.cu
35
)
36

Ke Bao's avatar
Ke Bao committed
37
target_include_directories(_kernels
38
39
40
41
    PRIVATE
        ${CMAKE_CURRENT_SOURCE_DIR}/src/sgl-kernel/csrc
        ${CUDA_INCLUDE_DIRS}
        ${TORCH_INCLUDE_DIRS}
42
43
        ${CUTLASS_DIR}/include
        ${CUTLASS_DIR}/tools/util/include
44
45
)

Ke Bao's avatar
Ke Bao committed
46
target_link_libraries(_kernels
47
48
49
50
51
52
    PRIVATE
        ${TORCH_LIBRARIES}
        Python3::Python
)

# Set common properties for both libraries
Ke Bao's avatar
Ke Bao committed
53
foreach(target _kernels)
54
55
56
57
58
59
60
61
    set_target_properties(${target} PROPERTIES
        CUDA_SEPARABLE_COMPILATION ON
        POSITION_INDEPENDENT_CODE ON
        CUDA_RESOLVE_DEVICE_SYMBOLS ON
        PREFIX ""
        SUFFIX ".so"
    )
endforeach()