Dispatcher

Active contributors: ezyang, bdhirsh

Purpose

The dispatcher is the central op-routing table that decides, for every operator call, which kernel to run. It is what makes "the same Python op works on CPU, CUDA, MPS, autograd, autocast, vmap, FakeTensor, traced symbolic execution, …" possible.

Almost every cross-cutting feature in PyTorch is implemented either as a dispatcher kernel (for some DispatchKey) or by interposing on the dispatcher.

Directory layout

Path	Contents
aten/src/ATen/core/dispatch/	The dispatcher itself
aten/src/ATen/core/dispatch/Dispatcher.h/.cpp	c10::Dispatcher, OperatorEntry, OperatorHandle
aten/src/ATen/core/dispatch/DispatchKeyExtractor.h	How dispatch keys are extracted from arguments
aten/src/ATen/core/op_registration/	The m.def/m.impl registration API
c10/core/DispatchKey.h	The enum of all dispatch keys
c10/core/DispatchKeySet.h	KeySet bitmap and priority arithmetic
aten/src/ATen/core/Library.h	TORCH_LIBRARY / TORCH_LIBRARY_IMPL macros
aten/src/ATen/core/PythonOpRegistrationTrampoline.cpp	Python-side op registration (torch.library)

Key abstractions

Type	File	Purpose
c10::DispatchKey	c10/core/DispatchKey.h	Enum: CPU, CUDA, Autograd, Autocast, Functionalize, …
c10::DispatchKeySet	c10/core/DispatchKeySet.h	64-bit bitmap of dispatch keys
c10::Dispatcher	aten/src/ATen/core/dispatch/Dispatcher.h	The singleton routing table
c10::OperatorEntry	aten/src/ATen/core/dispatch/OperatorEntry.h	Per-op kernel table indexed by DispatchKey
c10::OperatorHandle	aten/src/ATen/core/dispatch/Dispatcher.h	Handle returned by findOp
torch::Library	aten/src/ATen/core/Library.h	Builder used by TORCH_LIBRARY macros
at::FunctionalTensorWrapper	aten/src/ATen/FunctionalTensorWrapper.h	Functionalization key kernel

How it works

Dispatch keys are concerns

Each DispatchKey represents a concern. A simplified prioritized list (highest first):

PythonTLSSnapshot
PythonDispatcher
FuncTorchVmapMode / FuncTorchBatched / FuncTorchGradWrapper
Autocast{CPU,CUDA,XPU,MPS,…}
AutogradOther / Autograd{CPU,CUDA,XPU,MPS,…}
ZeroTensor
Negative
Conjugate
Functionalize
Sparse{CPU,CUDA,…} / SparseCsr{…}
Mkldnn{CPU,CUDA}
Quantized{CPU,CUDA,…}
Meta
CPU / CUDA / MPS / XPU / MTIA / Lazy / XLA / PrivateUse{1,2,3} / Vulkan / Metal

The full enum is ~140 entries; see c10/core/DispatchKey.h. Some are backend keys (CPU, CUDA), some are functionality keys (Autograd, Autocast, Functionalize, Sparse, Quantized), and some are alias keys that fan out (CompositeImplicitAutograd, CompositeExplicitAutograd, …).

Where keys come from

The keyset for a call is the union of:

1. The keysets carried by every input Tensor (each TensorImpl has a key_set_).
2. Ambient included keys from TLS — e.g., autograd is enabled, autocast is enabled, vmap is on, the Python dispatcher is recording.
3. Minus excluded keys from TLS — c10::AutoDispatchBelowAutograd, c10::InferenceMode, etc., remove keys for the duration of a scope.

DispatchKeyExtractor is generated per op based on which arguments are tensors.

Step semantics

Each kernel typically does its concern's work and then redispatches by calling at::redispatch::<op>(keyset_minus_self, args...). The dispatcher recomputes the highest priority key in the remaining set and jumps to that kernel. So a single user call to at::add may pass through Autograd → Functionalize → CUDA in sequence.

Some keys are fallthrough (no-op redispatch) for ops they don't care about. Some are alias keys (CompositeImplicitAutograd) that say "fill in this slot for every backend key unless overridden". Some are fallback (e.g., LegacyBatchedFallback.cpp) that handle every op generically by looping the inner dim.

Registration

