上周帮一个做视觉推理的团队排查一个很典型的问题：他们在 Python 层用好几个小算子拼出了一个带系数的激活函数，训练跑得挺好，结果一导 ONNX 就出问题——torch.onnx.export把这一串操作打成几十个零散节点，ONNX Runtime 跑起来比 PyTorch 还慢一截，显存占用也没降下来。这个场景我见过太多次了。PyTorch 自定义 C++ 拓展算子、转换 ONNX 模型、ORT 部署实现，这三件事听起来像三个独立动作，实际上是一条环环相扣的链路：算子在 PyTorch 侧怎么注册，直接决定了 ONNX 导出时能不能被识别；ONNX 里用什么 domain 和什么节点名，又决定了 ORT 加载时该怎么注册；而 ORT 侧 kernel 的输入输出类型和内存管理方式，最后又反过来约束你在 C++ 拓展里的写法。任何一环没想清楚，后面大概率要返工。

这篇就把这条链路完整走一遍。从为什么要写自定义 C++ 拓展算子的真实动机讲起，到TORCH_LIBRARY注册、CUDA 核函数的前向反向推导，再到 custom symbolic 的注册方式、动态 shape 的导出参数、ORT CustomOpDomain 的 C++ 实现和 Python 加载，最后落到数値对齐和性能验证。适合已经会用 PyTorch 训练、但对“训练完怎么高效部署到推理引擎”这一步还没打通的同学，也适合手里有一堆性能不达标的模型、想从算子粒度上做优化的同学。示例算子我选 SiLU（也就是常说的 Swish），数学不复杂、前向反向都能手推、融合收益明显，作为模板最容易迁移到你自己真正的目标算子上。

1. 这条链路到底解决什么问题

1.1 为什么 Python 层写不出来

先说说我为什么不推荐一开始就在 Python 里“拼算子”。很多人的第一反应是：我的自定义激活函数不就是x * torch.sigmoid(x)吗，一行代码的事，干嘛要费劲写 C++。这话在训练阶段成立，在推理阶段就不成立了。Python 层每调用一次算子，都要经过一次 Python 解释器到 C++ 后端的跨界跳转，一次 kernel launch，一次显存读写。当你把这些小算子拼在一个 Transformer block 里，重复调用几百次，跨界的开销和反复读写中间张量的开销就会变得很显眼。

更关键的是图结构问题。Python 层拼出来的算子，在 PyTorch 的 FX/TorchScript 图里是一串独立节点，导出 ONNX 之后依然是一串独立节点。ORT 虽然有自己的图优化器，但很多融合规则依赖节点模式匹配，一旦你的模式稍微偏一点，或者跨了 fusion 的边界，它就融不起来。而如果你把这段逻辑写成一个真正的 C++ 拓展算子，PyTorch 图里只有一个节点，ONNX 图里也只有一个节点，融合是在你手里完成的，不依赖推理引擎的优化器心情。

还有一类场景是纯 Python 根本没法做的，比如需要精细控制显存布局的算子、需要调用 CUDA 特定指令（如 warp shuffle、异步拷贝）的算子、需要和外部硬件库对接的算子。这些必须下沉到 C++/CUDA 层。所以我个人的判断标准很简单：如果一个操作在模型前向里被调用超过几十次，或者它内部有可以被合并的多次显存访问，就值得做成自定义 C++ 拓展算子。

1.2 三段式分工：训练态、中间态、推理态

整条链路我习惯按三个状态来切分，每个状态负责不同的事，边界清楚之后，调试起来就不容易互相甩锅。

训练态的核心诉求是“可微”和“可跑在 GPU 上”。这意味着你不仅要写前向核函数，还要写反向核函数，并且要在TORCH_LIBRARY_IMPL里把 autograd 接上，否则反向传播到你这个节点就断了。

中间态的核心诉求是“图上的节点名和 domain 必须可被 ONNX 识别”。纯 C++ 拓展算子，ONNX 是不知道它是什么的，你必须通过register_custom_op_symbolic给它注册一个映射规则，告诉导出器“当你看到myops::swish这个节点时，把它翻译成 ONNX 图里的什么”。这里有个分岔：翻译成 ONNX 原生算子的组合，还是翻译成带自定义 domain 的单节点。这两条路后面的部署成本完全不同，第 4 节会详细展开。

