install_dcu.sh

#!/bin/bash
set -e
# default backend
DEV="cuda"

# parse --dev argument
while [[ "$#" -gt 0 ]]; do
    case $1 in
        --dev) DEV="$2"; shift ;;
        *) echo "Unknown parameter passed: $1"; exit 1 ;;
    esac
    shift
done
export DEV_BACKEND="$DEV"
echo "Selected backend: $DEV_BACKEND"

# 清理构建目录和旧的分发文件
rm -rf build
rm -rf dist
rm -rf *.egg-info
rm -rf ktransformers/ktransformers_ext/build
rm -rf ktransformers/ktransformers_ext/cuda/build
rm -rf ktransformers/ktransformers_ext/cuda/dist
rm -rf ktransformers/ktransformers_ext/cuda/*.egg-info
echo "初始化Git子模块..."
git submodule update --init --recursive

export CMAKE_BUILD_PARALLEL_LEVEL=32

echo "构建ktransformers wheel包"

mkdir -p dist

KTRANSFORMERS_FORCE_BUILD=TRUE pip wheel . -w dist --no-build-isolation --no-deps
if [[ "$DEV_BACKEND" == "cuda" ]]; then
    echo "Installing custom_flashinfer for CUDA backend"
    pip install third_party/custom_flashinfer/
fi

echo "生成的wheel包位于："
ls -l dist/*.whl

echo "构建成功！wheel包已生成在dist目录"