第5章 Dispatcher：多分派的工程实现

"This folder contains the c10 dispatcher. This dispatcher is a single point through which we are planning to route all kernel calls."

—— aten/src/ATen/core/dispatch/README.md 第一句

5.1 Dispatcher 是什么

第 1 章 §1.4 我们提过，每一次张量算子（如 a + b）都要经过一个全局的 dispatcher 单例。这一章我们彻底拆开它的内部 —— 看看 PyTorch 怎么用 aten/src/ATen/core/dispatch/ 下 三千来行核心代码 撑起整个深度学习生态的多态性。

让我们先复习问题：当用户写 c = a + b，PyTorch 至少要回答这些问题：

1. a、b 是 CPU 还是 CUDA？走哪个 backend kernel？
2. dtype 是 fp32 / bf16 / int8 / 量化？走哪个数值实现？
3. 是不是 sparse / nested / 量化张量？要不要走专门的稀疏算子？
4. 是不是要记反向图（autograd）？
5. 是不是在 inference_mode 里？autograd 路径要不要绕开？
6. 是不是在 vmap / grad / jvp 里？要不要做函数变换？
7. 是不是在 autocast 上下文里？要不要把 fp32 转 bf16？
8. 是不是在 torch.compile 的 trace 阶段？要不要用 FakeTensor 假执行？
9. 是不是有用户注册的 __torch_dispatch__ mode？
10. 是不是有用户注册的 __torch_function__ 子类？

其中第 2 条（dtype）不归 dispatcher 管 —— 它在选中的 kernel 内部由 AT_DISPATCH_* 宏（aten/src/ATen/Dispatch.h）按 ScalarType 再分一次；第 10 条（__torch_function__）在进入 C++ dispatcher 之前的 Python 绑定层就处理了（§5.7.2）。其余每个问题都对应一个或几个 DispatchKey（或改变线程局部的 include / exclude 集合）。Dispatcher 的工作就是把这些 keys 综合起来，挑出最高优先级的那个，调对应的 kernel。

graph TB
    Call["a + b"] --> Extract["DispatchKeyExtractor<br/>从张量元数据 + TLS 算 DispatchKeySet"]
    Extract --> Lookup["OperatorEntry.lookup(ks)<br/>dispatchTable_[ks.idx]"]
    Lookup --> Kernel["KernelFunction"]
    Kernel --> Invoke["调 unboxed 实现"]
    Invoke --> Result["返回张量"]
    Kernel -.可选.-> Boxed["boxed 实现 (Python / vmap / fallback)"]
    style Extract fill:#dbeafe,stroke:#3b82f6
    style Lookup fill:#fef3c7,stroke:#f59e0b,stroke-width:2px
    style Kernel fill:#dcfce7,stroke:#22c55e

整条路径中最关键的优化是：lookup 不是 hashmap、不是 vtable、不是 if-else，而是一个数组下标。这是 PyTorch 能把上百个 key 的多分派做到"和一次普通函数调用同量级"的根本原因。

5.2 数据结构：Dispatcher / OperatorEntry / KernelFunction

打开 aten/src/ATen/core/dispatch/Dispatcher.h:71：

class TORCH_API Dispatcher final {
 public:
    static Dispatcher& singleton();

    template <class Return, class... Args>
    Return call(const TypedOperatorHandle<Return(Args...)>& op,
                Args... args) const;

    void callBoxed(const OperatorHandle& op, Stack* stack) const;
    Return redispatch(...) const;

    RegistrationHandleRAII registerImpl(
        OperatorName op_name,
        std::optional<DispatchKey> dispatch_key,
        KernelFunction kernel, ...);
    ...

 private:
    // 算子本体的所有者：每个算子一个 OperatorDef（内含 OperatorEntry）
    std::list<OperatorDef> operators_;
    // 名字 → handle 的查找表
    LeftRight<ska::flat_hash_map<OperatorName, OperatorHandle>> operatorLookupTable_;
    // 每个运行时 key 一个"兜底"kernel
    std::array<impl::AnnotatedKernel, num_runtime_entries> backendFallbackKernels_;
    ...
};

——Dispatcher 单例持有：

• operators_：算子本体，一个 std::list<OperatorDef>（用 list 是因为 OperatorHandle 里存的就是指向元素的迭代器，必须保证插入不失效）
• operatorLookupTable_：OperatorName → OperatorHandle 的查找表。LeftRight 是 PyTorch 自家的"读优先"并发数据结构，读的一方不必等锁
• backendFallbackKernels_：每个运行时 key 的"兜底"kernel —— 当某算子在某后端没注册具体实现时用

每个算子对应一个 OperatorEntry（aten/src/ATen/core/dispatch/OperatorEntry.h:70），它的核心字段：

class OperatorEntry final {
 private:
    OperatorName name_;
    std::optional<AnnotatedSchema> schema_;

    // 热路径：定长数组，按 DispatchKeySet 索引
    std::array<KernelFunction, c10::num_runtime_entries> dispatchTable_;

    DispatchKeyExtractor dispatchKeyExtractor_;

    // 注册历史：hashmap，按 DispatchKey 索引，每个 key 可能有多次注册
    ska::flat_hash_map<DispatchKey, std::list<AnnotatedKernel>> kernels_;
    ...
};

注释里专门解释了为什么有两个数据结构（aten/src/ATen/core/dispatch/OperatorEntry.h:243-273）：

We do not combine dispatchTable and kernels into one hash map because kernels is a larger data structure and accessed quite infrequently while dispatchTable is accessed often and should be kept small to fit into CPU caches.

—— dispatchTable_ 是热路径，kernels_ 是冷路径。前者每次算子调用都要访问、必须 cache friendly；后者只在注册/反注册时访问。它们的同步规则：

dispatchTable_[k] = kernels_[k].front()   // 取最新注册的

每次有人 m.impl("aten::add", &my_add) 时，新 kernel 进入 kernels_[CPU].front()，然后 dispatchTable_[CPU] 同步更新到这一个。这就是为什么 Jupyter 里反复执行同一段 TORCH_LIBRARY_IMPL 不会出错 —— 后注册的覆盖前注册的。

5.2.1 注册接口的三层 API

PyTorch 提供三种公开注册 API，用法不同，底层 m.def 进入 Dispatcher::registerDef、m.impl 进入 Dispatcher::registerImpl：

API	用途	典型场景
TORCH_LIBRARY(ns, m)	声明一个新算子库（一次声明 schema）	第一方算子或自定义算子集
TORCH_LIBRARY_IMPL(ns, key, m)	给 (库, key) 注册具体实现	NVIDIA / AMD 等后端给同一算子注册不同 device 实现
TORCH_LIBRARY_FRAGMENT(ns, m)	在已有库追加更多算子	大型库分成多个 .cpp 文件，每个文件用 fragment 注册

举一个真实例子，PyTorch 自己注册 ATen 库的 add 算子：

// 生成结果的形状（真实代码在构建目录的 RegisterSchema.cpp、RegisterCPU_0.cpp / RegisterCUDA_0.cpp、
// VariableType_*.cpp 等分片文件里）
TORCH_LIBRARY(aten, m) {
    m.def("add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor");
}

TORCH_LIBRARY_IMPL(aten, CPU, m) {
    m.impl("add.Tensor", TORCH_FN(wrapper_CPU_add_Tensor));
}

TORCH_LIBRARY_IMPL(aten, CUDA, m) {
    m.impl("add.Tensor", TORCH_FN(wrapper_CUDA_add_Tensor));
}

TORCH_LIBRARY_IMPL(aten, Autograd, m) {
    m.impl("add.Tensor", TORCH_FN(VariableType::add_Tensor));
}

