windows下安装大模型加速包flash-attn

大家好，欢迎来到IT知识分享网。

Flash Attention

Flash-attn（Flash Attention），作为一种用于神经网络模型的注意力机制，在处理序列数据时有显著的优势，可以提高计算效率并减少内存使用，特别是Transformer模型。它通过优化数据布局和计算流程，降低了注意力计算的内存访问开销，在本地部署大模型可以减轻GPU的开销，加快推理速度。项目地址 ：
https://github.com/Dao-AILab/flash-attention

然而，官方发布的版本只支持linux，没有windows包，如果要在windows下安装，要么使用docker,要么自己编译，而自己编译耗时很长，而且也会遇到各种问题。这里我使用的是Brian Dashore编译好的whl. 下载地址：
https://github.com/bdashore3/flash-attention/releases

要注意选择对应的版本，首先查看自己的cuda版本，在CMD 里输入：nvcc –version

查看cuda版本

我的本地环境如下：

系统 windows10 ,python版本11,cuda版本12.1

下载对应的版本为：flash_attn-2.5.9.post1+
cu122torch2.3.1cxx11abiFALSE-cp311-cp311-win_amd64.whl

flash-attn版本

下载到本地后执行，pip安装：

pip install d:\downloads\flash_attn-2.5.9.post1+cu122torch2.3.1cxx11abiFALSE-cp311-cp311-win_amd64.whl

注意：安装完成后，如果你以前已经安装过PyTorch，需要重新安装一遍，否则运行会报错。

打开网址
https://pytorch.org/get-started/locally/ 选择对应的版本

安装pytorch

根据pytorch.org给出的提示命令，在本地执行安装：

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

安装完成，写一个简单的脚本检验一下。

test.py

import flash_attn print(flash_attn.__version__)

运行结果如果正确输出版本号，则表示安装成功。

检验flash_attn版本