1. 为什么需要自定义算子1.1 PyTorch 算子的两种来源我一直觉得很多人对“自定义算子”这个词的第一反应是“这是那些做框架的人才会碰的东西”。实际上 PyTorch 每天在用的所有功能无论是torch.add还是conv2d本质上都能归成两类一类是官方已经编译进 C 核心的算子另一类是用 Python 组合出来的高层模块。官方算子覆盖了 95% 的日常需求可一旦你开始做模型部署、性能优化、或者写一篇新论文里的特殊函数官方算子经常不够用这时候你就会想自己动手写一个。自定义算子的含义其实很朴素我们自己写一段 C 代码或者 CUDA 代码把它编译成一个 Python 可以调用的函数。这个函数和torch.sigmoid长得一模一样能接收 Tensor、返回 Tensor差别在于它的计算逻辑由你来定义。这意味着你可以把 Python 层很难表达的循环、内存访问模式、或者某个第三方 C 库原封不动地接进 PyTorch 的计算图里。我这里说的“吃透”不是只讲 API 怎么调而是想把背后的机制拆出来PyTorch 怎么把你的 C 函数变成 Python 能导入的模块CUDA kernel 怎么在 GPU 上被调度反向传播怎么接到 autograd 上只有把这三件事弄明白你才能在一张白纸上设计自己的算子而不是照着官方示例改个名字就完事。1.2 什么场景必须自己写算子判断依据其实很简单遇到下面四种情况你就要认真考虑上自定义算子。第一种是“计算逻辑没法用 PyTorch 原生算子拼出来”。比如说你要做某种自定义的稀疏索引操作或者实现一个拉普拉斯算子的变体需要逐个遍历邻域像素并做条件更新用 Python 写for循环在 GPU 上会直接卡死用torch操作东拼西凑又很别扭这种情况写成 CUDA kernel 才是正路。第二种是“能拼出来但性能差得离谱”。最典型的就是多个逐元素算子串联比如y a * b c如果拆成mul和add两个 kernelGPU 要启动两次、读写两遍内存带宽成本直接翻倍。如果把它融合成一个 CUDA kernel一次遍历全做完速度可能提升 20% 甚至 50% 以上。我在做图像滤波的时候遇到过类似问题一个简单的权重滤波算子拆开写处理一张 4K 图像每次都要多等两个 kernel launch融合之后肉眼可见的流畅。第三种是“要复用一个已有的 C/CUDA 库”。比如某些工业软件里的 Halcon 滤波核或者自研相机 SDK 已经提供了高效的 C 函数你不想把它翻译成 Python更不想用 ctypes 一点点抠内存直接从 C 扩展里包一层 Tensor 封装是最高效的接法。第四种是“需要自定义反向传播逻辑”。PyTorch 的 autograd 已经把所有标准算子都定义了梯度但如果你的前向过程有分段不可导点或者你想在反向里注入某种额外的梯度变换就必须自己控制反向函数。自定义算子里的backward让你拥有绝对主动权。这里我给一个实用建议不要为了自定义而自定义。如果你的需求能用torch.compile或者torch.vmap解决优先用这些高层工具。它们的开发成本低而且能自动做算子级优化。只有当你需要极致的控制权时才决定走 C/CUDA 这条路。这个判断会帮你省掉大量调试时间。2. 动手前的工具链准备2.1 理解 PyTorch 内置的扩展机制PyTorch 官方早就把“自定义算子”的通道给铺好了核心就是torch.utils.cpp_extension这个模块。它本质上是一个编译器包装器拿到你的 C 源文件和 CUDA 源文件之后会调用编译器把它们编译成动态库再用 pybind11 把函数导出成 Python 可调用的对象。你不需要手动写完整的setup.py和 pybind11 绑定虽然自己写也可以但官方模块已经处理了大部分坑。这套机制里其实还有底层注册逻辑叫TORCH_LIBRARY。通过它可以在 PyTorch 内部注册一个带命名空间的算子让 PyTorch 的调度器认识你。官方推荐的现代写法是TORCH_LIBRARY(my_ops, m) { m.def(custom_add, custom_add); }然后用torch.ops.my_ops.custom_add在 Python 端访问。很多老教程喜欢用PYBIND11_MODULE直接导出普通函数差别在于PYBIND11_MODULE导出的函数不参与 PyTorch 的 op 调度也没有办法被torch.jit或者torch.compile完整追踪。如果你只是给自己跑实验怎么搞都行如果要进正式模型或者部署链路建议一开始就用TORCH_LIBRARY的姿势。2.2 版本匹配与依赖安装PyTorch 自定义 C/CUDA 算子最怕的就是版本错位。CUDA 编译器必须和 PyTorch 编译时使用的 CUDA 版本兼容不然会链接到一半报undefined reference to cudaMemcpy或者显卡驱动不支持跑起来直接报CUDA driver version is insufficient。建议先跑一下python -c import torch; print(torch.version.cuda)拿到当前 PyTorch 对应的 CUDA 版本号。如果你使用的是 conda 安装的 PyTorch通常自带配套 CUDA 库安装 CUDA Toolkit 时选择相近版本即可。Windows 下还有一个隐藏依赖Microsoft Visual C Redistributable。很多人奇怪为什么装了个 Python 包还要管 VC 运行库因为 PyTorch 的 C 扩展在 Windows 上依赖 MSVC 编译器和运行时库。建议提前安装 Visual Studio 2019 或 2022尤其要注意勾选“使用 C 的桌面开发”工作负载否则命令行里找不到cl.exe。至于 Redistributable它一般会随 VS 一起装好单独下载也有用但真正编译时你需要的是编译器不是运行库。把这两者分开理解能少翻很多帖子。我在 Ubuntu 上工作的时间比较长那边主要用 g。但说实话Ubuntu 上的坑也不比 Windows 少最典型的是系统 gcc 和 PyTorch 官方预编译用的 gcc 版本不一致后面编译时你会看到一大堆莫名其妙的符号找不到。我的建议是尽量用系统默认的编译器版本或者参考 PyTorch 官方文档里的支持矩阵。避坑的核心就是一句话让你的编译环境和 PyTorch 的编译环境保持同频。2.3 从零搭建一个最小可编译工程我平时写自定义算子文件夹结构非常简单。无论你是要用setup.py还是load_inline都有几样东西是必备的一个.cpp文件负责定义 C 函数和绑定一个.cu文件负责写 CUDA kernel一个setup.py负责告诉编译器怎么组织和输出。目录大致长这样my_op/ ├── setup.py ├── my_add.cpp ├── my_add_kernel.cu └── test.pysetup.py最简单的写法如下from torch.utils.cpp_extension import CUDAExtension, BuildExtension from setuptools import setup setup( namemy_add, ext_modules[ CUDAExtension( namemy_add, sources[my_add.cpp, my_add_kernel.cu], extra_compile_args{ cxx: [-O3], nvcc: [-O3] } ) ], cmdclass{build_ext: BuildExtension} )然后执行python setup.py build_ext --inplace。如果顺利会在文件夹里生成一个.pyd或.so文件之后就能import my_add。这里我想强调一下extra_compile_args如果你是第一次编译先别加花里胡哨的-arch标志让 PyTorch 自己根据torch.cuda.get_device_capability()决定。编译通过以后再慢慢调-gencodearchcompute_80,codesm_80这类参数这样能少踩很多“编译错误一半是熄火一半是警告”的坑。3. 核心实战C/CUDA 算子怎么落地3.1 用 C 扩展写一个带反向的算子我拿一个最简单但完整的例子说明实现z (x y) * w。这在 PyTorch 里当然一行就写完了但它的反向要回传三个梯度很适合用来展示 autograd 对接。先看 C 端定义#include torch/extension.h torch::Tensor my_forward(torch::Tensor x, torch::Tensor y, torch::Tensor w) { return (x y) * w; } std::vectortorch::Tensor my_backward(torch::Tensor grad_output, torch::Tensor x, torch::Tensor y, torch::Tensor w) { auto grad_w grad_output * (x y); return {grad_output * w, grad_output * w, grad_w}; }然后需要注册成一个 autograd 函数。在 PyTorch 的 C API 里可以使用torch::autograd::Function派生类class MyCombinedOp : public torch::autograd::FunctionMyCombinedOp { public: static torch::Tensor forward(AutogradContext* ctx, torch::Tensor x, torch::Tensor y, torch::Tensor w) { ctx-save_for_backward({x, y, w}); return my_forward(x, y, w); } static std::vectortorch::Tensor backward(AutogradContext* ctx, std::vectortorch::Tensor grad_output) { auto saved ctx-get_saved_variables(); auto x saved[0]; auto y saved[1]; auto w saved[2]; return my_backward(grad_output[0], x, y, w); } }; torch::Tensor my_combined_op(torch::Tensor x, torch::Tensor y, torch::Tensor w) { return MyCombinedOp::apply(x, y, w); } TORCH_LIBRARY(my_ops, m) { m.def(my_combined_op, my_combined_op); }注意forward和backward都是静态成员函数apply负责把参数传进去并触发 autograd 记录。save_for_backward保存前向用过的变量反向时调用get_saved_variables取出来。可能你会问为什么不直接在 Python 里写torch.autograd.Function原因很简单因为这里我们在演示 C 侧完成整个闭环这样编译后的算子不依赖 Python 解释器的 Function 机制调度性能更高也更接近原生 op。写好这段之后最重要的就是注册。上面代码用了TORCH_LIBRARY这是 PyTorch 1.10 推荐的方式。需要注意my_ops是命名空间Python 端调用时使用torch.ops.my_ops.my_combined_op。如果你希望让 Python 直接import my_add然后调用普通函数也可以用PYBIND11_MODULE但正如前面说的那样 autograd 不会自动帮你记录反向。通常的做法是要么用Function包一层并导出普通函数要么用TORCH_LIBRARY注册Python 端再包一个torch.autograd.Function定义反向。我用上面的写法是为了展示纯 C autograd实际项目中你完全可以根据需求混用。3.2 CUDA kernel 的写法与线程分配真正大规模计算时我们希望算子跑在 GPU 上。CUDA 端的基本模式是先写一个__global__函数作为 kernel再在 host 函数里设置 grid 和 block 大小最后把 Tensor 的 data_ptr 传给 kernel。以最简单的逐元素加法为例__global__ void vector_add_kernel(const float* a, const float* b, float* out, int n) { int idx blockIdx.x * blockDim.x threadIdx.x; if (idx n) { out[idx] a[idx] b[idx]; } } torch::Tensor vector_add_cuda(torch::Tensor a, torch::Tensor b) { auto out torch::empty_like(a); auto n a.numel(); const int threads 256; const int blocks (n threads - 1) / threads; vector_add_kernelblocks, threads( a.data_ptrfloat(), b.data_ptrfloat(), out.data_ptrfloat(), n); return out; }这段代码里最关键的两个参数是threads和blocks。threads我固定为 256是因为对大多数 GPU 来说一个 block 里 128 到 512 个线程通常能获得不错的占用率。blocks的计算方法是一个经典公式(元素总数 线程数 - 1) / 线程数这样能保证即使最后一个 block 不满也不会漏掉元素。kernel 内部通过if (idx n)做边界检查防止越界访问。真实工程里我强烈建议再加上 CUDA 错误检查。比如内核启动后立刻调用cudaGetLastError()或者AT_CUDA_CHECK(cudaDeviceSynchronize())不然一个越界访问可能不会立刻报错而是等到下一帧才给你一个莫名其妙的illegal memory access。另外a.data_ptrfloat()返回的是指针必须确保传入的 tensor 在 GPU 上、类型是 float、内存连续。一般我会在函数开头写三个检查TORCH_CHECK(a.is_cuda(), input must be CUDA tensor); TORCH_CHECK(a.scalar_type() torch::kFloat, input must be float32); TORCH_CHECK(a.is_contiguous(), input must be contiguous);这三个检查会避免 90% 的奇怪崩溃。如果发现输入不连续你可以直接用.contiguous()拷贝一份再传进去只是要记住这会增加一次内存拷贝性能敏感时要提前在 Python 侧保证数据布局。3.3 Autograd 怎么接管梯度有人问前向写完了反向怎么自动接上其实原理很简单。torch::autograd::Function派生类重写了forward和backward两个接口。当你调用apply的时候PyTorch 的 autograd 引擎会创建一个节点放在计算图里节点里保存了你通过save_for_backward留下的变量。反向传播时引擎从最终输出往输入方向走每经过一个节点就调用对应节点的backward把上游梯度传进去拿到回传给前向输入的梯度。我上面那个例子反向实现的是链式法则。z (x y) * w设g是上游梯度那么对x的梯度是g * w对y的梯度也是g * w对w的梯度是g * (x y)。可以看到我在my_backward里就按这个公式算。这里有个非常重要的细节backward返回的std::vectortorch::Tensor顺序必须和forward的输入顺序一致。也就是说forward接受的第一个参数 x、第二个 y、第三个 w那么backward返回的三元组就分别对应 x、y、w 的梯度顺序千万不能乱。我刚开始写的时候就在这儿翻过车梯度不报错但对不上号最后只能靠打印每个变量的梯度检查。另外需要注意自动微分对传进去的 Tensor 默认会做版本控制如果前向过程中你对保存的变量做了原地修改梯度计算会报错。如果你是优化内存的狂热爱好者用了set_data或者copy_这类操作一定要想清楚会不会破坏 autograd 的版本计数器。3.4 Python 端加载与调用如果你不想搞一整个 setup.py 工程PyTorch 提供了load_inline这个神器可以直接在 Python 脚本里把 C/CUDA 源码编译成一个模块。它的典型用法是from torch.utils.cpp_extension import load_inline cuda_source #include torch/extension.h #include cuda_runtime.h __global__ void add_kernel(const float* a, const float* b, float* out, int n) { int idx blockIdx.x * blockDim.x threadIdx.x; if (idx n) out[idx] a[idx] b[idx]; } torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b) { auto out torch::empty_like(a); const int threads 256; const int blocks (a.numel() threads - 1) / threads; add_kernelblocks, threads(a.data_ptrfloat(), b.data_ptrfloat(), out.data_ptrfloat(), a.numel()); return out; } module load_inline( nameinline_add_op, cuda_sources[cuda_source], functions[add_cuda], with_cudaTrue, extra_cuda_cflags[-O3], verboseFalse, ) result module.add_cuda(torch.randn(1000, devicecuda), torch.randn(1000, devicecuda))load_inline会临时在系统缓存目录里生成一个工程并编译好处是适合快速试验坏处是每次编译都要等几十秒而且不好控制版本依赖。如果你想长期维护一个算子我建议还是用标准 setup.py把源码放在项目里方便版本管理。我在实验阶段用load_inline验证通过以后立刻迁移到独立工程这样后面加测试和 benchmark 都很方便。4. 性能优化与调试4.1 从正确到高效访存优化三板斧很多人第一次写完自定义算子发现比 PyTorch 原生还慢就开始怀疑人生。其实绝大多数时候不是 C 的问题而是访存模式没写好。GPU 上的计算能力远超内存带宽所以逐元素算子基本上都是内存密集型。你写的 kernel 每多读一遍内存性能就掉一截优化访存是自定义算子性能提升的大头。第一板斧是“减少全局内存读写”。这个很好理解比如y (a * b) (a * c)如果写成两次 kernel就注定要把a读两遍、中间结果写一遍。正确做法是融合成一个 kernel只读一次a算完直接写出去。这种融合只靠 Python 层很难控制但在自定义 operator 里你可以自由调度。第二板斧是“向量化访问”。现代 GPU 支持float4这种 128 位宽度的加载一次能拿 4 个 float。把普通的标量循环改成float4读入可以有效减少指令数和请求次数。常见做法是把float*转成float4*不过要注意总线程数和向量宽度的整除关系通常要处理余数部分。这个优化对带宽受限的算子提升非常可观有时能从 60% 的利用率涨到 85% 以上。第三板斧是“使用 grid-stride loop”。当你的 tensor 元素数远超 GPU 能同时启动的线程数时与其让一个线程只处理一个元素不如让每个线程循环处理多个间隔的元素形成一个网格跨步循环。这样做可以减少 block 数量分摊 kernel 启动开销同时也有助于更好地利用 GPU 上不同的 SM 调度。最简单的形式是把idx gridDim.x * blockDim.x写进循环直到越界。这个模式我几乎在所有 elementwise kernel 里都用代码量没增加多少实测却很稳定。4.2 调试工具与常见坑自定义算子调试比纯 Python 痛苦主要是我们突然从 Python 的友好报错切换到了原生世界的空指针、段错误和 CUDA 错误码。我调试的固定套路是先加cudaDeviceSynchronize()兜底再开CUDA_LAUNCH_BLOCKING1。前者确保 kernel 同步后者让 CUDA 的异步错误变成同步错误这样崩溃位置能更准确。然后在关键位置用std::cerr 打印变量形状、数值、指针确认数据流没断。如果出现illegal memory access首选工具是compute-sanitizer。在终端跑compute-sanitizer --tool memcheck python test.py它能报告具体是哪个 kernel 哪一行访问非法。虽然跑起来会慢很多但找内存问题比盲猜效率高多了。还有一个容易被忽略的是“kernel 启动失败但不报错”比如blocks参数过大超过机器上限或者threads不是 32 的倍数。CUDA 有个规则就是 blockDim.x 最大 1024且应该是 warp size 的整数倍你用 100 或者 300 这种值看起来能跑实际上可能直接在启动阶段挂掉。所以安全做法是TORCH_CHECK(threads 1024)并且固定用 128、256 或 512。4.3 和现有 PyTorch 生态的融合自定义算子最终是要被人调用的。在 Python 层你可以做一个torch.nn.Module把它包起来也可以做一个普通函数。如果希望它参与 autograd就用Function包一层或者依靠自定义 forward 里返回能触发 C autograd 的节点。更复杂的场景是torch.compile它会尝试把 Python 操作图降低到计算图如果你自定义的算子没有注册进 dispatch 体系很可能被当成黑盒保留下这样计算图优化就不能穿越这个算子。所以如果要塞进一个需要被compile的模型尽量走TORCH_LIBRARY注册路线。不过按我实际踩坑的经验torch.compile对自定义 C 算子的支持还在快速演进。如果你只是想在一个稳定模型里跑推理更稳妥的方案是直接用torch.no_grad加普通 function 调用不要强行交给torch.compile去改。等确认框架版本的支持程度之后再升级集成方式也不迟。5. 常见问题与排查技巧实录5.1 版本不匹配与 ABI 问题自定义算子最让人头疼的就是“换个环境就编译不过”。我遇到过最典型的是 PyTorch 用 conda 装CUDA Toolkit 用系统装两个版本相差一个主版本。编译的时候 nvcc 用了 Toolkit 11.8但运行时 PyTorch 动态库链接的是 CUDA 12 的 runtime结果一执行就报CUDA driver version is insufficient或者干脆符号找不到。我的经验是torch.version.cuda显示什么你的 CUDA Toolkit 就尽量用什么哪怕你的显卡驱动已经支持更新的版本也别轻易混用。混合版本带来的收益远小于排查成本。还有 ABI 问题。如果你自己编译的 PyTorch 和预编译的 PyTorch 编译器版本不同一个用 GCC 9一个用 GCC 11链接时会出现一堆undefined symbol: _ZN2at4TensorC1ERKS_这类符号找不到的报错。解决办法也很简单尽量使用和 PyTorch 官方预编译环境相同的编译器。Linux 下官方一般用 GCC 7.3 到 9.3 的版本范围Windows 下要求 MSVC 版本不低于 2015。所以在折腾自定义算子之前先检查编译日志里提示的编译器版本再决定要不要升级。5.2 Windows 下 MSVC 的幺蛾子Windows 下的坑比 Linux 多得多。最常见的是error: identifier cl is undefined或者cl.exe not found。这多半是因为你在普通 cmd 或者 PowerShell 里直接跑python setup.py而cl.exe只在 Visual Studio 的开发者命令提示符里才有。解决方法是开始菜单里搜索“x64 Native Tools Command Prompt”打开后进入项目目录再执行编译命令。需要注意的是Python 默认是 64 位所以必须用 x64 的命令行如果用 x86 的命令行链接时会出现平台冲突。另一个坑是 Redistributable 的版本。经常有人把 VC 运行库装好但编译器还是缺失因为 Redistributable 只提供 DLL编译期工具链需要完整的 Visual Studio Build Tools。如果你用的是 VS Code 或者 JetBrains千万别以为编辑器能写代码就能编译C 扩展始终需要 MSVC 编译器。我在 Windows 上还遇到过.pyd生成到子目录、import找不到模块的情况检查一下build_ext --inplace的日志确认输出路径是否和脚本位置一致。实在不行就手动把生成的.pyd复制到项目根目录。5.3 多卡分布式里的自坑当你在DataParallel或者DistributedDataParallel里使用自定义 CUDA 算子时风险主要来自设备指针和 device guard。CUDA 的 kernel 是按 device 区分的如果你的函数在某个设备上分配了 Tensor却没把当前设备切换过去就可能访问错误的内存。解决方法是使用const c10::cuda::OptionalCUDAGuard device_guard(device_of(x));保证 kernel 在当前 tensor 所在的设备上运行。另一个坑是 stream。PyTorch 默认使用各自的 stream 来做异步操作如果你的自定义 kernel 是在 CUDA default stream 上启动而前后操作都在当前 PyTorch stream 上就可能出现未定义顺序表现为结果时对时错。安全起见在 host 函数里调用 CUDAGuard 的同时还需要用at::cuda::CUDAStream将当前 kernel 调度到与 PyTorch 一致的 stream 上。简单做法是auto stream at::cuda::getCurrentCUDAStream(); vector_add_kernelblocks, threads, 0, stream.stream()(...);这个细节如果漏掉单卡跑得好好的多卡一上就开始随机错误排查起来非常磨人。所以我建议大家把 device guard 和 stream 两件事写进自己的 kernel 模板里不要等出了问题再补。下面是一个常见问题速查表整理了我经常遇到的现象和对应的解法现象常见原因优先排查方向cl.exe not found没在 MSVC 命令行里编译使用 x64 Native Tools Command Promptundefined symbol编译器或 ABI 不匹配检查 gcc/MSVC 版本与 PyTorch 一致性illegal memory access越界、非连续、错误 device加 TORCH_CHECK用 compute-sanitizer多卡结果时好时坏未设置 CUDA stream使用当前 stream 启动 kernel5.4 ONNX 导出自定义算子进阶如果你的自定义算子最终要部署到 ONNX Runtime或者要转换成其他推理引擎直接导出是导不出来的因为 ONNX 没有你的自定义节点。PyTorch 提供了一个 symbolic 注册机制把自定义算子映射到既有 ONNX op或者自定义 domain 下的 op。你需要写一个 symbolic 函数告诉 TorchScript 在导出时如何表示这个节点。以最简单的加法为例如果自定义算子本质上就是 Add你可以这么写from torch.onnx import register_custom_op_symbolic def my_add_symbolic(g, x, y): return g.op(Add, x, y) register_custom_op_symbolic(my_ops::my_combined_op, my_add_symbolic, opset_version17)更复杂的算子比如拉普拉斯算子这种本地滤波通常会在自定义 domainai.onnx.contrib下生成一个CustomOp节点然后由推理后端去实现对应 kernel。导出时一般需要传入custom_opsets{ai.onnx.contrib: 1}。这里要特别注意算子版本的约定onnx op 版本、PyTorch 版本、后端实现版本三者的对应关系。我在实践中建议导出前先在小模型上试跑确认节点格式和 ONNX Runtime 支持的 opset 匹配否则部署时还是容易踩到unsupported operator的坑。6. 进阶战术融合算子与性能基线6.1 为什么要融合前面讲了很多理论最后我用一个具体例子说清楚“融合”的收益。假设我们有三个输入a、b、c要计算d a * b c。如果拆成两条 PyTorch kernel第一步tmp a * b第二步d tmp c内存读写次数是读a、读b、写tmp再读tmp、读c、写d一共至少 6 次全局内存操作。如果写一个融合 kernel一次循环里同时读a、b、c算出a*bc后直接写d全局内存操作是 4 次少了三分之一。这种差距在超大 tensor 上尤其明显因为内存带宽就是瓶颈。写法其实很简单就是把两个 kernel 的逻辑合并到一个__global__函数里__global__ void fused_mul_add_kernel(const float* a, const float* b, const float* c, float* d, int n) { int idx blockIdx.x * blockDim.x threadIdx.x; if (idx n) { d[idx] a[idx] * b[idx] c[idx]; } }编译器通常会把它优化成一条 FMA 指令计算吞吐也更高。更复杂的融合例如 softmax 融合 scale 和 masked fill收益会更明显。不过要注意融合后不能再从中间结果做别的分支因为你省掉的就是那份中间结果。这是设计层面需要权衡的。6.2 性能基线的正确打开方式写算子一定要测准不然根本不知道优化有没有效果。我推荐用 PyTorch 自带的torch.profiler配合torch.cuda.Event做 baseline。torch.cuda.Event的用法很简单start torch.cuda.Event(enable_timingTrue) end torch.cuda.Event(enable_timingTrue) start.record() result fused_op(a, b, c) end.record() torch.cuda.synchronize() print(ffused: {start.elapsed_time(end):.3f} ms)注意一定要torch.cuda.synchronize()不然异步执行会报一个极不合理的时间。除了测总耗时我还会用torch.profiler看 kernel 的 achieved occupancy、memory throughput、dram bytes read 这类指标判断瓶颈究竟在计算还是访存。工具输出的数据里memory throughput 接近峰值而 compute throughput 很低说明算子已经带宽饱和再优化计算没有意义。这时候能优化的方向只有减少访存比如用更低精度、压缩数据或者改用 shared memory 做数据复用。测 baseline 时还要注意排除掉第一次调用时的 lazy initialization。PyTorch 的 CUDA context 启动、算子第一次加载都有一笔一次性开销所以正式测试前先跑几十次 warm-up让 GPU 内核都准备好再记录后面若干次的时间。这个过程没人爱做但只有这么做才能得到可信的数字否则你会被第一轮的 200ms 欺骗误以为自己的算子写得很差。我还会把每次测到的耗时都存进一个列表三次取中位数而不是平均值因为 GPU 上偶尔会有背景任务干扰中位数更稳定。我个人折腾下来最深的感受是自定义算子的难点从来不在写 kernel而在于你愿不愿意把 Python 以外的工具链弄明白。当你第一次绕过 GIL、直接在 C 层把 Tensor 逻辑跑通再回头看nn.Sequential就像看说明书。下一步你可以试试把我这里的最小示例改成原子操作的累加或者给算子加一个 workspace 分配能力就是这么慢慢长出来的。
1. 为什么需要自定义算子1.1 PyTorch 算子的两种来源我一直觉得很多人对“自定义算子”这个词的第一反应是“这是那些做框架的人才会碰的东西”。实际上 PyTorch 每天在用的所有功能无论是torch.add还是conv2d本质上都能归成两类一类是官方已经编译进 C 核心的算子另一类是用 Python 组合出来的高层模块。官方算子覆盖了 95% 的日常需求可一旦你开始做模型部署、性能优化、或者写一篇新论文里的特殊函数官方算子经常不够用这时候你就会想自己动手写一个。自定义算子的含义其实很朴素我们自己写一段 C 代码或者 CUDA 代码把它编译成一个 Python 可以调用的函数。这个函数和torch.sigmoid长得一模一样能接收 Tensor、返回 Tensor差别在于它的计算逻辑由你来定义。这意味着你可以把 Python 层很难表达的循环、内存访问模式、或者某个第三方 C 库原封不动地接进 PyTorch 的计算图里。我这里说的“吃透”不是只讲 API 怎么调而是想把背后的机制拆出来PyTorch 怎么把你的 C 函数变成 Python 能导入的模块CUDA kernel 怎么在 GPU 上被调度反向传播怎么接到 autograd 上只有把这三件事弄明白你才能在一张白纸上设计自己的算子而不是照着官方示例改个名字就完事。1.2 什么场景必须自己写算子判断依据其实很简单遇到下面四种情况你就要认真考虑上自定义算子。第一种是“计算逻辑没法用 PyTorch 原生算子拼出来”。比如说你要做某种自定义的稀疏索引操作或者实现一个拉普拉斯算子的变体需要逐个遍历邻域像素并做条件更新用 Python 写for循环在 GPU 上会直接卡死用torch操作东拼西凑又很别扭这种情况写成 CUDA kernel 才是正路。第二种是“能拼出来但性能差得离谱”。最典型的就是多个逐元素算子串联比如y a * b c如果拆成mul和add两个 kernelGPU 要启动两次、读写两遍内存带宽成本直接翻倍。如果把它融合成一个 CUDA kernel一次遍历全做完速度可能提升 20% 甚至 50% 以上。我在做图像滤波的时候遇到过类似问题一个简单的权重滤波算子拆开写处理一张 4K 图像每次都要多等两个 kernel launch融合之后肉眼可见的流畅。第三种是“要复用一个已有的 C/CUDA 库”。比如某些工业软件里的 Halcon 滤波核或者自研相机 SDK 已经提供了高效的 C 函数你不想把它翻译成 Python更不想用 ctypes 一点点抠内存直接从 C 扩展里包一层 Tensor 封装是最高效的接法。第四种是“需要自定义反向传播逻辑”。PyTorch 的 autograd 已经把所有标准算子都定义了梯度但如果你的前向过程有分段不可导点或者你想在反向里注入某种额外的梯度变换就必须自己控制反向函数。自定义算子里的backward让你拥有绝对主动权。这里我给一个实用建议不要为了自定义而自定义。如果你的需求能用torch.compile或者torch.vmap解决优先用这些高层工具。它们的开发成本低而且能自动做算子级优化。只有当你需要极致的控制权时才决定走 C/CUDA 这条路。这个判断会帮你省掉大量调试时间。2. 动手前的工具链准备2.1 理解 PyTorch 内置的扩展机制PyTorch 官方早就把“自定义算子”的通道给铺好了核心就是torch.utils.cpp_extension这个模块。它本质上是一个编译器包装器拿到你的 C 源文件和 CUDA 源文件之后会调用编译器把它们编译成动态库再用 pybind11 把函数导出成 Python 可调用的对象。你不需要手动写完整的setup.py和 pybind11 绑定虽然自己写也可以但官方模块已经处理了大部分坑。这套机制里其实还有底层注册逻辑叫TORCH_LIBRARY。通过它可以在 PyTorch 内部注册一个带命名空间的算子让 PyTorch 的调度器认识你。官方推荐的现代写法是TORCH_LIBRARY(my_ops, m) { m.def(custom_add, custom_add); }然后用torch.ops.my_ops.custom_add在 Python 端访问。很多老教程喜欢用PYBIND11_MODULE直接导出普通函数差别在于PYBIND11_MODULE导出的函数不参与 PyTorch 的 op 调度也没有办法被torch.jit或者torch.compile完整追踪。如果你只是给自己跑实验怎么搞都行如果要进正式模型或者部署链路建议一开始就用TORCH_LIBRARY的姿势。2.2 版本匹配与依赖安装PyTorch 自定义 C/CUDA 算子最怕的就是版本错位。CUDA 编译器必须和 PyTorch 编译时使用的 CUDA 版本兼容不然会链接到一半报undefined reference to cudaMemcpy或者显卡驱动不支持跑起来直接报CUDA driver version is insufficient。建议先跑一下python -c import torch; print(torch.version.cuda)拿到当前 PyTorch 对应的 CUDA 版本号。如果你使用的是 conda 安装的 PyTorch通常自带配套 CUDA 库安装 CUDA Toolkit 时选择相近版本即可。Windows 下还有一个隐藏依赖Microsoft Visual C Redistributable。很多人奇怪为什么装了个 Python 包还要管 VC 运行库因为 PyTorch 的 C 扩展在 Windows 上依赖 MSVC 编译器和运行时库。建议提前安装 Visual Studio 2019 或 2022尤其要注意勾选“使用 C 的桌面开发”工作负载否则命令行里找不到cl.exe。至于 Redistributable它一般会随 VS 一起装好单独下载也有用但真正编译时你需要的是编译器不是运行库。把这两者分开理解能少翻很多帖子。我在 Ubuntu 上工作的时间比较长那边主要用 g。但说实话Ubuntu 上的坑也不比 Windows 少最典型的是系统 gcc 和 PyTorch 官方预编译用的 gcc 版本不一致后面编译时你会看到一大堆莫名其妙的符号找不到。我的建议是尽量用系统默认的编译器版本或者参考 PyTorch 官方文档里的支持矩阵。避坑的核心就是一句话让你的编译环境和 PyTorch 的编译环境保持同频。2.3 从零搭建一个最小可编译工程我平时写自定义算子文件夹结构非常简单。无论你是要用setup.py还是load_inline都有几样东西是必备的一个.cpp文件负责定义 C 函数和绑定一个.cu文件负责写 CUDA kernel一个setup.py负责告诉编译器怎么组织和输出。目录大致长这样my_op/ ├── setup.py ├── my_add.cpp ├── my_add_kernel.cu └── test.pysetup.py最简单的写法如下from torch.utils.cpp_extension import CUDAExtension, BuildExtension from setuptools import setup setup( namemy_add, ext_modules[ CUDAExtension( namemy_add, sources[my_add.cpp, my_add_kernel.cu], extra_compile_args{ cxx: [-O3], nvcc: [-O3] } ) ], cmdclass{build_ext: BuildExtension} )然后执行python setup.py build_ext --inplace。如果顺利会在文件夹里生成一个.pyd或.so文件之后就能import my_add。这里我想强调一下extra_compile_args如果你是第一次编译先别加花里胡哨的-arch标志让 PyTorch 自己根据torch.cuda.get_device_capability()决定。编译通过以后再慢慢调-gencodearchcompute_80,codesm_80这类参数这样能少踩很多“编译错误一半是熄火一半是警告”的坑。3. 核心实战C/CUDA 算子怎么落地3.1 用 C 扩展写一个带反向的算子我拿一个最简单但完整的例子说明实现z (x y) * w。这在 PyTorch 里当然一行就写完了但它的反向要回传三个梯度很适合用来展示 autograd 对接。先看 C 端定义#include torch/extension.h torch::Tensor my_forward(torch::Tensor x, torch::Tensor y, torch::Tensor w) { return (x y) * w; } std::vectortorch::Tensor my_backward(torch::Tensor grad_output, torch::Tensor x, torch::Tensor y, torch::Tensor w) { auto grad_w grad_output * (x y); return {grad_output * w, grad_output * w, grad_w}; }然后需要注册成一个 autograd 函数。在 PyTorch 的 C API 里可以使用torch::autograd::Function派生类class MyCombinedOp : public torch::autograd::FunctionMyCombinedOp { public: static torch::Tensor forward(AutogradContext* ctx, torch::Tensor x, torch::Tensor y, torch::Tensor w) { ctx-save_for_backward({x, y, w}); return my_forward(x, y, w); } static std::vectortorch::Tensor backward(AutogradContext* ctx, std::vectortorch::Tensor grad_output) { auto saved ctx-get_saved_variables(); auto x saved[0]; auto y saved[1]; auto w saved[2]; return my_backward(grad_output[0], x, y, w); } }; torch::Tensor my_combined_op(torch::Tensor x, torch::Tensor y, torch::Tensor w) { return MyCombinedOp::apply(x, y, w); } TORCH_LIBRARY(my_ops, m) { m.def(my_combined_op, my_combined_op); }注意forward和backward都是静态成员函数apply负责把参数传进去并触发 autograd 记录。save_for_backward保存前向用过的变量反向时调用get_saved_variables取出来。可能你会问为什么不直接在 Python 里写torch.autograd.Function原因很简单因为这里我们在演示 C 侧完成整个闭环这样编译后的算子不依赖 Python 解释器的 Function 机制调度性能更高也更接近原生 op。写好这段之后最重要的就是注册。上面代码用了TORCH_LIBRARY这是 PyTorch 1.10 推荐的方式。需要注意my_ops是命名空间Python 端调用时使用torch.ops.my_ops.my_combined_op。如果你希望让 Python 直接import my_add然后调用普通函数也可以用PYBIND11_MODULE但正如前面说的那样 autograd 不会自动帮你记录反向。通常的做法是要么用Function包一层并导出普通函数要么用TORCH_LIBRARY注册Python 端再包一个torch.autograd.Function定义反向。我用上面的写法是为了展示纯 C autograd实际项目中你完全可以根据需求混用。3.2 CUDA kernel 的写法与线程分配真正大规模计算时我们希望算子跑在 GPU 上。CUDA 端的基本模式是先写一个__global__函数作为 kernel再在 host 函数里设置 grid 和 block 大小最后把 Tensor 的 data_ptr 传给 kernel。以最简单的逐元素加法为例__global__ void vector_add_kernel(const float* a, const float* b, float* out, int n) { int idx blockIdx.x * blockDim.x threadIdx.x; if (idx n) { out[idx] a[idx] b[idx]; } } torch::Tensor vector_add_cuda(torch::Tensor a, torch::Tensor b) { auto out torch::empty_like(a); auto n a.numel(); const int threads 256; const int blocks (n threads - 1) / threads; vector_add_kernelblocks, threads( a.data_ptrfloat(), b.data_ptrfloat(), out.data_ptrfloat(), n); return out; }这段代码里最关键的两个参数是threads和blocks。threads我固定为 256是因为对大多数 GPU 来说一个 block 里 128 到 512 个线程通常能获得不错的占用率。blocks的计算方法是一个经典公式(元素总数 线程数 - 1) / 线程数这样能保证即使最后一个 block 不满也不会漏掉元素。kernel 内部通过if (idx n)做边界检查防止越界访问。真实工程里我强烈建议再加上 CUDA 错误检查。比如内核启动后立刻调用cudaGetLastError()或者AT_CUDA_CHECK(cudaDeviceSynchronize())不然一个越界访问可能不会立刻报错而是等到下一帧才给你一个莫名其妙的illegal memory access。另外a.data_ptrfloat()返回的是指针必须确保传入的 tensor 在 GPU 上、类型是 float、内存连续。一般我会在函数开头写三个检查TORCH_CHECK(a.is_cuda(), input must be CUDA tensor); TORCH_CHECK(a.scalar_type() torch::kFloat, input must be float32); TORCH_CHECK(a.is_contiguous(), input must be contiguous);这三个检查会避免 90% 的奇怪崩溃。如果发现输入不连续你可以直接用.contiguous()拷贝一份再传进去只是要记住这会增加一次内存拷贝性能敏感时要提前在 Python 侧保证数据布局。3.3 Autograd 怎么接管梯度有人问前向写完了反向怎么自动接上其实原理很简单。torch::autograd::Function派生类重写了forward和backward两个接口。当你调用apply的时候PyTorch 的 autograd 引擎会创建一个节点放在计算图里节点里保存了你通过save_for_backward留下的变量。反向传播时引擎从最终输出往输入方向走每经过一个节点就调用对应节点的backward把上游梯度传进去拿到回传给前向输入的梯度。我上面那个例子反向实现的是链式法则。z (x y) * w设g是上游梯度那么对x的梯度是g * w对y的梯度也是g * w对w的梯度是g * (x y)。可以看到我在my_backward里就按这个公式算。这里有个非常重要的细节backward返回的std::vectortorch::Tensor顺序必须和forward的输入顺序一致。也就是说forward接受的第一个参数 x、第二个 y、第三个 w那么backward返回的三元组就分别对应 x、y、w 的梯度顺序千万不能乱。我刚开始写的时候就在这儿翻过车梯度不报错但对不上号最后只能靠打印每个变量的梯度检查。另外需要注意自动微分对传进去的 Tensor 默认会做版本控制如果前向过程中你对保存的变量做了原地修改梯度计算会报错。如果你是优化内存的狂热爱好者用了set_data或者copy_这类操作一定要想清楚会不会破坏 autograd 的版本计数器。3.4 Python 端加载与调用如果你不想搞一整个 setup.py 工程PyTorch 提供了load_inline这个神器可以直接在 Python 脚本里把 C/CUDA 源码编译成一个模块。它的典型用法是from torch.utils.cpp_extension import load_inline cuda_source #include torch/extension.h #include cuda_runtime.h __global__ void add_kernel(const float* a, const float* b, float* out, int n) { int idx blockIdx.x * blockDim.x threadIdx.x; if (idx n) out[idx] a[idx] b[idx]; } torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b) { auto out torch::empty_like(a); const int threads 256; const int blocks (a.numel() threads - 1) / threads; add_kernelblocks, threads(a.data_ptrfloat(), b.data_ptrfloat(), out.data_ptrfloat(), a.numel()); return out; } module load_inline( nameinline_add_op, cuda_sources[cuda_source], functions[add_cuda], with_cudaTrue, extra_cuda_cflags[-O3], verboseFalse, ) result module.add_cuda(torch.randn(1000, devicecuda), torch.randn(1000, devicecuda))load_inline会临时在系统缓存目录里生成一个工程并编译好处是适合快速试验坏处是每次编译都要等几十秒而且不好控制版本依赖。如果你想长期维护一个算子我建议还是用标准 setup.py把源码放在项目里方便版本管理。我在实验阶段用load_inline验证通过以后立刻迁移到独立工程这样后面加测试和 benchmark 都很方便。4. 性能优化与调试4.1 从正确到高效访存优化三板斧很多人第一次写完自定义算子发现比 PyTorch 原生还慢就开始怀疑人生。其实绝大多数时候不是 C 的问题而是访存模式没写好。GPU 上的计算能力远超内存带宽所以逐元素算子基本上都是内存密集型。你写的 kernel 每多读一遍内存性能就掉一截优化访存是自定义算子性能提升的大头。第一板斧是“减少全局内存读写”。这个很好理解比如y (a * b) (a * c)如果写成两次 kernel就注定要把a读两遍、中间结果写一遍。正确做法是融合成一个 kernel只读一次a算完直接写出去。这种融合只靠 Python 层很难控制但在自定义 operator 里你可以自由调度。第二板斧是“向量化访问”。现代 GPU 支持float4这种 128 位宽度的加载一次能拿 4 个 float。把普通的标量循环改成float4读入可以有效减少指令数和请求次数。常见做法是把float*转成float4*不过要注意总线程数和向量宽度的整除关系通常要处理余数部分。这个优化对带宽受限的算子提升非常可观有时能从 60% 的利用率涨到 85% 以上。第三板斧是“使用 grid-stride loop”。当你的 tensor 元素数远超 GPU 能同时启动的线程数时与其让一个线程只处理一个元素不如让每个线程循环处理多个间隔的元素形成一个网格跨步循环。这样做可以减少 block 数量分摊 kernel 启动开销同时也有助于更好地利用 GPU 上不同的 SM 调度。最简单的形式是把idx gridDim.x * blockDim.x写进循环直到越界。这个模式我几乎在所有 elementwise kernel 里都用代码量没增加多少实测却很稳定。4.2 调试工具与常见坑自定义算子调试比纯 Python 痛苦主要是我们突然从 Python 的友好报错切换到了原生世界的空指针、段错误和 CUDA 错误码。我调试的固定套路是先加cudaDeviceSynchronize()兜底再开CUDA_LAUNCH_BLOCKING1。前者确保 kernel 同步后者让 CUDA 的异步错误变成同步错误这样崩溃位置能更准确。然后在关键位置用std::cerr 打印变量形状、数值、指针确认数据流没断。如果出现illegal memory access首选工具是compute-sanitizer。在终端跑compute-sanitizer --tool memcheck python test.py它能报告具体是哪个 kernel 哪一行访问非法。虽然跑起来会慢很多但找内存问题比盲猜效率高多了。还有一个容易被忽略的是“kernel 启动失败但不报错”比如blocks参数过大超过机器上限或者threads不是 32 的倍数。CUDA 有个规则就是 blockDim.x 最大 1024且应该是 warp size 的整数倍你用 100 或者 300 这种值看起来能跑实际上可能直接在启动阶段挂掉。所以安全做法是TORCH_CHECK(threads 1024)并且固定用 128、256 或 512。4.3 和现有 PyTorch 生态的融合自定义算子最终是要被人调用的。在 Python 层你可以做一个torch.nn.Module把它包起来也可以做一个普通函数。如果希望它参与 autograd就用Function包一层或者依靠自定义 forward 里返回能触发 C autograd 的节点。更复杂的场景是torch.compile它会尝试把 Python 操作图降低到计算图如果你自定义的算子没有注册进 dispatch 体系很可能被当成黑盒保留下这样计算图优化就不能穿越这个算子。所以如果要塞进一个需要被compile的模型尽量走TORCH_LIBRARY注册路线。不过按我实际踩坑的经验torch.compile对自定义 C 算子的支持还在快速演进。如果你只是想在一个稳定模型里跑推理更稳妥的方案是直接用torch.no_grad加普通 function 调用不要强行交给torch.compile去改。等确认框架版本的支持程度之后再升级集成方式也不迟。5. 常见问题与排查技巧实录5.1 版本不匹配与 ABI 问题自定义算子最让人头疼的就是“换个环境就编译不过”。我遇到过最典型的是 PyTorch 用 conda 装CUDA Toolkit 用系统装两个版本相差一个主版本。编译的时候 nvcc 用了 Toolkit 11.8但运行时 PyTorch 动态库链接的是 CUDA 12 的 runtime结果一执行就报CUDA driver version is insufficient或者干脆符号找不到。我的经验是torch.version.cuda显示什么你的 CUDA Toolkit 就尽量用什么哪怕你的显卡驱动已经支持更新的版本也别轻易混用。混合版本带来的收益远小于排查成本。还有 ABI 问题。如果你自己编译的 PyTorch 和预编译的 PyTorch 编译器版本不同一个用 GCC 9一个用 GCC 11链接时会出现一堆undefined symbol: _ZN2at4TensorC1ERKS_这类符号找不到的报错。解决办法也很简单尽量使用和 PyTorch 官方预编译环境相同的编译器。Linux 下官方一般用 GCC 7.3 到 9.3 的版本范围Windows 下要求 MSVC 版本不低于 2015。所以在折腾自定义算子之前先检查编译日志里提示的编译器版本再决定要不要升级。5.2 Windows 下 MSVC 的幺蛾子Windows 下的坑比 Linux 多得多。最常见的是error: identifier cl is undefined或者cl.exe not found。这多半是因为你在普通 cmd 或者 PowerShell 里直接跑python setup.py而cl.exe只在 Visual Studio 的开发者命令提示符里才有。解决方法是开始菜单里搜索“x64 Native Tools Command Prompt”打开后进入项目目录再执行编译命令。需要注意的是Python 默认是 64 位所以必须用 x64 的命令行如果用 x86 的命令行链接时会出现平台冲突。另一个坑是 Redistributable 的版本。经常有人把 VC 运行库装好但编译器还是缺失因为 Redistributable 只提供 DLL编译期工具链需要完整的 Visual Studio Build Tools。如果你用的是 VS Code 或者 JetBrains千万别以为编辑器能写代码就能编译C 扩展始终需要 MSVC 编译器。我在 Windows 上还遇到过.pyd生成到子目录、import找不到模块的情况检查一下build_ext --inplace的日志确认输出路径是否和脚本位置一致。实在不行就手动把生成的.pyd复制到项目根目录。5.3 多卡分布式里的自坑当你在DataParallel或者DistributedDataParallel里使用自定义 CUDA 算子时风险主要来自设备指针和 device guard。CUDA 的 kernel 是按 device 区分的如果你的函数在某个设备上分配了 Tensor却没把当前设备切换过去就可能访问错误的内存。解决方法是使用const c10::cuda::OptionalCUDAGuard device_guard(device_of(x));保证 kernel 在当前 tensor 所在的设备上运行。另一个坑是 stream。PyTorch 默认使用各自的 stream 来做异步操作如果你的自定义 kernel 是在 CUDA default stream 上启动而前后操作都在当前 PyTorch stream 上就可能出现未定义顺序表现为结果时对时错。安全起见在 host 函数里调用 CUDAGuard 的同时还需要用at::cuda::CUDAStream将当前 kernel 调度到与 PyTorch 一致的 stream 上。简单做法是auto stream at::cuda::getCurrentCUDAStream(); vector_add_kernelblocks, threads, 0, stream.stream()(...);这个细节如果漏掉单卡跑得好好的多卡一上就开始随机错误排查起来非常磨人。所以我建议大家把 device guard 和 stream 两件事写进自己的 kernel 模板里不要等出了问题再补。下面是一个常见问题速查表整理了我经常遇到的现象和对应的解法现象常见原因优先排查方向cl.exe not found没在 MSVC 命令行里编译使用 x64 Native Tools Command Promptundefined symbol编译器或 ABI 不匹配检查 gcc/MSVC 版本与 PyTorch 一致性illegal memory access越界、非连续、错误 device加 TORCH_CHECK用 compute-sanitizer多卡结果时好时坏未设置 CUDA stream使用当前 stream 启动 kernel5.4 ONNX 导出自定义算子进阶如果你的自定义算子最终要部署到 ONNX Runtime或者要转换成其他推理引擎直接导出是导不出来的因为 ONNX 没有你的自定义节点。PyTorch 提供了一个 symbolic 注册机制把自定义算子映射到既有 ONNX op或者自定义 domain 下的 op。你需要写一个 symbolic 函数告诉 TorchScript 在导出时如何表示这个节点。以最简单的加法为例如果自定义算子本质上就是 Add你可以这么写from torch.onnx import register_custom_op_symbolic def my_add_symbolic(g, x, y): return g.op(Add, x, y) register_custom_op_symbolic(my_ops::my_combined_op, my_add_symbolic, opset_version17)更复杂的算子比如拉普拉斯算子这种本地滤波通常会在自定义 domainai.onnx.contrib下生成一个CustomOp节点然后由推理后端去实现对应 kernel。导出时一般需要传入custom_opsets{ai.onnx.contrib: 1}。这里要特别注意算子版本的约定onnx op 版本、PyTorch 版本、后端实现版本三者的对应关系。我在实践中建议导出前先在小模型上试跑确认节点格式和 ONNX Runtime 支持的 opset 匹配否则部署时还是容易踩到unsupported operator的坑。6. 进阶战术融合算子与性能基线6.1 为什么要融合前面讲了很多理论最后我用一个具体例子说清楚“融合”的收益。假设我们有三个输入a、b、c要计算d a * b c。如果拆成两条 PyTorch kernel第一步tmp a * b第二步d tmp c内存读写次数是读a、读b、写tmp再读tmp、读c、写d一共至少 6 次全局内存操作。如果写一个融合 kernel一次循环里同时读a、b、c算出a*bc后直接写d全局内存操作是 4 次少了三分之一。这种差距在超大 tensor 上尤其明显因为内存带宽就是瓶颈。写法其实很简单就是把两个 kernel 的逻辑合并到一个__global__函数里__global__ void fused_mul_add_kernel(const float* a, const float* b, const float* c, float* d, int n) { int idx blockIdx.x * blockDim.x threadIdx.x; if (idx n) { d[idx] a[idx] * b[idx] c[idx]; } }编译器通常会把它优化成一条 FMA 指令计算吞吐也更高。更复杂的融合例如 softmax 融合 scale 和 masked fill收益会更明显。不过要注意融合后不能再从中间结果做别的分支因为你省掉的就是那份中间结果。这是设计层面需要权衡的。6.2 性能基线的正确打开方式写算子一定要测准不然根本不知道优化有没有效果。我推荐用 PyTorch 自带的torch.profiler配合torch.cuda.Event做 baseline。torch.cuda.Event的用法很简单start torch.cuda.Event(enable_timingTrue) end torch.cuda.Event(enable_timingTrue) start.record() result fused_op(a, b, c) end.record() torch.cuda.synchronize() print(ffused: {start.elapsed_time(end):.3f} ms)注意一定要torch.cuda.synchronize()不然异步执行会报一个极不合理的时间。除了测总耗时我还会用torch.profiler看 kernel 的 achieved occupancy、memory throughput、dram bytes read 这类指标判断瓶颈究竟在计算还是访存。工具输出的数据里memory throughput 接近峰值而 compute throughput 很低说明算子已经带宽饱和再优化计算没有意义。这时候能优化的方向只有减少访存比如用更低精度、压缩数据或者改用 shared memory 做数据复用。测 baseline 时还要注意排除掉第一次调用时的 lazy initialization。PyTorch 的 CUDA context 启动、算子第一次加载都有一笔一次性开销所以正式测试前先跑几十次 warm-up让 GPU 内核都准备好再记录后面若干次的时间。这个过程没人爱做但只有这么做才能得到可信的数字否则你会被第一轮的 200ms 欺骗误以为自己的算子写得很差。我还会把每次测到的耗时都存进一个列表三次取中位数而不是平均值因为 GPU 上偶尔会有背景任务干扰中位数更稳定。我个人折腾下来最深的感受是自定义算子的难点从来不在写 kernel而在于你愿不愿意把 Python 以外的工具链弄明白。当你第一次绕过 GIL、直接在 C 层把 Tensor 逻辑跑通再回头看nn.Sequential就像看说明书。下一步你可以试试把我这里的最小示例改成原子操作的累加或者给算子加一个 workspace 分配能力就是这么慢慢长出来的。