推理态的核心诉求是“ORT 能加载并正确执行”。如果走自定义 domain 路线，你需要在 ORT 侧用Ort::CustomOpBase派生一个算子实现，编译成共享库，然后在创建 Session 之前通过SessionOptions::AddCustomOpDomain注册进去。这一步的输入输出类型声明必须和模型里的类型严格对齐，差一个float和double就会在加载阶段直接抛异常。

2. 工程骨架与依赖版本锁定

2.1 目录结构

我一般在项目根目录下开一个独立目录放拓展，不跟模型代码混在一起，方便单独编译和单测。结构大致长这样：

project/ ├── ops/ │ ├── setup.py │ ├── swish.cpp # 算子声明、TORCH_LIBRARY 注册 │ ├── swish_cuda.cu # CUDA 核函数 │ ├── swish_cpu.cpp # CPU fallback │ └── ort_plugin/ │ ├── swish_ort.cc # ORT CustomOp 实现 │ └── CMakeLists.txt ├── export_onnx.py └── infer_ort.py

把 PyTorch 拓展和 ORT 插件分开放的原因是它们的编译工具链不一样。PyTorch 拓展用torch.utils.cpp_extension或者setup.py里的CUDAExtension就够了，会自动带上 PyTorch 的头文件和链接参数；而 ORT 插件我倾向用 CMake 单独编，因为它只需要onnxruntime_cxx_api.h和 ORT 的库，不依赖 PyTorch，编出来的产物更干净，也方便放到没有 PyTorch 环境的推理机器上。

2.2 三个必须提前确认的环境细节

第一个是CXX11_ABI。PyTorch 官方预编译包默认是_GLIBCXX_USE_CXX11_ABI=0，如果你自己的拓展用 g++ 默认参数编（也就是 ABI=1），链接时会出现一大堆undefined symbol，报错信息里全是std::__cxx11::basic_string之类的符号。解决办法是在setup.py的extra_compile_args里显式加-D_GLIBCXX_USE_CXX11_ABI=0，和 PyTorch 对齐。

第二个是CUDA 版本与 PyTorch 编译版本的一致性。torch.version.cuda打印出来的版本必须和你本地nvcc --version的主版本号一致，小版本可以有差异，但 11.x 和 12.x 混用大概率编不过或者运行时报非法指令。

第三个是ORT 的版本与 ONNX opset 的匹配关系。比如 opset 17 需要 ORT 1.13 以上，opset 19 需要 ORT 1.15 以上。如果你的模型里用到了比较新的算子，导出时把 opset 压太低会直接报“不支持的算子”，压太高又可能 ORT 加载不了。我一般固定在 opset 17，兼容性最稳。

提示：编译前先跑一遍python -c "import torch; print(torch.__version__, torch.version.cuda, torch._C._GLIBCXX_USE_CXX11_ABI)"，把这三个值贴到 README 里，团队换机器时能省掉大量“为什么你能编我不能编”的扯皮。

3. C++ 拓展算子的编写与 PyTorch 接入

3.1 从组合算子到融合核函数

先看清楚我们要融合的是什么。SiLU 的数学形式是y = x * sigmoid(x)，拆开看是三个操作：sigmoid 计算、一次乘法、以及中间张量的分配。组合写法在 PyTorch 里会启动至少两个 kernel，中间还要写回一次显存。融合之后，一次 kernel launch 读一次 x，算完 sigmoid 直接乘回去，写一次 y，中间结果全在寄存器里。在 batch 较大、通道数较多的场景下，这个访存次数的差异会直接体现在端到端延迟上。

CPU 侧的 fallback 我也建议一起写，不是为了性能，是为了让同一份代码在没有 GPU 的 CI 机器上也能跑通单测。下面是我常用的写法，先用TORCH_CHECK把参数前置条件卡住，再分发到具体实现：

// swish.cpp #include <torch/extension.h> namespace myops { torch::Tensor swish_cpu(torch::Tensor x); torch::Tensor swish_cuda(torch::Tensor x); torch::Tensor swish_backward_cuda(torch::Tensor grad_out, torch::Tensor x); torch::Tensor swish(torch::Tensor x) { TORCH_CHECK(x.scalar_type() == torch::kFloat32, "swish 目前只支持 float32，收到 ", x.scalar_type()); auto xc = x.contiguous(); if (xc.is_cuda()) { return swish_cuda(xc); } return swish_cpu(xc); } torch::Tensor swish_backward(torch::Tensor grad_out, torch::Tensor x) { TORCH_CHECK(x.is_cuda(), "反向目前只实现了 CUDA 版本"); return swish_backward_cuda(grad_out.contiguous(), x.contiguous()); } } // namespace myops

