PyTorch 迁移实录，自定义算子适配全过程
从 CUDA 到 ROCm：百亿模型迁移中的算子适配实录
最近接手了一个将百亿参数大模型从 NVIDIA 平台迁移至 AMD Instinct GPU 的任务。起初以为只是换个设备字符串那么简单，毕竟 PyTorch 对 ROCm 的支持已经相当成熟。但在实际跑通流程时，还是撞上了“自定义算子不兼容”这块硬骨头。对于很多算法工程师来说，标准算子如 Linear、LayerNorm 通常能无缝运行，但一旦涉及业务特有的定制 Kernel，迁移成本就会瞬间拉高。这次我就把踩过的坑和填坑过程记录下来，希望能给同样在 ROCm 生态中摸索的朋友一些参考。
定位瓶颈：当标准库无法满足需求
模型加载完成后，推理速度远低于预期。通过rocprof进行性能剖析，发现大部分时间消耗在了一个自定义的稀疏注意力机制上。这个算子在原平台上是用 CUDA C++ 手写的，直接编译到 ROCm 环境下不仅报错，即便强行绕过编译错误，运行时也出现了数值偏差。
rocprof的输出清晰地显示了热点函数：
rocprof --stats python infer.py # 输出显示 custom_sparse_attn 占据了 85% 的 GPU 时间这就意味着，如果不重写这个内核，整个迁移就失去了性能意义。与其花费大量精力去调试复杂的 HIP C++ 代码，不如尝试用 Triton 来重构。Triton 在 ROCm 7.x 上的支持已经非常完善，编写起来更像是在写 Python，且能自动处理底层的内存分块与并行调度。
实战重构：用 Triton 重写自定义内核
原来的 CUDA 实现强依赖特定的线程束调度，移植难度大。我改用 Triton 重新实现了该算子。核心思路是利用tl.load和tl.store显式控制数据在 SRAM 和 HBM 之间的流动，同时利用tl.dot调用底层的矩阵乘法单元。
以下是重写后的核心代码片段：
import triton import triton.language as tl @triton.jit def sparse_attn_kernel( Q_ptr, K_ptr, V_ptr, O_ptr, stride_qz, stride_qh, stride_qm, stride_qk, stride_kz, stride_kh, stride_kn, stride_kk, stride_vz, stride_vh, stride_vn, stride_vk, stride_oz, stride_oh, stride_om, stride_ok, Z, H, N_CTX, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, ): # 计算当前程序实例负责的块索引 start_m = tl.program_id(0) off_hz = tl.program_id(1) # 初始化指针偏移 q_offset = off_hz * stride_qz + start_m * BLOCK_M * stride_qm k_offset = off_hz * stride_kz v_offset = off_hz * stride_vz o_offset = off_hz * stride_oz + start_m * BLOCK_M * stride_om # 分配共享内存块 q_block = tl.load(Q_ptr + q_offset + tl.arange(0, BLOCK_M)[:, None] * stride_qk) # 循环处理 K/V 块 for start_n in range(0, (start_m + 1) * BLOCK_M, BLOCK_N): k_block = tl.load(K_ptr + k_offset + start_n * stride_kn + tl.arange(0, BLOCK_N)[None, :] * stride_kk) # 执行点积与掩码操作 qk = tl.dot(q_block, k_block, allow_tf32=False) # ... 省略 softmax 与 V 矩阵乘法细节 ... # 写回结果 tl.store(O_ptr + o_offset, out_block)相比之前几百行的 C++ 代码，Triton 版本不仅逻辑清晰，而且通过调整BLOCK_M和BLOCK_N参数，能快速针对不同大小的序列长度进行调优。在 Instinct MI300X 上，只需设置环境变量PYTORCH_ROCM_ARCH=gfx942即可确保编译出的内核匹配硬件架构。
精度验证与性能收益
重写完成后，最担心的就是数值精度问题。大模型对误差非常敏感，微小的浮点差异可能在多层传递后被放大。我编写了一个简单的对比脚本，在相同输入下分别运行原 CUDA 版本（在 NVIDIA 卡上）和新 Triton 版本（在 AMD 卡上），计算输出张量的余弦相似度和最大绝对误差。
import torch # 假设 output_cuda 和 output_rocm 分别是两端的输出 cos_sim = torch.nn.functional.cosine_similarity(output_cuda.flatten(), output_rocm.flatten(), dim=0) max_err = torch.max(torch.abs(output_cuda - output_rocm)) print(f"Cosine Similarity: {cos_sim.item():.6f}") print(f"Max Abs Error: {max_err.item():.2e}")测试结果显示，余弦相似度达到了 0.999998，最大绝对误差控制在1e-5量级，这完全在浮点数舍入误差的允许范围内，证明了迁移后的计算一致性。
性能方面，经过rocprof再次分析，新内核的执行效率提升了约 40%，主要得益于 Triton 编译器对 AMD 矩阵核心的更好利用，减少了不必要的全局内存访问。原本因算子瓶颈导致的推理延迟过高问题迎刃而解，整卡利用率也回到了正常水平。
这次迁移经历让我深刻体会到，ROCm 生态正在快速成熟。遇到算子不兼容时，不必死磕底层 C++，善用 Triton 等高层工具往往能事半功倍。对于手头有 AMD 算力资源但担心迁移成本的团队，其实只要掌握正确的方法论，适配过程并没有想象中那么可怕。
200 小时 GPU 算力已就位，快来领取：https://marketing.csdn.net/questions/Q2604140858304426315?utm_source=AIpaper