[CPU] Add CMakeLists.txt for sgl-kernel (#6115)

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}
+)

sgl-kernel/csrc/cpu/gemm.h

@@ -1,7 +1,8 @@
 #pragma once
-
 #include <ATen/native/CPUBlas.h>
 
+#include "common.h"
+
 // amx-bf16
 #define TILE_M 16
 #define TILE_N 16

sgl-kernel/csrc/cpu/interface.cpp

@@ -1,5 +1,5 @@
 #include <ATen/record_function.h>
-#include <torch/extension.h>
+#include <torch/all.h>
 
 #include "shm.h"
 

sgl-kernel/csrc/cpu/torch_extension_cpu.cpp

@@ -14,7 +14,7 @@ limitations under the License.
 ==============================================================================*/
 
 #include <ATen/ATen.h>
-#include <torch/extension.h>
+#include <torch/all.h>
 #include <torch/library.h>
 
 #include "shm.h"

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"]

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",
 ]