x.contiguous()这一句经常被忽略。如果你的输入是从别的地方 view 出来的非连续张量，直接按线性地址遍历会读到错位的数据，数值看起来“没错但就是差一点点”，查起来非常折磨。花一次连续性检查的开销，换来的是长时间的安心。

3.2 前向与反向的推导

前向没什么好说的，y = x / (1 + exp(-x))。反向稍微推一下，设s = sigmoid(x)，那么y = x * s，对 x 求导：

dy/dx = s + x * s * (1 - s) = s * (1 + x * (1 - s))

把这个式子记牢，反向核函数里直接算，不要重新算一遍 sigmoid 就浪费了。CUDA 核函数我写成这样：

// swish_cuda.cu #include <torch/extension.h> #include <ATen/cuda/CUDAContext.h> #include <cuda_runtime.h> __global__ void swish_fwd_kernel(const float* __restrict__ x, float* __restrict__ y, int64_t n) { int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; if (i < n) { float v = x[i]; y[i] = v / (1.0f + expf(-v)); } } __global__ void swish_bwd_kernel(const float* __restrict__ grad_out, const float* __restrict__ x, float* __restrict__ grad_x, int64_t n) { int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; if (i < n) { float v = x[i]; float s = 1.0f / (1.0f + expf(-v)); grad_x[i] = grad_out[i] * s * (1.0f + v * (1.0f - s)); } } torch::Tensor swish_cuda(torch::Tensor x) { auto y = torch::empty_like(x); int64_t n = x.numel(); const int threads = 256; const int blocks = (n + threads - 1) / threads; auto stream = at::cuda::getCurrentCUDAStream(); swish_fwd_kernel<<<blocks, threads, 0, stream>>>( x.data_ptr<float>(), y.data_ptr<float>(), n); return y; }

注意at::cuda::getCurrentCUDAStream()这一行。不用当前 stream 的话，你的 kernel 会跑到默认 stream 上，和 PyTorch 其他算子的 stream 不同步，结果就是偶尔正确、偶尔错一个数，多卡训练时尤其明显。踩过一次之后就再也不会忘了。

3.3 用 TORCH_LIBRARY 注册与编译验证

注册方式有两种，一种是老式的PYBIND11_MODULE直接绑函数，另一种是新式的TORCH_LIBRARY。我强烈推荐后者，因为只有用TORCH_LIBRARY注册的算子，才能被 TorchScript 和 ONNX 导出器识别为图上的一个 dispatcher 节点，PYBIND11_MODULE绑出来的只是个 Python 函数，导出时会被追踪成一堆基本操作的组合，等于白做。

// 在 swish.cpp 末尾追加 TORCH_LIBRARY(myops, m) { m.def("swish(Tensor x) -> Tensor"); m.def("swish_backward(Tensor grad_out, Tensor x) -> Tensor"); } TORCH_LIBRARY_IMPL(myops, CUDA, m) { m.impl("swish", &myops::swish); m.impl("swish_backward", &myops::swish_backward); } TORCH_LIBRARY_IMPL(myops, CPU, m) { m.impl("swish", &myops::swish); }

m.def里的类型签名必须写死Tensor，不能省。这个 schema 字符串是导出器识别节点的依据，写错了会在导出阶段报 schema 解析失败。

编完之后立刻验证注册是否成功，别拖到导出再查：

import torch torch.ops.load_library("./build/libmyops.so") x = torch.randn(4, 8, device="cuda", requires_grad=True) y = torch.ops.myops.swish(x) y.sum().backward() ref = x * torch.sigmoid(x) print(torch.allclose(y, ref.detach(), atol=1e-6)) # 期望 True print(x.grad is not None) # 期望 True

前向和反向都过一遍再往下走。我见过有人只测了前向就往下推，等到训练时才发现反向没接上，梯度全是 None，白跑了一晚上。

4. ONNX 导出与自定义算子的符号化

4.1 为什么直接 export 会失败

