Windows 11 下再次成功本地编译 Flash-Attention 2.8.3 并生成自定义 Wheel(RTX 3090 sm_86 专属版)

Windows 11 下成功本地编译 Flash-Attention 2.8.3 并生成自定义 Wheel(RTX 3090 sm_86 专属版)

发布时间:2025年12月29日

Flash Attention 2.6.3 在 Windows 上编译成功复盘笔记
Windows 下成功编译 Flash Attention 2.8.3 (flash-attn /flash_attn)个人复盘记录
Windows 11 下 Z-Image-Turbo 完整部署与 Flash Attention 2.8.3 本地编译复盘

大家好,今天再次试练并分享一个在 Windows 11 + RTX 3090 + PyTorch 2.9.1+cu130 环境下成功本地编译 Flash-Attention 2.8.3 并生成 wheel 文件的完整过程。

官方 Dao-AILab/flash-attention 到 2025 年 12 月底最新版本仍是 v2.8.3(2025 年 8 月发布),没有官方 Windows wheel。Windows 编译仍是实验性,但社区有很多成功案例,尤其在 RTX 30/40 系列上。

我之前尝试过社区预编译 wheel,但为了完美匹配自己的环境(避免任何潜在的 kernel 不匹配),最终选择本地源码编译,并指定 sm_86 架构,生成专属 wheel 备份。

最终结果:编译成功!生成的 wheel 支持 RTX 3090 的所有 Flash-Attention 功能(head_dim 64/128/256 前向+反向,fp16/bf16 等),测试通过,且可以永久保存备用。

我的环境

  • OS:Windows 11 专业工作站版
  • GPU:NVIDIA RTX 3090 (Compute Capability 8.6, 24GB VRAM)
  • CPU:Intel Core Ultra 9 285K
  • RAM:128 GB
  • Python:3.10.18 (本地.venv 虚拟环境)
  • PyTorch:2.9.1+cu130
  • CUDA Toolkit:13.1 (完整安装)
  • Visual Studio:2022 Professional (带 C++ 桌面开发工作负载)
  • ninja:已安装 (pip install ninja)

编译前准备

Dao-AILab/flash-attention:快速且内存高效的精确注意力

  1. 克隆源码并更新子模块(必须 git clone,不能下载 zip)

    复制代码
    git clone https://github.com/Dao-AILab/flash-attention.git
    cd flash-attention
    git pull
    git submodule update --init --recursive
  2. 切换到 main 分支(可选,但推荐获取最新小修复)

    复制代码
    git checkout main

    (如果 git pull 失败,是因为网络问题,代码已是最新可跳过)

关键环境变量设置(在 CMD 中设置,确保生成 sm_86 kernel)

使用 Developer Command Prompt for VS 2022(以管理员运行)打开 CMD,然后:

安装必要工具

复制代码
pip install ninja wheel build packaging

set FLASH_ATTN_CUDA_ARCHS=86          :: 只针对 RTX 3090,编译更快、文件更小
:: 如果想兼容未来 4090,可设 set FLASH_ATTN_CUDA_ARCHS=86;89

set MAX_JOBS=8                        :: 并行编译数,根据 RAM 调整(我 128GB RAM 可高设)
set TORCH_CUDA_ARCH_LIST=8.6          :: 额外保险,确保 Torch 识别架构
set NVCC_THREADS=2                    :: 限制 nvcc 线程,避免卡死
set DISTUTILS_USE_SDK=1               :: Windows 编译必要
set CUDA_HOME=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.1  :: 你的 CUDA Toolkit 路径(可选,可根据实际进行调整)

编译并生成 wheel

复制代码
python -m build --wheel --no-isolation
  • 这个命令会自动使用 ninja 加速。
  • 编译过程会输出大量日志,包括 nvcc 编译 .cu 文件。
  • 此次我的机器编译时间约 20-40 分钟(单架构 86 更快)。
  • 日志中会有很多 warning(如 #221-D floating-point value does not fit),这些是正常现象(softmax 和 mask 中的大负数处理),不影响功能。
  • 还有 MSVC 预处理器警告,也可忽略。

成功标志

  • 最后输出 Successfully built flash_attn-2.8.3-cp310-cp310-win_amd64.whl
  • wheel 文件位于 dist/ 目录下,文件名会带你的环境标识(如 +cu130torch2.9)

备份 wheel

复制代码
copy dist\flash_attn-*.whl E:\Downloads\Other\

以后在新环境直接:

复制代码
pip install E:\Downloads\Other\flash_attn-xxx.whl

验证安装

卸载旧版后安装新 wheel:

复制代码
pip uninstall flash-attn -y
pip install dist\flash_attn-*.whl

然后运行测试:

输入 Python 进入 python 环境

复制代码
import torch
from flash_attn import flash_attn_func

# head_dim=64
q = torch.randn(2, 32, 1024, 64, device='cuda', dtype=torch.float16)
k = torch.randn(2, 32, 1024, 64, device='cuda', dtype=torch.float16)
v = torch.randn(2, 32, 1024, 64, device='cuda', dtype=torch.float16)
out = flash_attn_func(q, k, v, dropout_p=0.0, causal=False)
print(out.shape)  # 成功

# head_dim=128 / 256 同理,全通过

注意事项 & 经验总结

  • 必须用 Developer Command Prompt,否则链接错误。
  • 子模块更新不能忘,否则编译失败。
  • 单架构 (86) 编译更快,wheel 更小;多架构兼容性更好。
  • RTX 3090 上 v2.8.3 已支持所有功能(包括 head_dim=256 backward,但 dropout>0 时会 fallback,这是消费级卡硬件限制)。
  • 如果网络问题无法 git pull,直接用现有代码编译即可(main 已是最新的)。
  • 社区还有预编译 wheel(如 sunsetcoder 或 lldacing 的 HF 仓库),但本地编译最保险,完美匹配环境。

这个自定义 wheel 就是我的"黄金版本"备份,以后重装环境一键恢复,彻底告别 kernel 不匹配的烦恼!

如果你们也在 Windows 下折腾 Flash-Attention,欢迎交流~有问题可以留言,我会尽量解答。

相关推荐
模型时代5 分钟前
Anthropic明确拒绝在Claude中加入广告功能
人工智能·microsoft
夕小瑶8 分钟前
OpenClaw、Moltbook爆火,算力如何48小时内扩到1900张卡
人工智能
一枕眠秋雨>o<11 分钟前
透视算力:cann-tools如何让AI性能调优从玄学走向科学
人工智能
wdfk_prog19 分钟前
[Linux]学习笔记系列 -- [drivers][I2C]I2C
linux·笔记·学习
那个村的李富贵24 分钟前
昇腾CANN跨行业实战:五大新领域AI落地案例深度解析
人工智能·aigc·cann
猫头虎27 分钟前
如何解决 OpenClaw “Pairing required” 报错:两种官方解决方案详解
网络·windows·网络协议·macos·智能路由器·pip·scipy
集简云-软件连接神器27 分钟前
技术实战:集简云语聚AI实现小红书私信接入AI大模型全流程解析
人工智能·小红书·ai客服
松☆27 分钟前
深入理解CANN:面向AI加速的异构计算架构
人工智能·架构
rainbow72424428 分钟前
无基础学AI的入门核心,从基础工具和理论开始学
人工智能
子榆.32 分钟前
CANN 与主流 AI 框架集成:从 PyTorch/TensorFlow 到高效推理的无缝迁移指南
人工智能·pytorch·tensorflow