Operator registration and dispatch: how PyTorch selects a kernel算子注册与调度：PyTorch 如何选择内核

One call, several implementations一次调用，多种实现

The previous installment explained how a tensor locates its elements. Now consider torch.mm(a, b): the operands carry a device and a dispatch key set, yet the Python call does not name a CPU or CUDA function. The missing connection is an operator-specific dispatch table. Registration fills its entries; each invocation selects an entry using its arguments and execution context.上一期解释了张量怎样定位元素。现在看 torch.mm(a, b)：输入携带设备信息和调度键集合，但 Python 调用没有指定 CPU 或 CUDA 函数。把两者接起来的是每个算子各自的调度表。注册负责填充表项，每次调用再根据参数与执行上下文选择入口。

import torch

a = torch.tensor([[1., 2.]], requires_grad=True)
b = torch.tensor([[3.], [4.]])
y = torch.mm(a, b)
print(y.item(), type(y.grad_fn).__name__)
# 11.0 MmBackward0

Here the numerical computation returns 11, while a separate wrapper records how that value depends on a. A dispatcher entry called a “kernel” can be such a wrapper; it does not necessarily multiply any numbers. Backward computation happens later, when requested.这里数值计算得到 11，另一层包装代码记录结果对 a 的依赖。调度表中称为“kernel”的入口也可能是这种包装层，它未必执行乘加运算。真正的反向计算发生在随后请求 backward 时。

Source baseline: PyTorch 2.10.0, commit 449b17684101, as in parts 1 and 2. The examples use ordinary dense, real-valued tensors in eager mode, without autocast, tensor subclasses, or transforms. CPU outputs below were verified with 2.10.0+cpu; the CUDA route is a source-based illustration, not a GPU measurement.源码基线：PyTorch 2.10.0，提交 449b17684101，与前两期一致。示例限定为 eager 模式下的普通稠密实数张量，不启用 autocast、张量子类或变换。下文 CPU 输出已在 2.10.0+cpu 中验证；CUDA 路径依据源码绘制，不代表 GPU 实测。

Schema, registration, implementation接口定义、注册与实现

In ATen, native_functions.yaml supplies schemas and backend declarations to code generation. For dense inputs, functional mm delegates to the structured mm.out implementation. Generated code manages the functional result and connects it to backend code; the Python caller need not supply out=. The following excerpt keeps only the relevant declarations.在 ATen 中，native_functions.yaml 为代码生成提供接口和后端声明。对于稠密输入，函数式 mm 委托给结构化的 mm.out 实现。生成代码负责函数式结果的准备，并连接后端代码；Python 调用者无需传入 out=。下面摘取与此相关的声明。

- func: mm(Tensor self, Tensor mat2) -> Tensor
  structured_delegate: mm.out

# mm.out backend declarations:
CPU: mm_out_cpu
CUDA: mm_out_cuda

native_functions.yaml: mm and mm.out declarationsnative_functions.yaml：mm 与 mm.out 声明

An extension can register a schema with TORCH_LIBRARY and a CPU implementation with TORCH_LIBRARY_IMPL. This C++ fragment requires a PyTorch extension build and an existing twice_cpu function; it is not a standalone program. The Python version in section 6 runs without compiling an extension.扩展可以通过 TORCH_LIBRARY 注册接口，再通过 TORCH_LIBRARY_IMPL 注册 CPU 实现。下面的 C++ 片段需要 PyTorch 扩展构建环境，以及已有的 twice_cpu 函数；它不是独立程序。第 6 节提供无需编译扩展即可运行的 Python 版本。

TORCH_LIBRARY(feng_dispatch, m) {
    m.def("twice(Tensor x) -> Tensor");
}
TORCH_LIBRARY_IMPL(feng_dispatch, CPU, m) {
    m.impl("twice", twice_cpu);
}

torch/library.h: operator definition and backend registration APIstorch/library.h：算子定义与后端注册接口

