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

“If you don’t know how it works, you don’t really know how to use it.”

—— Edward Yang，PyTorch 核心维护者，《PyTorch internals》博客

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。所有 AI 工程师每天都在写这种代码，几乎没人停下来问：这一行 a + b 到底发生了什么？

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

flowchart TD
    P[Python 用户代码<br/>c = a + b] --> P1[Python 前端层<br/>Tensor.__add__]
    P1 --> P2[C++ 桥接层<br/>pybind11 / TORCH_PYTHON_API]
    P2 --> P3[Dispatcher 层<br/>at::add 入口 / OperatorEntry]
    P3 --> P4[ATen 算子层<br/>structured_delegate / TensorIterator]
    P4 --> P5[后端 Kernel 层<br/>add_stub @ CPU/CUDA/MPS/...]
    P5 --> P6[内存与执行<br/>CUDA Caching Allocator + Stream]

    P3 -.记录反向图.-> AG[Autograd 层<br/>VariableType_2.cpp:add]
    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__
<slot wrapper '__add__' of 'torch._C.TensorBase' objects>

注意两件事：

1. a 的类型是 torch.Tensor —— 这是一个纯 Python 类
2. 但它的 __add__ 是一个 slot wrapper，对象指向 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（v2.11 实测行号）：

class Tensor(torch._C.TensorBase):
    """Wrapper around a torch._C.TensorBase ... allowing Python users
    to add custom subclasses ..."""
    ...

为什么要分两层？因为：

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

Python 的 + 操作符在 CPython 解释器里会优先查找类型的 tp_add slot —— 而 torch._C.TensorBase 在 C 层注册了这个 slot，直接跳进 C++。Python 层就这样被”绕过去”，但用户感受不到 —— 这是 PyTorch 用户层友好与 C++ 层高性能并存的关键设计。

为什么这种设计如此重要

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

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

• C 层用 CPython 的 PyTypeObject 接口注册 THPVariableType，把 tp_add / tp_mul / tp_call 等 slot 直接指向 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++ 代码大致是：

// 伪代码（去除了 dispatch key 处理与 trace）
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++ 桥接 —— pybind11 与 TensorBase

