PyTorch 训练框架内核深度解析

第22章自定义算子与 C++ 扩展

作者杨艺韬 · 16,045 字 · 发布于 2026-04-28 · 更新于 2026-09-25

"Reasons why you may want to create a custom op include:"

"Wrapping a third-party library or custom kernel to work with PyTorch subsystems like Autograd."

"Preventing torch.compile/export/FX tracing from peeking inside your function."

—— torch.library.custom_op 的 docstring（torch/_library/custom_ops.py:63-66）。第二条常被忽略：custom_op 是不透明的

本章要点

推荐 torch.library.custom_op 装饰器：一个 API 注册算子 + 接入 dispatcher / autograd / torch.compile
register_fake 给 FakeTensor 路径：torch.compile / FSDP 在 trace 时需要"shape 推导而不真算"
register_autograd 加反向规则：用类似 autograd.Function.backward 的写法
C++ 扩展走 TORCH_LIBRARY：性能敏感时手写 C++ / CUDA kernel，注册进 dispatcher 后经 torch.ops 调用，不需要 pybind11（§22.4）
custom_op 对编译栈是黑盒：torch.compile / export 不会 trace 进函数体。要让 Inductor 看见里面的 Triton kernel，用 torch.library.triton_op + wrap_triton
替代老 API：autograd.Function 还能用，但 torch.compile 兼容性差，新代码用 custom_op

22.1 何时需要自定义算子

PyTorch 的 native_functions.yaml 里有 2674 条算子 schema 声明（v2.11 实测），但仍有缺口：

新硬件指令：自家芯片有特殊指令，想用就要写 kernel 包成 PyTorch op
新算子：论文里某个新激活函数、特殊归一化，PyTorch 还没收
性能极致：某段热路径手写 kernel，把本来要多次往返 HBM 的一串 ATen 算子合成一趟
第三方库集成：外部 CUDA / Triton kernel 想暴露成 torch op，享受 autograd / compile / profiler

如何让自家 kernel 像内置算子一样工作 —— autograd 自动反向、torch.compile 能编译、profiler 能看到、FSDP 能正确处理 —— 是本章主题。

22.2 现代标配：`torch.library.custom_op`

现代推荐写法：

import torch

@torch.library.custom_op("mylib::mymul", mutates_args=())
def my_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    return x * y

@my_mul.register_fake
def _(x, y):
    # FakeTensor 路径: 只返回正确 shape/dtype 的 empty tensor
    return torch.empty_like(x)

def my_mul_backward(ctx, grad):
    x, y = ctx.saved_tensors
    return grad * y, grad * x

def my_mul_setup_context(ctx, inputs, output):
    x, y = inputs
    ctx.save_for_backward(x, y)

my_mul.register_autograd(my_mul_backward, setup_context=my_mul_setup_context)

——这一段做了三件事：

@custom_op 把 my_mul 注册到 dispatcher 的 mylib::mymul schema
register_fake 给 FakeTensor 路径提供 shape 推导
register_autograd 给反向规则

之后 my_mul(x, y) 就像内置算子一样工作。

torch/library.py 在 v2.11 提供的模块级注册点大致是这些（各自的定义行号见 §22.21 的表）：custom_op（从 torch/_library/custom_ops.py:51 re-export）、triton_op、wrap_triton、register_kernel、register_autocast、register_fake、register_autograd、register_torch_dispatch、register_vmap、opcheck、以及 impl / define / infer_schema / Library / get_ctx / get_kernel / fallthrough_kernel / impl_abstract（register_fake 的旧名）。每个对应 PyTorch 的一个子系统 —— 这份清单本身就是"算子要接哪些生态"的检查表。

提醒一句：torch/library.py:29 的 __all__ 只列了其中 15 个，register_autograd / register_kernel / opcheck 三个不在里面（它们仍是可以直接调的公开函数，只是没进 __all__，from torch.library import * 带不出来）。

22.2.1 schema 字符串

"mylib::mymul" 是命名空间 + 算子名。mutates_args=() 表示"不修改任何输入"（如果修改了 x，要写 mutates_args=("x",)）。完整 schema 由 PyTorch 从函数 type hint 自动推导：

mylib::mymul(Tensor x, Tensor y) -> Tensor

如果你的 op 改了输入张量，schema 用 Tensor(a!) x 标记 alias。这套语法第 6 章 §6.2 讲过。

22.2.2 `register_fake` 的角色

FakeTensor 在第 5 章 §5.7 与第 13 章 AOTAutograd 出现过。几乎所有现代 PyTorch 高级特性都依赖 fake 路径：

torch.compile 用它做 graph capture
FSDP 用它做 lazy init / shape 推导
export 用它做 torch.export(model)
meta tensor (无数据张量) 也走 fake

所以没注册 fake 函数的自定义算子在 torch.compile 下直接报错（v2.11 实跑：Dynamo 抛 TorchRuntimeError，内层是 There was no fake impl registered for <CustomOpDef(...)>，不是退回 eager 的 graph break）。register_fake 不是可选 —— 现代代码必须有。opcheck 的 test_faketensor（§22.7）专门查这一项。

fake 函数只允许调用 shape 操作（empty_like / zeros / view / 算 shape），不能做实际数值计算。第 6 章 §6.4.3 警告过这条。

22.2.3 `register_autograd`：反向规则

register_autograd 接受两个函数：backward 和 setup_context。语义与 autograd.Function 类似，但分开成两步：

setup_context(ctx, inputs, output)：保存反向需要的张量（在 forward 完成后调用）
backward(ctx, *grads)：算反向

PyTorch 内部把这套包成 autograd Node，与第 7 章讲的 XxxBackward0 完全等价。自定义算子的反向图与内置算子的反向图无差别，能被 autograd Engine（第 8 章）调度、被 AOTAutograd（第 13 章）capture。

22.3 Triton kernel 作为 custom_op 的实现

用 Triton 写 kernel 时有两种包法，语义完全不同，选错会白白丢掉 fusion 机会。先看常见的那种 —— 用 custom_op 直接包：

import triton
import triton.language as tl

