PyTorch Internals: Overview

This file is the index for the concepts/pytorch-internals/ folder. It lists planned and written subtopic notes, organizes them by theme, and collects the canonical references for the field. Use it to decide what to write next without needing to re-survey the landscape.

📋 Notes in This Folder

🗺️ Subtopic Map

🧱 Tensor Model

Yang (2019): The three-layer object hierarchy. Tensor is a thin Python-side handle (ref-counted pointer); TensorImpl holds all metadata (dtype, device, strides, sizes, AutogradMeta); Storage owns the raw byte buffer. Multiple TensorImpls can point to the same Storage, which is how views are implemented without copying data.

Yang (2019): Physical memory layout for a contiguous 2D tensor. Elements are stored row-major; the stride tuple [ncols, 1] encodes that stepping along dim-0 advances by ncols elements and stepping along dim-1 advances by 1.

Yang (2019): A transposed or sliced view reuses the same Storage but changes the stride/offset fields in TensorImpl. No data is copied; is_contiguous() returns False when strides do not satisfy the row-major ordering invariant.

⚙️ Dispatch & Ops

🔄 Autograd

🚀 Compilation (TorchDynamo / TorchInductor)

🧠 Memory Management

🔧 Custom Operators

🔗 Dependency Graph

flowchart TD
    TS["Tensor & Storage
tensor-storage.md"]
    D["Dispatcher
dispatcher.md"]
    AE["Autograd Engine
autograd-engine.md"]
    TC["torch.compile
torch-compile/overview.md"]
    MM["Memory Management
memory-management.md"]
    CO["Custom Ops
custom-ops.md"]

    TS --> D
    D --> AE
    D --> CO
    AE --> TC
    D --> TC
    TS --> MM
    AE --> CO

🚶 How PyTorch Works: A Brief Walkthrough

This section traces a minimal training step end-to-end through PyTorch’s internal machinery, grounding the subtopic map above in a concrete execution sequence.

The program:

x = torch.randn(4, 8, requires_grad=True, device="cuda")
y = torch.sin(x)
loss = y.sum()
loss.backward()
# x.grad is now populated

Step 1 — Tensor creation

torch.randn(...) allocates a Storage (a raw byte buffer) via CUDACachingAllocator, then constructs a TensorImpl pointing into that Storage with strides [8, 1] and offset 0. The Python-side Tensor object is a thin reference-counted handle to this TensorImpl. Because requires_grad=True, the TensorImpl also carries an AutogradMeta struct that marks it as a leaf variable whose gradient will accumulate into x.grad.

Yang (2019): The Tensor → TensorImpl → Storage ownership chain. The Tensor Python object is a thin ref-counted handle; TensorImpl holds all metadata including strides, dtype, and optional AutogradMeta; Storage is the raw byte buffer allocated by CUDACachingAllocator.

Step 2 — Dispatching torch.sin(x)

Calling torch.sin(x) crosses the Python/C++ boundary into the C10 dispatcher. The dispatcher computes a DispatchKeySet bitmask by inspecting the tensor arguments and thread-local state. For a CUDA tensor with requires_grad=True outside a torch.no_grad() block, the highest-priority active key is Autograd. The dispatcher finds the kernel registered at Autograd for aten::sin and calls it.

The Autograd kernel is auto-generated from native_functions.yaml + derivatives.yaml at build time.

Yang (2020): The C10 dispatch table. Each operator (aten::sin, aten::add, …) has a row of function pointers, one per dispatch key (CPU, CUDA, Autograd, Tracing, …). The dispatcher selects the highest-priority key present in the computed DispatchKeySet and jumps to that cell.

Yang (2020): How the DispatchKeySet bitmask is computed. It is the union of (1) the keys contributed by each tensor argument (device/layout bits), (2) a thread-local include set (e.g. Tracer when inside torch.jit.trace), and (3) a global default set — minus a thread-local exclude set (e.g. Autograd excluded inside torch.no_grad()).

Step 3 — Autograd forward wrapper

The Autograd-keyed kernel for sin is a wrapper, not the math. It:

1. Saves inputs needed for the backward (here, x itself, since \(\frac{d}{dx}\sin x = \cos x\) needs \(x\)).
2. Re-dispatches by excluding Autograd from the DispatchKeySet, causing the dispatcher to fall through to the next key — CUDA. This runs the actual CUDA sin kernel, producing the output tensor y.
3. Wraps the output: sets y.grad_fn = SinBackward, a Node subclass whose next_edges_ point to x’s AccumulateGrad node. This links y into the autograd DAG.

This two-pass structure — Autograd key builds the tape, CUDA key does the math — is repeated for every differentiable op.

Yang (2019): The Variable abstraction (now unified into Tensor) carries an AutogradMeta payload inside TensorImpl. AutogradMeta holds the grad_fn pointer (the backward Node), the accumulated .grad tensor for leaf variables, and the requires_grad flag.

Yang (2020): The two-pass dispatch sequence for a differentiable op. The Autograd key fires first — its kernel saves inputs and prepares the grad_fn — then re-dispatches with Autograd excluded, falling through to the CUDA kernel that performs the actual computation. The output tensor is then wrapped to carry the newly constructed grad_fn.

Step 4 — Tape after the forward pass

After y = torch.sin(x) and loss = y.sum(), the autograd DAG looks like:

flowchart LR
    L["loss
(scalar)"]
    S["SumBackward
grad_fn of loss"]
    SB["SinBackward
grad_fn of y
saved: x"]
    AG["AccumulateGrad
leaf node for x"]

    L --> S
    S --> SB
    SB --> AG

Each Node stores: - next_edges_: pointers to the Nodes of its inputs (the backward DAG edges). - Saved tensors: intermediate values needed to compute the JVP (here, x inside SinBackward). - A sequence_nr_: a monotonically increasing integer used by the engine to schedule execution priority.

Step 5 — loss.backward()

Engine::execute() performs a reverse topological traversal of the DAG using a thread-pool-backed priority queue:

1. Seeds the queue with (SumBackward, upstream_grad=tensor(1.0)).
2. Pops SumBackward, calls its apply(tensor(1.0)), which returns ones_like(y) — the gradient of sum w.r.t. its input. Enqueues (SinBackward, ones_like(y)).
3. Pops SinBackward, calls apply(ones_like(y)). The JVP formula from derivatives.yaml is \(\cos(x) \cdot \text{grad\_output}\). It loads the saved x, computes cos(x) * ones_like(y), and sends the result to AccumulateGrad.
4. AccumulateGrad adds the incoming gradient into x.grad (allocating it on first call).

The engine dispatches each apply() call through the C10 dispatcher again — the gradient ops (cos, pointwise multiply) go through exactly the same Autograd → CUDA path as the forward ops, but because they are themselves not tracked (we are inside the backward), the Autograd key is disabled (via AutoGradMode::set_enabled(false)) and dispatch falls directly to CUDA.

Step 6 — Memory throughout

The CUDACachingAllocator services every tensor allocation in the above sequence — x’s storage, y’s storage, intermediate gradient buffers — without ever calling cudaMalloc after the first few warmup iterations. Freed tensors return their blocks to the pool; the next allocation performs a best-fit lookup in \(O(1)\) (small pool) or \(O(\log n)\) (large pool tree). The SinBackward node holds a reference to the saved x tensor, preventing its block from being reclaimed until the backward completes.

The full picture

flowchart TD
    PY["Python: torch.sin(x)"]
    DISP["C10 Dispatcher
compute DispatchKeySet"]
    AUTO["Autograd kernel
(save inputs, wrap output)"]
    REDISP["Re-dispatch (Autograd excluded)"]
    CUDA["CUDA kernel
actual math"]
    TAPE["Autograd DAG
SinBackward node added"]
    BWD["loss.backward()
Engine::execute()"]
    JVP["SinBackward::apply()
cos(x) * grad_out"]
    ACCUM["AccumulateGrad
x.grad += result"]
    ALLOC["CUDACachingAllocator
services all tensor allocs"]

    PY --> DISP
    DISP -->|"highest key: Autograd"| AUTO
    AUTO --> REDISP
    REDISP -->|"highest key: CUDA"| CUDA
    CUDA --> TAPE
    TAPE --> BWD
    BWD --> JVP
    JVP --> ACCUM
    ALLOC -.->|"allocates buffers"| CUDA
    ALLOC -.->|"allocates grad buffers"| JVP

Each box in this diagram corresponds to a subsystem with a planned deepdive note. The walkthrough above is the thread that ties them together.

📚 Potential Deepdives

The table below lists specific internal subsystems worth dedicated deep-dive notes, beyond the planned cluster files. These are not yet scheduled but represent natural next steps as the cluster matures.

📖 Master References