pytorch/pytorch

Architecture

PyTorch is a layered system. The Python API at the top is mostly thin — almost every numerical operation crosses into C++ within microseconds, lands on the dispatcher, and ends up in a kernel selected by tensor device, dtype, and various TLS state. This page walks through the layers from top to bottom and shows the main data and control flow paths.

Layered view

graph TB
    subgraph Python["Python layer (torch/)"]
        torch_init["torch/__init__.py"]
        nn["torch.nn"]
        optim["torch.optim"]
        compile["torch.compile<br/>(_dynamo, _inductor, _functorch)"]
        dist["torch.distributed"]
        export["torch.export, torch.jit"]
    end
    subgraph Bind["Python <-> C++ bindings (torch/csrc/)"]
        Module["Module.cpp init"]
        AutogradPy["autograd Python binding"]
        DispatcherPy["dispatcher binding"]
    end
    subgraph CPP["C++ runtime"]
        Autograd["autograd engine<br/>torch/csrc/autograd/"]
        Dispatcher["c10 dispatcher<br/>aten/src/ATen/core/dispatch/"]
        ATen["ATen ops<br/>aten/src/ATen/native/*"]
        JIT["JIT / TorchScript<br/>torch/csrc/jit/"]
        DistC["c10d distributed<br/>torch/csrc/distributed/"]
    end
    subgraph Core["c10 core (c10/)"]
        TensorImpl
        Storage
        Allocator
        DispatchKey
    end
    subgraph Backend["Backend kernels"]
        CPU["aten/src/ATen/native/cpu/"]
        CUDA["aten/src/ATen/native/cuda/"]
        MPS["aten/src/ATen/native/mps/"]
        XPU["aten/src/ATen/native/xpu/"]
        Vendor["cuDNN / MKL / NCCL / MIOpen"]
    end

    Python -->|pybind11| Bind
    Bind --> Autograd
    Bind --> Dispatcher
    Autograd --> Dispatcher
    Dispatcher --> ATen
    ATen --> Core
    ATen --> Backend
    Backend --> Vendor
    DistC --> Core
    JIT --> Dispatcher

The dispatcher is the spine

Every public ATen op is registered through a small set of macros (TORCH_LIBRARY, TORCH_LIBRARY_IMPL, m.impl) into the c10::Dispatcher defined under aten/src/ATen/core/dispatch/. The dispatcher is a key-indexed table: for each operator schema it stores a row of "kernel function pointers" indexed by DispatchKey, declared in c10/core/DispatchKey.h. Calls move through the table left-to-right, with each key handling its concern (autograd taping, autocast, vmap batching, functionalization, named tensors, conjugate, ZeroTensor, …) before "redispatching" to the next key. The eventual physical kernel (CPU/CUDA/MPS/XPU/…) is reached at the end.

This is the single most important pattern to internalise. Almost every cross-cutting feature — autograd, autocast, AMP, fake tensors, vmap, functionalization, torch.compile's dynamo/aot — is either implemented as a dispatch key, registered as a kernel for one, or interposes on the dispatcher itself. See Systems / Dispatcher for details and Primitives / Dispatch keys for the key-by-key reference.

Op definition pipeline

graph LR
    YAML["aten/src/ATen/native/native_functions.yaml<br/>+ derivatives.yaml"] -->|torchgen| GenH[Generated headers<br/>build/aten/src/ATen/Operators.h]
    YAML -->|torchgen| GenAuto[Generated autograd code<br/>torch/csrc/autograd/generated/]
    YAML -->|torchgen| GenPy[Generated Python bindings<br/>torch/csrc/autograd/generated/python_*.cpp]
    Native["aten/src/ATen/native/*.cpp"] --> Lib[libtorch_cpu.so]
    GenH --> Lib
    GenAuto --> Lib
    GenPy --> PyExt[_C extension module]
    Lib --> PyExt

A new operator typically requires:

A schema and dispatch declaration in aten/src/ATen/native/native_functions.yaml.
A reference implementation in aten/src/ATen/native/<file>.cpp (and CUDA/MPS specializations under cuda/, mps/, etc.).
A derivative declaration in tools/autograd/derivatives.yaml if the op is differentiable.
Test coverage (often via OpInfo in torch/testing/_internal/common_methods_invocations.py).

