vrr/kvcache-ai-ktransformers
2
0
Fork 0
mirror of https://github.com/kvcache-ai/ktransformers.git synced 2025-09-10 23:34:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix compile, add abi check to setup.py

Browse source
This commit is contained in:
Atream 2025-04-08 06:18:30 +00:00
parent f4ae7c85ed
commit 9dd24ecd72
6 changed files with 17 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

12
setup.py
View file

@ -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
)
)