那些 wrapper_<DispatchKey>_<name>_<overload> 的名字不是我编的，是 torchgen 拼出来的 —— add.Tensor 是 structured 算子，torchgen/dest/register_dispatch_key.py:809 里那行 prefix=f"wrapper_{self.backend_index.dispatch_key}_" 再接上 torchgen/api/native.py 的 name()（算子名 + _ + overload 名）就是命名规则本身（非 structured 算子走 :312，前缀里还会多嵌一段 overload 名）；本机 libtorch_cpu.so 里能直接搜到 wrapper_CPU_add_Tensor 这个符号。代码生成器从 native_functions.yaml 读出每个算子的 schema 和 dispatch 表，自动展开成上面的注册块。第 6 章会拆代码生成器。

5.2.2 alias keys 与"组合实现"

注意 Autograd 这个 key 看起来是个普通后端，但其实它是 alias key —— 它在内部展开成 AutogradCPU + AutogradCUDA + AutogradHIP + ...。同样的 alias 还有 CompositeImplicitAutograd、CompositeExplicitAutograd 等。

为什么要有 alias key？因为 PyTorch 团队不想为每个 (functionality, backend) 笛卡尔积都写一份代码。CompositeImplicitAutograd 的语义是 "用其他算子实现的算子，不需要专门的 autograd" —— 比如 torch.square(x) 内部就是 at::pow(self, 2)（aten/src/ATen/native/UnaryOps.cpp:772），pow 自带 autograd，整体 autograd 自动正确，不需要专门写 SquareBackward —— 实跑 torch.square(x).grad_fn 得到的就是 PowBackward0。

注册 alias key 时，dispatcher 在 computeDispatchTableEntry 阶段把它展开到所有对应的 runtime keys。这是为什么 dispatchTable_ 的查找永远是 O(1) —— alias 解析在注册期完成，运行期不付代价。

5.3 主路径：Dispatcher::call 的全旅程

打开 aten/src/ATen/core/dispatch/Dispatcher.h:776：

template <class Return, class... Args>
C10_ALWAYS_INLINE_UNLESS_MOBILE Return Dispatcher::call(
    const TypedOperatorHandle<Return(Args...)>& op,
    Args... args) const {

    auto dispatchKeySet =
        op.operatorDef_->op.dispatchKeyExtractor()
            .template getDispatchKeySetUnboxed<Args...>(args...);

    const KernelFunction& kernel =
        op.operatorDef_->op.lookup(dispatchKeySet);

    // ... profiler 切片 ...

    return kernel.template call<Return, Args...>(
        op, dispatchKeySet, std::forward<Args>(args)...);
}

三步：

1. getDispatchKeySetUnboxed：扫描所有 Tensor 参数的 key_set_，OR 起来，再按 thread-local 的 include / exclude 集合增减（inference_mode、functorch 变换、TorchDispatchMode 等都通过它们生效）
2. lookup：把 keyset 转成 dispatch table 索引，取出 KernelFunction
3. kernel.call(...)：调用 KernelFunction 持有的具体实现

让我们逐一拆开。

5.3.1 DispatchKeyExtractor：从参数算 keyset

DispatchKeyExtractor（aten/src/ATen/core/dispatch/DispatchKeyExtractor.h）是一个针对每个算子定制的对象，它知道 哪些参数是 Tensor、哪些不是。运行期它只对 Tensor 参数读 key_set_：

// 简化版（省略了按后端区分 fallthrough 掩码的分支）
template <class... Args>
DispatchKeySet getDispatchKeySetUnboxed(const Args&... args) const {
    auto ks = detail::multi_dispatch_key_set(args...);   // 折叠所有 Tensor 参数的 key_set
    // computeDispatchKeySet 内部：((ks | tls.included_) - tls.excluded_) & nonFallthroughKeys_
    return impl::computeDispatchKeySet(ks, nonFallthroughKeys_);
}

tls_local_dispatch_key_set() 返回 thread-local 的"包含 / 排除"对：

• included 是当前线程"应该附加上的"keys（如 torch.compile trace 时附加 Functionalize）
• excluded 是当前线程"应该屏蔽的"keys（如 torch.inference_mode() 排除 Autograd）

最终的 keyset 是 (参数 keysets ∪ tls.included) ∩ ¬tls.excluded，再与该算子的非 fallthrough 掩码相与（§5.6）。这条计算让 dispatcher 同时处理"张量自带"（如所在设备、布局、是否 inference tensor）和"上下文决定"（如 inference_mode、autocast）两种来源的 dispatch 行为。

举个具体例子：用户写 with torch.inference_mode(): a + b，且 a.requires_grad=True：

• 参数 keysets：a.key_set_ = {CUDA, ADInplaceOrView, AutogradCUDA, AutocastCUDA}。这里的 CUDA 本身就是 Dense 功能位 + CUDA 后端位的组合；AutogradCUDA、ADInplaceOrView 是非 inference 张量创建时就带上的，与 requires_grad 无关（本机 CPU 版实跑 torch._C._dispatch_keys(torch.randn(2)) 得到 CPU, ADInplaceOrView, AutogradCPU, AutocastCPU，不开 requires_grad 也一样）
• tls.excluded：InferenceMode 把 Autograd 相关 key 加入 excluded（同时把 ADInplaceOrView 从 included 里摘掉）；Autocast* 在没开 autocast 时本来就在 excluded 里
• 最终 keyset：Autograd 与 Autocast 被减掉，剩下 {CUDA, ADInplaceOrView}，外加 TLS included 里常驻的 BackendSelect；对 add.Tensor 来说这两个 key 都是 fallthrough（§5.6），会被掩码滤掉

dispatcher lookup 这个 keyset，找到的就是 aten::add 的 CUDA backend kernel —— autograd 那一层从 keyset 层面就被摘掉了。

注意 torch.no_grad() 不是这个机制。它走的是另一条路：GradMode::set_enabled(false) 只翻 AutogradState 里的一个 TLS 布尔（c10/core/GradMode.cpp:9），完全不动 dispatch key。所以 no_grad() 下 keyset 里的 AutogradCUDA 照样在、autograd kernel 照样被 lookup 到并执行 —— 只是它里面的 compute_requires_grad() 第一行就 if (!GradMode::is_enabled()) return false;（torch/csrc/autograd/functions/utils.h:60），于是不建 grad_fn、直接 redispatch 下去。这正是第 1 章 §1.7 说的"no_grad 仍然走 autograd dispatch，只是跳过反向图构建"。

更进一步：inference_mode() 两件事一起做 —— 既改 TLS（把 Autograd 放进 excluded、把 ADInplaceOrView 从 included 里拿掉、同时关掉 GradMode），又让在它作用域里新建的张量成为 inference tensor，本身就不带 Autograd key。c10/core/InferenceMode.h:14 的 Note [Expected TLS state in InferenceMode] 把这两条 TLS 不变式写得很清楚。真正的差别不在"改不改 TLS"，而在 ADInplaceOrView 那一层也被摘掉了 —— 这也是为什么 inference tensor 在 InferenceMode 之外做 inplace 会直接报错（见第 2 章 §2.15 的 version counter）。

5.3.2 DispatchKeyExtractor 的"参数选择"优化

不是每个算子的所有参数都参与 dispatch。比如 tensor.add(scalar)，第二个参数是 Scalar 而不是 Tensor —— 它的 dtype 不影响 dispatch。DispatchKeyExtractor 在算子注册时根据 schema 算出 "哪些参数位的 keyset 要被 OR"，把它编码成一个 bitmap：

// 简化版
class DispatchKeyExtractor {
    c10::utils::bitset dispatch_arg_indices_reverse_;
    DispatchKeySet nonFallthroughKeys_;
    ...
};