A schema is not a device-transfer rule. A CPU registration does not make CUDA inputs run on the CPU automatically. It also does not define a derivative: an implementation and an autograd rule have different responsibilities.接口定义不是设备迁移规则。注册 CPU 实现不会让 CUDA 输入自动转到 CPU 上运行，也不会因此得到求导规则：数值实现与自动求导分别承担不同职责。

The key set is more than a device调度键集合不只有设备

For tensor arguments, the extractor unions their key sets. It then includes and excludes keys from thread-local state (TLS), and applies the operator’s mask. The mask removes fallthrough entries so they need not become ordinary kernel calls. For this tensor-only example, the selection can be written as:对于张量参数，提取器先合并它们的键集合，再应用线程局部状态（TLS）中的包含集合与排除集合，最后使用算子自己的掩码。掩码会过滤 fallthrough 表项，避免把这些直接放行的入口当作普通内核来调用。对于这里仅含张量参数的示例，选择过程可写为：

I and E are the TLS included and excluded sets; M is the operator mask. This is a set model of the selection, not the bit layout of DispatchKeySet. The representation separates backend and functionality components; readable names such as AutogradCPU should not be mistaken for an independent Boolean switch for every feature.I 与 E 分别为 TLS 的包含集合与排除集合，M 为算子掩码。这是选择过程的集合模型，不是 DispatchKeySet 的实际位布局。内部表示区分后端与功能分量，AutogradCPU 这样的可读名称并不意味着每项功能都对应一个彼此独立的布尔开关。

return (((ks | local.included_) - local.excluded_) & key_mask);

DispatchKeyExtractor.h: computeDispatchKeySet and tensor key collectionDispatchKeyExtractor.h：computeDispatchKeySet 与张量键提取

DispatchKeySet.h: backend and functionality representationDispatchKeySet.h：后端与功能的集合表示

The highest-priority eligible key chooses a table slot for this operator. Registration aliases such as Autograd and CompositeImplicitAutograd help populate multiple runtime slots; they are not extra stages every invocation must visit. A composite implementation expresses an operator using other operators, whose calls can dispatch again.优先级最高的有效键决定访问这个算子的哪一个表项。Autograd、CompositeImplicitAutograd 等注册别名用于填充多个运行时表项，并非每次调用都必须经过的额外阶段。组合实现用其他算子表达当前运算，其中的算子调用可以再次触发调度。

OperatorEntry.cpp: computeDispatchTableEntryWithDebugOperatorEntry.cpp：computeDispatchTableEntryWithDebug

Follow the wrapper down to the backend沿包装层进入后端

The diagram shows the complete route for the example above. The numbered arrows describe control flow, not elapsed time. Playback highlights one stage at a time; the result and all stages remain readable without playing. The comparison controls change one execution condition while keeping the matrix values fixed.下图默认显示上述示例的完整路径。编号和箭头表示控制流，不表示耗时。播放时逐步高亮各阶段；不播放也能直接阅读完整流程和结果。比较控件保持矩阵数值不变，只改变执行条件。

The wrapper must avoid selecting itself recursively. Native autograd code generation uses a reduced key set for redispatch, and guards the appropriate lower-level call. Dispatcher::redispatch accepts the supplied key set as-is; it does not recompute the initial argument/TLS selection. This is how a layer can finish its own work and continue to another implementation of the same operator.包装层必须避免递归选中自己。原生 Autograd 的代码生成会为 redispatch 缩减键集合，并为相应的下层调用设置 guard。Dispatcher::redispatch 直接使用传入的键集合，不重新进行最初的参数与 TLS 选择。这样，一层代码处理完自己的职责后，就能进入同一算子的另一层实现。

Dispatcher.h: Dispatcher::redispatchDispatcher.h：Dispatcher::redispatch

gen_variable_type.py: after_autograd_keyset in native wrapper generationgen_variable_type.py：原生包装代码生成中的 after_autograd_keyset

An Autograd key does not guarantee a backward graph有 Autograd 键，不等于记录了反向图

