PyTorch 训练框架内核深度解析

第1章 架构总览：一次张量加法的全旅程

作者 杨艺韬 · 14,128 字 · 发布于 · 更新于

"Most of these functionalities have a single handler assigned to them, making them runtime keys. ... A few functionalities are allowed to be customizable per backend."

—— c10/core/DispatchKey.h:162，本章第三层要拆的那颗心脏，它的设计说明就写在枚举定义里

本章要点

  • 一行 c = a + b 在 PyTorch 内部要穿过 七层抽象：Python 前端 → C++ 桥接 → Dispatcher → ATen 算子 → 后端 Kernel → Autograd → (可选) torch.compile
  • torch.Tensor 是 双层结构 —— Python 层提供灵活、torch._C.TensorBase 提供 C++ 性能，__add__ 这类魔法方法直接由 C 扩展类型的方法表提供
  • Dispatcher 是 PyTorch 的心脏：用 DispatchKeySet bitmap + redispatch 模式实现 Autograd / Functionalize / vmap / FakeTensor 等"洋葱式"中间层
  • ATen 算子靠 YAML + 代码生成：native_functions.yaml 与 derivatives.yaml 是 PyTorch 算子注册的"宪法"
  • 小张量上 CPU 侧的调度开销会盖过 GPU 上真正的算术时间，这正是 torch.compile 要消灭的目标 —— 大张量场景下调度开销可忽略
  • 全专栏 22 个后续章节都对应到本章铺出的源码地图上的某个具体位置

1.1 引子：一行最普通的代码

你写下这一行代码：

import torch

a = torch.randn(1024, 1024, device='cuda')
b = torch.randn(1024, 1024, device='cuda')
c = a + b

按下回车，几毫秒后 GPU 上多出了一个 1024×1024 的张量 c。（没有 GPU 的话把 device 改成 'cpu'：本章讲的 Python 前端、Dispatcher、ATen 这几层完全一样，只是最后落到 CPU kernel。）所有 AI 工程师每天都在写这种代码，几乎没人停下来问：这一行 a + b 到底发生了什么？

让我们用一整章来回答这个问题。这一章是全专栏的地图：把这一次加法走过的每一条路径标出来，后面 22 章就是把每一条路径分别走深。

flowchart TD
    P[Python 用户代码<br/>c = a + b] --> P1[Python 前端层<br/>Tensor.__add__]
    P1 --> P2[C++ 桥接层<br/>CPython C API 绑定 / THPVariable_add]
    P2 --> P3[Dispatcher 层<br/>at::_ops::add_Tensor 入口 / OperatorEntry]
    P3 --> P4[ATen 算子层<br/>structured_delegate / TensorIterator]
    P4 --> P5[后端 Kernel 层<br/>add_kernel @ CPU/CUDA，MPS 另有实现]
    P5 --> P6[内存与执行<br/>CUDA Caching Allocator + Stream]
    P3 -.记录反向图.-> AG[Autograd 层<br/>生成的 VariableType_N.cpp::add_Tensor]
    P3 -.编译路径.-> TC[torch.compile<br/>Dynamo+AOTAutograd+Inductor]
    style P3 fill:#fef3c7,stroke:#f59e0b,stroke-width:2px
    style P4 fill:#dbeafe,stroke:#3b82f6
    style AG fill:#fce7f3,stroke:#ec4899
    style TC fill:#dcfce7,stroke:#22c55e

七层下潜，每一层都对应后面的一个具体章节。读完这一章，你将拥有一张 PyTorch 的"心智地图"。

1.2 第一层：Python 前端的"魔法"

回到代码：

c = a + b

a + b 是 Python 的二元加法操作符。Python 解释器会把它翻译成对 a.__add__(b) 的调用。所以第一个问题是：a 是什么类型？它的 __add__ 怎么写的？

打开 Python 交互式：

>>> type(a)
<class 'torch.Tensor'>
>>> type(a).__add__
<method '__add__' of 'torch._C.TensorBase' objects>

注意两件事：

  1. a 的类型是 torch.Tensor —— 这是一个纯 Python 类
  2. 但它的 __add__ 是一个 C 实现的方法描述符，归属于 torch._C.TensorBase —— 这是一个用 C++ 实现的扩展类型

也就是说，torch.Tensor 是一个"双层结构"：

graph TB
    UC[用户看到的<br/>torch.Tensor]
    BC[C++ 扩展类型<br/>torch._C.TensorBase]
    UC -- 继承 --> BC
    UC --> M1[Python 实现的方法<br/>__repr__ / __format__ / 等]
    BC --> M2[C++ 实现的方法<br/>__add__ / matmul / sum / 等]
    style BC fill:#dbeafe,stroke:#3b82f6,stroke-width:2px

这个双层结构定义在 torch/_tensor.py:110：

# NB: If you add a new method to Tensor, you must update
# torch/_C/__init__.pyi.in to add a type annotation for your method;
# otherwise, it will not show up in autocomplete.
class Tensor(torch._C.TensorBase):
    _is_param: bool
    ...

为什么要分两层？因为：

  • C++ 层（torch._C.TensorBase）需要追求性能：每一个张量操作可能在毫秒甚至微秒级被调用，绝不能走 Python 的属性查找慢路径
  • Python 层（torch.Tensor）需要灵活：用户可以继承、可以 monkey patch、可以打 hook、可以注册 __torch_function__

Python 的 + 操作符在 CPython 里走 PyNumber_Add，查的是类型的 nb_add slot。torch._C.TensorBase 本身并没有填这个 slot（torch/csrc/autograd/python_variable.cpp 里 THPVariableType 的 tp_as_number 是 nullptr），而是把 __add__ 作为普通 C 方法放进方法表（生成模板 tools/autograd/templates/python_variable_methods.cpp 里的 variable_methods[]，指向 TypeError_to_NotImplemented_<THPVariable_add>）；torch.Tensor 作为 Python 子类被创建时，CPython 发现 MRO 上有 __add__，就给它的 nb_add 装上一个通用转发函数，调用时按名字找到这个 C 方法、直接跳进 C++。整条路上没有执行任何一行 Python 字节码，用户也感受不到这层转接 —— 这是 PyTorch 用户层友好与 C++ 层高性能并存的关键设计。

为什么这种设计如此重要

把"双层结构"拆开来想，会发现它其实在解一道框架设计师必须面对的两难题：你既想让框架对用户表现得"像 Python 一样灵活"，又必须让张量操作"快到接近 C++ 手写"。这两件事天然矛盾 —— Python 的属性查找、方法解析、对象包装都有可观开销，高频的张量操作承受不起。

PyTorch 的解法是 "让用户看到 Python 子类，让运行时看到 C 扩展类型"。具体做法是：

  • C 层用 CPython 的 PyTypeObject 接口注册 THPVariableType（Python 里叫 torch._C.TensorBase），把 __add__ / __mul__ 等魔法方法和大量普通方法都以 C 函数指针的形式挂进方法表，调用时不执行 Python 代码
  • Python 层用 class Tensor(torch._C.TensorBase) 包了一层，让用户能 isinstance(x, torch.Tensor)、能继承 torch.Tensor 写自定义子类、能写 torch.Tensor.foo = ... 打 monkey patch
  • 用户不知不觉地在使用一个 C 扩展类型，性能上享受 C 速度，体验上享受 Python 灵活

这个设计的代价是 PyTorch 源码里有大量"Python ↔ C++ 数据互转"的胶水代码 —— THPVariable_Wrap / THPVariable_Unpack / THPObjectPtr 这一类的工具函数。它们的存在不是冗余，而是这个"双层"哲学必须付出的工程税。第 6 章会把这部分胶水代码生成的来龙去脉摊开。

谁实现了 __add__

在 PyTorch 源码里搜 __add__ 的实现，会发现它不是手写的 —— 它是代码生成出来的。生成的源头是 tools/autograd/templates/python_variable_methods.cpp + tools/autograd/gen_python_functions.py，最终产物在编译目录里的 python_variable_methods.cpp。

简化看，生成出来的 __add__ C++ 代码大致是：

// 伪代码（省略了 PythonArgParser 解析、__torch_function__ 检查与释放 GIL）
static PyObject* THPVariable_add(PyObject* self_, PyObject* args, PyObject* kwargs) {
    HANDLE_TH_ERRORS
    Tensor& self = THPVariable_Unpack(self_);  // 把 PyObject 解出 C++ Tensor
    auto other = parse_tensor_arg(args, 0);    // 解参数
    auto alpha = parse_scalar_arg(args, 1, 1); // 默认 alpha=1
    return THPVariable_Wrap(self.add(other, alpha));  // 调 C++ 入口
    END_HANDLE_TH_ERRORS
}

注意最后一行：self.add(other, alpha) —— 这是 C++ 层的 Tensor::add 方法。Python 端的 __add__ 已经把活儿交给了 C++ 端的 Tensor::add。我们的旅程进入第二层。

延伸阅读：第 6 章会详细讲 PyTorch 的代码生成体系。tools/autograd/ 这个目录是 PyTorch 工程化的精华，把数千个算子的 Python 包装、autograd 反向、TorchScript 注册等等批量生成出来。

1.3 第二层：C++ 桥接 —— Python 绑定与 TensorBase

