没合适的组合wheel包，就自行编译flash_attn吧

applebomb2025-07-09 13:48

不知道是不是也有很多小伙伴像我一样，经常一装机就把最新的cuda安装上去了。等装到组件时，才发现cuda和其它组件各种不适配。这次把nvcc直接装到12.8了，然后发现flash attn没有合适的wheel，就只能自己编译了。下面记录下我的编译过程：

我的环境是ubuntu 24.04, ubuntu 22下也是一样的。

设置环境：

export CUDA_HOME=/usr/local/cuda-12.8 # 替换 <version> 为你的CUDA版本号

export PATH= $CUDA_HOME/bin:$ PATH

export CPATH= $CUDA_HOME/include:$ CPATH # 对于GCC, Clang等C/C++编译器

export LD_LIBRARY_PATH= $CUDA_HOME/lib64:$ LD_LIBRARY_PATH

安装编译器：

pip install ninja

安装编译工具链：

apt-get update && sudo apt-get install -y nvidia-cuda-toolkit

pip install --upgrade pip setuptools wheel

下载 flash atten源码：

准备编译：

cd flash-attention

python setup.py bdist_wheel

开始编译：

MAX_JOBS=1 python setup.py install

上面为的MAX_JOBS根据你的内存来设置，我的32G内存，默认3个JOB会OOM，被kill导致编译不过。所以就设置了1个并发编译