PyTorch Dispatcher介绍

PyTorch中的Dispatcher（调度器）是内核层的核心算子调度中枢，是连接上层算子API和下层不同实现的路有组件，负责根据输入的特征自动选择最合适的算子执行逻辑，是PyTorch兼容多设备、多特性的核心支撑

为什么需要Dispatcher

PyTorch需要适配大量不同的运行场景：

1. 不同设备：CPU、CUDA、MPS（苹果芯片）、XLA（TPU）、各类NPU等
2. 不同数据类型/张量格式：普通稠密张量、稀疏张量、量化张量、FP16/FP32/FP64等
3. 不同功能特性：自动求导（Autograd）、vmap（批量映射）、torch.compile、TorchScript等

如果没有统一调度，每个算子都要写大量if-else判断适配不同场景，代码会极度冗余且难以扩展，Dispatcher就是把所有算子的调度逻辑统一抽离出来的组件；

而有了Dispatcher后，上层API和下层实现分离，新增设备/特性只需要注册对应实现即可，不需要修改上层API，比如第三方硬件厂商适配NPU只需要注册对应Dispatch Key的Kernel，不需要修改PyTorch核心代码。且同一个API自动适配不同场景，用户不需要关心底层实现细节，使用成本极低。

核心工作原理

Dispatcher 的运行逻辑可以概括为："提前注册 + 动态分派" 两个环节：

1. 算子注册

   所有算子的不同实现（即kernel）都会提前注册到Dispatcher的全局注册表，每个实现都会绑定一个Dispatcher Key，用来标识这个实现适用的场景：

   • 设备类Key：CPU / CUDA / MPS等，标识适用的硬件
   • 功能类Key：AutogradCPU / AutogradCUDA（支持自动求导）、QuantizedCPU（量化）、SparseCPU（稀疏张量）等，标识支持的特性

   比如torch.add算子会同时注册CPU版、CUDA版、支持自动求导的CPU/CUDA版、量化版等多个实现，每个实现对应不同的Dispatch Key

2. 动态分派

   用户调用算子时，Dispatcher会自动完成以下流程：

   • 收集输入张量的所有属性：设备类型、是否开启求导、是否是稀疏/量化张量、数据类型等，生成对应的Dispatch Key集合
   • 在全局注册表里匹配优先级最高、最适配当前场景的算子实现
   • 调用该实现返回结果