Tensor::add 这个 C++ 方法定义在哪里？通过 IDE 跳转或者 grep，会看到一份生成的头文件 TensorBody.h —— 源码树里只有模板 aten/src/ATen/templates/TensorBody.h（第 93 行就是下面这个类），真正的 aten/src/ATen/core/TensorBody.h 在构建目录里，内容形如：

class TORCH_API Tensor: public TensorBase {
    ...
    Tensor add(const Tensor& other, const Scalar& alpha = 1) const {
        return at::_ops::add_Tensor::call(*this, other, alpha);
    }
    ...
};

注意几点：

  1. Tensor 继承自 TensorBase（不是 Python 的 TensorBase，而是另一个 C++ 类，定义在 aten/src/ATen/core/TensorBase.h）
  2. Tensor::add 是一个薄包装，立刻调用了 at::_ops::add_Tensor::call
  3. at::_ops::add_Tensor 这个东西是哪里来的？也是代码生成的，由 aten/src/ATen/native/native_functions.yaml 通过 torchgen 工具生成

我们打开 native_functions.yaml，搜 add.Tensor：

- func: add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
  device_check: NoCheck   # TensorIterator
  structured_delegate: add.out
  variants: function, method
  dispatch:
    SparseCPU, SparseCUDA, SparseMPS, SparseMeta: add_sparse
    SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: add_sparse_csr
    MkldnnCPU: mkldnn_add
    ZeroTensor: add_zerotensor
    NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_add_Tensor
  tags: [core, pointwise]

——aten/src/ATen/native/native_functions.yaml 中 add.Tensor 条目

这一段 YAML 是 PyTorch 算子注册的灵魂。它声明了：

  • 这是一个名叫 add.Tensor 的算子，有一个 self 张量、一个 other 张量、一个 alpha 标量，返回一个张量
  • 它是 structured_delegate —— 把真正的实现委托给 add.out（一个带 out 参数的版本）
  • 可以作为函数（torch.add）或者方法（tensor.add）调用
  • 对几个特殊后端（Sparse、MKLDNN、嵌套张量），有专门的实现

torchgen 在编译时把这一段 YAML 翻译成 C++ 代码。其中 at::_ops::add_Tensor::call 是这样的：

// 编译后生成的代码（精简），路径：build/aten/src/ATen/Operators_2.cpp
// 模板见 torchgen/gen.py 的 ComputeOperators；按 root_name 哈希分片，add 落在第 2 片
Tensor add_Tensor::call(const Tensor& self, const Tensor& other, const Scalar& alpha) {
    // 原文拆成 create_add_Tensor_typed_handle() 一个辅助函数，这里合并展示
    static auto op = c10::Dispatcher::singleton()
        .findSchemaOrThrow("aten::add", "Tensor")   // (name, overload_name)
        .typed<Tensor(const Tensor&, const Tensor&, const Scalar&)>();
    return op.call(self, other, alpha);
}

——Dispatcher::singleton() 这一行是命运的转折点。我们的旅程从此进入 PyTorch 真正的"心脏"：Dispatcher。

1.4 第三层：Dispatcher —— PyTorch 的"中央调度"

所有 PyTorch 算子，最终都要经过一个全局的 c10::Dispatcher 单例，由它来决定这次调用应该走哪个具体实现。

为什么需要 Dispatcher

考虑 a + b 这次调用，下面所有问题都需要 Dispatcher 回答：

  1. a 是 CPU 张量还是 CUDA 张量？走 CPU 实现还是 CUDA kernel？
  2. a 是 float32 还是 bfloat16？走 fp32 算法还是 bf16 算法？
  3. 当前是不是在 loss.backward() 之前？要不要记录反向图？
  4. 当前是不是在 torch.no_grad() 上下文里？要不要跳过 autograd？
  5. 当前是不是在 torch.compile 的 trace 里？要不要走符号执行？
  6. 当前是不是在 vmap / functorch.grad 里？要不要做批处理变换？
  7. 当前是不是在量化场景？要不要走量化路径？
  8. a 是不是稀疏张量？要不要走 SparseCSR 实现？

每一个问题都对应一个 DispatchKey（分派键）。Dispatcher 收到一次调用时，会从张量的属性、当前线程上下文、全局状态里推导出一组 keys，然后在 keys 上按优先级排序，挑出最高优先级的 key 对应的实现，调过去。

Dispatcher 的核心数据结构

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

class TORCH_API Dispatcher final {
 private:
  struct OperatorDef final { impl::OperatorEntry op; /* ... */ };

  // 算子表本体：每个算子一个 OperatorDef（内含 OperatorEntry）
  std::list<OperatorDef> operators_;
  // 名字 → OperatorHandle 的查找表，读多写少，用 LeftRight 做无锁读
  LeftRight<ska::flat_hash_map<OperatorName, OperatorHandle>> operatorLookupTable_;
  // 后端 fallback 表：按运行时 dispatch 表下标索引
  std::array<impl::AnnotatedKernel, num_runtime_entries> backendFallbackKernels_;

 public:
  C10_ALWAYS_INLINE static Dispatcher& singleton();   // 全局单例

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

每一个算子（如 add.Tensor）在运行时对应一个 OperatorEntry，里面保存了它在不同 DispatchKey 下的实现：

// aten/src/ATen/core/dispatch/OperatorEntry.h:70
class TORCH_API OperatorEntry final {
    OperatorName name_;

    // 运行时真正用的分派表：下标由 DispatchKeySet 算出
    std::array<KernelFunction, c10::num_runtime_entries> dispatchTable_;

    // 按 key 查表，取出这次要调的 kernel（OperatorEntry.h:182）
    const KernelFunction& lookup(DispatchKeySet ks) const;
    ...
};

DispatchKeySet 是一个 64 位 bitmap：低几位表示后端（CPU、CUDA……），其余位表示"功能"（Dense、Autograd、Python……），两者组合才定出一个具体的 key（c10/core/DispatchKey.h 开头的注释讲了这套布局，第 3 章 §3.5 细拆）。当一个张量参与到调用里，PyTorch 会从张量自身的 key_set_ 里拿到它的 keys（普通 CPU 张量自带 CPU、ADInplaceOrView、AutogradCPU、AutocastCPU，与 requires_grad 无关），再并上线程局部的 included 集、减去 excluded 集（如是否在 vmap 里、是否在 inference_mode 里），得到本次调用的最终 DispatchKeySet。

一次调用的 Dispatcher 流程

sequenceDiagram
    autonumber
    participant U as 用户调用 a.add(b)
    participant D as Dispatcher::call()
    participant E as OperatorEntry
    participant K as KernelFunction
    participant I as 实际实现 (CUDA add_kernel)
    U->>D: at::_ops::add_Tensor::call(a, b, 1)
    D->>D: 收集 DispatchKeySet:<br/>张量 a/b 的 keys ∪ 线程上下文 keys
    D->>E: lookup(keySet)
    E->>E: 取最高优先级的 key (如 AutogradCUDA)
    E->>D: 返回 KernelFunction 指针
    D->>K: invoke(boxed_args)
    K->>I: VariableType::add_Tensor(ks, a, b, 1)<br/>记录反向图、redispatch 到 CUDA
    I->>D: 重新进 dispatcher<br/>这次去掉 Autograd key
    D->>K: 取 CUDA key 的 KernelFunction
    K->>I: 生成的 ufunc_add_CUDA → add_kernel<br/>真正的 CUDA kernel 调用
    I-->>U: 返回结果张量

注意第 7-9 步：第一次进 dispatcher 命中的不是 CUDA kernel，而是 Autograd key。Autograd 的实现做完"记录反向图"的工作后，自己又重新进了一次 dispatcher（这叫 redispatch），这次的 keySet 已经把 Autograd key 移除，于是命中 CUDA kernel。

这种"递归 dispatch"的设计极其精妙：

  • Autograd、Functionalization、AMP、vmap 都是用同样的 redispatch 模式实现的"中间层"
  • 每一层只关心自己的事，做完就 redispatch 到下一层
  • 整个调用链像剥洋葱：外面的层先做、剥掉、再交给里面的层

延伸阅读：第 5 章会详细讲 Dispatcher 的实现细节，包括 DispatchKeySet 的位序、key 的优先级、boxed/unboxed 调用、kernel 注册的几种方式（TORCH_LIBRARY / TORCH_LIBRARY_IMPL / Library::impl）。

Dispatcher 与本系列其他专栏的对照

如果你读过本系列的《Tokio 源码深度解析》或《Rust 编译器与运行时揭秘》专栏，这种"通过 key 选实现"的模式应该不陌生：

系统 类似机制 关键点
PyTorch Dispatcher DispatchKeySet → KernelFunction 张量属性 + 线程状态决定
Rust Trait Resolver trait + 类型 → impl 编译期决定
Tokio Runtime TaskWaker → Future poll 运行时决定
C++ 虚函数 this 指针 → vtable 运行时单分派

PyTorch 的 dispatcher 是多分派 + 优先级 + bitmap 的组合，比 C++ 单分派强大得多。它是 PyTorch 能在不破坏 API 的前提下不断引入新特性（FunctionalTensor、Lazy Tensor、FakeTensor……）的根本原因 —— 每一个新特性都是新增一个 DispatchKey，而不是改算子的实现签名。

Dispatcher 的演进史

值得停下来回顾的一段历史：PyTorch 早期（1.0 之前）并没有真正意义上的"通用 dispatcher"。那时算子的分派散落在各处 —— 有的用 if (tensor.is_cuda()) 直接判断，有的用 Type 抽象（一个早已废弃的旧概念）。autograd 也是一个独立的 wrapper 层，包在每个算子外面。

这种设计在 PyTorch 1.0 之后开始难以维持：

