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

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

- 作者：David（道雾轩）
- 专栏：PyTorch 源码漫游（https://daiw.org/manual/pytorch-internals.md）
- 最后更新：2026-10-01
- 原文：https://daiw.org/manual/pytorch-internals/dispatch-journey
- 转载与引用：请注明出处并附原文链接（https://daiw.org/about/copyright）

# 一次 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` 走：

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

```mermaid
flowchart TD
  py["Python: a + b"] --> aten["C++: at::add"]
  aten --> disp{"dispatcher<br/>算出 DispatchKeySet"}
  disp -->|"最高优先级键"| k1["查分发表 → 对应内核"]
  k1 --> run["执行内核"]
```

## 精髓：分发键是「一层层」的，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 键这一层：进来的算子先被它拦下、按规则把输入转成低精度，再往下传。

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

## 源码在哪

- 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](https://daiw.org/manual/pytorch-internals/operators-and-kernels)

## 参考文献

- [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)。