mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-04-19 22:09:10 +00:00
46b0f36980
* [feat]: Enhance CPU feature detection and support for AVX512 extensions - Added cmake/DetectCPU.cmake for automatic CPU feature detection. - Updated CMakeLists.txt to include auto-detection logic for AVX512 features. - Modified install.sh to include new AVX512_VBMI option for FP8 MoE. - Enhanced _cpu_detect.py to support progressive matching of CPU variants. - Created scripts/check_cpu_features.py for manual CPU feature checks. - Updated setup.py to reflect changes in CPU variant building and environment variables. * [fix](kt-kernel): Add conditional inclusion of FP8 MoE for AVX512 BF16 support * [chore](kt-kernel): update project version to 0.5.0 in CMakeLists.txt and version.py
121 lines
3.9 KiB
Python
Executable File
121 lines
3.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""
|
||
CPU feature detection script for kt-kernel.
|
||
|
||
This script checks if your CPU supports the required instruction sets for FP8 MoE:
|
||
- AVX512F (foundation)
|
||
- AVX512_BF16 (BF16 dot product)
|
||
- AVX512_VNNI (VNNI instructions)
|
||
- AVX512_VBMI (byte permutation)
|
||
|
||
Usage:
|
||
python3 scripts/check_cpu_features.py
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
|
||
|
||
def check_cpuinfo():
|
||
"""Check CPU features via /proc/cpuinfo."""
|
||
try:
|
||
with open("/proc/cpuinfo", "r") as f:
|
||
cpuinfo = f.read().lower()
|
||
return cpuinfo
|
||
except FileNotFoundError:
|
||
return None
|
||
|
||
|
||
def main():
|
||
print("=" * 70)
|
||
print("KT-Kernel CPU Feature Detection")
|
||
print("=" * 70)
|
||
print()
|
||
|
||
cpuinfo = check_cpuinfo()
|
||
|
||
if cpuinfo is None:
|
||
print("❌ /proc/cpuinfo not found (not on Linux?)")
|
||
print(" Cannot detect CPU features automatically.")
|
||
sys.exit(1)
|
||
|
||
# Extract CPU model
|
||
for line in cpuinfo.split("\n"):
|
||
if "model name" in line:
|
||
model = line.split(":")[1].strip()
|
||
print(f"CPU Model: {model}")
|
||
break
|
||
print()
|
||
|
||
# Check AMX support
|
||
print("AMX Support (Intel Sapphire Rapids+):")
|
||
amx_flags = ["amx_tile", "amx_int8", "amx_bf16"]
|
||
amx_status = {}
|
||
for flag in amx_flags:
|
||
has_flag = flag in cpuinfo
|
||
amx_status[flag] = has_flag
|
||
status = "✅" if has_flag else "❌"
|
||
print(f" {status} {flag.upper()}")
|
||
|
||
has_amx = all(amx_status.values())
|
||
print(f"\n Overall AMX Support: {'✅ YES' if has_amx else '❌ NO'}")
|
||
print()
|
||
|
||
# Check AVX512 support
|
||
print("AVX512 Support (required for FP8 MoE):")
|
||
avx512_flags = ["avx512f", "avx512_bf16", "avx512_vnni", "avx512_vbmi"]
|
||
avx512_status = {}
|
||
for flag in avx512_flags:
|
||
has_flag = flag in cpuinfo
|
||
avx512_status[flag] = has_flag
|
||
status = "✅" if has_flag else "❌"
|
||
flag_desc = {
|
||
"avx512f": "AVX512F (foundation)",
|
||
"avx512_bf16": "AVX512_BF16 (BF16 dot product)",
|
||
"avx512_vnni": "AVX512_VNNI (VNNI instructions)",
|
||
"avx512_vbmi": "AVX512_VBMI (byte permutation)",
|
||
}
|
||
print(f" {status} {flag_desc.get(flag, flag.upper())}")
|
||
|
||
has_avx512_full = all(avx512_status.values())
|
||
print(f"\n Overall AVX512 Support: {'✅ YES' if has_avx512_full else '❌ NO'}")
|
||
|
||
if not has_avx512_full and avx512_status["avx512f"]:
|
||
missing = [f for f in avx512_flags if not avx512_status[f]]
|
||
print(f" ⚠️ Warning: AVX512F detected but missing: {', '.join(missing)}")
|
||
print(f" kt-kernel will fall back to AVX2 mode")
|
||
print()
|
||
|
||
# Check AVX2 support
|
||
print("AVX2 Support (fallback):")
|
||
has_avx2 = "avx2" in cpuinfo
|
||
status = "✅" if has_avx2 else "❌"
|
||
print(f" {status} AVX2")
|
||
print()
|
||
|
||
# Recommendation
|
||
print("=" * 70)
|
||
print("Recommendation:")
|
||
print("=" * 70)
|
||
if has_amx:
|
||
print("✅ Your CPU supports AMX - you can use the highest performance mode!")
|
||
print(" Build with: -DKTRANSFORMERS_CPU_USE_AMX_AVX512=ON -DKTRANSFORMERS_CPU_USE_AMX=ON")
|
||
elif has_avx512_full:
|
||
print("✅ Your CPU supports full AVX512 (F/BF16/VNNI/VBMI) - FP8 MoE will work!")
|
||
print(" Build with: -DKTRANSFORMERS_CPU_USE_AMX_AVX512=ON")
|
||
elif avx512_status.get("avx512f", False):
|
||
print("⚠️ Your CPU has AVX512F but missing required extensions.")
|
||
print(" FP8 MoE will NOT work. kt-kernel will fall back to AVX2 mode.")
|
||
print(" Missing extensions:", ", ".join([f for f in avx512_flags if not avx512_status.get(f, False)]))
|
||
elif has_avx2:
|
||
print("ℹ️ Your CPU supports AVX2 only - basic compatibility mode.")
|
||
print(" FP8 MoE will NOT be available, but other features will work.")
|
||
else:
|
||
print("❌ Your CPU does not support the minimum required instruction set (AVX2).")
|
||
print(" kt-kernel may not work on this system.")
|
||
print()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|