Open-Source Wikis

/

PyTorch

/

Primitives

/

Dispatch keys

pytorch/pytorch

Dispatch keys

A reference page for the c10::DispatchKey enum (c10/core/DispatchKey.h). Most cross-cutting features in PyTorch are implemented as kernels for one of these keys. For how the dispatcher uses keys see Systems / Dispatcher.

The full enum has ~140 entries. Below are the categories and the most commonly-encountered keys.

Categories

Category Examples
Backend (real device) CPU, CUDA, HIP, XPU, MPS, MTIA, XLA, Lazy, Vulkan, Metal, IPU, PrivateUse1..3
Backend (sparse) SparseCPU, SparseCUDA, SparseCsrCPU, SparseCsrCUDA, SparseMps, SparseXPU
Backend (mkldnn) MkldnnCPU, MkldnnCUDA
Backend (quantized) QuantizedCPU, QuantizedCUDA, QuantizedXPU, QuantizedMeta
Backend (nested) NestedTensorCPU, NestedTensorCUDA, NestedTensorMeta
Functionality Autograd*, Autocast*, Functionalize, Conjugate, Negative, ZeroTensor, FuncTorchBatched, FuncTorchVmapMode, FuncTorchGradWrapper, Python, PythonTLSSnapshot, PythonDispatcher, Meta
Composite (alias) CompositeImplicitAutograd, CompositeExplicitAutograd, CompositeExplicitAutogradNonFunctional, Autograd, Sparse, SparseCsr, Quantized, NestedTensor

Priority

The dispatcher iterates the keyset highest-first. Approximate priority order (from DispatchKeySet.cpp):

PythonTLSSnapshot          ← captures Python TLS at top
PythonDispatcher           ← Python-implemented op overrides
FuncTorchVmapMode          ← vmap level marker
Functionalize              ← rewrite in-place / view to functional
Named                      ← named tensor checking
ConjugateNegative          ← lazy conj/neg
ZeroTensor                 ← lazy zero
ADInplaceOrView            ← autograd in-place / view tracking
AutogradOther / AutogradCPU / AutogradCUDA / AutogradXPU / AutogradMPS / ...
Tracer                     ← JIT tracer
AutocastCPU / AutocastCUDA / AutocastMPS / AutocastXPU
FuncTorchBatched           ← vmap batch rule
FuncTorchGradWrapper       ← grad/jvp wrapper
Sparse* / SparseCsr* / Mkldnn* / Quantized* / NestedTensor*
Meta                       ← shape-only kernel
CPU / CUDA / HIP / XPU / MPS / MTIA / Vulkan / Metal / Lazy / XLA / IPU / PrivateUse1..3

The "include" set picks the next key to dispatch to; "exclude" sets remove keys for the duration of a scope (e.g., c10::AutoDispatchBelowAutograd excludes Autograd*).

The most-encountered keys

Autograd*

One per backend that has an autograd flavour (AutogradCPU, AutogradCUDA, AutogradMPS, AutogradXPU, AutogradXLA, …) plus AutogradOther as a fallback. Generated kernels for these keys are in torch/csrc/autograd/generated/VariableType_*.cpp.

Autocast*

One per backend (AutocastCPU, AutocastCUDA, AutocastMPS, AutocastXPU). Kernels in aten/src/ATen/autocast_mode.cpp.

Functionalize

The single key that powers AOT autograd's "make this graph functional" pass. Kernels in aten/src/ATen/FunctionalizeFallbackKernel.cpp and FunctionalTensorWrapper.cpp.

Conjugate / Negative / ZeroTensor

Lazy markers that defer the corresponding op until materialization is forced. Implemented as fallback kernels in aten/src/ATen/ConjugateFallback.cpp, etc.

FuncTorch*

Three keys (FuncTorchBatched, FuncTorchVmapMode, FuncTorchGradWrapper) that implement the torch.func transforms. Kernels in aten/src/ATen/functorch/.

Python and PythonTLSSnapshot

Power __torch_dispatch__ and TorchDispatchMode. The Python key is added when a tensor with __torch_dispatch__ is in the args; the dispatcher then routes to a Python-side fallback that calls the user's hook.

Meta

A "device" with no real data — meta tensors carry only shape/dtype/device. Kernels for the Meta key (often called meta kernels) compute output shape/dtype without doing math. FakeTensor mode redirects ops here when no real compute is wanted.

Composite alias keys

CompositeImplicitAutograd is a key that expands to "the same kernel for every backend that doesn't override it". Used for ops with a single device-agnostic implementation (e.g., relu = max(0, x)). CompositeExplicitAutograd is a similar alias for ops that compose other ops but need explicit autograd handling.

How to figure out what dispatched

Set TORCH_SHOW_DISPATCH_TRACE=1 and PyTorch will print every dispatcher hop:

[call] op=aten::add, key set={CUDA, AutogradCUDA}, ...
[redispatch] op=aten::add, key set={CUDA}, kernel=...

This is the single most useful debugging tool when a kernel selection surprises you.

Where to look

File Purpose
c10/core/DispatchKey.h The enum
c10/core/DispatchKeySet.h KeySet + priority arithmetic
c10/core/DispatchKeySet.cpp Priority table implementation
aten/src/ATen/core/dispatch/Dispatcher.cpp Lookup logic

Built by Factory AutoWiki from public repository content. It is a generated preview for codebase exploration, not source-maintained documentation.

Dispatch keys – PyTorch wiki | Factory