  • 特性数爆炸：从 fp32/fp64 两种到 fp32/fp64/fp16/bf16/int8/int4 多种，从 CPU/CUDA 到 CPU/CUDA/XPU/MPS/HPU/TPU 多种后端，从纯 eager 到 functorch/vmap/grad/jvp 等函数变换 —— 排列组合呈爆炸增长
  • 特性正交化失败：每加一个新特性都要在每个算子里手动加一段 if-else，复用极差，修一个 bug 往往要改一大片文件
  • 后端供应商压力：NVIDIA、AMD、Intel、华为都希望把自家硬件接进 PyTorch，但每次合并都要动核心代码

1.x 时期那次 dispatcher 重构就是为了解决这个问题：把所有这些"特性维度"统一成 DispatchKey，用 bitmap + 优先级做组合，functionality 与 backend 拆成两个正交轴（第 3 章 §3.5 会拆这个布局）。这是 PyTorch 历史上最重要的一次内部重构之一，没有它就没有今天的 functorch、torch.compile、FakeTensor 等生态。

这次重构留下的注释是源码里最珍贵的资料。打开 c10/core/DispatchKey.h（v2.11 下全文 745 行），你会发现 enum class DispatchKey 要到第 136 行才开始 —— 前面一百多行、以及枚举项之间穿插的大段注释，把 DispatchKey 的设计原则、key 的优先级排序、为什么某些 key 要放在某个位置全部讲清楚了。这份注释就是 dispatcher 的设计文档，比任何博客都权威。第 3 章 §3.5 拆它的位布局，第 5 章拆它在查表时怎么用。

1.5 第四层：ATen 算子层 —— structured_delegate 与 TensorIterator

经过 Dispatcher 后，调用最终落到 ATen 层的真正实现。回到 add.Tensor 的 YAML：

- func: add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
  structured_delegate: add.out

structured_delegate: add.out 这句话告诉 torchgen：add.Tensor 自己不实现，去找 add.out 的实现，分配输出张量后调它。这是 PyTorch 用来减少代码重复的机制 —— 一个数学操作通常有 out= 版本（add(a, b, out=c)）和无 out 版本（add(a, b)）；structured_delegate 让无 out 版本自动委托给 out 版本，省掉重复实现。

add.out 的真正实现

继续在 YAML 里搜 add.out：

- func: add.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!)
  device_check: NoCheck   # TensorIterator
  structured: True
  structured_inherits: TensorIteratorBase
  ufunc_inner_loop:
    Generic: add (AllAndComplex, BFloat16, Half, ComplexHalf)
    ScalarOnly: add (Bool)
  dispatch:
    SparseCPU, SparseMeta: add_out_sparse_cpu
    SparseCUDA: add_out_sparse_cuda
    SparseMPS: add_out_sparse_mps
    SparseCsrCPU, SparseCsrMeta: add_out_sparse_compressed_cpu
    SparseCsrCUDA: add_out_sparse_compressed_cuda
    MkldnnCPU: mkldnn_add_out
    MPS: add_out_mps
    MTIA: add_out_mtia
  tags: pointwise

——aten/src/ATen/native/native_functions.yaml:577 的 add.out 条目

这次的 YAML 揭示了几件大事：

  • structured: True —— 这是一个"结构化算子"，遵循 PyTorch 的现代算子框架
  • structured_inherits: TensorIteratorBase —— 它使用 TensorIterator 来处理"逐元素操作"
  • ufunc_inner_loop —— 它的内层循环是一个 ufunc，针对各种数据类型（fp32/bf16/fp16/complex）由代码生成

TensorIterator：逐元素操作的"引擎"

PyTorch 中所有逐元素操作（add、sub、mul、div、relu、sigmoid、log、…）都建立在 TensorIterator 这套基础设施之上。它的工作是：

  1. 接收输入张量和输出张量
  2. 处理广播：把 [1024, 1] + [1, 1024] 自动扩展到 [1024, 1024]
  3. 处理类型提升：bf16 + fp32 → 提升到 fp32
  4. 处理内存连续性：判断是不是 contiguous，要不要 reshape
  5. 把所有"配置好的迭代"交给底层的 kernel

简化的调用链：

// aten/src/ATen/native/BinaryOps.cpp:434 —— sub.out 的结构化实现
TORCH_IMPL_FUNC(sub_out) (
    const Tensor& self, const Tensor& other, const Scalar& alpha, const Tensor& result
) {
    add_stub(device_type(), *this, -alpha);
    TORCH_INTERNAL_ASSERT(result.scalar_type() == output().dtype());
}

这里特意拿 sub_out 举例：add.out 因为声明了 ufunc_inner_loop，它的 TORCH_IMPL_FUNC(ufunc_add_CPU) / TORCH_IMPL_FUNC(ufunc_add_CUDA) 是 torchgen 生成到构建目录的 UfuncCPU_add.cpp / UfuncCUDA_add.cu 里的，源码树里手写不到；而 sub.out 手写在 BinaryOps.cpp 里，形状完全一样，正好当样板看。

*this 是 TensorIteratorBase&，已经把广播、类型提升、连续性都处理好了。剩下的工作是调 add_stub —— 这是一个 dispatch stub，会按 device_type 分派到 CPU 或者 CUDA 的真正 kernel。

ufunc 内层循环：数学公式在哪里

最后，"加法到底是怎么算的"这件事本身，藏在 aten/src/ATen/native/ufunc/add.h:14：

namespace at::native::ufunc {

template <typename T>
C10_HOST_DEVICE C10_ALWAYS_INLINE T add(T self, T other, T alpha)
    __ubsan_ignore_undefined__ {
    return self + alpha * other;
}

#if !defined(__CUDACC__) && !defined(__HIPCC__)
template <typename T>
C10_ALWAYS_INLINE Vectorized<T> add(Vectorized<T> self, Vectorized<T> other,
                                     Vectorized<T> alpha) {
    return vec::fmadd(other, alpha, self);
}
#endif

} // namespace at::native::ufunc

——三件事值得停下来看一遍：

  1. return self + alpha * other —— 这就是 PyTorch add 的数学定义。alpha 是默认 1 的缩放系数，所以 a + b 实际上是 a + 1 * b
  2. C10_HOST_DEVICE —— 这个宏让函数同时能在 CPU 和 GPU（CUDA / HIP）上编译
  3. CPU 路径用 vec::fmadd（fused multiply-add） —— Vectorized<T> 本身就是按 AVX2/AVX512/NEON 等指令集分别实现的 SIMD 封装，这一版重载直接对整条向量做乘加

这一段不到 10 行的代码，会被代码生成器在编译时模板展开成几十个版本：每个 dtype（fp32/fp16/bf16/i32/i64/complex）一份，CPU 与 CUDA 各一套（MPS 不走 ufunc，用 YAML 里单独登记的 add_out_mps）。CPU 侧与 CUDA 侧的 add kernel 都是 ufunc codegen 生成到构建目录的（UfuncCPUKernel_add.cpp / UfuncCUDA_add.cu），源码树里没有对应文件 —— 比如 aten/src/ATen/native/cpu/BinaryOpsKernel.cpp 里只有 add_clamp 之类的邻居，没有 add 本身；数学定义就是上面的 aten/src/ATen/native/ufunc/add.h。

广播与类型提升：在你看不见的地方

回到一个真实场景：你写下 a + b，但 a.shape == [1024, 1]、b.shape == [1, 1024]、a.dtype == bf16、b.dtype == fp32。这次加法会发生什么？

正确答案是：

  1. 形状广播：[1024, 1] 与 [1, 1024] 广播为 [1024, 1024]
  2. 类型提升：bf16 与 fp32 提升为 fp32
  3. 输出 dtype：c.dtype == fp32，shape 是 [1024, 1024]
  4. 内存：a 和 b 不会被复制，只是 strides 被重新解释成"广播视图"（这是 CUDA 上的行为：kernel 边读边转 dtype；在 CPU 上，compute_types 会先把 bf16 的 a 转成一份 fp32 临时张量）
  5. 输出：分配一块新的 1024 × 1024 × 4 = 4MB fp32 显存做 c

这一切都不是用户写的，全部由 TensorIterator 在背后完成。它用一个内部状态机：

// aten/src/ATen/TensorIterator.cpp:1493 TensorIteratorBase::build 的骨架（精简）
void TensorIteratorBase::build(TensorIteratorConfig& config) {
    populate_operands(config);    // 收集输入/输出 operand
    mark_outputs();               // 标出哪些 operand 是输出、哪些读写兼有
    compute_mem_overlaps(config); // 检查输出与输入是否内存重叠
    compute_names(config);        // 命名张量的名字传播
    compute_shape(config);        // 算广播后的形状
    mark_resize_outputs(config);  // 标记需要 resize 的输出
    compute_types(config);        // 类型提升 (promote_types) 与公共 device
    if (!fast_set_up(config)) {   // 连续、同形状的快路径能直接跳过下面几步
        compute_strides(config);      // 把所有输入的 strides 调整到广播形状
        reorder_dimensions();         // 重排维度以改善合并
        allocate_or_resize_outputs(); // 分配 c
        if (!is_meta_) coalesce_dimensions(); // 合并可合并的维度以加速循环
    }
    // ... 取 data_ptr、初始化 view_offsets_
}