torch.onnx.export的工作机制是先追踪一遍模型前向，拿到一张 TorchScript 图，然后按节点逐个翻译成 ONNX 节点。当它遇到一个myops::swish的节点时，会去查自己的符号表里有没有对应的翻译规则。查不到，就抛UnsupportedOperatorError，或者在某些版本里直接把节点拆解掉，产生一张结构诡异的图。

所以第一步必须告诉导出器“这个节点该翻译成什么”。这个动作叫注册 custom symbolic，本质是提供一个 Python 函数，接收 ONNX 图构建器和输入节点，返回一个 ONNX 节点。

4.2 register_custom_op_symbolic 的写法

这里有两个方案，我按适用场景讲清楚。

方案 A：降级为原生算子组合。如果自定义算子的逻辑能用 ONNX 原生算子等价表达（SiLU 就属于这种，Mul(x, Sigmoid(x))就完事了），那最省事。好处是 ORT 原生就支持，不需要注册任何插件，任何版本的 ORT 都能跑。代价是融合度回到了组合算子的水平，前面 C++ 拓展做的融合白做了。适合那种“自定义算子只是为了训练方便，推理时可以接受不融合”的场景。

import torch.onnx def swish_symbolic(g, x): return g.op("Mul", x, g.op("Sigmoid", x)) torch.onnx.register_custom_op_symbolic("myops::swish", swish_symbolic, 17)

方案 B：原样保留为自定义 domain 节点。这是真正体现“自定义算子”价值的路线，导出的图里依然只是一个节点，融合效果原样保留。ORT 侧需要配合注册 kernel。

import torch.onnx def swish_symbolic(g, x): return g.op("myops::Swish", x) torch.onnx.register_custom_op_symbolic("myops::swish", swish_symbolic, 17) torch.onnx.export( model, dummy_input, "model_custom.onnx", opset_version=17, custom_opsets={"myops": 1}, input_names=["x"], output_names=["y"], dynamic_axes={"x": {0: "batch", 1: "seq"}}, )

关键在custom_opsets={"myops": 1}这一行。不加的话，导出器不会把myops这个 domain 的 opset 版本写进模型文件，ORT 加载时会抱怨 domain 版本不匹配。这个参数在很多教程里被漏掉了，导致加载阶段报一个很隐晦的错。

导出之后务必用onnx.checker过一遍，再用 onnx 的可视化工具确认图里节点确实只有一个：

import onnx model = onnx.load("model_custom.onnx") onnx.checker.check_model(model) print([(n.domain, n.op_type) for n in model.graph.node]) # 期望看到 ('myops', 'Swish') 而不是一堆 Mul/Sigmoid

4.3 动态 shape 与 opset 的取舍

dynamic_axes这参数我必须多说两句。如果你的推理场景 batch size 是变化的，或者序列长度是变化的，不声明动态轴，导出的模型就把 shape 固化成了 dummy input 的尺寸。ORT 加载后你喂一个 batch 更大的输入，它会直接报Got invalid dimensions for input，而且报错信息里的期望值就是你的 dummy 尺寸，看起来像是模型坏了。

声明动态轴之后，还要注意 ORT 侧的图优化行为。有些优化规则需要知道 shape 上界，如果你在dynamic_axes里只给了轴名不给范围，ORT 会按“无穷大”处理，某些 fuser 会因此放弃优化。解决办法是导出完成后用onnxruntime.tools.make_dynamic_shape_fixed或者手动在模型里补上dim_param的范围信息，如果性能不达标，可以再权衡。

opset 的选择我一般遵循一个原则：宁可选一个 ORT 支持得早、算子集已经稳定下来的版本。opset 17 是我目前用得最顺的，覆盖了绝大多数常见算子，同时 ORT 1.13 以上都支持。如果你的自定义算子用到了 18 以上才有的特性，那再往上抬，但要同步升级 ORT 版本。

注意：register_custom_op_symbolic的第三个参数是 opset 版本，这个值必须小于等于你导出时用的opset_version，否则注册不生效，导出时依然会报找不到翻译规则。

5. ONNX Runtime 侧的算子注册与部署

5.1 自定义算子的 C++ 实现

ORT 的自定义算子实现分三块：一个继承自Ort::CustomOpBase的算子描述类，一个真正干活的 Kernel 类，以及一个注册到CustomOpDomain的动作。