这张位图由 makeBitsetForDispatchArgs(schema) 在注册 schema 时从 FunctionSchema 直接算出，供 boxed 路径（getDispatchKeySetBoxed）在 IValue 栈上按位取参数；unboxed 路径则靠 multi_dispatch_key_set 的模板重载，在编译期就只对 Tensor / optional<Tensor> / Tensor 列表（以及 Generator）这类参数取 key_set，其余类型是空操作。两条路径都让"5 个参数中有 2 个 Tensor"的算子的 keyset 计算只读 2 个张量，不浪费时间在 scalar / int 参数上。

5.3.3 OperatorEntry::lookup：O(1) 数组索引

aten/src/ATen/core/dispatch/OperatorEntry.h:182：

const KernelFunction& lookup(DispatchKeySet ks) const {
    const auto idx = ks.getDispatchTableIndexForDispatchKeySet();
    if (C10_UNLIKELY(idx == -1)) {
        reportError(ks.highestPriorityTypeId());
    }
    const auto& kernel = dispatchTable_[idx];
    if (C10_UNLIKELY(!kernel.isValidUnboxed())) {
        if (!kernel.isValid()) {
            reportError(ks.highestPriorityTypeId());
        }
    }
    return kernel;
}

—— 核心就是 dispatchTable_[idx]。idx 怎么算？

// c10/core/DispatchKeySet.h:488
int getDispatchTableIndexForDispatchKeySet() const {
    // 1. 取最高优先级的 functionality（把低 16 位的 backend 段移开再找最高位）
    auto functionality_idx =
        DispatchKeySet(repr_ >> num_backends).indexOfHighestBit();
    // 2. 查预计算好的 (offset, mask) 表
    auto offset_and_mask = offsetsAndMasks()[functionality_idx];
    // 3. 在 mask 允许的范围内取最高优先级的 backend
    auto backend_idx =
        DispatchKeySet((repr_ & offset_and_mask.mask) >> 1).indexOfHighestBit();
    return offset_and_mask.offset + backend_idx;
}

回顾第 3 章 §3.5：DispatchKeySet 是 64-bit bitmap，分两段（低 16 位 backend + 其上的 functionality 位）。这里做的是"最高位 + 查偏移表"：先取最高的 functionality 位，用它去查一张预计算的 offsetsAndMasks 表拿到该 functionality 在运行时表里的起始偏移和 backend 掩码；如果这个 functionality 不是"可按后端定制"的，掩码就是 0，backend_idx 恒为 0，直接返回 offset。

注意表的尺寸不是 48 × 16 = 768，而是 num_runtime_entries = 48 + 6 × 15 = 138 —— 只有 6 个 functionality（Dense / Quantized / Sparse / SparseCsr / NestedTensor / AutogradFunctionality）需要为每个 backend 各留一格，其余 42 个各占一格就够。这套"偏移表"正是为了不把 768 格全开出来而存在的。

为什么不直接用一条 __builtin_clzl 当索引？因为 backend 与 functionality 是两个正交维度，不能简单 concat。这种"两段最高位 + 偏移表 → 一维索引"是 PyTorch dispatcher 的核心算法 —— 既保留了多维 dispatch 的语义，又把查找压到一次数组访问。

5.3.4 dispatchTable_ 的尺寸账

num_runtime_entries 是编译期常量，定义在 c10/core/DispatchKey.h:585：

constexpr uint16_t num_runtime_entries = num_functionality_keys +
    (numPerBackendFunctionalityKeys() * (num_backends - 1));

代入 v2.11 的数字：48 + 6 × (16 - 1) = 138。也就是说每个 OperatorEntry 的 dispatchTable_ 是 138 个 KernelFunction 槽位；一个 KernelFunction 是一个 BoxedKernel、两个裸指针再加一个 tokens_ 智能指针，在 x86-64 上共 40 字节，整张表约 5.4 KB。

aten/src/ATen/native/native_functions.yaml 里声明了 2674 个 - func: 条目，实际进 dispatcher 的算子还更多 —— 本机 CPU 版 torch 2.11.0 里 torch._C._dispatch_get_all_op_names() 返回 3692 个。乘上每张约 5.4 KB 的表，光 dispatchTable_ 就是 20 MB 量级 —— 远超 L1/L2，但真实进程实际只反复走其中几十到几百个算子，热的那几张表会稳稳待在 cache 里。

注意一个细节：dispatchTable_ 是按值嵌入 OperatorEntry 的，不是指针引用。这种"扁平内存"布局让访问 dispatchTable_ 不需要额外指针跳转，代价是 OperatorEntry 本身不小。所以 Dispatcher 把真正的 OperatorEntry 放在 std::list<OperatorDef> 里稳定持有，查找表 operatorLookupTable_ 只存轻量的 OperatorHandle。这是性能与灵活性之间又一个精细的权衡。

理解这个内存账，你能预判 PyTorch 为什么对"加一个新 functionality key"那么谨慎 —— 每加一个 single-key，所有 OperatorEntry 的 dispatchTable_ 都要扩张。

5.3.5 KernelFunction::call：模板化调用

KernelFunction（aten/src/ATen/core/boxing/KernelFunction.h:90）持有一个 boxed kernel 加两个类型擦除的函数指针（字段在 :290-292，其后还有一个管理失效令牌的 tokens_）：

class KernelFunction final {
    ...
 private:
    BoxedKernel boxed_kernel_func_;
    void* unboxed_kernel_func_;       // 类型擦除的指针
    void* sym_unboxed_kernel_func_;   // 支持符号 shape 的版本
};

调用时根据 Return(Args...) 模板参数把 unboxed_kernel_func_ 重新解释成具体类型（aten/src/ATen/core/boxing/KernelFunction_impl.h，简化）：

template <class Return, class... Args>
Return KernelFunction::call(...) const {
    if (C10_LIKELY(unboxed_kernel_func_ != nullptr)) {
        using ActualSignature = Return(OperatorKernel*, DispatchKeySet, Args...);
        auto* f = reinterpret_cast<ActualSignature*>(unboxed_kernel_func_);
        return (*f)(boxed_kernel_func_.getFunctor(), dispatchKeySet,
                    std::forward<Args>(args)...);
    }
    // 没有 unboxed 实现时退回 boxed：把参数装进 IValue 栈再调
    return impl::BoxedKernelWrapper<Return(Args...)>::call(
        boxed_kernel_func_, opHandle, dispatchKeySet, std::forward<Args>(args)...);
}

快路径上没有虚函数、没有 std::function、没有 hashmap —— 就是一次 reinterpret_cast 加一次函数指针间接调用，和调一个普通 C 函数指针同量级。

5.3.6 LeftRight 并发结构：让 dispatch 路径完全不阻塞写

Dispatcher::operatorLookupTable_（不是 operators_）是一个 LeftRight<flat_hash_map<...>>（c10/util/LeftRight.h）。这是 PyTorch 自家造的 "reader-prioritized concurrent map"：

• 同时维护两份 hashmap "L" 和 "R"
• 读线程不等锁（wait-free）：通过原子计数器记录自己正在读哪一份 —— 所以每次读仍有一次原子写，文件开头的注释明说 "Reads still incur an atomic write"
• 写线程持 mutex，在后台那份上更新、切换前后台，等读者离开旧的那份后再把它也更新一遍（写要执行两次）

这个数据结构来自论文《Left-right: A concurrency control technique with wait-free population oblivious reads》。要注意它服务的是"按名字找算子"（Dispatcher::findOp / findSchema），并不在 Dispatcher::call 的热路径上：生成的 at::add 等函数只在第一次调用时按名字查一次，拿到 TypedOperatorHandle 后存进函数内的静态变量（torchgen/gen.py:686 的 static auto op = create_{name}_typed_handle();），之后每次 call 都直接经 op.operatorDef_ 访问 OperatorEntry。LeftRight 让算子注册（极少发生）走慢路径，按名字查找走不等锁的读路径。

