How PyTorch Executes: Dispatcher, Tensor Subclasses, and Devices

Trace torch.add(a, b) where a requires grad and both live on the CPU:

Four facts carry this module.

1. Every tensor carries a dispatch key set. A CPU float tensor in 2.8 reports DispatchKeySet(CPU, ADInplaceOrView, AutogradCPU, AutocastCPU) (private torch._C._dispatch_keys, run locally). It is the same with or without requires_grad: the Autograd kernel always runs and decides at runtime whether to record anything.
2. Thread-local state adds and removes keys. no_grad, autocast, Python modes and functorch transforms edit a thread-local include/exclude set; they never touch tensors.
3. The highest-priority key's kernel runs and may redispatch. Autograd records its node, then calls the operator again with Autograd keys masked, and the chain ends at a backend kernel.
4. Each layer costs host time per op. Measured locally with torch.utils.benchmark (Apple M-series, PyTorch 2.8, 4 threads): torch.add on two 4-element tensors 0.56 µs; with one input requiring grad 0.99 µs; under an empty Python dispatch mode 2.1 µs; a 1M-element add 92 µs. For small tensors the framework is the workload: small ops are overhead-bound, which module 21 measures and module 23 attacks with fusion and CUDA graphs.

Keys are an enum in c10/core/DispatchKey.h; a later value means higher priority. The ones a practitioner meets, highest first:

The order has consequences. Python sits below Autograd, so a dispatch mode sees ops after autograd recorded them, including backward kernels (fences 1–2). Functionalize is also below Autograd: compilers differentiate first, then remove mutation. Only PythonTLSSnapshot is near the top, and it only saves state.

Operator schemas

Every aten operator is declared in native_functions.yaml (2,650 - func: entries at v2.8.0), from which codegen builds the C++ API, bindings and registrations. Addition:

- func: add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
  structured_delegate: add.out
  tags: [core, pointwise]
- func: add_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> Tensor(a!)
  structured_delegate: add.out
- func: add.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!)
  structured: True

Overload names (Tensor, Scalar, out) separate signatures sharing a Python name; torch.add(x, 2.0) resolves to add.Scalar. Tensor(a!) declares that the argument is mutated and the return aliases it; functionalization, export and opcheck trust these annotations. The out= variant is the one implemented: structured_delegate: add.out makes the functional and in-place forms reuse its shape function and loop, so ATen writes each kernel once (module 03 has the user-facing out= rules). The core tag marks Core ATen IR, the opset export targets (module 28).

In Python, torch.ops.aten.add is an overload packet and torch.ops.aten.add.Tensor an OpOverload: a callable handle to one schema that skips overload resolution but still dispatches (0.96 µs locally vs 0.56 µs for torch.add). Modes, FX graphs, decomposition tables and torch.library all name operators this way: the "aten op" is the shared unit of modules 24 and 25.

The first hook runs before C++. Every torch.* function, Tensor method and torch.nn.functional function checks whether an argument or an active TorchFunctionMode overrides __torch_function__ (torch.overrides, Extending PyTorch); if so it calls __torch_function__(func, types, args, kwargs) and returns its result. Logging nn.Linear(4, 3) records ["linear"] (fence 1).

So a function mode sees the public API as written: F.linear, Tensor.__add__, F.scaled_dot_product_attention. It runs before autograd, autocast and C++ decompositions, so it never sees addmm and never sees backward. That is the right layer for duck-typed tensors, for API rewriting (swap every SDPA call for your kernel), and for compiler front ends: Dynamo respects function modes, as the function-modes recipe shows. Helpers: has_torch_function(args) (the fast check the bindings run), handle_torch_function, get_overridable_functions(). A classmethod __torch_function__ on a subclass differs from a mode only in scope: it fires when an argument is that subclass, a mode fires for every call inside its with block.