Kernel 类的Compute方法里，最关键的几个 ORT C API 是：KernelContext_GetInput拿输入，KernelContext_GetOutput拿输出（注意是 ORT 帮你分配好内存的，你只管往里写），GetTensorData拿只读指针，GetTensorMutableData拿可写指针。下面是 CPU 版的 SiLU 实现：

#include <onnxruntime_cxx_api.h> #include <cmath> #include <vector> struct SwishKernel { explicit SwishKernel(const OrtApi& api) : api_(api) {} void Compute(OrtKernelContext* context) { const OrtValue* input = api_.KernelContext_GetInput(context, 0); OrtTensorTypeAndShapeInfo* info = nullptr; api_.GetTensorTypeAndShape(input, &info); size_t dim_count = 0; api_.GetDimensionsCount(info, &dim_count); std::vector<int64_t> dims(dim_count); api_.GetDimensions(info, dims.data(), dim_count); api_.ReleaseTensorTypeAndShapeInfo(info); int64_t total = 1; for (auto d : dims) total *= d; OrtValue* output = api_.KernelContext_GetOutput( context, 0, dims.data(), dim_count); const float* x = api_.GetTensorData<float>(input); float* y = api_.GetTensorMutableData<float>(output); for (int64_t i = 0; i < total; ++i) { y[i] = x[i] / (1.0f + std::exp(-x[i])); } } const OrtApi& api_; };

这里有个坑得提前说：KernelContext_GetOutput返回的内存是 ORT 的 Arena 分配的，你绝对不能用delete或者free去释放它，也不能把它塞进一个长期存活的std::shared_ptr里，否则会在 Session 析构时触发 double free。只管往里写数据，生命周期交给 ORT 管。

5.2 编译成共享库并在 Python 里加载

算子描述类要显式声明输入输出类型和归属的 domain：

struct SwishOp : Ort::CustomOpBase<SwishOp, SwishKernel> { void* CreateKernel(const OrtApi& api, const OrtKernelInfo*) const { return new SwishKernel(api); } const char* GetName() const override { return "Swish"; } const char* GetExecutionProviderType() const override { return "CPUExecutionProvider"; } size_t GetInputTypeCount() const override { return 1; } ONNXTensorElementDataType GetInputType(size_t) const override { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; } size_t GetOutputTypeCount() const override { return 1; } ONNXTensorElementDataType GetOutputType(size_t) const override { return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; } }; static SwishOp g_swish_op; static Ort::CustomOpDomain g_domain("myops"); static bool g_inited = [] { g_domain.Add(&g_swish_op); return true; }(); extern "C" Ort::CustomOpDomain* get_custom_op_domain() { return &g_domain; }

extern "C"这个出口函数很重要。ORT 的 Python 端register_custom_ops_library会通过dlsym去找一个叫RegisterCustomOps或者约定名字的符号，不同版本约定略有差异。如果你的共享库是自己写的加载逻辑，用extern "C"导出更稳，避免 C++ name mangling 导致符号找不到。

CMake 配置大致是这样，重点是 include 路径指向 ORT 的 release 包自带的头文件目录，链接时不需要链 ORT 的库（因为算子实现本身不依赖 ORT 的符号，运行时由宿主进程提供）：

cmake_minimum_required(VERSION 3.18) project(swish_ort_plugin CXX) set(CMAKE_CXX_STANDARD 17) set(ORT_INCLUDE_DIR "/path/to/onnxruntime/include") add_library(swish_ort SHARED swish_ort.cc) target_include_directories(swish_ort PRIVATE ${ORT_INCLUDE_DIR}) set_target_properties(swish_ort PROPERTIES OUTPUT_NAME "myops_ort")

Python 侧加载：

import onnxruntime as ort so = ort.SessionOptions() so.register_custom_ops_library("./libmyops_ort.so") # Windows 下是 .dll sess = ort.InferenceSession( "model_custom.onnx", sess_options=so, providers=["CPUExecutionProvider"], )

如果加载成功，InferenceSession构造完成就会打印出模型的输入输出信息。如果加载失败，最典型的报错是No Op registered for Swish with domain_version of 1，这说明三件事之一：so 没加载上、domain 名字对不上、或者 kernel 名字对不上。排查顺序按这个来，先确认 so 的路径和加载返回值，再打印一下g_domain里注册的算子名。

