mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-03-15 02:47:22 +00:00
307 lines
9.7 KiB
Bash
Executable File
307 lines
9.7 KiB
Bash
Executable File
#!/usr/bin/env bash
|
||
set -euo pipefail
|
||
|
||
usage() {
|
||
cat <<EOF
|
||
Usage: $0 [SUBCOMMAND] [BUILD_OPTIONS]
|
||
|
||
Two-step installation in one file. Choose a subcommand:
|
||
|
||
SUBCOMMANDS:
|
||
deps Install system prerequisites only
|
||
build Build and install kt-kernel (no dependency install)
|
||
all Run deps then build (default when no subcommand)
|
||
-h, --help Show this help message
|
||
|
||
BUILD_OPTIONS (for "build" or "all"):
|
||
(none) Auto-detect CPU and configure automatically (recommended)
|
||
--manual Skip auto-detection, use manual configuration (see below)
|
||
--no-clean Do not delete local build/ before building (default cleans)
|
||
|
||
AUTO-DETECTION (Default):
|
||
The script will automatically detect your CPU capabilities and configure:
|
||
- If AMX instructions detected → NATIVE + AMX=ON
|
||
- Otherwise → NATIVE + AMX=OFF
|
||
|
||
MANUAL CONFIGURATION:
|
||
Use --manual flag and set these environment variables before running:
|
||
|
||
CPUINFER_CPU_INSTRUCT - CPU instruction set
|
||
Options: NATIVE, AVX512, AVX2, FANCY
|
||
CPUINFER_ENABLE_AMX - Enable Intel AMX support
|
||
Options: ON, OFF
|
||
|
||
Manual configuration examples:
|
||
|
||
┌─────────────────────────────────────────────────────────────────────────┐
|
||
│ Configuration │ Use Case │
|
||
├──────────────────────────────────┼──────────────────────────────────────┤
|
||
│ NATIVE + AMX=ON │ Best performance on AMX CPUs │
|
||
│ AVX512 + AMX=OFF │ AVX512 CPUs without AMX │
|
||
│ AVX2 + AMX=OFF │ Older CPUs or maximum compatibility │
|
||
└──────────────────────────────────┴──────────────────────────────────────┘
|
||
|
||
Example manual build:
|
||
export CPUINFER_CPU_INSTRUCT=AVX512
|
||
export CPUINFER_ENABLE_AMX=OFF
|
||
$0 build --manual
|
||
|
||
Advanced option (for binary distribution):
|
||
FANCY - AVX512 with full extensions for Ice Lake+/Zen 4+
|
||
Use this when building pre-compiled binaries to distribute.
|
||
|
||
Optional variables (with defaults):
|
||
CPUINFER_BUILD_TYPE=Release Build type (Debug/RelWithDebInfo/Release)
|
||
CPUINFER_PARALLEL=8 Number of parallel build jobs
|
||
CPUINFER_VERBOSE=1 Verbose build output (0/1)
|
||
|
||
EOF
|
||
exit 1
|
||
}
|
||
|
||
install_dependencies() {
|
||
echo "Checking and installing system dependencies..."
|
||
|
||
# Determine if we need to use sudo
|
||
SUDO=""
|
||
if [ "${EUID:-0}" -ne 0 ]; then
|
||
if command -v sudo &> /dev/null; then
|
||
SUDO="sudo"
|
||
else
|
||
echo "Warning: Not running as root and sudo not found. Package installation may fail."
|
||
echo "Please run as root or install sudo."
|
||
fi
|
||
fi
|
||
|
||
if command -v conda &> /dev/null; then
|
||
echo "Installing cmake via conda..."
|
||
conda install -y cmake
|
||
else
|
||
echo "Warning: conda not found. Skipping cmake installation via conda."
|
||
echo "Please install conda or manually install cmake."
|
||
fi
|
||
|
||
# Detect OS type
|
||
if [ -f /etc/os-release ]; then
|
||
. /etc/os-release
|
||
OS=$ID
|
||
elif [ -f /etc/debian_version ]; then
|
||
OS="debian"
|
||
elif [ -f /etc/redhat-release ]; then
|
||
OS="rhel"
|
||
else
|
||
echo "Warning: Unable to detect OS type. Skipping dependency installation."
|
||
return 0
|
||
fi
|
||
|
||
# Install dependencies based on OS
|
||
case "$OS" in
|
||
debian|ubuntu|linuxmint|pop)
|
||
echo "Detected Debian-based system. Installing libhwloc-dev and pkg-config..."
|
||
$SUDO apt update
|
||
$SUDO apt install -y libhwloc-dev pkg-config
|
||
;;
|
||
fedora|rhel|centos|rocky|almalinux)
|
||
echo "Detected Red Hat-based system. Installing hwloc-devel and pkgconfig..."
|
||
$SUDO dnf install -y hwloc-devel pkgconfig || $SUDO yum install -y hwloc-devel pkgconfig
|
||
;;
|
||
arch|manjaro)
|
||
echo "Detected Arch-based system. Installing hwloc and pkgconf..."
|
||
$SUDO pacman -S --noconfirm hwloc pkgconf
|
||
;;
|
||
opensuse*|sles)
|
||
echo "Detected openSUSE-based system. Installing hwloc-devel and pkg-config..."
|
||
$SUDO zypper install -y hwloc-devel pkg-config
|
||
;;
|
||
*)
|
||
echo "Warning: Unsupported OS '$OS'. Please manually install libhwloc-dev and pkg-config."
|
||
;;
|
||
esac
|
||
}
|
||
|
||
# Function to detect CPU features
|
||
detect_cpu_features() {
|
||
local has_amx=0
|
||
|
||
if [ -f /proc/cpuinfo ]; then
|
||
# Check for AMX support on Linux
|
||
if grep -q "amx_tile\|amx_int8\|amx_bf16" /proc/cpuinfo; then
|
||
has_amx=1
|
||
fi
|
||
elif [ "$(uname)" = "Darwin" ]; then
|
||
# macOS doesn't have AMX (ARM or Intel without AMX)
|
||
has_amx=0
|
||
fi
|
||
|
||
echo "$has_amx"
|
||
}
|
||
|
||
build_step() {
|
||
# Parse build-only flags from arguments to this function
|
||
local MANUAL_MODE=0
|
||
local CLEAN_BUILD=1
|
||
while [[ $# -gt 0 ]]; do
|
||
case "$1" in
|
||
--manual) MANUAL_MODE=1; shift ;;
|
||
--no-clean) CLEAN_BUILD=0; shift ;;
|
||
-h|--help) usage ;;
|
||
*) break ;;
|
||
esac
|
||
done
|
||
|
||
# Clean local build directory to ensure a fresh CMake/configure
|
||
local REPO_ROOT
|
||
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||
if [[ "$CLEAN_BUILD" -eq 1 ]]; then
|
||
if [[ -d "$REPO_ROOT/build" ]]; then
|
||
echo "Cleaning previous build directory: $REPO_ROOT/build"
|
||
rm -rf "$REPO_ROOT/build"
|
||
fi
|
||
else
|
||
echo "Skipping clean of $REPO_ROOT/build (requested by --no-clean)"
|
||
fi
|
||
|
||
# Check for multi-variant build mode (Docker environment)
|
||
if [ "${CPUINFER_BUILD_ALL_VARIANTS:-0}" = "1" ]; then
|
||
echo "=========================================="
|
||
echo "Building ALL CPU variants (AMX/AVX512/AVX2)"
|
||
echo "=========================================="
|
||
echo ""
|
||
echo "This will build three variants in a single wheel:"
|
||
echo " - AMX variant (Intel Sapphire Rapids+)"
|
||
echo " - AVX512 variant (Intel Skylake-X/Ice Lake+)"
|
||
echo " - AVX2 variant (maximum compatibility)"
|
||
echo ""
|
||
echo "Runtime CPU detection will automatically select the best variant."
|
||
echo ""
|
||
|
||
export CPUINFER_FORCE_REBUILD=1
|
||
export CPUINFER_BUILD_TYPE=${CPUINFER_BUILD_TYPE:-Release}
|
||
export CPUINFER_PARALLEL=${CPUINFER_PARALLEL:-8}
|
||
|
||
echo "Building with:"
|
||
echo " CPUINFER_BUILD_ALL_VARIANTS=1"
|
||
echo " CPUINFER_BUILD_TYPE=$CPUINFER_BUILD_TYPE"
|
||
echo " CPUINFER_PARALLEL=$CPUINFER_PARALLEL"
|
||
echo ""
|
||
|
||
pip install . -v
|
||
return 0
|
||
fi
|
||
|
||
if [ "$MANUAL_MODE" = "0" ]; then
|
||
# Auto-detection mode
|
||
echo "=========================================="
|
||
echo "Auto-detecting CPU capabilities..."
|
||
echo "=========================================="
|
||
echo ""
|
||
|
||
HAS_AMX=$(detect_cpu_features)
|
||
|
||
if [ "$HAS_AMX" = "1" ]; then
|
||
echo "✓ AMX instructions detected"
|
||
export CPUINFER_CPU_INSTRUCT=NATIVE
|
||
export CPUINFER_ENABLE_AMX=ON
|
||
echo " Configuration: NATIVE + AMX=ON (best performance)"
|
||
echo ""
|
||
echo " ⚠️ Note: If you plan to use LLAMAFILE backend, use manual mode:"
|
||
echo " export CPUINFER_CPU_INSTRUCT=AVX512 # or AVX2/FANCY"
|
||
echo " export CPUINFER_ENABLE_AMX=OFF"
|
||
echo " ./install.sh build --manual"
|
||
else
|
||
echo "ℹ AMX instructions not detected"
|
||
export CPUINFER_CPU_INSTRUCT=NATIVE
|
||
export CPUINFER_ENABLE_AMX=OFF
|
||
echo " Configuration: NATIVE + AMX=OFF"
|
||
fi
|
||
|
||
echo ""
|
||
echo "To use manual configuration instead, run: $0 build --manual"
|
||
echo ""
|
||
else
|
||
# Manual mode - validate user configuration (no exports)
|
||
if [ -z "$CPUINFER_CPU_INSTRUCT" ] || [ -z "$CPUINFER_ENABLE_AMX" ]; then
|
||
echo "Error: Manual mode requires CPUINFER_CPU_INSTRUCT and CPUINFER_ENABLE_AMX to be set."
|
||
echo ""
|
||
usage
|
||
fi
|
||
|
||
# Validate CPUINFER_CPU_INSTRUCT
|
||
case "$CPUINFER_CPU_INSTRUCT" in
|
||
NATIVE|FANCY|AVX512|AVX2)
|
||
;;
|
||
*)
|
||
echo "Error: Invalid CPUINFER_CPU_INSTRUCT='$CPUINFER_CPU_INSTRUCT'"
|
||
echo "Must be one of: NATIVE, FANCY, AVX512, AVX2"
|
||
exit 1
|
||
;;
|
||
esac
|
||
|
||
# Validate CPUINFER_ENABLE_AMX
|
||
case "$CPUINFER_ENABLE_AMX" in
|
||
ON|OFF)
|
||
;;
|
||
*)
|
||
echo "Error: Invalid CPUINFER_ENABLE_AMX='$CPUINFER_ENABLE_AMX'"
|
||
echo "Must be either: ON or OFF"
|
||
exit 1
|
||
;;
|
||
esac
|
||
|
||
# Warn about problematic configuration
|
||
if [ "$CPUINFER_CPU_INSTRUCT" = "NATIVE" ] && [ "$CPUINFER_ENABLE_AMX" = "OFF" ]; then
|
||
HAS_AMX=$(detect_cpu_features)
|
||
if [ "$HAS_AMX" = "1" ]; then
|
||
echo "⚠️ WARNING: NATIVE + AMX=OFF on AMX-capable CPU may cause compilation issues!"
|
||
echo " Recommended: Use AVX512 or AVX2 instead of NATIVE when AMX=OFF"
|
||
echo ""
|
||
read -p "Continue anyway? (y/N) " -n 1 -r
|
||
echo
|
||
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
|
||
exit 1
|
||
fi
|
||
fi
|
||
fi
|
||
|
||
# Close MANUAL_MODE conditional
|
||
fi
|
||
|
||
# Set defaults for optional variables
|
||
export CPUINFER_BUILD_TYPE=${CPUINFER_BUILD_TYPE:-Release}
|
||
export CPUINFER_PARALLEL=${CPUINFER_PARALLEL:-8}
|
||
export CPUINFER_VERBOSE=${CPUINFER_VERBOSE:-1}
|
||
|
||
echo "Building kt-kernel with configuration:"
|
||
echo " CPUINFER_CPU_INSTRUCT=$CPUINFER_CPU_INSTRUCT"
|
||
echo " CPUINFER_ENABLE_AMX=$CPUINFER_ENABLE_AMX"
|
||
echo " CPUINFER_BUILD_TYPE=$CPUINFER_BUILD_TYPE"
|
||
echo " CPUINFER_PARALLEL=$CPUINFER_PARALLEL"
|
||
echo " CPUINFER_VERBOSE=$CPUINFER_VERBOSE"
|
||
echo ""
|
||
|
||
pip install . -v
|
||
}
|
||
|
||
# Subcommand dispatcher: default to "all"
|
||
SUBCMD="all"
|
||
if [[ $# -gt 0 ]]; then
|
||
case "$1" in
|
||
deps|build|all) SUBCMD="$1"; shift ;;
|
||
-h|--help) usage ;;
|
||
*) SUBCMD="build" ;; # backward compatibility: flags-only => build
|
||
esac
|
||
fi
|
||
|
||
case "$SUBCMD" in
|
||
deps)
|
||
install_dependencies
|
||
;;
|
||
build)
|
||
build_step "$@"
|
||
;;
|
||
all)
|
||
install_dependencies
|
||
build_step "$@"
|
||
;;
|
||
esac
|