5.4 boxed vs unboxed：一个算子的两副面孔

PyTorch 每个算子都有 两套调用约定：

graph LR
    subgraph Unboxed["Unboxed Call (快路径)"]
        U1["编译期已知 Return / Args"]
        U2["参数直接传寄存器"]
        U3["无 IValue 构造/析构，无类型检查"]
    end
    subgraph Boxed["Boxed Call (慢路径但通用)"]
        B1["参数装 IValue 栈"]
        B2["运行期解析 schema"]
        B3["每个参数构造/析构一次 IValue"]
    end
    Use1["at::add(a, b) C++ API"] --> Unboxed
    Use2["torch.ops.aten.add.Tensor(a, b) Python"] --> Boxed
    Use3["TorchDispatchMode 拦截"] --> Boxed
    Use4["vmap / FakeTensor"] --> Boxed
    Use5["JIT / TorchScript"] --> Boxed
    style Unboxed fill:#dcfce7,stroke:#22c55e
    style Boxed fill:#fef3c7,stroke:#f59e0b

unboxed 用 C++ 模板把参数类型在编译期固定，调用时不需要做类型检查 —— 像 C++ 普通函数调用一样快。boxed 把参数装成 c10::IValue（一个类型擦除的"任意 PyTorch 值"容器），通过一个统一签名 void(*)(OperatorHandle&, Stack*) 调用。

为什么需要 boxed？因为某些场景编译期不知道签名：

• Python 端：torch.ops.aten.add.Tensor(a, b) 在 Python 没有静态类型
• vmap / FakeTensor：要拦截所有算子用同一段代码处理，不能为每个算子写一份模板代码
• TorchScript / JIT：序列化的 IR 节点没有 C++ 静态信息
• Profiler / 钩子：要观察所有算子，必须用统一接口

PyTorch 用 C++ 模板在注册 kernel 时自动实例化 unboxed → boxed 适配器和反向适配器，让两套协议在同一个 KernelFunction 里共存。源码里 aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h（unboxed kernel 包成 boxed）与 aten/src/ATen/core/boxing/impl/boxing.h 的 BoxedKernelWrapper（unboxed 调用落到只有 boxed 实现的 kernel）就是这套适配。

5.4.1 box / unbox 的代价

box 一个 Tensor 要构造一个 IValue、给 TensorImpl 加一次引用计数，unbox 时再反过来一遍；参数越多，这份固定成本乘得越多。所以 PyTorch 一直在 hot path 上避免 box —— 一旦掉到 boxed 路径，每个算子调用都要多付这一串构造/析构。

但 boxed 路径不可避免，因为某些 mode（如 TorchDispatchMode）必须看到所有参数才能拦截。PyTorch 的策略是：

• 默认走 unboxed（编译期类型已知时）
• 遇到需要 box 的 mode 时再 box（如 dispatch 到 Python key 时）
• box 后下游沿 boxed 接口继续 redispatch，直到命中的 kernel 由自动生成的 boxed → unboxed 适配器把参数从栈上拆出来再调

这种"按需 box"避免了在绝大多数不需要 box 的场景里付代价。torch.compile 的 AOTAutograd 阶段靠 FakeTensorMode 等 Python Mode trace 计算图（第 13 章），这一阶段走的就是 boxed 路径；但 trace 只在编译时发生，运行时执行的是编译产物，不再逐算子回到 Python Mode。

5.4.2 IValue：boxed 路径的"通用容器"

c10::IValue（aten/src/ATen/core/ivalue.h）是 PyTorch 自己设计的类型擦除变体类型 —— 类似 std::variant。它能装：

• Tensor / Storage / Scalar
• int64_t / double / bool / std::string
• List<...> / Dict<...> / Tuple<...>
• Future / Stream / Generator / PyObject*
• ……TORCH_FORALL_TAGS 里共 28 种 tag

内部实现是一个 16 字节结构（8 字节 payload + Tag）：

struct IValue {
    ...
    Payload payload;                 // 8 字节, 类型擦除
    Tag tag{IValue::Tag::None};      // uint32_t 枚举, 标记类型
};

box 一个 Tensor：

IValue value(tensor);   // 在 payload 里就地构造一个 at::Tensor（拷贝 => TensorImpl 引用计数 +1）
                        // tag = Tensor

unbox 一个 Tensor：

Tensor t = value.toTensor();   // 检查 tag, 取出 payload 里的 Tensor（拷贝出来再 +1）

这种轻量级类型擦除让 box+invoke+unbox 的全流程只是"几次小对象构造 + 一次间接调用"，比起真正的动态类型系统已经很省。它是 PyTorch 跨 C++ / Python / TorchScript / JIT 边界的"通用货币"。

5.5 redispatch：洋葱式中间层的关键

第 1 章 §1.7 我们用 a + b 的 autograd 路径预告过 redispatch 模式：autograd 这一层"做点事"（创建 grad_fn），然后重新进 dispatcher，这次把 Autograd key 排除掉。让我们看具体实现。

每个 autograd 包装函数（由 tools/autograd/gen_variable_type.py 生成；它按 num_shards=5 把结果摊到 torch/csrc/autograd/generated/VariableType_0.cpp ~ VariableType_4.cpp）大致长这样：

// 简化的 VariableType::add (生成代码)
Tensor add_Tensor(c10::DispatchKeySet ks, const Tensor& self, const Tensor& other,
                  const Scalar& alpha) {
    // 1. 需要求导时才创建反向节点
    std::shared_ptr<AddBackward0> grad_fn;
    if (compute_requires_grad(self, other)) {
        grad_fn = std::make_shared<AddBackward0>();
        grad_fn->alpha = alpha;
        // 设置 saved tensors / inputs
    }

    // 2. redispatch 到下一层：把当前 keyset 与 after_autograd_keyset 求交，
    //    等价于"把 Autograd 及其以上的所有 key 都剔掉"
    auto result = at::redispatch::add(
        ks & c10::after_autograd_keyset, self, other, alpha);

    // 3. 把 grad_fn 挂到 result
    if (grad_fn) {
        set_history(result, grad_fn);
    }
    return result;
}

第 2 步的 at::redispatch::add 调用 Dispatcher::redispatch（声明在 aten/src/ATen/core/dispatch/Dispatcher.h:198，下面是 :835 处的定义，省略了 trace 分支）：

template <class Return, class... Args>
Return Dispatcher::redispatch(
    const TypedOperatorHandle<Return(Args...)>& op,
    DispatchKeySet currentDispatchKeySet,
    Args... args) const {
    const KernelFunction& kernel = op.operatorDef_->op.lookup(currentDispatchKeySet);
    return kernel.template call<Return, Args...>(
        op, currentDispatchKeySet, std::forward<Args>(args)...);
}

注意区别：call 会重新算 keyset，redispatch 直接用调用方传入的 keyset。这种区分是为了让中间层精确控制下一层走什么 key。

c10::after_autograd_keyset 是预先算好的常量（c10/core/DispatchKeySet.h:731），代表"autograd 这一层之下还剩哪些 key"；同一处还有一个 after_ADInplaceOrView_keyset（:735）给视图/inplace 记账那一层用。与当前 keyset 求交之后，lookup 命中的就是 backend kernel（如 aten::add 的 CUDA 实现），而不是又一次命中 VariableType::add 死循环。

5.5.1 redispatch 在哪些地方用

PyTorch 有大量"中间层"用 redispatch 实现：

中间层	做的事	减去的 key
Autograd	创建反向节点	Autograd*
Functionalize	把 inplace 改成纯函数式	Functionalize
AutocastCUDA	fp32 → fp16 / bf16 自动转换	AutocastCUDA
Tracer (老 JIT)	记录到 IR	Tracer
Python (TorchDispatchMode)	Python 拦截	Python
FuncTorch (vmap / grad / jvp)	函数变换	FuncTorchXxx