每一步都是一次小型的"算子前奏"。理解 TensorIterator 就理解了 PyTorch 逐元素操作的"通用前置流水线"。"广播为什么不复制内存"在第 2 章 §2.7 用 stride = 0 解释；TensorIterator 自身的状态机本专栏不再单独展开，想深挖可以直接读 aten/src/ATen/TensorIterator.cpp。

延伸阅读：第 6 章会拆开整个 ATen 代码生成体系，理解为什么 PyTorch 用 YAML 而不是 C++ 模板来声明算子。

1.6 第五层：后端 Kernel —— CUDA Stream 与显存

add.out 在 CUDA 上最终落到 ufunc codegen 生成的具体 kernel：

// add 的 CUDA kernel 是**生成**出来的：ufunc codegen（torchgen/dest/ufunc.py 的
// compute_ufunc_cuda）读 aten/src/ATen/native/ufunc/add.h（下方 ufunc::add 就在那里），
// 生成到构建目录的 UfuncCUDA_add.cu —— 源码树里没有这个文件。
// 下面是生成结果的精简形态。
template <typename scalar_t>
struct CUDAFunctor_add {   // 另有 CUDAFunctorOnSelf_add / CUDAFunctorOnOther_add，处理一侧是 CPU 标量的情况
    using opmath_t = at::opmath_type<scalar_t>;
    opmath_t alpha_;
    __device__ scalar_t operator()(scalar_t self, scalar_t other) {
        return ufunc::add(static_cast<opmath_t>(self), static_cast<opmath_t>(other), alpha_);
    }
};

void add_kernel(TensorIteratorBase& iter, const at::Scalar& alpha) {
    AT_DISPATCH_SWITCH(iter.common_dtype(), "ufunc_add_CUDA", /* 每个 dtype 一个 case */
        AT_DISPATCH_CASE(at::ScalarType::Float, [&]() {
            using opmath_t = at::opmath_type<scalar_t>;
            gpu_kernel(iter, CUDAFunctor_add<scalar_t>(alpha.to<opmath_t>()));
        })
        /* ... Double / Half / BFloat16 / Complex ... */
    );
}

// 登记进 add_stub 的 CUDA 表项，供 sub_out 这类手写算子经 add_stub 找到它
REGISTER_DISPATCH(add_stub, &add_kernel)

// 结构化实现本体，同样由 codegen 生成；CUDA 上直接调 add_kernel，不再经 add_stub 二次分派
TORCH_IMPL_FUNC(ufunc_add_CUDA)(const Tensor& self, const Tensor& other,
                                const Scalar& alpha, const Tensor& out) {
    add_kernel(*this, alpha);
}

gpu_kernel 是 PyTorch CUDA 逐元素算子的模板基础设施，它做的事是：

  1. 算 launch config：根据元素总数算出 grid 和 block 大小
  2. launch CUDA kernel：把 functor 实例化进一个模板化的 __global__ kernel
  3. 写入当前 CUDA stream：每一个 launch 都排进当前线程的"当前 stream"（没切换过就是 default stream）

gpu_kernel 定义在 aten/src/ATen/native/cuda/Loops.cuh，真正发射 kernel 的代码在 aten/src/ATen/native/cuda/CUDALoops.cuh，大致是：

// 示意：CUDALoops.cuh 中 launch_vectorized_kernel 的骨架（省略了向量宽度选择与 ROCm 分支）
template <typename func_t, typename array_t>
static inline void launch_vectorized_kernel(int64_t N, const func_t& f, array_t data) {
    auto stream = at::cuda::getCurrentCUDAStream();
    int bws = tws * num_threads();              // 每个 block 处理的元素数
    int64_t grid = (N + bws - 1) / bws;
    vectorized_elementwise_kernel<vec_size, func_t, array_t>
        <<<grid, num_threads(), 0, stream>>>(N, f, data);
}

注意 getCurrentCUDAStream() —— 这是 PyTorch 异步执行的关键。所有 kernel launch 都进了同一个 stream（默认情况下），CUDA 保证 stream 内的指令按顺序执行，但与 CPU 是异步的。这就是为什么 c = a + b 在 Python 端立即返回，而不是等 GPU 算完。

显存从哪里来

a + b 要产生新张量 c，需要在 GPU 上分配一块 1024 × 1024 × 4 = 4MB 的显存。这次分配不会调 cudaMalloc —— 它走的是 PyTorch 自己的 CUDA Caching Allocator：

// c10/cuda/CUDACachingAllocator.cpp（极简示意，非源码原文）
// 张量分配走 NativeCachingAllocator::allocate → DeviceCachingAllocator::malloc
Block* DeviceCachingAllocator::malloc(size_t orig_size, cudaStream_t stream) {
    // 1. 按大小选 small / large 池，在"同一 stream"的空闲 block 里 best-fit 查找
    // 2. 找不到才调 cudaMalloc 申请新 segment（失败时先释放缓存再重试）
    // 3. block 记下所属 stream；张量释放后 block 回到池子，留给同一 stream 复用
}

PyTorch 设计这个 caching allocator 的原因是：cudaMalloc / cudaFree 是昂贵的驱动调用，还会带来隐式同步，每个张量都直接调它们性能极差。Caching allocator 把 free 的 block 留在用户态池子里循环利用，把这部分开销压到极低。

延伸阅读：第 4 章会专门拆解 CUDA Caching Allocator —— 它的 block 管理、流绑定、defragmentation、expandable_segments 模式、以及为什么 torch.cuda.empty_cache() 常常不如预期。

CUDA Stream：异步执行的全部秘密

a + b 在 Python 端瞬间返回，但实际上 GPU 此刻还在干活。这背后是 PyTorch 的 异步执行模型 —— 每一次 kernel launch 都是把"指令"压进 CUDA stream 队列，CPU 立刻继续执行后面的代码。

这就引出几个常被混淆的问题：

问题 1：怎么知道结果什么时候真正算完？

PyTorch 不需要你显式查询。下一次你访问 c.cpu()、c.item()、c.numpy()、print(c) 时，PyTorch 会自动同步当前 stream（阻塞式的 D2H 拷贝走 c10/cuda/CUDAFunctions.h 的 memcpy_and_sync：cudaMemcpyAsync 之后显式 cudaStreamSynchronize）。如果你只是 c = a + b 然后又 d = c * 2，这两个 kernel 都进同一个 stream，CUDA 保证按顺序执行 —— 你不需要也不应该手动 torch.cuda.synchronize()。

问题 2：那什么时候才需要 synchronize()？

最常见的是两种情况：

  1. 跨 stream：你显式用 with torch.cuda.stream(s): 切到了非默认 stream，要等它的工作完成
  2. 想测时间：你要测 GPU 上某段代码的 wall clock time，必须用 cuda.Event 或者 sync 后用 time.time()

把这两条记牢，能避免大多数"不必要的 synchronize"。

问题 3：异步执行带来什么风险？

最经典的坑是 错误延迟报告。如果你的 CUDA kernel 越界访问了显存，错误不会在那次 kernel launch 时报，而是在下一次 sync 时（可能是几十次操作之后）才报，报错的栈指向的往往是一个无辜的后续调用 —— PyTorch 自己也在错误信息后面附了一句提示："CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect."（c10/cuda/CUDAMiscFunctions.cpp）。调试时需要设 CUDA_LAUNCH_BLOCKING=1 强制每次 launch 都同步，才能定位到真正出错的 kernel。这是 PyTorch / CUDA 用户每年都会被坑一次的经典陷阱。

1.7 第六层：Autograd 的"暗影"

我们前面有意忽略了一件事：a + b 如果 a.requires_grad=True，那 c 也会有反向图。这次 dispatch 实际上走的是 AutogradCUDA → CUDA 的两段式：

sequenceDiagram
    autonumber
    participant U as 用户代码
    participant D1 as Dispatcher (第一次)
    participant V as VariableType::add_Tensor
    participant D2 as Dispatcher (redispatch)
    participant K as CUDA add_kernel
    U->>D1: at::add(a, b, 1)
    D1->>D1: keySet 里有 AutogradCUDA + CUDA
    D1->>V: 命中最高优先级 AutogradCUDA
    V->>V: 创建 AddBackward 节点<br/>记录 grad_fn
    V->>D2: at::redispatch::add(<br/>excluding AutogradCUDA)
    D2->>K: 命中 CUDA key, 调真正 kernel
    K-->>V: 返回 c
    V->>V: 把 c.grad_fn = AddBackward
    V-->>U: 返回带反向图的 c

VariableType::add_Tensor 这一层是自动生成的，源头是 tools/autograd/derivatives.yaml：

- name: add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
  self: handle_r_to_c(self.scalar_type(), grad)
  other: handle_r_to_c(other.scalar_type(), maybe_multiply(grad, alpha.conj()))

——tools/autograd/derivatives.yaml 中 add.Tensor 条目

这一段 YAML 描述了 add 的反向规则：d(add)/d(self) = grad，d(add)/d(other) = grad × alpha。tools/autograd/gen_autograd_functions.py 会把它翻译成 C++ 反向函数：