The torchgen pipeline (torchgen/) reads the YAML and emits all the boilerplate: dispatcher registrations, autograd backwards, Python argument parsers, and stub headers. See Systems / torchgen.

Eager execution path

A typical user call like c = a @ b goes:

sequenceDiagram
    participant User
    participant PyBinding as Python binding<br/>(generated)
    participant Autograd as Autograd kernel
    participant Disp as Dispatcher
    participant Native as native::matmul
    participant CUDA as CUDA kernel

    User->>PyBinding: a @ b
    PyBinding->>Autograd: at::matmul(a, b)
    Autograd->>Autograd: record op + inputs<br/>build MatmulBackward
    Autograd->>Disp: redispatch to AutogradCUDA next keys
    Disp->>Native: native::matmul (CUDA dispatch)
    Native->>CUDA: cublasGemm
    CUDA-->>User: result Tensor with grad_fn

The autograd kernel runs before the device kernel because the Autograd dispatch keys sit higher in the priority order than physical-device keys.

Compile path (`torch.compile`)

graph LR
    User[user code] -->|torch.compile| Dynamo[Dynamo<br/>torch/_dynamo]
    Dynamo -->|symbolic FX graph| AotAutograd[AOT Autograd<br/>torch/_functorch/_aot_autograd]
    AotAutograd -->|forward + backward graph| Inductor[Inductor<br/>torch/_inductor]
    Inductor -->|generates| Triton[Triton kernels]
    Inductor -->|generates| CppCode[C++ kernels]
    Triton --> GPU[GPU]
    CppCode --> CPU[CPU]

Dynamo intercepts Python bytecode, traces a region into FX, and hands it to AOT Autograd to produce joint forward + backward graphs. Inductor lowers those graphs to Triton (GPU) and C++ (CPU). See Features / torch.compile, Systems / dynamo, Systems / inductor.

Distributed training

torch.distributed (Python in torch/distributed/, C++ in torch/csrc/distributed/c10d/) layers on top of process groups (NCCL, Gloo, UCC, MPI). On top of process groups sit FSDP/FSDP2, DTensor, pipeline parallelism, and TorchElastic for multi-node launching. See Systems / distributed.

Repository layout map

Directory	Purpose
`c10/`	Header-only core: `Tensor`/`Storage` impls, dispatch keys, scalar types
`aten/src/ATen/`	Tensor ops, dispatcher, native kernels
`aten/src/ATen/native/`	The vast majority of ATen op implementations
`torch/`	Python package (top-level `__init__.py` is ~111K lines)
`torch/csrc/`	Python <-> C++ bindings, autograd engine, JIT, distributed C++
`torch/_dynamo/`	Bytecode-level graph capture
`torch/_inductor/`	Compiler backend (Triton, C++)
`torch/_functorch/`	AOT Autograd, vmap, function transforms
`torch/distributed/`	Process groups, FSDP, DTensor, pipelining, RPC
`torch/onnx/`	ONNX export
`torch/jit/`	TorchScript frontend
`torchgen/`	Code generator for ops, autograd, Python bindings
`tools/`	Build helpers, autograd YAML, codegen scripts
`caffe2/`	Legacy Caffe2 code (slowly being deleted; serialization still lives here)
`functorch/`	C++ side of functorch (vmap/grad transforms)
`test/`	Test suite (~10K+ test files)
`third_party/`	Submodules: pybind11, googletest, sleef, fbgemm, kineto, cutlass, …

Build system entry points

From source: pip install -e . -v --no-build-isolation (per CLAUDE.md and README.md).
CMake root: CMakeLists.txt (~58K lines).
Variables: setup.py (~60K lines) defines the user-facing build options (USE_CUDA, USE_DISTRIBUTED, BUILD_TEST, MAX_JOBS, DEBUG, …).
Lint runner: .lintrunner.toml configures clang-tidy, mypy, ruff, flake8, and dozens of project-specific linters; spin lint is the user-facing wrapper (see .spin/).
Bazel/Buck: BUCK.oss, *.bzl, build_variables.bzl describe internal Meta builds.

Where to read next

For the dispatcher in detail: Systems / dispatcher.
For an end-to-end autograd walkthrough: Systems / autograd.
For how torch.compile lowers to kernels: Systems / inductor.
For tensor internals: Primitives / Tensor.

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