先放结论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 编译时真正用到的主要是这几样hipccHIP 编译器入口负责把.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组重新打开 WSLsudo 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.3TritonSageAttention 部分内核的编译和调度使用 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 [archamd64 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_ARCHgfx1201ROCM_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_sizeimport 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, devicecuda, dtypetorch.float16) k torch.randn(batch, heads, seq_len, head_dim, devicecuda, dtypetorch.float16) v torch.randn(batch, heads, seq_len, head_dim, devicecuda, dtypetorch.float16) def bench(fn, q, k, v, warmup10, iters50): # 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(fSDPA median: {sdpa_time * 1000:.2f} ms) print(fSageAttention median: {sage_time * 1000:.2f} ms) print(fSpeedup: {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 环境下用上面这个脚本测试了不同序列长度。结论是这样的序列长度 4096batch size 116 头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(fMax 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 foundHIP 头文件路径没设置或没安装确认HIP_PATH指向包含include/hip的目录gcnArch not supported或invalid device function目标架构没指定或指定错误检查PYTORCH_ROCM_ARCH9070XT 一般设为gfx1201librocblas.so: cannot open shared object file运行库路径不完整export LD_LIBRARY_PATH$ROCM_HOME/lib:$LD_LIBRARY_PATHTriton is not installed or version too lowPyTorch 自带 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 不能跑而是版本之间没有对齐。