PyTorch自定义算子开发：从C++实现到AOTInductor部署的完整流程

在深度学习模型迭代过程中，原生算子往往无法覆盖所有业务场景 —— 无论是融合多个标准操作以减少 Kernel 启动开销，还是接入专用硬件的定制指令，自定义算子（Custom Operator）都是性能优化的关键路径。PyTorch 提供了从 C++ 实现到 Python 绑定的完整扩展机制，本文将梳理从算子开发到 AOTInductor 部署的工程实践。

1. 算子实现：CPU 与 CUDA 双版本

自定义算子的核心实现通常需要同时提供 CPU 和 CUDA 两个版本，PyTorch 的运行时将根据输入张量的设备类型自动分发到对应实现。

CPU 实现相对直接，主要关注算法正确性和边界条件处理。使用 TORCH_CHECK 进行输入校验是推荐做法，可在调试阶段快速定位设备不匹配问题。

CUDA 实现则需要考虑 GPU 并行架构的特性。一个典型的 CUDA 算子实现包含以下要素：

• 线程配置：通常选择 256 作为线程块大小（constexpr int kThreads = 256），这是基于现代 GPU 的 warp 大小（32）和 occupancy 优化的经验值
• 数据类型分发：使用 AT_DISPATCH_FLOATING_TYPES_AND2 宏支持 Float、Half、BFloat16 等多种精度，避免为每种类型重复编写内核代码
• 设备张量操作：通过 torch::tensor 创建设备端张量存储形状和步长信息，避免在 kernel 中频繁访问主机内存
• 错误检查：C10_CUDA_KERNEL_LAUNCH_CHECK() 用于捕获 CUDA 内核启动错误

块数计算遵循标准模式：(numel + kThreads - 1) / kThreads，确保所有元素都被覆盖。

2. 注册机制：TORCH_LIBRARY 宏体系

PyTorch 的算子注册采用分层宏设计，将 schema 定义与实现分离：

TORCH_LIBRARY：用于定义算子的 schema（签名），声明输入输出类型。例如 m.def("identity_conv_op(Tensor x) -> Tensor") 定义了一个接受张量、返回张量的算子。

TORCH_LIBRARY_IMPL：用于注册具体实现，通过模板参数指定命名空间和设备类型（CPU/CUDA）。这种分离设计允许在不修改 schema 的情况下扩展新后端（如 ROCm、XPU）。

对于有状态的算子（如包含可学习参数的层），PyTorch 提供了 torch::CustomClassHolder 基类。自定义类需要在 TORCH_LIBRARY 中通过 m.class_<T>() 注册，并显式声明构造函数、方法以及序列化支持（def_pickle）。__obj_flatten__ 方法的实现对于 torch.export 的追踪至关重要。

3. Python 互操作：Fake 实现与符号追踪

将 C++ 共享库暴露给 Python 运行时只需一步：torch.ops.load_library("libcustom_ops.so")。然而，要让自定义算子兼容 torch.export 和 torch.compile，还需要额外的抽象层。

Fake 实现（Abstract Implementation） 是 torch.export 进行符号追踪时的 "替身"。由于追踪阶段不执行实际计算，PyTorch 需要知道算子的输出形状和类型。通过 @torch.library.register_fake 装饰器，可以为自定义算子提供轻量级的 Python 实现，通常只需返回 torch.empty_like(x)。

对于自定义类，需使用 @register_fake_class 装饰器提供对应的 Fake 类，确保 __obj_flatten__ 和 __obj_unflatten__ 方法与 C++ 实现保持一致。这一层抽象是连接动态图与静态图编译的关键桥梁。

4. 部署流程：AOTInductor 与 C++ 推理

完成算子开发和 Python 端验证后，下一步是将模型编译为可部署格式。AOTInductor（Ahead-of-Time Inductor）允许将 PyTorch 模型预编译为 .pt2 包，支持 Python 和 C++ 双端推理。

导出阶段：使用 torch.export.export(model, args) 捕获计算图，此时自定义算子会显示为 torch.ops.my_ops.identity_conv_op 或 torch.ops.higher_order.call_torchbind（针对自定义类）。

编译阶段：torch._inductor.aoti_compile_and_package 将导出后的程序编译为优化后的二进制，自定义算子的实现会被打包到共享库中。

C++ 推理：在纯 C++ 环境中加载 .pt2 包时，需要通过 dlopen 显式加载自定义算子共享库。值得注意的是，这种方式不依赖 pybind11 或 libpython，仅需链接 LibTorch 和 CUDA 运行时。AOTIModelPackageLoader 会自动解析模型结构并绑定到已注册的算子实现。

5. 工程 Checklist

基于上述流程，自定义算子开发的完整路径可归纳为：

实现层

• CPU 实现：完成算法逻辑，添加 TORCH_CHECK 校验
• CUDA 实现：配置 256 线程块，实现类型分发宏，添加内核启动检查
• 自定义类（可选）：继承 CustomClassHolder，实现序列化方法

注册层

• 使用 TORCH_LIBRARY 定义 schema
• 使用 TORCH_LIBRARY_IMPL 注册 CPU/CUDA 实现
• 自定义类注册构造函数、方法及 pickle 支持

Python 层

• 实现 Fake 算子（register_fake）
• 实现 Fake 类（register_fake_class），保持 flatten/unflatten 一致
• 验证 torch.export 导出成功

部署层

• AOTInductor 编译通过
• Python 端推理验证（torch._inductor.aoti_load_package）
• C++ 端推理验证（dlopen 加载共享库 + AOTIModelPackageLoader）

自定义算子的开发虽然增加了代码维护成本 —— 需要同步更新 CPU/CUDA 双版本、保持 Fake 实现与真实实现语义一致 —— 但对于性能敏感场景，这种投入通常是值得的。特别是在需要融合多个小算子、接入第三方库或支持新硬件时，掌握完整的扩展开发流程是 AI 系统工程师的必备技能。

参考来源

• Lei Mao, "PyTorch Custom Operation", 2026-05-10. https://leimao.github.io/blog/PyTorch-Custom-Operation/

ai-systems