Torch and xformers versions for 5090

快速结论：用户在升级到 NVIDIA RTX 5090（Blackwell 架构）后，Stable Diffusion WebUI Forge 无法启动。优先排查 PyTorch 和 xformers 是否支持 sm_120（Blackwell 的 CUDA 算力），并考虑临时禁用 xformers 或从源码编译适配版本。

问题场景

用户将 GPU 从旧型号升级至 NVIDIA RTX 5090，重新安装了 torch、torchaudio、torchvision 和 xformers，但 Stable Diffusion WebUI Forge 仍然无法正常工作。其他用户报告在 ComfyUI、Kohya’s GUI 等工具中也遇到类似问题。

报错原文

torch                        2.7.1+cu128
torchaudio               2.7.1+cu128
torchvision               0.22.1+cu128
xformers                  0.0.32.dev1073
# WebUI Forge still not working after reinstall

原因分析

RTX 5090 采用 Blackwell 架构，其 CUDA 算力为 sm_120。当前 PyTorch 官方预编译的 xformers 二进包通常未包含对 sm_120 的支持，导致 xformers 无法正常加载或运行。此外，Flash Attention 3 在 Blackwell 上可能存在兼容性问题。禁用 xformers 或从源码编译适配版本可以解决。

环境排查

确认显卡型号为 RTX 5090 及 CUDA 算力是否为 sm_120（可通过 nvidia-smi --query-gpu=compute_capability --format=csv 验证）。
确认 CUDA 驱动版本 ≥ 12.8（用户使用的 cu128 轮子对应 CUDA 12.8）。
确认 Python 版本符合 PyTorch nightly 要求（建议 3.10+）。
确认系统已安装 CUDA Toolkit 12.8（包含 nvcc 编译器）。

解决步骤

临时禁用 xformers 验证：运行 WebUI Forge 或 ComfyUI 时添加 --disable-xformers 参数。如果可运行，说明问题出在 xformers 兼容性。
升级 PyTorch 至 nightly 版本：安装或更新至符合 CUDA 12.8 的 nightly 构建，例如用户在评论中使用的：
```
torch 2.9.0.dev20250812+cu128
torchvision 0.24.0.dev20250812+cu128
torchaudio 2.8.0.dev20250812+cu128
```
从源码编译 xformers（可优先尝试）：
a. 卸载现有 xformers：pip uninstall -y xformers
b. 从 GitHub 主线源码安装：pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers
c. 如果编译失败，可尝试改用预编译轮子：pip install xformers==0.0.31.post1 --index-url https://download.pytorch.org/whl/cu128
针对 Kohya’s GUI 等训练工具的完整 Blackwell 环境配置：
a. 为 RTX 5090 创建独立的虚拟环境（如 .venv_blackwell）。
b. 安装 CUDA 12.8 Toolkit 并确保 nvcc 可用。
c. 克隆 xformers 源码及其子模块（cutlass, flash-attention, composable_kernel），设置 TORCH_CUDA_ARCH_LIST=12.0 进行编译。
d. 在 xformers 源码中手动禁用 Flash Attention 3：编辑 dispatch.py，将 _USE_FLASH_ATTENTION_3 设为 False。
e. 修改训练脚本，跳过 VAE 中的 xformers 相关操作（sm_120 不支持的部分）。

验证方法

启动 WebUI Forge 或 ComfyUI，观察是否正常加载模型并执行推理/训练。如果使用 --disable-xformers 可以运行，而启用 xformers 仍旧报错，说明仍需完善编译步骤。可在命令行中运行 python -c "import xformers; print(xformers.__version__)" 确认 xformers 已安装且无导入错误。