从 CUDA 到 HIP,用 HIPify 工具迁移大模型代码实战
从 CUDA 到 HIP:一次真实的算子迁移手记
最近手头有个项目,原本是基于 NVIDIA GPU 跑的,因为成本考量,老板拍板要迁到 AMD Instinct MI300X 上。刚开始我也头大,毕竟在 CUDA 舒适区待久了,突然面对 ROCm 生态,心里多少有点打鼓。但实际跑下来发现,只要工具用对,路数摸清,这事儿并没想象中那么难。今天就想和大家聊聊,我是怎么利用HIPify把一段关键的 PyTorch 自定义算子从 CUDA 迁移到 HIP 的,顺便分享点踩坑经验。
为什么选择 HIPify 作为切入点?
很多从 NVIDIA 转投 AMD 平台的工程师,第一反应是“重写代码”。其实真没必要。AMD 官方提供的HIPify工具集(包括hipify-clang和hipify-perl)已经非常成熟,尤其是在 ROCm 7.x 版本更新后,它对 C++ 新特性的支持度大幅提升。
我的策略很明确:能自动转的绝不手写。HIPify的核心作用就是扫描你的 CUDA 源码,把cudaMalloc变成hipMalloc,把<<< >>>核函数启动语法转换成 HIP 风格。对于大部分标准算子,它甚至能做到“一键转换,编译通过”。但这只是第一步,真正的挑战往往藏在自动转换后的细节里。
实战:一个 PyTorch 自定义算子的迁移过程
这次迁移的目标是一个用于加速注意力机制中特定掩码操作的自定义 CUDA Kernel。原始代码大概长这样:
// original_cuda_kernel.cu__global__voidmasked_add_kernel(float*input,float*mask,float*output,intn){intidx=blockIdx.x*blockDim.x+threadIdx.x;if(idx<n){if(mask[idx]>0.5f){output[idx]=input[idx]+1.0f;}else{output[idx]=input[idx];}}}voidlaunch_kernel(float*d_in,float*d_mask,float*d_out,intsize){intthreads=256;intblocks=(size+threads-1)/threads;masked_add_kernel<<<blocks,threads>>>(d_in,d_mask,d_out,size);cudaDeviceSynchronize();}运行hipify-clang original_cuda_kernel.cu后,生成的代码大部分看起来没问题:
// hipified_kernel.cpp#include<hip/hip_runtime.h>__global__voidmasked_add_kernel(float*input,float*mask,float*output,intn){intidx=blockIdx.x*blockDim.x+threadIdx.x;if(idx<n){if(mask[idx]>0.5f){output[idx]=input[idx]+1.0f;}else{output[idx]=input[idx];}}}voidlaunch_kernel(float*d_in,float*d_mask,float*d_out,intsize){intthreads=256;intblocks=(size+threads-1)/threads;hipLaunchKernelGGL(masked_add_kernel,dim3(blocks),dim3(threads),0,0,d_in,d_mask,d_out,size);hipDeviceSynchronize();}看着挺美,直接编译却报了错,或者运行时结果不对。问题出在哪?内存管理逻辑。
在 CUDA 代码中,我们习惯性地假设显存分配和释放是严格配对的,但在迁移到 AMD 平台时,特别是结合 PyTorch 的 Tensor 接口时,直接操作原始指针容易出问题。HIPify 不会帮你检查逻辑错误。我不得不手动修正了内存分配部分,确保使用hipMalloc分配的内存生命周期与 PyTorch 的torch::Tensor托管内存不冲突。
更重要的是,在 ROCm 7.x 下,某些原子操作或共享内存的使用方式需要微调。比如,如果原代码用了__syncthreads(),在复杂的块调度下,AMD 架构可能需要更明确的内存围栏指令来保证一致性。这部分没法完全依赖工具,必须人工 Review。
进阶优化:引入 TileLang 编写高性能 Kernel
自动转换能保证“跑通”,但要想“跑得快”,还得针对 AMD 架构做深度优化。这时候我就用上了TileLang。
TileLang 是个挺有意思的项目,它旨在简化张量程序的编写。相比于直接写晦涩的 HIP C++ 代码,TileLang 提供了一种更高级的抽象,让你能更专注于数据流和计算逻辑,而不是被线程索引搞得晕头转向。
在重构那个掩码算子时,我尝试用 TileLang 重写了核心循环。它自动处理了分块(Tiling)和数据预取,生成的代码在 MI300X 上的 HBM 带宽利用率明显高于 naive 的 HIP 版本。虽然目前 TileLang 还处在快速迭代期,但对于这种规则密集的矩阵操作,它的产出效率非常高。
# 伪代码示例:TileLang 风格的描述@tilelang.kerneldefoptimized_masked_add(input,mask,output):# 自动处理 block 和 thread 映射foriintile_range(N):ifmask[i]>0.5:output[i]=input[i]+1.0这种写法不仅可读性强,而且编译器后端能更好地针对gfx942(MI300 系列架构)进行指令调度优化。
避坑指南:如何在 Github 寻找靠谱的 Triton 分支
除了手写 Kernel,很多时候我们依赖Triton来写高性能算子。但众所周知,Triton 原生是对 CUDA 的,要在 ROCm 上跑,得找特定的分支。
在 Github 上搜 “Triton ROCm” 会跳出一堆库,千万别随便拉一个就用。我的筛选原则很简单:
- 看 Commit 时间:如果最后更新是半年前,直接 pass。ROCm 版本迭代快,旧分支在新驱动上大概率报 “illegal instruction”。
- 查 Issue 闭环:搜索关键词
gfx942或MI300,看看作者是否在积极修复相关报错。如果一个库的 Issue 里全是未解决的段错误,哪怕 Star 再多也别碰。 - 验依赖链条:确认它依赖的 PyTorch 版本是否与你当前的 ROCm 7.x 环境匹配。
目前社区里几个活跃的 Triton ROCm 分支维护得不错,配合PYTORCH_ROCM_ARCH环境变量指定正确的架构代码,基本能覆盖大部分主流算子需求。记得在编译前清理掉旧的 build 缓存,否则很容易因为残留的 CUDA 对象文件导致链接失败。
写在最后
从 CUDA 到 HIP,本质上不是推倒重来,而是一次技术栈的平滑演进。HIPify帮我们解决了 80% 的语法转换工作,剩下的 20% 则需要我们对底层内存模型和架构特性有更深的理解。配合像 TileLang 这样的新工具,以及社区里活跃的 Triton 分支,在 AMD 平台上构建高性能的大模型推理服务已经完全可行。
如果你也在经历类似的迁移,别被初期的编译报错吓退。多看看日志,多查查 Github 上的最新 Issue,很多时候问题就藏在某个不起眼的配置项里。一旦跑通了第一个算子,后面的路就会顺畅很多。
200小时GPU算力已就位,快来领取:https://marketing.csdn.net/questions/Q2604140858304426315?utm_source=AIpaper
