1. 从一次算子调用说起dispatcher 到底在忙什么很多人在读完 dispatcher 的第一篇源码笔记之后脑子里会留下一个模糊的印象它是个“分发器”负责把算子调用路由到正确的后端实现上。这个印象没错但太粗了。真正让我决定写第二篇的契机是某次在调试一个自定义算子时发现同一个torch.add调用在不同条件下竟然走了完全不同的代码路径而且中间还穿插了好几次看起来“多余”的查表和装箱操作。当时我就意识到如果只停留在“分发”这个层面很多性能问题和行为差异根本解释不通。这篇内容就是想把 dispatcher 的第二层逻辑讲透它不只是一个路由器更像是一个带有多级缓存、支持动态注册、并且能在运行时根据张量属性做决策的调度中枢。适合已经读过第一篇、对 PyTorch 源码结构有基本了解但还想搞清楚“为什么这么设计”“性能开销到底在哪”“自定义扩展时该注意什么”的读者。如果你正在做算子扩展、性能调优或者单纯想理解 PyTorch 动态图底层的调度机制这篇应该能帮你省下不少翻源码的时间。我个人的习惯是读源码不能只看调用链还要看数据结构。dispatcher 的核心其实就藏在几个关键结构体里把它们的生命周期搞清楚整个调度过程就透明了。下面我会按照“设计思路 → 核心结构 → 实操验证 → 问题排查”的顺序来展开中间会穿插一些我实际调试时用到的技巧和踩过的坑。2. 整体设计思路为什么需要两级分发2.1 从“算子重载”到“调度表”的演进逻辑早期 PyTorch 的算子实现相对直接每个算子在不同后端CPU、CUDA上各写一份调用时通过类型判断来跳转。这种做法的好处是直观坏处是扩展性极差。每加一个后端就要改所有算子的分发逻辑每加一个算子就要在所有后端上实现一遍。更麻烦的是像自动微分、混合精度、量化这些功能本质上也是对算子行为的“包装”如果分发逻辑写死在调用点这些功能就很难优雅地插入。dispatcher 的设计目标就是把这些横切关注点从算子实现中剥离出来。它的核心思路是算子调用不再直接绑定到某个具体实现而是先经过一个统一的调度层由调度层根据运行时信息决定走哪条路径。这个运行时信息包括张量的设备类型、布局、数据类型以及当前是否处于自动微分、是否开启了某些全局模式等。这里的关键设计决策是“两级分发”。第一级是算子级分发根据算子名找到对应的调度表第二级是内核级分发根据张量属性从调度表中选出具体的内核函数。为什么要分两级因为算子名是静态的编译期就能确定而张量属性是动态的只有运行时才知道。把静态部分和动态部分分开可以让静态部分提前构建好索引减少运行时的查找开销。2.2 调度键dispatcher 的“语言”dispatcher 用来描述运行时信息的抽象叫做DispatchKey。你可以把它理解成一种标签每个张量身上都带着若干标签调度表也是按标签来组织的。常见的 DispatchKey 包括CPU、CUDA、AutogradCPU、AutogradCUDA、QuantizedCPU等等。注意这里有个容易混淆的点AutogradCPU和CPU是两个不同的键前者表示“这个算子需要在 CPU 上做自动微分处理”后者表示“这个算子在 CPU 上的普通实现”。DispatchKey 的集合是有限且预先定义好的它们之间还有优先级关系。当一次调用同时匹配多个键时dispatcher 会按照优先级从高到低选择。这个优先级顺序不是随便定的而是根据功能层次来排比如自动微分相关的键通常优先于普通后端键因为自动微分需要在后端实现之上做包装。理解这个优先级是排查“为什么我的算子走了意料之外的路径”这类问题的关键。提示DispatchKey 的完整列表在源码的c10/core/DispatchKey.h中建议打印出来贴在显示器旁边调试时对照着看会快很多。2.3 缓存机制为什么第二次调用快得多如果每次算子调用都要遍历所有 DispatchKey 做匹配性能肯定扛不住。dispatcher 的优化手段是引入调度缓存。第一次调用某个算子时dispatcher 会根据输入张量的 DispatchKey 集合计算出一个“调度结果”然后把结果缓存起来。后续调用如果 DispatchKey 集合相同直接查缓存即可。这个缓存的设计有几个细节值得注意。首先缓存的键不是单个 DispatchKey而是 DispatchKey 的集合用位掩码表示因为一个张量可能同时带有多个键。其次缓存是分算子维护的每个算子有自己的缓存表。最后缓存会在调度表发生变化时失效比如动态注册了新内核或者某些全局模式被切换。理解这一点很重要因为如果你在运行时动态修改了调度表却没看到生效很可能就是缓存没刷新。3. 核心结构拆解调度表与内核的绑定关系3.1 OperatorEntry每个算子的“档案袋”在 dispatcher 内部每个算子对应一个OperatorEntry结构。这个结构里存放了算子的元信息比如 schema、调度表、以及缓存。调度表本质上是一个从 DispatchKey 到内核函数的映射但实际实现比简单的哈希表复杂因为它需要支持“回退”机制。所谓回退是指当某个 DispatchKey 没有直接对应的内核时dispatcher 会尝试找一个“更通用”的键来兜底。比如某个算子在AutogradCUDA上没有专门实现但AutogradCPU上有而自动微分的逻辑与设备无关那么就可以回退到AutogradCPU的实现。这种回退链是预先定义好的不是运行时动态搜索的所以效率很高。我在阅读这部分源码时最大的收获是理解了“调度表不是扁平的而是有层次的”。每一层对应一个 DispatchKey层与层之间有明确的覆盖和回退关系。这种设计让算子实现可以只关注自己关心的那一层其他层交给回退机制处理。3.2 KernelFunction类型擦除的代价与收益内核函数在存储时并不是以具体类型保存的而是经过了一层类型擦除。具体来说所有内核都被包装成一个统一的函数指针类型参数通过一个Stack结构传递。这个Stack可以理解为一个类型安全的变体数组里面存放着输入输出张量和其他参数。类型擦除的好处是调度表可以用统一的方式存储和调用内核不需要为每种签名单独处理。坏处是每次调用都要做一次“装箱”和“拆箱”把具体参数打包进 Stack内核再从 Stack 里取出参数。这个开销在大多数场景下可以接受但在极端性能敏感的场景下比如小算子的高频调用就会显现出来。我实测过一个简单的逐元素加法在 CPU 上走 dispatcher 的路径比直接调用底层内核大约多出 15% 到 20% 的开销。这个数字不是绝对的取决于参数数量和缓存命中情况但足以说明类型擦除不是免费的。PyTorch 团队显然也意识到了这一点所以在一些热点路径上做了特殊优化比如对常见签名提供快速通道。3.3 调度键的优先级与回退链DispatchKey 的优先级顺序是 dispatcher 行为的核心规则之一。我把它总结成一句话功能越“外围”的键优先级越高。比如自动微分、量化、追踪这些功能都是在基础计算之上做包装所以它们的键排在前面。而纯后端键CPU、CUDA排在后面作为最终的执行层。这个顺序带来的一个实际影响是当你注册一个自定义内核时如果注册的键优先级不对可能会被其他键“截胡”。比如你注册了一个CPU键的内核但输入张量同时带有AutogradCPU键那么实际走的是自动微分的路径你的内核可能根本不会被调用。这种情况在调试时很容易让人困惑因为从表面看你明明注册了 CPU 实现为什么没生效答案就是优先级在起作用。注意注册自定义内核时务必确认目标 DispatchKey 的优先级位置。如果希望内核在所有情况下都生效可能需要注册多个键或者使用“兜底”键。4. 实操验证用代码观察调度过程4.1 环境准备与调试开关要观察 dispatcher 的实际行为光看源码不够还得动手跑。PyTorch 提供了一些环境变量和调试接口可以打印调度信息。我常用的几个# 打印算子调度详情输出量较大建议配合 grep 使用 export TORCH_SHOW_DISPATCH_TRACE1 # 打印自动微分相关的调度信息 export TORCH_LOGSautograd设置好之后运行一个简单的算子调用就能在终端看到调度路径。输出里会包含每次调用的 DispatchKey 集合、选中的内核、以及是否命中缓存。第一次运行时缓存未命中会看到完整的查找过程第二次运行同样的调用就会看到缓存命中的提示。我建议用一个最小的脚本做实验比如只包含一个torch.add这样输出不会太乱。等熟悉了输出格式再逐步增加复杂度比如加入自动微分、切换设备等。4.2 一个最小复现观察缓存命中与失效下面这段代码可以用来观察缓存行为import torch # 第一次调用缓存未命中 a torch.randn(3, 3) b torch.randn(3, 3) c a b # 第二次调用相同形状和设备应该命中缓存 d a b # 改变设备如果有 CUDA 的话缓存键变化 if torch.cuda.is_available(): a_cuda a.cuda() b_cuda b.cuda() e a_cuda b_cuda在开启调度追踪的情况下第一次a b会打印出完整的 DispatchKey 匹配过程第二次则直接显示缓存命中。切换到 CUDA 后DispatchKey 集合变了缓存自然失效又会走一次完整查找。这个实验很直观地展示了缓存的作用范围它只对相同 DispatchKey 集合的调用有效。4.3 自定义算子注册的完整流程注册自定义算子并观察它在 dispatcher 中的表现是理解调度机制的最好方式。下面是一个简化但完整的流程import torch from torch.library import Library # 创建一个自定义库 my_lib Library(my_ops, DEF) # 定义算子 schema my_lib.define(my_add(Tensor a, Tensor b) - Tensor) # 注册 CPU 实现 torch.library.impl(my_lib, my_add, CPU) def my_add_cpu(a, b): return a b # 注册 CUDA 实现如果有 CUDA torch.library.impl(my_lib, my_add, CUDA) def my_add_cuda(a, b): return a b # 调用 x torch.randn(2, 2) y torch.randn(2, 2) z torch.ops.my_ops.my_add(x, y)注册完成后用调度追踪观察my_add的调用路径。你会看到 dispatcher 先根据算子名找到my_ops::my_add的 OperatorEntry然后根据输入张量的 DispatchKey 选择 CPU 或 CUDA 内核。如果输入张量带有自动微分信息还会看到AutogradCPU键被优先匹配自动微分层会在你的内核之上做包装。这里有个容易踩的坑注册内核时指定的键必须与张量实际携带的键匹配。比如你只注册了CPU但输入张量是AutogradCPU那么自动微分层会先接管然后可能回退到你的 CPU 内核也可能因为缺少自动微分规则而报错。具体行为取决于算子是否被标记为可微分。5. 常见问题与排查技巧实录5.1 内核未生效从优先级和缓存入手“我注册了内核但调用时没走我的实现”是最高频的问题。排查思路可以按以下顺序排查步骤检查内容常见原因1确认 DispatchKey 是否匹配张量实际携带的键与注册的键不一致2检查优先级是否被覆盖更高优先级的键如 Autograd先接管3确认缓存是否失效动态注册后缓存未刷新4检查算子名是否拼写正确命名空间或算子名错误导致注册到错误位置我遇到过一次典型情况注册了CPU内核但输入张量因为开启了梯度追踪而带有AutogradCPU键自动微分层优先匹配后发现没有对应的自动微分规则直接报错而不是回退到 CPU 内核。解决办法是同时注册自动微分规则或者用torch.no_grad()关闭梯度追踪。5.2 性能异常类型擦除与缓存未命中如果发现某个算子调用比预期慢可以从两个方向排查。第一是缓存未命中表现为每次调用都走完整查找流程。常见原因是输入张量的 DispatchKey 集合频繁变化比如在 CPU 和 CUDA 之间来回切换或者梯度追踪状态不稳定。解决办法是尽量保持调用环境一致避免在热点路径上频繁切换设备或梯度模式。第二是类型擦除开销表现为小算子高频调用时开销占比明显。这个问题的根源是 Stack 的装箱拆箱缓解办法包括合并小算子减少调用次数、使用融合算子、或者在性能关键路径上直接调用底层内核绕过 dispatcher。不过最后一种做法会失去 dispatcher 带来的灵活性需要权衡。5.3 动态注册后的行为变化动态注册内核后调度表会更新但缓存不会自动全部失效。PyTorch 的处理方式是给调度表打一个版本号缓存条目记录版本号版本号变化时缓存失效。这个机制大多数时候工作正常但在多线程环境下如果注册和调用并发进行可能会出现短暂的窗口期调用方看到的是旧缓存。这种情况在实际使用中很少遇到但在做插件化扩展时需要留意。提示如果动态注册后行为不符合预期可以先手动清空缓存再试。虽然 PyTorch 没有公开的清空接口但重启进程是最简单粗暴的验证方式。5.4 自动微分与 dispatcher 的交互陷阱自动微分层是 dispatcher 中最复杂的部分之一。它通过AutogradCPU、AutogradCUDA等键插入调度链在基础内核之上包装梯度计算逻辑。这里有两个容易出问题的地方。第一是梯度模式与调度键的对应关系。开启梯度追踪时张量会带上Autograd相关的键关闭时则不带。如果你在自定义算子中手动处理了梯度但没有正确注册自动微分规则可能会出现梯度计算错误或静默失败。第二是原地操作的特殊处理。原地操作在自动微分中需要特殊标记否则会破坏计算图。dispatcher 通过Autograd键下的特殊内核来处理这些情况自定义算子如果涉及原地修改需要格外小心。6. 从调度机制看扩展与调优的边界理解 dispatcher 的调度机制之后很多之前觉得“玄学”的现象都有了合理解释。比如为什么某些自定义算子在特定条件下性能骤降为什么动态注册的内核有时不生效为什么自动微分和普通调用走的路径完全不同。这些问题的答案都藏在 DispatchKey 的优先级、缓存机制和类型擦除这几个核心设计里。我在实际扩展算子时总结出一条经验尽量让自定义算子的 DispatchKey 集合保持稳定。如果输入张量的键集合频繁变化缓存命中率会下降调度开销就会上升。具体做法包括在热点路径上固定设备和梯度模式、避免混用不同布局的张量、把动态注册集中在初始化阶段完成而不是运行时反复注册。另一个体会是dispatcher 的设计虽然增加了调用开销但它带来的扩展性收益远大于成本。没有这套机制PyTorch 不可能在保持核心稳定的同时快速支持这么多后端和功能。读源码的意义不在于记住每个结构体的字段而在于理解这些设计决策背后的权衡。当你下次遇到调度相关的问题时能快速定位到是优先级问题、缓存问题还是类型擦除问题这篇内容的目的就达到了。最后分享一个调试小技巧在怀疑调度路径有问题时先用TORCH_SHOW_DISPATCH_TRACE1跑一遍把输出保存到文件然后对照 DispatchKey 列表逐行分析。大多数情况下问题在输出里已经写得很清楚了只是需要知道该看哪一行。