The second hook is the Python key. A TorchDispatchMode (_python_dispatch.py) receives an OpOverload and plain tensors. The same Linear(4, 3) forward plus .sum().backward(), logged at this level (fence 1, run locally): forward aten.t.default, aten.addmm.default; then sum, ones_like (gradient seed), expand (SumBackward0), one mm for grad_W (the input needs no grad, so no second mm), sum.dim_IntList for grad_b, and t, view and detach bookkeeping. F.linear became t + addmm through a composite decomposition you cannot see from the function level.

Counting FLOPs including backward, or swapping every kernel for a shape-only one (FakeTensor), needs dispatch level; changing what an API call means before autograd needs function level. Inside either handler the mode is popped, so calling func(...) does not recurse; return NotImplemented to defer to the next handler. The DebugMode recipe is a ready-made dispatch mode that logs ops with tensor hashes for diffing runs, but it requires PyTorch 2.10+; torch.utils._debug_mode does not exist in 2.8 (checked locally), so write fence 1's mode instead.

A mode intercepts a scope; a subclass intercepts its own instances and carries state. DTensor (local shard plus placements), FakeTensor (metadata plus fake device) and torchao's quantized tensors (packed ints plus scales) are all wrapper subclasses: the outer tensor has size, stride, dtype and device but no storage, and holds real inner tensors.

class ScaledInt8Tensor(torch.Tensor):
    @staticmethod
    def __new__(cls, qdata, scale, orig_dtype):
        r = torch.Tensor._make_wrapper_subclass(cls, qdata.shape, dtype=orig_dtype,
                                                device=qdata.device)
        r.qdata, r.scale = qdata, scale                  # int8 payload, fp32 scale
        return r
    def __tensor_flatten__(self):
        return ["qdata", "scale"], self.dtype            # inner tensor attrs, static context
    @staticmethod
    def __tensor_unflatten__(inner, orig_dtype, outer_size, outer_stride):
        return ScaledInt8Tensor(inner["qdata"], inner["scale"], orig_dtype)
    @classmethod
    def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
        if func is torch.ops.aten.mm.default:
            a, w = args
            return a @ (w.qdata.to(a.dtype) * w.scale)   # dequantize, then matmul
        return NotImplemented

_make_wrapper_subclass creates the outer tensor without storage, so the reported dtype (bf16) can differ from the payload (int8). torch.Tensor.as_subclass instead reuses an existing tensor's storage: fine for adding methods, useless when the logical tensor is not one dense buffer. __tensor_flatten__ / __tensor_unflatten__ tell Dynamo, AOTAutograd and export how to split the subclass into plain tensors and rebuild it; without them torch.compile graph-breaks and torch.export fails (module 24). Autograd runs above the subclass, so gradients flow with no extra work and come back as the subclass, which the lab's LoggingTensor demonstrates.

Sections 20.6–20.8 are not executed in this course's CI (CPU only); they follow CUDA semantics. Hardware: GPU architecture.

Asynchronous launch. The CUDA backend kernel enqueues work on the current stream (an in-order device queue) and returns; the host dispatches the next op while the GPU computes. Errors surface at a later sync (hence CUDA_LAUNCH_BLOCKING=1), and timing without a sync measures enqueue time.

torch.cuda.set_sync_debug_mode("warn") (docs) warns at every synchronizing call; "error" turns it into a test failure.

The usual offender is loss.item() every step: the host blocks until the GPU drains its queue, so the next step's kernels are not enqueued while the GPU finishes, and the device idles in the gap. Two fixes keep the logging: accumulate on device (running += loss.detach()) and call .item() every N steps, or copy the loss into a pinned host buffer with non_blocking=True, record an event, and read the value one step later after event.synchronize().

Streams and events. One stream is ordered; across streams, order explicitly:

# not executed in this course's CI (CPU only)
side = torch.cuda.Stream()
side.wait_stream(torch.cuda.current_stream())   # side starts after queued default work
with torch.cuda.stream(side):
    y = part(x)                                 # overlaps default-stream work