There are three ways to register a kernel:

1. Codegen — aten/src/ATen/native/native_functions.yaml plus torchgen produces RegisterCPU.cpp, RegisterCUDA.cpp, etc. that call m.impl("op_name", kernel).
2. TORCH_LIBRARY and TORCH_LIBRARY_IMPL — explicit C++ macros in aten/src/ATen/core/Library.h. Example:

   TORCH_LIBRARY_IMPL(aten, Autograd, m) {
     m.impl("matmul", autograd::matmul);
   }

3. Python torch.library — torch/library.py and torch/_library/. A user-facing API that calls PythonOpRegistrationTrampoline to register kernels backed by Python callables.

Performance

The dispatcher is on the hot path of every op call. It is heavily optimized: most lookups are a DispatchKeySet::highestPriorityTypeId() (a __builtin_clz), one indirect call, no allocations. Boxed (IValue-based) and unboxed (typed) call paths exist; the unboxed path is the fast one.

Functionality keys worth knowing

Key	What it does	Where its kernels live
Autograd*	Records the op on the tape, builds the backward node, redispatches	torch/csrc/autograd/generated/
Autocast*	Up/down-casts dtype for mixed-precision ops	aten/src/ATen/autocast_mode.cpp
Functionalize	Rewrites in-place / view ops into pure functional ones	aten/src/ATen/FunctionalizeFallbackKernel.cpp, FunctionalTensorWrapper.cpp
FuncTorchBatched	vmap	aten/src/ATen/functorch/
Conjugate	Lazy .conj()	aten/src/ATen/ConjugateFallback.cpp
Negative	Lazy unary -	(similar pattern)
ZeroTensor	Tensors known to be all zero (used in some autograd shortcuts)	aten/src/ATen/ZeroTensorFallback.cpp
Sparse*	Sparse tensor ops	aten/src/ATen/native/sparse/
Quantized*	Quantized ops	aten/src/ATen/native/quantized/
Meta	Shape-only (no data) device	Embedded in structured kernels
Python	Python-implemented ops	torch/_library/, torch/library.py
PythonTLSSnapshot	Captures Python TLS for retracing	aten/src/ATen/PythonTorchFunctionTLS.cpp

Integration points

• ATen owns the dispatcher and most kernels. See ATen.
• Autograd is just dispatcher kernels for the Autograd* keys, generated from derivatives.yaml. See Autograd.
• AOT Autograd uses the Functionalize key + FakeTensor mode + PythonDispatcher to symbolically execute graphs. See AOT Autograd.
• torch.compile sits on top of the dispatcher; Inductor emits Triton/C++ that bypass it for known shapes.
• Out-of-tree backends (XLA, MTIA, Lazy) register kernels for their own backend keys via TORCH_LIBRARY_IMPL.

Entry points for modification

• To add a dispatch key: edit c10/core/DispatchKey.h, update DispatchKeySet.cpp priorities, and update aten/src/ATen/core/dispatch/DispatchKeyExtractor.cpp if needed. Many of the alias-key fan-outs need updating too — this is rarely a small change.
• To add a kernel for an existing key: prefer the dispatch: table in native_functions.yaml, or write a TORCH_LIBRARY_IMPL block in your file.
• To trace what the dispatcher does at runtime, set TORCH_SHOW_DISPATCH_TRACE=1 — every op call prints its dispatch path.

Key source files

File	Purpose
aten/src/ATen/core/dispatch/Dispatcher.h	Dispatcher API and OperatorHandle
aten/src/ATen/core/dispatch/Dispatcher.cpp	Dispatcher implementation
aten/src/ATen/core/dispatch/OperatorEntry.h/.cpp	Per-op kernel table
aten/src/ATen/core/dispatch/DispatchKeyExtractor.h	How keysets are computed from args
aten/src/ATen/core/Library.h	TORCH_LIBRARY macros
c10/core/DispatchKey.h	Enum of dispatch keys
c10/core/DispatchKeySet.h	KeySet bitmap and priority logic
aten/src/ATen/PythonTorchFunctionTLS.cpp	Python TLS dispatch hook
aten/src/ATen/FunctionalTensorWrapper.cpp	Functionalize key
aten/src/ATen/autocast_mode.cpp	Autocast key

For the per-key reference, see Primitives / Dispatch keys.