Tensor::add 这个 C++ 方法定义在哪里？通过 IDE 跳转或者 grep，会看到一份生成的头文件 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++ 类，定义在 c10/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_4.cpp
Tensor add_Tensor::call(const Tensor& self, const Tensor& other, const Scalar& alpha) {
    static auto op = c10::Dispatcher::singleton()
        .findSchemaOrThrow("aten::add.Tensor", "")
        .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:159：

class TORCH_API Dispatcher final {
public:
    static Dispatcher& singleton();   // 全局单例

    // 算子注册表：算子名 → OperatorEntry
    LeftRight<ska::flat_hash_map<OperatorName, OperatorHandle>> operators_;

    // 后端 fallback 表：DispatchKey → 兜底实现
    KernelFunction backendFallbackKernels_[num_backends];

    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:88
class TORCH_API OperatorEntry {
    OperatorName name_;

    // 一个 OperatorEntry 对应多个 (DispatchKey, KernelFunction) 注册
    std::array<KernelFunction, num_dispatch_keys> dispatchTable_;

    // 计算给定输入张量集合的 DispatchKeySet
    KernelFunction lookup(DispatchKeySet ks) const;
    ...
};

DispatchKeySet 是一个 64 位 bitmap，每一位对应一个 DispatchKey。当一个张量参与到调用里，PyTorch 会从张量自身的 key_set_ 里拿到它的 keys，再 OR 上线程局部状态里的 keys（如 autograd 是否启用、是否在 vmap 里、是否在 compile trace 里），得到本次调用的最终 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(a, b, 1)<br/>记录反向图、redispatch 到 CUDA
    I->>D: 重新进 dispatcher<br/>这次去掉 Autograd key
    D->>K: 取 CUDA key 的 KernelFunction
    K->>I: at::native::add_out(...)<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 是多分派 + 优先级 + bitmap 的组合，比 C++ 单分派强大得多。它是 PyTorch 能在不破坏 API 的前提下不断引入新特性（FunctionalTensor、SymbolicTensor、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 要改 50 个文件
• 后端供应商压力：NVIDIA、AMD、Intel、华为都希望把自家硬件接进 PyTorch，但每次合并都要动核心代码

2019-2020 年 Edward Yang 主导的 dispatcher 重构（俗称 “Dispatcher V2”）就是为了解决这个问题。新的 dispatcher 把所有这些”特性维度”统一成 DispatchKey 这个一维概念，然后用 bitmap + priority 机制做组合。这是 PyTorch 历史上最重要的一次内部重构之一，没有它就没有今天的 functorch、torch.compile、FakeTensor 等生态。

这次重构留下的注释是源码里最珍贵的资料。打开 c10/core/DispatchKey.h，你会看到顶上有 200 多行精心写就的注释，把 DispatchKey 的设计原则、key 的优先级排序、为什么某些 key 要放在某个位置 —— 全部讲得清清楚楚。这是 Edward Yang 留给后来者的”设计文档”，比任何博客都权威。第 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
    MkldnnCPU: mkldnn_add_out

——aten/src/ATen/native/native_functions.yaml 中 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（精简版）
TORCH_IMPL_FUNC(add_out) (
    const Tensor& self, const Tensor& other, const Scalar& alpha, const Tensor& result
) {
    add_stub(device_type(), *this, alpha);
}

*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） —— 编译器会进一步把它向量化为 SIMD 指令（AVX2/AVX512/NEON）

这一段不到 10 行的代码，会被代码生成器在编译时模板展开成几十个版本：每个 dtype（fp32/fp16/bf16/i32/i64/complex）一份，每个后端（CPU/CUDA/MPS）一份。最终 aten/src/ATen/native/cpu/BinaryOpsKernel.cpp 与 aten/src/ATen/native/cuda/BinaryAddSubKernel.cu 里能看到这些展开后的实现。

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

回到一个真实场景：你写下 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 被重新解释成”广播视图”
5. 输出：分配一块新的 1024 × 1024 × 4 = 4MB fp32 显存做 c

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

// aten/src/ATen/TensorIterator.cpp 内部流程的简化
TensorIterator::build() {
    compute_shape();              // 算广播后的形状
    compute_strides();            // 把所有输入的 strides 调整到广播形状
    compute_types();              // 类型提升 (promote_types)
    allocate_or_resize_outputs(); // 分配 c
    coalesce_dimensions();        // 把可合并的维度合并以加速循环
}

每一步都是一次小型的”算子前奏”。理解 TensorIterator 就理解了 PyTorch 逐元素操作的”通用前置流水线”。第 6 章会专门有一节讲它的内部状态机，把”广播为什么不复制内存""类型提升的优先级表怎么定”这些经典问题彻底说清。

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

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

add_stub 这个分派桩，最终会落到 CUDA 的具体 kernel：

// aten/src/ATen/native/cuda/BinaryAddSubKernel.cu（精简）
void add_kernel_cuda(TensorIteratorBase& iter, const Scalar& alpha_scalar) {
    AT_DISPATCH_ALL_TYPES_AND_COMPLEX(iter.common_dtype(), "add_cuda/sub_cuda", [&]() {
        gpu_kernel_with_scalars(iter, [=]GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t {
            return ::at::native::ufunc::add(a, b, alpha_scalar.to<scalar_t>());
        });
    });
}

REGISTER_DISPATCH(add_stub, &add_kernel_cuda);

gpu_kernel_with_scalars 是 PyTorch CUDA 算子的模板基础设施，它做的事是：

1. 算 launch config：根据元素总数算出 grid 和 block 大小
2. launch CUDA kernel：把 lambda 编译成一个 CUDA __global__ 函数
3. 写入当前 CUDA stream：每一个 launch 都绑定到当前线程的 default CUDA stream

gpu_kernel_with_scalars 的内部实现（在 aten/src/ATen/native/cuda/Loops.cuh 里）大致是：

template <typename func_t>
void gpu_kernel(TensorIteratorBase& iter, func_t f) {
    ...
    int64_t grid = (numel + block_size - 1) / block_size;
    auto stream = at::cuda::getCurrentCUDAStream();
    elementwise_kernel<<<grid, block_size, 0, stream>>>(numel, f, ...);
}

注意 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（精简）
void* CUDACachingAllocator::raw_alloc(size_t size) {
    // 1. 在缓存池里找一个合适大小的 block
    Block* block = find_or_create_block(size);
    // 2. 如果缓存里没有，调 cudaMalloc 申请新 block
    // 3. 标记 block 给某个 stream 使用，stream 析构时回收到池子
    return block->ptr;
}

PyTorch 设计这个 caching allocator 的原因是：cudaMalloc 是设备级同步操作，每次调用都会触发整设备的 sync，性能极差。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（这是 cudaMemcpyAsync 在 D2H 方向上隐含的同步语义）。如果你只是 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()

把这两条记牢，能避免 90% 的”不必要的 synchronize”。

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

最经典的坑是 错误延迟报告。如果你的 CUDA kernel 越界访问了显存，错误不会在那次 kernel launch 时报，而是在下一次 sync 时（可能是几十次操作之后）才报，并且报错信息是”some random kernel failed”。调试时需要设 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
    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 这一层是自动生成的，源头是 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 怎么用 work-stealing 多线程跑后向传播。autograd 是 PyTorch 区别于 TensorFlow 静态图的核心所在。

“记录反向图”是怎么记录的

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

具体的开销由几部分组成：

• 一次 DispatchKeySet 计算（几十纳秒）
• 一次 dispatcher lookup 和 vtable 跳转（一两百纳秒）
• 创建 AddBackward0 节点（几百纳秒，包含一次小内存分配）
• 把节点挂上 c.grad_fn（几十纳秒）
• 一次 redispatch（再来一遍 dispatcher 查找）

加起来在张量加法这种”廉价算子”上，autograd 包装的总开销在 1-2 微秒量级。对于一个张量包含百万元素的 GEMM 运算（几毫秒），autograd 的开销可以忽略不计；但如果你写的代码里有大量小张量操作（如循环里逐元素加法），autograd 开销就会明显。

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

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

性能差距在小算子上能到 2x。第 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/>命中即跳过 Python 解释器]

    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 3.11+ 的 PEP 523 帧评估钩子拦截到 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 都成立）：直接跳过 Python 解释器，直接调编译好的 kernel。这就是为什么 torch.compile 能在很多场景拿到 30%-200% 的提速。

延伸阅读：第 12-15 章是本书的”编译器栈四章”，分别拆 Dynamo、AOTAutograd、Inductor、CUDA Graph 协同。这是 PyTorch 团队最近 4 年里最聪明的工程师在花最多心思做的事，也是本书的技术高峰。

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

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

具体来说，torch.compile 在内部仍然走 dispatcher、仍然调 ATen 算子、仍然用 CUDA Caching Allocator —— 它做的事是把”用户代码 → 算子调用”这一段的开销拿掉（消除 Python 解释器、消除 dispatcher 查找、消除 kernel launch 的小算子合并），但调用真正的算子内核这件事和 eager 没有区别。

这意味着：

• 你可以 混用 @torch.compile 装饰的函数和普通 PyTorch 代码，张量在两者间自由流动
• 编译失败时（如遇到 graph break）会自动 fallback 到 eager —— 用户代码不需要改
• ATen 算子的优化（如 FlashAttention-2 的引入）自动同时惠及 eager 和 compile 两条路径
• 自定义算子（用 TORCH_LIBRARY 注册）也能自动被 torch.compile 识别和编译（前提是注册了 meta 函数）

理解这一点很重要：你不需要”二选一”，而是把 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[pybind11 / THPVariable<br/>解 PyObject → Tensor]
        TB[at::_ops::add_Tensor::call]
    end

    subgraph Core["Dispatcher + ATen"]
        DP[Dispatcher 单例<br/>DispatchKeySet 决策]
        AT[at::native::add_out<br/>structured + TensorIterator]
        UF[ufunc::add<br/>self + alpha * other]
    end

    subgraph Backend["后端 Kernel"]
        CU[CUDA add_kernel_cuda<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

对应到本书章节：

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
│
├── 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（已 freeze 维护）
│   ├── 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/                      #   量化（PT2E / FX graph mode）
│
└── tools/                       # 构建与代码生成工具
    ├── autograd/                #   autograd 与 Python 包装的代码生成器
    │   ├── derivatives.yaml     #     反向规则
    │   ├── gen_autograd_functions.py
    │   └── gen_python_functions.py
    └── codegen/                 #   torchgen：算子代码生成器

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

把这张图打印出来或者贴在屏幕边上，每读一章对照看一遍。读到第 10 章时，你应该已经能凭记忆说出”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 的元类设计。

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()

这段 8 行代码，背后调用了几十次 ATen 算子。让我们粗略追一下：

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 → matmul + add
    Disp->>Auto: 记录 MmBackward + AddmmBackward
    Disp->>Aten: at::native::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_flash_attention
    Note over Aten: 命中 FlashAttention-2 内核
    Disp->>Aten: at::native::flash_attention
    Aten->>CUDA: launch FlashAttn kernel
    CUDA-->>Py: attn

    Py->>Disp: proj(attn): aten::linear
    Py->>Disp: x + proj_out: aten::add (我们 §1 的主角!)
    Disp-->>Py: y

    Py->>Disp: y.sum(): aten::sum
    Py->>Disp: backward(): autograd Engine 启动
    Note over Auto: 沿 grad_fn 链反向<br/>先 SumBackward → AddBackward<br/>→ MmBackward → SDPABackward<br/>→ NativeLayerNormBackward
    Auto->>CUDA: launch 一系列反向 kernel
    CUDA-->>Py: 完成 backward

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

• 至少 6 个 ATen 算子 被调用（norm / linear×2 / chunk / SDPA / add）
• 每个算子都走过 dispatcher 一次（共 6+ 次）
• 有 autograd 的算子还多一次 redispatch（再 6 次）
• backward 又是另一组 6+ 次的反向算子调用
• 所有计算都进同一个 CUDA stream 异步执行
• 中间张量的内存全部由 caching allocator 管理

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

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

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

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

# 1. 在 debug 模式编译 PyTorch（需要时间，可选）
DEBUG=1 python setup.py develop

# 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 c10::Dispatcher::callBoxed
(gdb) continue
# 回到第一个终端按 enter，gdb 会在 dispatcher 命中时停下
(gdb) bt    # 看完整调用栈

如果你把整条调用栈打印下来贴在终端旁边，会看到本章描述的每一层都赫然在列：从 THPVariable_add 到 at::_ops::add_Tensor::call，到 c10::Dispatcher::call，到 VariableType::add，到 redispatch，到 at::native::structured_add_out，到 add_kernel_cuda，到 gpu_kernel。每一层都对应你电脑 CPU 上某段被 JIT 编译进二进制的 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(...) → at::structured_add_out::impl(...) 的 C++ 栈帧 —— 这正是本章描述的层。

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

import torch
op = torch.ops.aten.add.Tensor
print(op.default)               # 默认实现
print(op.default._schema)       # schema (类型签名)
# OperatorHandle 提供了 dumpKernelTable 等内省接口（需要 debug 模式）

到这里，你已经从”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.1-§1.5 同样结构的调用图。

高级（实现一个 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 的耗时大约是这样分布的（H100 GPU 上的典型数值）：

——所以在如此小的张量上，真正算加法的时间比所有调度开销加起来短 50 倍。这是 PyTorch（以及所有 eager 框架）的”小张量诅咒”：调度开销支配一切。

这就是 torch.compile 的根本动机：把 Python + dispatcher + autograd 的多次跳转合并成一段编译好的二进制，把上表的前六行从 4000 ns 压到接近 0。在小算子密集的工作负载（如 CUDA Graph 加 transformer 推理）中，torch.compile 能拿到 1.5-3x 的加速 —— 不是因为算子算得快了，而是因为调度开销被消灭了。

这个分解还能解释另一件事：为什么 PyTorch 在大算子上和静态图框架的性能差距小？因为大算子（如 4096×4096 的 GEMM）的 GPU 实际计算时间能到几十毫秒，调度开销 5 微秒占比 0.01%，可以忽略。所以 PyTorch 在大模型训练（每个 op 都很大）上和 JAX/TF 几乎没有性能差距，只在大量小算子的场景才有显著劣势。

理解这一点，你就能更精确地判断：自己的代码该不该用 torch.compile？ 看 profile 里小算子（kernel time < 100 us）的累计时长 —— 如果占总时长 30% 以上，编译收益大；如果都是大算子，编译收益小。

1.15 几个常见误解的源头

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

“tensor.cuda() 到底是不是同步的？” —— 是异步的（H2D 拷贝排进 stream），但下次访问 CPU 数据时会自动等待，所以用户感觉不到。这是 §1.5 CUDA Stream 一节的内容。

“为什么 view() 不能用，要用 reshape()？” —— view 只在张量内存连续时有效，因为它本质上是改 strides 不改 storage；非连续张量必须先 contiguous() 复制一份。这是第 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.Tensor': 1, 'aten::relu': 1, 'aten::sum': 1}

这段代码的魔法在于 TorchDispatchMode —— 它在 dispatcher 层注册了一个”全局拦截器”，每一个进入 dispatcher 的算子都会被它先看一眼，再决定要不要往下走。

它的实现机制就是利用了 §1.3 讲过的 dispatch key：TorchDispatchMode 对应一个特殊的 DispatchKey（叫 Python），当进入 with 块时这个 key 被压入 thread-local 状态，dispatcher 在每次调用时优先命中 Python key 走到用户的 __torch_dispatch__ 方法。

这套 Mode 系统是 PyTorch 元编程能力的巅峰。它支持的能力远不止”计数”，还包括：

• FakeTensorMode：让张量”假执行” —— 只走 shape/dtype 推导，不真做计算（torch.compile 用它做 graph capture）
• FunctionalMode：自动把 inplace 操作改写成纯函数式（AOTAutograd 用它做 functionalization）
• ProxyMode：把每次算子调用记录成 FX Graph 节点（torch.compile 的 trace 阶段用它）
• AutocastMode：自动把 fp32 张量转 fp16/bf16（这就是 torch.cuda.amp.autocast 的内部实现）
• NoGradMode 与 InferenceMode：你天天用的 torch.no_grad() 也是 Mode

所有这些”看起来魔法的功能”，本质上都是在 dispatcher 上注册了一个 mode key。当你理解了 Mode 系统，你会发现 PyTorch 的扩展性几乎是无限的 —— 任何”对所有算子做点什么”的需求，都可以写一个 5-10 行的 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.dispatch vs dispatch_table：C++ 单例 vs Python 模块 vs 一个数据结构。分别在第 5 章三处出现
• Tensor vs TensorBase vs TensorImpl：用户类 vs C++ 基类 vs 内部实现类。第 2 章会拆三件套
• TORCH_LIBRARY vs TORCH_LIBRARY_IMPL vs TORCH_LIBRARY_FRAGMENT：声明算子 vs 实现算子 vs 在已有库里加更多算子。第 6 章会用到
• “算子（op）” vs “kernel” vs “stub”：算子是接口（如 add.Tensor），kernel 是后端实现（如 add_kernel_cuda），stub 是 dispatcher 桩（如 add_stub）。三者层层包裹

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

1.19 延伸阅读

• 官方源码导读：Edward Yang 的博客《PyTorch internals》（2019）—— 仍然是 PyTorch 整体架构最权威的导读
• PyTorch Developer Conference 2024 KeyNote：Soumith Chintala 关于 PyTorch 演进的回顾
• PyTorch RFC 仓库：github.com/pytorch/rfcs —— 大特性的设计讨论都在这里
• 本系列《vLLM 内核探秘》第 1 章：从推理引擎角度看张量调用链，与本章可以对照阅读
• PEP 523：Python 帧评估 API 的提案文档 —— torch.compile 拦截字节码的底层机制
• c10/core/DispatchKey.h 顶部注释：Edward Yang 亲手写的 200 行设计文档，是理解 dispatcher 最权威的一手资料，比任何博客都细致。建议在读第 5 章前先把这份注释从头到尾读一遍 —— 你会发现本章 §1.3 的全部内容都是它的”读后感”
• tools/codegen/README.md：torchgen 代码生成器的使用说明，是第 6 章的预习材料
• PyTorch GitHub Discussions 的 “Internals” 标签：核心维护者在那里回答内部机制问题，是除了源码之外最权威的信息源
• 本系列《vLLM 内核探秘》第 1-2 章：从推理引擎角度看 PyTorch dispatcher 的调用模式，与本章可对照阅读，互为补集