From 9dd24ecd724301361a479658e50b5a05537def53 Mon Sep 17 00:00:00 2001
From: Atream <boxin-zhang@qq.com>
Date: Tue, 8 Apr 2025 06:18:30 +0000
Subject: [PATCH] fix compile, add abi check to setup.py

---
 CMakeLists.txt                          | 25 -------------------------
 csrc/balance_serve/CMakeLists.txt       |  1 +
 csrc/balance_serve/kvc2/CMakeLists.txt  |  5 +++--
 csrc/balance_serve/sched/CMakeLists.txt |  1 +
 csrc/ktransformers_ext/CMakeLists.txt   |  1 +
 setup.py                                | 12 +++++++++++-
 6 files changed, 17 insertions(+), 28 deletions(-)
 delete mode 100644 CMakeLists.txt

diff --git a/CMakeLists.txt b/CMakeLists.txt
deleted file mode 100644
index 56e3642..0000000
--- a/CMakeLists.txt
+++ /dev/null
@@ -1,25 +0,0 @@
-cmake_minimum_required(VERSION 3.10)
-
-find_package(PythonInterp REQUIRED)
-
-execute_process(
-    COMMAND ${PYTHON_EXECUTABLE} -c "
-import torch
-print(torch.compiled_with_cxx11_abi())
-"
-    OUTPUT_VARIABLE TORCH_CXX11_ABI
-    OUTPUT_STRIP_TRAILING_WHITESPACE
-)
-
-if(TORCH_CXX11_ABI MATCHES "True|False")
-    if(TORCH_CXX11_ABI)
-        message(STATUS "PyTorch was compiled with CXX11 ABI: ON")
-        set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=1")
-    else()
-        message(STATUS "PyTorch was compiled with CXX11 ABI: OFF")
-        set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=0")
-    endif()
-else()
-    message(WARNING "Could not determine if PyTorch was compiled with CXX11 ABI. Defaulting to OFF.")
-    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D_GLIBCXX_USE_CXX11_ABI=0")
-endif()
diff --git a/csrc/balance_serve/CMakeLists.txt b/csrc/balance_serve/CMakeLists.txt
index 6eb00bd..efabd7d 100644
--- a/csrc/balance_serve/CMakeLists.txt
+++ b/csrc/balance_serve/CMakeLists.txt
@@ -13,6 +13,7 @@ set(CMAKE_CXX_STANDARD 20)
 # set(CMAKE_CXX_FLAGS "-Og -march=native -Wall -Wextra -g -fPIC")
 # set(CMAKE_BUILD_TYPE "Debug")
 set(CMAKE_CXX_FLAGS "-O3 -march=native -Wall -Wextra -fPIC")
+add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=${_GLIBCXX_USE_CXX11_ABI})
 set(CMAKE_BUILD_TYPE "Release")
 
 file(GLOB_RECURSE FMT_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/*.hpp" "${CMAKE_CURRENT_SOURCE_DIR}/*.h")
diff --git a/csrc/balance_serve/kvc2/CMakeLists.txt b/csrc/balance_serve/kvc2/CMakeLists.txt
index b939b80..0c84af2 100644
--- a/csrc/balance_serve/kvc2/CMakeLists.txt
+++ b/csrc/balance_serve/kvc2/CMakeLists.txt
@@ -11,8 +11,9 @@ set(CMAKE_CXX_STANDARD 20)
 # set(CMAKE_CXX_FLAGS "-Og -march=native -Wall -Wextra -Wpedantic -g -fsanitize=address")
 # set(CMAKE_CXX_FLAGS "-march=native -Wall -Wextra -Wpedantic -g")
 # set(CMAKE_CXX_FLAGS "-fPIC -O3 -ffast-math -march=native -Wall -Wextra -g")
-# set(CMAKE_BUILD_TYPE "Release")
-set(CMAKE_BUILD_TYPE "Debug")
+set(CMAKE_BUILD_TYPE "Release")
+add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=${_GLIBCXX_USE_CXX11_ABI})
+# set(CMAKE_BUILD_TYPE "Debug")
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 set(BUILD_TEST OFF)
 set(BUILD_PYTHON_EXT OFF)
diff --git a/csrc/balance_serve/sched/CMakeLists.txt b/csrc/balance_serve/sched/CMakeLists.txt
index 6b5dcda..bad6162 100644
--- a/csrc/balance_serve/sched/CMakeLists.txt
+++ b/csrc/balance_serve/sched/CMakeLists.txt
@@ -1,5 +1,6 @@
 set(CMAKE_CXX_FLAGS "-Og -march=native -Wall -Wextra -g -fPIC")
 # set(CMAKE_CXX_FLAGS "-O3 -march=native -Wall -Wextra -fPIC")
+add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=${_GLIBCXX_USE_CXX11_ABI})
 
 set(UTILS_DIR ${CMAKE_CURRENT_SOURCE_DIR}/utils)
 
diff --git a/csrc/ktransformers_ext/CMakeLists.txt b/csrc/ktransformers_ext/CMakeLists.txt
index b273dd3..774d82d 100644
--- a/csrc/ktransformers_ext/CMakeLists.txt
+++ b/csrc/ktransformers_ext/CMakeLists.txt
@@ -6,6 +6,7 @@ set(CMAKE_CXX_STANDARD 17)
 
 
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -ffast-math -fopenmp")
+add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=${_GLIBCXX_USE_CXX11_ABI})
 set(CMAKE_BUILD_TYPE "Release")
 # set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -ffast-math -fopenmp")
 # set(CMAKE_BUILD_TYPE "Debug")
diff --git a/setup.py b/setup.py
index 6eccef2..c0a3a17 100644
--- a/setup.py
+++ b/setup.py
@@ -25,6 +25,7 @@ import urllib.request
 import urllib.error
 from pathlib import Path
 from packaging.version import parse
+import torch
 import torch.version
 from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
 from setuptools import setup, Extension
@@ -281,6 +282,13 @@ class CMakeExtension(Extension):
         print(name, sourcedir)
         self.sourcedir = sourcedir
 
+def get_cmake_abi_args(cmake_args):
+    if torch.compiled_with_cxx11_abi():
+        cmake_args.append("-D_GLIBCXX_USE_CXX11_ABI=1")
+    else:
+        cmake_args.append("-D_GLIBCXX_USE_CXX11_ABI=0")
+    return cmake_args
+
 class CMakeBuild(BuildExtension):
 
     def build_extension(self, ext) -> None:
@@ -318,6 +326,8 @@ class CMakeBuild(BuildExtension):
             cmake_args += ["-DKTRANSFORMERS_USE_ROCM=ON"]
         else:
             raise ValueError("Unsupported backend: CUDA_HOME, MUSA_HOME, and ROCM_HOME are not set.")
+        
+        cmake_args = get_cmake_abi_args(cmake_args)
         # log cmake_args
         print("CMake args:", cmake_args)
 
@@ -470,4 +480,4 @@ setup(
     version=VersionInfo().get_package_version(),
     cmdclass={"bdist_wheel":BuildWheelsCommand ,"build_ext": CMakeBuild},
     ext_modules=ext_modules
-)
+)
\ No newline at end of file