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

Update setup.py

Browse source
This commit is contained in:
Johnny 2025-05-06 19:14:47 +02:00 • committed by GitHub
parent 7530491f5b
commit 6661724b1c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 34 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

52
setup.py
View file

@ -168,32 +168,48 @@ class VersionInfo:
return "avx2" return "avx2"
else: else:
print("Using native cpu instruct") print("Using native cpu instruct")
if sys.platform.startswith("linux"): if sys.platform.startswith("linux"):
with open('/proc/cpuinfo', 'r', encoding="utf-8") as cpu_f: with open('/proc/cpuinfo', 'r', encoding="utf-8") as cpu_f:
cpuinfo = cpu_f.read() cpuinfo = cpu_f.read()
flags_line = [line for line in cpuinfo.split( if platform.machine() == "aarch64":
'\n') if line.startswith('flags')][0] # Adapt this part based on GH200's /proc/cpuinfo
flags = flags_line.split(':')[1].strip().split(' ') for line in cpuinfo.split('\n'):
# fancy with AVX512-VL, AVX512-BW, AVX512-DQ, AVX512-VNNI if line.startswith('Features'):
for flag in flags: features_line = line
if 'avx512bw' in flag: features = features_line.split(':')[1].strip().split(' ')
return 'fancy' if 'sve' in features: # Example: Scalable Vector Extension
for flag in flags: return 'sve' # Or a custom label
if 'avx512' in flag: elif 'neon' in features:
return 'avx512' return 'neon'
for flag in flags: else:
if 'avx2' in flag: print("Using generic Arm CPU instructions")
return 'avx2' return 'native_arm' # Or a default Arm label
raise ValueError( print("Warning: Could not find 'Features' line in /proc/cpuinfo on aarch64. Using native.")
"Unsupported cpu Instructions: {}".format(flags_line)) return 'native' # Fallback for aarch64 if 'Features' not found
else: # Assume x86-like if not aarch64
for line in cpuinfo.split('\n'):
if line.startswith('flags'):
flags_line = line
flags = flags_line.split(':')[1].strip().split(' ')
if 'avx512bw' in flags:
return 'fancy'
elif 'avx512' in flags:
return 'avx512'
elif 'avx2' in flags:
return 'avx2'
raise ValueError(
"Unsupported cpu Instructions: {}".format(flags_line))
print("Warning: Could not find 'flags' line in /proc/cpuinfo on x86-like. Using native.")
return 'native' # Fallback for x86-like if 'flags' not found
elif sys.platform == "win32": elif sys.platform == "win32":
from cpufeature.extension import CPUFeature from cpufeature.extension import CPUFeature
if CPUFeature.get("AVX512bw", False): if CPUFeature.get("AVX512bw", False):
return 'fancy' return 'fancy'
if CPUFeature.get("AVX512f", False): elif CPUFeature.get("AVX512f", False):
return 'avx512' return 'avx512'
if CPUFeature.get("AVX2", False): elif CPUFeature.get("AVX2", False):
return 'avx2' return 'avx2'
raise ValueError( raise ValueError(
"Unsupported cpu Instructions: {}".format(str(CPUFeature))) "Unsupported cpu Instructions: {}".format(str(CPUFeature)))