这个命令:
确保使用正确的 CUDA 12.6 工具链
设置必要的 CUDA 环境变量
包含了常见的 GPU 架构支持
利用你的128核心进行并行编译
# 清理之前的安装
proxychains4 pip uninstall -y flash-attn
# 获取 CUDA 路径
CUDA_PATH=$(dirname $(dirname $(which nvcc)))
# 使用 proxychains4 安装
CUDA_HOME=$CUDA_PATH \
TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9;9.0" \
MAKEFLAGS="-j128" \
CMAKE_BUILD_PARALLEL_LEVEL=128 \
CMAKE_GENERATOR="Ninja" \
CFLAGS="-march=native -O3" \
CXXFLAGS="-march=native -O3" \
proxychains4 pip install flash-attn --no-build-isolation
# 如果上面失败,可以尝试这个备选命令(从源码安装)
CUDA_HOME=$CUDA_PATH \
TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9;9.0" \
MAKEFLAGS="-j128" \
CMAKE_BUILD_PARALLEL_LEVEL=128 \
CMAKE_GENERATOR="Ninja" \
CFLAGS="-march=native -O3" \
CXXFLAGS="-march=native -O3" \
proxychains4 pip install git+https://github.com/Dao-AILab/flash-attention.git