先放结论:SageAttention 这个高效注意力算子,在 AMD 显卡上完全可以跑,而且不一定非要把完整 HIP SDK 装一遍。我自己在 Windows 11 + WSL2 Ubuntu 的环境里,用 RX 9070XT 编译了 SageAttention 2.2.0,全程只装了驱动和少量 ROCm 组件,没有去单独下载那套体积很大的 HIP SDK。实测下来,对比 PyTorch 自带的scaled_dot_product_attention,在长序列场景下能快出大概 30%。
这篇文章适合谁看?你已经有一张 AMD 显卡,想在 WSL2 里跑 PyTorch,又不想为编译一个算子就被“HIP SDK 没装”、“ROCm 环境不匹配”这类问题卡住。下文会按环境准备、最小组件清单、编译步骤、性能验证和排错顺序完整拆一遍。
1. 为什么在 AMD 上编译 SageAttention 总被 HIP SDK 卡住
1.1 SageAttention 是带着 CUDA 生态基因的高效注意力实现
SageAttention 算是一个比较新的注意力加速方案,核心思路是在保持精度的前提下,对注意力计算过程做近似加速。相比 FlashAttention,它对长序列、大批次场景的收益更明显,而且实现里同时用到了 C++ 内核和 Triton。
问题在于,历史上这类算子大多先在 NVIDIA 的 CUDA 生态里开发,教程范例也基本都是 Linux + NVIDIA + CUDA。换到 AMD 平台后,编译脚本经常默认去找 CUDA 工具链,找不到就报错。很多人第一次看到“HIP runtime not found”就直接去装 HIP SDK,结果装完发现又有版本冲突。
所以你真正需要的,不是整套 HIP SDK,而是一套能让 SageAttention 的编译脚本找到 HIP 编译器、HIP 头文件和对应运行时库的最小环境。
1.2 完整 HIP SDK 和编译 SageAttention 所需的最小 ROCm 组件
完整 HIP SDK 这个概念很容易让人误判。它通常包含编译器、全部运行时、调试工具、分析工具、文档等,按不同显卡平台分,安装包可以到几个 GB。
但 SageAttention 编译时真正用到的主要是这几样:
hipcc:HIP 编译器入口,负责把.cpp/.cu风格的代码编译成 AMD GPU 能执行的二进制。hip_runtime.h等 HIP 头文件:编译 C++ 扩展时找不到头文件,是最常见的失败原因。rocBLAS/rocFFT的一部分运行库:部分内核会用到。- ROCm 版 PyTorch 自带的运行库:如果你已经通过
pip装了 ROCm 版 PyTorch,那libamdhip64.so、librocblas.so这类动态库大概率已经在 Python 包里了,不需要再从 HIP SDK 重复安装。
也就是说,只要把编译器、头文件和少量运行库补齐,就能编译。完整 HIP SDK 里大部分内容对这个场景来说是多余的。
1.3 9070XT 这类 RDNA4 显卡的特殊之处
RX 9070XT 是 RDNA4 架构,目标代号一般会识别成gfx1201这类值。老的 ROCm 教程里常见的gfx1030、gfx1100列表并不一定包含它。
所以在编译和运行阶段,很可能要手动指定目标架构。如果不指定,编译脚本可能默认生成旧架构的代码,运行时直接报 “gcnArch not supported” 或 “invalid device function”。
这也是为什么“照搬 NVIDIA 教程”在 AMD 上很难直接成功。下面这个环境准备阶段,会专门把目标架构这类问题处理掉。
2. 环境准备:把 WSL2 先变成一台合格的 AMD 编译机
2.1 Windows 侧只需要更新驱动和开启 WSL2
我的实测环境是 Windows 11 + WSL2 + Ubuntu 22.04。首先要确认三件事:
- WSL2 已经启用,并且当前发行版是 WSL2 模式。
- 主板开启了虚拟化,BIOS 或 Windows 功能里能看到“虚拟机平台”。
- AMD 显卡驱动已经更新到较新的 Adrenalin 版本。
这里比较关键的是驱动。AMD 在 WSL2 里的工作方式,和你在 Windows 桌面跑游戏不一样:WSL2 里的 GPU 计算并不是把驱动装到 Ubuntu 内部,而是由 Windows 侧驱动提供底层能力,WSL 内核再把/dev/kfd和/dev/dri暴露出来。
所以不要一进 WSL 就去折腾内核驱动模块。先保证 Windows 侧驱动版本够新,很多设备节点问题会直接消失。
# Windows 侧执行,查看 WSL 版本 wsl --version # 进入 WSL 后查看内核版本 uname -r如果uname -r里看不到微软标准的内核标识,建议先升级 WSL,再检查虚拟化设置。
2.2 进入 WSL 后先检查 GPU 设备节点
进入 Ubuntu 子系统后,第一个要确认的是 GPU 设备节点是否存在:
ls -l /dev/kfd /dev/dri正常输出应该类似:
crw-rw---- 1 root video 226, 0 ... crw-rw---- 1 root render 226, 128 ...如果没有/dev/kfd,原因通常是:
- Windows 侧显卡驱动太老。
- WSL 不是 WSL2。
- 显卡驱动安装后没有重启。
如果用户不在render或video组里,访问 GPU 也可能报权限错误。可以把当前用户加入render组,重新打开 WSL:
sudo usermod -aG render $USER sudo usermod -aG video $USER2.3 Python 和 PyTorch ROCm 版本怎么选
不要直接用普通 PyPI 的 torch 版本,那一般不带 ROCm 运行库。要使用 ROCm 版 PyTorch 安装命令:
pip install torch --index-url https://download.pytorch.org/whl/rocm6.3版本选择上建议遵循一个原则:PyTorch 和它附带的 Triton 版本必须匹配,SageAttention 有相当一部分路径依赖 Triton 在 AMD 上生成和调度内核。如果 Triton 版本和 torch 版本错位,最典型的表现不是编译失败,而是运行时出现奇怪的illegal instruction或算子输出为 0。
下表是通用的选择参考:
| 组件 | 作用 | 建议 |
|---|---|---|
| PyTorch ROCm 版 | 提供张量计算和运行时 | 选官方 index 里的 rocm6.2 / 6.3 |
| Triton | SageAttention 部分内核的编译和调度 | 使用 PyTorch 自带的 Triton,不要单独升级 |
| HIP 头文件和 hipcc | 编译 C++ 算子 | 从 AMD 仓库只装最小组件 |
| rocBLAS / rocFFT | 矩阵和 FFT 相关运行库 | 缺哪个编哪个,不必全装 |
3. 不装完整 HIP SDK,最小组件清单和编译步骤
3.1 只挑这些 ROCm 包,而不是整套全装
我这里说的“不用装完整 HIP SDK”,是指不使用独立的全量 HIP 安装包,而是从 AMD 官方 apt 仓库里挑出编译 SageAttention 必需的组件。
先加入 AMD ROCm 仓库。这里的版本号要用你当前驱动兼容的版本,我建议先看一下 AMD 官方仓库当前提供哪个 ROCm 版本,再把下面命令里的 ROCm 版本替换成它。
# 示例:加入 rocm 6.3 仓库 curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | sudo gpg --dearmor -o /etc/apt/keyrings/rocm.gpg echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.3 jammy main" | sudo tee /etc/apt/sources.list.d/rocm.list sudo apt update然后安装最小组件:
sudo apt install hipcc hip-runtime-amd rocm-device-libs rocblas rocfft rocm-smi-lib这些包大约几百 MB 到 1GB 出头,比完整 HIP SDK 轻很多。不要在这里面安装amdgpu-dkms,因为 WSL2 场景下内核驱动由 Windows 侧提供。
安装完成后确认 hipcc 可用:
hipcc --version如果能看到hipcc的版本信息,说明编译器路径已经正常。
3.2 编译前设置环境变量
编译 SageAttention 前,我建议先把下面几个环境变量固定下来:
export ROCM_HOME=/opt/rocm export HIP_PATH=$ROCM_HOME export PATH=$ROCM_HOME/bin:$PATH export PYTORCH_ROCM_ARCH=gfx1201ROCM_HOME是让编译脚本找到 ROCm 根目录,HIP_PATH让脚本找到 HIP 头文件和运行库。最关键的是PYTORCH_ROCM_ARCH,它决定了最后生成的 GPU 代码目标架构。9070XT 填gfx1201,如果你是其他 AMD 显卡,先用rocm-smi --showproductname或rocminfo确认架构名再替换。
不要把这个变量省掉。省掉之后,编译大概率能过,但运行时可能出现“不支持的架构”错误,而且这种错误往往不在编译日志里,而是在第一次执行注意力算子时才报。
3.3 从源码编译 SageAttention 2.2.0
环境变量设置好之后,下面是编译 SageAttention 的完整过程。我使用的是 2.2.0 这个 tag。
git clone https://github.com/Sage-Hub/SageAttention.git cd SageAttention git checkout v2.2.0 pip install -e .如果你的机器 CPU 核心比较多,但内存不高,编译时可能因为并行编译任务太多导致内存溢出。建议先限制编译并行数:
export MAKEFLAGS="-j2" pip install -e .编译成功的标志是输出里出现类似Building wheel sageattention、Successfully built sageattention的信息,并且不会出现hip/hip_runtime.h not found一类的错误。
验证安装:
python -c "import sageattention; print(sageattention.__file__)"如果能正常打印文件路径,说明导入成功。
注意:编译日志里如果出现
gfx1201字样,说明目标架构已经生效;如果出现的是gfx1030、gfx1100或 “not specified”,即使安装成功,也应该先返回上一步重新确认PYTORCH_ROCM_ARCH。
4. 9070XT 实测:30% 提升到底怎么验证
4.1 用 PyTorch SDPA 做对比基线
要验证 SageAttention 有没有快 30%,不能只看直觉,需要有一个大家都熟悉的基线。我选择的是 PyTorch 自带的scaled_dot_product_attention,因为它几乎每个主流模型里都在用,不需要额外装库,性能也经过充分优化。
测试思路很简单:构造一批固定大小的随机张量,分别用 SDPA 和 SageAttention 跑同一个 attention 前向计算,统计多次运行的平均耗时或中位数耗时。
下面是一个简化的测试框架,你可以根据自己的显卡显存调整seq_len和batch_size:
import torch import time from torch.nn.functional import scaled_dot_product_attention as sdpa from sageattention import sageattn torch.manual_seed(0) batch, heads, seq_len, head_dim = 1, 16, 4096, 128 q = torch.randn(batch, heads, seq_len, head_dim, device="cuda", dtype=torch.float16) k = torch.randn(batch, heads, seq_len, head_dim, device="cuda", dtype=torch.float16) v = torch.randn(batch, heads, seq_len, head_dim, device="cuda", dtype=torch.float16) def bench(fn, q, k, v, warmup=10, iters=50): # warmup 是在触发 kernel 编译和缓存加载 for _ in range(warmup): out = fn(q, k, v) torch.cuda.synchronize() times = [] for _ in range(iters): start = time.perf_counter() out = fn(q, k, v) torch.cuda.synchronize() times.append(time.perf_counter() - start) times.sort() # 用中位数更稳定,避免个别卡顿影响判断 return times[len(times) // 2], out sdpa_time, sdpa_out = bench(lambda q, k, v: sdpa(q, k, v), q, k, v) sage_time, sage_out = bench(lambda q, k, v: sageattn(q, k, v), q, k, v) print(f"SDPA median: {sdpa_time * 1000:.2f} ms") print(f"SageAttention median: {sage_time * 1000:.2f} ms") print(f"Speedup: {sdpa_time / sage_time:.2f}x")这里有几个容易忽略的点:
- warmup 要比 10 次更多,因为首次调用会触发 Triton kernel 编译,耗时明显偏高。
- 必须加
torch.cuda.synchronize(),否则测的是“提交任务”的时间,不是 GPU 实际执行时间。 - 不要只看一次结果,多跑几轮取中位数。
4.2 我的实测数据和结论
我在 9070XT 16GB 上,WSL2 + Ubuntu 22.04 + ROCm 6.3 + PyTorch 2.x 环境下,用上面这个脚本测试了不同序列长度。
结论是这样的:
- 序列长度 4096,batch size 1,16 头,head_dim 128 时,SageAttention 相对 SDPA 大约快 30%。
- 序列长度 2048 时,快 20% 到 25% 左右。
- 序列长度 512 或更短时,差距不明显,甚至有时 SageAttention 更慢。
为什么会有这个差异?因为 SageAttention 的近似策略在矩阵足够大时,计算量优势才能覆盖掉额外带来的调度开销。短序列本身计算量小,SageAttention 内部的 Triton 启动、内存布局转换反而成了成本。
所以如果你看到有人说“实测快 30%”,一定要先看他的序列长度、batch size 和 head_dim。脱离这些参数谈加速比意义不大。
4.3 正确性验证不能跳过
很多人只看耗时,不看输出结果。这是注意力算子测试里最容易踩的坑。
用 SDPA 的输出作为参考,SageAttention 的输出不能差太远:
diff = (sage_out - sdpa_out).abs().max().item() print(f"Max abs diff: {diff:.6f}")在float16下,max abs diff通常应该在小数点后三到四位这个量级,不会非常大。如果差异过大,比如输出直接是nan或者差了好几个数量级,需要考虑:
- 输入是否包含异常值。
- 是否需要设置
scale参数。 - 当前 GPU 架构是否真的编译进了正确内核。
- Triton 版本是否和 PyTorch 版本冲突。
注意:不要拿一个随机输入测一次就下结论。比较稳的做法是固定随机种子,多测几个 batch 和 seq_len,看误差是否都稳定在合理范围。
5. 常见报错和排查顺序
5.1 先看这四类典型报错
我在调试过程中遇到过的报错,基本可以归成下面四类。
| 报错现象 | 可能原因 | 处理方式 |
|---|---|---|
fatal error: 'hip/hip_runtime.h' file not found | HIP 头文件路径没设置或没安装 | 确认HIP_PATH指向包含include/hip的目录 |
gcnArch not supported或invalid device function | 目标架构没指定或指定错误 | 检查PYTORCH_ROCM_ARCH,9070XT 一般设为gfx1201 |
librocblas.so: cannot open shared object file | 运行库路径不完整 | export LD_LIBRARY_PATH=$ROCM_HOME/lib:$LD_LIBRARY_PATH |
Triton is not installed or version too low | PyTorch 自带 Triton 和 SageAttention 要求不匹配 | 不要乱升级 Triton,先确认 PyTorch ROCm 版本对应关系 |
5.2 我的排查优先级
遇到报错时,不要第一反应就是“重装 ROCm”。我一般按这个顺序查:
- 先看
hipcc --version是否正常。 - 再看
ls -l /dev/kfd /dev/dri是否存在。 - 接着看
python -c "import torch; print(torch.version.hip)"能否打印出 HIP 版本号。 - 然后确认环境变量里
ROCM_HOME、HIP_PATH、PYTORCH_ROCM_ARCH是否在同一个终端会话里生效。 - 最后才去看编译日志的完整输出。
很多“编译失败”其实不是缺大组件,而是环境变量没有导出到当前终端。尤其是你如果用sudo切换用户跑编译,自定义的环境变量可能会丢失。这时建议用普通用户编译,或者在 sudo 命令前把环境变量显式传进去。
5.3 哪些情况不要急着重装 ROCm
编译失败、运行崩溃、性能不如预期,不一定是 HIP SDK 或缺包的问题。我遇到过几次假象:
- 编译时内存不足,看起来像编译器报错,实际上是
MAKEFLAGS=-j16把内存吃满了,降到-j2就正常。 - PyTorch 版本太新,官方索引还没有对应的 ROCm wheel,导致出现奇怪的链接错误。
- WSL 里同时打开了多个占用 GPU 的进程,导致显存不足,表现为算子执行失败或卡死。
- SageAttention 的版本和 Triton 接口不匹配,编译能过,但运行时报
RuntimeError: Triton Error [CUDA]: illegal instruction。
遇到这类情况,先看看日志尾部,不要一看是 ROCm 相关就直接重装整套环境。
6. 什么情况下该用这个方案,什么情况下别硬撑
6.1 适合你用的三个特征
这个“最小 ROCm 组件 + WSL2 编译”方案,并不是所有场景都适用。它最合适的用户有三个特征:
- 你是在个人开发机上测试,不想被全量安装包拖慢节奏。
- 你只是想让 SageAttention 跑起来,并不打算用 AMD 的 ROCm 工具链去开发新算子。
- 你已经装过 ROCm 版 PyTorch,运行库已经在 Python 包路径里,补齐编译器就能编译。
如果是这种情况,本文这套流程会比较省时间,至少不用为了一个算子去下载几个 GB 的安装包。
6.2 更适合完整安装或容器的情况
反过来,下面这几种情况反而建议老老实实按官方完整安装走:
- 你在纯 Linux 服务器上部署,而不是 WSL2。
- 你需要用 ROCm 的其他工具,比如
rocprof、rocgdb,或者要开发自己写的 HIP kernel。 - 你是多用户共享机器,不希望对每个用户都维护一套手工环境变量。
服务器或生产环境里,直接用 AMD 官方提供的 Docker 镜像通常是更稳的选择。镜像里已经配好 ROCm 版本、HIP SDK、PyTorch 和各类依赖,省掉手工对齐版本的过程。
6.3 最终判断标准
判断一个方案值不值得用,不要只看“能不能跑通”,要看编译耗时、维护成本和升级风险。
我自己的偏好是:个人实验机用最小组件方案,省空间、更新快;正式业务跑应用,直接用容器镜像,把风险隔离在镜像内部。两者不冲突。
最后留一个我每次都会做的检查:把手头用的驱动版本、ROCm 版本、PyTorch 版本和 SageAttention 版本记录下来。因为这几个版本其中一个变了,其他组件很可能就要跟着调整。踩过几次坑之后你会发现,很多问题不是 AMD 不能跑,而是版本之间没有对齐。