Two independent questions matter: which implementation is selected? and should it record a gradient history? For the ordinary tensors here, requires_grad=False leaves AutogradCPU present. no_grad() disables gradient recording, while the wrapper can still be selected. Its gradient check returns false and computation proceeds without an attached backward node.这里有两个独立问题：选中了哪个实现？以及是否记录梯度历史？对于本文的普通张量，requires_grad=False 不会移除 AutogradCPU。no_grad() 关闭梯度记录，但仍可能选中包装层；包装层中的梯度检查返回 false，计算继续进行，只是不关联反向节点。

functions/utils.h: compute_requires_grad checks GradModefunctions/utils.h：compute_requires_grad 检查 GradMode

inference_mode() additionally changes dispatch state: Autograd keys are excluded, and newly created inference tensors lack the usual autograd/view-tracking keys. The following experiment creates fresh tensors inside each context. The private torch._C._dispatch_keys helper is a version-specific diagnostic, not a stable public API.inference_mode() 还会改变调度状态：排除 Autograd 键，并且新建的 inference 张量不带通常的求导与视图追踪键。下面的实验在每个上下文内部创建新张量。私有辅助函数 torch._C._dispatch_keys 仅用于当前版本的诊断，不是稳定的公开 API。

InferenceMode.h: TLS invariants and autograd exclusionInferenceMode.h：TLS 不变量与 Autograd 排除规则

import torch
from contextlib import nullcontext
print('version:', torch.__version__)
for name, context, requires_grad in [
    ('grad', nullcontext(), True),
    ('plain', nullcontext(), False),
    ('no_grad', torch.no_grad(), True),
    ('inference', torch.inference_mode(), True),
]:
    with context:
        a = torch.tensor([[1., 2.]], requires_grad=requires_grad)
        b = torch.tensor([[3.], [4.]])
        y = torch.mm(a, b)
        print(name, 'AutogradCPU' in str(torch._C._dispatch_keys(a)),
              torch.is_grad_enabled(), y.item(),
              type(y.grad_fn).__name__)

version: 2.10.0+cpu
grad True True 11.0 MmBackward0
plain True True 11.0 NoneType
no_grad True False 11.0 NoneType
inference False False 11.0 NoneType

The columns after the mode are: AutogradCPU present on the input, gradient mode enabled, numerical result, and backward-node type. All four results equal 11; only the first records MmBackward0. This experiment does not measure speed, and disabling recording does not change the matrix-product formula.模式名之后依次为：输入是否携带 AutogradCPU、是否启用梯度模式、数值结果、反向节点类型。四种模式都得到 11，只有第一种记录 MmBackward0。本实验没有测量速度，关闭记录也不会改变矩阵乘法公式。

Register an operator and its derivative注册一个算子及其导数

A small CPU operator makes the separation observable. In a fresh Python process, define feng_dispatch::twice, register a CPU function, then register its reverse-mode derivative. The CPU function returns a new tensor, consistent with a schema that declares no mutation or aliasing. Keep the Library object alive for the registrations to remain available.一个小型 CPU 算子可以实际展示这种分工。在新的 Python 进程中定义 feng_dispatch::twice，注册 CPU 函数，再注册反向求导规则。CPU 函数返回新张量，与未声明修改或别名的接口一致。需要保留 Library 对象，使注册项持续有效。

import torch

lib = torch.library.Library('feng_dispatch', 'DEF')
lib.define('twice(Tensor x) -> Tensor')
@torch.library.impl(lib, 'twice', 'CPU')
def twice_cpu(x):
    return x * 2

def backward(ctx, grad_output):
    return grad_output * 2

torch.library.register_autograd('feng_dispatch::twice', backward)
@torch.library.register_fake('feng_dispatch::twice')
def twice_fake(x):
    return torch.empty_like(x)

