从 CUDA 到 ROCm：attention 这一步到底差在哪？
把 Stable Diffusion 搬到 AMD Instinct™ MI50 之前，我先用 PyTorch 2.1 官方镜像跑了一遍 512×512 的默认配置，结果nvidia-smi换成rocm-smi那一刻，显存直接飙到 7.8 GB——比同卡 CUDA 环境多出 1.8 GB。
罪魁祸首是torch.nn.MultiheadAttention在 ROCm 上默认走BMM + softmax + BMM的朴素实现，中间激活值全留在显存里；而 CUDA 路径早就默认打开xFormers的memory_efficient_attention，把 QK^T 切块后算一块丢一块，峰值显存直接腰斩。
MI50 只有 16 GB，1024×1024 图一跑就 OOM，根本玩不动。于是目标很明确：把 CUDA 上省 3 GB 的那条 attention 捷径，在 ROCm 上原样复刻出来。
hipMemory 效率实测：一张图看懂差距
先用 micro-benchmark 把问题量化。下面这段脚本分别在 CUDA 与 ROCm 上跑 4096-token、head-dim 64 的 self-attention，batch=8，测的是峰值显存与吞吐。
dockerrun--gpusall-itpytorch:2.1-cuda python benchmark_attn.py# CUDAdockerrun--device/dev/kfd--device/dev/dri-itrocm/pytorch:rocm5.7_ubuntu22.04_py3.9_pt2.1 python benchmark_attn.py# ROCm| 平台 | 峰值显存 | 吞吐 (token/s) |
|---|---|---|
| CUDA + xFormers | 4.1 GB | 6 800 |
| ROCm 默认 | 7.9 GB | 5 200 |
显存差 3.8 GB，吞吐还掉 24 %，难怪 1024 分辨率图直接炸。
把横轴换成序列长度，纵轴是显存占用，曲线斜率 ROCm 默认路径是 CUDA 的 1.9 倍——这就是典型的 O(N²) 激活膨胀。
两行环境变量，让 xFormers 在 ROCm 上复活
xFormers 的memory_efficient_attention其实早就做了 HIP 后端，只是编译开关默认 OFF。把下面两行写进 Dockerfile，就能在 MI50 上打开：
ENV FORCE_CUDA=0 ENV ROCM_HOME=/opt/rocm RUN pip install -v -U git+https://github.com/facebookresearch/xformers@main \ --no-binary :all: \ -C cmake.define.XFORMERS_ENABLE_HIP=ON \ -C cmake.define.PYTORCH_HOME=/opt/conda/lib/python3.9/site-packages/torch编译 8 分钟，wheel 体积 112 MB，比 CUDA 版大 20 %，因为把ck_tile的 HIP kernel 全静态编进去了。装完再跑一次 benchmark，显存直接掉到 4.3 GB，吞吐回到 6 500 token/s，基本打平 CUDA。
1024×1024 文生图：显存监控曲线
把模型换成 Stable Diffusion v2-1-768，batch=1，steps=50，Euler a，prompt 用默认的 “a high-tech rocm robot, ultra detailed”。
用radeontop每 500 ms 采样显存，曲线如下（单位 GB）：
16 | ╱╲ 14 | ╱ ╲___ 12 | ╱ ╲ 10 |___╱ ╲___ 8 | ╲ 6 | ╲ 4 |________________╲____ 0 10 20 30 40 50 step峰值 13.1 GB，比不开 xFormers 的 16.4 GB 省下 3.3 GB，刚好把 MI50 从悬崖边拉回来。生成一张 1024×1024 图 47 秒，只比 CUDA 慢 3 秒，完全可接受。
多阶段 Dockerfile：把环境压成 SaaS 镜像
做 SaaS 不可能让用户现场编译 8 分钟，因此把编译阶段拆到 builder 镜像，再把 wheel 拷到运行时，最终镜像从 5.8 GB 压到 2.1 GB。
# ---- builder stage ---- FROM rocm/pytorch:rocm5.7_ubuntu22.04_py3.9_pt2.1 as builder ENV FORCE_CUDA=0 ROCM_HOME=/opt/rocm RUN apt-get update && apt-get install -y git cmake ninja-build RUN pip wheel -v -U git+https://github.com/facebookresearch/xformers@main \ --no-binary :all: \ -w /dist \ -C cmake.define.XFORMERS_ENABLE_HIP=ON # ---- runtime stage ---- FROM rocm/pytorch:rocm5.7_ubuntu22.04_py3.9_pt2.1 COPY --from=builder /dist/*.whl /tmp/ RUN pip install /tmp/*.whl && rm -rf /tmp/* /root/.cache COPY app.py /app/ WORKDIR /app CMD ["python", "-u", "app.py"]把app.py做成 Gradio 界面，用户下拉选分辨率，后台自动设置--attention-type memory_efficient，云镜像跑在 AMD 开发者云g4ad.xlarge（MI50） 实例，按小时 1.2 元，比同档 CUDA 机型便宜 35 %，显存还更宽裕。
避坑小结
- PyTorch 版本要对齐：ROCm 5.7 必须配 PT2.1，PT2.0 的 hipBLAS 缺符号，import 就崩。
- xFormers 版本要追新：0.0.22 旧分支没有 HIP kernel，一定用 main。
- 编译内存别省：ninja 默认并行 job 数拉满，8 核 16 G 内存会 OOM，加
CMAKE_BUILD_PARALLEL_LEVEL=4可稳。 - 运行时权限：Docker 里
/dev/kfd要挂，非 root 用户还得加video组，否则hipInit报错 1002。
把这几颗雷排掉，Stable Diffusion 在 ROCm 上就真正“吃得下、跑得动、省得出”。下回再聊怎么把这套 attention 优化反哺到 vLLM，让大模型推理也白嫖这 3 GB 显存红利。
200小时GPU算力已就位，快来领取：https://marketing.csdn.net/questions/Q2604140858304426315?utm_source=AIpaper