整个 dispatcher 体系就是这些中间层叠在一起，每一层做完自己那份工作后 redispatch 到下一层。像剥洋葱：

graph LR
    User["a + b"] --> L1["FuncTorch (vmap/grad)"]
    L1 -- redispatch --> L2["AutocastCUDA"]
    L2 -- redispatch --> L3["Autograd"]
    L3 -- redispatch --> L4["Functionalize"]
    L4 -- redispatch --> L5["Python TorchDispatchMode"]
    L5 -- redispatch --> L6["Backend Kernel<br/>(CUDA add_kernel)"]
    style L6 fill:#dcfce7,stroke:#22c55e,stroke-width:2px

每一层都是可选的 —— 没启用时它不在 keyset 里、不会被 lookup 命中。这种"叠加式"设计让 PyTorch 能在不破坏现有 API 的前提下不断增加新中间层（vmap、torch.compile、TorchDispatchMode 全部是这样加进来的）。

5.5.2 一个具体追踪：vmap(grad(f))(x) 在 dispatcher 里走什么

举一个能把 redispatch 价值讲清的例子。functorch.vmap(functorch.grad(f))(x) 里 f 的每个算子进 dispatcher 时，keyset 里除了后端 key（如 CUDA）和 Autograd*，还会有 functorch 加进 TLS 的 FuncTorchDynamicLayerFrontMode / FuncTorchDynamicLayerBackMode，以及 vmap 对应的 FuncTorchBatched；若还有用户 TorchDispatchMode，另有 Python。

但 functorch 的处理顺序不是简单按 key 优先级一路吃下去。aten/src/ATen/functorch/DynamicLayer.cpp 的 NOTE: [functorch front and back key fallbacks] 写得很清楚：每个变换对应一组 key（vmap → FuncTorchBatched / FuncTorchVmapMode，grad → Autograd{Backend} / ADInplaceOrView）；优先级最高的 FuncTorchDynamicLayerFrontMode 先命中，取出变换栈栈顶那一层，改写 TLS 让只有这一层的 key 生效，再 redispatch；这一层做完后由 FuncTorchDynamicLayerBackMode 接住，转回 FrontMode 处理下一层。对 vmap(grad(f)) 大致是：

1. FuncTorchDynamicLayerFrontMode → 取栈顶的 grad 层，只放行它的 key
2. AutogradCUDA → 创建反向节点 → redispatch 减 Autograd*
3. FuncTorchDynamicLayerBackMode → 回到 FrontMode，轮到 vmap 层
4. FuncTorchBatched → 按 batch rule 处理批维
5. （如有用户 mode）Python → 调 __torch_dispatch__
6. CUDA → 真正调 CUDA kernel

一次算子调用要经过这么多层，每层都是"算 keyset → 查表 → 调 kernel"。开销当然不是零，但每一层都做了关键的事，且互不干扰 —— 这就是组合性 (composability) 在工程上的真实表达。

如果用 try-catch 或者全局 mutator 实现这套叠加，会有大量耦合 bug；用 dispatcher + redispatch 模式让每一层都是自包含的。

5.6 Fallthrough：让某些 key 无开销跳过

有时候我们想注册一个 key 但啥也不做，让 dispatcher 自动跳到下一个 key。比如：

• BackendSelect：绝大多数算子用不着它，它的 backend fallback 就是 fallthrough（aten/src/ATen/core/BackendSelectFallbackKernel.cpp）
• ADInplaceOrView：只有 inplace / view 算子需要记账，其余算子走 fallthrough（aten/src/ATen/core/VariableFallbackKernel.cpp:103）
• Autocast*：不在 autocast 名单里的算子走 fallthrough（aten/src/ATen/autocast_mode.cpp）

PyTorch 的解法是 CppFunction::makeFallthrough() / KernelFunction::makeFallthrough() —— 它返回的 KernelFunction 只是把 boxed 函数指针设成一个哨兵函数 fallthrough_kernel，而这个函数根本不该被执行：

// aten/src/ATen/core/boxing/KernelFunction.cpp:13
void fallthrough_kernel(OperatorKernel*, const OperatorHandle&, DispatchKeySet, Stack*) {
  TORCH_INTERNAL_ASSERT(0,
    "fallthrough_kernel was executed but it should have been short-circuited by the dispatcher. "
    ...);
}

真正的跳过发生在算 keyset 那一步。每当某个槽位的 kernel 更新，OperatorEntry 就调 dispatchKeyExtractor_.setOperatorHasFallthroughForKey(dispatch_key, dispatchTable_[dispatch_ix].isFallthrough())（aten/src/ATen/core/dispatch/OperatorEntry.cpp:485），把 fallthrough 的 key 从这个算子的 nonFallthroughKeys_ 掩码里去掉。运行期 computeDispatchKeySet 最后与这个掩码相与（§5.3.1），fallthrough 的 key 直接从 keyset 里消失，lookup 一次就落到下一层真正的 kernel，根本不进 fallthrough 函数。

顺带对比一下为什么 inference_mode() 比 no_grad() 快 —— no_grad() 只关 GradMode 布尔，autograd kernel 照样被 lookup 到、照样跑一遍；inference_mode() 直接让 inference 张量不带 Autograd key、dispatcher lookup 都不进 autograd 槽位。

dispatcher 的"零成本"哲学：能在编译期 / 注册期消除的开销，绝不留到运行期。

5.6.1 backend fallback 的兜底机制

Dispatcher::backendFallbackKernels_ 是另一个"兜底"层。当某算子在某 backend 完全没注册任何 kernel（也没有可用的 alias kernel）时，dispatcher 退回到这个 backend 的全局 fallback。

例子：用户想给 aten::add 在某个自定义 backend 上"反正先走 CPU"，可以注册一个 backend fallback（自定义后端通常用 PrivateUse1 这个 key）：

TORCH_LIBRARY_IMPL(_, PrivateUse1, m) {
    m.fallback(torch::CppFunction::makeFromBoxedFunction<&my_fallback_to_cpu>());
}

这样所有未在 PrivateUse1 注册的算子都会进 my_fallback_to_cpu，里面可以做"把张量搬到 CPU 算完再搬回"的退路。这是新硬件后端"逐步迁移"的常用模式 —— 先全部 fallback 到 CPU，能跑就行；然后逐个把 hot 算子重写成 native 实现。

第 22 章自定义算子会演示完整的 backend 接入流程。

5.6.2 fallthrough 与"包含 / 排除"集合的关系

回到 §5.3.1 提过的 tlsLocalDispatchKeySet：它分 included 和 excluded 两部分。这套机制和 fallthrough 怎么共同工作？

• included 把某个 key 强制加入 keyset（例如 torch.compile 的某些阶段强制加入 Functionalize）
• excluded 从 keyset 中减去某个 key（例如 torch.inference_mode() 减去 Autograd）

如果某 key 是 fallthrough 注册的，注册期就把它从该算子的 nonFallthroughKeys_ 掩码里去掉；运行期算 keyset 时，先做 (参数 ∪ included) − excluded，再与这个掩码相与。fallthrough 和 included/excluded 是两个独立的开关：前者按算子、在注册期确定，后者按线程、在运行期随上下文变化；它们最终都作用在"算 keyset"这一步，查表本身仍是一次数组访问。

理解这两条路径的分工，是写自定义 mode 时的基础。

5.6.3 inference_mode 的"无开销"奥秘

让我们把 §5.6 与第 1 章 §1.7 提到的"为什么 inference_mode 比 no_grad 快"完整解开：