@triton.jit
def my_kernel(x_ptr, y_ptr, out_ptr, n: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * 128 + tl.arange(0, 128)
    mask = offsets < n
    x = tl.load(x_ptr + offsets, mask)
    y = tl.load(y_ptr + offsets, mask)
    tl.store(out_ptr + offsets, x * y, mask)

@torch.library.custom_op("mylib::triton_mul", mutates_args=())
def triton_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    n = x.numel()
    grid = lambda meta: (triton.cdiv(n, 128),)
    my_kernel[grid](x, y, out, n)
    return out

这里有个广为流传的误解要纠正：torch.compile 不会 inline 进 custom_op 的函数体。custom_op 的 docstring 把"Preventing torch.compile/export/FX tracing from peeking inside your function"列为它的用途之一 —— 它就是设计成黑盒的。编译器只知道"调这个 op、shape 由 fake 函数给出"，不会把里面的 Triton kernel 与周围算子 fuse。

想让编译器看进去，用 torch.library.triton_op + torch.library.wrap_triton（torch/_library/triton.py:297 / :504）：

下面是读者示例，不是 PyTorch 源码原文——triton_mul 是示例自取的算子名。

和上面那份 custom_op 写法比，差别只有两处：装饰器换成 triton_op，kernel 启动那一行套上 wrap_triton。

from torch.library import triton_op, wrap_triton

@triton_op("mylib::triton_mul", mutates_args={})
def triton_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    n = x.numel()
    # 函数体里调 triton kernel 必须套一层 wrap_triton
    wrap_triton(my_kernel)[(triton.cdiv(n, 128),)](x, y, out, n)
    return out

它的 docstring 把差别讲得很直白：

custom_op treats custom operators as opaque (torch.compile and torch.export.export will never trace into them), but triton_op makes the implementation visible to these subsystems, allowing them to optimize the triton kernel(s).

约束是 triton_op 的函数体只能由 PyTorch 认识的算子和 triton kernel 组成，且每个 triton kernel 调用都要经过 wrap_triton。

选择规则：

场景	用哪个
实现就是 1 个或多个 triton kernel，希望被 Inductor 一起优化	`triton_op` + `wrap_triton`
实现是 C++/CUDA kernel、第三方库、或不希望被 trace 进去	`custom_op`
只是想在 eager 里调 triton kernel、不需要它表现得像内置算子	两个都不用，直接调（docstring 里明说这是最简单的做法）

22.4 C++ / CUDA 扩展

性能极敏感时手写 C++（含 CUDA）。流程：

写 .cpp 文件，用 TORCH_LIBRARY 注册算子
写 setup.py 用 torch.utils.cpp_extension.CUDAExtension
python setup.py install 编译成 .so
Python 端加载编出来的 .so 即可

C++ 端：

// my_ops.cpp
#include <torch/extension.h>
#include <torch/library.h>

at::Tensor my_mul_cpu(const at::Tensor& x, const at::Tensor& y) {
    return x * y;
}

// 定义在 my_kernel.cu 里（由 nvcc 编译）: 分配 out, 用 <<<grid, block>>> launch CUDA kernel
// .cpp 由主机编译器编译, 不认 <<<>>> 语法, 所以这里只放声明
at::Tensor my_mul_cuda(const at::Tensor& x, const at::Tensor& y);

TORCH_LIBRARY(mylib, m) {
    m.def("mymul(Tensor x, Tensor y) -> Tensor");
}

TORCH_LIBRARY_IMPL(mylib, CPU, m) {
    m.impl("mymul", my_mul_cpu);
}

TORCH_LIBRARY_IMPL(mylib, CUDA, m) {
    m.impl("mymul", my_mul_cuda);
}

setup.py：

from setuptools import setup
from torch.utils.cpp_extension import CUDAExtension, BuildExtension

setup(
    name='mylib',
    ext_modules=[CUDAExtension('mylib', ['my_ops.cpp', 'my_kernel.cu'])],
    cmdclass={'build_ext': BuildExtension},
)

加载后 Python 端：

import torch
# my_ops.cpp 只用 TORCH_LIBRARY 注册、没有 PYBIND11_MODULE, 不能当 Python 模块 import;
# 用 load_library 把 .so 载入进程, 静态注册随之生效（`import torch.ops.mylib` 在 v2.11 实跑报 ModuleNotFoundError）
torch.ops.load_library("<setup.py 产出的 mylib*.so 路径>")
out = torch.ops.mylib.mymul(x, y)

C++ 扩展是芯片厂商接 PyTorch 的标准路径 —— 在 cpp 端用自家 SDK 写 kernel，注册到 dispatcher 的 PrivateUse1 key（§22.9）。

22.5 老 API：`autograd.Function`

老写法仍然支持：

class MyMul(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, y):
        ctx.save_for_backward(x, y)
        return x * y

    @staticmethod
    def backward(ctx, grad):
        x, y = ctx.saved_tensors
        return grad * y, grad * x

out = MyMul.apply(x, y)

简单直接，调试方便。但缺点：

torch.compile 会 trace 进 forward / backward：上面这种纯 ATen 组成的 MyMul 在 v2.11 用 torch._dynamo.explain 实跑是 0 次 graph break，fullgraph=True 也能过；但函数体里一旦有 Dynamo 处理不了的东西（第三方 C/CUDA 扩展、data_ptr 等），就会 graph break 或报错
没有 schema：它不是 dispatcher 里的算子，没有 torch.ops.xxx 名字，export 出来的图里也不会留下一个独立节点
没有 fake 实现可挂：函数体若是不透明 kernel，compile / export 就没法推 shape

如果实现就是一串 PyTorch 算子：autograd.Function 够用，compile 也能 trace 进去。如果实现是不透明的 kernel（第三方库 / C++ / CUDA）且要进 compile / export：用 torch.library.custom_op。

顺带说一句：register_autograd 的 docstring 明确写着 backward 的语义"与 torch.autograd.Function.backward 相同"、ctx 就是同一个 ctx 对象。所以从 autograd.Function 迁到 custom_op，反向那段代码基本可以原样搬 —— 只是把 forward 里 ctx.save_for_backward 那部分拆成独立的 setup_context。

第 7 章 §7.8.1 我们对比过两套接口，结论一致。

22.6 完整集成检查清单

写一个生产级自定义算子，要做的事：

flowchart TB
    Op[custom_op 装饰器]
    Op --> Fake[register_fake<br/>shape 推导]
    Op --> Auto[register_autograd<br/>反向规则]
    Op --> Cpu[CPU kernel<br/>register_kernel device='cpu']
    Op --> Cuda[CUDA kernel<br/>register_kernel device='cuda']
    Fake --> Compile[✓ torch.compile 兼容]
    Cpu --> Eager[✓ eager 路径]
    Cuda --> Eager
    Auto --> Eng[✓ autograd Engine]
    Style0[注册到 dispatcher<br/>自动获得]
    Op --> Style0
    Style0 --> Disp[dispatch 调度]
    Style0 --> Prof[profiler 自动看到]
    Style0 --> Fsdp[FSDP / DDP 兼容]
    style Op fill:#fef3c7,stroke:#f59e0b,stroke-width:2px

清单：

✅ @custom_op 装饰器声明
✅ register_fake 给每个 op
✅ register_autograd 如果可微
✅ register_kernel(..., 'cpu') + register_kernel(..., 'cuda') 各自实现
✅ 写单元测试用 torch.library.opcheck 自动验证（PyTorch 提供的算子合规性检查）

22.7 `opcheck`：自定义算子的合规性测试矩阵

torch/library.py:1632 的 torch.library.opcheck 是自定义算子的"质保检查"：

import torch
from torch.library import opcheck

opcheck(my_mul, args=(x, y))       # 默认跑 4 项

默认跑的是 _OPCHECK_DEFAULT_UTILS（torch/library.py:1624）里那 4 项：

测试	按 docstring，检查什么
`test_schema`	schema 声明的 mutation / aliasing 与实现是否一致 —— schema 说会 mutate 就必须真 mutate，说返回新 Tensor 就不能返回已有张量或它的 view。注意它查的不是 dtype / shape
`test_autograd_registration`	op 支持训练时，反向公式是否通过 `register_autograd` 或 `DispatchKey::Autograd` 正确注册。它查的是"注册方式对不对"，不是数值梯度对不对
`test_faketensor`	有没有 FakeTensor kernel，以及真张量跑出来的 metadata（sizes / strides / dtype / device）与 fake 跑出来的是否一致
`test_aot_dispatch_dynamic`	e2e：eager 与 torch.compile 的输出（含梯度）是否一致。docstring 说它是 `test_faketensor` 的超集，还顺带验 functionalization 与反向的 fake 兼容性

ALL_TEST_UTILS（torch/testing/_internal/optests/generate_tests.py:150）里还有第 5 项 test_aot_dispatch_static（静态 shape 版），但它只在 DEPRECATED_DEFAULT_TEST_UTILS 里 —— 不在默认列表，别照着老教程去传它。

想验数值梯度得自己另外调 torch.autograd.gradcheck，opcheck 不做这件事。

生产级自定义算子必须 opcheck 通过，并且 docstring 给了三条使用建议：用多组有代表性的输入反复调；op 支持 autograd 就用 requires_grad=True 的输入调；支持多设备就每个设备各调一遍。

22.8 `Library` 低级 API

@custom_op 是高级糖，底层是 torch.library.Library（torch/library.py:68）。它提供更细粒度的算子注册：

from torch.library import Library

# 创建一个 library (类似 C++ 端的 TORCH_LIBRARY)
lib = Library("mylib", "DEF")

# 注册 schema (没有实现, 等下注册)
lib.define("mymul(Tensor x, Tensor y) -> Tensor")

# 给特定 dispatch key 注册实现
lib.impl("mymul", lambda x, y: x * y, "CPU")
lib.impl("mymul", my_cuda_kernel, "CUDA")
lib.impl("mymul", my_meta_kernel, "Meta")    # FakeTensor 也是 Meta key

第二个参数 "DEF" 是 library 的 kind，源码里只接受这三个值（torch/library.py:90 的校验）：

DEF：创建一个新 namespace 来注册新算子
IMPL：给已有算子加新 dispatch key 实现（要给 aten:: 的内置算子挂 backend kernel 就用它）
FRAGMENT：往已有 namespace 里追加新算子

_reserved_namespaces（:58）目前只有 "prim" 一个成员：用 DEF 或 FRAGMENT 去开 prim 会直接报错（:93）。

@custom_op 装饰器内部就是构造 Library 然后调 define / impl。直接用 Library 时你能精确控制每个 dispatch key 的实现 —— 适合需要"给 PrivateUse1 注册新 backend"等高级场景。

22.9 PrivateUse1：国产芯片接入完整路径

PyTorch 给厂商扩展自家硬件留了 3 个 dispatch key：PrivateUse1 / PrivateUse2 / PrivateUse3（第 3 章 §3.5）。完整接入流程：

# 1. 给 PrivateUse1 起个有意义的名字
torch.utils.rename_privateuse1_backend("npu")
# 之后用户可以写 tensor.to('npu') 而非 'privateuseone'

# 2. 给 PrivateUse1 注册所有 ATen 算子的实现
@torch.library.impl("aten::add.Tensor", "PrivateUse1")
def npu_add(self, other, alpha=1):
    # 调你家硬件 SDK 的 add kernel
    return _npu_runtime.add(self, other, alpha)

# ... 给几百个常用算子各注册一个 impl ...

# 3. 提供 generate_methods_for_privateuse1_backend 让 tensor.npu() 等方法可用
torch.utils.generate_methods_for_privateuse1_backend()

torch/utils/backend_registration.py:20 的 rename_privateuse1_backend 把 PrivateUse1 重命名，:350 的 generate_methods_for_privateuse1_backend 自动给 Tensor 添加 .npu() / .is_npu 这类方法。这套 API 让芯片厂商可以做出完整 PyTorch 体验而不修改主仓代码。

实际工作量：给 PyTorch 全部算子各写一个 backend impl 不现实，但 torchgen/gen_backend_stubs.py（第 6 章 §6.9）能从一份「目标算子列表 YAML」自动生成 stub 代码，厂商只需要填实现细节，并且可以只覆盖真正用到的那批算子。

torch_npu（华为）、torch_mlu（寒武纪）等厂商扩展都走这条路，代码开源，可以直接当模板读。

22.10 `allow_in_graph` 与 `disable`：torch.compile 的两个逃生口

写自定义算子时常遇到 Dynamo 不会 trace 的代码（如调了第三方 C 扩展、动态行为太复杂）。PyTorch 提供两个装饰器作为逃生口：

@torch.compiler.allow_in_graph（torch/compiler/__init__.py:72）：

下面是读者示例，不是 PyTorch 源码原文——my_special_function 是示例自取的函数名。

装饰器加在哪个函数上，Dynamo 就不再往那个函数体里 trace。

@torch.compiler.allow_in_graph
def my_special_function(x, y):
    # Dynamo 不 trace 这个函数体, 把整个调用直接写进 graph
    # 但后端 (AOTAutograd + Inductor) 仍会 trace 进来, 函数体必须对它们可 trace
    return some_lib.do_magic(x, y)

效果：Dynamo 看到调用 my_special_function(x, y) 时，不展开内部、把这次调用作为一个节点写进 FX Graph。但它只对前端 Dynamo 不透明：docstring 里的对照表写得很清楚，allow_in_graph 在后端（AOTAutograd + Inductor）那一栏是 "trace inside"，只有 custom op 才对整条编译栈都是 "opaque callable"。docstring 还直接说：想在整个 tracing 过程中把函数当黑盒，别用这个 API，去写 custom op。

@torch.compiler.disable（torch/compiler/__init__.py:249，也可写 torch._dynamo.disable）：

下面是读者示例，不是 PyTorch 源码原文——my_complex_logic 是示例自取的函数名。

disable 的语义比 allow_in_graph 更狠：不是把函数当不透明 op 收进图，而是直接 graph break。

@torch.compiler.disable
def my_complex_logic(x):
    # Dynamo 看到这个调用直接 graph break, 退回 eager
    if x.sum() > 0:
        return some_python_heavy_logic(x)
    else:
        return another_branch(x)

效果：Dynamo 在调用处触发 graph break，整段函数用 eager 跑，break 之后再开始新 trace。

两者关键区别：

装饰器	Dynamo 行为	适合场景
`allow_in_graph`	不进函数体、整个调用写进 graph（后端仍 trace 进去）	函数对 AOTAutograd / Inductor 是可 trace 的，只是 Dynamo 有 bug 或看不进去
`disable`	触发 graph break，退回 eager	函数有复杂 Python 逻辑（动态控制流 / 大量 dict 操作 / print），不希望 Dynamo 浪费时间分析

实际工程里：

写自定义 Triton kernel + register_fake：用 custom_op（§22.2），不需要这两个装饰器
集成第三方 C/CUDA 扩展（如 FlashAttention v1 的私有 wrapper）：要对整条编译栈当黑盒，用 custom_op + register_fake，而不是 allow_in_graph（后端会 trace 进去，碰到不透明 kernel 就过不去）
训练循环里的 logging / metric reporting 函数：用 disable 让 Dynamo 不要试图分析

torch/_dynamo/decorators.py 还提供更细的开关：disallow_in_graph（强制某 op 触发 graph break）、mark_static_address（声明 tensor 地址不会变，让 CUDA Graph 能复用）等。生产代码里写自定义算子的 escape hatch，理解这套装饰器家族能让你优雅处理"compile 不动"的边角情况。

22.11 inplace 与多输出算子的注册

前面的例子都是 mutates_args=()、单个 Tensor 输出的"纯函数"。两种特殊形态要换写法：

inplace 算子（mutate input）：

@torch.library.custom_op("mylib::add_inplace_", mutates_args=("x",))
def add_inplace_(x: torch.Tensor, y: torch.Tensor) -> None:
    x.add_(y)
    # 没有 return: v2.11 实跑推出的 schema 是 mylib::add_inplace_(Tensor(a0!) x, Tensor y) -> ()

mutates_args=("x",) 让 schema 里 x 标上 alias 与写入标记 Tensor(a0!)。functionalize（§13.4）看到这个标记后会改写调用：它读不懂函数体，所以不会把它展开成 x + y，而是把这次调用包进 auto_functionalized_v2 这个高阶算子 —— 在 x 的副本上调原 op，再把结果当作新值传下去（v2.11 用 aot_function 实跑，前向图里就是 torch.ops.higher_order.auto_functionalized_v2(torch.ops.mylib.add_inplace_.default, ...)）。这是 v2.x 让 inplace 自定义算子与 compile 共存的关键。

mutates_args 漏写了实际被修改的参数 → 隐蔽 bug：torch.compile 假设无副作用、生成的 kernel 不会复制 x，运行时 x 被修改但 graph 看不到 → 后续算子拿到错的 x。

多输出算子：

@torch.library.custom_op("mylib::topk_with_idx", mutates_args=())
def topk_with_idx(x: torch.Tensor, k: int) -> tuple[torch.Tensor, torch.Tensor]:
    values, indices = torch.topk(x, k)
    return values, indices

@topk_with_idx.register_fake
def _(x, k):
    new_shape = list(x.shape)
    new_shape[-1] = k
    return torch.empty(new_shape, dtype=x.dtype), torch.empty(new_shape, dtype=torch.int64)

返回 Tuple[Tensor, ...] 时 schema 自动是 -> (Tensor, Tensor)。fake 函数也返回 tuple。

inplace + 多输出组合：

@torch.library.custom_op("mylib::layernorm_inplace", mutates_args=("x", "running_mean"))
def layernorm_inplace(
    x: torch.Tensor,
    running_mean: torch.Tensor,
    weight: torch.Tensor,
) -> torch.Tensor:
    # 修改 x 与 running_mean, 返回新 tensor
    ...

复杂场景里这套语法要小心写。schema 错了 → AOTAutograd 会在 trace 时报"functionalize 失败"。opcheck 内置 functionalize 检查能在 commit 前发现这类问题（§22.7）。

22.12 register_kernel：每个 device 单独注册

@custom_op 的函数体是 op 的默认实现，注册在 CompositeExplicitAutograd key 上（torch/_library/custom_ops.py:366）。如果你想为不同 device 写专门 kernel，用 register_kernel（也可以在装饰器上直接写 device_types="cuda"）：

下面是读者示例，不是 PyTorch 源码原文——mymul 是示例自取的算子名。

看两处就够：装饰器函数体是默认实现，@mymul.register_kernel("cuda") 再给 CUDA 单独挂一份。

@torch.library.custom_op("mylib::mymul", mutates_args=())
def mymul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    # 默认实现 (eager 路径用)
    return x * y

@mymul.register_kernel("cuda")
def _(x, y):
    # CUDA 专用: 调 Triton kernel
    out = torch.empty_like(x)
    grid = (triton.cdiv(x.numel(), 128),)
    my_triton_kernel[grid](x, y, out, x.numel())
    return out

@mymul.register_kernel("cpu")
def _(x, y):
    # CPU 专用: 调 OpenMP kernel
    out = torch.empty_like(x)
    my_cpp_extension.cpu_mul(x, y, out)
    return out

@mymul.register_kernel("xpu")
def _(x, y):
    # Intel XPU 专用
    return x * y    # 通用 fallback

dispatcher（第 5 章）根据 input device 自动路由。这套机制让"一个 op 多 backend"不需要写 dispatch 逻辑、PyTorch 帮你做。

以 FlashAttention 为例，按这套机制可以这样分工（示意）：

默认实现：调用 F.scaled_dot_product_attention（fallback）
CUDA：调自家 CUDA kernel（Hopper / Ampere 各一份）
CPU：调 PyTorch 通用 attention（性能差但能跑）

理解 register_kernel 让你看到自定义算子的"多后端"不需要复杂代码 —— 装饰器 + dispatcher 自动协作。

22.13 JIT 加载 C++ 扩展：开发期免编译

§22.4 用 setup.py 编译 C++ 扩展，每次改完要重新 build。开发期更方便的方式是 torch.utils.cpp_extension.load：

import torch.utils.cpp_extension as cpp_ext

cpp_ext.load(
    name='my_ops',
    sources=['my_ops.cpp', 'my_kernel.cu'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '-arch=sm_90'],
    is_python_module=False,   # §22.4 的 my_ops.cpp 只有 TORCH_LIBRARY、没有 PYBIND11_MODULE
    verbose=True,
)

# .so 已作为普通动态库载入进程, 算子走 torch.ops 调
out = torch.ops.mylib.mymul(x, y)

load 内部：

把 sources 编译成 .so（首次要等完整编译）
缓存到用户缓存目录下的 torch_extensions/（Linux 上通常是 ~/.cache/torch_extensions/，可用 TORCH_EXTENSIONS_DIR 改）
后续相同 sources 跳过编译、直接加载
改了 source 自动重编

适合开发场景：写 / 改 / 测的循环里不用每次跑 setup.py install。

进阶：load_inline 让你直接传 C++ source 字符串、不用文件：

my_ops = cpp_ext.load_inline(
    name='inline_ops',
    cpp_sources='''
        torch::Tensor add_one(torch::Tensor x) {
            return x + 1;
        }
    ''',
    functions=['add_one'],
)

适合写小 demo / unit test。生产代码仍用 setup.py + .so 文件（避免每次进程启动都编译）。

实战：研究迭代算法时，load_inline + Jupyter notebook 让你能像写 Python 一样快速迭代 C++ kernel。这套工程便利极大降低了"写 C++ 扩展"的心智门槛。

22.14 ABI 兼容性：跨 PyTorch 版本的痛点

C++ 扩展编译出的 .so 对 PyTorch 版本敏感。原因：

libtorch C++ ABI 不冻结：PyTorch 团队在 v2.x 多次重构内部 API
CUDA Toolkit 版本：编译用 12.4、运行时 12.5+ OK；但 12.4 → 11.8 不行
Compiler ABI：扩展与 PyTorch 用的 C++ 标准库或 _GLIBCXX_USE_CXX11_ABI 设置不一致时，加载会报 undefined symbol（本专栏在本机用 libc++ 编的 TORCH_LIBRARY 扩展去加载 libstdc++ 构建的 torch 2.11 wheel，就报了 undefined symbol: _ZN5torch7LibraryC1...NSt3__112basic_string...）

典型症状是加载 .so 时报 ImportError: undefined symbol: _ZN3c10...（_ZN3c10 是 c10:: 命名空间的 C++ mangled 前缀）—— 编译时链接的那个符号在当前 PyTorch 里已经改名或消失了。

解决方案：

1. 锁版本 + per-version build

# 用户安装时根据 PyTorch 版本下载对应 wheel
pip install my-extension==0.1.0+pt2.6
pip install my-extension==0.1.0+pt2.4

每个 PyTorch 主版本编一份 wheel。

2. 编译环境对齐

TORCH_CUDA_ARCH_LIST 控制为哪些 GPU 架构生成代码（torch/utils/cpp_extension.py:1397 有说明），编译机与目标机的编译器、CUDA Toolkit、_GLIBCXX_USE_CXX11_ABI 设置都要与构建该 PyTorch wheel 时一致：

TORCH_CUDA_ARCH_LIST="8.0 9.0" python setup.py bdist_wheel

3. JIT load (§22.13)

绕过 ABI 问题：用户机器现场编 → 自动用当前 PyTorch 的 ABI。代价是首次启动慢。

4. 用 ABI-stable 头文件

torch/csrc/stable/（§22.28）就是为这条痛点准备的：只用它暴露的类型与函数，.so 就不再绑死内部 ABI。

实战：广泛分发的开源 PyTorch 扩展基本都在 CI 里维护一张"PyTorch 版本 × CUDA 版本 × Python 版本"的 build matrix，生产部署时把 PyTorch 与扩展版本一起锁死。这是写 C++ 扩展躲不掉的工程税 —— 纯 Python + Triton（§22.3）根本不产生 .so，也就完全没有这个问题。

22.15 Composite Implicit Autograd：算子的 decomposition

PyTorch 内置 op 有几类 autograd 处理方式：

Autograd Key	含义
Autograd	显式注册反向规则（如 `mm`，在 `derivatives.yaml` 里写死反向）
CompositeImplicitAutograd	op 内部调其他 op，autograd 自动追踪（不需要写反向）
CompositeExplicitAutograd	一份对所有 backend 通用的实现，但 autograd 不会穿透它，反向要另行显式给出
AutogradPrivateUse1	厂商自家硬件的 autograd 实现

这里要纠正一个非常常见的误解：@torch.library.custom_op 包出来的算子不是 CompositeImplicitAutograd，也不会"自动追踪"函数体里的可微算子。

看源码：torch/_library/custom_ops.py:366 把设备实现注册到 CompositeExplicitAutograd，:641 把 autograd 实现注册到 Autograd key。也就是说 custom_op 对 autograd 是不透明的 —— 它的函数体不进反向图。register_autograd 的 docstring 开头就说得很明白：

In order for an operator to work with autograd, you need to register a backward formula.

所以下面这段代码是错的：

@torch.library.custom_op("mylib::my_attention", mutates_args=())
def my_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
    scores = q @ k.transpose(-2, -1)
    attn = scores.softmax(-1)
    return attn @ v
# ✗ 以为 autograd 会穿过 mm + softmax + mm 自动追踪 —— 不会。
#   这个 op 对 autograd 是黑盒, 不 register_autograd 就求不了梯度: v2.11 实跑, backward 时报
#   RuntimeError: Trying to backward through mylib.my_attention.default but no autograd formula was registered

custom_op 的用途本来就是"把一段实现包成不透明的原子算子"（它的 docstring 把 "Preventing torch.compile/export/FX tracing from peeking inside your function" 列为用途之一），所以：

只要用了 custom_op，需要梯度就必须 register_autograd，函数体里是什么无关紧要。

反过来，如果你的实现就是一串可微的 ATen 算子、并且不需要它表现成一个原子算子，那正确做法是根本别包 custom_op —— 直接写普通 Python 函数，autograd 自然穿过去。这也是 triton_op docstring 里那句建议的同一个意思：能不包就不包，只有需要它"表现得像内置算子"时才包。

CompositeImplicitAutograd 是 ATen 内置算子的一种注册方式（native_functions.yaml 里不写 dispatch: 或标 CompositeImplicitAutograd 的那些，第 6 章），不是 custom_op 的默认行为。

用了 Triton kernel / C++ kernel 时更是必须 register_autograd：

@torch.library.custom_op("mylib::triton_attention", mutates_args=())
def triton_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
    # Triton kernel 内部 op autograd 看不到
    return my_triton_kernel(q, k, v)

# 必须显式注册反向
def backward(ctx, grad_out):
    q, k, v = ctx.saved_tensors
    return triton_backward_kernel(q, k, v, grad_out)

所以判断规则其实很简单：包了 custom_op 就得自己给反向；不想写反向就别包。

CompositeImplicitAutograd 这条路径仍然重要，只是它属于 ATen 内置算子那一侧：大量内置 op 靠"用其他 op 拼出来 + autograd 自动穿透"来省掉手写反向，让 PyTorch 几千算子的反向规则维护成本可控（第 6 章）。理解两侧的差别，才不会把内置算子的规则套到自定义算子上。

22.16 Triton autotune：让 kernel 自动找最优配置

写 Triton kernel 时关键参数（block size / num_warps / num_stages）需要为每个硬件 / shape 调优。手动调耗时，Triton 内置 autotune 自动搜索：

import triton
import triton.language as tl

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE': 128}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_SIZE': 256}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_SIZE': 256}, num_warps=8, num_stages=3),
        triton.Config({'BLOCK_SIZE': 512}, num_warps=8, num_stages=4),
        # ... 更多配置
    ],
    key=['n'],    # n 不同时重新选 config
)
@triton.jit
def my_kernel(x_ptr, y_ptr, out_ptr, n: tl.constexpr,
              BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n
    x = tl.load(x_ptr + offsets, mask)
    y = tl.load(y_ptr + offsets, mask)
    tl.store(out_ptr + offsets, x * y, mask)

工作机制：

第一次某个 n 调用时，autotune 跑所有 configs、测每个的 GPU 时间
选最快的 config
缓存到 (kernel, n) → best_config 映射
后续相同 n 直接用 best_config

代价：某个 key 第一次出现时要把所有 config 各跑一遍测时间，这段延迟与 config 数成正比。生产代码在 warmup 阶段把常见 shape 都过一遍触发 autotune，之后稳态运行命中缓存。

进阶：

prune_configs_by 让你写自定义函数过滤掉 illegal config（如 BLOCK_SIZE 太大超 shared memory）
reset_to_zero 让某些 input 在每次 autotune trial 后清零（避免累积副作用）
do_bench 自定义 benchmark 函数

现代 SOTA kernel 的工程实质就在这里：不是手写一个完美 kernel，而是定义搜索空间 + 让机器搜。同一段 Triton 源码要在不同代 GPU、不同 shape 上都接近硬件峰值，靠人手调参数不现实。Inductor 的 max_autotune 做的也是同一件事（第 14 章）。

22.17 vmap × custom_op：批量化的自动支持

vmap（functorch / torch.func.vmap）让 op 自动批量化：

def add(x, y):
    return x + y

batched_add = torch.func.vmap(add)
# batched_add 接受 [B, ...] 输入, 内部 batched 算 add

PyTorch 内置 op 的 vmap 规则已经写好。自定义 op 默认 vmap 会失败：

@torch.library.custom_op("mylib::mymul", mutates_args=())
def mymul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    return x * y

torch.func.vmap(mymul)(x, y)
# v2.11 实跑: 能出结果, 但打警告 "There is a performance drop because we have not yet
# implemented the batching rule for mylib::mymul" —— 退回 BatchedFallback, 按 batch 逐个调原 op

想要真正批量化的实现，需要 register_vmap（docstring 的措辞是 "you may need to register a vmap implementation"）：

@mymul.register_vmap
def _(info, in_dims, x, y):
    # in_dims: 输入 tensor 沿哪个维度 batch
    # 实现: 把 vmap 输入展开成 normal call
    x_dim, y_dim = in_dims
    if x_dim is not None and y_dim is None:
        y = y.unsqueeze(x_dim).expand_as(x)
    elif y_dim is not None and x_dim is None:
        x = x.unsqueeze(y_dim).expand_as(y)
    out = mymul(x, y)
    out_dim = x_dim if x_dim is not None else y_dim
    return out, out_dim

实战工作量：复杂 op 的 vmap rule 常常比 forward 还难写。务实做法是先不实现（v2.11 下 vmap 仍能走逐个调用的慢速 fallback，见上面的实跑），在文档里说明没有批量化实现，等真有人需要再补。register_vmap（torch/library.py:1351）随时可以后加，不影响已有代码。

理解 vmap 的存在让你知道 PyTorch 的"自动批量化"也是抽象层 + 各 op 单独支持。custom_op 想完整融入 PyTorch 生态需要 fake / autograd / vmap / dispatch 多层注册。

22.18 自定义 op 注册到 Inductor lowering

torch.compile 看到自定义 op 时，默认走 fallback 路径：直接调用原 op、不与周围算子 fuse。如果你想让 Inductor 真正编译你的 op（fuse 到 Triton kernel 里），用 register_lowering：

# 注意这是私有 API (torch/_inductor/lowering.py:529), 没有兼容性保证
from torch._inductor.lowering import register_lowering
from torch._inductor.ir import Pointwise

@register_lowering(torch.ops.mylib.mymul)
def mymul_lowering(x, y):
    # 返回 Inductor IR (Pointwise)
    return Pointwise.create(
        device=x.get_device(),
        dtype=x.get_dtype(),
        inner_fn=lambda idx: x.make_loader()(idx) * y.make_loader()(idx),
        ranges=x.get_size(),
    )

效果：torch.compile 看到 mymul(a, b) + c 时，不是"调 mymul kernel + 调 add kernel"，而是直接把 mymul 的语义编译进同一个 fused Triton kernel —— 真正的 op fusion。

适用：

简单算子（pointwise / reduction）：写 lowering 让 Inductor 优化
复杂算子（attention / GEMM）：保留 fallback，让 Inductor 当黑盒处理

PyTorch 内置 ATen op 大多有 lowering 或 decomposition（其余的同样走 fallback），自定义 op 默认没有。写 lowering 是性能极致场景才做的工作：GEMM / attention 这类本身就是一整个 kernel 的算子，fuse 进周围算子没什么可捞的；而 pointwise 算子留在 fallback 里意味着每次都要完整读写一遍显存，写 lowering 的收益直接体现在省下的 HBM 往返上。

但注意 torch._inductor.lowering 是私有 API（下划线开头），跨版本没有兼容承诺。如果只是想让 Triton kernel 被 Inductor 优化，优先考虑 §22.3 的 triton_op —— 那是公开路径。

理解 lowering 让你看 Inductor 不是"魔法编译器"，是 lowering registry 驱动的代码生成器。

22.19 完整 FlashAttention 接入路径

把全章话题合起来看 FlashAttention 这种 SOTA op 怎么完整接入 PyTorch：

graph TB
    FA[FlashAttention CUDA kernel]
    FA --> CO[custom_op 装饰器<br/>schema: q, k, v -> out]
    CO --> Fake[register_fake<br/>shape 推导]
    CO --> Auto[register_autograd<br/>反向 = 另一个 FA backward kernel]
    CO --> Cuda[register_kernel cuda<br/>调实际 CUDA kernel]
    CO --> Cpu[register_kernel cpu<br/>调用 fallback SDPA]
    Fake --> Compile[torch.compile 兼容]
    Auto --> Backward[autograd Engine 调度反向]
    Cuda --> Eager[eager 路径]
    style FA fill:#fef3c7,stroke:#f59e0b,stroke-width:2px
    style Compile fill:#dcfce7

代码骨架：

import torch
import torch.nn.functional as F
from torch.library import opcheck

@torch.library.custom_op("mylib::flash_attention", mutates_args=())
def flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
    # 默认实现 (CPU fallback)
    return F.scaled_dot_product_attention(q, k, v)

@flash_attention.register_kernel("cuda")
def _(q, k, v):
    # 调真实 CUDA kernel (FlashAttention v3)
    return _flash_attn_v3.forward(q, k, v)

@flash_attention.register_fake
def _(q, k, v):
    return torch.empty_like(q)

def fa_backward(ctx, grad_out):
    q, k, v, out, lse = ctx.saved_tensors    # lse = log-sum-exp, FA 内部产物
    grad_q, grad_k, grad_v = _flash_attn_v3.backward(grad_out, q, k, v, out, lse)
    return grad_q, grad_k, grad_v

def fa_setup_context(ctx, inputs, output):
    q, k, v = inputs
    out, lse = output_with_lse(q, k, v)    # 实际场景里 forward 输出 lse
    ctx.save_for_backward(q, k, v, out, lse)

flash_attention.register_autograd(fa_backward, setup_context=fa_setup_context)

# 测试
opcheck(flash_attention, args=(torch.randn(2, 8, 1024, 64, device='cuda'),) * 3)

部署后用法与内置 op 完全一致：

out = torch.ops.mylib.flash_attention(q, k, v)
# 或者 monkey-patch F.scaled_dot_product_attention 让全局透明用 FA

理解这套接入让你看到：一个外部 CUDA kernel 经过 custom_op 的四件套（schema / fake / autograd / per-device kernel）之后，就和内置算子没有区别 —— autograd 认它、compile 认它、profiler 认它。custom_op 是 PyTorch 生态吸纳新算子的标准接口。

22.20 自家 AI 芯片完整接入 PyTorch 的工程

国产 AI 芯片厂商把硬件接进 PyTorch 是一项大型系统工程。完整路径：

第 1 阶段：基础 backend

# 1. 注册 PrivateUse1 → "npu"
torch.utils.rename_privateuse1_backend("npu")

# 2. 实现 device guard / stream / event 抽象
class NPUStream(...): ...
class NPUEvent(...): ...

# 3. tensor.npu() / tensor.is_npu 等 method
torch.utils.generate_methods_for_privateuse1_backend()

# 4. 其余的 device guard / allocator / stream 注册走 C++ 侧
#    (c10::impl::DeviceGuardImplRegistrar 等), 没有统一的 Python 入口

第 2 阶段：算子实现

# 给最常用的几百个算子各写 NPU impl
# 用 codegen 减少手写代码
@torch.library.impl("aten::add.Tensor", "PrivateUse1")
def npu_add(self, other, alpha=1):
    return _npu_runtime.add(self, other, alpha)

@torch.library.impl("aten::mm", "PrivateUse1")
def npu_mm(self, mat2):
    return _npu_runtime.gemm(self, mat2)

# ... 几百个算子 ...

第 3 阶段：CommunicationBackend (NCCL 替代)

class NPUCommBackend(ProcessGroup):
    def allreduce(self, tensors, opts): ...
    def allgather(self, output_tensors, input_tensors, opts): ...
    # ... 实现完整 c10d ProcessGroup 接口 ...

torch.distributed.Backend.register_backend("hccl", create_npu_comm)

第 4 阶段：编译栈集成

# 给 torch.compile 注册自家 backend (torch/_dynamo/backends/registry.py:84)
from torch._dynamo import register_backend

@register_backend
def npu_compiler(fx_graph, example_inputs):
    # 调自家编译器把 fx_graph 编译成 NPU binary
    return npu_compile(fx_graph, example_inputs)

# 用法
@torch.compile(backend="npu_compiler")
def model(x): ...

第 5 阶段：训练 / 推理生态

FSDP / DDP 适配（用 hccl backend）
AMP / bf16 支持
safetensors / DCP 集成
profile + Kineto 自家 backend

五个阶段里，算子实现的绝对工作量最大（要覆盖的算子数以百计），但编译集成的不确定性最高（要把自家 graph compiler 塞进 Dynamo/Inductor 的接口里）。真实排期高度依赖已有的 SDK 成熟度，没有可搬运的人月数字。

torch_npu（华为）、torch_mlu（寒武纪）等厂商扩展都走过这条路，代码开源，可以逐阶段对照着读。PrivateUse1 + torch.library 是第三方芯片参与 PyTorch 生态的核心入口，全程不需要 fork 主仓代码。

理解这条路径让你看 PyTorch 不是 NVIDIA 专属，是真正"硬件中立"的开放生态。

22.21 自定义算子的演进时间线

与其背版本号，不如把 v2.11 这一刻 torch/library.py 暴露的注册点当成地图（其中 register_autograd / register_kernel / opcheck 不在 __all__ 里，见 §22.2）—— 每个都对应"接入某个子系统"这一件事：

注册点	接哪个子系统	源码位置
`custom_op`	dispatcher（把函数变成不透明算子）	`torch/_library/custom_ops.py:51`
`triton_op` + `wrap_triton`	Inductor（让 triton kernel 可被优化）	`torch/_library/triton.py:297` / `:504`
`register_fake`（旧名 `impl_abstract`）	FakeTensor → compile / export / meta	`torch/library.py:962` / `:781`
`register_autograd`	autograd 引擎	`:1138`
`register_kernel`	各 device 的具体实现	`:795`
`register_autocast`	AMP（第 20 章 §20.20）	`:864`
`register_vmap`	`torch.func.vmap`	`:1351`
`register_torch_dispatch`	tensor subclass / TorchDispatchMode	`:1269`
`opcheck`	合规性测试	`:1632`
`Library` / `define` / `impl`	低层注册（`@custom_op` 的底座）	`:68` 起

整体趋势很清楚：

早期是 autograd.Function（只接 autograd）与 TORCH_LIBRARY C++ 宏（只接 dispatcher）两条互不相通的路
v2.x 收敛成 torch.library 一个 Python 命名空间，一个子系统一个 register_*
新加的注册点（register_autocast、register_vmap、triton_op）都是在补"某个子系统还没打通"的缺口

理解这张表你就知道：自定义 op 的"完整度"不是一个布尔值，而是"接通了几个子系统"。

22.22 常见 bug 排查 cheat sheet

实战写自定义 op 遇到的报错与解法：

报错	根因	解决
`RuntimeError: ... shape mismatch` 在 compile 但不在 eager	fake 函数 shape 推导错	检查 fake 返回 shape 是否与真实 kernel 一致
`Expected at most 0 ... got X`	schema 字符串与函数签名不匹配	type hint 改对 / schema 显式
`mutates_args` 错	functionalize 假设无副作用、kernel 实际 mutate	加正确 `mutates_args=("x",)`
`There was no fake impl registered for <CustomOpDef(...)>`	自定义算子没注册 fake，compile 直接报错（§22.2.2）	补 `register_fake`
Dynamo `Unsupported: ...` graph break	Dynamo 看不进这段代码	包成 custom_op + `register_fake`；若后端能 trace、只是 Dynamo 看不进，可用 `allow_in_graph`（§22.10）
`torch.autograd.gradcheck` 失败	反向规则数值不对	逐项对照公式；注意 opcheck 不查数值梯度，gradcheck 要自己调
Inductor 对该 op 走 fallback、没有 fuse	custom_op 本来就是黑盒（§22.3）	换 `triton_op`（公开路径），或写 `register_lowering`（私有 API）
ABI undefined symbol	C++ 扩展与 PyTorch 版本不匹配	重新编译 / 用 JIT load
vmap 时警告 `we have not yet implemented the batching rule`	没有 vmap 规则，退回逐个调用的慢速 fallback	需要性能时 register_vmap
autograd 反向时 saved_tensors 是 None	`setup_context` 没保存	在 `setup_context` 里调 `ctx.save_for_backward`
`opcheck test_aot_dispatch_dynamic` fail	fake 函数没处理 SymInt 输入	fake 里所有 shape 操作改用 SymInt-friendly API

这几类报错覆盖了写自定义 op 时的绝大多数坑，遇到先按表对号入座能少走很多弯路。

22.23 export 与自定义算子

torch.export（§12.32）把 model 导成 ExportedProgram，给部署用。自定义算子在 export 路径的处理：

@torch.library.custom_op("mylib::flash_attention", mutates_args=())
def flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: ...

class MyModel(nn.Module):
    def forward(self, x):
        q, k, v = split(x)
        return torch.ops.mylib.flash_attention(q, k, v)

# 导出
exported = torch.export.export(MyModel(), example_inputs)

# ExportedProgram 内部的 fx graph 含 mylib::flash_attention 节点
print(exported.graph)

ExportedProgram 内部用 op 的完整 fqn（mylib::flash_attention）记录，而不是 inline op body。部署时：

AOTI：编出的 .so 里保留对 mylib::flash_attention 的调用，运行时仍调原 kernel
ExecuTorch：让 op 走 delegate 到目标硬件
ONNX：自定义 op 不在 ONNX 标准算子集里 → 需要走 custom domain 或提供 translation，见第 19 章 §19.9.3

为让自定义 op 能 export：

必须有 register_fake（export 用 FakeTensor 跑）
schema 要稳定（不能动态加参数）
不能有 graph break（complex Python logic）

自家算子如果将来要部署，一开始就按 export 兼容写：schema 稳定、fake 完整、不在 op 边界处引入 Python 控制流。"导得出 vs 导不出"是判断一个 op 工程完成度的硬指标。

22.24 自定义 op 性能调优 flow

写完一个 custom_op 跑通后，通常发现"比预期慢"。调优流程：

flowchart TD
    Slow[op 慢]
    Slow --> P1[1. profile 看 op 在 trace 里占多少]
    P1 --> Q1{是 op 内部慢, 还是 op 外部 dispatch 慢?}
    Q1 -->|op 内部| Q2{kernel 是否 launch 多次?}
    Q1 -->|dispatch 多| FUSE[让 op 接受更大 input<br/>减少 dispatch 次数]
    Q2 -->|是| Bundle[bundle 多次小 launch 成一次大 launch]
    Q2 -->|否| Q3{Tensor Core 利用率?}
    Q3 -->|低| Align[shape padding 到 16 倍数<br/>+ 用 fp16/bf16]
    Q3 -->|高| MB[memory bound<br/>看能否减少 read/write]
    Slow --> P2[2. 看 with torch.compile 是否能 fuse]
    P2 --> Lower[实现 register_lowering<br/>让 op 进入 fusion]
    Slow --> P3[3. 比对竞品 baseline]
    P3 --> Algo[换更优算法<br/>FA v2 → v3 → ...]
    style P1 fill:#fef3c7
    style P2 fill:#dcfce7
    style P3 fill:#dbeafe

一个典型的两轮排查形态（以自家写的 fused RMSNorm 为例）：

第 1 轮 —— profile 显示这个 op 占的时间远超它的算术量。看 trace 发现 op 内部 launch 了好几个 kernel（mean / rsqrt / scale 各一个），每个都要完整读写一遍输入。这是没真正 fuse：包了 custom_op 不等于 kernel 合一，函数体里调几次 ATen 就是几个 kernel。改成一个 Triton kernel 把几步合到一趟读写里。

第 2 轮 —— kernel 数对了但仍慢于成熟实现。看硬件 metric 发现 occupancy 偏低，说明 BLOCK_SIZE / num_warps 没调对。挂上 @triton.autotune（§22.16）让它自己搜。

关键是 profile 驱动：每一步都先看数据定位根因，再动手改。第 21 章那套方法论在这里原样适用。

22.25 multi-level dispatch：算子的多层 fallback

dispatcher（第 5 章 §5.3）按 keyset 里优先级最高的 key 查表调用算子：效果上是从最具体的 device kernel 退到最通用的 fallback。custom_op 也参与这套机制。下图画的是这个效果；实际上 alias key 在注册期就已展开进 dispatchTable_（§5.2.2），运行期仍是一次查表，没有逐级搜索。

graph TB
    Call[mymul x y]
    Call --> D[dispatcher]
    D --> D1{x is on CUDA?}
    D1 -->|是| K1[找 CUDA impl]
    K1 -->|找到| RunCuda[运行 CUDA kernel]
    K1 -->|没有| K2[找 CompositeExplicitAutograd]
    K2 -->|找到| RunComp[运行默认实现]
    K2 -->|没有| Fail[报错: 没注册]
    D1 -->|是 PrivateUse1| KP[找 PrivateUse1 impl]
    KP -->|找到| RunNpu[运行 NPU kernel]
    KP -->|没有| K2
    style RunCuda fill:#dcfce7
    style RunComp fill:#fef3c7
    style RunNpu fill:#dbeafe

dispatch key 的优先级由 c10/core/DispatchKey.h 里的枚举顺序决定（第 5 章），大致是"功能性 key 在前、backend key 在后"：Autograd 这类先跑并 redispatch 下去，最后落到具体 backend。

对 custom_op 来说需要记住的只有两条：

register_kernel("cuda") 注册的是 CUDA 这个 backend key
装饰器函数体注册在 CompositeExplicitAutograd 上（torch/_library/custom_ops.py:366），它是一个 alias key，会展开成所有 backend —— 所以没写专门 kernel 的 device 会落到它

于是就有了这条 fallback 链：

下面是读者示例，不是 PyTorch 源码原文——mymul 是示例自取的算子名：

@torch.library.custom_op("mylib::mymul", mutates_args=())
def mymul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    return x * y    # 默认 (CompositeExplicitAutograd)

@mymul.register_kernel("cuda")
def _(x, y):
    return my_cuda_kernel(x, y)    # CUDA fast path

# 没注册 cpu impl?
# CPU input 调 mymul → 找不到 CPU impl → 走默认 (composite) → x * y

CPU 用户能跑（虽然慢），CUDA 用户走快路径，不需要在 op 里写任何 if device ==。

理解这套分派让你看 PyTorch 的扩展性：一个 op 可以为 N 个 device 注册 N 份实现，dispatcher 按 key 优先级自动选。这是单一 codebase 支持多种硬件的工程基础。

22.26 SOTA op 接入示例：开源生态中的 5 个典型 case

把全章话题落到具体例子，5 个开源 SOTA op 的接入方式：

1. FlashAttention (Tri Dao)

路径：CUDA kernel → flash_attn Python wrapper → torch.library 注册
全套：fake / autograd / per-device kernel
PyTorch 内置的 F.scaled_dot_product_attention 有自己的 flash 后端实现（在主仓内，不是引第三方包）

2. xformers

路径：CUDA + Triton kernel → 自家 wrapper → 部分注册成 PyTorch op
不全用 custom_op（早于 v2.4 出现），有些走 autograd.Function
v2.x 时代逐步迁到 custom_op

3. Liger Kernel (Linkedin)

路径：纯 Triton kernel（fused RMSNorm / GeGLU / RoPE 等）
全 Python：@triton.jit kernel + torch.autograd.Function 包装（如 Liger 仓库里 RMSNorm 用的 LigerRMSNormFunction），并没有用 custom_op
"纯 Triton + Python、不产生 .so"这条路线的代表

4. bitsandbytes (8-bit / 4-bit ops)

路径：自家 CUDA kernel → C++ extension
部分注册成 PyTorch op，部分仍是普通函数

5. Apex (NVIDIA)

路径：CUDA kernel + setup.py 编译 .so
老一代实践，许多 op 是 autograd.Function
其中不少功能后来被 PyTorch 内置取代

观察：

新项目倾向 Triton + torch.library：省掉整条 C++ 编译与 ABI 链路（§22.14）
老项目逐步迁移：能力被主仓吸收，或者从 autograd.Function 迁到 custom_op
追求极致性能的团队仍写 CUDA：Triton 覆盖不到的指令级优化只能手写

理解这些案例让你看到 PyTorch 自定义 op 生态的全貌：研究项目 → Triton + Python，生产 SOTA → CUDA + 完整 custom_op，硬件厂商 → PrivateUse1 完整接入。每条路径有自己的 trade-off。

22.27 functorch 高阶变换：grad / jacrev / vmap 组合

functorch（v1.13 起随 PyTorch 主仓发布，v2.0 起以 torch.func 形式提供）提供"函数变换"：把可微函数变成它的梯度、Jacobian、Hessian 等。custom_op 与这些 transform 的兼容性要分开看 —— 先看一段 v2.11 实跑的结果：

import torch
from torch.func import grad, jacrev, vmap

@torch.library.custom_op("mylib::squared", mutates_args=())
def squared(x: torch.Tensor) -> torch.Tensor:
    return x ** 2

@squared.register_fake
def _(x):
    return torch.empty_like(x)

def squared_backward(ctx, grad_out):
    x, = ctx.saved_tensors
    return 2 * x * grad_out

def squared_setup(ctx, inputs, output):
    ctx.save_for_backward(inputs[0])

squared.register_autograd(squared_backward, setup_context=squared_setup)

# 普通 autograd 没问题
x = torch.tensor(3.0, requires_grad=True)
g, = torch.autograd.grad(squared(x), x, create_graph=True)    # tensor(6.)
gg, = torch.autograd.grad(g, x)                                # tensor(2.): 反向本身可微, 二阶也行

# torch.func 的 grad / jacrev 不行
grad(squared)(torch.tensor(3.0))
# RuntimeError: In order to use an autograd.Function with functorch transforms
# (vmap, grad, jvp, jacrev, ...), it must override the setup_context staticmethod. ...
jacrev(squared)(torch.tensor([1.0, 2.0, 3.0]))    # 同一个 RuntimeError

# vmap 可以 (没注册 register_vmap 时走慢速 fallback, §22.17)
vmap(squared)(torch.tensor([1.0, 2.0]))            # tensor([1., 4.])

原因在 register_autograd 背后的实现：torch/_library/autograd.py 的 make_autograd_impl 用 type(name, (autograd.Function,), {...}) 现造一个 autograd.Function 子类，只给了 forward / backward 两个 staticmethod，setup_context 是在 forward 内部调用的，类上并没有覆盖 setup_context。而 torch/autograd/function.py 的 apply 在 functorch 变换激活时检查到这一点，就抛上面那个错。所以在 v2.11 上，custom_op + register_autograd 能配合普通 autograd（含高阶导），但不能配合 torch.func.grad / jacrev / jvp；vmap 走的是 register_vmap 那条独立的路。

特殊情况：

需要 torch.func.grad / jacrev 的场景：别包 custom_op，用普通 Python 函数，或者写一个覆盖了 setup_context 的 autograd.Function
高阶导数（torch.autograd.grad(..., create_graph=True) 再求导） 要求反向函数自身可微 —— backward 里调的 op 都得是可微 op，不能返回 detach 过的值
register_autograd 的 backward 与 setup_context 必须 traceable：docstring 明确要求"不得直接访问 Tensor.data_ptr、不得依赖或修改全局状态"。真需要不可 trace 的反向，就把它再包成一个 custom_op 在 backward 里调

实战：生产 LLM 训练几乎不用 jacrev / hessian（模型太大算不动），torch.func 主要出现在二阶优化、物理模拟、per-sample gradient 这类场景。多数自定义 op 写对 register_autograd 就够。

理解 functorch 兼容性让你看 custom_op 的"完整生态接入"含义 —— 不只是 forward + backward，函数变换这一层在 v2.11 上还有缺口。

22.28 ABI-stable C++ 扩展：v2.7+ 实验性新路径

§22.14 提了 ABI 兼容性是 C++ 扩展的痛点。PyTorch 为此引入了一套 ABI-stable 的 C++ 头文件（torch/csrc/stable/library.h 首见于 v2.7.0，tensor.h 于 v2.8.0、ops.h 于 v2.9.0 加入），v2.11 的 torch/csrc/stable/ 目录下有 library.h / ops.h / tensor.h / device.h / accelerator.h 等。注册要换成它自己的一套宏 —— STABLE_TORCH_LIBRARY / STABLE_TORCH_LIBRARY_IMPL，kernel 用 TORCH_BOX 包成 boxed 形式（都在 torch/csrc/stable/library.h）；写法照 v2.11 树里 test/cpp_extensions/libtorch_agn_2_10_extension/csrc/my_shape.cpp 这类测试扩展：

#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/ops.h>
#include <torch/csrc/stable/tensor.h>

using torch::stable::Tensor;

// stable API 不暴露内部数据结构; v2.11 的 ops.h 里没有 mul, 这里用它提供的 matmul
Tensor my_matmul(Tensor x, Tensor y) {
    return torch::stable::matmul(x, y);
}

STABLE_TORCH_LIBRARY(mylib, m) {
    m.def("my_matmul(Tensor x, Tensor y) -> Tensor");
}

STABLE_TORCH_LIBRARY_IMPL(mylib, CompositeExplicitAutograd, m) {
    m.impl("my_matmul", TORCH_BOX(&my_matmul));
}

本专栏在本机实测过这一段：用 libc++ 编出 .so，torch.ops.load_library 载入 libstdc++ 构建的 torch 2.11 wheel 后，torch.ops.mylib.my_matmul 正常出结果 —— 而 §22.14 那种普通 TORCH_LIBRARY 扩展在同样条件下加载即报 undefined symbol。边界上只剩 C 接口，C++ 标准库不一致也就不再要紧。

思路：

不暴露 libtorch 的内部 C++ 类型，只给一层稳定的 torch::stable::Tensor 等封装，靠 stableivalue_conversions.h 在边界上做转换
API 面刻意收窄：只覆盖常用操作，复杂需求仍要回到不稳定的内部 API
目标是让同一份 .so 能跨 PyTorch 版本加载，不必每次升级都重编

这套 API 还在演进中，具体覆盖范围以你所用版本的 torch/csrc/stable/ops.h 为准 —— 直接读那个头文件比读任何教程都可靠。

适用场景：长期维护、面向大量用户分发的开源扩展。自家研究 prototype 用普通 C++ 扩展就行，反正跟着自己的 PyTorch 版本走。

更彻底的规避方式仍然是 §22.3 那条：纯 Python + Triton，根本不产生 .so，也就没有 ABI 问题。

22.29 distributed 训练里的 custom op

custom_op 在分布式训练里要注意：

1. collective 尽量留在 op 外面，用 functional API

import torch.distributed as dist
from torch.distributed.distributed_c10d import _resolve_process_group

# 能跑, 但不推荐: collective 藏在 custom_op 函数体里
@torch.library.custom_op("mylib::ring_attention", mutates_args=())
def ring_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, group_name: str) -> torch.Tensor:
    out = q @ k.transpose(-2, -1)
    dist.all_reduce(out, group=_resolve_process_group(group_name))   # ← 编译器看不见这次通信
    return out @ v

