从 CUDA 到 HIP：一次真实的 PyTorch 自定义算子迁移实录

最近手头有个老项目，里面塞了不少为了加速推理而手写的 CUDA Kernel。随着 AMD Instinct MI300X 这类大显存卡性价比的凸显，老板提议把部分推理任务迁到 AMD 平台上跑。起初我心里也打鼓，毕竟过去几年一直泡在 NVIDIA 生态里，对 ROCm 的印象还停留在“配置麻烦、算子缺失”的阶段。但实际折腾了一周下来，发现 ROCm 7.x 的变化确实很大，尤其是 Triton 编译器的原生支持，让很多原本需要重写 C++ 的代码可以直接复用逻辑。今天就把这次从 CUDA 到 HIP 的迁移过程复盘一下，重点聊聊自动化工具的局限性和手动修复算子的真实坑点。

HIPify 自动化转换：爽快感与“假象”

迁移的第一步自然是能省则省。AMD 官方提供的 hipify 工具链（主要是 hipify-python 和 hipify-clang）确实是神器。对于项目中那些标准的 PyTorch API 调用，比如 torch.cuda.is_available() 或者常规的 Tensor 操作，HIPify 几乎能一键完成替换。

我直接在项目根目录执行了转换命令：

hipify-python -p . --output-directory ./hip_ported

脚本跑完后，扫了一眼生成的代码，cuda 关键字基本都变成了 hip，头文件引用也自动修正了。这时候很容易产生一种“已经迁移成功”的错觉。但千万别急着编译，自动化只能解决语法层面的映射，它不懂你的业务逻辑，更处理不了那些深度依赖 NVIDIA 特定硬件特性的自定义算子。

踩坑实录：架构代码不匹配的“非法指令”

真正的挑战来自项目中几个核心的自定义 Attention 算子。这些算子是用 Triton 编写的，原本在 A100 上跑得飞起。在 AMD 环境下，ROCm 7.x 虽然已经支持 Triton，但对底层架构的敏感度极高。

第一次尝试运行迁移后的代码时，程序直接崩溃，报错信息非常晦涩：Illegal instruction (core dumped)。查了半天日志才发现，问题出在编译时的架构指定上。NVIDIA 的编译器通常能自动识别当前显卡，但在 ROCm 环境下，必须显式告诉编译器目标架构是什么。

解决方法是在运行 Python 脚本前，强制导出环境变量。假设你使用的是 MI300X（架构代号 gfx942），必须执行：

export PYTORCH_ROCM_ARCH="gfx942"
export HSA_OVERRIDE_GFX_VERSION="9.4.2"

如果这一步漏掉，或者填成了旧款显卡的 gfx90a，编译出的二进制指令集就会与当前硬件不匹配，导致运行时直接崩掉。这个坑我足足花了半天才定位到，建议大家在做任何自定义算子迁移前，先用 rocminfo 确认好自己的架构代号。

手动修复：当 Triton 遇到 ROCm

搞定环境变量后，第二个拦路虎出现了。部分复杂的 Triton Kernel 在编译时报错，提示某些原子操作（atomic ops）在 HIP 后端未定义。这是因为虽然 Triton 支持 ROCm，但部分高级特性在 HIP 后端的实现还在完善中，不像 CUDA 后端那么成熟。

针对这个问题，我有两个解决方案：

1. 降级算子逻辑：检查报错的 Kernel 代码，发现是因为使用了特定的 atomic_add 变体。通过修改 Triton 代码，将其拆解为标准的加载 - 计算 - 存储流程，虽然牺牲了一点理论峰值性能，但保证了兼容性。
2. 回退到 HIP C++：对于实在无法用 Triton 兼容的极端优化算子，我选择直接用 HIP C++ 重写。好在 hipify-clang 能把大部分 CUDA C++ 代码转成 HIP C++，我只需要手动修补几个内联汇编部分。

以下是一个典型的修复片段，展示了如何处理精度转换中的兼容性问题：

# 原始 CUDA/Triton 逻辑可能隐式依赖 FP16 的特定行为
# 在 ROCm 下，建议显式指定数据类型以避免歧义

import torch
import triton
import triton.language as tl

@triton.jit
def scaled_add_kernel(
    x_ptr, y_ptr, out_ptr,
    scale: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    
    # 显式加载并转换类型，增强在 HIP 后端的稳定性
    x = tl.load(x_ptr + offs).to(tl.float32)
    y = tl.load(y_ptr + offs).to(tl.float32)
    
    res = x + y * scale
    tl.store(out_ptr + offs, res.to(tl.float16))

这段代码在 NVIDIA 平台上可能省略 .to(tl.float32) 也能跑，但在 ROCm 7.x 上，显式的类型转换能有效避免编译器优化带来的数值偏差或编译失败。

验证与性能调优

代码跑通只是第一步，性能才是关键。在 MI300X 上重新 benchmark 后，我发现显存带宽利用率比预期要高，这得益于 ROCm 7.x 对 HBM3 的调度优化。不过，初次运行的吞吐量只有 NVIDIA 平台的 80% 左右。

经过 profiling 发现，瓶颈在于 Kernel 启动开销。通过调整 Triton 的 num_warps 和 num_stages 参数，使其更贴合 AMD GPU 的 Wavefront 机制（AMD 的线程束概念），最终将推理延迟压回了与 NVIDIA 相当的水平。这个过程让我深刻体会到，跨平台迁移不是简单的“翻译代码”，而是需要对新硬件架构有深入理解后的“重新调优”。

写在最后

这次迁移经历打破了我对 AMD 生态的刻板印象。ROCm 7.x 加上成熟的 Triton 支持，已经足以承载生产级的自定义算子需求。当然，过程中难免会遇到文档语焉不详的地方，需要结合源码和社区 Issue 去摸索，但这种掌控底层细节的过程本身也是技术成长的机会。

如果你也想尝试将现有的 CUDA 项目迁移到 AMD 平台，或者需要大规模算力来验证迁移效果，不妨利用现成的资源动手试试。

200 小时 GPU 算力已就位，快来领取：https://marketing.csdn.net/questions/Q2604140858304426315?utm_source=AIpaper