x = torch.tensor([1., 3.], requires_grad=True)
y = torch.ops.feng_dispatch.twice(x)
y.sum().backward()
print('value:', y.tolist())
print('gradient:', x.grad.tolist())
print('gradcheck:', torch.autograd.gradcheck(torch.ops.feng_dispatch.twice,
      (torch.tensor([1., 3.], dtype=torch.double, requires_grad=True),)))
print(torch.library.opcheck(torch.ops.feng_dispatch.twice, (x,)))

value: [2.0, 6.0]
gradient: [2.0, 2.0]
gradcheck: True
{'test_schema': 'SUCCESS', 'test_autograd_registration': 'SUCCESS',
 'test_faketensor': 'SUCCESS', 'test_aot_dispatch_dynamic': 'SUCCESS'}

For y = 2x, an incoming gradient g becomes 2g. No input values need saving for this derivative. The fake implementation returns only output metadata without reading tensor values. opcheck checks registration contracts; gradcheck independently compares derivatives with finite differences. Passing these sample checks is not a proof for all shapes and dtypes.对于 y = 2x，传入梯度 g 变为 2g，不需要为这个导数保存输入数值。fake 实现只提供输出元信息，不读取张量数值。opcheck 检查注册约定，gradcheck 另用有限差分比较导数。通过这些样例检查不等于证明所有形状和数据类型都正确。

torch.library.register_autograd · torch.library.opchecktorch.library.register_autograd：注册反向规则 · torch.library.opcheck：检查注册约定

In torch/_library/autograd.py, the generated custom-op wrapper checks gradient mode and the input requirements, then calls a forward path that redispatches below Autograd. Without a registered derivative, a custom operator should not be assumed differentiable merely because its CPU implementation calls differentiable PyTorch operations. For this example, deleting the autograd registration removes the explicitly supplied rule.在 torch/_library/autograd.py 中，生成的自定义算子包装层检查梯度模式与输入需求，再通过前向路径绕过 Autograd 重新调度。不能仅因为 CPU 实现调用了可求导的 PyTorch 运算，就认定未注册导数的自定义算子已经具有正确的求导支持。本例删除 Autograd 注册后，就失去了显式提供的求导规则。

torch/_library/autograd.py: make_autograd_impl and redispatchtorch/_library/autograd.py：make_autograd_impl 与 redispatch

Only CPU numerical execution is registered here. A fake implementation does not supply a CUDA kernel, and this example is not a performance optimization over x * 2. It isolates registration, execution, and differentiation so that each can be checked separately.这里仅注册了 CPU 数值实现。fake 实现不能替代 CUDA 内核，这个例子也不是对 x * 2 的性能优化。它把注册、执行与求导分开，以便分别检查。

A small C++ dispatcher you can inspect用小型 C++ 调度器看清控制流

This standalone teaching model has one operator, a priority order, and a function-pointer table. Its Autograd wrapper removes its own key before redispatch. The second call uses the same key set with recording disabled; the last call deliberately lacks a CUDA implementation. It models control flow only: no tensors, automatic differentiation engine, or GPU code are implemented. The output below was verified locally and with Godbolt GCC 14.2.这个独立教学模型只有一个算子、一组优先级和一张函数指针表。Autograd 包装层先移除自己的键，再重新调度。第二次调用保持相同键集合但关闭记录；最后一次调用故意缺少 CUDA 实现。它只模拟控制流，没有实现张量、自动求导引擎或 GPU 代码。下方输出已在本地与 Godbolt GCC 14.2 上核验。

#include <array>
#include <iostream>
#include <stdexcept>

using Keys = unsigned;
constexpr Keys CPU = 1, CUDA = 2, Autograd = 4;
struct Context { bool grad_enabled; bool requires_grad; };
using Kernel = double (*)(Keys, Context, double);
double dispatch(Keys keys, Context context, double x);

double cpu(Keys, Context, double x) {
    std::cout << "CPU: twice\n";
    return 2 * x;
}
double autograd(Keys keys, Context context, double x) {
    const bool record = context.grad_enabled && context.requires_grad;
    std::cout << "Autograd: record=" << record << '\n';
    const double result = dispatch(keys & ~Autograd, context, x);
    if (record) std::cout << "Attach backward rule\n";
    return result;
}
struct Entry { Keys key; Kernel kernel; };
const std::array<Entry, 3> table{{
    {Autograd, autograd}, {CUDA, nullptr}, {CPU, cpu}
}};

