FlashAttention 安装学习笔记

在安装和编译 FlashAttention 的过程中，遇到如下问题，进行记录，以防后面再次踩坑

CUDA 环境变量设置

问题描述

在编译 flash-attn 的时候，系统找不到 CUDA 相关头文件或编译器，表现为：

• nvcc not found
• cannot find cuda.h

排查过程

• 运行 ls /usr/local/cuda/include 确认是否存在 cusparse.h、cusparse_v2.h 等头文件
• 执行 which nvcc 或 ps -ef | grep nvcc 检查 nvcc 是否在路径中

解决办法

需要在 环境变量 中显式指定 CUDA 路径：

export CUDA_HOME=/usr/local/cuda
export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH

这样 nvcc 就能被找到，头文件和库路径也能被正常识别

CUDA 运行时与 PyTorch 编译版本不对应

使用conda安装pytorch及其几个附属包的指令：conda install pytorch torchvision torchaudio cpuonly -c pytorch

问题描述

• nvcc --version 显示 CUDA 12.4
• nvidia-smi 显示驱动支持 CUDA 12.2
• conda 安装的 torch 版本默认为 +cu124，导致和本地驱动版本对不上

排查过程

• 使用 nvcc --version 和 nvidia-smi 对比 CUDA 版本
• 检查pytorch是否正确安装以及cuda是否可用：进入python命令行，import torch, torch.cuda.is_available()，查看返回是否为True

解决办法

如果上面排查的返回为True，就算nvcc --version返回的cuda支持版本和系统实际的cuda版本不一致，那也没问题，只要能正常调用就可以

找不到 pybind11 文件

问题描述

编译时出现：

fatal error: pybind11/pybind11.h: No such file or directory

排查过程

• 检查 conda list | grep pybind11，发现并没有安装。

解决办法

安装 pybind11：

pip install pybind11

额外依赖：packaging 和 ninja

在编译过程中还需要两个辅助包：

1. packaging：用于 setup.py 判断环境版本

   pip install packaging

2. ninja：更快的编译工具，比 make 高效

   pip install ninja

如果不提前安装，编译 FlashAttention 时会报错提示找不到对应模块。

安装验证

官方的Git仓库中包含了示例代码，在benchmark文件夹中，在flash-attn编译安装完成后，测试是否安装成功，可以运行 benchmark/benchmark_flash_attn.py，结果显示：

• FlashAttention 前向 + 反向比 PyTorch 标准 Attention 提速 8-9 倍
• 能够正常运行并输出 benchmark 结果 → 安装成功