• no_grad()：只翻 GradMode 这个 TLS 布尔，dispatch key 一个都不动。autograd kernel 照样被 lookup 到、照样执行，只是里面 compute_requires_grad() 直接返回 false，省下建 grad_fn 与 saved tensor 的那部分工作，然后照常 redispatch 到 backend。所以 keyset 计算、lookup、autograd 那一跳的函数调用全都还在
• inference_mode()：TLS 上把 Autograd 放进 excluded 的同时，还把 ADInplaceOrView 从 included 里摘掉；并且在它作用域里新建的张量本身就不带 Autograd* key。少走的不只是 autograd 那一层，还有视图/inplace 记账那一层

差异在小算子密集的推理循环里能被 profiler 量出来，具体多少要看你的算子构成。inference_mode() 的代价是张量"结构上"已经不参与 autograd，不能再用 requires_grad_(True) 救回来。所以它是"明确知道这段代码不反向"时的最优选项。第 7 章 autograd 章会再讲这两种 mode 的实现细节。

5.7 TorchDispatchMode：Python 端的 dispatcher 拦截

第 1 章 §1.16 预告过 TorchDispatchMode —— 一个 Python 端可以在 dispatcher 层拦截所有算子的接口。它的实现是 Python 这个 DispatchKey。

工作机制：

sequenceDiagram
    autonumber
    participant U as Python with CountOps()
    participant TLS as thread-local
    participant D as Dispatcher
    participant P as PythonKernel
    participant B as Backend
    U->>TLS: __enter__: 把 Python key 加入 included
    Note over U,TLS: 进入 with 块
    U->>D: a + b (任意算子)
    D->>D: keyset = ... ∪ {Python}
    D->>D: lookup → PythonKernel
    D->>P: 调 boxed Python kernel
    P->>P: 调 user.__torch_dispatch__(op, types, args, kwargs)
    P->>D: user 内部调 op(*args) 重新进入 dispatcher
    D->>D: keyset 不再带 Python (mode 已暂时出栈)
    D->>B: lookup → Backend kernel
    B-->>P: 返回结果
    P-->>U: 返回结果
    U->>TLS: __exit__: 把 Python key 从 included 移除

注意第 7、8 步：调用 __torch_dispatch__ 之前，C++ 侧的 StashTorchDispatchModeGuard（torch/csrc/utils/torch_dispatch_mode.h）把当前 mode 从 mode 栈上暂时弹出，栈空时 Python key 也随之从 TLS included 里移除；所以 Python kernel 内部再调 op(*args) 时不会又落回同一个 mode，避免无限递归。

这套机制让用户能用 ~10 行 Python 代码做：

• 算子计数（第 1 章 §1.16 的例子）
• 假执行（FakeTensorMode 实现，torch.compile 用）
• inplace 改纯函数式（FunctionalTensorMode 实现）
• 把每次算子调用记录成 FX 节点（ProxyTorchDispatchMode 实现）
• ……

所有"对每个算子做点什么"的需求，都可以用 Mode 实现。这是 PyTorch 最灵活的扩展机制之一，也是为什么 torch.compile 能在不破坏 eager API 的前提下优雅切入 —— 它的 AOTAutograd 阶段就是 FakeTensorMode / FunctionalTensorMode / ProxyTorchDispatchMode 这几个 Mode 的组合。

5.7.1 实现一个真正有用的 Mode

举个有信息量的例子：写一个 MemoryLeakDetector Mode，记录每次 alloc 但不记录 free，最终找出"只 alloc 不 free"的张量来源：

import torch, traceback
from torch.utils._python_dispatch import TorchDispatchMode

class MemoryLeakDetector(TorchDispatchMode):
    def __init__(self):
        self.allocations = {}     # tensor_id -> stack trace

    def __torch_dispatch__(self, op, types, args=(), kwargs=None):
        result = op(*args, **(kwargs or {}))
        # 只记录 op 名字含 "empty" 或 "zeros" 的（典型分配）
        if 'empty' in str(op._schema.name) or 'zeros' in str(op._schema.name):
            if isinstance(result, torch.Tensor):
                self.allocations[id(result)] = traceback.format_stack()
        return result

    def report(self):
        print(f'Live allocations: {len(self.allocations)}')

短短 15 行就实现了一个张量分配追踪器。无 Mode 的等价物需要 monkey-patch torch.empty、torch.zeros 等 N 个 API，且会被绕过 —— Mode 因为在 dispatcher 层拦截，任何路径来的算子都逃不掉。

5.7.2 Mode 与 __torch_function__ 的区别

新手容易混淆 TorchDispatchMode 与 __torch_function__：

• __torch_function__ 在 Python 端拦截。挂在 Tensor 子类上，每次 Python 调 torch.add(...) 等 API 时被检查
• TorchDispatchMode 在 C++ dispatcher 层拦截。Python kernel 通过 Python DispatchKey 命中

差别大致：

特性	__torch_function__	TorchDispatchMode
拦截点	Python API	C++ dispatcher
看到的算子粒度	torch.add 这种"友好"名字	aten::add.Tensor 这种 schema-精确名字
开销	每次 Python call 多一次属性查找	每次 dispatch 多一次 boxed call
能拦截 C++ 内部调用吗	不能	能
能拦截 autograd 反向	不能（反向走 C++）	能

要做"假执行 / FX trace / 全局观察"这类需求，用 TorchDispatchMode。要做"自定义子类让用户用起来像 Tensor"这类需求，用 __torch_function__。第 9 章 nn.Module 章会再讨论这两条路径。

5.7.3 一个被低估的应用：分布式调试

写过分布式训练的人都知道 NCCL hang / OOM 之类的故障极难诊断，因为某个 rank 上某个算子突然慢下来，其他 rank 在 AllReduce 上等成一片。

用 TorchDispatchMode 可以做一个 RankAwareLogger：每个 rank 拦截所有算子，记录"算子名 + 张量 shape + 时间戳"到本地文件。当 hang 发生后，把所有 rank 的日志拼起来对比，就能精确看到"哪个 rank 在哪个算子上和别人不同"。

这种调试模式是"用 30 行 Python 做出实用诊断工具"的典型例子。

5.7.4 stacked Mode：嵌套 with 块

TorchDispatchMode 支持嵌套：

with FakeTensorMode():
    with FunctionalTensorMode():
        out = f(x)         # 两个 mode 都看得到

实现是维护一个线程局部的 mode stack（c10/core/impl/TorchDispatchModeTLS.h），每次进 dispatcher 命中 Python key 时按栈顶 mode 调它的 __torch_dispatch__。调用期间当前 mode 暂时出栈，所以它内部再调 op(*args) 时由下一个 mode 接管。这种"用栈做嵌套"的模式让 mode 可以任意组合，是 torch.compile 多个 trace 阶段联动的关键。

5.8 dispatcher 到底有多快

本机没有 GPU，也拿不到可引用的官方基准，所以这一节不给伪精确的纳秒数 —— 但各条路径的相对量级是由代码结构决定的、可以直接读出来：

路径	每次调用要做的事	相对量级
纯 unboxed call（无 autograd、无 profiler）	几次位运算算 keyset + 一次数组索引 + 一次函数指针间接调用	基准
含 autograd 的一次 redispatch	上面那套走两遍，外加建一个 AddBackward0 节点（含一次堆分配）	明显高于基准
Python 端调 torch.add	再加 CPython 的方法解析、PythonArgParser 参数解析、THPVariable_Wrap	更高
含 TorchDispatchMode 拦截	参数全部 box 成 IValue，回到 Python 执行用户代码，再 unbox	最高

关键在于结构上的量级差，而不是某个具体数字：unboxed 那条路上没有虚函数、没有哈希、没有分支树，只有"位运算 + 数组下标 + 间接调用"。这就是 dispatcher 在工程上的成就：让"多分派"几乎和直接函数调用一样便宜。要量自己机器上的数字，用 torch.profiler 或 timeit 对比开/关 requires_grad、开/关 mode 的同一段循环。

