kvcache-ai/ktransformers
1
0
Fork 0
mirror of https://github.com/kvcache-ai/ktransformers.git synced 2026-04-29 18:51:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

[feat]: Automatically detect whether blis is installed on amd cpus (#1702)

Browse Source
This commit is contained in:
Jianwei Dong
2025-12-11 14:25:36 +08:00
committed by GitHub
parent 6431888928
commit c65febe05c
1 changed files with 18 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
kt-kernel/setup.py
View File

@@ -18,7 +18,7 @@ Environment knobs (export before running pip install .):
CPUINFER_CPU_INSTRUCT=FANCY One of: NATIVE|FANCY|AVX512|AVX2 (maps to CMake flags) CPUINFER_CPU_INSTRUCT=FANCY One of: NATIVE|FANCY|AVX512|AVX2 (maps to CMake flags)
CPUINFER_ENABLE_AMX=OFF ON/OFF -> -DKTRANSFORMERS_CPU_USE_AMX CPUINFER_ENABLE_AMX=OFF ON/OFF -> -DKTRANSFORMERS_CPU_USE_AMX
CPUINFER_ENABLE_MLA=OFF ON/OFF -> -DKTRANSFORMERS_CPU_MLA CPUINFER_ENABLE_MLA=OFF ON/OFF -> -DKTRANSFORMERS_CPU_MLA
CPUINFER_ENABLE_AMD=OFF ON/OFF -> -DKTRANSFORMERS_CPU_MOE_AMD CPUINFER_ENABLE_BLIS=OFF ON/OFF -> -DKTRANSFORMERS_CPU_MOE_AMD
CPUINFER_ENABLE_KML=OFF ON/OFF -> -DKTRANSFORMERS_CPU_USE_KML CPUINFER_ENABLE_KML=OFF ON/OFF -> -DKTRANSFORMERS_CPU_USE_KML
CPUINFER_ENABLE_AVX512=OFF ON/OFF -> -DKTRANSFORMERS_CPU_USE_AMX_AVX512 CPUINFER_ENABLE_AVX512=OFF ON/OFF -> -DKTRANSFORMERS_CPU_USE_AMX_AVX512
CPUINFER_BLIS_ROOT=/path/to/blis Forward to -DBLIS_ROOT CPUINFER_BLIS_ROOT=/path/to/blis Forward to -DBLIS_ROOT
@@ -53,6 +53,7 @@ from setuptools import setup, Extension
from setuptools.command.build_ext import build_ext from setuptools.command.build_ext import build_ext
import shutil import shutil
# ------------------------- # -------------------------
# Env parsing helpers # Env parsing helpers
# ------------------------- # -------------------------
@@ -90,6 +91,7 @@ def _forward_str_env(cmake_args: list[str], env_name: str, cmake_flag: str) -> b
print(f"-- Forward {env_name} -> -D{cmake_flag}={v}") print(f"-- Forward {env_name} -> -D{cmake_flag}={v}")
return True return True
################################################################################ ################################################################################
# Helpers # Helpers
################################################################################ ################################################################################
@@ -296,19 +298,19 @@ class CMakeBuild(build_ext):
# Vendor / feature specific toggles # Vendor / feature specific toggles
# AMD MoE: explicit env overrides; otherwise default ON on AMD CPU # AMD MoE: explicit env overrides; otherwise default ON on AMD CPU
if not _forward_bool_env(cmake_args, "CPUINFER_ENABLE_AMD", "KTRANSFORMERS_CPU_MOE_AMD"): _forward_bool_env(cmake_args, "CPUINFER_ENABLE_BLIS", "KTRANSFORMERS_CPU_MOE_AMD")
if d.get("vendor") == "amd": # if d.get("vendor") == "amd":
auto_moe_kernel_ = True # auto_moe_kernel_ = True
cmake_args.append("-DKTRANSFORMERS_CPU_MOE_AMD=ON") # cmake_args.append("-DKTRANSFORMERS_CPU_MOE_AMD=ON")
print("-- Detected AMD CPU; enabling AMD MoE kernel (-DKTRANSFORMERS_CPU_MOE_AMD=ON)") # print("-- Detected AMD CPU; enabling AMD MoE kernel (-DKTRANSFORMERS_CPU_MOE_AMD=ON)")
_forward_str_env(cmake_args, "CPUINFER_BLIS_ROOT", "BLIS_ROOT") # _forward_str_env(cmake_args, "CPUINFER_BLIS_ROOT", "BLIS_ROOT")
# KML: explicit env overrides; otherwise default ON on ARM # KML: explicit env overrides; otherwise default ON on ARM
if not _forward_bool_env(cmake_args, "CPUINFER_ENABLE_KML", "KTRANSFORMERS_CPU_USE_KML"): _forward_bool_env(cmake_args, "CPUINFER_ENABLE_KML", "KTRANSFORMERS_CPU_USE_KML")
if d.get("vendor") == "arm": # if d.get("vendor") == "arm":
auto_moe_kernel_ = True # auto_moe_kernel_ = True
cmake_args.append("-DKTRANSFORMERS_CPU_USE_KML=ON") # cmake_args.append("-DKTRANSFORMERS_CPU_USE_KML=ON")
print("-- Detected ARM CPU; enabling KML (-DKTRANSFORMERS_CPU_USE_KML=ON)") # print("-- Detected ARM CPU; enabling KML (-DKTRANSFORMERS_CPU_USE_KML=ON)")
# AMX: explicit env overrides; else enable if detected # AMX: explicit env overrides; else enable if detected
if not _forward_bool_env(cmake_args, "CPUINFER_ENABLE_AMX", "KTRANSFORMERS_CPU_USE_AMX"): if not _forward_bool_env(cmake_args, "CPUINFER_ENABLE_AMX", "KTRANSFORMERS_CPU_USE_AMX"):
@@ -330,14 +332,16 @@ class CMakeBuild(build_ext):
# Auto-enable MOE kernel only when env explicitly turns on AMD or KML backend # Auto-enable MOE kernel only when env explicitly turns on AMD or KML backend
# (Do not enable purely on vendor auto-detection to avoid surprise behavior.) # (Do not enable purely on vendor auto-detection to avoid surprise behavior.)
amd_env = _env_get_bool("CPUINFER_ENABLE_AMD", None) amd_env = _env_get_bool("CPUINFER_ENABLE_BLIS", None)
kml_env = _env_get_bool("CPUINFER_ENABLE_KML", None) kml_env = _env_get_bool("CPUINFER_ENABLE_KML", None)
if amd_env or kml_env: if amd_env or kml_env:
auto_moe_kernel_ = True auto_moe_kernel_ = True
already_set = any("KTRANSFORMERS_CPU_MOE_KERNEL" in a for a in cmake_args) already_set = any("KTRANSFORMERS_CPU_MOE_KERNEL" in a for a in cmake_args)
if not already_set and auto_moe_kernel_: if not already_set and auto_moe_kernel_:
cmake_args.append("-DKTRANSFORMERS_CPU_MOE_KERNEL=ON") cmake_args.append("-DKTRANSFORMERS_CPU_MOE_KERNEL=ON")
print("-- Auto-enabling MOE kernel (-DKTRANSFORMERS_CPU_MOE_KERNEL=ON) because CPUINFER_ENABLE_AMD or CPUINFER_ENABLE_KML is ON") print(
"-- Auto-enabling MOE kernel (-DKTRANSFORMERS_CPU_MOE_KERNEL=ON) because CPUINFER_ENABLE_BLIS or CPUINFER_ENABLE_KML is ON"
)
# Friendly summary # Friendly summary
print( print(