Ubuntu的uv虚拟环境中安装Flash-attn

1、确保Ubuntu下已经安装了uv和cuda

uv --version

查看uv版本信息

nvcc --version

查看cuda版本信息

2、创建并激活虚拟环境

uv venv

source .venv/bin/activate

3、安装依赖库

uv pip install ninja packaging

4、安装CUDA 12.1版本的PyTorch

uv pip install torch==2.5.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

查看PyTorch的CUDA版本:

python -c "import torch; print(f'Torch version: {torch.__version__}'); print(f'CUDA version used by Torch: {torch.version.cuda}')"

5、下载FlashAttention的预编译wheel文件:

https://github.com/Dao-AILab/flash-attention/releases

(https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.5cxx11abiFALSE-cp310-cp310-linux_x86_64.whl)

cuda 12.x版本,torch 2.5版本,python 3.10版本

6、安装.whl文件

uv pip install flash_attn-2.8.3+cu12torch2.5cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

7、验证是否安装成功

python -c "import flash_attn; print(f'Flash-Attn version: {flash_attn.__version__}')"

若输出:

Flash-Attn version: 2.8.3

则表示Flash-Attn 2.8.3版本已经成功安装!

------------------------------------------

FlashAttention 通过利用 GPU 的片上高速缓存（SRAM）并融合计算步骤，显著加速了Transformer模型的训练和推理速度，同时大幅减少了显存占用。