5.8.1 dispatcher 开销的累积

单次很快，但训练里每个张量 op 都要进一次 dispatcher。一个简单的 transformer block：

Attention: norm + linear*3 (QKV) + matmul + softmax + matmul + linear (out)
           + reshape / transpose / residual add 若干 ≈ 12 算子
MLP:       linear + gelu + linear ≈ 3 算子
total ≈ 15 算子 / block

70 层 transformer ≈ 1000 次 dispatcher 调用 / forward；反向再来一遍，一个 step 就是几千次。

结论的形状很清楚：dispatcher 开销是常数 × 算子个数，与张量大小无关。所以大模型训练（每个 op 的 GPU 时间是毫秒量级）里它可以忽略；而小 batch 的自回归 decode（每步 GPU 时间可能只有几百微秒、算子却一个不少）里，它就会浮到 profile 的顶部。这正是推理引擎要么上 CUDA Graph、要么上 torch.compile 把整段调用链固化的原因。

5.8.2 减少 dispatcher 开销的几条路

如果你需要在小算子密集场景挤性能，几条思路：

1. torch.compile：把 forward 里的逐点 / 归约算子融合成少量 Triton / C++ kernel，dispatcher 调用大幅减少（matmul 等仍作为 extern kernel 经 ATen 调用）
2. torch.cuda.graph：把 dispatcher + kernel launch 压成一次 graph replay
3. 合并 Python 端调用：把 q = qkv[:, 0]; k = qkv[:, 1]; ... 改成 q, k, v = qkv.unbind(1)，Python 端少两次调用、autograd 只记一个 UnbindBackward0 节点（unbind 内部仍逐个调 select，dispatcher 调用次数并不减少）
4. 使用 fused 算子：torch.nn.functional.scaled_dot_product_attention 把 QKᵀ、scale、softmax、乘 V 这几次调用合成一个入口，满足条件时走融合 kernel
5. 禁用 autograd：在确定不反向的代码用 inference_mode() 而非 no_grad()，省掉 autograd 那一跳本身

第 21 章 profiler 章会教你怎么用 PyTorch 自带工具找出 dispatcher 开销占比高的代码段。

但一旦掉到 boxed 路径或者有 mode 干预，开销会明显上升。这也是 torch.compile 能省下框架开销的原因之一 —— 编译产物里 Inductor 生成的融合 kernel 直接启动，不再逐个算子走 dispatcher。

5.8.3 调试工具：TORCH_SHOW_DISPATCH_TRACE

PyTorch 提供一个环境变量 TORCH_SHOW_DISPATCH_TRACE=1，开启后 dispatcher 会把每一次 call / redispatch 打印到 stderr。注意它要求构建时带上 -DHAS_TORCH_SHOW_DISPATCH_TRACE（仓库根目录 setup.py 第 26 行的注释写了这一点），或者干脆用一个非 NDEBUG 的 debug 构建：

TORCH_SHOW_DISPATCH_TRACE=1 python -c "
import torch
a = torch.randn(2, 2, requires_grad=True)
b = torch.randn(2, 2)
c = a + b
"

输出（节选）：

[call] op=[aten::add.Tensor], key=[AutogradCPU]
 [redispatch] op=[aten::add.Tensor], key=[CPU]
...

格式由 aten/src/ATen/core/dispatch/Dispatcher.cpp:70 的 _print_dispatch_trace 决定：行首的缩进是 dispatch 嵌套深度，key=[...] 里打的是这次 keyset 的 highestPriorityTypeId()（即真正被选中的那一个 key），不是整个 keyset。你能清楚看到 redispatch 在哪些层级发生、最终命中哪个 backend kernel。这是诊断"为什么我的算子走错路径"的关键工具。

这个 trace 最适合的场景，是在自己编译的 debug 版 PyTorch 上排查"某个算子为什么走进了意料之外的 kernel"：它按调用顺序列出每一次 dispatch 命中的 key，哪一层 redispatch、哪一层落到 boxed fallback 一目了然，不必只凭最终的报错信息倒推。

dispatcher 是 PyTorch 大部分"奇怪行为"的真实源头，trace 是直接看清这些行为的最锋利工具。

再强调一次：这段 trace 代码被 aten/src/ATen/core/dispatch/Dispatcher.h:782 的 #if defined(HAS_TORCH_SHOW_DISPATCH_TRACE) || !defined(NDEBUG) 包着（另外三处在 :840 / :860 / :928，分别对应 redispatch 与 boxed 路径）—— 官方 release wheel 不带这个宏，环境变量设了也不会有输出。

写自定义算子（第 22 章）时，如果手头有带 trace 宏的构建，它是排查注册问题的好帮手。

5.8.4 历史回顾：dispatcher 重构的几次浪潮

PyTorch dispatcher 自身有清晰的演进史：

时期	特点
pre-1.0	没有统一的 dispatcher，算子分派散落在早已废弃的 Type 抽象与各处 if-else 里
1.x 早期	有了 DispatchKey + per-op kernel registration，但 functionality 与 backend 挤在同一个一维枚举里
1.x 中期（v1.10 起）	引入 __torch_dispatch__，Python 子类可以在 dispatcher 层拦截算子，随后演进出 TorchDispatchMode
1.x 后期（v1.12）	关键重构（PR #69633 "free up dispatch key space"）：DispatchKey 拆成 functionality + backend 两个维度，DispatchKeySet 变成本章讲的"后端位 + 功能位"布局
2.0+	torch.compile 上线，dispatcher 与 Dynamo / Inductor 配合
2.x	Compiled Autograd 让反向图也能被 Inductor 编译，dispatcher 在反向路径上的开销同样可以被编掉

这条演进的方向很清楚：能提前到编译期 / 注册期做的决定，就不留到运行期。c10/core/DispatchKey.h 里那些"TODO: put this in BackendComponents"的注释（FPGA、Vulkan、Metal 几个 key 上面就有），是这条整理工作还没做完的现场痕迹。

5.9 跨专栏关联

• 《Tokio 源码深度解析》专栏第 2 章 Future trait 与 poll 模型回顾：dispatcher 的 redispatch 模式与 Tokio future 的 poll 调用链有相似的"一层做事 + 让出 + 下一层接管"思想
• 《Rust 编译器与运行时揭秘》专栏第 7-8 章 Trait 静态分发与 Trait Object：Rust 的 trait 静态分发在编译期选定实现、trait object 在运行期经虚表单分派，PyTorch 的 dispatcher 则是运行期的"多分派"。前者零开销但需重新编译，后者有开销但允许动态注册 —— 各有所长
• 《MCP 协议设计与实现》专栏第 3 章 JSON-RPC 与消息格式：MCP server 的方法路由也是"name → handler"的多态映射，但 MCP 是按字符串方法名查找，PyTorch 是 enum 算出来的数组下标 —— 两者的性能要求本就不在一个档次
• 《vLLM 推理内核深度解析》专栏第 5 章 KV Cache 管理：vLLM 的场景比 PyTorch 窄得多（只有 LLM 推理、算子集合固定），所以它可以把大量"选实现"的决定提前到引擎初始化时定死，而不需要一套通用的运行期多分派 —— 和本章正好是一个反向的对照

5.10 dispatcher 设计的"通用启示"

如果你在自己的项目里设计一个"运行期多分派"系统，本章 PyTorch 的设计经验值得带走：

第一：bitmap + 优先级排序 + 数组下标查找 是处理 N 维多分派的经典套路。在维度有限、key 可以编号的场景里，其他选择（visitor + variant、虚函数链、hashmap）在性能上通常不如这个组合

第二：注册表与查找表分开：注册表（kernels_）允许复杂操作（重注册、查询、迭代），查找表（dispatchTable_）只为 hot path 服务，cache friendly