# 推荐: collective 写在 torch.compile 会 trace 的代码里, 用 functional collectives (§16.12)
import torch.distributed._functional_collectives as funcol

def ring_attention_block(q, k, v, group):
    out = q @ k.transpose(-2, -1)
    out = funcol.all_reduce(out, "sum", group)    # ← functional, 编译器看得见、能参与调度
    return out @ v

v2.11 用 gloo 单进程、aot_eager + fullgraph=True 实跑，两种写法都不报错：custom_op 的函数体本来就不被 trace；而 Dynamo trace 到普通代码里的 dist.all_reduce 时，会按 torch/distributed/_functional_collectives.py:1706 的 traceable_collective_remaps 自动换成 functional 版本。差别在可见性：藏进 custom_op 的通信对编译器是黑盒，没法与周围计算一起排布。

2. process_group 不能直接放 schema

ProcessGroup 不是 tensor，不能作为 op 输入。变通：用 group_name (str) 在 op 内部 lookup：

from torch.distributed.distributed_c10d import _resolve_process_group

@torch.library.custom_op("mylib::ring_attention", mutates_args=())
def ring_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, group_name: str) -> torch.Tensor:
    group = _resolve_process_group(group_name)
    ...

torch.distributed._functional_collectives 的 collective 也接受 group name 字符串，正好能配合这种"只能传可序列化标量"的 schema 约束。