double dispatch(Keys keys, Context context, double x) {
    for (const auto& entry : table) {
        if (!(keys & entry.key)) continue;
        if (!entry.kernel) throw std::runtime_error("Missing CUDA kernel");
        return entry.kernel(keys, context, x);
    }
    throw std::runtime_error("No eligible dispatch key");
}
int main() {
    for (const bool grad : {true, false}) {
        const double result = dispatch(CPU | Autograd, {grad, true}, 3.0);
        std::cout << "value=" << result << '\n';
    }
    try {
        dispatch(CUDA, {false, false}, 3.0);
    } catch (const std::runtime_error& error) {
        std::cout << error.what() << '\n';
    }
}

Autograd: record=1
CPU: twice
Attach backward rule
value=6
Autograd: record=0
CPU: twice
value=6
Missing CUDA kernel

The model throws at a missing selected entry. Falling through to an unrelated CPU implementation would silently change the execution contract. Real PyTorch also supports registered backend fallbacks and composite implementations; its table construction resolves those possibilities. That is different from retrying arbitrary lower-priority kernels after an error.模型在选中的入口缺失时抛错。如果直接改用无关的 CPU 实现，就悄悄改变了执行约定。真实 PyTorch 还支持已注册的后端 fallback 和组合实现，这些可能性由调度表构建规则处理；它不同于报错后随意尝试其他低优先级内核。

Inspect the table before blaming the kernel先看调度表，再定位内核问题

The following private diagnostics expose an operator’s schema and its computed dispatch table in the pinned build. Registration filenames and available entries depend on build options. A printed CUDA slot alone does not prove that a CPU-only installation can execute CUDA tensors.以下私有诊断接口可查看当前固定版本构建中的接口定义与计算后的调度表。注册文件名和可用表项取决于构建选项。即使打印出 CUDA 表项，也不能据此认定纯 CPU 安装能够执行 CUDA 张量。

print(torch.ops.aten.mm.default._schema)
print(torch._C._dispatch_dump_table("aten::mm"))
print(torch._C._dispatch_dump_table("feng_dispatch::twice"))

OperatorEntry.cpp: dumpComputedTable and registration provenanceOperatorEntry.cpp：dumpComputedTable 与注册来源

Dispatch does not search for the fastest algorithm by benchmarking candidates on each call. It selects a registered implementation from keys and context; that implementation may perform further dtype, layout, shape, or library-level algorithm selection. Tensor strides from part 2 therefore still matter after dispatch has chosen a backend.调度器不会在每次调用时测量候选算法，再选择最快者。它根据键与上下文选择已注册实现；具体实现内部还可能按数据类型、布局、形状或数值库规则进一步选择算法。因此，即便已经选定后端，第二期介绍的张量步幅仍会影响后续执行。

LinearAlgebra.cpp: addmm_impl_cpu_ layout handlingLinearAlgebra.cpp：addmm_impl_cpu_ 的布局处理

Check your understanding: does moving a tensor to CUDA create a backward rule?检验理解：把张量移到 CUDA，就会自动得到反向规则吗？

No. Device placement affects backend selection. The operator still needs a supported numerical implementation and appropriate autograd support. For the CPU-only custom operator above, neither its schema nor its fake implementation supplies CUDA arithmetic. Separately, a recorded backward node will later be scheduled by the autograd engine—the subject of part 4.不会。设备位置影响后端选择，算子仍需要受支持的数值实现和相应的求导支持。对于上面只注册 CPU 的自定义算子，接口定义和 fake 实现都不能提供 CUDA 数值计算。另一个独立问题是：已经记录的反向节点随后怎样被 Autograd 引擎调度，这将是第四期的主题。