一次 a + b 的旅程 · dispatcher 分发器 · PyTorch 源码漫游

你写下 a + b，PyTorch 怎么从 Python 一路找到「该在 CPU 还是 CUDA、该不该记梯度、要不要混合精度」的正确内核？答案是它的中枢神经——dispatcher。而 autograd、autocast 其实只是这条分发链上的一层拦截。

一次 a + b 的旅程 · dispatcher 分发器

数据结构清楚了，这一站我们让一个最简单的操作——a + b——动起来，跟着它从 Python 一路往下走。走完你会遇见 PyTorch 的中枢神经：dispatcher（分发器），并看清一个让无数人恍然大悟的事实：autograd 和混合精度，其实只是这条链上的「拦截层」。

一个 `+`，背后有几十种实现

先想想 dispatcher 要解决什么问题。add 这一个操作，底下有一大堆不同的实现：

a、b 在 CPU 上，还是 CUDA 上？内核不同。
float32 还是 int64？内核不同。
要不要记梯度(autograd)？要不要自动混合精度(autocast)？稀疏张量？量化张量？……

如果用一坨 if device == 'cuda' and dtype == float and requires_grad ... 去判断，那会是一场噩梦，而且每加一种后端就要改一次。PyTorch 需要一个干净、可扩展的「路由」机制——这就是 dispatcher。

旅程：从 Python 到内核

跟着 a + b 走：

Python 层：a + b 触发 Tensor.__add__，它穿过 Python/C++ 绑定(torch._C)，调到 C++ 的 at::add。
进入 dispatcher：at::add 不直接算，而是把请求交给 dispatcher。
算出「该走哪条路」：dispatcher 看 a、b 的属性——在哪个 device、什么 dtype、requires_grad 吗、autocast 开着吗——凑出一个 DispatchKeySet（分发键集合）。
查表、调内核：每个算子（如 aten::add）都有一张分发表(dispatch table)：分发键 → 内核函数指针。dispatcher 从键集里挑优先级最高的那个键，调它注册的内核。

图表加载中…

精髓：分发键是「一层层」的，autograd 就是其中一层

这里是全章的题眼。分发键不是「二选一」，而是有优先级、可以逐层往下传的。 高优先级的键先被调用，处理完自己那摊事，再**「重新分发(redispatch)」给下一个键。于是一次 add 的真实旅程，是穿过一串拦截层**：

a + b(requires_grad=True, autocast 开)
   │
   ▼  ① Autocast 键:把输入按需转成 fp16,然后 redispatch ↓
   ▼  ② Autograd 键:记一笔账(建计算图节点),然后 redispatch ↓
   ▼  ③ CUDA 键:真正在 GPU 上把加法算出来   ← 到这才落地计算

看懂这张图，你就解锁了 PyTorch 一个格外优雅的设计：

autograd（自动求导）和 autocast（自动混合精度），本质上都只是 dispatcher 链上的一层「拦截键」。 它们不是缝在每个算子里的特殊逻辑，而是统一地在分发链上插一层：轮到我时，我做我的事（记账 / 转精度），做完把球往下传。

这解释了两件你一直觉得「理所当然」的事：

为什么 autograd 是透明的？ 因为它是一层拦截：你写 a + b，Autograd 键先跑、悄悄把这步记进计算图，再让真正的加法内核去算。你什么都没多写，梯度的账就记上了。
为什么 with torch.autocast(): 能一键混合精度？ 因为它开启了 Autocast 键这一层：进来的算子先被它拦下、按规则把输入转成低精度，再往下传。

这是 PyTorch「可扩展」的秘密。 想加一种新后端（比如一种新硬件、一种新数值类型）？你不用改 add 本身，只需为你的分发键注册一套内核，再挂进这条分发链。autograd、autocast、量化、稀疏、meta 设备（只算形状不算数）……全是用同一套「加一层分发键」的方式做出来的。一个机制，撑起了整个生态的扩展。

源码在哪

dispatcher 的核心在 aten/src/ATen/core/dispatch/（Dispatcher.h、DispatchKeySet.h）。
分发键的定义在 c10/core/DispatchKey.h——你能在里面看到 CPU、CUDA、Autograd、Autocast、Sparse 等一长串键，以及它们的优先级顺序。
内核怎么注册进分发表，靠的是 TORCH_LIBRARY / TORCH_LIBRARY_IMPL 宏，以及下一章要讲的 native_functions.yaml。

dispatcher 负责把请求路由到正确的内核，但那个内核——真正在 CPU/GPU 上执行加法的代码——长什么样、写在哪、怎么同时供 CPU 和 CUDA 两套？这就是下一章。我们看看 PyTorch 的算子库 ATen、那个「算子花名册」native_functions.yaml，以及一份内核怎么服务多种数据类型。👉 算子与内核：ATen 与 native_functions.yaml

参考文献

[1] Edward Z. Yang. “Let's talk about the PyTorch dispatcher.” 2020. http://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/ —— dispatcher 与分发键的权威讲解，本章主要依据。
[2] PyTorch 文档：“Extending dispatcher for a new backend in C++.” https://pytorch.org/tutorials/advanced/extend_dispatcher.html —— 分发键与后端扩展。
[3] Paszke et al. “PyTorch: An Imperative Style, High-Performance Deep Learning Library.” NeurIPS 2019. arXiv:1912.01703.
源码：aten/src/ATen/core/dispatch/Dispatcher.h、c10/core/DispatchKey.h(pytorch/pytorch)。

一次 a + b 的旅程 · dispatcher 分发器