3. FSDP-2 / DTensor 协作

DTensor（§18.8）有 placement 概念，DTensor 不知道你的自定义算子该怎么分片。真正的扩展点是 torch.distributed.tensor.experimental.register_sharding（注意 experimental 命名空间）：

from torch.distributed.tensor.experimental import register_sharding

@register_sharding(torch.ops.mylib.ring_attention.default)
def _(q, k, v, group_name):
    # 返回该算子可接受的 (输出 placement, 输入 placement) 策略
    ...

内部还有一层更低的 register_prop_rule（torch/distributed/tensor/_ops/utils.py:65），DTensor 自己的内置算子规则走那条路，但它是私有 API。

实战：custom_op 要在 FSDP-2 / DTensor 模型里用，就得告诉 DTensor 它的分片语义，否则会因为找不到 sharding 策略而报错。

4. 别在 op 里反复建 communicator

如果 custom_op 内部直接调 NCCL（而不是走 PyTorch 的 ProcessGroup），务必复用同一个 communicator —— NCCL communicator 的创建是集合操作，各 rank 必须以完全一致的顺序建立，在算子里按需创建极易造成死锁（第 16 章 §16.8）。优先的做法是别自己建：拿现成的 ProcessGroup 或 functional collectives，把 communicator 的生命周期交给 PyTorch。

