最近在 AMD 9070XT 上跑大模型推理的朋友应该都碰到过同一个纠结想用 SageAttention 把注意力部分的耗时压下去结果一看官方文档全是 CUDA 生态的命令又听说 AMD 上编译要装整套 HIP SDK光依赖就能把人劝退。但实际踩完一遍就会发现这里存在一个被夸大很多年的误区编译 SageAttention 并不需要你把 AMD 那套几个 GB 的开发环境装全。更准确地说你只需要保证 HIP 编译器工具链和运行时可用而 PyTorch 的 ROCm wheel 里已经内置了大部分运行库。这篇文章就直接把这个过程拆开讲清楚。我会先说明“不用装 HIP SDK”到底在什么意义上成立再给出一套在 Ubuntu 22.04 AMD 9070XT 环境下照抄即可的编译流程最后附上验证脚本和常见报错排查清单。按标题说的跑通后 SageAttention 在注意力部分相对 PyTorch 自带 SDPA 在 9070XT 上大概有 30% 量级的提升但注意这里是“注意力内核本身”的提升整模型推理端到端提升会因显存带宽和算子占比而打折扣这部分后文也会详细解释。1. 这篇文章真正要解决的问题在 AMD 显卡上做深度学习推理用户通常会卡在两个地方。第一个是环境安装。很多教程会让用户去 AMD 官网下载“AMD ROCm”全家桶再按系统版本安装 hip-sdk、rocm-dev、llvm 等一大堆元包。结果是装完动辄占用十几 GB 磁盘中间任何一个依赖冲突都可能导致后面的编译失败。于是不少人的第一反应是算了还是用 NVIDIA 吧。第二个问题是编译方式。SageAttention 这种底层优化库官方示例默认面向 CUDA。AMD 用户把它拉到本地直接 pip install大概率会遇到gfx1201 is not supported、hipcc not found、找不到hiprtc这类报错。报错信息彼此之间看起来毫无关联但根子往往是同一个编译器没找到 HIP 工具链或者环境变量没有指到正确的位置。这篇文章要解决的就是这两个问题怎么用最小的依赖在 AMD 平台获得可用的 HIP 编译环境怎么正确编译 SageAttention并且让它在 9070XT 这类较新的 RDNA 架构显卡上真正跑起来怎么判断跑出来的性能是否符合预期以及如果报错第一步应该查哪里。如果你手里的卡是 AMD 6000 系、7000 系或 9000 系这篇文章同样适用。差别只在于 ROCm 版本支持和 gfx 架构标识核心思路一致。2. SageAttention 是什么为什么 AMD 用户值得折腾2.1 一句话理解 SageAttentionSageAttention 是一个注意力机制的高效实现目标是在尽量不损失精度的前提下让 Attention 的算得更快、占用显存更少。我们平时在 PyTorch 里直接写torch.nn.functional.scaled_dot_product_attentionPyTorch 会自动选择它内部的融合 attention 内核但这个内核在长序列场景下并不是最优的。SageAttention 通过把 Q、K 的相乘结果做低比特量化再用特殊算法补偿误差让注意力计算在同等精度下比标准实现快出一截。它和 FlashAttention 的关系可以这样理解项目FlashAttentionSageAttention原理方向通过分块和重计算降低显存访问通过低比特量化降低计算量和显存带宽精度控制依赖 fp16/bf16 累加用量化补偿机制保持精度硬件适配对 CUDA 最成熟ROCm 支持看版本CUDA 合并了 ROCm 构建AMD 可自行编译适合场景长文本推理、训练长文本推理、对显存占用敏感的场景对 AMD 用户来说SageAttention 最大的吸引力在于它不只是“能用”而是官方代码里有 ROCm 后端的支持路径。虽然官方 release 里没有直接给你一个适合所有 ROCm 版本的预编译包但自己编译的难度并没有想象中那么高。2.2 为什么它能快 30%先说清楚这 30% 是怎么来的。Attention 的计算本身分为两个阶段Q K^T得到注意力分数矩阵对分数矩阵做 softmax再乘 V。瓶颈在于需要把 Q、K、V 三块大矩阵从显存里反复读出来。SageAttention 的思路是把中间的浮点运算用低精度表示减少显存读写的数据量同时对量化误差做非对称补偿。显存带宽一旦降下来计算速度自然提升。在 9070XT 上它的显存带宽和 RDNA3 架构相比有提升所以量化带来的带宽收益会更明显。但你要注意30% 是“Attention 算子本身”的收益。如果整个模型里 Attention 只占 20% 的时间那么端到端可能也就快 6% 左右。这就是为什么很多人在网上发帖说“我测了 SageAttention怎么整模型没快多少”因为他们只看了端到端时延没有单独对比 attention 算子。所以本文的验证脚本也要单独测 Attention而不是直接用一个完整的 LLM。先把算子层面的收益验证出来再谈集成到推理管线里。2.3 什么情况下不值得编译如果你的模型全是短序列比如平均不到 512 token或者你只是偶尔跑一次测试编译 SageAttention 的收益有限。它更适合长序列、批量推理、显存吃紧的生产场景。下面这种情况更适合投入时间序列长度经常到 2048 甚至 8192 以上同时跑多个 batch显存已经比较紧张推理延迟敏感希望每个算子都快一点。3. 为什么说“不用装 HIP SDK”也能编译这句话需要先做一点语义澄清否则会有误导。我这里说的“不用装 HIP SDK”指的是不需要安装 AMD 官方那个rocm-hip-sdk元包。这个元包会把整套开发工具链、数学库、调试工具全部拉下来体积很大而且容易和系统里已有的 CUDA 或其他 GPU 库产生依赖冲突。那为什么可以不装它关键在 PyTorch 的 ROCm wheel。当你用官方命令安装torch的 ROCm 版本时wheel 内部会携带一批 HIP 运行库包括 hiprtc 等。也就是说你的 Python 环境里实际上已经有了能让 HIP 程序运行时链接的那部分库你缺的只是编译阶段的工具链。3.1 HIP SDK 与 HIP 运行时的区别术语上容易混淆先用一张表理清楚。组件作用是否需要完整安装HIP Runtime提供程序运行时所需的动态库如 libhiprtc.so通常已随 PyTorch ROCm wheel 打包HIP Compiler提供 hipcc 编译器把 HIP/CUDA 风格源码编译成 GPU 二进制需要但可以只安装最小工具链ROCm 内核驱动让操作系统识别 AMD GPU提供 /dev/kfd 等设备节点必须安装rocm-hip-sdk 元包把编译器、头文件、数学库、调试工具打包在一起的集合不需要换句话说你需要的不是“整个 SDK”而是“HIP 编译工具链 可用的运行时”。这正好对应了标题里的结论。只要用 AMD 官方或社区维护的包管理器安装一个轻量级工具链再配合 PyTorch ROCm wheel就足以编译 SageAttention。3.2 PyTorch ROCm wheel 究竟给了你什么PyTorch 官方在发布 ROCm 版 wheel 时会把lib目录下塞进一套 AMD 库这样用户拿到手后就不需要再独立安装 rocblas、hiprand、rocrand 等一堆依赖。你可以用下面这段代码验证当前环境里的 HIP 信息python - EOF import torch print(PyTorch version:, torch.__version__) print(HIP version:, torch.version.hip) print(GPU name:, torch.cuda.get_device_name(0)) print(GPU capability:, torch.cuda.get_device_capability(0)) EOF注意PyTorch 里把 AMD GPU 也统一叫torch.cuda这是历史命名原因看到cuda字样不要慌。如果输出里能正常打印 HIP 版本和显卡型号说明运行时没有问题。这时候进度已经完成了一半剩下的工作就是让编译工具链找到这个运行时。4. 环境准备与前置条件4.1 硬件与系统要求本文以如下环境为基准AMD Radeon RX 9070 XTRDNA4 架构gfx1201Ubuntu 22.04 或 24.0432GB 以上内存编译内核时会占用比较多内存16GB 可能吃力建议预留至少 20GB 磁盘空间如果你用的是 RX 7900 XTX、780M 核显等 RDNA3 甚至更老的架构流程基本一致只要把 ROCm 版本换成官方支持你那张卡的版本即可。4.2 安装 ROCm 基础环境这里说的是“最基础版本”目标是让系统能识别 GPU并提供 hipcc 编译器。不装完整 SDK。打开终端按顺序执行sudo apt update sudo apt install -y wget curl git python3-venv python3-pip ninja-build接下来从 AMD 官方软件源安装 ROCm 基础组件。不同系统的安装命令略有差异以 Ubuntu 22.04 为例sudo apt install -y rocm-hip-runtime hipcc rocm-dev解释一下这几个包rocm-hip-runtimeHIP 运行时环境包含各种 HIP 动态库hipccHIP 编译器本身rocm-devROCm 的头文件等开发组件。这套体积远小于完整rocm-hip-sdk元包但仍然足够编译 SageAttention。安装完成后把 ROCm 路径加入环境变量export ROCM_HOME/opt/rocm export HIP_PATH$ROCM_HOME export PATH$ROCM_HOME/bin:$PATH export LD_LIBRARY_PATH$ROCM_HOME/lib:$LD_LIBRARY_PATH验证一下hipcc --version rocm-smi如果hipcc能打印版本号rocm-smi能看到 GPU 温度、显存、风扇等状态说明基础环境已经就绪。4.3 准备 Python 虚拟环境我不建议直接在系统 Python 里装也不建议用 Anaconda base 环境硬闯。编译这种带 C/HIP 扩展的库虚拟环境隔离是最稳妥的。python3 -m venv ~/venvs/sage source ~/venvs/sage/bin/activate pip install --upgrade pip然后安装 PyTorch ROCm 版本pip install torch --index-url https://download.pytorch.org/whl/rocm6.3如果网络条件不好可以选择国内 PyTorch 镜像源这里以官方源为例。装完后验证。python - EOF import torch print(torch.__version__) print(torch.version.hip) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0)) EOF能打印True和显卡名称说明 PyTorch 能正确调用 ROCm可以进入下一步。5. 核心编译流程照抄版步骤这部分是全文重点。我在写的时候刻意把步骤拆得很细原因很简单SageAttention 编译失败的用户绝大多数不是因为缺什么高端技巧而是环境变量或者 Python 包顺序不对。5.1 设置编译环境变量每次打开新终端都要先导出环境变量。建议直接用export写入当前会话。export ROCM_HOME/opt/rocm export HIP_PATH$ROCM_HOME export PATH$ROCM_HOME/bin:$PATH export LD_LIBRARY_PATH$ROCM_HOME/lib:$LD_LIBRARY_PATH export TORCH_CXX_FLAGS-D__HIP_PLATFORM_HCC__最后一行TORCH_CXX_FLAGS的意思是告诉编译过程当前平台是 HIP 平台避免出现平台宏缺失导致的奇怪报错。如果确认你的显卡是 9070 XT可以额外指定export PYTORCH_ROCM_ARCHgfx1201指定显存架构可以避免编译出来的二进制包含太多不必要的指令变体。5.2 安装编译期依赖SageAttention 在构建时需要ninja和packaging。前者负责构建并行加速后者用于版本解析。pip install ninja packaging pybind115.3 安装 SageAttention最直接的安装命令pip install sageattention2.2.0等一下先不要急着敲回车。如果你确认当前环境比较干净直接安装确实可以。但我更推荐先用源码安装因为源码方式方便你定位问题。git clone https://github.com/thu-ml/SageAttention.git cd SageAttention pip install -v .安装日志会滚动输出编译过程。看到类似下面的行说明 HIP 编译器已经正常工作hipcc -O3 ... -c gemm_forward.cu如果这一步直接成功就可以跳到第 6 节验证效果。如果出现报错尤其是gfx1201相关的看 5.4。5.4 RDNA4 架构不支持时的处理较新的 ROCm 版本已经能直接识别 gfx1201但如果你使用的 ROCm 6.2 或更早版本编译时可能出现error: failed to load gfx1201...或者The specified target gfx1201 is not recognized by the compiler这种场景下常见的做法是把目标架构覆盖到 RDNA3 的 gfx1100。RDNA4 硬件可以执行 RDNA3 指令集性能会有少量损失但能正常跑。export HSA_OVERRIDE_GFX_VERSION11.0.0 export PYTORCH_ROCM_ARCHgfx1100设置后再重新编译一次。注意这个变量在运行时也要保留否则程序可能加载不了编译好的内核。5.5 编译成功后的检查编译完成后直接进入 Python 验证导入是否正常python - EOF from sageattention import sageattn print(sageattn import OK) EOF只要能打印sageattn import OK说明编译产物已经被正确加载。此时你也可以顺便看一下安装版本pip show sageattention如果显示版本 2.2.0说明安装完成。6. 运行结果与效果验证编译通过的兴奋劲过去后一定要做性能验证。不然你无法确定跑出来的结果到底是“正常的快”还是“凑合能用”。6.1 最小验证脚本下面这个脚本做两件事验证 SageAttention 与 PyTorch 自带 SDPA 的输出是否一致对比两者的耗时。把脚本保存为bench_attention.py。import time import torch import torch.nn.functional as F from sageattention import sageattn torch.manual_seed(0) batch_size 4 num_heads 32 seq_len 2048 head_dim 128 q torch.randn(batch_size, num_heads, seq_len, head_dim, devicecuda, dtypetorch.float16) k torch.randn(batch_size, num_heads, seq_len, head_dim, devicecuda, dtypetorch.float16) v torch.randn(batch_size, num_heads, seq_len, head_dim, devicecuda, dtypetorch.float16) # 预热 with torch.inference_mode(): for _ in range(5): _ F.scaled_dot_product_attention(q, k, v) _ sageattn(q, k, v, is_causalFalse) torch.cuda.synchronize() # 正确性 with torch.inference_mode(): out_ref F.scaled_dot_product_attention(q, k, v) out_sage sageattn(q, k, v, is_causalFalse) diff (out_ref.float() - out_sage.float()).abs().max().item() print(max abs diff:, diff) print(output shape:, out_sage.shape) print(output dtype:, out_sage.dtype) # 性能对比 def bench(fn, repeat50): for _ in range(5): fn() torch.cuda.synchronize() start time.perf_counter() for _ in range(repeat): fn() torch.cuda.synchronize() return (time.perf_counter() - start) / repeat * 1000 t_ref bench(lambda: F.scaled_dot_product_attention(q, k, v)) t_sage bench(lambda: sageattn(q, k, v, is_causalFalse)) print(fSDPA avg time: {t_ref:.3f} ms) print(fSageAttention avg time: {t_sage:.3f} ms) print(fspeedup: {t_ref / t_sage:.2f}x)6.2 预期结果说明正确性部分max abs diff通常会落在1e-1到1e-2量级这是因为 SageAttention 采用了低比特量化并非完全一致的浮点结果列表。性能部分在 9070 XT 上跑 2048 长度时经常能观察到1.2x到1.4x的算子级加速。这里对应标题中提到的“快 30%”量级。如果你的序列长度更长比如 4096 或 8192收益通常更明显因为显存带宽节省更多。另外请务必把脚本里的is_causalFalse换成正弦模式。写法上is_causalTrue会走不同的内核分支性能表现也会不一样。6.3 使用 Triton 后端的替代方案如果你不想忍受编译 C 扩展的曲折SageAttention 也提供了 Triton 后端。ROCm 版本的 PyTorch 会自带一个 pytorch-triton 的 AMD 分支所以可以直接这样用from sageattention import sageattn_triton out sageattn_triton(q, k, v, is_causalFalse)Triton 后端的编译是 JIT 的也就是运行时第一次调用时会自动编译不需要单独装 HIP SDK。这个方案对新手更友好但性能通常比手工编译的 C 内核略低。我的建议是先跑 Triton 后端确认业务逻辑正常再编译 C 内核并对比精度和性能哪个效果好线上就用哪个。7. 常见问题与排查思路把我在排查过程中最常遇到的几类问题整理成表方便你对照排查。问题现象可能原因排查方式解决方案编译时报gfx1201 is not supportedROCm 版本较老不认识 RDNA4 架构查看编译日志中的 target 行设置HSA_OVERRIDE_GFX_VERSION11.0.0并重装编译报hipcc not found没有安装 hipcc 或 PATH 没配好which hipccsudo apt install hipcc确保/opt/rocm/bin在 PATH编译报Cannot open include file: hip/hip_runtime.hROCM_HOME 未设置echo $ROCM_HOMEexportROCM_HOME/opt/rocm编译失败但日志乱码并行编译导致错误被吞掉使用pip install -v .查看完整日志先精确复现单文件编译错误再重试import sageattention报undefined symbol运行时没有找到 HIP 库ldd查看扩展 so 的依赖设置LD_LIBRARY_PATH/opt/rocm/lib运行时提示显存不足测试的 batch 或 seq_len 过大用rocm-smi查看显存占用调小 batch_size / seq_len精度比对差距大使用 fp32 或没有设置合适的缩放检查 dtype 是否 fp16/bf16将输入转为torch.float16或torch.bfloat16第一次跑很慢Triton JIT 编译尚未完成观察后续轮次耗时预热后再次计时一个容易忽略的点如果你之前安装过 CUDA-toolkit或者系统里存在/usr/local/cuda某些构建脚本可能会误把 CUDA 路径当成 HIP 路径。遇到这种场景我一般会在当前终端环境里这样隔离unset CUDA_HOME unset CUDA_PATH然后再编译。另外如果你在 Linux 上看到HIP_VISIBLE_DEVICES相关的报错是因为 AMD 复用了 HIP 的设备编号机制。可以用rocm-smi --showbus查看 GPU 在 ROCm 中的索引然后设置export HIP_VISIBLE_DEVICES08. 最佳实践与工程建议8.1 环境管理我强烈建议把 ROCm 编译环境固定成一份可复现的配置。SageAttention 这类底层库对工具链版本很敏感今天能编过不代表三个月后换一个 ROCm 小版本还能编过。你可以把下面的内容保存为setup_rocm_env.sh#!/usr/bin/env bash export ROCM_HOME/opt/rocm export HIP_PATH$ROCM_HOME export PATH$ROCM_HOME/bin:$PATH export LD_LIBRARY_PATH$ROCM_HOME/lib:$LD_LIBRARY_PATH export TORCH_CXX_FLAGS-D__HIP_PLATFORM_HCC__每次新建终端时source setup_rocm_env.sh这样可以避免反复踩环境变量丢失的坑。8.2 版本管理在项目里我建议你把requirements.txt写成这样torch2.5.1rocm6.3 ninja1.11.1 packaging23.0 sageattention2.2.0注意PyTorch 版本和 SageAttention 版本有对应关系。如果未来升级 PyTorch建议先在小规模环境验证再决定是否升级 SageAttention。8.3 性能调优如果发现 SageAttention 编译通过但速度提升不明显可以按下面顺序排查。第一确认注意力部分真的是瓶颈。可以先跑一次整体模型再跑一次只替换注意力算子的模型对比两者端到端时延。如果端到端没变化但算子确实快了说明你模型里其他算子占比太高SageAttention 的收益被摊薄了。第二确认显存带宽确实节省了。可以用rocm-smi监控功耗和显存运行频率。如果频率一直在高位但耗时没有下降可能需要检查是否 system memory 参与了数据交换。第三尝试更大序列长度。SageAttention 的量化优势在长序列下更明显。如果平均序列只有 512那收益可能只有几个百分点。8.4 集成到推理框架把 SageAttention 集成到推理管线时注意输入 layout。它默认接受的输入是(batch, heads, seq_len, head_dim)也就是 PyTorch 原生格式。但很多推理框架里KV cache 的 layout 是(batch, heads, head_dim, seq_len)或者反过来的。这种情况下最稳妥的做法是在调用sageattn之前把k和v转成框架需要的布局。虽然多了一次 permute 的开销但通常远小于注意力本身的时间节省。8.5 关于生产环境的提醒如果你是生产系统不要跳过精度验证。SageAttention 虽然有精度补偿但仍然是低比特量化实现。把输入换成长短差异很大的文本、或者加入大量 padding 后误差分布可能与基准测试不同。建议在生产环境增加一个开关ENABLE_SAGE os.environ.get(USE_SAGE_ATTENTION, 0) 1默认关闭先用 PyTorch 原生实现跑全量回归对比输出合理后再开启。这样就算出了精度问题也能快速回退。9. 总结与后续学习方向现在全文的核心结论已经清楚AMD 上编译 SageAttention 不一定要装庞大的 HIP SDK最小依赖就是 HIP 运行时加 hipcc 编译器核心工作量在环境变量和架构设置上只要 ROCM_HOME、HIP_PATH、PYTORCH_ROCM_ARCH 配好就没有想象中那么玄乎9070XT 上跑通后Attention 算子级大约能拿到 20% 到 30% 的提升但端到端提升还要看模型里注意力占比遇到报错按第 7 节的排查表走90% 的问题都能在十分钟内定位。接下来如果你想把这条路走得更深可以关注这三个方向。第一个是 ROCm 官方文档里关于 gfx 架构和 LLVM 编译目标的内容这决定了你能不能在更老的 ROCm 版本上运行新显卡。第二个是 SageAttention 论文里关于量化误差补偿的细节。只知道怎么编译还不够真正到了生产环境你需要能解释为什么你的业务场景里误差偏大以及可不可以调整量化参数。第三个是 Triton 内核的写法。SageAttention 的 Triton 后端给了很好的示例你可以把它当成学习 AMD 上算子开发的第一课逐步理解在 ROCm 平台上调优一个算子需要关注哪些性能指标。建议收藏这篇文章等你的环境装好之后照着第 5 节的步骤一步步跑比反复看报错日志要省时间得多。