torch.cuda.current_stream().wait_stream(side)   # rejoin before using y
x.record_stream(side)                           # x allocated on default, used on side
start, end = (torch.cuda.Event(enable_timing=True) for _ in range(2))
start.record(); model(x); end.record(); end.synchronize()
ms = start.elapsed_time(end)                    # device time between events

record_stream exists because the allocator tracks blocks per stream: without it, x's memory can be reused while side still reads it. Backward streams: each backward op runs on its forward op's stream, and backward() relates to surrounding code like any other op. Per the docs' "BC note: using grads on the default stream", 1.9 and earlier synced the default stream with all backward ops; now, if backward() runs under torch.cuda.stream(s) and grads are used outside, you must wait_stream(s) first, even on the default stream.

cudaMalloc/cudaFree are slow and can synchronize, so the caching allocator keeps freed blocks for reuse on the same stream. Hence two counters: memory_allocated() (live tensors) and memory_reserved() (held from the driver; what nvidia-smi shows). empty_cache() returns unused cached blocks to the driver; it frees no live tensor and gives this process no extra memory, so calling it every step only adds cudaMalloc cost.

Byte accounting and snapshots: module 22, CUDA memory usage.

Pinned memory. Pageable host memory can be moved by the OS, so the GPU's DMA engine cannot read it directly. An H2D copy from it is staged: the driver copies into a page-locked buffer, then DMAs, while the host waits (pin-memory tutorial). A pinned tensor is DMA'd directly, so x_pinned.to("cuda", non_blocking=True) returns immediately and overlaps compute; DataLoader(pin_memory=True) pins in a background thread (module 14). Pinning is itself a copy and takes RAM from the OS, so don't pin everything.

The D2H hazard. y.to("cpu", non_blocking=True) returns a CPU tensor before the copy finishes; reading it immediately can see stale data with no error. Record an event after the copy and synchronize on it first. H2D is safe because later kernels on the stream are ordered after the copy; the host is not on the stream.

On CPU-only builds none of this runs: torch.randn(2).is_pinned() is False, and on this module's Apple machine (MPS accelerator) pin_memory() raised a RuntimeError. Device-agnostic code should pin only when device.type == "cuda".

Not executed in this course's CI. In batch-1 decode each kernel may run for microseconds while the host needs microseconds to dispatch it (20.1), so the GPU idles between kernels. A CUDA graph records a launch sequence once and replays it with one launch (CUDA semantics, CUDA graphs for inference).

# not executed in this course's CI (CPU only)
static_x = torch.zeros(B, d_model, device="cuda")
s = torch.cuda.Stream(); s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):                      # warm up on a side stream
    for _ in range(3): model(static_x)
torch.cuda.current_stream().wait_stream(s)
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):                       # capture: recorded, not run
    static_y = model(static_x)
for batch in batches:
    static_x.copy_(batch)                       # write into the captured address
    g.replay()
    use(static_y)                               # output at the captured address

torch.cuda.make_graphed_callables(module, sample_args) graphs a callable's forward and backward and leaves the rest eager (partial capture). With AMP, capture forward, loss and backward but run scaler.step/update eagerly. Side streams in capture must branch from and rejoin the capturing stream. Whole-network capture with DDP needs NCCL ≥ 2.9.6; older NCCL cannot capture collectives. torch.compile(mode="reduce-overhead") automates this as CUDA graph trees (module 23).

torch.accelerator is the backend-neutral front door in 2.8; on this module's Apple laptop current_accelerator() returned device(type='mps').

device = (torch.accelerator.current_accelerator()
          if torch.accelerator.is_available() else torch.device("cpu"))
x = torch.randn(B, T, d_model, device=device)   # create on the target, not .to() afterwards

Creating tensors on the target device skips a host allocation and copy (a Performance Tuning Guide item; living page, accessed 2026-10-10).