理解分布式 custom_op 的这些约束，写"适配多卡"的自定义算子时才不会踩雷。生产 LLM 训练里 custom_op 与 FSDP / TP / PP 协作是真实需求（比如自家 attention 实现要兼容现有训练栈）。

22.30 推理引擎中的 custom_op：vLLM / SGLang 实例

LLM 推理引擎大量使用自定义 op。下面两段是示意性骨架（真实代码请以各项目仓库为准），用来说明推理场景与训练场景的注册差异：

变长 attention（形态类似 vLLM 的 flash attention backend）：

@torch.library.custom_op("myengine::flash_attn_varlen", mutates_args=())
def flash_attn_varlen(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
    cu_seqlens_q: torch.Tensor, cu_seqlens_k: torch.Tensor,
    max_seqlen_q: int, max_seqlen_k: int,
) -> torch.Tensor:
    return _flash_attn_v3.varlen_forward(...)

@flash_attn_varlen.register_fake
def _(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k):
    return torch.empty_like(q)

把"长度变化的 attention"包成 custom_op，torch.compile 才能把它 capture 进图（否则每次都 graph break），进而配合 piecewise CUDA Graph（§15.12.1）拿到低 CPU 开销。

paged attention（形态类似各引擎的 KV cache 读写算子）：

@torch.library.custom_op("myengine::paged_attn", mutates_args=("output",))
def paged_attn(
    output: torch.Tensor,    # mutates 输出 tensor (KV cache 持续累积)
    query: torch.Tensor,
    key_cache: torch.Tensor, value_cache: torch.Tensor,
    block_tables: torch.Tensor,
    seq_lens: torch.Tensor,
) -> None:
    _my_kernel.paged_attention(output, query, key_cache, value_cache,
                                block_tables, seq_lens)

