1. 从一次算子调用说起dispatcher 到底在干什么很多人在读 PyTorch 源码时第一次碰到dispatcher这个词都会有点懵。torch.add(a, b)这行代码看起来平平无奇但如果你在 Python 层打断点往里跟会发现它经过torch/_C/_VariableFunctions.pyi的声明、穿过 pybind11 的绑定层、进入 C 的at::add然后——就撞上了一堵墙。这堵墙就是 dispatcher。我在第一次读这块代码的时候最困惑的不是它有多复杂而是为什么需要这么一层东西。直接调对应的 kernel 不就行了吗后来踩了几次坑才明白PyTorch 要面对的场景远比“一个算子对应一个实现”复杂得多同一个addCPU 上有一种实现CUDA 上有另一种稀疏张量上又是另一种还有自动微分需要插入反向图节点还有vmap、functionalization、autocast这些变换需要拦截。如果每个算子都手写一堆if-else来判断走哪条路代码会彻底失控。dispatcher 就是来解决这个问题的。它本质上是一个运行时的算子分派中枢给定一个算子名和一组输入张量它根据这些张量的属性设备类型、布局、是否需要梯度等以及当前生效的各种变换状态决定最终应该调用哪个具体的 kernel 函数。这一篇是“谈谈 dispatcher”系列的第二部分重点放在dispatch key 的计算逻辑、dispatch table 的组织结构以及一次完整分派的执行路径上。第一部分我们聊了 operator schema 和注册机制如果你还没看建议先补一下否则这篇里提到的OperatorHandle、Dispatcher::findSchemaOrThrow这些会有点接不上。这篇适合已经能编译 PyTorch、想在源码层面搞清楚“一个算子调用到底经历了什么”的读者也适合正在做自定义算子扩展、被 dispatch 行为搞得莫名其妙的同学。2. Dispatch Key分派决策的核心依据2.1 DispatchKey 到底是什么为什么不用简单的枚举打开c10/core/DispatchKey.h你会看到一个巨大的枚举定义。粗略数一下光DispatchKey这个 enum class 就有几十个条目CPU、CUDA、HIP、XLA、MPS、SparseCPU、SparseCUDA、AutogradCPU、AutogradCUDA、AutocastCPU、AutocastCUDA、Functionalize、VmapMode、Python、BackendSelect、CompositeExplicitAutograd……第一次看会觉得很乱。这里有个关键设计点DispatchKey 不是简单的“设备类型”。它混合了好几个维度的信息后端/设备维度CPU、CUDA、MPS、XLA 等决定 kernel 跑在什么硬件上。布局维度Sparse、SparseCsr、NestedTensor 等决定张量的内存组织方式。变换维度Autograd、Autocast、Functionalize、VmapMode 等这些不是“设备”而是“当前有哪些变换层需要介入”。特殊用途Python用于 torch::deploy 之类的场景、BackendSelect用于需要动态决定后端的算子、CompositeExplicitAutograd复合算子的兜底。为什么要把这些混在一起因为 dispatcher 的决策逻辑是统一的它不关心你这个 key 代表的是设备还是变换它只关心“这个 key 有没有注册 kernel优先级如何”。这种统一带来一个巨大的好处变换可以像设备一样被“插入”到分派链中。比如 Autograd 就是一个 dispatch key当张量需要梯度时AutogradCPU 或 AutogradCUDA 会被加入 key setdispatcher 就会优先走 autograd 的 kernel由它来记录反向图然后再调用真正的计算 kernel。2.2 DispatchKeySet一个被精心设计的位集合单个 DispatchKey 好理解但实际分派时一个张量往往同时对应多个 key。比如一个需要梯度的 CUDA 稠密张量它的 key set 里至少有CUDA、AutogradCUDA、Dense、Autograd这些。dispatcher 需要在这些 key 里挑一个优先级最高的来执行。PyTorch 用DispatchKeySet来表示这个集合底层是一个 64 位实际是 128 位分两个 uint64的位图。每个 DispatchKey 对应一个 bit 位置。这个设计非常讲究// c10/core/DispatchKeySet.h 简化示意 class DispatchKeySet { uint64_t repr_[2]; public: bool has(DispatchKey k) const { return repr_[k / 64] (1ULL (k % 64)); } // ... };用位图的好处是集合运算并集、交集、判断包含都是 O(1) 的位操作而且可以一次性处理多个 key。这在热路径上非常关键——每次算子调用都要算一遍 key set如果这里慢整个框架都会慢。注意DispatchKeySet 的位序不是随便排的。DispatchKey.h里对每个 key 的数值有严格约定因为数值大小直接决定了优先级。数值越大的 key在分派时优先级越高。这个设计让“找最高优先级 key”变成了“找最高位的 1”可以用一条__builtin_clzll之类的指令搞定。2.3 优先级是怎么定的从 BackendSelect 到 Python优先级顺序不是拍脑袋定的它反映了“谁应该先看到这次调用”。我把它整理成一张表方便对照优先级层级代表 DispatchKey作用最高Python, BackendSelect需要动态决定后端或走 Python 特殊路径高Functionalize, VmapMode, Autocast变换层需要在计算前介入中Autograd记录反向图包裹计算低CPU, CUDA, MPS 等真正的后端 kernel兜底CompositeExplicitAutograd, Undefined复合算子或未注册时的处理这个顺序背后的逻辑是变换层要能“包裹”后端 kernel。Autograd 必须先于 CPU/CUDA 执行这样它才能记录下“这个操作发生在哪个后端上”然后调用后端 kernel 完成实际计算。如果反过来后端 kernel 先跑了Autograd 就没机会插入反向节点了。我在调试一个自定义算子时踩过一个坑我给算子注册了AutogradCPU的 kernel但忘了注册CPU的 kernel。结果在前向传播时一切正常因为 AutogradCPU 的 kernel 里我手动调了计算但一旦进入torch.no_grad()上下文AutogradCPU 被从 key set 里移除dispatcher 找不到任何可用的 kernel直接报 “could not run” 错误。这个错误的根源就是没理解 key set 是动态变化的。3. Dispatch Tablekernel 是怎么被组织和查找的3.1 从 OperatorHandle 到 DispatchTable理解了 key接下来要看 kernel 存在哪里。dispatcher 内部维护了一个巨大的映射从OperatorName比如aten::add到OperatorDef每个OperatorDef里又包含一个DispatchTable。DispatchTable的结构可以简化成// 简化示意非真实代码 struct DispatchTable { // 按 DispatchKey 索引的 kernel 函数指针数组 std::arrayKernelFunction, kNumDispatchKeys kernels; // 每个 key 对应的“兜底”kernel用于 fallthrough std::arrayKernelFunction, kNumDispatchKeys fallbacks; };这里有个关键概念叫fallthrough kernel。假设一个算子只注册了CUDA的 kernel没有注册AutogradCUDA。当一个需要梯度的 CUDA 张量调用它时dispatcher 在AutogradCUDA这个 key 上找不到 kernel就会走 fallthrough落到CUDA的 kernel 上。这个机制让算子注册变得灵活你不需要为每个可能的 key 组合都注册 kernel只需要注册你真正关心的那些。但 fallthrough 也有代价。如果 fallthrough 链太长每次分派都要沿着链往下找性能会受影响。PyTorch 的做法是在注册时就把 fallthrough 关系“压平”到 table 里运行时直接查表不需要动态遍历。3.2 注册时发生了什么从 TORCH_LIBRARY 到 table 填充看一个典型的注册代码TORCH_LIBRARY_IMPL(aten, CUDA, m) { m.impl(add.Tensor, TORCH_FN(add_cuda)); }这行代码在静态初始化阶段执行做的事情大致是找到aten::add.Tensor对应的OperatorDef。在它的DispatchTable里把CUDA这个位置填上add_cuda的函数指针。同时更新 fallthrough 关系如果AutogradCUDA没有显式注册它的 fallthrough 会指向CUDA。这里有个容易忽略的细节注册顺序会影响 fallthrough 的解析。如果先注册了AutogradCUDA再注册CUDAfallthrough 关系会在第二次注册时重新计算。PyTorch 在Dispatcher::registerImpl里做了处理保证最终 table 是一致的。但如果你在运行时动态注册比如通过torch.library就要小心顺序问题。实操心得用torch.library做自定义算子时建议把所有 key 的注册放在同一个TORCH_LIBRARY_IMPL块里或者至少保证CompositeExplicitAutograd的注册在具体后端之前。我遇到过因为注册顺序导致 fallthrough 指向错误 kernel 的情况排查了很久才发现是初始化顺序的问题。3.3 一个算子的 DispatchTable 长什么样以aten::add.Tensor为例在完整编译的 PyTorch 里它的 DispatchTable 大致包含这些条目DispatchKeyKernel说明AutogradCPUadd_autograd_cpu记录反向图调用 CPU kernelAutogradCUDAadd_autograd_cuda记录反向图调用 CUDA kernelCPUadd_cpu实际 CPU 计算CUDAadd_cuda实际 CUDA 计算SparseCPUadd_sparse_cpu稀疏张量加法SparseCUDAadd_sparse_cuda稀疏 CUDA 加法AutocastCPUadd_autocast_cpu自动混合精度处理CompositeExplicitAutogradadd_composite复合实现兜底这张表不是静态的它会随着你加载的扩展、注册的自定义算子而变化。理解这一点很重要dispatcher 的行为是运行时可变的这也是它比编译期分派更灵活的地方。4. 一次完整分派的执行路径4.1 从 at::add 到 kernel 调用逐步拆解现在把前面两块拼起来看一次完整的调用。假设执行torch.add(a, b)其中a和b都是需要梯度的 CUDA 张量。第一步Python 层通过 pybind11 进入 C 的at::add。这个函数是自动生成的内部会调用Dispatcher::singleton().callBoxed或callUnboxed。第二步dispatcher 根据输入张量计算 key set。对于a和b各自的 key set 取并集得到类似{CUDA, AutogradCUDA, Dense, Autograd}的集合。第三步从 key set 里找最高优先级的 key。按照优先级表AutogradCUDA比CUDA高所以先看AutogradCUDA有没有注册 kernel。有就用它。第四步调用AutogradCUDA的 kernel。这个 kernel 内部会做两件事记录反向图节点然后调用CUDA的 kernel 完成实际计算。第五步CUDA的 kernel 执行真正的 GPU 计算返回结果。整个过程听起来简单但实际代码里有很多优化。比如 key set 的计算不是每次重新算而是缓存在张量的TensorImpl里dispatch table 的查找用了两级索引先按算子 ID 再按 key避免字符串比较。4.2 变换层是怎么“插入”的以 Autograd 为例Autograd 的插入机制值得单独说。当AutogradCUDA的 kernel 被调用时它并不是简单地“记录一下然后转发”。它做的事情包括检查当前是否在torch.no_grad()上下文中。如果是直接转发到CUDAkernel不记录任何东西。检查输入张量是否真的需要梯度。如果不需要同样直接转发。如果需要梯度创建一个AddBackward节点把输入张量挂上去然后把输出张量的grad_fn指向这个节点。这里的关键是Autograd 的 kernel 和计算 kernel 是分开注册的。add_autograd_cuda和add_cuda是两个不同的函数。这种分离让 Autograd 可以独立于计算逻辑演进也让你可以在自定义算子时只注册计算 kernel让 Autograd 自动 fallthrough 过去当然这样就没有反向支持了。常见问题为什么我的自定义算子在前向时正常反向时报 “element 0 of tensors does not require grad”大概率是因为你只注册了CUDA或CPU的 kernel没有注册AutogradCUDA或AutogradCPU。dispatcher 在反向时找不到 autograd kernelfallthrough 到了计算 kernel而计算 kernel 不会创建反向节点。4.3 性能关键分派开销到底有多大很多人担心 dispatcher 会带来性能损失。实测下来一次分派的开销大约在几十到几百纳秒级别对于大算子比如矩阵乘法可以忽略但对于小算子比如逐元素加法在 CPU 上可能会有 5% 到 10% 的影响。PyTorch 做了几件事来降低开销key set 缓存张量的 key set 在创建时算好存在TensorImpl里后续调用直接读。dispatch table 压平fallthrough 关系在注册时解析好运行时直接查表。避免虚函数kernel 调用用的是函数指针不是虚函数减少间接跳转。boxed vs unboxedcallUnboxed路径避免了参数装箱比callBoxed快很多。如果你在做性能敏感的扩展建议用callUnboxed路径并且尽量把多个小算子合并成一个大算子减少分派次数。5. 常见问题与排查技巧实录5.1 “Could not run ‘aten::xxx’ with arguments from the ‘XXX’ backend”这是最常见的 dispatcher 报错。原因通常是你用的张量后端比如 MPS、XLA没有注册对应算子的 kernel而且没有 fallthrough 到其他后端。排查步骤确认报错信息里的算子名和后端名。用torch._C._dispatch_dump_table(aten::xxx)打印这个算子的 dispatch table看有哪些 key 注册了 kernel。如果目标后端确实没有 kernel要么注册一个要么把张量转到有 kernel 的后端比如.cpu()。避坑技巧有些算子在某些后端上是通过CompositeExplicitAutograd实现的这种实现会自动 fallthrough 到所有后端。如果你看到 table 里有CompositeExplicitAutograd但调用还是报错可能是 key set 里没有包含这个 key。检查一下张量的布局和设备类型。5.2 自定义算子注册后不生效这种情况通常有几个原因注册的 key 不对比如你注册了CPU但张量是 CUDA 的dispatcher 不会 fallthrough 到 CPU。注册顺序问题如果CompositeExplicitAutograd在具体后端之后注册可能会覆盖掉具体后端的 kernel。命名空间冲突自定义算子的名字和已有算子重名注册被忽略或覆盖。排查方法用torch._C._dispatch_dump_table确认 table 内容用torch._C._dispatch_find_schema_or_throw确认 schema 存在。5.3 Autograd 和 no_grad 下的行为差异前面提到过torch.no_grad()会从 key set 里移除 Autograd 相关的 key。这导致同一个算子在两种上下文下走不同的 kernel。如果你在自定义算子里依赖 Autograd 的行为比如手动创建反向节点在no_grad下就会失效。实操心得写自定义算子时如果需要在no_grad下也保持某些行为不要依赖 Autograd kernel而是把逻辑放在计算 kernel 里或者用at::AutoDispatchBelowAutograd之类的 guard 显式控制。5.4 常见问题速查表现象可能原因解决方法Could not run with arguments from backend后端无 kernel 且无 fallthrough注册 kernel 或转换后端自定义算子不生效key 不对/顺序问题/命名冲突dump table 确认调整注册反向报 does not require grad未注册 Autograd kernel注册 AutogradXXX kernelno_grad 下行为异常Autograd key 被移除逻辑移到计算 kernel分派性能差小算子频繁调用合并算子用 unboxed 路径6. 从源码里学到的设计思路读 dispatcher 源码最大的收获不是记住了多少函数名而是理解了一种用统一机制处理多维度分派的设计思路。把设备、布局、变换都抽象成 dispatch key用位图和优先级表来统一决策这个设计让 PyTorch 能够在不修改核心分派逻辑的前提下不断加入新的后端和变换层。我自己在做自定义算子扩展时会刻意模仿这种思路把“什么条件下走哪条路”抽象成 key把“具体怎么走”抽象成 kernel中间用一张表来连接。这样代码的可扩展性会好很多加新后端时只需要注册新 kernel不需要改分派逻辑。另外一点体会是dispatcher 的很多“坑”其实都源于对 key set 动态变化的理解不足。Autograd 会加 keyno_grad 会减 keyautocast 会加 keyvmap 会加 key。每次调用前key set 可能都不一样。理解了这一点很多看似诡异的行为就都能解释了。最后分享一个调试小技巧在 C 层可以用c10::Dispatcher::singleton().dispatchTable直接打印 table在 Python 层可以用torch._C._dispatch_dump系列函数。这两个工具在我排查 dispatch 问题时帮了大忙比单纯读代码快得多。