Refactor sgl-kernel build (#2642)

This commit is contained in:
Ke Bao
2024-12-30 18:07:01 +08:00
committed by GitHub
parent bdd2827a80
commit b02da24a5b
9 changed files with 108 additions and 126 deletions

View File

@@ -25,46 +25,29 @@ list(APPEND CMAKE_PREFIX_PATH "${TORCH_CMAKE_PATH}")
find_package(Torch REQUIRED)
# Warp Reduce library
add_library(warp_reduce SHARED
src/sgl-kernel/csrc/warp_reduce.cc
add_library(_kernels SHARED
src/sgl-kernel/csrc/warp_reduce_kernel.cu
)
target_include_directories(warp_reduce
PRIVATE
${CMAKE_CURRENT_SOURCE_DIR}/src/sgl-kernel/csrc
${CUDA_INCLUDE_DIRS}
${TORCH_INCLUDE_DIRS}
)
target_link_libraries(warp_reduce
PRIVATE
${TORCH_LIBRARIES}
Python3::Python
)
# TRT Reduce library
add_library(trt_reduce SHARED
src/sgl-kernel/csrc/trt_reduce.cc
src/sgl-kernel/csrc/trt_reduce_internal.cu
src/sgl-kernel/csrc/trt_reduce_kernel.cu
src/sgl-kernel/csrc/moe_align_kernel.cu
src/sgl-kernel/csrc/sgl_kernel_ops.cu
)
target_include_directories(trt_reduce
target_include_directories(_kernels
PRIVATE
${CMAKE_CURRENT_SOURCE_DIR}/src/sgl-kernel/csrc
${CUDA_INCLUDE_DIRS}
${TORCH_INCLUDE_DIRS}
)
target_link_libraries(trt_reduce
target_link_libraries(_kernels
PRIVATE
${TORCH_LIBRARIES}
Python3::Python
)
# Set common properties for both libraries
foreach(target warp_reduce trt_reduce)
foreach(target _kernels)
set_target_properties(${target} PROPERTIES
CUDA_SEPARABLE_COMPILATION ON
POSITION_INDEPENDENT_CODE ON