diff --git a/sgl-kernel/csrc/cpu/CMakeLists.txt b/sgl-kernel/csrc/cpu/CMakeLists.txt new file mode 100755 index 000000000..fd3c8aae3 --- /dev/null +++ b/sgl-kernel/csrc/cpu/CMakeLists.txt @@ -0,0 +1,87 @@ +cmake_minimum_required(VERSION 3.18 FATAL_ERROR) +project(sgl_kernel) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CXX_EXTENSIONS OFF) + +# Torch +find_package(Torch REQUIRED) +find_package(Python3 COMPONENTS Interpreter Development REQUIRED) + +execute_process( + COMMAND ${Python_EXECUTABLE} + -c "import torch; print(torch.utils.cmake_prefix_path)" + OUTPUT_VARIABLE TORCH_PY_PREFIX + OUTPUT_STRIP_TRAILING_WHITESPACE +) + +message(STATUS ${TORCH_PY_PREFIX}) +list(APPEND CMAKE_PREFIX_PATH ${TORCH_PY_PREFIX}/Torch) +find_package(Torch REQUIRED) + +include_directories( + ${TORCH_INCLUDE_DIRS} + ${TORCH_INSTALL_PREFIX}/include + ${Python3_INCLUDE_DIRS} + ${CMAKE_SOURCE_DIR}/csrc +) + +# Platform-specific library directory +if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|AMD64") + set(PLAT_LIB_DIR "/usr/lib/x86_64-linux-gnu") + elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64") + set(PLAT_LIB_DIR "/usr/lib/aarch64-linux-gnu") + elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "ppc64le|ppc64") + set(PLAT_LIB_DIR "/usr/lib/powerpc64le-linux-gnu") +else() + set(PLAT_LIB_DIR "/usr/lib/${CMAKE_SYSTEM_PROCESSOR}-linux-gnu") +endif() +link_directories(${PLAT_LIB_DIR}) + +set(SOURCES + ${CMAKE_CURRENT_SOURCE_DIR}/activation.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/bmm.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/decode.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/extend.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/gemm.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/gemm_int8.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/moe.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/moe_int8.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/norm.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/qkv_proj.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/topk.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/rope.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/interface.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/shm.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/torch_extension_cpu.cpp +) + +add_compile_options( + -O3 + -Wno-unknown-pragmas + -march=native + -fopenmp +) + +add_library(sgl_kernel_common_ops SHARED ${SOURCES}) + +target_link_libraries(sgl_kernel_common_ops + PRIVATE + ${TORCH_LIBRARIES} + ${Python3_LIBRARIES} + c10 +) + +set_target_properties(sgl_kernel_common_ops PROPERTIES + INSTALL_RPATH "$ORIGIN/../../torch/lib" + PREFIX "" + OUTPUT_NAME "sgl_kernel.common_ops" +) + +target_compile_definitions(sgl_kernel_common_ops PRIVATE TORCH_API_INCLUDE_EXTENSION_H) + +# Install +install(TARGETS sgl_kernel_common_ops + LIBRARY DESTINATION ${Python3_SITEARCH} +) diff --git a/sgl-kernel/csrc/cpu/gemm.h b/sgl-kernel/csrc/cpu/gemm.h index 87dfbd7bf..026e158a0 100644 --- a/sgl-kernel/csrc/cpu/gemm.h +++ b/sgl-kernel/csrc/cpu/gemm.h @@ -1,7 +1,8 @@ #pragma once - #include +#include "common.h" + // amx-bf16 #define TILE_M 16 #define TILE_N 16 diff --git a/sgl-kernel/csrc/cpu/interface.cpp b/sgl-kernel/csrc/cpu/interface.cpp index cc11c4928..633a3ada7 100644 --- a/sgl-kernel/csrc/cpu/interface.cpp +++ b/sgl-kernel/csrc/cpu/interface.cpp @@ -1,5 +1,5 @@ #include -#include +#include #include "shm.h" diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index 6b7cc1d39..02db7f61e 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -14,7 +14,7 @@ limitations under the License. ==============================================================================*/ #include -#include +#include #include #include "shm.h" diff --git a/sgl-kernel/pyproject_cpu.toml b/sgl-kernel/pyproject_cpu.toml new file mode 100644 index 000000000..626973a12 --- /dev/null +++ b/sgl-kernel/pyproject_cpu.toml @@ -0,0 +1,40 @@ +[build-system] +requires = [ + "scikit-build-core>=0.10", + "torch>=2.6.0", + "wheel", +] +build-backend = "scikit_build_core.build" + +[project] +name = "sgl-kernel" +version = "0.1.2.post1" +description = "Kernel Library for SGLang" +readme = "README.md" +requires-python = ">=3.9" +license = { file = "LICENSE" } +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: Apache Software License", + "Environment :: CPU" +] +dependencies = [] + +[project.urls] +"Homepage" = "https://github.com/sgl-project/sglang/tree/main/sgl-kernel" +"Bug Tracker" = "https://github.com/sgl-project/sglang/issues" + +[tool.wheel] +exclude = [ + "dist*", + "tests*", +] + +[tool.scikit-build] +cmake.source-dir = "csrc/cpu" +cmake.build-type = "Release" +minimum-version = "build-system.requires" + +wheel.py-api = "cp39" +wheel.license-files = [] +wheel.packages = ["python/sgl_kernel"] diff --git a/sgl-kernel/setup_cpu.py b/sgl-kernel/setup_cpu.py index 8de4fcf2d..6237011bc 100644 --- a/sgl-kernel/setup_cpu.py +++ b/sgl-kernel/setup_cpu.py @@ -64,6 +64,7 @@ sources = [ "csrc/cpu/topk.cpp", "csrc/cpu/interface.cpp", "csrc/cpu/shm.cpp", + "csrc/cpu/rope.cpp", "csrc/cpu/torch_extension_cpu.cpp", ]