注意 mutates_args=("output",) 让 output 显式标 inplace。这是推理引擎与训练 op 不同的地方：推理时 KV cache 持续累积、必须 inplace 写入，无法走纯函数路径。

实战经验：

推理引擎的 op 不需要 register_autograd：推理无反向，省工作
必须 register_fake：CUDA Graph 与 torch.compile 都需要
mutates_args 要写正确：KV cache mutation 必须显式标
register_kernel("cuda") 调真实 CUDA kernel；CPU fallback 可选

理解这套用法之后再去读推理引擎的源码就不会困惑：那些看起来神秘的 kernel 包装，本质都是本章讲的四件套。《vLLM 推理内核深度解析》专栏第 4 章与第 13 章有对应的具体实现。

22.31 算子注册的"产品哲学"

把全章合起来看，custom_op 接口的设计能读出几条取向：

1. "扩展是用户体验的一部分"

老 PyTorch（v1.x）扩展接口分散：autograd.Function、TORCH_LIBRARY、Library.impl()……每条路径覆盖一部分场景。结果：用户写自定义 op 痛苦、社区贡献 PR 质量参差不齐。

v2.x 收敛到 torch.library 一个命名空间：custom_op 定义算子，其余 register_* 各接一个子系统（§22.21 的表）。把"如何扩展"当成产品的核心 UX 来设计，是 PyTorch 从研究框架成熟为工业级 ML 平台的标志。