第三：alias 在注册期解析，运行期不付代价：把"逻辑组合"的复杂度推到注册期（一次性解析），运行期就是 O(1) 数组下标

第四：两套调用约定（boxed + unboxed）让"性能优先的常态路径"和"通用性优先的扩展路径"共存

第五：Mode 系统用一个特殊 key + thread-local stack 实现可叠加的运行期拦截器。这套模式的表达力很强

第六：让 hot path 上的检查在编译/注册期消除：fallthrough 掩码的预计算、alias key 的解析、DispatchKeyExtractor 的"哪些参数参与"全部在注册期完成，运行期只剩单数组访问

把这六条记下来，你设计任何"运行期多分派"系统都有参照。

5.11 一句话回顾整个 dispatcher

把整章的精华浓缩成一句：

PyTorch dispatcher = 把"张量元数据 + 线程上下文"通过 bitmap OR 算成 64-bit DispatchKeySet → 用 (functionality 最高位 → 查偏移) + (backend 最高位) 算出数组索引 → 取出 KernelFunction → 调它 → 多层中间层用 redispatch 减自己的 key 后递归命中下一层。

这一句话能在面试时让你 3 分钟讲清 PyTorch 的核心调度系统。

5.12 本章 takeaway

读完本章你应该能回答这些问题：

1. 为什么 PyTorch 的 dispatcher 能在上百个 key 里 O(1) 选中 kernel？ —— 因为它把 functionality × backend 两个维度的多分派压成单次数组下标
2. no_grad() 与 inference_mode() 的真实差异？ —— no_grad 只翻 GradMode 布尔、autograd kernel 照常执行，inference_mode 直接把 Autograd 与 ADInplaceOrView 从 keyset 里摘掉
3. box / unbox 的本质？ —— unbox 是 C++ 模板（编译期类型固定），box 是 IValue 栈（运行期动态）
4. redispatch 解决了什么？ —— 让 autograd / functionalize / autocast 等多个正交"中间层"可以叠加而不互相干扰
5. TorchDispatchMode 怎么实现的？ —— Python DispatchKey + 自动屏蔽自身防递归
6. dispatcher 的性能极限在哪？ —— unboxed 路径已经压到"位运算 + 数组下标 + 一次间接调用"，再往下就只剩把调用整段编掉（torch.compile / CUDA Graph）

如果你能在脑子里画出"a + b 在 dispatcher 里走过的几步" —— getDispatchKeySetUnboxed → lookup → KernelFunction::call → 可能的 redispatch —— 那本章的目标就达到了。后面 18 章的所有内容都建立在这套机制上：autograd 是 dispatcher 中间层、torch.compile 的 AOTAutograd 靠 dispatcher 上的 Mode 把算子 trace 成图、自定义算子是往 dispatcher 注册新 kernel。理解了 dispatcher，你就理解了 PyTorch 一半。

5.13 几条容易踩的坑

实战里最常见的几个 dispatcher 相关坑：

1. 误以为 __torch_function__ 子类会被 dispatcher 看到：实际上 __torch_function__ 只在 Python API 层拦截，dispatch 进 C++ 后子类信息丢失。要让 dispatcher 看到自定义类型，用 TorchDispatchMode 或自定义 backend。

2. 给 alias key 注册的 kernel 被具体 backend 覆盖：你给 CompositeExplicitAutograd 注册了一个实现，期望所有 backend 都用这个；但某个 backend 又注册了自己的具体实现，最终 dispatcher 用具体的 —— computeDispatchTableEntryWithDebug 里直接注册永远排在所有 alias 之前（aten/src/ATen/core/dispatch/OperatorEntry.cpp 的 "Alias Key Precedence" 注释）。如果你想"无条件"覆盖，只能在那个 backend key 上直接注册（后注册的覆盖先注册的）。

3. 以为 fallthrough 会在运行期"跳一下"：fallthrough 的 key 在算 keyset 时就被 nonFallthroughKeys_ 掩码滤掉了，dispatch trace 里根本看不到它，运行期 lookup 仍是一次 O(1)；调试时你会看到 dispatch trace 比预期短。反过来，如果 fallthrough_kernel 真被执行到，只会触发 TORCH_INTERNAL_ASSERT。

4. with torch.no_grad(): 包了 reentrant autograd 调用：某些算子内部会重新 enable autograd（如 checkpointing）。这种代码与外层 no_grad 互动有时反直觉，需要看 dispatch trace 才能确认。

5. __torch_dispatch__ 里错误地修改输入张量：dispatch mode 拦截到的算子参数有时候是用户原始张量、有时候是 wrapper 子类。盲目 inplace 修改可能污染外部状态。规范做法：永远 op(*args) 转发给底层 + 把返回值包回自己的子类，不直接改输入。

第 22 章自定义算子会列出更长的"自定义算子注册时容易踩的坑"清单。

5.14 一个练习：观察自己的 dispatch 链

import torch
torch.ops.aten.add.Tensor
# <OpOverload(op='aten.add', overload='Tensor')>

# 看注册的 dispatch keys
op = torch.ops.aten.add.Tensor
print(op._defined_in_python)
# 用 introspection API 看每个 key 注册了什么
for key in ['CPU', 'CUDA', 'Autograd', 'AutogradCPU', 'AutogradCUDA',
            'CompositeImplicitAutograd', 'CompositeExplicitAutograd']:
    try:
        print(f'{key}: registered =', op.has_kernel_for_dispatch_key(key))
    except Exception as e:
        print(f'{key}: {e}')

这段代码能让你看到 aten::add.Tensor 直接注册在哪些 key 上。在本机 CPU 版 torch 2.11.0 上实跑：CPU 为 True，CUDA 为 False（CPU wheel 里没有编译 CUDA kernel）；Autograd 为 True，AutogradCPU / AutogradCUDA 却是 False —— autograd 包装是注册在 alias key Autograd 上的独立 kernel，has_kernel_for_dispatch_key 只查直接注册，不看 alias 展开；两个 Composite* 也是 False，因为 add.Tensor 的默认实现挂在 CompositeExplicitAutogradNonFunctional 上。想看展开之后每个运行时 key 最终落到哪个 kernel，用 print(torch._C._dispatch_dump_table("aten::add.Tensor"))（会显示 AutogradCPU: ... [autograd kernel]、ADInplaceOrView: fallthrough ... [backend fallback] 等），或 print(torch._C._dispatch_dump("aten::add.Tensor")) 看注册历史。把这些信息和 dispatch trace 对照看，dispatcher 的全貌就在脑子里了。

5.14.1 进阶：用 PrivateUse1 接入一个"自己的后端"

为了真正掌握 dispatcher，可以试着用 PyTorch 预留的 PrivateUse1 key（C++ 的 TORCH_LIBRARY_IMPL(aten, PrivateUse1, m)，或 Python 的 torch.library.impl(..., "PrivateUse1")）给若干算子注册实现 —— functionality key 本身是编译期固定的枚举，不能在运行期新增。完整的练习可以参考 PyTorch 官方教程 "Extending dispatcher for a new backend in C++"（搜 PyTorch docs 即可）。

写完这个练习之后，你会真正体会到：dispatcher 不是黑盒、也不只是 PyTorch 内部用的工具，它是一个开放的多分派引擎，任何想在 PyTorch 生态里做自定义扩展的人都可以接入。这就是为什么国产 AI 芯片厂商能在不修改 PyTorch 主仓的情况下接入自家硬件 —— 全靠 dispatcher 这个公开扩展点。

下一章拆 ATen 代码生成体系：native_functions.yaml + derivatives.yaml + torchgen 怎么把数千个算子的 C++ / Python / autograd 包装一次性生成出来。本章讲的 dispatchTable_ 注册过程，就是代码生成器最后那一步往 dispatcher 喂数据。