// 编译生成的 torch/csrc/autograd/generated/Functions.cpp 简化版
struct AddBackward0 : public Node {
    variable_list apply(variable_list&& grads) override {
        auto& grad = grads[0];
        return { grad, maybe_multiply(grad, alpha.conj()) };
    }
    Scalar alpha;
};

当用户最后调用 loss.backward() 时，Autograd Engine 会从 loss.grad_fn 出发，沿着这些 AddBackward0 节点反向遍历，调每个节点的 apply 计算梯度。

延伸阅读：第 7、8 章是本专栏 Autograd 的两章，分别讲反向图怎么建、Engine 怎么用"每个非 CPU 设备一条常驻 worker 线程、CPU 上的活由调用 backward 的线程自己跑，各配一个 ReadyQueue 优先级队列"跑后向传播。autograd 是 PyTorch 区别于 TensorFlow 静态图的核心所在。

"记录反向图"是怎么记录的

有人会问：每一个张量加法都被 autograd 包了一层，那训练时这层包装的开销大不大？答案是：比你想象的小，但也不能忽视。

具体的开销由几部分组成（按量级从小到大）：

  • 一次 DispatchKeySet 计算 —— 几条位运算
  • 一次 dispatcher 查表 + 函数指针跳转
  • 一次 redispatch（再走一遍上面两步）
  • 把 AddBackward0 节点挂到 c.grad_fn 上 —— 若干指针赋值
  • 创建 AddBackward0 节点本身 —— 这一步最贵，因为含一次堆分配

这些都是常数开销：和张量多大无关。对一个百万元素的 GEMM（GPU 上几毫秒），它可以忽略不计；但如果代码里全是小张量逐元素操作，这份常数开销就会浮到 profile 的顶部。想量化它，用 torch.profiler 或 timeit 在你自己的机器上对比 requires_grad=True / False 两种情况 —— 具体数值和 CPU 型号、Python 版本强相关，不存在一个通用数字。

这就是为什么 PyTorch 提供了 torch.no_grad() 和 torch.inference_mode() 两层"关闭 autograd"的接口：

  • no_grad：仍然走 autograd dispatch，但跳过反向图构建（省下后两步）
  • inference_mode：完全绕过 autograd dispatch（连 dispatcher 那次跳转都省了），把张量标记为"不会再被反向求导"

在小算子密集的推理循环里，inference_mode 通常比 no_grad 更快一档；具体差多少要在自己的负载上测。第 5 章会拆开这两个 mode 的实现差异。

1.8 第七层：torch.compile 的"另一条路"

到这里，传统 eager 模式的旅程就走完了。但 PyTorch 2.x 提供了第二条路径。如果用户改写成：

@torch.compile
def add_fn(a, b):
    return a + b

c = add_fn(a, b)

这次的旅程完全不同：

flowchart TD
    U["add_fn(a, b)"] --> D[CPython 帧拦截<br/>PEP 523 frame eval]
    D --> Dy[TorchDynamo<br/>字节码逐条分析]
    Dy --> FX[FX Graph<br/>+ Guards<br/>+ Output Code]
    FX --> AOT[AOTAutograd<br/>functionalize + partition]
    AOT --> IND[TorchInductor<br/>Lowering → Triton DSL]
    IND --> Tri[Triton 编译<br/>→ PTX → CUBIN]
    Tri --> Run[运行编译后的 kernel<br/>命中即跳过原函数的逐条解释]
    style D fill:#fef3c7
    style Dy fill:#fef3c7,stroke:#f59e0b,stroke-width:2px
    style AOT fill:#dbeafe,stroke:#3b82f6,stroke-width:2px
    style IND fill:#dcfce7,stroke:#22c55e,stroke-width:2px

第一次调用时，TorchDynamo 通过 CPython 的 PEP 523 帧求值钩子（3.6 起可用，v2.11 要求 Python ≥ 3.10）拦截到 add_fn 的字节码，逐条解析，把每个 Python 操作翻译成 FX Graph 的节点；同时记录一组 guards（对输入张量的形状、dtype、device 的假设）。

然后 AOTAutograd 接手，把 FX Graph 传给一个"假"前向，跟踪算子的 functionalized 形式，再用 min-cut partitioning 把正向和反向切分成两个图。

最后 TorchInductor 把 ATen IR 一路 lower 到 Triton DSL，再用 OpenAI 的 Triton 编译器生成 PTX → CUBIN 二进制。

第二次调用时（如果 guards 都成立）：不再逐条解释原函数的字节码，而是执行 Inductor 生成的包装代码、直接 launch 编译好的 kernel。省下的是逐算子的 Python 解释、dispatcher 查找和多次 kernel launch —— 提速幅度完全取决于负载里"小算子占比"有多高，官方的 TorchBench / HuggingFace / TIMM 三套基准就是用来量化这件事的。

延伸阅读：第 12-15 章是本专栏的"编译器栈四章"，分别拆 Dynamo、AOTAutograd、Inductor、CUDA Graph 协同。这是 PyTorch 自 2.0 以来的主线方向，也是本专栏的技术高峰。

两条路径并不冲突，而是互补

很多人误以为 torch.compile 是"另一个 PyTorch"，是和 eager 模式互斥的选择。事实正好相反：torch.compile 是建立在 eager 模式之上的、可选的加速层。

具体来说，torch.compile 在内部仍然依赖 dispatcher 与 ATen 算子（trace 时靠它们得到 ATen 级的图，GEMM、SDPA 这类 extern kernel 运行时也仍调 ATen 实现）、仍然用 CUDA Caching Allocator —— 它做的事是把"用户代码 → 算子调用"这一段的开销拿掉（减少 Python 解释、减少 dispatcher 查找），并把逐元素类小算子融合成新生成的 Triton kernel、减少 kernel launch。

这意味着：

  • 你可以 混用 @torch.compile 装饰的函数和普通 PyTorch 代码，张量在两者间自由流动
  • 遇到 graph break 时，断点处的代码会回到 eager 执行 —— 用户代码不需要改
  • ATen 算子的优化（如 FlashAttention-2 的引入）自动同时惠及 eager 和 compile 两条路径
  • 自定义算子（用 TORCH_LIBRARY 注册）也能被 torch.compile 捕获进图（前提是注册了 fake/meta 函数）—— 但默认是不透明的黑盒，编译器不会 inline 进它的函数体、也不会把它与周围算子 fuse（第 22 章 §22.3）

理解这一点很重要：你不需要"二选一"，而是把 torch.compile 当作一种可选的加速，写代码时按 eager 思路写，跑得慢的地方加 @torch.compile 试试 —— 心智负担很小。

1.9 整张地图：把七层放在一起看

把这一整章的内容压缩成一张图：

graph TB
    subgraph Python["Python 层"]
        PY[c = a + b]
        PT[torch.Tensor 类<br/>双层结构]
    end
    subgraph Bind["C++ 桥接"]
        PB[THPVariable_add / PythonArgParser<br/>解 PyObject → Tensor]
        TB[at::_ops::add_Tensor::call]
    end
    subgraph Core["Dispatcher + ATen"]
        DP[Dispatcher 单例<br/>DispatchKeySet 决策]
        AT[structured add.out<br/>ufunc_add_CPU/CUDA + TensorIterator]
        UF[ufunc::add<br/>self + alpha * other]
    end
    subgraph Backend["后端 Kernel"]
        CU[CUDA add_kernel（生成）<br/>gpu_kernel + Stream]
        CP[CPU add_kernel（生成）<br/>SIMD vec::fmadd]
        MP[MPS / XLA / 其他]
    end
    subgraph Side["副作用层"]
        AG[Autograd<br/>记录反向图]
        AL[CUDA Caching Allocator<br/>分配 c 的显存]
    end
    subgraph Compile["编译路径(可选)"]
        TC[torch.compile]
        DY[Dynamo + FX]
        AO[AOTAutograd]
        IN[Inductor → Triton]
    end
    PY --> PT
    PT --> PB
    PB --> TB
    TB --> DP
    DP --> AG
    AG --> DP
    DP --> AT
    AT --> UF
    UF --> CU
    UF --> CP
    UF --> MP
    CU --> AL
    PY -. 加上编译装饰器 .-> TC
    TC --> DY
    DY --> AO
    AO --> IN
    IN --> CU
    style DP fill:#fef3c7,stroke:#f59e0b,stroke-width:3px
    style AT fill:#dbeafe,stroke:#3b82f6,stroke-width:2px
    style AG fill:#fce7f3,stroke:#ec4899,stroke-width:2px
    style AL fill:#fce7f3,stroke:#ec4899,stroke-width:2px
    style IN fill:#dcfce7,stroke:#22c55e,stroke-width:2px

对应到本专栏章节：

层 本章节 后续展开
Python 前端 §1.2 第 9 章 nn.Module
C++ 桥接 / 代码生成 §1.3 第 6 章 ATen 代码生成
Dispatcher §1.4 第 5 章 Dispatcher 详解
ATen / TensorIterator §1.5 第 6 章 算子注册
后端 Kernel §1.6 第 4 章 Caching Allocator
Tensor / Storage (未展开) 第 2-3 章 张量层
Autograd §1.7 第 7-8 章 Autograd 详解
torch.compile §1.8 第 12-15 章 编译器栈

1.10 PyTorch 源码地图：从层到目录

把上面七层架构对应到 PyTorch 源码的目录结构，能得到一份导航图。后面所有章节都会反复引用这张图：