2. "fake / shape inference 是底座"

v2.x 把 fake 函数从可选变成"几乎必填"。这看起来增加了用户负担，实际是强制让所有 op 都能进入编译路径。否则 LLM 时代 torch.compile 会被零散的 op 不兼容拖累。

效果上，要求 op 注册时就声明 fake，等于把"能进编译路径"前置到了注册这一步。

3. "Triton 作为一等公民"

v1.x 时代要写高性能自定义 kernel，基本得走 C++ + CUDA + setup.py。v2.x 不但让 Triton kernel 能被包成算子，还专门给了 triton_op / wrap_triton 让编译器看得进去（§22.3）—— 这不是"支持一下"，是把 Triton 当成与 ATen 并列的实现语言。社区里也出现了 Liger Kernel 这类纯 Triton 的高性能算子库。

4. "PrivateUse1 给硬件中立"

不绑死单一厂商。提供完整的 backend extension API，让第三方芯片能在不 fork 主仓的前提下接进来（§22.9、§22.20）。这条决策让 PyTorch 在 NVIDIA 之外的硬件上仍是可选项。

理解这些设计决策让你看自定义 op 接口不只是"技术 API"，而是 PyTorch 对"开放生态"的具体兑现。

22.32 一段实战脚本：从零到生产 op

把全章的步骤合并成一个完整的实战脚本，写一个 fused "GeLU + Linear" op：

import torch
import triton
import triton.language as tl