5.3 会话配置与执行提供器

生产环境部署时，SessionOptions 里有两个参数值得花时间调：intra_op_num_threads和graph_optimization_level。前者控制单个算子内部并行度，后者控制图优化强度。

对于自定义算子占比较大的模型，intra_op_num_threads不要调得太高。因为你的自定义 kernel 内部可能已经是多线程并行的（如果你在 kernel 里自己开了 OpenMP 或者 std::thread），再叠加 ORT 的线程池，会出现严重的线程争用，实际反而变慢。我一般的做法是先用默认值跑一遍 baseline，再把intra_op_num_threads设成物理核数的一半，观察吞吐变化，取拐点。

graph_optimization_level建议开到ORT_ENABLE_ALL，但如果你发现某些优化规则把你的自定义节点错误地融合掉了（这种情况罕见但不是没有），可以退回ORT_ENABLE_EXTENDED。这个参数在 Python 里的写法是so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL。

如果要在 GPU 上跑自定义算子，GetExecutionProviderType就得返回CUDAExecutionProvider，并且 kernel 的 Compute 里要自己处理 stream 同步。这块我建议一开始先用 CPU 版本跑通全链路，业务逻辑验证无误之后，再花时间做 CUDA 版本的 kernel，否则你同时要排查数值问题、类型问题、stream 问题，定位起来极其痛苦。

6. 数值对齐与性能验证

6.1 三层数值对齐

我一直坚持做三层对齐，而不是只对比首尾。第一层是 PyTorch 原生算子 vs C++ 拓展算子，第二层是 PyTorch 模型输出 vs ONNX 模型（用 onnxruntime 直接跑，或者用 onnx 的 reference 实现），第三层是 ONNX 导出端 vs 目标机器的 ORT 加载端。

第一层的代码前面给过了，用torch.allclose加绝对误差 1e-6 不过就说明核函数有 bug。第二层的对比如下：

import numpy as np import onnxruntime as ort import torch x = torch.randn(2, 128, 512, device="cuda") with torch.no_grad(): ref = model(x).cpu().numpy() so = ort.SessionOptions() so.register_custom_ops_library("./libmyops_ort.so") sess = ort.InferenceSession("model_custom.onnx", so, providers=["CPUExecutionProvider"]) out = sess.run(None, {"x": x.cpu().numpy()})[0] np.testing.assert_allclose(out, ref, rtol=1e-4, atol=1e-5)

FP32 场景下用rtol=1e-4, atol=1e-5是合理的宽容度。如果差得远，优先查三个地方：contiguous有没有加、dtype 有没有在导出时被隐式转换成 double、以及 ORT custom op 里GetDimensions的顺序和 PyTorch 的维度定义是否一致（PyTorch 是 NCHW，ONNX 也是 NCHW，但有些老代码里习惯性用 NHWC 做转换，这个坑不小）。

第三层的对比分机器来做。用同一份 onnx 文件，在两台机器分别跑，输出落到 numpy 数组，逐元素相减取最大绝对值。差异大于 1e-5 就说明有平台相关的实现差异，常见于用了expf之外的特殊函数，或者 FP32 累加顺序不同。对于自定义算子这种边界清晰的情况，一般不会有什么差异，这个检查更多是为了防止 onnx 文件本身在传输过程中损坏。

6.2 性能数据与瓶颈定位

我给一个实测过的对比数据，场景是 batch=8、序列长度 512、隐藏维度 768 的单层前向，跑 100 次取平均，硬件是一张消费级显卡，输入输出都固定。

这个数据不代表所有场景，但能看出来一个规律：融合带来的收益主要在访存次数上，时间收益大概在 15% 到 25% 之间，显存收益在 10% 左右。如果你的模型里自定义算子占比很小，整体收益会被稀释到个位数，这时候就要评估一下投入产出比。

定位瓶颈的时候，我习惯用 ORT 的 profiling 接口导出 Chrome Trace，然后看时间柱状图：

so.enable_profiling = True sess = ort.InferenceSession("model_custom.onnx", so, providers=["CPUExecutionProvider"]) sess.run(None, {"x": x.cpu().numpy()}) prof_file = sess.end_profiling() print(prof_file)