pytorch/
├── c10/                         # 核心抽象层（Caffe2 + ATen 的并集名）
│   ├── core/                    #   Tensor / Storage / Device / DType / DispatchKey
│   ├── cuda/                    #   CUDA Caching Allocator、Stream、Event
│   ├── util/                    #   工具：intrusive_ptr、SmallVector、ArrayRef
│   └── macros/                  #   平台相关宏：C10_HOST_DEVICE / TORCH_API（v2.11 的定义实体在 torch/headeronly/macros/，这里只剩转发头）
│
├── aten/src/ATen/               # 算子库（A Tensor library）
│   ├── core/                    #   Dispatcher、OperatorEntry、Library 注册接口
│   ├── native/                  #   算子的"原生"实现（CPU 与设备无关版本）
│   │   ├── *.cpp                #     算子的入口函数（structured impl 与 dispatch stub）
│   │   ├── *.h                  #     算子声明、stub 声明
│   │   ├── cpu/                 #     CPU SIMD kernel 实现
│   │   ├── cuda/                #     CUDA kernel 实现
│   │   ├── mps/                 #     Apple Silicon MPS 实现
│   │   ├── sparse/              #     稀疏张量实现
│   │   ├── nested/              #     嵌套张量实现
│   │   └── ufunc/               #     ufunc 内层循环（数学公式）
│   └── native/native_functions.yaml   # 算子声明的 "宪法"
│
├── torch/                       # Python 前端
│   ├── csrc/                    #   Python ↔ C++ 桥接 (THPVariable 等)
│   │   ├── autograd/            #     C++ autograd 引擎
│   │   ├── api/                 #     C++ Frontend (libtorch)
│   │   ├── distributed/         #     ProcessGroup / NCCL 集成
│   │   └── jit/                 #     旧 TorchScript（2.5 起标记 deprecated）
│   ├── nn/                      #   nn.Module、各种 Layer
│   ├── optim/                   #   优化器
│   ├── utils/data/              #   DataLoader / Dataset
│   ├── distributed/             #   DDP / FSDP / DTensor（Python 侧）
│   ├── _dynamo/                 #   TorchDynamo（CPython 帧拦截）
│   ├── _functorch/              #   AOTAutograd / vmap / functional transforms
│   ├── _inductor/               #   TorchInductor（编译器后端）
│   └── ao/                      #   量化（eager / FX graph mode，整体已 deprecated；PT2E 已迁 torchao）
│
├── torchgen/                    # 算子代码生成器（读 native_functions.yaml 生 C++）
│   ├── gen.py                   #   生成入口：Operators / RegisterXXX / TensorBody 等
│   ├── api/、dest/              #   各种目标代码的签名与生成规则（含 ufunc）
│   └── packaged/                #   打包进 wheel 的 YAML 副本（构建时由 setup.py 拷入，源码树里没有）
│
└── tools/                       # 构建与代码生成工具
    └── autograd/                #   autograd 与 Python 包装的代码生成器
        ├── derivatives.yaml     #     反向规则
        ├── templates/           #     VariableType.cpp / python_variable_methods.cpp 等模板
        ├── gen_autograd_functions.py
        └── gen_python_functions.py

这张图把全专栏的章节定位都标好了：

章节 主要涉及目录
第 2-3 章（张量层） c10/core/
第 4 章（CUDA Allocator） c10/cuda/CUDACachingAllocator.cpp
第 5 章（Dispatcher） aten/src/ATen/core/dispatch/
第 6 章（ATen 代码生成） torchgen/ + aten/src/ATen/native/native_functions.yaml
第 7-8 章（Autograd） torch/csrc/autograd/ + tools/autograd/
第 9 章（nn.Module） torch/nn/
第 10-11 章（优化器、数据） torch/optim/ + torch/utils/data/
第 12 章（Dynamo） torch/_dynamo/
第 13 章（AOTAutograd） torch/_functorch/aot_autograd.py
第 14 章（Inductor） torch/_inductor/
第 16-18 章（分布式） torch/distributed/ + torch/csrc/distributed/
第 19 章（序列化） torch/serialization.py + torch/distributed/checkpoint/
第 20 章（量化） torch/ao/quantization/
第 21 章（Profiler） torch/profiler/ + torch/csrc/profiler/

把这张图打印出来或者贴在屏幕边上，每读一章对照看一遍。读到第 18 章时，你应该已经能凭记忆说出"DDP 在哪里、FSDP 在哪里、Inductor lowering 在哪里" —— 那时这张地图就内化成你的肌肉记忆了。

1.11 PyTorch 的设计哲学：六条线索

七层架构看完，你大概会问：这些层的设计为什么都长成这样？背后有没有统一的哲学？

我把 PyTorch 设计中反复出现的六条线索总结如下，它们是理解每一章的"先验"：

1. Pythonic First（Python 原生优先）

PyTorch 不是"绑定到 Python 的 C++ 库"，而是"看起来像 Python 库的 C++ 引擎"。所有 API 设计都先确保 Python 端用起来自然 —— tensor[0] 能用、for t in tensor: 能用、tensor.shape 是 tuple、tensor.dtype == torch.float32 能比较。

这条线索指导了第 1 章的"双层 Tensor 结构"和第 9 章 nn.Module 用 __setattr__ 拦截属性赋值的设计。

2. Dynamic by Default, Static when Needed（默认动态，按需静态）

动态图是默认行为，静态图（torch.compile / torch.jit.script / torch.export）是可选优化。这与 TF 1.x 的"先静态再加 eager"是相反的演进方向。

这条线索指导了第 7-8 章 autograd 的 "前向时偷偷建图" 设计，以及第 12-15 章 torch.compile 的"非侵入式编译"。

3. Composable Transforms（可组合的变换）

像 vmap、grad、jvp、functionalize 这些"对函数做变换"的能力，必须能组合起来用 —— vmap(grad(f)) 应该工作。

这条线索指导了第 5 章 dispatcher 的"中间层"设计 —— 每一种变换都是一个 dispatch key，可以叠加。

4. Backend Agnostic at the Op Boundary（在算子边界做后端无关）

PyTorch 的算子接口不绑定 CPU/CUDA/MPS。新增一个后端只需要实现算子的 dispatch 表项，不需要改 Python 前端。

这条线索指导了第 5-6 章的 dispatcher + 代码生成体系。

5. Compile-time Codegen Over Runtime Polymorphism（编译期生成代码胜过运行时多态）

PyTorch 大量用 YAML + 代码生成（native_functions.yaml、derivatives.yaml）而不是 C++ 模板魔法。代码生成的好处：可调试、可读、出错信息友好；缺点：一次编译时间长。

这条线索贯穿第 6 章。

6. Backwards Compatibility is Sacred（向后兼容神圣不可侵犯）

公开 API 一旦发布就不破坏。源码里大量"保留兼容层"，是为了让用户的旧代码继续能跑。

这条线索是为什么 PyTorch 源码总是看起来"代码量随版本只增不减"。

把这六条线索记在心里，再读 PyTorch 的设计讨论和 RFC 时，大多能找到它背后对应的是哪一条取舍。

1.12 把章节内容串起来：一次真实的 Transformer Block

把这章学到的所有东西串起来，看看一次真实的 Transformer Block 调用会发生什么。代码：

import torch
import torch.nn as nn

class Block(nn.Module):
    def __init__(self, d=4096):
        super().__init__()
        self.qkv = nn.Linear(d, 3 * d)
        self.proj = nn.Linear(d, d)
        self.norm = nn.LayerNorm(d)

    def forward(self, x):
        h = self.norm(x)
        q, k, v = self.qkv(h).chunk(3, dim=-1)
        attn = torch.nn.functional.scaled_dot_product_attention(q, k, v)
        return x + self.proj(attn)

block = Block().cuda()
x = torch.randn(32, 1024, 4096, device='cuda', requires_grad=True)
y = block(x)
y.sum().backward()

这段代码的 forward + backward，背后调用了几十次 ATen 算子。让我们粗略追一下（CUDA 上的走法；CPU 上的算子序列相同，只是 kernel 换成 CPU 实现）：

sequenceDiagram
    participant Py as Python forward
    participant Disp as Dispatcher
    participant Auto as Autograd
    participant Aten as ATen kernel
    participant CUDA as CUDA stream
    Py->>Disp: norm(x): aten::native_layer_norm
    Disp->>Auto: 记录 NativeLayerNormBackward
    Disp->>Aten: at::native::layer_norm
    Aten->>CUDA: launch layer_norm kernel
    CUDA-->>Py: h
    Py->>Disp: qkv(h): aten::linear → view + t + addmm
    Disp->>Auto: 记录 AddmmBackward0（外加 view/t 的反向节点）
    Disp->>Aten: addmm
    Aten->>CUDA: launch cuBLAS gemm
    CUDA-->>Py: qkv_out
    Py->>Disp: chunk(3): aten::split
    Note over Disp: split 不复制内存<br/>只产生 3 个 view
    Disp-->>Py: q, k, v (views)
    Py->>Disp: SDPA: aten::scaled_dot_product_attention
    Note over Aten: q/k/v 是 3 维 fp32，<br/>不满足 fused kernel 要 4 维的前提<br/>→ 回退 math 实现
    Disp->>Aten: bmm + _safe_softmax + bmm
    Aten->>CUDA: launch gemm / softmax kernel
    CUDA-->>Py: attn
    Py->>Disp: proj(attn): aten::linear
    Py->>Disp: x + proj_out: aten::add (本章的主角!)
    Disp-->>Py: y
    Py->>Disp: y.sum(): aten::sum
    Py->>Disp: backward(): autograd Engine 启动
    Note over Auto: 沿 grad_fn 链反向<br/>先 SumBackward0 → AddBackward0<br/>→ AddmmBackward0 → BmmBackward0 / SafeSoftmaxBackward0<br/>→ SplitBackward0 → AddmmBackward0<br/>→ NativeLayerNormBackward0
    Auto->>CUDA: launch 一系列反向 kernel
    CUDA-->>Py: 完成 backward