A meta kernel runs only shape logic (for structured ops, the real kernel's own meta function) and allocates nothing; module 10 builds models on meta this way. An out-of-tree backend, per the dispatcher and extend-dispatcher tutorials, registers kernels for PrivateUse1 and AutogradPrivateUse1 and renames it with torch.utils.rename_privateuse1_backend, with no core changes.

Two cuDNN knobs (backends): cudnn.benchmark = True autotunes convolution algorithms per shape (good for fixed shapes, bad when they change every step); cudnn.deterministic = True restricts to deterministic algorithms. Convolutions on cuDNN tensor cores and on oneDNN prefer NHWC, so model.to(memory_format=torch.channels_last) plus channels-last inputs avoids layout transposes inside each conv (module 03 derives the strides; the speedup is hardware-dependent and not measured here). TF32 flags are in module 17.

CPU kernels parallelize at two levels (CPU threading). Intra-op splits one op across a pool (OpenMP on standard builds) sized by torch.set_num_threads or OMP_NUM_THREADS. Inter-op runs independent ops concurrently (TorchScript fork) on a pool sized by set_num_interop_threads, settable once, before inter-op work starts. torch.__config__.parallel_info() here reported OpenMP with 4 intra-op and 2 inter-op threads on 10 hardware threads. torch.backends.cpu.get_cpu_capability() names the vector ISA kernels dispatch to: "AVX2"/"AVX512" on x86, "NO AVX" on this arm64 Mac. oneDNN (torch.backends.mkldnn) supplies optimized x86 conv/matmul kernels; it reported unavailable on this arm64 build.

The tuning guide's CPU items are system configuration:

The common failure is oversubscription: several workers each defaulting to all cores. Every lab here calls torch.set_num_threads(min(4, os.cpu_count())): stable CI timing on 4-vCPU runners, no oversubscription when tests share a process, comparable local numbers. Tests and fence 4 restore it.

Fence 1 proves a dispatch mode sees aten overloads in forward and backward while a function mode sees only linear. It does not prove the full op list is version-stable; only addmm (forward) and mm (backward) are asserted.

import torch
import torch.nn as nn
from torch.overrides import TorchFunctionMode
from torch.utils._python_dispatch import TorchDispatchMode

class DispatchLog(TorchDispatchMode):
    def __init__(self):
        super().__init__()
        self.ops = []
    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
        self.ops.append(str(func))
        return func(*args, **(kwargs or {}))

class FunctionLog(TorchFunctionMode):
    def __init__(self):
        super().__init__()
        self.ops = []
    def __torch_function__(self, func, types, args=(), kwargs=None):
        self.ops.append(getattr(func, "__name__", str(func)))
        return func(*args, **(kwargs or {}))

torch.manual_seed(0)
layer = nn.Linear(4, 3)
x = torch.randn(2, 4)

with DispatchLog() as dlog:
    y = layer(x)
    n_forward = len(dlog.ops)
    y.sum().backward()
forward_ops, after_ops = dlog.ops[:n_forward], dlog.ops[n_forward:]
assert "aten.addmm.default" in forward_ops, forward_ops
assert "aten.mm.default" in after_ops, after_ops          # grad_W in backward
assert "aten.mm.default" not in forward_ops

with FunctionLog() as flog:
    layer(x)
assert "linear" in flog.ops and not any("addmm" in op for op in flog.ops), flog.ops
print(f"dispatch forward {forward_ops}; {len(after_ops)} ops after; function level {flog.ops}")

Fence 2 proves autograd is a dispatch layer: under no_grad the forward dispatches the same aten ops but records no graph, so no backward can follow. It also proves torch.ops.aten.add.Tensor is a callable handle to one schema with alpha semantics. It does not inspect key sets (private API).

import torch
import torch.nn as nn
from torch.utils._python_dispatch import TorchDispatchMode

class DispatchLog(TorchDispatchMode):
    def __init__(self):
        super().__init__()
        self.ops = []
    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
        self.ops.append(str(func))
        return func(*args, **(kwargs or {}))

torch.manual_seed(0)
layer = nn.Linear(4, 3)
x = torch.randn(2, 4)
with torch.no_grad(), DispatchLog() as log:
    y = layer(x)
assert y.grad_fn is None and not y.requires_grad
assert "aten.addmm.default" in log.ops
assert "aten.mm.default" not in log.ops

a, b = torch.randn(5), torch.randn(5)
out = torch.ops.aten.add.Tensor(a, b, alpha=2)
torch.testing.assert_close(out, a + 2 * b)
assert "Tensor" in torch.ops.aten.add.overloads()
schema = str(torch.ops.aten.add.Tensor._schema)
assert schema.startswith("aten::add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1)"), schema
print(f"no_grad ops {log.ops}; schema {schema}")

Fence 3 proves meta runs shape logic only: a (1000, 1000) @ (1000, 10) matmul yields a correctly shaped meta tensor without allocating the 3.8 MiB input, and reading data raises. It does not prove every op has a meta kernel (module 25 covers register_fake).

import torch

a = torch.empty(1000, 1000, device="meta")
b = torch.empty(1000, 10, device="meta")
c = a @ b
assert c.shape == (1000, 10) and c.device.type == "meta"
assert c.dtype == torch.float32
try:
    c.tolist()
    raised = None
except (RuntimeError, NotImplementedError) as exc:
    raised = type(exc).__name__
assert raised is not None, "meta tensors must not expose data"
print(f"meta matmul shape {tuple(c.shape)} on {c.device}; data access raised {raised}")

Fence 4 proves the intra-op thread count round-trips and is restored, and that the capability query returns a string. It asserts no specific ISA, since x86 CI and arm64 differ.

import torch

saved = torch.get_num_threads()
try:
    torch.set_num_threads(1)
    assert torch.get_num_threads() == 1
    a = torch.randn(64, 64, generator=torch.Generator().manual_seed(0))
    single = a @ a
finally:
    torch.set_num_threads(saved)
assert torch.get_num_threads() == saved
torch.testing.assert_close(a @ a, single, rtol=1e-5, atol=1e-5)
cap = torch.backends.cpu.get_cpu_capability()
assert isinstance(cap, str) and len(cap) > 0
print(f"threads restored to {saved}; cpu capability {cap!r}")

./code/20_dispatch_tracer.py: (a) OpCounter(TorchDispatchMode) counts aten overloads per phase (a flag set before .backward() and opt.step()) for one SGD step of a 16 → 32 → 4 ReLU MLP with cross-entropy, batch 8; (b) LoggingTensor, a wrapper subclass (_make_wrapper_subclass, __tensor_flatten__/__tensor_unflatten__, and a __torch_dispatch__ that unwraps, runs, logs and rewraps), runs the same step on wrapped parameters; (c) both are compared with a plain step.

python content/courses/pytorch/code/20_dispatch_tracer.py
python site/test_pytorch_20_dispatch_tracer.py

Tested with Python 3.11 and PyTorch 2.8.0 on CPU. One local PyTorch 2.8 CPU run on an Apple M-series laptop (4 threads) produced:

One addmm per Linear forward; three backward mm (grad_W2, grad_h, grad_W1) and none for the input, which needs no grad. SGD took its per-parameter loop (add_ × 4) rather than _foreach_add_, because 2.8 defaults to foreach only on CUDA-class devices. Cross-entropy is _log_softmax + nll_loss_forward, invisible from the function level.

Scope. Names and counts are PyTorch 2.8 CPU kernels; other versions and backends decompose differently (3-D linear becomes mm + add via views). The test asserts structure only: two forward addmm, ≥ 3 backward matmul-like ops, an add-family optimizer op, subclass outputs and grads equal to the plain run, JSON-serialisable output, threads restored. The script took 1.9 s wall including import torch; the test's work took 0.4–0.7 s locally.