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

Add ABI args to compile KTransformersOps

Browse source
This commit is contained in:
zydxt 2025-06-03 19:53:33 +08:00
parent dcba29b291
commit 0dce748549
1 changed files with 9 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

18
setup.py
View file

@ -465,12 +465,12 @@ class CMakeExtension(Extension):
print(name, sourcedir) print(name, sourcedir)
self.sourcedir = sourcedir self.sourcedir = sourcedir
def get_cmake_abi_args(cmake_args): def get_compile_abi_args(compile_args):
if torch.compiled_with_cxx11_abi(): if torch.compiled_with_cxx11_abi():
cmake_args.append("-D_GLIBCXX_USE_CXX11_ABI=1") compile_args.append("-D_GLIBCXX_USE_CXX11_ABI=1")
else: else:
cmake_args.append("-D_GLIBCXX_USE_CXX11_ABI=0") compile_args.append("-D_GLIBCXX_USE_CXX11_ABI=0")
return cmake_args return compile_args
class CMakeBuild(BuildExtension): class CMakeBuild(BuildExtension):
@ -512,7 +512,7 @@ class CMakeBuild(BuildExtension):
else: else:
raise ValueError("Unsupported backend: CUDA_HOME, MUSA_HOME, and ROCM_HOME are not set and XPU is not available.") raise ValueError("Unsupported backend: CUDA_HOME, MUSA_HOME, and ROCM_HOME are not set and XPU is not available.")
cmake_args = get_cmake_abi_args(cmake_args) cmake_args = get_compile_abi_args(cmake_args)
# log cmake_args # log cmake_args
print("CMake args:", cmake_args) print("CMake args:", cmake_args)
@ -603,12 +603,12 @@ if CUDA_HOME is not None or ROCM_HOME is not None:
], ],
extra_compile_args={ extra_compile_args={
'cxx': ['-O3', '-DKTRANSFORMERS_USE_CUDA'], 'cxx': ['-O3', '-DKTRANSFORMERS_USE_CUDA'],
'nvcc': [ 'nvcc': get_compile_abi_args([
'-O3', '-O3',
# '--use_fast_math', # '--use_fast_math',
'-Xcompiler', '-fPIC', '-Xcompiler', '-fPIC',
'-DKTRANSFORMERS_USE_CUDA', '-DKTRANSFORMERS_USE_CUDA',
] ])
} }
) )
elif MUSA_HOME is not None: elif MUSA_HOME is not None:
@ -627,11 +627,11 @@ elif MUSA_HOME is not None:
], ],
extra_compile_args={ extra_compile_args={
'cxx': ['force_mcc'], 'cxx': ['force_mcc'],
'mcc': [ 'mcc': get_compile_abi_args([
'-O3', '-O3',
'-DKTRANSFORMERS_USE_MUSA', '-DKTRANSFORMERS_USE_MUSA',
'-DTHRUST_IGNORE_CUB_VERSION_CHECK', '-DTHRUST_IGNORE_CUB_VERSION_CHECK',
] ])
} }
) )
elif torch.xpu.is_available(): #XPUExtension is not available now. elif torch.xpu.is_available(): #XPUExtension is not available now.