Flash Attention 注意力机制python包的安装

一、简介：

flash_attn 是一种专为神经网络中的注意力机制（Attention Mechanism）优化的库，旨在减少显存使用并提升训练和推理速度。传统的注意力机制计算会在大型矩阵操作中消耗大量内存和计算资源，而 flash_attn 通过改进计算流程，使其能够以更少的内存开销完成注意力计算。

在安装flash attention包中，经常需要提前安装CUTLASS包 (CUDA Templates for Linear Algebra Subroutines and Solvers)，他们都是深度学习框架（如 PyTorch 和 TensorFlow）的底层加速模块。

以下安装均在ubuntu上运行

二、安装 CUTLASS 库

1.克隆 CUTLASS 源代码：

git clone https://github.com/NVIDIA/cutlass.git
cd cutlass

（如果下载困难，可在作者主页免费下载）

2.安装依赖项

确保已经安装了 CUDA Toolkit 和 CMake（3.12 或更高版本）。可以通过以下命令安装 CMake

sudo apt update
sudo apt install cmake

3.编译 CUTLASS

（1）创建一个用于构建的目录（如 build）

mkdir build
cd build

（2）使用 CMake 配置项目

cmake .. -DCUTLASS_ENABLE_CUBLAS=ON -DCUTLASS_ENABLE_CUTENSOR=ON

（3）编译 CUTLASS：

make -j$(nproc)

这个命令会使用所有可用的 CPU 内核来加快编译过程。

4.验证编译结果

编译完成后，可以通过运行一些测试来验证 CUTLASS 是否正常工作：

ctest

这将运行所有测试，检查 CUTLASS 的功能。

5.将 CUTLASS 添加到编译路径

如果需要将 CUTLASS 用于其他项目，可以将 CUTLASS 的 include 目录路径添加到环境变量中，以便其他项目在编译时可以找到 CUTLASS。

在 .bashrc 文件中添加以下内容：

export CUTLASS_PATH=/path/to/cutlass
export CPLUS_INCLUDE_PATH=$CUTLASS_PATH/include:$CPLUS_INCLUDE_PATH

然后，刷新环境变量：

source ~/.bashrc

三、安装flash-attn

直接通过pip安装可能会存在超时报错，建议从官网下载源码编译安装

1.克隆 flash-attn ：

git clone https://github.com/Dao-AILab/flash-attention.git
cd flash-attention

（如果下载困难，可在作者主页免费下载）

2.构建 .whl 文件：

pip wheel . 

3.安装 .whl 文件

pip install flash_attn-<version>.whl

将 <version> 替换为实际的版本号，以匹配生成的 .whl 文件名。

例如，如果在构建过程中生成了一个名为 flash_attn-0.2.8-cp39-cp39-linux_x86_64.whl 的文件，那么安装命令应为：

pip install flash_attn-0.2.8-cp39-cp39-linux_x86_64.whl