160 lines
4.7 KiB
Plaintext
160 lines
4.7 KiB
Plaintext
cmake_minimum_required(VERSION 3.24 FATAL_ERROR)
|
|
project(sgl_kernel LANGUAGES CXX)
|
|
|
|
# Cmake
|
|
set(CMAKE_CXX_STANDARD 17)
|
|
set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
|
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
|
|
set(CMAKE_SHARED_LIBRARY_PREFIX "")
|
|
|
|
set(CMAKE_COLOR_DIAGNOSTICS ON)
|
|
set(CMAKE_VERBOSE_MAKEFILE ON CACHE BOOL "ON")
|
|
|
|
# Python / Torch
|
|
find_package(Python COMPONENTS Interpreter Development.Module ${SKBUILD_SABI_COMPONENT} REQUIRED)
|
|
|
|
execute_process(
|
|
COMMAND ${Python_EXECUTABLE} -c "import torch; print(torch.utils.cmake_prefix_path)"
|
|
OUTPUT_VARIABLE TORCH_PY_PREFIX
|
|
OUTPUT_STRIP_TRAILING_WHITESPACE
|
|
)
|
|
|
|
set(Torch_DIR "${TORCH_PY_PREFIX}/Torch")
|
|
list(APPEND CMAKE_PREFIX_PATH "${TORCH_PY_PREFIX}/Torch")
|
|
find_package(Torch REQUIRED)
|
|
|
|
execute_process(
|
|
COMMAND ${Python_EXECUTABLE} -c "import torch; print(int(torch._C._GLIBCXX_USE_CXX11_ABI))"
|
|
OUTPUT_VARIABLE TORCH_CXX11_ABI
|
|
OUTPUT_STRIP_TRAILING_WHITESPACE
|
|
)
|
|
if(TORCH_CXX11_ABI STREQUAL "0")
|
|
add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=0)
|
|
else()
|
|
add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=1)
|
|
endif()
|
|
|
|
# ROCm/HIP
|
|
enable_language(HIP)
|
|
find_package(hip REQUIRED CONFIG)
|
|
|
|
# Determine AMDGPU target from environment variable or default to gfx942
|
|
set(AMDGPU_TARGET_ENV "$ENV{AMDGPU_TARGET}")
|
|
|
|
if(AMDGPU_TARGET_ENV)
|
|
# Use environment variable if specified
|
|
set(AMDGPU_TARGETS "${AMDGPU_TARGET_ENV}")
|
|
message(STATUS "Using AMDGPU_TARGET from environment: ${AMDGPU_TARGETS}")
|
|
else()
|
|
# Default to gfx942 only
|
|
set(AMDGPU_TARGETS "gfx942")
|
|
message(STATUS "AMDGPU_TARGET not set, defaulting to gfx942")
|
|
endif()
|
|
|
|
# Set HIP architectures
|
|
set(CMAKE_HIP_ARCHITECTURES ${AMDGPU_TARGETS})
|
|
|
|
# FP8 macro selection
|
|
# Always define HIP_FP8_TYPE_FNUZ=1 (for gfx942 and host compilation)
|
|
# Additionally define HIP_FP8_TYPE_E4M3=1 when building for gfx950
|
|
# The existing utils.h logic will pick the right one based on architecture
|
|
set(SGL_FP8_MACROS "-DHIP_FP8_TYPE_FNUZ=1")
|
|
|
|
if(AMDGPU_TARGETS MATCHES "gfx950")
|
|
list(APPEND SGL_FP8_MACROS "-DHIP_FP8_TYPE_E4M3=1")
|
|
message(STATUS "Multi-arch build: Enabling both HIP_FP8_TYPE_FNUZ (gfx942) and HIP_FP8_TYPE_E4M3 (gfx950)")
|
|
elseif(AMDGPU_TARGETS MATCHES "gfx942")
|
|
message(STATUS "Single-arch build: Enabling HIP_FP8_TYPE_FNUZ for gfx942")
|
|
else()
|
|
message(FATAL_ERROR "Unsupported AMDGPU_TARGET '${AMDGPU_TARGETS}'. Expected 'gfx942' or 'gfx950' or both.")
|
|
endif()
|
|
|
|
# TopK dynamic smem bytes
|
|
# Dynamic shared-memory budget for the TopK kernels.
|
|
# - gfx942 (MI300/MI325): LDS is typically 64KB per workgroup -> keep dynamic smem <= ~48KB
|
|
# (leaves room for static shared allocations in the kernel).
|
|
# - gfx95x (MI350): LDS is larger (e.g. 160KB per CU) -> allow the original 128KB dynamic smem.
|
|
if(AMDGPU_TARGET_ONE STREQUAL "gfx942")
|
|
math(EXPR SGL_TOPK_DYNAMIC_SMEM_BYTES "48 * 1024")
|
|
else()
|
|
math(EXPR SGL_TOPK_DYNAMIC_SMEM_BYTES "32 * 1024 * 4")
|
|
endif()
|
|
|
|
set(SGL_TOPK_MACROS "-DSGL_TOPK_DYNAMIC_SMEM_BYTES=${SGL_TOPK_DYNAMIC_SMEM_BYTES}")
|
|
|
|
# Paths / includes
|
|
set(PROJ_ROOT ${CMAKE_CURRENT_LIST_DIR})
|
|
set(SGL_INCLUDE_DIRS
|
|
${PROJ_ROOT}/include
|
|
${PROJ_ROOT}/include/impl
|
|
${PROJ_ROOT}/csrc
|
|
${TORCH_INCLUDE_DIRS}
|
|
)
|
|
|
|
# Platform-specific library directory
|
|
set(PLAT_LIB_DIR "/usr/lib/x86_64-linux-gnu")
|
|
link_directories(${PLAT_LIB_DIR})
|
|
|
|
# Sources
|
|
set(SOURCES
|
|
${PROJ_ROOT}/csrc/allreduce/custom_all_reduce.hip
|
|
${PROJ_ROOT}/csrc/allreduce/deterministic_all_reduce.hip
|
|
${PROJ_ROOT}/csrc/allreduce/quick_all_reduce.hip
|
|
${PROJ_ROOT}/csrc/common_extension_rocm.cc
|
|
${PROJ_ROOT}/csrc/elementwise/activation.hip
|
|
${PROJ_ROOT}/csrc/elementwise/pos_enc.hip
|
|
${PROJ_ROOT}/csrc/elementwise/topk.hip
|
|
${PROJ_ROOT}/csrc/grammar/apply_token_bitmask_inplace_hip.hip
|
|
${PROJ_ROOT}/csrc/kvcacheio/transfer.hip
|
|
${PROJ_ROOT}/csrc/moe/moe_align_kernel.hip
|
|
${PROJ_ROOT}/csrc/moe/moe_topk_softmax_kernels.hip
|
|
${PROJ_ROOT}/csrc/moe/moe_topk_sigmoid_kernels.hip
|
|
${PROJ_ROOT}/csrc/speculative/eagle_utils.hip
|
|
)
|
|
set_source_files_properties(
|
|
${SOURCES}
|
|
PROPERTIES
|
|
LANGUAGE HIP
|
|
)
|
|
|
|
# Compile / Link flags
|
|
add_compile_options($<$<COMPILE_LANGUAGE:CXX>:-O3>)
|
|
|
|
set(SGL_HIP_FLAGS
|
|
-DNDEBUG
|
|
-DOPERATOR_NAMESPACE=sgl_kernel
|
|
-O3
|
|
-std=c++17
|
|
-DENABLE_BF16
|
|
-DENABLE_FP8
|
|
${SGL_FP8_MACROS}
|
|
-Wno-pass-failed
|
|
-Wundefined-internal
|
|
${SGL_TOPK_MACROS}
|
|
)
|
|
|
|
# Python extension
|
|
Python_add_library(common_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${SOURCES})
|
|
target_include_directories(common_ops PRIVATE ${SGL_INCLUDE_DIRS})
|
|
|
|
# Apply per-language flags
|
|
target_compile_options(common_ops PRIVATE
|
|
$<$<COMPILE_LANGUAGE:HIP>:${SGL_HIP_FLAGS}>
|
|
)
|
|
|
|
target_link_libraries(common_ops PRIVATE
|
|
${TORCH_LIBRARIES}
|
|
hip::device
|
|
hip::host
|
|
hiprtc
|
|
amdhip64
|
|
)
|
|
|
|
target_link_options(common_ops PRIVATE
|
|
"SHELL:-Wl,-rpath,'\$ORIGIN/../../torch/lib'"
|
|
)
|
|
|
|
install(TARGETS common_ops
|
|
LIBRARY DESTINATION sgl_kernel
|
|
)
|