注意到了吗？这一次 forward + backward 中：

  • 6 处模块 / 函数调用（norm / linear×2 / chunk / SDPA / add），拆到 ATen 层是十几个算子（view、t、addmm、split、bmm、_safe_softmax……）
  • 每个 ATen 算子都至少走过 dispatcher 一次
  • 有 autograd 的算子还多一次 redispatch
  • backward 又是另一组数量相当的反向算子调用
  • 所有计算都进同一个 CUDA stream 异步执行
  • 中间张量的内存全部由 caching allocator 管理

这就是你每天 model.forward() 背后实际发生的事。如果你能在脑子里画出这张图，那这个专栏的目标已经实现了一半。

剩下的一半，就是把每个箭头里的细节展开 —— 这就是后面 22 章。

1.13 一个练习：自己跑一次"全旅程"

读完这一章，可以做一个验证理解的练习：用 gdb 给 PyTorch 的 dispatcher 入口下断点，亲眼看到 a + b 怎么穿过这些层。

# 1. 在 debug 模式编译 PyTorch（需要时间，可选）
DEBUG=1 python -m pip install -e . -v --no-build-isolation

# 2. 启动 Python 并 attach gdb
python -c "
import torch, os
print('PID:', os.getpid())
input('press enter')

a = torch.randn(4, 4, device='cuda', requires_grad=True)
b = torch.randn(4, 4, device='cuda', requires_grad=True)
c = a + b
c.sum().backward()
"

# 3. 在另一个终端
gdb -p <PID>
(gdb) break at::_ops::add_Tensor::call
(gdb) continue
# 回到第一个终端按 enter，gdb 会在 a + b 进入 dispatcher 入口时停下
(gdb) break cudaLaunchKernel   # 再在 kernel 发射处下一个断点
(gdb) continue
(gdb) bt    # 看从 THPVariable_add 到 kernel launch 的完整调用栈

如果你把整条调用栈打印下来贴在终端旁边，会看到本章描述的每一层都赫然在列：从 THPVariable_add 到 at::_ops::add_Tensor::call，到 c10::Dispatcher::call，到 VariableType::add_Tensor，到 at::_ops::add_Tensor::redispatch，到生成的 TORCH_IMPL_FUNC(ufunc_add_CUDA)，到 add_kernel，到 gpu_kernel。每一层都对应一段预先编译进 libtorch 二进制的 C++ 代码。

如果不想编 debug 版（很慢），也有一个更轻的练习：用 TORCH_SHOW_CPP_STACKTRACES=1 让 PyTorch 在抛 C++ 异常时打印完整 stack：

TORCH_SHOW_CPP_STACKTRACES=1 python -c "
import torch
a = torch.randn(3, 3)
b = torch.randn(4, 4)  # 故意搞错形状
c = a + b              # 会触发 broadcast 失败
"

异常信息里你会看到形如 at::infer_size_dimvector(...) → at::TensorIteratorBase::compute_shape(...) 的 C++ 栈帧 —— 这正是本章 §1.5 描述的那一层。再往下翻，wrapper_CPU_add_Tensor、at::_ops::add_Tensor::redispatch、VariableType::add_Tensor、at::_ops::add_Tensor::call、THPVariable_add 依次在列，正好是 §1.2–§1.4 这条链倒过来读（即使张量没有 requires_grad，也照样先经过 Autograd 这一层）。

还有一个不依赖任何环境变量的方法 —— 在 Python 里直接打印 dispatcher 的状态：

import torch
op = torch.ops.aten.add.Tensor  # 已经是具体的重载（OpOverload）
print(op)                       # aten.add.Tensor
print(op._schema)               # schema (类型签名)
# 每个 dispatch key 上注册了哪个 kernel、注册在哪个生成文件里
print(torch._C._dispatch_dump_table("aten::add.Tensor"))

到这里，你已经从"PyTorch 用户"变成了"知道 PyTorch 内部地形"的人。剩下的 22 章，是把这张地形图上每一座山、每一条河都走一遍。

三个分级练习

读完本章的练习不只一个。按难度分三级：

初级（读懂一段代码）：

打开 aten/src/ATen/native/native_functions.yaml，找到一个你常用的算子（比如 relu、matmul、softmax），看它的 YAML 声明，搞清楚它有没有 structured、有没有 structured_delegate、dispatch 表里有哪些后端、是不是 tags: pointwise。

中级（追一条调用链）：

挑 torch.matmul(a, b)，从 Python 端开始一路追：Tensor.__matmul__ → aten::matmul → ... 一路追到 cuBLAS 的 gemm 调用。把每一层的文件名记下来，画成一张和本章 §1.2-§1.6 同样结构的调用图。

高级（实现一个 Mode）：

写一个 TimingMode，继承 TorchDispatchMode，对每一个进入 dispatcher 的算子记录耗时（用 torch.cuda.Event），with 块结束后打印 top 10 慢算子。这能让你真切体验 Mode 系统的强大 —— 几十行代码替代了一个完整的 profiler。

如果你能完成这三级练习，本专栏后面的内容你都能轻松吃下。

1.14 时间都去哪儿了：一次加法的性能分解

理解架构只是第一步，工程师真正关心的是性能。让我们对一次 c = a + b 做精细的性能分解，让你下次看 profiler 时心里有数。

假设 a、b 都是 [1024, 1024] 的 fp32 CUDA 张量。一次 a + b 的 CPU 侧 wall time 由下面这些环节拼成，按经验的量级排序（具体数值高度依赖 CPU 型号、Python 版本、是否 debug 构建，本专栏不给伪精确的纳秒数）：

阶段 性质 量级
Python 解释器：__add__ 解析、参数打包 常数开销，与张量大小无关 最大的一块
PythonArgParser 解参数、THPVariable_Unpack 解包 常数开销 次之
Autograd 包装：建 AddBackward0 + redispatch 常数开销，含一次堆分配 次之
TensorIterator build：广播 / 类型提升 / 分配输出 常数开销为主 次之
Dispatcher 算 keySet + 查表 几条位运算 + 一次数组索引 小
CUDA Caching Allocator 分配 c 命中缓存时是纯用户态操作 小
CUDA kernel launch（写进 stream） 一次 driver 调用 小
GPU 实际计算：1024×1024 fp32 逐元素加法 与张量大小成正比 在这个尺寸上远小于上面之和

关键结论不在于某一行是多少纳秒，而在于前七行全是常数开销、最后一行才随数据量增长。在 1024×1024 这种"小到只有 4MB"的张量上，CPU 侧调度开销支配一切 —— 这是 PyTorch（以及所有 eager 框架）的"小张量诅咒"。

想在自己的机器上量出这张表，用 torch.profiler.profile(activities=[CPU, CUDA]) 跑一个只含 a + b 的循环，看 self_cpu_time_total 与 self_device_time_total 的比值（旧名 self_cuda_time_total 在 v2.11 已标 deprecated） —— 第 21 章会详细讲怎么读这份 profile。

这也是 torch.compile 的根本动机：把 Python + dispatcher + autograd 的多次跳转合并成一段编译好的二进制，把上表的常数开销压到接近 0。在小算子密集的工作负载里，加速来源不是"算子算得更快"，而是调度开销被消灭了。

这个分解还能解释另一件事：为什么 PyTorch 在大算子上和静态图框架的性能差距小？因为一个 4096×4096 的 GEMM，GPU 计算时间是毫秒量级，几微秒的调度开销占比可以忽略。所以在大模型训练（每个 op 都很大）上，eager 调度开销不再是瓶颈，性能主要取决于 kernel 本身；eager 模式的劣势集中在大量小算子的场景。

理解这一点，你就能更精确地判断：自己的代码该不该用 torch.compile？ 看 profile 里小算子的累计时长占比 —— 占比越高，编译收益越大；如果时间都花在少数几个大 GEMM 上，编译收益就有限。

1.15 几个常见误解的源头

理解架构之后，回头看一些 PyTorch 用户长期纠结的问题，会发现答案都藏在前面这几层里：

"tensor.cuda() 到底是不是同步的？" —— 默认（non_blocking=False）是同步的：H2D 拷贝走 memcpy_and_sync，拷完才返回；只有源张量在 pinned memory 上并传 non_blocking=True 时，拷贝才只排进 stream、立即返回。这是 §1.6 CUDA Stream 一节的延伸。

"为什么 view() 不能用，要用 reshape()？" —— view 本质上是改 strides 不改 storage，所以要求新形状能用现有 strides 表达（连续张量总能满足）；表达不了时 view 报错，reshape 则会退回复制一份。这是第 2 章张量层的内容。