# 第 1 步: Triton kernel
@triton.autotune(
    configs=[
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 32}, num_warps=4),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 32}, num_warps=4),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64}, num_warps=8),
    ],
    key=['M', 'N', 'K'],
)
@triton.jit
def fused_gelu_linear_kernel(
    x_ptr, w_ptr, b_ptr, out_ptr,
    M, N, K,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    # ... 实现 GeLU(X @ W + b) 的 fused kernel ...
    # 略 (完整实现 30+ 行)

# 第 2 步: custom_op 包装
@torch.library.custom_op("mylib::fused_gelu_linear", mutates_args=())
def fused_gelu_linear(
    x: torch.Tensor, w: torch.Tensor, b: torch.Tensor,
) -> torch.Tensor:
    M, K = x.shape
    K2, N = w.shape
    assert K == K2

    out = torch.empty(M, N, device=x.device, dtype=x.dtype)
    # autotune 选中的 BLOCK_M / BLOCK_N 不固定, grid 必须按选中的 config 算
    grid = lambda meta: (triton.cdiv(M, meta['BLOCK_M']), triton.cdiv(N, meta['BLOCK_N']))
    fused_gelu_linear_kernel[grid](x, w, b, out, M, N, K)
    return out

# 第 3 步: register_fake
@fused_gelu_linear.register_fake
def _(x, w, b):
    M, K = x.shape
    K2, N = w.shape
    return torch.empty(M, N, device=x.device, dtype=x.dtype)

# 第 4 步: register_autograd
def fgl_backward(ctx, grad_out):
    x, w, b = ctx.saved_tensors
    # ... 实现反向 ...
    grad_x = (grad_out * gelu_grad(x @ w + b)) @ w.T
    grad_w = x.T @ (grad_out * gelu_grad(x @ w + b))
    grad_b = grad_out.sum(0)
    return grad_x, grad_w, grad_b

def fgl_setup(ctx, inputs, output):
    ctx.save_for_backward(*inputs)

fused_gelu_linear.register_autograd(fgl_backward, setup_context=fgl_setup)

# 第 5 步: opcheck 验证
from torch.library import opcheck
x = torch.randn(64, 256, device='cuda', requires_grad=True)
w = torch.randn(256, 128, device='cuda', requires_grad=True)
b = torch.randn(128, device='cuda', requires_grad=True)
opcheck(fused_gelu_linear, args=(x, w, b))    # 通过

# 第 6 步: 集成到模型
class FusedFFN(torch.nn.Module):
    def __init__(self, dim, hidden):
        super().__init__()
        self.w = torch.nn.Parameter(torch.randn(dim, hidden))
        self.b = torch.nn.Parameter(torch.randn(hidden))

    def forward(self, x):
        return torch.ops.mylib.fused_gelu_linear(x, self.w, self.b)

# 第 7 步: torch.compile 验证
model = FusedFFN(256, 1024).cuda()
compiled = torch.compile(model)
out = compiled(x)
loss = out.sum()
loss.backward()    # 反向自动调 fgl_backward, fused 进 inductor graph

整套约 80 行 Python（不含 Triton kernel 实现）。工作的顺序比工期估算更有价值：

写 Triton kernel，先在 eager 里跑通 forward 并与参考实现对数值
register_fake + register_autograd，跑 opcheck 把 schema / autograd / fake / compile 四条路一次性验通
需要 vmap 或 Inductor fusion 时再补 register_vmap / 换 triton_op
挂 @triton.autotune，用 profiler（第 21 章）确认瓶颈真的消失
集成进模型，与 baseline 对齐精度
补单测，把 opcheck 放进 CI

关键是第 2 步不要往后拖 —— 先跑通 opcheck 再调性能，否则很容易在一个 schema 就写错的算子上白调半天。

这条流程也说明了为什么本章要放在专栏靠后：写一个生产级 custom_op，等于把前面所有章节的知识各用一遍。

22.33 自定义 op 的版本兼容性策略

随着 PyTorch / 自家库迭代，自定义 op 的 schema 可能变化。生产代码必须考虑兼容性：

1. schema 演进的安全规则：

修改	是否 break 兼容
新增 op	不 break（旧代码不调用就行）
新增 op 的可选参数（带默认值）	不 break
新增 op 的必选参数	break（旧代码不传新参数报错）
重命名 op	break
改 input dtype	break（schema 校验失败）
改 output shape 推导	隐性 break（compile 后行为变）

实战做法：

新功能加可选参数：def my_op(x, y, *, optional_flag: bool = False) -> Tensor
deprecated 老 op，加新 op：保留 mylib::v1_op，新 ckpt 用 mylib::v2_op
schema 重大变化：bumping namespace（mylib::op → mylib_v2::op）

2. 与 PyTorch 版本的兼容：

import torch
if torch.__version__ >= "2.4":
    @torch.library.custom_op("mylib::myop", mutates_args=())
    def myop(x: torch.Tensor) -> torch.Tensor:
        ...
else:
    # v2.4 之前的 fallback 写法
    class MyOp(torch.autograd.Function):
        ...

或用 try/except 兜底：

try:
    from torch.library import custom_op
except ImportError:
    # 老 PyTorch 没有这个 API
    custom_op = None

3. ckpt 兼容性：

如果 op 是 model 的一部分，model state_dict 没区别（op 的实现不在 state_dict 里）。但用户代码必须能 import 到 op——升级时确保自家 op 库一并升级。

4. 渐进式 deprecation：

import warnings

@torch.library.custom_op("mylib::old_op", mutates_args=())
def old_op(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    warnings.warn(
        "mylib::old_op is deprecated, use mylib::new_op instead",
        DeprecationWarning, stacklevel=2,
    )
    return new_op(x, y)

让用户有时间迁移，几个月后正式删除。

理解这些策略让你写自定义 op 时考虑"长期维护"，不是只考虑 v1。生产 op 一旦上线就要支持多年（用户的 ckpt 还在用），向前 / 向后兼容性是必修课。

22.34 op 注册的内部数据结构：从 schema 到 dispatcher

把全章话题落到底层数据结构。@torch.library.custom_op("mylib::mymul", ...) 在 PyTorch 内部最终落到几张表：

graph TB
    Decorator[custom_op 装饰器]
    Decorator --> Lib[Library 对象 mylib<br/>Python 层 wrapper]
    Lib --> CppLib[C++ Library<br/>持有 schema list]
    CppLib --> Dispatcher[Dispatcher 全局表<br/>OperatorHandle]
    Dispatcher --> Schema["schema string<br/>mymul(Tensor, Tensor) -> Tensor"]
    Dispatcher --> Kernels{各 dispatch key 实现表}
    Kernels --> CPU[CPU: lambda x,y: x*y]
    Kernels --> CUDA[CUDA: triton_kernel]
    Kernels --> Auto[AutogradCUDA: 自动包 backward]
    Kernels --> Fake[Meta/FakeTensor: register_fake fn]
    Decorator --> AutogradReg[autograd info<br/>setup_context + backward fn]
    AutogradReg --> Auto
    style Dispatcher fill:#fef3c7
    style Kernels fill:#dcfce7

具体源码位置（v2.11）：

CustomOpDef：torch/_library/custom_ops.py（torch/library.py 只是把它 re-export）
Python Library：torch/library.py:68
C++ Library：torch/library.h（TORCH_LIBRARY / TORCH_LIBRARY_IMPL 宏都在这里）
Dispatcher 与 OperatorHandle：都在 aten/src/ATen/core/dispatch/Dispatcher.h（没有单独的 OperatorHandle.h）

调用 mymul(x, y) 的内部路径：

Python 调 torch.ops.mylib.mymul(x, y)
C++ OperatorHandle.callBoxed(stack)
Dispatcher 查 dispatch key set（input device + autograd state + ...）
先命中 Autograd key 上 register_autograd 装出来的 wrapper（torch/_library/custom_ops.py:641 的 autograd_impl 注册的那个），它负责建反向 Node 并 redispatch
redispatch 后落到具体 backend key：有 register_kernel("cuda") 就用它，没有就落到装饰器函数体所在的 CompositeExplicitAutograd
该 kernel 里调你写的 Triton / CUDA kernel
反向时 autograd Engine（第 8 章）调度那个 Node，最终执行 register_autograd 传入的 backward 函数

每一步都用第 5 章讲过的同一套机制 —— 自定义 op 与内置 op 走完全相同的 dispatch 路径。这就是"扩展与内置无差别"（§22.37 第一条）的具体含义：不是"看起来像"，是字面意义上的同一条代码路径。

22.35 几条工程经验

1. 优先用 torch.library.custom_op：替代老的 autograd.Function 与手写 TORCH_LIBRARY 宏；实现是 Triton kernel 且想被 Inductor 优化时改用 triton_op（§22.3）

2. torch.library.opcheck(my_op, args) 是合规性测试：自动检查 fake / autograd / schema 等是否一致。生产 op 必跑

3. 用 Python 写新算子，Triton kernel + triton_op 是好组合：性能、灵活性都不错，Inductor 还看得进去（§22.3）；只要一个黑盒时再用 custom_op 包

4. mutates_args= 一定写正确：错了 functionalize 会出问题、torch.compile 编译错代码

5. 不要在 fake 函数里做实际计算：会让 torch.compile / FSDP 内存爆 / 性能崩

6. C++ 扩展跨 PyTorch 版本要重编：libtorch 内部 ABI 不保证版本兼容。要长期分发就看 torch/csrc/stable/（§22.28），或者干脆不写 C++

7. PrivateUse1 是国产芯片接入路径：注册成新 backend 而非新算子，让所有现有算子都能跑

8. torch.compiler.allow_in_graph 给某些函数特殊白名单：Dynamo 看不进、但 AOTAutograd / Inductor 能 trace 的函数用它绕过 Dynamo（它只对 Dynamo 是黑盒，要对整条编译栈当黑盒得用 custom_op）；纯 Python 逻辑（logging / metric）用 torch.compiler.disable 直接 break 掉（§22.10）

9. 推理引擎用的 op 不需要 register_autograd：推理无反向，省一步工作。但 register_fake 仍必须

10. 别指望 custom_op 会被 compile 展开：它按设计就是黑盒。要 fusion 就上 triton_op，要极致就写 lowering（私有 API）

11. 想让编译器看见 collective，就把它留在被 compile 的代码里、用 functional API：torch.distributed._functional_collectives；compile 区域里直接写的 dist.all_reduce 会被 Dynamo 自动换成 functional 版本，但藏进 custom_op 函数体就成了编译器看不见的黑盒（§22.29）

12. 写 Triton kernel 通常配上 @triton.autotune：让 BLOCK_SIZE / num_warps 自动搜索，避免手调

22.36 跨专栏关联

第 5 章 dispatcher：自定义 op 注册的底层机制
第 6 章 ATen 代码生成：内置 op 是 codegen，自定义 op 是 register —— 两条路殊途同归
第 7 章 autograd：register_autograd 与 autograd.Function.backward 等价语义
第 12-14 章编译器栈：fake 函数让自定义 op 进入编译路径，register_lowering 让 op 真正被 Inductor fuse 而非走 fallback
第 16 章 ProcessGroup：分布式训练里 custom_op 与 functional collectives 的协作
第 18 章 FSDP-2 / DTensor：DTensor placement 与 custom_op 的 dispatch 协作
第 21 章 Profiler：opcheck 与 profile 共同保证 op 正确性 + 性能符合预期

22.37 设计启示

PyTorch 自定义算子接口的核心思想：

第一：让"扩展"与"内置"无差别：自定义 op 一旦注册就和 torch.add 走同一条 dispatch 路径（§22.34）。上层特性不需要为它开特例

第二：fake 函数是高级特性的入场券：编译、export、meta 全都建在"不真算也能推出 shape"之上。要求每个 op 提供 fake 看似增加负担，实际是把"所有 op 都要能进编译路径"这个判断前置到注册时

第三：多种 device 各注册一份 kernel：PrivateUse1 给国产芯片厂商完整的扩展能力，不需要 fork PyTorch 主仓，让硬件中立性成为生态扩展的基础设施

第四：用装饰器替代宏 / Python 替代 C++：现代 API 让"写自定义 op"从需要 C++ + 宏的工程任务，降级到 10 行 Python 装饰器。这种"降低门槛 + 保留性能"的设计让 PyTorch 自定义 op 的生态得以壮大

第五：一个子系统一个注册点：register_fake / register_autograd / register_autocast / register_vmap / register_torch_dispatch 各接一个维度（§22.21 的表）。算子的"完整度"不是布尔值而是覆盖率 —— 想用上哪个特性，就补哪个 hook。这种正交设计让 PyTorch 能不断加新子系统而不推倒重来

第六：opcheck 把"扩展正确性"自动化：schema 的 mutation 声明、autograd 的注册方式、fake 与真实的 metadata 一致性、eager 与 compile 的输出一致性 —— 四项自动查（§22.7）。有这层质量基础设施，社区贡献的算子才可能保持水准

22.38 跨章呼应：自定义 op 是这个专栏的"集大成"

把全章合起来看，自定义 op 几乎需要全专栏前面所有章节的知识：

写自定义 op 时用到	对应章节
schema / IValue / ATen	§6 ATen 代码生成
dispatcher 注册	§5 dispatcher
TensorImpl / Storage	§2 Tensor 数据结构
autograd Function / Engine	§7-8 autograd
AOTAutograd functionalize	§13 AOT Autograd
FakeTensor / register_fake	§5.7 + §13
Inductor lowering / fusion	§14 Inductor
torch.compile 协作	§12-15 编译栈
AMP custom_fwd	§20.20
FSDP / DTensor / collective	§16-18 分布式
profile + opcheck	§21 Profiler

写一个生产级 custom_op = 整个专栏的综合实践。这就是为什么把它放在最后一章（除 23 章哲学收束外）—— 它是检验前面知识掌握程度的"期末考试"。

新人写自定义 op 卡在哪一步，对应回去复习对应章节。这是本专栏的内部 cross-reference 网络的最后一环。

下一章是收官章 —— 拆 PyTorch 整体设计哲学与未来演进，把 22 章的内容串成一条主线，看从 Tensor 到 custom_op 这条 trace 上 PyTorch 团队留下了什么共通的设计原则。