From 7301a39b13c769e3b9eac38f8e08e6c22018a799 Mon Sep 17 00:00:00 2001
From: Yineng Zhang <me@zhyncs.com>
Date: Sun, 1 Dec 2024 23:55:19 +0800
Subject: [PATCH] fix: resolve CodeQL cpp issue (#2305)

---
 docker/Dockerfile.dev     |  8 ++++++-
 sgl-kernel/CMakeLists.txt | 47 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 sgl-kernel/CMakeLists.txt

diff --git a/docker/Dockerfile.dev b/docker/Dockerfile.dev
index a0ff89b64..8bd052899 100644
--- a/docker/Dockerfile.dev
+++ b/docker/Dockerfile.dev
@@ -4,7 +4,6 @@ FROM lmsysorg/sglang:latest
 # Install development tools and utilities
 RUN apt-get update && apt-get install -y \
     gdb \
-    cmake \
     ninja-build \
     vim \
     tmux \
@@ -50,6 +49,13 @@ RUN curl -L https://github.com/clangd/clangd/releases/download/18.1.3/clangd-lin
     && cp -r clangd_18.1.3/lib/* /usr/local/lib/ \
     && rm -rf clangd_18.1.3 clangd.zip
 
+# Install CMake
+RUN curl -L https://cmake.org/download/#:~:text=cmake%2D3.31.1%2Dlinux%2Dx86_64.tar.gz -o cmake.tar.gz \
+    && tar -xzf cmake.tar.gz \
+    && cp -r cmake-3.31.1-linux-x86_64/bin/* /usr/local/bin/ \
+    && cp -r cmake-3.31.1-linux-x86_64/share/* /usr/local/share/ \
+    && rm -rf cmake-3.31.1-linux-x86_64 cmake.tar.gz
+
 # Add yank script
 COPY --chown=root:root <<-"EOF" /usr/local/bin/yank
 #!/bin/bash
diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt
new file mode 100644
index 000000000..c635b75c3
--- /dev/null
+++ b/sgl-kernel/CMakeLists.txt
@@ -0,0 +1,47 @@
+cmake_minimum_required(VERSION 3.18)
+project(sgl-kernel LANGUAGES CXX CUDA)
+
+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)
+
+find_package(PythonInterp 3 REQUIRED)
+find_package(PythonLibs 3 REQUIRED)
+
+execute_process(
+    COMMAND ${PYTHON_EXECUTABLE} -c "import torch; print(torch.utils.cmake_prefix_path)"
+    OUTPUT_VARIABLE TORCH_CMAKE_PATH
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+)
+
+message(STATUS "PYTHON_EXECUTABLE: ${PYTHON_EXECUTABLE}")
+message(STATUS "TORCH_CMAKE_PATH: ${TORCH_CMAKE_PATH}")
+
+list(APPEND CMAKE_PREFIX_PATH "${TORCH_CMAKE_PATH}")
+
+find_package(Torch REQUIRED)
+
+include_directories(${PYTHON_INCLUDE_DIRS})
+
+add_library(warp_reduce SHARED
+    src/sgl-kernel/csrc/warp_reduce.cc
+    src/sgl-kernel/csrc/warp_reduce_kernel.cu
+)
+
+target_include_directories(warp_reduce PRIVATE
+    ${CUDA_INCLUDE_DIRS}
+    ${TORCH_INCLUDE_DIRS}
+)
+
+target_link_libraries(warp_reduce PRIVATE
+    ${TORCH_LIBRARIES}
+    ${PYTHON_LIBRARIES}
+)
+
+set_target_properties(warp_reduce PROPERTIES
+    CUDA_SEPARABLE_COMPILATION ON
+)