打开 Chrome 的chrome://tracing，把 json 拖进去，能非常直观地看到每个节点的耗时。如果发现自定义算子的耗时占比远超预期，大概率是 kernel 里的访存模式不理想，比如用 stride 访问代替了线性访问，或者 thread block 的大小没调好。256是我常用的起点，非元素级算子的 batching 策略就要根据实际数据形状来定。

提示：benchmark 一定要用time.perf_counter配合多次预热，前几次跑出来的数字往往包含内存分配和 JIT 的开销，不代表稳态性能。我一般是预热 10 次、跑 100 次取中位数，这样数据比较有说服力。

7. 高频坑位速查与实战经验

7.1 编译链接类

这一类的错误大多出现在import torch或者torch.ops.load_library之后立刻崩溃，报错信息往往一屏看不完。

TORCH_CUDA_ARCH_LIST这个参数我要单独提一句。如果你在 A100 上编的拓展拿到 3090 上跑，可能会出现no kernel image is available for execution on the device。根本原因是编译时按 A100 的算力生成了 cubin，而 3090 的算力代号不同。解决办法是编译前显式设置export TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9"，把目标机器可能用到的算力都覆盖上，代价是编译时间变长、产物变大，但换来的是部署时的省心。

7.2 导出类

导出阶段的问题通常不致命，但会让你怀疑人生，因为图看起来能跑，就是结果不对。

最常见的是节点没有真正被替换成自定义算子。你注册了 symbolic，也导出了，但打开 onnx 一看，里面是一串Mul/Sigmoid。原因往往是注册的名字和m.def里的名字对不上——m.def("swish(...)")注册出来的是myops::swish，而register_custom_op_symbolic里你可能写成了myops.swish或者myops::Swish。大小写和分隔符都要严格对齐，这是硬性要求。

第二个是动态 shape 导出了但没生效。用onnx.load读进来，打印model.graph.input[0].type.tensor_type.shape，如果 dim 里的dim_param全是空的，说明dynamic_axes没被正确处理。检查一下 dummy input 是不是和你传给dynamic_axes的 key 对应，常见的错误是input_names=["input"]但dynamic_axes里写成了"x"。

第三个是opset 版本不匹配导致的算子分解。有些算子在低 opset 里不存在，导出器会尝试用其他算子拼出来，拼出来的结果可能和你的原意不同。比如 LayerNormalization 在 opset 17 之前是分解为 ReduceMean/Sub/Pow 一长串的，如果你在导出后看到这些节点，说明 opset 版本低了。

7.3 运行时类

运行时的坑主要集中在 ORT 侧的自定义算子加载。

No Op registered for Swish with domain_version of 1这个报错的排查顺序我总结为三步。第一步，register_custom_ops_library的路径对不对，用os.path.exists确认一下；第二步，共享库里的extern "C"符号能不能被找到，用nm -D libmyops_ort.so | grep -i domain看一下；第三步，CustomOpDomain的名字和模型里节点的 domain 是不是完全一样，注意这里区分大小写，myops和MyOps是两个不同的 domain。

另一个高频问题是输出是 NaN。自定义算子里出现 NaN，最常见的来源是0 / 0或者inf - inf。SiLU 这个式子里，expf(-x)在 x 很接近 0 时没问题，但如果你的输入被量化过或者有极大值，expf可能溢出成inf，然后x / inf = 0，反向传播时又变成inf * 0 = NaN。稳妥一点的做法是在核函数里对极端值做 clamp，或者把计算式换成数值更稳定的形式。我一般会加一个断言，在前向结束时检查isfinite，一旦出现非有限值就打出前几个元素的原始值，方便快速定位。

还有一个很隐蔽的问题：多线程或多实例下的状态冲突。有些实现为了省事，把中间缓存放在了 kernel 类的成员变量里，结果多个 Session 或者多线程并发时互相踩数据。正确做法是 kernel 内部的临时变量全部用栈上局部变量，或者每次 Compute 里现分配。OrtKernelContext提供的输出内存足够你用，不需要自己维护 cache。

最后分享一个我在项目里逐渐收敛出来的习惯：把自定义算子的 forward 单测、导出后的 onnx 结构校验、ORT 加载后的数值对齐，写成一个脚本里的三个函数，CI 每次提交都跑一遍。这套东西写起来不到 200 行，但它拦住过的回归问题至少有一次直接避免了一次线上事故。自定义算子是那种“改一行、崩一片”的东西，靠人眼 review 实在不太可靠。