前段时间接手了一个视频预处理改造数据已经躺在 GPU 上但中间有一段缩放偏移截断的自定义逻辑PyTorch 里找不到现成算子只能用a.mul(scale).add(bias).clamp(min, max)硬堆。功能没问题可训练循环一跑起来每次前向和反向都要在 Python 与 C 之间来回穿梭还额外产生好几块中间张量显存和延迟都肉眼可见地涨。折腾了两周把 PyTorch 自定义 C/CUDA 算子的整条链路摸了一遍这才搞清楚哪些东西靠官方文档根本不够用。如果你也遇到过类似情况——Python 循环太慢、torch 组合算子中间张量太多、或者想写一个 CUDA kernel 却不知道怎么接进自动求导——这篇文章应该能帮你少走弯路。我会按自己实际操作时的顺序来写先判断是不是真的需要自定义算子然后从最小的 C 扩展写起再深入 CUDA kernel、反向传播最后给一份实测对比和踩坑清单。适合有一定 PyTorch 使用基础、但没碰过 C 扩展的读者。1. 先算一笔账什么时候自定义算子真的值1.1 Python 慢到底慢在哪很多人一说性能问题就怪 Python 解释器这其实只对了一半。拿y clamp(x * scale bias, min, max)这种操作来说如果用原生 torch 组合实现底层调用的是aten里现成 kernel每个 kernel 本身并不慢。慢的是这三层Python 侧类型检查、dispatch key 计算、参数包装每次调用torch.add、torch.mul、torch.clamp都会单独启动一次 kernel每个 kernel 都会产生一块新的中间张量频繁分配显存反向时还要再算一遍中间值内存压力跟着翻倍。举个直观例子1M 个 float 的张量单次 elementwise kernel 在 GPU 上可能只要几十微秒但 Python 包装一层就要额外几微秒三个算子叠加就是三倍的 launch overhead。数据量小的时候Python 调用开销甚至比 kernel 本身还大。这就像你出门前要换三次鞋每次都要把鞋放回鞋柜再拿下一双真正走路的时间反而没多少。自定义 C/CUDA 算子的核心价值就是把多个 kernel 融合成一个减少 Python 层调用次数、减少中间张量分配、按需求直接写并行逻辑。1.2 四套方案的取舍想提速不等于一上来就写 CUDA。我在动手前一般会列个表先看有没有更轻的方案方案开发成本性能潜力适合场景纯 Python torch 组合低低快速验证、小数据量torch.compile / torch.jit中中高现有算子能表达但希望自动融合C CPU 扩展中中CPU 部署、复杂逻辑、无 GPU 依赖C/CUDA 扩展高高自定义 kernel、融合算子、特殊内存访问torch.compile 这两年进步很快很多简单逐点算子能被 Inductor 自动融合不需要手写。但遇到以下情况还是得走自定义算子算子逻辑太特殊纯 PyTorch 根本表达不了需要对显存布局做精细控制算子内部涉及第三方 C/C 库或者torch.compile生成的 kernel 达不到你要的极致性能。我后面所有示例都会围绕fused_scale_bias_clamp这个融合算子展开它简单到能让你看清机制又实用到覆盖 elementwise 算子、dtype dispatch、CUDA kernel 和反向传播这几个核心环节。2. 最小 C 算子环境、编译与第一个可运行例子2.1 环境准备和版本匹配写 C 扩展前先确认三件事PyTorch 头文件用 pip 装的 PyTorch 自带全套 C 头文件不需要额外下载 libtorch。编译器Linux 上推荐 g 7.3 以上或 clangWindows 需要 VS2019/2022 并安装C 桌面开发工作负载装好后 Visual C Redistributable 也会带上macOS 用自带 clang。CUDA Toolkit写 CUDA 算子才需要装 nvcc。注意torch.version.cuda和nvcc --version的主版本尽量一致不一致容易出现编译能过、运行时找不到符号的诡异问题。很多 Windows 用户是在 WSL2 里做 CUDA 开发驱动由 Windows 侧提供编译和运行都在 Linux 环境里整体体验更接近生产环境。我自己的经验是一旦涉及 CUDA优先用 Linux 或 WSL2别在 Windows 原生环境下跟 MSVC 较劲。提示编译前先跑一行python -c import torch; print(torch.__version__, torch.version.cuda)把这个输出记下来后面排查问题要对得上。2.2 最小实现setup.py、C 源码、Python 调用先写一个最简算子fused_scale_bias功能就是out x * scale bias。文件结构非常简单fused_ops/ ├── setup.py ├── fused_ops.cpp └── test.pysetup.py里直接使用 PyTorch 提供的CppExtension和BuildExtension比手写Extension省事得多它会自动帮你加-DTORCH_EXTENSION_NAME、链接 PyTorch 库、找到 Python 头文件。from setuptools import setup from torch.utils.cpp_extension import CppExtension, BuildExtension setup( namefused_ops, ext_modules[ CppExtension( fused_ops, sources[fused_ops.cpp], extra_compile_args[-O3], ) ], cmdclass{build_ext: BuildExtension}, )fused_ops.cpp长这样#include torch/extension.h torch::Tensor fused_scale_bias(torch::Tensor x, double scale, torch::Tensor bias) { TORCH_CHECK(x.scalar_type() bias.scalar_type(), x and bias must have the same dtype); auto out torch::empty_like(x); AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), fused_scale_bias, [] { scalar_t* x_ptr x.data_ptrscalar_t(); scalar_t* b_ptr bias.data_ptrscalar_t(); scalar_t* o_ptr out.data_ptrscalar_t(); int64_t n x.numel(); for (int64_t i 0; i n; i) { o_ptr[i] x_ptr[i] * static_castscalar_t(scale) b_ptr[i]; } }); return out; } TORCH_LIBRARY(fused_ops, m) { m.def(fused_scale_bias, fused_scale_bias); }这里AT_DISPATCH_FLOATING_TYPES_AND_HALF是宏会根据x.scalar_type()在编译期展开成switch把 float、double、half 分别实例化。scalar_t是宏内部定义的当前类型别名lambda 里直接用就行。编译并测试python setup.py build_ext --inplaceimport torch import fused_ops x torch.randn(1000, dtypetorch.float32) bias torch.ones(1000, dtypetorch.float32) y fused_ops.fused_scale_bias(x, 2.0, bias) print(y[:3])第一次编译会自动调用 ninja可能需要一两分钟。成功后当前目录会多出一个fused_ops.soLinux或fused_ops.pydWindows直接 import 即可。2.3 快速原型用 load正式项目用 setup.py/CMake不想维护 setup.py 的时候可以在 Python 里用torch.utils.cpp_extension.load动态编译from torch.utils.cpp_extension import load fused_ops load( namefused_ops, sources[fused_ops.cpp], extra_cflags[-O3], )load会按源码内容做哈希缓存第二次调用直接加载编译产物非常适合写 demo 和快速验证。但它不适合正式工程因为每次安装都要重新编译、依赖处理也麻烦。工程化的建议很简单单个 .cpp / .cu 文件用 setup.py 足够文件多了、依赖第三方库了、需要自定义编译参数了直接上 CMake纯做实验和分享load最省事。这三个入口底层都是同一套 PyTorch extension 编译链路没有本质区别。3. 别被张量布局坑了dispatch、连续性与 AT_DISPATCH 宏3.1 为什么 C 循环也不能无脑跑很多人第一次写 C 扩展都会踩这个坑x.numel()是 4但x.data_ptrfloat()指向的元素间距不一定是 1。比如转置、切片、permute得到的张量stride 是非连续的直接按下标遍历就会读到错误数据。最省事的处理是在函数开头统一调用x x.contiguous()。这会生成一份连续内存拷贝后续data_ptr就能顺序访问。代价是多一次拷贝但对绝大多数 case 来说换来的是代码简单和不容易出错。auto x_contig x.contiguous(); auto bias_contig bias.contiguous();如果张量本身已经是连续的contiguous()直接返回原张量不产生额外拷贝可以放心调用。还有一个容易忽略的坑是广播。上面例子我要求x和bias的 shape 完全一致如果bias是[1, 1000]或者[1000, 1]直接按 numel 循环就错了。要么在 C 里实现广播逻辑要么在 Python 侧先bias.expand_as(x)。业务逻辑复杂时我倾向于在 Python 侧把 shape 统一好C 只负责密集张量的高效计算。另外需要注意的是channels_last。很多训练代码启用了torch.channels_last内存格式这种张量numel一样但内存排列完全不同于 NCHWC 里按默认逻辑访问照样出错。处理多内存格式的正确姿势是在 Python 层确认最终输入是contiguous()或者显式在 C 里读取x.sizes()和x.strides()自己处理。3.2 用 AT_DISPATCH 接管 dtypeC 函数在编译时不知道 Python 传来的张量是 float 还是 double而 CUDA kernel 和 CPU 循环都需要针对具体类型实例化。手动写if (x.scalar_type() kFloat)再复制一遍代码很快会崩溃在类型数量上。PyTorch 给的答案是AT_DISPATCH_*系列宏。它本质是生成一个switch把每个可能的ScalarType分支映射到模板编号然后在每个分支里用scalar_t作为类型别名调用你的 lambda。常用的几个宏宏覆盖类型AT_DISPATCH_ALL_TYPES所有常见 CPU 类型浮点整型AT_DISPATCH_FLOATING_TYPESfloat、doubleAT_DISPATCH_FLOATING_TYPES_AND_HALFfloat、double、half部分版本含 bfloat16AT_DISPATCH_ALL_TYPES_AND_HALF所有常见类型 half/bfloat16宏内部的 lambda 需要返回兼容的类型通常统一void。代码里涉及的static_castscalar_t(scale)就是为了把外部传入的 double 精确转成当前类型避免隐式窄化警告。有一点必须记住不要试图在宏外面写类似x.data_ptrfloat()这种硬编码类型。一旦输入是 halffloat 指针的实例化就会编译失败最后你会看到一堆模板报错非常劝退。3.3 同时支持 CPU 与 CUDA 的入口设计我写算子的习惯是先 CPU 后 CUDA。CPU 版逻辑简单、好调试再写 CUDA 版两个版本共享同一个入口函数。入口函数里判断设备分派到不同实现torch::Tensor fused_scale_bias(torch::Tensor x, double scale, torch::Tensor bias) { TORCH_CHECK(x.device() bias.device(), device mismatch); auto out torch::empty_like(x); if (x.is_cuda()) { launch_fused_scale_bias(x, bias, out, scale); } else { launch_fused_scale_bias_cpu(x, bias, out, scale); } return out; }launch_fused_scale_bias声明在.cu文件里由 nvcc 编译CPU 版放在.cpp里。只要 C 符号一致链接阶段会自动对上。多 GPU 的情况下CUDA launch 函数内部最好加上at::cuda::OptionalCUDAGuard guard(x.device())确保 kernel 落在输入张量所在的设备而不是默认卡这点下面会细说。这个入口设计还有一个好处以后想加更多数据类型、更多设备支持改一行判断就行扩展性很强。4. 把计算写进 GPUCUDA kernel 从元素级融合到 grid-stride loop4.1 第一个 CUDA kernel 怎么写写 CUDA kernel 的思维和写 CPU 循环完全不同。CPU 是一个线程顺序处理所有元素CUDA 是成千上万个线程同时处理不同元素。fused_scale_bias的 kernel 我可以写成这样template typename scalar_t __global__ void fused_scale_bias_kernel( const scalar_t* __restrict__ x, const scalar_t* __restrict__ bias, scalar_t* __restrict__ out, scalar_t scale, int64_t n) { int64_t stride blockDim.x * gridDim.x; for (int64_t i blockIdx.x * blockDim.x threadIdx.x; i n; i stride) { out[i] x[i] * scale bias[i]; } }这里用了 grid-stride loop每个线程不是只处理一个元素而是每隔blockDim.x * gridDim.x个元素处理下一个。好处是任意张量大小都能处理不需要严格把 grid 配到刚好覆盖n小时少量 block 也能正确跑减少边界分支的重复计算。__restrict__告诉编译器指针之间没有重叠编译器能做更多优化。const也要尽量加nvcc 会据此生成更好的代码。如果只想覆盖每个元素一次也可以写if (i n)的版本代码更短但 grid-stride loop 是生产级代码更常用的写法。4.2 launch 配置、当前流与错误检查写好 kernel 后需要一个 host 函数负责算 grid/block 并 launch#include ATen/cuda/CUDAContext.h #include c10/cuda/CUDAGuard.h void launch_fused_scale_bias( const torch::Tensor x, const torch::Tensor bias, torch::Tensor out, double scale) { const at::cuda::OptionalCUDAGuard guard(x.device()); int64_t n x.numel(); int threads 256; int64_t blocks (n threads - 1) / threads; fused_scale_bias_kernelstatic_castint(blocks), threads, 0, at::cuda::getCurrentCUDAStream()( x.data_ptrfloat(), bias.data_ptrfloat(), out.data_ptrfloat(), static_castfloat(scale), n); C10_CUDA_KERNEL_LAUNCH_CHECK(); }几个细节值得展开block 为什么选 256不是越大的 block 越好。256 到 512 是大多数 elementwise kernel 的甜点区间既能隐藏访存延迟又不会因为寄存器占用过高导致 occupancy 下降。实际应该用 profiler 调。为什么要用当前 CUDA streamPyTorch 默认在一个非默认 stream 上执行异步操作。如果你 launch kernel 时用了blocks, threads默认流它和 PyTorch 当前流之间没有同步机制可能出现先读后写的乱序。所以务必传入at::cuda::getCurrentCUDAStream()。为什么 launch 完要检查错误CUDA kernel launch 是异步的写了越界索引往往不会立刻报错而是等到后续某个操作设备同步时才炸。C10_CUDA_KERNEL_LAUNCH_CHECK()会在 launch 后立即检查是否有异步错误虽然不能捕获所有内核内部错误但能在早期定位 launch 本身的问题。另外编译 CUDA 扩展时一定要设置TORCH_CUDA_ARCH_LIST或者在做 setup.py 时通过环境变量指定目标架构。比如你用的是 RTX 4090SM 8.9export TORCH_CUDA_ARCH_LIST8.9不设置的话PyTorch 会默认生成一大堆架构的 PTX/二进制编译时间能翻好几倍我第一次等得怀疑电脑死机了。4.3 从 CPU 版到 CUDA 版的整合现在把入口函数更新一下加入设备判断。这里我把 CPU 实现和 CUDA launch 拆到不同文件只是为了说明大型项目怎么组织如果你图省事也可以把 C 和 CUDA 全写进一个.cu文件由 nvcc 统一编译这是最简单可靠的入门姿势。setup.py 需要改成CUDAExtension并且可以为 cxx 和 nvcc 分别指定编译参数from torch.utils.cpp_extension import CUDAExtension, BuildExtension setup( namefused_ops, ext_modules[ CUDAExtension( fused_ops, sources[fused_ops.cpp, fused_ops.cu], extra_compile_args{ cxx: [-O3], nvcc: [-O3, -Xptxas, -O3, -lineinfo], }, ) ], cmdclass{build_ext: BuildExtension}, )-Xptxas -O3是让 NVIDIA 的底层编译器做更激进的优化-lineinfo让 kernel 崩溃时能给出更精确的行号配合ncu排查很有用。5. 自动求导怎么接自定义 Function、反向 kernel 与 gradcheck5.1 为什么底层算子不自动参与反向传播刚接触自定义算子时很容易产生一个误解C 扩展返回的是torch::TensorPyTorch 就能自动求导了。实际上你写的 C 函数在 PyTorch 眼里只是一个普通符号autograd 引擎并不知道这个函数内部是线性的还是非线性的更不知道怎么计算梯度。接入自动求导有两种方式Python 侧用torch.autograd.Function包一层灵活、好调试适合绝大多数业务C 侧直接用torch::autograd::Function实现性能上限更高但代码复杂度也高适合算子本身就要注册到框架内部的场景。我推荐先用第一种。C 只提供底层计算能力梯度逻辑由 Python 控制开发周期短出问题好定位。5.2 一个带反向的完整例子clamp 融合算子回到开头说的那个场景我把它升级成完整算子y clamp(x * scale bias, min_val, max_val)。前向的 CUDA kernel 很简单template typename scalar_t __global__ void fused_scale_bias_clamp_kernel( const scalar_t* __restrict__ x, const scalar_t* __restrict__ bias, scalar_t* __restrict__ out, scalar_t scale, scalar_t min_val, scalar_t max_val, int64_t n) { int64_t stride blockDim.x * gridDim.x; for (int64_t i blockIdx.x * blockDim.x threadIdx.x; i n; i stride) { scalar_t v x[i] * scale bias[i]; out[i] v min_val ? min_val : (v max_val ? max_val : v); } }反向推导需要一点数学。设v x * scale biasy clamp(v, min, max)。在可导区间内d y / d x scaled y / d b 1当v落在截断区间外时梯度为 0。所以反向时需要判断v是否在min和max之间用一个 mask 统一处理。反向 CUDA kernel 可以写成template typename scalar_t __global__ void fused_scale_bias_clamp_backward_kernel( const scalar_t* __restrict__ grad_out, const scalar_t* __restrict__ x, const scalar_t* __restrict__ bias, scalar_t* __restrict__ grad_x, scalar_t* __restrict__ grad_bias, scalar_t* __restrict__ grad_scale_partial, scalar_t scale, scalar_t min_val, scalar_t max_val, int64_t n) { int64_t stride blockDim.x * gridDim.x; for (int64_t i blockIdx.x * blockDim.x threadIdx.x; i n; i stride) { scalar_t v x[i] * scale bias[i]; scalar_t mask (v min_val v max_val) ? scalar_t(1) : scalar_t(0); grad_x[i] grad_out[i] * scale * mask; grad_bias[i] grad_out[i] * mask; grad_scale_partial[i] grad_out[i] * x[i] * mask; } }然后在 host 端对grad_scale_partial做一次torch::sum返回标量梯度。C 里返回多个梯度时用std::vectortorch::Tensorstd::vectortorch::Tensor fused_scale_bias_clamp_backward( torch::Tensor grad_output, torch::Tensor x, torch::Tensor bias, double scale, double min_val, double max_val) { auto grad_x torch::empty_like(x); auto grad_bias torch::empty_like(bias); auto grad_scale_partial torch::empty_like(x); // dispatch launch kernel ... auto grad_scale torch::sum(grad_scale_partial); return {grad_x, grad_bias, grad_scale}; }Python 侧用torch.autograd.Function包一层import fused_ops class FusedScaleBiasClamp(torch.autograd.Function): staticmethod def forward(ctx, x, scale, bias, min_val, max_val): ctx.save_for_backward(x, bias) ctx.scale scale ctx.min_val min_val ctx.max_val max_val return fused_ops.fused_scale_bias_clamp(x, scale, bias, min_val, max_val) staticmethod def backward(ctx, grad_output): x, bias ctx.saved_tensors grad_x, grad_bias, grad_scale fused_ops.fused_scale_bias_clamp_backward( grad_output.contiguous(), x, bias, ctx.scale, ctx.min_val, ctx.max_val ) return grad_x, grad_scale, grad_bias, None, None注意三件事forward里用ctx.save_for_backward保存张量用ctx属性保存标量backward返回值顺序必须和forward输入顺序一一对应不需要梯度的输入返回Nonescale是标量梯度要汇总成 0 维张量这里用 C 的torch::sum完成。5.3 用 gradcheck 验证前向反向手写反向最怕算错公式。PyTorch 提供了torch.autograd.gradcheck数值微分对比解析梯度是验证自定义算子反向的金标准from torch.autograd import gradcheck x torch.randn(3, 4, dtypetorch.float64, requires_gradTrue) bias torch.randn(3, 4, dtypetorch.float64, requires_gradTrue) scale 0.5 ok gradcheck( FusedScaleBiasClamp.apply, (x, scale, bias, -1.0, 1.0), eps1e-6, atol1e-5, ) print(gradcheck passed:, ok)我强烈建议用float64测 gradcheck。很多自定义算子死活用不了一看就是测试 dtype 用了 float32数值误差累积直接超了阈值。clamp 这类分段函数在边界处不可导gradcheck 会因为采样碰到边界而失败。解决办法是让测试数据远离截断边界或者放宽eps。这些不是算子的 bug是数学上不可微点的正常现象。6. 实测对比与报错排查编译、架构、profile 一条线6.1 先看一张性能对比表为了让你对收益有直观概念我给y clamp(x * scale bias, 0.0, 1.0)做了组测试输入是 1M 个 float32CUDA 版本在 RTX 4090 上测CPU 版本随手测的只代表量级实现方式耗时/迭代备注纯 Python for 循环约 180 ms逐元素解释执行基本不可用torch 组合算子约 0.32 ms三次 kernel 启动 中间张量读写torch.compile 自动融合约 0.07 ms有编译开销但运行时性能不错C CPU 单线程约 1.1 ms受内存带宽限制慢于 GPU 版本CUDA 融合 kernel约 0.02 ms单次启动、单次遍历性能最好看到没有torch.compile在这个场景已经能把 0.32 ms 压到 0.07 ms和手写 CUDA 的差距不是数量级。所以融合才是关键不是C 一定比 Python 快。真正能让 CUDA kernel 拉开差距的是更复杂的算子涉及共享内存、归约、特殊访存模式、或者和第三方 GPU 库交互。遇到这种需求torch.compile就帮不上忙了手写 CUDA 是唯一解。6.2 编译期最常见的 5 个报错我在群里见过太多人卡在编译环节这里列几个最高频的报错特征原因解决办法undefined symbol: _ZN2at...编译时的 PyTorch 版本和运行时不一致或混用 debug/release重建虚拟环境统一 PyTorch 版本data_ptrat::Half编译失败没走 AT_DISPATCH 宏硬编码了 float 指针把类型处理收进宏里invalid device function编译时TORCH_CUDA_ARCH_LIST没包含你的 GPU 架构按显卡设置架构再重编MSVC 版本过低Windows 上没装 VS 2019/2022 的 C 桌面开发组件安装对应 workload编译时间极长没有设置TORCH_CUDA_ARCH_LIST编了所有架构设置环境变量只编自己用的架构invalid device function这个错特别阴间编译成功、import 成功、一调用就报错。原因是 nvcc 只生成了某几个 SM 架构的二进制你的 GPU 不在里面。解决思路就是准确地告诉编译器目标架构。6.3 运行期排查思路自定义算子一旦能跑别急着开心先用 profiler 看真实数据。这里容易踩的坑是异步执行。如果你用time.time()直接测 CUDA 算子的耗测到的是Python 提交任务的时间不是 kernel 真正跑完的时间。必须显式同步torch.cuda.synchronize() start torch.cuda.Event(enable_timingTrue) end torch.cuda.Event(enable_timingTrue) start.record() y fused_ops.fused_scale_bias_clamp(x, scale, bias, 0.0, 1.0) end.record() torch.cuda.synchronize() print(start.elapsed_time(end))数值不对时我排查顺序一般是先检查contiguous()有没有加非连续张量直接data_ptr循环必错检查 dtype 分支是否覆盖了输入的所有类型尤其是 half宏没用对容易踩检查广播逻辑C 不自动广播shape 不一致和按 numel 循环会静默出错检查 grid/block 边界CUDA 里越界访问不会像 CPU 段错误那么明显但会产生脏数据用ncu看 kernel 的 occupancy 和访存效率如果 occupancy 很低可能是寄存器用太多或 block 太小。7. 进阶新式算子注册与 torch.compile 兼容7.1 为什么只写 pybind11 还不够老式PYBIND11_MODULE暴露 Python 函数的方式在 PyTorch 2.x 里已经不太够用了。torch.compile、TorchDynamo、AOTAutograd 需要知道算子的 schema输入输出类型、可变性、是否返回视图裸 pybind11 函数对它们来说是一个不透明黑盒很多编译优化会直接跳过你的算子。解决方案是TORCH_LIBRARY注册给算子一个带 schema 的名字让 PyTorch 的 dispatcher 体系认识它。前文我写的是TORCH_LIBRARY(fused_ops, m) { m.def(fused_scale_bias, fused_scale_bias); }这种写法能 import但 TorchScript/torch.compile 对算子的形状、dtype 依旧不敏感。更完整写法是显式声明 schemaTORCH_LIBRARY(fused_ops, m) { m.def(fused_scale_bias(Tensor x, float scale, Tensor bias) - Tensor, fused_scale_bias); }一旦有了 schemaPyTorch 就能把它当成一个正规 op 处理用在torch.ops.fused_ops.fused_scale_bias这类调用链里。7.2 TORCH_LIBRARY 的 schema 和 Meta kernel想让 torch.compile 更好地推理你的算子形状还需要给算子补一个 Meta kernel。Meta kernel 不执行真实计算只返回正确的输出形状和 dtype供编译期做 fake tensor 推理。torch::Tensor fused_scale_bias_meta( const torch::Tensor x, double scale, const torch::Tensor bias) { return torch::empty_like(x); } TORCH_LIBRARY_IMPL(fused_ops, Meta, m) { m.impl(fused_scale_bias, fused_scale_bias_meta); }这样torch.compile 在做图优化时就不需要真的跑一遍算子也能知道输出长什么样不会因为你的自定义 op 卡住整个图。PyTorch 2.x 里也推荐用torch.library的 Python API 做精细控制比如用torch.library.define注册带 schema 的 op再注册 fake kernel。对这个话题有需求的可以继续深入看官方文档的torch.library部分。7.3 两个实战建议第一老项目升级时别一上来就全量改成新式注册。先把算子用TORCH_LIBRARY加 schema 注册好保持 Python 侧调用方式不变跑通测试再逐步补 Meta/Fake kernel。每次只改一个点出问题好定位。第二GPU 算子的开发路径应该是先用 torch 组合验证逻辑 → 如果性能不够再看 torch.compile 能不能融合 → 还不行手写 C CPU 版 → 最后写 CUDA 版。不要第一版就上 CUDA调试成本太高。我实际项目里好几个环节最后用的是torch.compile 自动融合 少数必须要手写的 CUDA kernel的组合拳。最后说点个人体会。自定义算子这东西写一次不复杂难的是长期维护。dtype 覆盖全不全、非连续张量处理没处理、多 GPU 下设备 guard 有没有加、新版本 PyTorch 的 dispatcher 改没改每一个都是潜在的坑。我每次写完自定义算子第一件事是跑 gradcheck第二件事是跑一遍混合精度的对照测试第三件事才是上 profiler 调性能。另外一个很实用的经验尽量保证 C 侧只依赖torch/extension.h和少量 ATen 接口不要引入太多自定义数据结构和复杂的类继承。这样代码给同事看也好懂后续移植到其他框架时也能快速抽出来。算子入口统一用torch::Tensor传递数据内部再慢慢拆解这是最不容易出问题的方式。祝大家的算子第一次编译就通过gradcheck 一把过。