"为什么 loss.backward() 调一次后就不能再调？" —— 因为 autograd Engine 默认在反向后释放反向图（retain_graph=False），节省内存。这是第 7-8 章 Engine 的内容。

"为什么 cuda.empty_cache() 之后显存看上去还是占着？" —— 因为 caching allocator 把 free 的 block 留在用户态池子里，nvidia-smi 看到的是 PyTorch 进程持有的总显存，不是真正在用的。这是第 4 章 Allocator 的内容。

"为什么 torch.compile 第一次跑特别慢？" —— 因为 Dynamo + AOTAutograd + Inductor + Triton 编译本身要花时间，模型越大越久。第二次起命中编译缓存就快了。这是第 12-15 章编译器栈的内容。

"为什么 DDP 训练有时候会报 'one of the variables needed for gradient computation has been modified by an inplace operation'？" —— 因为 inplace 操作可能破坏 autograd 反向图依赖的中间值。这是第 7 章 autograd 与第 17 章 DDP 的交集。

每一个常见 PyTorch 错误，都对应本专栏的某一章。按错误回查章节，是用本专栏最高效的方式。

1.16 一个高级特性的预告：Mode 系统

最后，作为本章的"彩蛋"，给你看一个 PyTorch 高级特性 —— Mode 系统。它只用一两页就能讲完，但它把前面所有层的设计串起来，是理解 PyTorch 抽象能力的最好范例。

考虑一个真实需求：我想统计自己代码里每个张量操作被调用了多少次，不改任何用户代码。

在传统框架里这几乎做不到 —— 你要么修改 torch.add 的源码，要么写一个全局 monkey patch，两者都很丑。但 PyTorch 提供了一个优雅的接口：

import torch
from torch.utils._python_dispatch import TorchDispatchMode

class CountOps(TorchDispatchMode):
    def __init__(self):
        self.counts = {}

    def __torch_dispatch__(self, op, types, args=(), kwargs=None):
        op_name = str(op._schema.name)
        self.counts[op_name] = self.counts.get(op_name, 0) + 1
        return op(*args, **(kwargs or {}))   # 真正调用底层算子

with CountOps() as counter:
    a = torch.randn(100, 100, device='cuda')
    b = torch.randn(100, 100, device='cuda')
    c = (a + b).relu().sum()

print(counter.counts)
# {'aten::randn': 2, 'aten::add': 1, 'aten::relu': 1, 'aten::sum': 1}
# （_schema.name 不带重载名；要区分 add.Tensor / add.Scalar，用 str(op)）

这段代码的魔法在于 TorchDispatchMode —— 它在 dispatcher 层注册了一个"全局拦截器"，每一个穿过 Autograd 层、将要落到后端 kernel 的算子都会被它看一眼，再决定要不要往下走。

它的实现机制就是利用了 §1.4 讲过的 dispatch key：TorchDispatchMode 对应一个特殊的 DispatchKey（叫 Python），当进入 with 块时这个 key（连同用于保存 TLS 快照的 PythonTLSSnapshot）被加进 thread-local 的 included 集（c10/core/impl/TorchDispatchModeTLS.cpp）。Python key 的优先级排在 Autograd 之下、后端 key 之上，所以 dispatcher 先走完 Autograd，redispatch 下来时命中 Python key，转进用户的 __torch_dispatch__ 方法 —— 这也是为什么 Mode 看到的是 autograd 之后的算子。

这套 Mode 系统是 PyTorch 最灵活的扩展机制之一。它支持的能力远不止"计数"，还包括：

  • FakeTensorMode：让张量"假执行" —— 只走 shape/dtype 推导，不真做计算（torch.compile 用它做 graph capture）
  • FunctionalTensorMode：自动把 inplace 操作改写成纯函数式（AOTAutograd 用它做 functionalization）
  • ProxyTorchDispatchMode：把每次算子调用记录成 FX Graph 节点（torch.compile 的 trace 阶段用它）
  • autocast：自动把 fp32 张量转 fp16/bf16 —— 它不是 TorchDispatchMode，torch.amp.autocast('cuda') 是把 AutocastCUDA 这类 dispatch key 从线程局部的 excluded 集里放出来（aten/src/ATen/autocast_mode.cpp 的 set_autocast_enabled），由 C++ 注册的 kernel 做类型转换
  • InferenceMode：torch.inference_mode() 用线程局部的"排除 key"开关把整条 autograd 路径从 DispatchKeySet 里摘掉（torch.no_grad() 不动 dispatch key，只翻 GradMode 这个 TLS 布尔，见 §1.7）

这些"看起来魔法的功能"，本质上都落在 dispatcher 这一层 —— 或注册一个 mode key，或往 DispatchKeySet 里增删一个功能 key。当你理解了 Mode 系统，你会发现 PyTorch 的扩展性非常强 —— 许多"对所有算子做点什么"的需求，几十行以内的 Mode 就能实现。

第 5 章会把 Mode 系统的源码完整剖析，并教你写自己的 Mode。

1.17 本章小结

  • 七层架构：Python 前端 → C++ 桥接 → Dispatcher → ATen → 后端 Kernel → Autograd → torch.compile
  • torch.Tensor 是双层结构：Python 层提供灵活性，torch._C.TensorBase 提供 C++ 性能
  • 算子注册靠 YAML + 代码生成：native_functions.yaml 是 PyTorch 算子的"宪法"，torchgen 把它翻译成数千个 C++ 与 Python 包装
  • Dispatcher 是 PyTorch 的心脏：用 DispatchKeySet bitmap 决定每次调用走哪个实现，支持 redispatch 实现"洋葱式"中间层（Autograd / Functionalize / vmap / FakeTensor）
  • TensorIterator 处理逐元素操作的所有杂事：广播、类型提升、连续性、内存遍历
  • CUDA Caching Allocator 取代 cudaMalloc：把"alloc/free"做成用户态池子操作，避免设备同步
  • autograd 不是单独的层，而是 dispatcher 中间层：靠 redispatch 与 grad_fn 实现
  • torch.compile 是另一条路：Dynamo 拦截 CPython 帧、AOTAutograd 切分图、Inductor 生成 Triton —— 这是 PyTorch 2.x 的灵魂

下一章把这张架构图的最底层抽出来单独看：Tensor / Storage / TensorImpl 这"三件套"是怎么把"一段连续内存"变成"用户能在 Python 里写 a + b 的张量对象"的。理解它是看懂后续所有章节（dispatcher / autograd / compile）的前提 —— 因为所有这些子系统操作的都是 Tensor。

1.18 不要被名字误导：几个容易混淆的概念

最后澄清几个本章涉及但容易混淆的术语，免得后面章节出现时你卡住：

  • torch.add vs Tensor.add vs aten::add.Tensor：第一个是 Python 函数，第二个是 Tensor 方法，第三个是 ATen 算子注册名。三者通过代码生成对应到同一个底层实现 —— 不要被名字差异吓到，本质是同一个函数的不同包装
  • c10::Dispatcher vs torch._C._dispatch_* vs dispatchTable_：C++ 单例 vs Python 侧的内省函数族（如 §1.13 用到的 _dispatch_dump_table）vs OperatorEntry 里的定长分派数组。第 5 章都会出现
  • Tensor vs TensorBase vs TensorImpl：用户类 vs C++ 基类 vs 内部实现类。第 2 章会拆三件套
  • TORCH_LIBRARY vs TORCH_LIBRARY_IMPL vs TORCH_LIBRARY_FRAGMENT：声明算子 vs 实现算子 vs 在已有库里加更多算子。第 5 章 §5.2 会用到
  • "算子（op）" vs "kernel" vs "stub"：算子是接口（如 add.Tensor），kernel 是后端实现（如 ufunc codegen 生成的 add_kernel），stub 是按 device_type 二次分派的桩（如 add_stub，它不经过 dispatcher）。三者层层包裹

把这些区分记牢，源码里 grep 这些关键字就能精准定位，不会再迷路。

1.19 延伸阅读

  • 官方源码导读：Edward Yang 的博客《PyTorch internals》（2019）—— 仍然是 PyTorch 整体架构最权威的导读
  • PyTorch Conference 历年主题演讲：核心开发者对 PyTorch 演进的回顾（录像在 PyTorch 官方 YouTube 频道）
  • PyTorch RFC 仓库：github.com/pytorch/rfcs —— 大特性的设计讨论都在这里
  • 《vLLM 推理内核深度解析》专栏第 1 章：推理引擎的架构全景，其 CUDA/C++ 算子以 PyTorch 自定义算子（torch.ops._C_cache_ops.* 等）注册给 Python 调用，与本章可以对照阅读
  • PEP 523：Python 帧评估 API 的提案文档 —— torch.compile 拦截字节码的底层机制
  • c10/core/DispatchKey.h 的注释：745 行的文件里，枚举定义前后穿插着上百行设计说明，是理解 dispatcher 最权威的一手资料，比任何博客都细致。建议在读第 5 章前先把它从头到尾读一遍 —— 你会发现本章 §1.4 的不少说法都能在那里找到出处
  • aten/src/ATen/native/README.md 与 aten/src/ATen/core/op_registration/README.md：算子怎么声明、怎么注册的一手说明，是第 6 章的预习材料
  • PyTorch 开发者论坛 dev-discuss.pytorch.org：核心开发者在那里发设计说明、开发日志（compiler、distributed、Tensor subclass 等分类），是除了源码之外最直接的一手信息源