Warmup Guide — PyTorch & TensorFlow Framework Internals

Who this is for: Someone who knows Python basics and wants to reach principal ML engineer level — able to contribute to PyTorch, write CUDA kernels, debug framework internals, and reason about TensorFlow's compilation model. No prior ML, C++, or GPU experience assumed. Nothing is skipped. Every concept is built from the ground up.

Table of Contents

Part I — The Mathematical Substrate

• Chapter 1: Tensors — From Linear Algebra to Memory
• Chapter 2: Derivatives and the Chain Rule — The Mathematics of Learning
• Chapter 3: Jacobians — Why Backpropagation Works at Scale
• Chapter 4: Automatic Differentiation — Two Modes, One Winner

Part II — How Computers Execute Code

• Chapter 5: CPython — How Python Actually Runs
• Chapter 6: NumPy — Arrays, Memory Layout, and Vectorization
• Chapter 7: C++ and pybind11 — The Bridge to Performance

Part III — PyTorch Architecture

• Chapter 8: ATen, c10, and libtorch — The C++ Foundation
• Chapter 9: Tensors in PyTorch — Storage, Strides, and Memory Management
• Chapter 10: The Dispatcher — PyTorch's Universal Routing System
• Chapter 11: The Autograd Engine — Building and Traversing the Backward Graph
• Chapter 12: torch.autograd.Function — Custom Differentiable Operations

Part IV — Graph Capture and Compilation

• Chapter 13: torch.fx — Intermediate Representation and Graph Transformation
• Chapter 14: torch.compile — The Full Compilation Stack
• Chapter 15: GPU Architecture — From Silicon to CUDA Kernels
• Chapter 16: Custom C++ Operators — TORCH_LIBRARY

Part V — TensorFlow

• Chapter 17: TF1 to TF2 — Two Execution Philosophies
• Chapter 18: tf.function — AutoGraph, Tracing, and ConcreteFunctions
• Chapter 19: GradientTape, XLA, and MLIR

Part VI — Numerical Precision

• Chapter 20: IEEE 754 — How Computers Store Numbers
• Chapter 21: FP16, BF16, and Mixed Precision Training

Part VII — Synthesis and Practice

• Chapter 22: How Everything Connects — The Principal Engineer's Map
• Chapter 23: Lab Guidance — Building minigrad and warp_softmax
• Chapter 24: The Interview Gauntlet — Complete Answers

References

Part I — The Mathematical Substrate

Chapter 1: Tensors — From Linear Algebra to Memory

1.1 What Is a Tensor, Mathematically?

The word "tensor" comes from Latin tendere — to stretch. In mathematics, a tensor is a generalization of scalars, vectors, and matrices to any number of dimensions.

Start with the simplest case:

• A scalar is a single number: 5.0. It has no dimensions — rank 0.
• A vector is an ordered list of numbers: [1.0, 2.0, 3.0]. It has one dimension — rank 1. A vector of length 3 lives in 3D space. Each number describes magnitude along one axis.
• A matrix is a 2D grid of numbers, indexed by row and column. Shape (3, 4) means 3 rows and 4 columns — rank 2.
• A rank-3 tensor is a 3D array. Shape (batch, height, width) for a batch of grayscale images. You can picture it as a stack of matrices.
• A rank-4 tensor extends to (batch, channels, height, width) for RGB images.

The generalization: a rank-\(n\) tensor is indexed by \(n\) integers and holds a number at each index position.

Scalar:  5.0                                 shape ()
Vector:  [1, 2, 3]                           shape (3,)
Matrix:  [[1, 2], [3, 4], [5, 6]]            shape (3, 2)
Rank-3:  [[[1,2],[3,4]], [[5,6],[7,8]]]      shape (2, 2, 2)

Why tensors matter for ML: A neural network's weight matrix is a rank-2 tensor. A batch of images is rank-4. The activations inside a transformer attention layer are rank-3 (batch × sequence_length × d_model). Everything in deep learning lives in tensor space. The framework's entire job is to manipulate tensors efficiently.

1.2 Tensors as Multi-Linear Maps (The Deep Math)

In formal mathematics, a tensor is a multi-linear map — a function that takes multiple vectors as inputs and produces a number, and is linear in each of its inputs separately.

A matrix M of shape (m, n) is a rank-2 tensor. It maps a vector v of length n to a vector Mv of length m. The mapping is linear: M(αv + βw) = αMv + βMw.

Why does this matter for engineers? Because the transpose rule for backpropagation — the fact that the gradient of y = Ax with respect to x is A^T grad_y — comes directly from this multi-linear structure. If you know what kind of object a tensor is mathematically, you can derive its gradient mechanically instead of looking it up.

1.3 Tensors in Memory: The Physical Reality

When you write torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), what exists in your computer's RAM?

Physical memory (a flat array of bytes):
  [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]   ← 6 float32s = 24 bytes, contiguous

The tensor's 2D shape is an illusion maintained by metadata:

The strides are the key. To access element at row r, column c, the index into the flat memory is:

flat_index = storage_offset + r × strides[0] + c × strides[1]
           = 0 + r × 3 + c × 1

This formula works for any shape and any strides. Operations like .T (transpose), .reshape(), and slicing x[1:, ::2] only change the metadata — they do not copy memory. Two tensors can point to the same physical bytes with different shapes and strides.

x = torch.tensor([[1, 2, 3], [4, 5, 6]])   # shape (2, 3), strides (3, 1)
y = x.T                                       # shape (3, 2), strides (1, 3) — no copy
z = x[::2, :]                                # shape (1, 3), strides (6, 1) — no copy

Consequence you must internalize: when you call x.contiguous(), PyTorch checks if the strides match the canonical row-major (C-contiguous) layout where each row is stored consecutively. If they don't — because you transposed or sliced — PyTorch allocates a new buffer and copies data into the canonical layout. This is a hidden allocation cost that appears in profiles as unexpected memory spikes.

1.4 The Storage Object

The physical memory block is owned by a Storage object (in C++: at::StorageImpl). Multiple tensors can share the same Storage with different offsets/strides.

x = torch.arange(12).reshape(3, 4)
y = x[1, :]        # row 1 — shares storage, offset = 4 * sizeof(int) = 16 bytes

# Same underlying memory:
print(x.storage().data_ptr() == y.storage().data_ptr())  # True
print(y.storage_offset())  # 4  (4 elements offset into shared storage)

This sharing is why in-place operations are dangerous: modifying y modifies x because they share a Storage.

Chapter 2: Derivatives and the Chain Rule — The Mathematics of Learning

2.1 What Is a Derivative?

A derivative measures rate of change. For a function \(f(x)\), the derivative at a point \(x_0\) is:

\[ f'(x_0) = \lim_{h \to 0} \frac{f(x_0 + h) - f(x_0)}{h} \]

In plain English: if you nudge the input by a tiny amount \(h\), the derivative tells you how much the output changes per unit of nudge. If \(f'(x_0) = 6\), then nudging \(x\) by 0.001 changes \(f\) by approximately 0.006.

Geometrically: the derivative is the slope of the tangent line to the curve \(f(x)\) at the point \(x_0\).

For ML: the derivative of the loss with respect to a weight tells you — if I increase this weight slightly, does the loss go up or down, and by how much? This is exactly the information gradient descent needs.

2.2 Partial Derivatives — Functions of Many Variables

A neural network loss depends on millions of weights simultaneously. For a function \(f(x, y, z, \ldots)\), the partial derivative with respect to \(x\) is written \(\frac{\partial f}{\partial x}\) and measures the rate of change of \(f\) when only \(x\) changes, holding all other inputs fixed.

Example: \(f(x, y) = x^2 y + 3y\)

\[ \frac{\partial f}{\partial x} = 2xy \qquad \frac{\partial f}{\partial y} = x^2 + 3 \]

The gradient of \(f\) is the vector of all partial derivatives:

\[ \nabla f = \left[\frac{\partial f}{\partial x},\ \frac{\partial f}{\partial y}\right] \]

For a neural network with weights \(W_1, W_2, b_1, b_2, \ldots\), the gradient of the loss is the vector:

\[ \nabla_\theta \mathcal{L} = \left[\frac{\partial \mathcal{L}}{\partial W_1},\ \frac{\partial \mathcal{L}}{\partial W_2},\ \frac{\partial \mathcal{L}}{\partial b_1},\ \ldots\right] \]

Each entry tells you how sensitive the loss is to that weight. Gradient descent subtracts a small multiple of this gradient from the weights:

\[ \theta \leftarrow \theta - \eta \nabla_\theta \mathcal{L} \]

where \(\eta\) (eta) is the learning rate. This moves the weights in the direction that reduces loss.

2.3 The Chain Rule — The Engine of Everything

The chain rule is a rule for differentiating composed functions — functions of functions. This is exactly what a neural network is: a composition of layers.

Single-variable form: if \(z = f(y)\) and \(y = g(x)\), then:

\[ \frac{dz}{dx} = \frac{dz}{dy} \cdot \frac{dy}{dx} \]

Read this as: "how much does \(x\) affect \(z\)?" equals "how much does \(y\) affect \(z\)?" times "how much does \(x\) affect \(y\)?".

Derivation from first principles: Using the limit definition:

\[ \frac{dz}{dx} = \lim_{h \to 0} \frac{f(g(x+h)) - f(g(x))}{h} \]

Let \(\Delta y = g(x+h) - g(x)\). When \(h \to 0\), \(\Delta y \to 0\) as well (assuming \(g\) is continuous). So:

\[ = \lim_{\Delta y \to 0} \frac{f(g(x) + \Delta y) - f(g(x))}{\Delta y} \cdot \lim_{h \to 0} \frac{g(x+h) - g(x)}{h} = f'(g(x)) \cdot g'(x) \]

Multi-variable form: if \(z = f(x, y)\) where \(x = g(t)\) and \(y = h(t)\), then:

\[ \frac{dz}{dt} = \frac{\partial f}{\partial x} \cdot \frac{dx}{dt} + \frac{\partial f}{\partial y} \cdot \frac{dy}{dt} \]

This "sum over all paths" form is exactly how autograd handles a variable like a that is used in multiple downstream operations — it sums up the gradient contributions from each usage.

Concrete example — trace d = (a × b) + a:

a = 2.0, b = 3.0
c = a × b = 6.0        (∂c/∂a = b = 3,  ∂c/∂b = a = 2)
d = c + a  = 8.0        (∂d/∂c = 1,      ∂d/∂a = 1)

Chain rule to find ∂d/∂a (two paths: direct through add, and indirect through mul):

∂d/∂a = ∂d/∂c × ∂c/∂a   +   ∂d/∂a (direct)
       =    1  ×   3     +       1
       = 4

Verify: d = (a × b) + a = ab + a, so ∂d/∂a = b + 1 = 3 + 1 = 4. ✓

This sum-over-paths structure is not a coincidence — it is exactly what the backward pass computes when it accumulates gradients into leaf tensors.

Chapter 3: Jacobians — Why Backpropagation Works at Scale

3.1 The Jacobian Matrix

When a function maps a vector to a vector — \(\mathbf{y} = f(\mathbf{x})\) where \(\mathbf{x} \in \mathbb{R}^n\) and \(\mathbf{y} \in \mathbb{R}^m\) — the derivative is no longer a single number. It is a matrix called the Jacobian:

\[ J = \frac{\partial \mathbf{y}}{\partial \mathbf{x}} = \begin{bmatrix} \frac{\partial y_1}{\partial x_1} & \cdots & \frac{\partial y_1}{\partial x_n} \ \vdots & \ddots & \vdots \ \frac{\partial y_m}{\partial x_1} & \cdots & \frac{\partial y_m}{\partial x_n} \end{bmatrix} \]

Entry \(J_{ij} = \frac{\partial y_i}{\partial x_j}\) — how much output \(i\) changes when input \(j\) changes.

Example: linear layer \(\mathbf{y} = W\mathbf{x}\) where \(W\) is \((m \times n)\). Then \(y_i = \sum_j W_{ij} x_j\), so \(\frac{\partial y_i}{\partial x_j} = W_{ij}\). The Jacobian is exactly \(W\). The gradient of \(\mathbf{y}\) with respect to \(\mathbf{x}\) IS the weight matrix.

3.2 Why Backpropagation Is a Sequence of Jacobian-Vector Products

The chain rule in matrix form: given compostion \(f \circ g\) mapping \(\mathbf{x} \to \mathbf{y} \to \mathbf{z}\):

\[ \frac{\partial \mathbf{z}}{\partial \mathbf{x}} = \frac{\partial \mathbf{z}}{\partial \mathbf{y}} \cdot \frac{\partial \mathbf{y}}{\partial \mathbf{x}} = J_f \cdot J_g \]

For a neural network with loss \(\mathcal{L}\) (a scalar), backprop computes:

\[ \frac{\partial \mathcal{L}}{\partial \mathbf{x}} = \left(\frac{\partial \mathcal{L}}{\partial \mathbf{y}}\right)^T J_g \]

This is a vector-Jacobian product (VJP). The vector is the upstream gradient (already computed); the Jacobian is the local derivative of this operation. We never form the full Jacobian matrix — we compute the product directly.

For a Linear layer y = Wx:

• The Jacobian with respect to x is W
• The VJP (upstream gradient g times Jacobian) is g @ W or equivalently W.T @ g

This is why grad_x = grad_y @ W.T in every linear layer's backward pass. It's not magic — it's the VJP of the Jacobian W.

3.3 Why This Justifies Reverse Mode

Forward mode computes Jacobian-vector products — how outputs change per input direction. To get gradients of a scalar loss w.r.t. all inputs, you'd need one forward pass per input dimension.

Reverse mode computes vector-Jacobian products — how the scalar loss flows back through each operation. One reverse pass gives you gradients w.r.t. all inputs simultaneously.

For neural networks: millions of inputs (weights), one output (loss). Reverse mode is \(O(\text{forward pass cost})\) regardless of how many parameters you have. Forward mode would be \(O(\text{parameters} \times \text{forward pass cost})\) — completely infeasible for 7-billion-parameter models.

Chapter 4: Automatic Differentiation — Two Modes, One Winner

4.1 What Is Automatic Differentiation?

Automatic differentiation (autodiff) is not symbolic differentiation (like a CAS doing algebra) and not numerical differentiation (finite differences). It is the mechanical application of the chain rule to the actual program that computes the function, tracking derivatives alongside values.

Three approaches exist:

PyTorch's autograd IS automatic differentiation. It achieves exact derivatives (up to floating-point rounding) in O(forward pass cost) time.

4.2 The Execution Tape

Reverse-mode autodiff works by recording a tape (also called the Wengert list) of all operations executed during the forward pass. The tape stores:

• The operation performed (add, mul, matmul, relu, …)
• The inputs that were consumed
• Any saved intermediate values needed for the backward

After the forward pass, the tape is traversed in reverse to compute gradients.

Forward (recording):
  a = 2.0 (leaf)
  b = 3.0 (leaf)
  c = a * b    → tape: [MUL, inputs=(a,b), saved=(a=2.0, b=3.0)]
  d = c + a    → tape: [ADD, inputs=(c,a)]
  loss = d

Backward (replay in reverse):
  d.backward(grad=1.0):
    ADD.backward(grad=1.0)  → sends 1.0 to c, 1.0 to a
    MUL.backward(grad=1.0)  → grad_a += 1.0*b=3.0, grad_b = 1.0*a=2.0
  
  a.grad = 1.0 (from ADD) + 3.0 (from MUL) = 4.0
  b.grad = 2.0

PyTorch does not literally have a "tape list" — it uses a DAG of Function objects attached to tensors. But the concept is identical.

4.3 Forward Mode vs Reverse Mode — The Full Story

Forward mode propagates derivatives from inputs to outputs alongside the forward computation. Each variable carries a "tangent" value representing its derivative w.r.t. a chosen input direction.

For f: R^n → R^m with n inputs and m outputs:

• Forward mode costs O(n) × forward pass — one pass per input
• Efficient when n < m (few inputs, many outputs)

Reverse mode propagates gradients from the output back to inputs.

• Reverse mode costs O(m) × forward pass — one pass per output
• Efficient when m < n (many inputs, one output) — exactly the ML case

Neural network training: n = millions of parameters, m = 1 (the loss scalar). Reverse mode is literally a million times more efficient than forward mode for this setting. This is the entire reason deep learning is tractable.

JAX supports both modes (jax.jvp for forward, jax.vjp for reverse) as first-class citizens. PyTorch only exposes reverse mode through autograd but supports forward mode via torch.autograd.functional.jvp.

Part II — How Computers Execute Code

Chapter 5: CPython — How Python Actually Runs

5.1 Python Is Not Executed Directly

When you run python script.py, Python does NOT interpret the text of your file line by line. It compiles it first.

Step 1 — Compilation to bytecode: Python's parser turns your source code into an AST (Abstract Syntax Tree), then the compiler turns the AST into bytecode — a sequence of simple, stack-based instructions. This bytecode is stored in .pyc files in __pycache__/.

def add(a, b):
    return a + b

Compiled bytecode (shown via dis module):

  LOAD_FAST    0 (a)    # push local variable 'a' onto the stack
  LOAD_FAST    1 (b)    # push local variable 'b' onto the stack
  BINARY_ADD           # pop two values, push their sum
  RETURN_VALUE         # pop and return the top of the stack

Step 2 — The CPython interpreter loop: CPython (the standard Python implementation, written in C) executes bytecode in a giant C while loop called ceval.c. For each bytecode instruction, it dispatches to a C function that performs the operation.

The interpreter loop is why Python is "slow" — for every a + b in your code, CPython:

1. Fetches the BINARY_ADD instruction
2. Pops two PyObject* pointers from the stack
3. Looks up the __add__ method on the left object
4. Calls through multiple C function layers
5. Allocates a new Python object for the result
6. Pushes the pointer back onto the stack

This overhead is ~100-1000× more expensive than the equivalent C code doing the same addition.

5.2 Why This Matters for ML Frameworks

ML training loops repeat billions of simple arithmetic operations. If each operation paid the CPython interpreter overhead, training would be impossibly slow. The solution: batch operations into kernels.

When you write y = torch.matmul(a, b):

1. CPython executes ONE bytecode instruction (CALL_FUNCTION)
2. This calls a C++ function registered via pybind11
3. The C++ function dispatches to a highly optimized BLAS routine (Intel MKL or cuBLAS)
4. The kernel runs billions of multiplications and additions without returning to CPython
5. A single Python object (the output tensor) is returned

The Python layer is thin glue. The heavy lifting is in C++/CUDA.

5.3 The GIL — Global Interpreter Lock

CPython has a fundamental limitation called the GIL (Global Interpreter Lock). Only one Python thread can execute Python bytecode at a time. The GIL is a mutex that prevents concurrent Python execution.

Why does this exist? CPython's memory management (reference counting) is not thread-safe. Making it thread-safe would require fine-grained locks on every object, which would be slow. The GIL is a coarse lock that serializes all Python execution.

Impact on ML: Python-level parallelism doesn't help. But ML frameworks work around this by releasing the GIL during C++ computation:

// PyTorch's pattern for long-running C++ calls:
{
  pybind11::gil_scoped_release release;   // release the GIL
  // This C++ code runs without holding Python's lock.
  // Other Python threads can run while we're computing.
  result = do_expensive_computation();
}
// GIL is re-acquired here automatically

During a CUDA kernel launch or a matmul call, Python threads are free to run (e.g., data loading threads) because the GIL is released. This is the mechanism behind PyTorch's DataLoader with num_workers > 0.

5.4 Why TorchDynamo Needs to Understand Bytecode

torch.compile uses TorchDynamo to capture computation graphs. Dynamo works by installing a frame evaluation callback in CPython — a hook that intercepts bytecode execution before it runs.

When Python is about to execute a @torch.compile-decorated function, Dynamo's callback fires. Dynamo reads the raw bytecode instructions (LOAD_FAST, BINARY_MULTIPLY, CALL_FUNCTION, etc.) and converts them into a graph of tensor operations. To do this, it must understand:

• Which bytecode instructions produce tensor values
• Which instructions are Python-only (don't involve tensors)
• Where control flow (jumps, conditionals) happens

This is why torch.compile is more powerful than torch.fx.symbolic_trace — it operates at the bytecode level where it can see everything, not at the Python function level where it only sees what you call explicitly.

Chapter 6: NumPy — Arrays, Memory Layout, and Vectorization

6.1 Why NumPy Exists

Python lists are arrays of Python objects. A list of 1000 floats stores 1000 PyObject* pointers, each pointing to a heap-allocated Python float. The floats are scattered in memory. Iterating over them requires following 1000 pointer indirections.

NumPy arrays store raw values contiguously in C-allocated memory — no Python objects, no pointers. A NumPy array of 1000 float64s is exactly 8000 bytes in one contiguous block. Accessing element 500 is a single memory read at offset 500 × 8 = 4000 bytes. No pointer indirection.

This contiguity is what enables vectorization: modern CPUs have SIMD (Single Instruction, Multiple Data) instructions (SSE, AVX, AVX-512) that operate on 4, 8, or 16 floats simultaneously. A well-optimized loop over a contiguous float array uses these instructions automatically.

6.2 Broadcasting — The Rules You Must Internalize

Broadcasting is NumPy's mechanism for applying operations between arrays of different shapes without copying. The rules:

1. If arrays have different numbers of dimensions, the smaller is padded with 1s on the left
2. Two dimensions are compatible if they are equal OR one of them is 1
3. The output shape is the element-wise maximum of the input shapes

a shape: (     3, 1)
b shape: (  4, 3, 5)
───────────────────
         (  4, 3, 5)  ← output: a is repeated along dim 0 and 1, b along dim 2

Why this matters for autograd: when you compute a + b and a was broadcast, the gradient of the loss w.r.t. a must be summed over the broadcast dimensions — because the same a values contributed to multiple output positions. This is _unbroadcast in the Lab 01 solution:

def _unbroadcast(grad, target_shape):
    # Sum over dimensions that were broadcast during forward
    ...

Forgetting to unbroadcast is one of the most common bugs in custom autograd implementations.

6.3 Vectorization vs Python Loops

# Slow: Python loop, one operation per iteration
result = []
for i in range(len(a)):
    result.append(a[i] * 2.0)

# Fast: vectorized, operates on whole array at once
result = a * 2.0  # single NumPy/PyTorch call → one C kernel

The fast version:

1. Calls into C code
2. Loops in C (100-200× faster than Python loops)
3. May use SIMD to process 8 floats per CPU clock cycle
4. Has no Python object allocation per element

For 1 million elements: the Python loop might take 500ms; the vectorized version might take 0.5ms — 1000× faster.

Chapter 7: C++ and pybind11 — The Bridge to Performance

7.1 Why PyTorch Is Written in C++

C++ gives two things Python cannot: predictable memory management and zero-overhead abstractions.

In C++, you control when memory is allocated and freed. You can allocate a tensor buffer once and reuse it without the garbage collector deciding to reclaim it at an inopportune moment. In deep learning, where a single forward pass may allocate and free gigabytes of GPU memory, this control is critical.

C++ also compiles to native machine code. The compiler can inline function calls, unroll loops, and use SIMD instructions automatically. Python cannot do any of this because it must interpret bytecode at runtime.

PyTorch's core (libtorch) is entirely C++. The Python interface is generated via pybind11.

7.2 What pybind11 Does

pybind11 is a header-only C++ library that exposes C++ classes and functions to Python. It bridges the two type systems:

// C++ function
at::Tensor my_relu(const at::Tensor& input) {
    return input.clamp(0);
}

// pybind11 module definition
PYBIND11_MODULE(my_module, m) {
    m.def("relu", &my_relu,
          py::arg("input"),
          "Apply ReLU activation");
}

After building this as a .so file, Python can import and call it:

import my_module
y = my_module.relu(x)  # calls C++ function with Python tensor

pybind11 handles the conversion between Python objects (PyObject*) and C++ objects (at::Tensor) transparently. When you pass a PyTorch tensor from Python to the C++ function, pybind11 unwraps the Python object to get the underlying at::Tensor C++ object.

7.3 RAII — How C++ Manages Memory

C++ does not have garbage collection. Instead, it uses RAII (Resource Acquisition Is Initialization): resources (memory, file handles, GPU buffers) are tied to C++ object lifetimes. When an object goes out of scope, its destructor automatically releases resources.

{
    at::Tensor t = torch::randn({1024, 1024});  // GPU memory allocated
    // ... use t ...
}  // t goes out of scope → destructor runs → GPU memory freed (returned to pool)

This deterministic cleanup is why PyTorch can manage GPU memory so precisely. Python's garbage collector is non-deterministic — you don't know when del tensor will actually free memory. C++ destructors fire immediately when the object leaves scope.

Part III — PyTorch Architecture

Chapter 8: ATen, c10, and libtorch — The C++ Foundation

8.1 The Library Layers

PyTorch's C++ codebase is structured as three layers:

libtorch (the full C++ distribution)
    ├── ATen (A Tensor Library) — tensor operations
    └── c10 (Core 10 / Caffe2 + Torch)  — foundational types
            ├── Dispatcher
            ├── OperatorRegistry
            ├── IValue (universal value type)
            └── Storage, Device, Dtype

c10 is the foundation. It provides:

• c10::Dispatcher — the universal routing system (Chapter 10)
• c10::IValue — a type-erased container that holds any value (tensor, int, string, list, dict) — used to pass arguments through the dispatch system
• c10::Device — cpu:0, cuda:0, cuda:1, meta, etc.
• c10::ScalarType — float32, float16, int8, etc.
• c10::Storage / c10::StorageImpl — the physical memory block (Chapter 9)

ATen (A Tensor Library) sits on top of c10 and provides:

• at::Tensor — the C++ tensor class (the actual object Python's torch.Tensor wraps)
• ~1600 tensor operations defined in the aten:: namespace (aten::add, aten::mm, aten::relu, etc.)
• The schema system — each operation has a typed schema string used by the dispatcher

libtorch is the distributable package containing compiled ATen + c10 + autograd, exposing a C++ API (torch::Tensor, torch::nn::Module, etc.) that mirrors PyTorch's Python API.

8.2 The aten:: Namespace — Where Operations Live

Every PyTorch operation ultimately corresponds to a function in the aten:: namespace. When you call torch.relu(x) in Python:

1. Python's torch.relu is a Python wrapper generated from the ATen operator registry
2. It calls at::relu(x) in C++
3. at::relu is registered with the dispatcher under the key aten::relu
4. The dispatcher selects the correct implementation based on the tensor's dispatch keys

You can see the full schema of any operation:

print(torch.ops.aten.relu)
# aten::relu(Tensor self) -> Tensor

The schema string is not just documentation — it is parsed by the dispatcher at registration time to understand argument types, mutability, and aliasing. This is how PyTorch's compiler can reason about operations without running them.

Chapter 9: Tensors in PyTorch — Storage, Strides, and Memory Management

9.1 The C++ Object Behind torch.Tensor

Python's torch.Tensor is a thin Python wrapper around at::Tensor (C++). The C++ object contains:

struct TensorImpl {
    c10::Storage   storage_;      // the physical memory block
    int64_t        storage_offset_; // offset into storage in elements
    IntArrayRef    sizes_;         // shape, e.g., [3, 4]
    IntArrayRef    strides_;       // strides, e.g., [4, 1]
    c10::ScalarType dtype_;        // float32, int8, etc.
    c10::Device    device_;        // cpu, cuda:0, meta
    AutogradMeta*  autograd_meta_; // grad_fn, requires_grad, etc. (null if not needed)
    DispatchKeySet key_set_;       // bit-field of active dispatch keys
};

The DispatchKeySet is a 64-bit integer where each bit represents an active dispatch key. When the dispatcher receives an operation, it inspects the key_set_ of all input tensors (bitwise OR) to find the highest-priority active key.

9.2 The Caching Allocator — Why GPU Memory Isn't Really Freed

Calling cudaMalloc (NVIDIA's API to allocate GPU memory) is expensive: it may require the GPU driver to find a free region, update page tables, and sometimes synchronize with the GPU. A single cudaMalloc can take hundreds of microseconds — orders of magnitude longer than a fast CUDA kernel.

PyTorch's caching allocator (c10/cuda/CUDACachingAllocator.cpp) avoids this by maintaining a pool of already-allocated GPU memory:

When you create a CUDA tensor:
  1. Request size S from the allocator
  2. Allocator checks its free list for a block ≥ S
  3. If found: return that block (no cudaMalloc call — ~1μs)
  4. If not found: call cudaMalloc to expand the pool (~200μs)

When the tensor is deallocated:
  1. Reference count drops to zero
  2. Destructor runs: marks block as "free" in the allocator's pool
  3. Does NOT call cudaFree
  4. The memory stays in the pool for future allocations

Why torch.cuda.empty_cache() exists: it flushes the pool back to the OS via cudaFree. You'd call this before allocating a very large model to reclaim fragmented small blocks. Under normal operation, never call it — the caching is intentional.

The consequence you'll see in profiling: CUDA memory usage grows until it stabilizes. The allocator is holding memory it has already allocated to avoid future cudaMalloc overhead. This is not a leak — it's intentional pooling.

Chapter 10: The Dispatcher — PyTorch's Universal Routing System

10.1 Why the Dispatcher Exists

Before the dispatcher, PyTorch had a simpler design: each tensor operation had a set of if/else branches choosing the implementation based on tensor type:

// Old approach (simplified pseudocode):
Tensor relu(const Tensor& self) {
    if (self.is_cuda()) {
        return relu_cuda(self);
    } else if (self.is_cpu()) {
        return relu_cpu(self);
    } else {
        throw std::runtime_error("Unknown device");
    }
}

This broke every time someone wanted to add:

• A new device (ROCm GPU, XLA, NPU)
• A new feature (autograd, quantization, AMP)
• A new transformation (functorialize mutations, meta tensors)

Each addition required modifying every operation. With 1600+ operations, this was unmaintainable.

The dispatcher solves this with a table-driven routing system where:

• Operations are registered once with a schema
• Implementations are registered independently for each (operation, dispatch_key) pair
• The dispatcher selects the implementation at runtime without any if/else in the operation itself

10.2 The DispatchKey Enum

A DispatchKey is an integer in an enum. PyTorch has ~50 dispatch keys. The dispatcher maintains a table of size (num_operators × num_keys) mapping each pair to a function pointer.

The most important keys in priority order:

Highest priority:
  Functionalize   — converts in-place/aliasing ops to pure functional form.
                    TorchScript and torch.compile require functional semantics.
                    Without this, tracking aliasing through graph transformations
                    is intractable.

  InferenceMode   — completely disables gradient tracking.
                    Faster than no_grad for inference-only code.

  AutogradMeta    — handles FakeTensors (tensors with only shape/dtype, no data).
                    Used during torch.compile's shape analysis phase.

  Autograd        — THE gradient-recording layer. Before calling any compute
                    kernel, wraps the call to record it in the autograd graph.
                    Sets grad_fn on the output tensor.

  AutocastCPU/    — AMP (Automatic Mixed Precision). Intercepts calls and
  AutocastCUDA      promotes/demotes dtypes (e.g., matmul in float16, norms
                    in float32) without user code changes.

  CUDA            — NVIDIA GPU kernels. cuBLAS, cuDNN, custom CUDA code.
  CPU             — CPU kernels. Intel MKL, Eigen, vectorized C++.

  Meta            — Shape-only kernels. No data, no computation. Returns
                    tensors with correct shape/dtype but no storage.

Lowest priority:
  BackendSelect   — Selects which device backend to use (CPU vs CUDA vs ...)
                    based on input tensor devices.

10.3 A Complete Dispatch Trace

Call: torch.relu(x) where x is a CUDA tensor with requires_grad=True.

x.key_set_ has bits set for: {CUDA, Autograd}.

Dispatcher checks keys from highest to lowest:
  Functionalize? Bit not set. Skip.
  InferenceMode? Bit not set. Skip.
  AutogradMeta? Bit not set. Skip.
  Autograd? Bit IS set → call Autograd implementation.

Autograd implementation of relu:
  1. Calls torch::autograd::generated::ReluBackward::apply_with_saved(...)
  2. This wraps the call: records that relu ran, creates ReluBackward0 grad_fn
  3. Re-dispatches DOWN, excluding the Autograd key:
     new_key_set = x.key_set_ & ~DispatchKeySet{Autograd}
     → {CUDA}
  4. Calls CUDA implementation of relu

CUDA implementation of relu:
  1. Launches CUDA kernel: relu_kernel<<<grid, block>>>(x.data_ptr<float>(), ...)
  2. Returns output tensor

Back in Autograd:
  output.grad_fn = ReluBackward0  (records: "relu was applied to x")
  return output

Total: one Python call → one Autograd interception (sets grad_fn) → one CUDA kernel launch.

10.4 Registering a New Device Backend

This is the extension point for NPU vendors like Qualcomm:

// Register a new PrivateUse1 key for a hypothetical NPU
TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
    m.impl("relu", &npu_relu_impl);
    m.impl("mm", &npu_mm_impl);
    // ... etc. for all supported ops
}

A tensor on the NPU device has PrivateUse1 in its key_set_. The dispatcher routes to npu_relu_impl automatically. The Autograd key still sits above it, so gradient recording works automatically unless explicitly disabled.

The ordering matters: PrivateUse1 is below Autograd. If a vendor incorrectly registered above Autograd, gradient recording would be silently bypassed — a catastrophic bug that would only surface during training.

Chapter 11: The Autograd Engine — Building and Traversing the Backward Graph

11.1 The Data Structures

PyTorch's autograd engine (in torch/csrc/autograd/) is built around these C++ classes:

torch::autograd::Node (also called Function): represents one differentiable operation in the graph. Stores:

• next_edges_: a list of Edge objects pointing to upstream nodes
• input_metadata_: shape/dtype info for gradient validation
• Virtual method apply(variable_list&& inputs) — the backward function

torch::autograd::Edge: a directed edge in the backward graph. Contains:

• function: pointer to the upstream Node
• input_nr: which input of that Node this edge corresponds to

torch::autograd::AccumulateGrad: a leaf Node that accumulates gradients into a leaf tensor's .grad field. Every requires_grad=True leaf tensor has one.

AutogradMeta: attached to every tensor that participates in autograd. Contains:

• grad_fn_: the Node that produced this tensor (null for leaves)
• grad_: accumulated gradient (null until backward is called)
• requires_grad_: whether this tensor needs gradients
• output_nr_: which output of grad_fn_ this tensor is

11.2 DAG Construction During Forward Pass

When the Autograd dispatch key intercepts an operation, it:

1. Calls the compute kernel (CPU/CUDA) to get the output values
2. Creates a new Node subclass specific to this operation (e.g., MmBackward0 for matmul)
3. Saves any tensors the backward will need (via save_for_backward)
4. Sets next_edges_ to point to the grad_fn_ of each input (or AccumulateGrad for leaves)
5. Attaches this new Node as the grad_fn_ of the output tensor

After forward pass of:
  c = a @ b    (where a, b require_grad)

Memory layout:
  a → AutogradMeta { grad_fn=null, requires_grad=true }
  b → AutogradMeta { grad_fn=null, requires_grad=true }
  c → AutogradMeta { grad_fn=MmBackward0 }
         MmBackward0 {
           next_edges: [AccumulateGrad(a), AccumulateGrad(b)]
           saved_tensors: [a, b]   ← for computing grad_a = grad_c @ b.T
         }

11.3 The Backward Pass: Engine::execute()

c.backward() calls torch::autograd::Engine::execute(). The C++ engine:

1. Creates a priority queue (or thread pool) of nodes to process
2. Initializes the root node (c.grad_fn = MmBackward0) with grad = ones_like(c)
3. Processes nodes in reverse topological order:
   • For each node, calls node.apply(incoming_gradients) → returns gradient tuples
   • For each output gradient, sends it to the corresponding next_edge.function
   • When a function receives gradients from ALL its outputs (tracked by a counter), it's ready to run
4. When AccumulateGrad runs, adds incoming gradient to leaf.grad

The readiness counter is what implements topological ordering. Each node tracks how many downstream outputs still need to send it a gradient. A node runs only when this counter reaches zero — i.e., all downstream contributions have arrived.

Initial state:
  MmBackward0.ready_count = 1  (one consumer: the starting loss gradient)

After loss.backward(grad=1.0):
  Engine sends 1.0 to MmBackward0
  MmBackward0.ready_count decrements to 0 → runs now

MmBackward0.apply([grad_c]):
  grad_a = grad_c @ b.T    (saved b from forward)
  grad_b = a.T @ grad_c    (saved a from forward)
  Returns [grad_a, grad_b]

Engine sends grad_a to AccumulateGrad(a) → a.grad += grad_a
Engine sends grad_b to AccumulateGrad(b) → b.grad += grad_b

11.4 Why torch.no_grad() Is Not Just "Disabling Gradients"

When you enter torch.no_grad(), PyTorch sets GradMode::is_enabled() to false. The Autograd dispatch key's code checks this flag before creating Function objects:

// In the autograd dispatch implementation:
if (!GradMode::is_enabled()) {
    // Skip all autograd recording, call compute kernel directly
    return at::AutoDispatchBelowADInplaceOrView::redispatch(...);
}

What is NOT allocated when no_grad() is active:

• No Node objects (no MmBackward0, etc.)
• No AutogradMeta on output tensors (or stripped-down version)
• No saved tensors (the inputs to the backward function)
• No next_edges_ (the DAG edges)

For a large transformer forward pass, this can be gigabytes of saved activations that are simply never allocated. The speedup comes from: (1) no memory allocation for the graph, (2) no reference counting overhead for saved tensors, (3) no function dispatch overhead for the autograd wrapper.

Chapter 12: torch.autograd.Function — Custom Differentiable Operations

12.1 When and Why to Write One

PyTorch knows how to differentiate through all built-in operations. But sometimes you need to:

• Implement an operation that doesn't exist in PyTorch (a custom NPU intrinsic)
• Override the default backward with a more numerically stable or efficient version
• Implement a non-standard gradient (e.g., straight-through estimator for quantization)

torch.autograd.Function lets you inject a custom (forward, backward) pair into the autograd graph as a single opaque node.

12.2 The API, Fully Explained

class MySoftmax(torch.autograd.Function):

    @staticmethod
    def forward(ctx, x, dim):
        # ctx is a FunctionContext object. It is the storage for the backward.
        # It is NOT a Python object you create — it's provided by the engine.
        #
        # RULE: Never store tensors directly as ctx attributes.
        #       Use ctx.save_for_backward() for tensors.
        #       Use ctx.dim = dim for non-tensor metadata.
        
        x_max = x.max(dim=dim, keepdim=True).values
        x_shift = x - x_max                  # numerically stable: all values ≤ 0
        exp_x = x_shift.exp()
        sum_exp = exp_x.sum(dim=dim, keepdim=True)
        out = exp_x / sum_exp
        
        ctx.save_for_backward(out)            # saves output (not input) for backward
        ctx.dim = dim                         # saves non-tensor metadata as attribute
        return out

    @staticmethod
    def backward(ctx, grad_output):
        # grad_output: the gradient of the loss w.r.t. the output of forward()
        #              shape matches the output of forward()
        
        out, = ctx.saved_tensors              # retrieves saved tensors
        dim = ctx.dim
        
        # Softmax Jacobian-vector product:
        # ∂L/∂x_i = out_i × (∂L/∂y_i - Σ_j (∂L/∂y_j × out_j))
        dot = (grad_output * out).sum(dim=dim, keepdim=True)
        grad_x = out * (grad_output - dot)
        
        # RULE: return one gradient per argument to forward().
        # forward() took (ctx, x, dim). dim is a Python int (non-differentiable).
        # Return grad_x for x, and None for dim.
        return grad_x, None

# Usage:
result = MySoftmax.apply(x, dim=-1)   # Note: .apply(), not direct call

12.3 The Memory Management Rule: Why save_for_backward

If you write ctx.my_tensor = x instead of ctx.save_for_backward(x):

1. PyTorch's memory-efficient gradient checkpointing cannot track this tensor
2. The version counter check (detecting in-place mutations) is bypassed — you may compute wrong gradients silently if x is modified in-place after the forward
3. The tensor is not freed when the backward graph is freed (memory leak)

ctx.save_for_backward registers the tensor with the engine. The engine:

• Tracks the version counter at save time
• Checks the version counter at backward time (raises error if tensor was modified in-place)
• Frees the saved tensors as soon as backward() processes this node (unless retain_graph=True)

12.4 The Straight-Through Estimator — A Real Production Use Case

Quantization maps a float value to a small integer (e.g., float32 → int8). The quantize-dequantize operation:

def quantize(x, scale, zero_point):
    return torch.clamp(torch.round(x / scale) + zero_point, 0, 255)

The gradient of round() is zero almost everywhere (it's a staircase). This means no gradient flows back through quantized layers — training with quantization would not work.

The Straight-Through Estimator (STE): pretend in the backward pass that quantization didn't happen — pass gradients straight through as if the operation were an identity. This is physically inaccurate but empirically works for training quantized models.

class QuantizeSTE(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, scale, zero_point):
        # Real quantization in forward
        x_int = torch.clamp(torch.round(x / scale) + zero_point, 0, 255)
        return (x_int - zero_point) * scale   # dequantize

    @staticmethod
    def backward(ctx, grad_output):
        # Straight-through: gradient passes as if forward were identity
        return grad_output, None, None

This pattern is used in every QAT (Quantization-Aware Training) framework — AIMET, PyTorch native QAT, and custom NPU training toolchains. You will implement exactly this in Phase 03.

Part IV — Graph Capture and Compilation

Chapter 13: torch.fx — Intermediate Representation and Graph Transformation

13.1 What Is an Intermediate Representation?

In compiler design, an intermediate representation (IR) is a data structure that sits between the source language (Python) and the target (native code). Good IRs are:

• Analyzable: easy to traverse and inspect programmatically
• Transformable: easy to modify (insert, delete, reorder nodes)
• Lowerable: can be compiled/interpreted efficiently

Famous IRs: LLVM IR (used by Clang/Rust), JVM bytecode (used by Java/Kotlin), WebAssembly. torch.fx is PyTorch's ML-specific IR.

13.2 The Problem FX Solves

Before FX, if you wanted to fuse conv2d → batch_norm → relu into one kernel:

• You'd have to modify PyTorch's C++ source code
• Or hook into Python's metaclass system in fragile ways
• There was no clean programmatic way to inspect and modify a model's computation

FX provides a first-class Python object representing the computation — a GraphModule containing a Graph of Node objects. You can write Python code that traverses the graph, finds patterns, and rewrites them.

13.3 How Symbolic Tracing Works — The Proxy Mechanism

torch.fx.symbolic_trace(model) does something clever: it runs the model's forward() function, but instead of passing real tensors, it passes Proxy objects.

A Proxy is a Python object that:

• Behaves like a tensor (has .shape, .dtype, .device)
• Records every operation performed on it instead of executing anything
• Returns a new Proxy representing the output of each operation

class Proxy:
    def __add__(self, other):
        # Don't compute! Record that addition happened.
        node = self.graph.create_node(
            op='call_function',
            target=operator.add,
            args=(self, other)
        )
        return Proxy(node, self.graph)
    
    def __matmul__(self, other):
        # Same pattern for all operations
        node = self.graph.create_node(...)
        return Proxy(node, self.graph)

When model.forward(proxy_input) executes:

• self.linear(proxy_input) → records call_module[linear]
• torch.relu(intermediate_proxy) → records call_function[relu]
• return proxy.sum() → records call_method[sum]

After the function returns, the graph object contains all recorded nodes. FX wraps it in a GraphModule that re-executes the graph using real tensors.

13.4 The Four Node Types and Their Fields

Every node has:

• node.op: one of placeholder, call_function, call_module, call_method, output
• node.name: a unique string identifier
• node.target: what is being called (a function, module name, or method name)
• node.args: positional arguments (may reference other nodes)
• node.kwargs: keyword arguments
• node.users: set of downstream nodes that consume this node's output

# After tracing a model:
for node in graph.nodes:
    if node.op == 'call_function' and node.target == torch.relu:
        # Access predecessors:
        for arg in node.args:
            if isinstance(arg, Node):
                print(f"{arg.name} → relu → {node.name}")
        # Access successors:
        for user in node.users:
            print(f"relu output used by: {user.name}")

13.5 Transforming Graphs: The Pattern-Replace Loop

The core graph transformation pattern:

def insert_quantization_observers(model):
    traced = torch.fx.symbolic_trace(model)
    
    # Collect nodes to transform (can't modify graph while iterating)
    nodes_to_wrap = []
    for node in traced.graph.nodes:
        if node.op == 'call_module':
            module = traced.get_submodule(node.target)
            if isinstance(module, (torch.nn.Conv2d, torch.nn.Linear)):
                nodes_to_wrap.append(node)
    
    # Apply transformations
    for node in nodes_to_wrap:
        with traced.graph.inserting_after(node):
            # Insert an observer node after each conv/linear
            observer_name = node.name + '_observer'
            traced.add_module(observer_name, MinMaxObserver())
            new_node = traced.graph.call_module(observer_name, (node,))
        
        # Replace all uses of node with new_node
        node.replace_all_uses_with(new_node)
        # Move new_node to immediately after node in the graph
        new_node.args = (node,)
    
    traced.recompile()  # regenerate Python code from modified graph
    return traced

This pattern is exactly how AIMET, PyTorch's quantize_fx, and torch.compile's optimization passes work.

13.6 FX's Fundamental Limitations

Limitation 1: Data-dependent control flow is opaque.

def forward(x):
    n = x.shape[0]          # Python int — OK
    if x.sum() > 0:          # x.sum() returns a Proxy, > 0 raises TraceError
        return x * 2
    return x + 1

The Proxy for x.sum() is a symbolic object with no concrete value. Python's if statement calls bool() on it, which raises TraceError: Tried to use FX tracing to trace a dynamic value.

Limitation 2: Shapes are concrete at trace time.

If you trace with x.shape = (4, 128), shape information is baked into the graph. Running on x.shape = (8, 128) may work (if no shapes are used as constants) or silently produce wrong results (if shapes were used in computation).

Why torch.compile avoids these limitations: TorchDynamo operates at the bytecode level. It sees the raw COMPARE_OP bytecode instruction for x.sum() > 0 and handles it differently — it evaluates the condition with the actual tensor (breaking the graph if needed) rather than trying to symbolically trace through it.

Chapter 14: torch.compile — The Full Compilation Stack

14.1 What Problem torch.compile Solves

A pure eager PyTorch model executes one operation at a time:

1. Python calls relu(x) → dispatch → CUDA kernel launch
2. Python calls linear(x) → dispatch → cuBLAS call
3. Python calls layer_norm(x) → dispatch → another CUDA kernel

Each kernel launch has overhead: Python→C++ dispatch (~1μs), CUDA launch latency (~5μs), GPU memory round-trip for intermediate results. For a transformer layer with 100+ operations, these overheads accumulate.

torch.compile captures a sequence of operations as a static graph and compiles it to optimized native code:

• Fuses relu + linear into one kernel that reads input once
• Eliminates intermediate tensor allocations between fused ops
• Schedules computation to maximize GPU utilization

14.2 TorchDynamo — Bytecode Interception

TorchDynamo installs a Python frame evaluation callback via CPython's _PyEval_EvalFrameDefault override API (PEP 558). This is a CPython hook that lets C code intercept bytecode execution.

When Python is about to execute a decorated function:

Normal Python execution:
  Python → CPython eval loop → execute bytecodes

With TorchDynamo:
  Python → Dynamo's callback fires
         → Dynamo reads bytecodes using dis/bytecode analysis
         → Classifies each instruction: tensor op? Python op? control flow?
         → Extracts tensor operations into an FX graph
         → For unsupported ops: breaks graph, falls back to Python
         → Compiles captured graph via AOTAutograd + Inductor
         → Installs guards (fast runtime checks)
         → On future calls: checks guards, dispatches to compiled artifact

Guards are the key innovation. When Dynamo captures a graph, it records every assumption it made:

• "x.shape[0] == 4"
• "x.dtype == torch.float32"
• "model.training == False"

These become fast Python checks on each call:

# Generated guard code (simplified):
def check_guards(x, model):
    return (x.shape[0] == 4 and
            x.dtype == torch.float32 and
            not model.training)

If any guard fails, Dynamo recompiles for the new configuration. The compiled artifacts are cached by guard configuration, so "x.shape[0] == 4 and float32" and "x.shape[0] == 8 and float32" are two separate compiled graphs, both cached.

14.3 AOTAutograd — Capturing the Backward Graph

By default, PyTorch builds the backward graph lazily during the forward pass. For compilation, we need the backward graph ahead of time so it can be optimized along with the forward.

AOTAutograd (Ahead-Of-Time Autograd) traces both forward and backward simultaneously:

1. Runs the forward with FakeTensors (no data, just shapes/dtypes)
2. Runs autograd on the traced forward to symbolically derive the backward
3. Produces a joint FX graph containing both forward and backward as one graph
4. Applies functionalization — converts all in-place ops and aliasing ops into pure functional equivalents (required for correct compilation)

The joint graph is then split into a forward graph (returned to the caller) and a backward graph (stored for when .backward() is called).

14.4 TorchInductor — Code Generation

Inductor takes the FX graph from AOTAutograd and generates code:

For NVIDIA GPUs: generates Triton code. Triton is a Python-like DSL for writing GPU kernels. It handles thread indexing, memory coalescing, and warp management automatically — you write scalar code; Triton vectorizes it.

# What Inductor might generate for relu + add:
@triton.jit
def relu_add_kernel(x_ptr, y_ptr, out_ptr, n: tl.constexpr):
    pid = tl.program_id(0)
    block = tl.arange(0, n)
    x = tl.load(x_ptr + pid * n + block)
    y = tl.load(y_ptr + pid * n + block)
    out = tl.maximum(x, 0) + y          # fused relu + add in one pass
    tl.store(out_ptr + pid * n + block, out)

Fusion: instead of launching two kernels (one for relu, one for add), Inductor fuses them into one. The intermediate result (after relu) stays in GPU registers — no round-trip to global memory.

For CPUs: generates C++ with OpenMP parallelism and AVX vectorization hints.

14.5 The Meta Device — Shape Analysis Without Data

Before generating any kernel, Inductor must know the output shape of every operation. It uses the Meta device for this:

A Meta tensor has shape, dtype, device = "meta", but no storage — no memory allocated, no bytes of data. Operations on Meta tensors execute "Meta kernels" that compute output shapes without performing any arithmetic.

x = torch.empty(4, 128, device='meta')   # 0 bytes allocated
y = torch.nn.Linear(128, 64)(x)          # Meta kernel: (4,128) @ (128,64) → (4,64)
print(y.shape)   # torch.Size([4, 64])
print(y.numel()) # 256 — but no data exists

Every custom operator MUST register a Meta kernel. Without one, torch.compile cannot propagate shapes and will raise:

NotImplementedError: Could not run 'myops::my_op' with Meta key.

The Meta kernel need only return at::empty_like(input) for element-wise ops, or compute the correct output shape for reshape/reduction ops.

Chapter 15: GPU Architecture — From Silicon to CUDA Kernels

15.1 Why GPUs Exist: The Throughput vs Latency Tradeoff

A modern CPU (e.g., Intel Xeon with 32 cores) is optimized for latency — executing a single complex instruction stream as fast as possible. It has:

• Large caches (L1: 32KB, L2: 256KB, L3: 16MB per core)
• Branch prediction hardware (avoids stalls from if/else)
• Out-of-order execution (reorders instructions for efficiency)
• Speculative execution (pre-executes likely branches)

This complexity lets a CPU execute a 10,000-instruction sequential program in microseconds. But it scales poorly to parallel workloads.

A GPU (e.g., NVIDIA A100) is optimized for throughput — executing thousands of identical operations simultaneously. It has:

• 6912 CUDA cores (simple integer ALUs)
• 432 Tensor Cores (4×4 matrix multiply units)
• Minimal caching (L1: 192KB/SM, L2: 40MB shared)
• No branch prediction, no out-of-order execution per core
• Massive memory bandwidth: 2 TB/s (vs CPU's ~100 GB/s)

The tradeoff: a GPU cannot run a sequential 10,000-instruction program faster than a CPU. But running the same operation on 1 million data points is 100× faster on GPU because all 1 million operations run simultaneously.

15.2 The SM — Streaming Multiprocessor

The A100 has 108 Streaming Multiprocessors (SMs). Each SM is an independent processing unit containing:

One SM:
  ├── 4 Warp Schedulers     (each schedules 32-thread warps)
  ├── 64 INT32 units        (integer arithmetic)
  ├── 64 FP32 units         (floating-point arithmetic)
  ├── 32 FP64 units         (double precision)
  ├── 4 Tensor Cores        (4×4 matrix multiply in one cycle)
  ├── 192 KB L1 cache / shared memory (configurable split)
  └── Register file: 256 KB (65,536 × 32-bit registers)

An SM can hold up to 2048 threads in flight simultaneously ("resident threads"). It switches between warps each clock cycle — when one warp stalls waiting for memory, the SM immediately executes another warp. This latency hiding is how GPUs tolerate their slower memory system.

15.3 The CUDA Execution Model: SIMT

CUDA uses a programming model called SIMT (Single Instruction, Multiple Threads). You write code for one thread; the hardware runs that same code on thousands of threads simultaneously.

Thread:  The smallest unit. Runs one instance of your kernel.
Warp:    32 threads that execute the SAME instruction simultaneously.
         This is the fundamental dispatch unit of the GPU.
Block:   1–1024 threads grouped together. Threads in a block:
           - Run on the same SM
           - Share the SM's shared memory
           - Can synchronize with __syncthreads()
Grid:    Many blocks. Blocks run independently across all SMs.

When you launch a kernel:

my_kernel<<<grid_dim, block_dim, shared_mem, stream>>>(args...);

grid_dim × block_dim total threads are created. The GPU scheduler distributes blocks across SMs. Within each SM, threads are grouped into warps of 32.

15.4 Warp Execution — Why 32 Threads Must Stay Synchronized

A warp's 32 threads share one program counter. At each clock cycle, the warp executes one instruction, and all 32 threads execute it in parallel (on different data).

Warp divergence: if threads in a warp take different branches:

if (thread_id % 2 == 0) {
    // Even threads do this
    x = x * 2.0f;
} else {
    // Odd threads do this
    x = x + 1.0f;
}

The GPU CANNOT execute both branches simultaneously. It:

1. Executes x = x * 2.0f with even threads active, odd threads masked off
2. Executes x = x + 1.0f with odd threads active, even threads masked off

Result: 2× longer execution. Each half of the warp idles while the other executes. Minimizing warp divergence is a primary GPU optimization concern.

15.5 The Memory Hierarchy — The Critical Performance Factor

The golden rule: every read from global memory costs ~800 cycles. If your kernel reads each piece of data once and computes many operations on it, global memory cost is amortized. If your kernel reads data once, computes one operation, and writes back — you're memory-bound: the bottleneck is memory, not compute.

FlashAttention's breakthrough was recognizing that the standard attention algorithm is memory-bound (reads Q, K, V matrices multiple times) and restructuring the computation to keep everything in L1/shared memory using tiling.

15.6 Memory Coalescing — Why Access Pattern Matters

When 32 threads in a warp each read one float from global memory, the GPU issues a memory transaction. If the 32 addresses are contiguous (consecutive elements of an array), the transaction fetches one 128-byte cache line — one transaction total.

If the 32 addresses are strided (e.g., every other element) or random (scattered), the GPU must issue multiple transactions — 2× or 32× slower respectively.

Coalesced access (good):
  Thread 0: address 0×4 = float at byte 0
  Thread 1: address 1×4 = float at byte 4
  Thread 2: address 2×4 = float at byte 8
  ...
  Thread 31: address 31×4 = float at byte 124
  → One 128-byte transaction covers all 32 reads.

Strided access (bad):
  Thread 0: address 0×8 = float at byte 0
  Thread 1: address 1×8 = float at byte 8
  ...every other float, stride=2
  → Two transactions, 50% bandwidth wasted

This is why matrix transposition before matmul is sometimes applied — to ensure the memory access pattern in the CUDA kernel is coalesced.

15.7 Warp Shuffle Instructions — Register-Level Communication

Normally, threads share data via shared memory (write → barrier → read, ~5 cycles). Warp shuffle instructions (__shfl_*) let threads exchange register values directly within a warp without touching memory — ~1 cycle.

Available shuffles:

__shfl_sync(mask, val, src_lane)     // all threads get val from lane src_lane
__shfl_up_sync(mask, val, delta)     // get val from lane (current - delta)
__shfl_down_sync(mask, val, delta)   // get val from lane (current + delta)
__shfl_xor_sync(mask, val, laneMask) // get val from lane (current XOR laneMask)

mask = 0xffffffff means all 32 lanes participate.

The butterfly reduction using __shfl_xor_sync:

Goal: find the sum of values across all 32 threads.

Round 1 (laneMask=16):
  Lane 0 gets lane 16's value. Lane 0 now holds sum of elements 0 and 16.
  Lane 1 gets lane 17's value. Lane 1 now holds sum of elements 1 and 17.
  ... (simultaneously for all 32 pairs)

Round 2 (laneMask=8):
  Lane 0 gets lane 8's value. Lane 0 now holds sum of elements 0, 8, 16, 24.
  ...

Round 3 (laneMask=4): Each thread holds sum of 8 elements.
Round 4 (laneMask=2): Each thread holds sum of 16 elements.
Round 5 (laneMask=1): Each thread holds sum of all 32 elements.

5 rounds = log₂(32) = 5 communication steps. Every thread simultaneously holds the global sum. No shared memory, no barriers.

15.8 The warp_softmax Kernel Step by Step

Lab 02's kernel computes softmax over a 2D tensor [rows, vocab_size] where each warp handles one row:

Tensor: [batch × seq_len, vocab_size] = [rows, V]
  Lane k handles elements: k, k+32, k+64, ..., k+32*(V/32-1)
  Each lane processes V/32 elements (1000 for vocab=32000)

Step 1 — Find row max (prevents exp() overflow):
  Each lane: local_max = max(input[k], input[k+32], ..., input[k+32*(V/32-1)])
  Butterfly reduction: all lanes now hold row_max

Step 2 — Compute exp(x - row_max) and accumulate sum:
  Each lane: for i in range(k, V, 32):
               val = exp(input[i] - row_max)
               output[i] = val
               local_sum += val
  Butterfly reduction: all lanes now hold row_sum

Step 3 — Normalize:
  Each lane: inv_sum = 1.0 / row_sum   (multiply cheaper than divide)
             for i in range(k, V, 32): output[i] *= inv_sum

Three passes over the data. All intermediate values (row_max, partial sums) live in registers — no shared memory needed. This is why the kernel is fast: minimal memory traffic, maximum parallelism.

Chapter 16: Custom C++ Operators — TORCH_LIBRARY

16.1 The Two-Level Registration System

PyTorch separates what an operator is from how it's implemented on each device. This separation is the dispatcher's core design.

Level 1: Schema registration — defines the operator's name and type signature, visible to all backends:

TORCH_LIBRARY(myops, m) {
    // Schema language:
    //   "namespace::name(arg_type arg_name, ...) -> return_type"
    // Arg types: Tensor, Tensor?, Tensor[], int, int[], float, bool,
    //            ScalarType, Device, str, Scalar
    m.def("warp_softmax(Tensor self, int dim) -> Tensor");
}

This creates torch.ops.myops.warp_softmax in Python. The operator exists in the registry. Calling it without an implementation registered for the current dispatch key will error.

Level 2: Implementation registration — binds implementations to specific dispatch keys:

// CPU implementation
TORCH_LIBRARY_IMPL(myops, CPU, m) {
    m.impl("warp_softmax", &warp_softmax_cpu);   // function pointer
}

// CUDA implementation  
TORCH_LIBRARY_IMPL(myops, CUDA, m) {
    m.impl("warp_softmax", &warp_softmax_cuda);
}

// Meta implementation — REQUIRED for torch.compile
TORCH_LIBRARY_IMPL(myops, Meta, m) {
    m.impl("warp_softmax", &warp_softmax_meta);   // returns empty_like(input)
}

16.2 How Registration Actually Works at Runtime

TORCH_LIBRARY and TORCH_LIBRARY_IMPL expand to static initializers — C++ code that runs when the .so file is loaded. They call into c10::Dispatcher::singleton() to register the operator:

// What TORCH_LIBRARY expands to (simplified):
static auto registration = c10::RegisterOperators()
    .op("myops::warp_softmax(Tensor self, int dim) -> Tensor",
        c10::DispatchKey::CPU, &warp_softmax_cpu)
    .op("myops::warp_softmax(Tensor self, int dim) -> Tensor",
        c10::DispatchKey::CUDA, &warp_softmax_cuda);

When your .so is loaded via torch.ops.load_library("warp_softmax_ext.so"), these static initializers run, and the operator is available immediately.

16.3 The Two Autograd Paths

Path A — Composite (automatic differentiation):

If your forward kernel is built from differentiable ATen ops (like the CPU kernel in Lab 02 that uses .amax(), .exp(), .sum(), /), PyTorch differentiates through them automatically. The Autograd key uses CompositeImplicitAutograd — it traces through your implementation to find the gradient.

Path B — Explicit autograd (manual backward):

If your CUDA kernel is not built from differentiable ATen ops, you must register an explicit backward:

torch::autograd::Function<SoftmaxBackward>::apply = [](
    const at::Tensor& grad_output,
    const at::Tensor& output) -> at::Tensor {
    // Softmax backward: grad_input = output * (grad - (grad * output).sum(-1, true))
    auto dot = (grad_output * output).sum(-1, true);
    return output * (grad_output - dot);
};

TORCH_LIBRARY_IMPL(myops, Autograd, m) {
    m.impl("warp_softmax", torch::autograd::autogradNotImplementedFallback());
    // OR: m.impl("warp_softmax", &warp_softmax_autograd);
}

16.4 OpCheck — The Complete Validation Suite

torch.library.opcheck tests your operator against all dispatch keys automatically:

import torch

torch.ops.load_library("warp_softmax_ext.so")

# Test with a simple input
x = torch.randn(4, 128, dtype=torch.float32, requires_grad=True)
if torch.cuda.is_available():
    x = x.cuda()

results = torch.library.opcheck(
    torch.ops.myops.warp_softmax,
    args=(x, -1),
    kwargs={},
)
# Raises if any dispatch key implementation is wrong

What OpCheck verifies:

1. CPU/CUDA correctness: output matches reference (F.softmax) within 1e-5
2. Autograd correctness: torch.autograd.gradcheck on differentiable inputs
3. Meta kernel: shape propagation with device='meta' tensors
4. FakeTensor mode: compatibility with torch.compile's shape tracking
5. Out-of-place invariant: operator does not modify its inputs in-place unexpectedly

If OpCheck passes, your operator is safe in all contexts: eager mode, training, inference, torch.compile, quantization pipelines.

Part V — TensorFlow

Chapter 17: TF1 to TF2 — Two Execution Philosophies

17.1 The TF1 Philosophy: Define-Then-Run

TensorFlow 1.x was designed by Jeff Dean and the Google Brain team in 2015, influenced by Theano and the needs of large-scale production ML.

The core insight: if you separate graph definition from graph execution, you can:

• Ship the graph (a serialized Protocol Buffer file) to mobile devices, servers, or TPUs
• Optimize the graph globally before running it (constant folding, CSE, operation fusion)
• Distribute it across machines without re-serializing Python code
• Cache the compiled form and reuse it for millions of inference requests

The penalty: you cannot inspect intermediate values during execution. The graph runs entirely in C++ without Python touching it.

# TF1 workflow
import tensorflow as tf

# PHASE 1: Build the graph (nothing executes)
x = tf.placeholder(tf.float32, shape=[None, 784], name='input')
W = tf.Variable(tf.random.normal([784, 256]), name='weights')
b = tf.Variable(tf.zeros([256]), name='bias')
h = tf.nn.relu(tf.matmul(x, W) + b, name='hidden')
loss = tf.reduce_mean(tf.square(h - target), name='loss')
optimizer = tf.train.AdamOptimizer(lr=0.001)
train_op = optimizer.minimize(loss)

# PHASE 2: Execute (everything runs in C++)
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for step in range(1000):
        batch_x, batch_y = get_batch()
        loss_val, _ = sess.run([loss, train_op], {x: batch_x, target: batch_y})

The graph is a Protocol Buffer — Google's binary serialization format. The graph can be saved with tf.saved_model.save() and loaded on any device with the TF runtime, including mobile phones (TFLite) and Google TPUs.

17.2 TF2's Eager Execution: Define-By-Run

Python programmers found TF1's two-phase model confusing and hard to debug. TF2 (released 2019) switched to eager execution by default:

# TF2 eager — operations execute immediately
import tensorflow as tf

W = tf.Variable(tf.random.normal([784, 256]))
b = tf.Variable(tf.zeros([256]))
x = tf.constant(batch_x)               # a real tensor, not a placeholder

h = tf.nn.relu(x @ W + b)              # runs NOW, returns a real tensor
print(h.shape)                          # (batch_size, 256) — immediate
print(h.numpy()[:2])                    # can inspect values directly

Like PyTorch, operations execute immediately and return concrete tensors. Debugging is trivial — add print() statements anywhere.

The cost: no static optimization, no graph export. Each Python call has interpreter overhead.

The bridge: @tf.function re-introduces graph compilation as an opt-in feature, giving you the best of both worlds.

Chapter 18: tf.function — AutoGraph, Tracing, and ConcreteFunctions

18.1 What Happens When Python Hits @tf.function

The first time you call a @tf.function-decorated function, TF runs a multi-stage pipeline:

Stage 1 — AutoGraph transformation: TF rewrites your Python source code. Python control flow structures are converted to TF graph ops:

This transformation happens at the Python source level — TF uses Python's inspect module to get the source code, then rewrites it using the gast library (Generic AST).

Stage 2 — Tracing: The AutoGraph-transformed function runs with SymbolicTensor objects (analogous to PyTorch's Proxy objects). Every TF op call records a node in a FuncGraph — TF's internal graph representation.

Stage 3 — Shape inference: TF propagates shapes through all nodes in the FuncGraph. Operations that don't know their output shape at trace time are marked with unknown dimensions.

Stage 4 — Compilation to ConcreteFunction: The FuncGraph is compiled into a ConcreteFunction — a C++ execution artifact bound to a specific input signature. XLA may further compile this to hardware-specific machine code.

Stage 5 — Caching: The ConcreteFunction is cached. Subsequent calls with matching signatures (same shapes and dtypes) route to the cached artifact — no Python execution, no tracing overhead.

18.2 The Complete List of Retrace Triggers

A retrace is expensive (re-runs AutoGraph + tracing + compilation). You must understand every trigger:

Trigger 1: New input shape (without input_signature)

@tf.function
def f(x): return x * 2

f(tf.ones([3]))           # trace 1: shape (3,)
f(tf.ones([4]))           # RETRACE: shape (4,) is new

Prevention: @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.float32)])

Trigger 2: New input dtype

f(tf.ones([3], dtype=tf.float32))  # trace 1
f(tf.ones([3], dtype=tf.float64))  # RETRACE: different dtype

Trigger 3: Change in Python non-tensor arguments

@tf.function
def f(x, n_heads):          # n_heads is Python int, not tf.Tensor
    return tf.split(x, n_heads, axis=-1)

f(x, n_heads=4)   # trace 1: n_heads=4 baked in as compile-time constant
f(x, n_heads=8)   # RETRACE: n_heads=8 requires new graph structure

AutoGraph needs Python values at trace time to build the graph (how many splits to create). Changing a Python value creates a different graph.

Trigger 4: Changed Python closure

threshold = 0.5
@tf.function
def f(x): return x > threshold  # captures Python variable 'threshold'

threshold = 0.9
f(x)  # RETRACE: closure value changed

Trigger 5: tf.Variable created inside the function (error, not retrace)

@tf.function
def bad():
    v = tf.Variable(0.0)  # FORBIDDEN after first trace
    return v
# Second call raises: ValueError: tf.function only supports tf.Variables
# created on the first call.

Prevention pattern: always create Variables outside @tf.function:

v = tf.Variable(0.0)        # created once, outside
@tf.function
def good(x):
    v.assign_add(x)          # use the existing Variable
    return v

18.3 AutoGraph's Transformation in Detail

Let's trace exactly what AutoGraph does to a loop:

# Original Python:
@tf.function
def sum_to_n(n):
    total = tf.constant(0)
    for i in tf.range(n):      # tf.range, not Python range!
        total = total + i
    return total

AutoGraph rewrites this to (approximately):

def sum_to_n(n):
    total = tf.constant(0)
    # Loop body as a separate function
    def loop_body(i, total):
        total = total + i
        return total
    # TF while_loop handles the actual iteration in C++
    total = tf.while_loop(
        cond=lambda i, _: i < n,
        body=loop_body,
        loop_vars=[tf.constant(0), total]
    )[1]
    return total

tf.while_loop is a TF graph op that encodes the loop entirely within the graph. The loop runs in C++ without returning to Python. This is how TF achieves Python-free execution of loops.

Note: if you use Python range(n) instead of tf.range(n), AutoGraph unrolls the loop at trace time — creating n nodes in the graph. For variable n, this fails. Always use tf.range for dynamic loops inside @tf.function.

Chapter 19: GradientTape, XLA, and MLIR

19.1 GradientTape — TF's Explicit Backward Recording

PyTorch records gradients implicitly (every operation on a requires_grad tensor is automatically recorded). TF2 uses explicit recording via tf.GradientTape:

x = tf.Variable([1.0, 2.0, 3.0])   # Variables are automatically watched

with tf.GradientTape() as tape:
    # Everything inside this block is recorded
    y = x ** 2                        # element-wise square
    loss = tf.reduce_sum(y)           # sum: loss = x[0]² + x[1]² + x[2]²

# Compute gradients AFTER the tape context closes
grads = tape.gradient(loss, x)
print(grads)   # [2.0, 4.0, 6.0]  — ∂loss/∂xᵢ = 2xᵢ

tf.Variable objects are automatically watched. For tf.Tensor constants, you must explicitly watch:

with tf.GradientTape() as tape:
    tape.watch(constant_tensor)
    loss = f(constant_tensor)
grad = tape.gradient(loss, constant_tensor)

By default, a tape can only be used once — the recorded operations are freed after gradient() is called. Use tape = tf.GradientTape(persistent=True) to call gradient() multiple times (like PyTorch's retain_graph=True).

19.2 XLA — Accelerated Linear Algebra

XLA is Google's compiler for linear algebra operations. When you annotate with @tf.function(jit_compile=True) or use tf.device('/GPU:0') with XLA enabled, the FuncGraph is lowered to HLO (High-Level Operations) and compiled by XLA.

HLO is XLA's IR — a set of ~100 primitive operations (dot product, element-wise, reduce, reshape, select, etc.). Every TF op eventually lowers to one or more HLO ops.

XLA performs aggressive optimizations on HLO:

• Operation fusion: merges adjacent element-wise ops into one kernel
• Buffer allocation: assigns memory to each HLO op's output, reusing buffers when possible
• Loop tiling: tiles computations to fit in cache, improving memory locality
• Layout optimization: chooses the best tensor layout (row-major vs column-major) for each op to minimize memory transpose costs

XLA then lowers HLO to machine code via LLVM for CPU or NVPTX for NVIDIA GPUs (or the TPU backend for Google's hardware). The generated code is highly optimized, often outperforming manually written CUDA kernels.

19.3 MLIR — Multi-Level Intermediate Representation

MLIR (Multi-Level Intermediate Representation) is Google's framework for building compilers, introduced in 2020. TF now uses MLIR as a bridge between TF's high-level graph and XLA's HLO.

Why multiple levels? Different optimization passes work at different abstraction levels:

• At the TF graph level: common subexpression elimination, shape inference, dead code elimination
• At the XLA HLO level: fusion, buffer assignment, layout optimization
• At the machine code level: register allocation, instruction scheduling

MLIR allows each optimization to operate at the right abstraction level using dialects — pluggable extensions that define new operations and types. TF defines:

• tf dialect: TF-level operations
• mhlo dialect: the HLO equivalent in MLIR
• linalg dialect: linear algebra operations with explicit loop structure
• memref dialect: explicit memory operations

The compilation pipeline is a sequence of lowering passes:

tf dialect → mhlo dialect → linalg dialect → affine dialect → LLVM IR → machine code

Each lowering step progressively descends toward machine code while applying level-appropriate optimizations. This multi-level approach is why TF on TPUs and GPUs often outperforms manually tuned libraries — the compiler has complete knowledge of the computation at every level.

Part VI — Numerical Precision

Chapter 20: IEEE 754 — How Computers Store Numbers

20.1 Why Computers Can't Represent All Real Numbers

Computers store everything in binary (base 2). A 32-bit integer can represent exactly 2³² values. But real numbers are infinite — there are infinitely many values between 0 and 1 alone.

Floating-point is the compromise: represent a wide range of values approximately, with the precision concentrated near 1.0 and thinning out at extreme values (very large or very small).

20.2 The IEEE 754 Standard — Bit by Bit

IEEE 754 defines how floating-point numbers are stored. Every float is represented as:

\[ \text{value} = (-1)^{\text{sign}} \times 1.\text{mantissa}_2 \times 2^{\text{exponent} - \text{bias}} \]

For float32 (32 bits total):

Bit 31      | Bits 30-23   | Bits 22-0
sign (1 bit)| exponent (8) | mantissa (23)

• Sign: 0 = positive, 1 = negative
• Exponent: stored with a bias of 127. Actual exponent = stored_value − 127. Range: 2^(0-127) = 2^-127 to 2^(255-127) = 2^128
• Mantissa: the fractional part after the leading 1 (which is implicit — not stored). 23 bits gives about 7 decimal digits of precision.

Example: how is −6.5 stored?

1. Sign: negative → bit 31 = 1
2. 6.5 in binary: 110.1 = 1.101 × 2^2
3. Exponent: 2 + 127 = 129 = 10000001 in binary
4. Mantissa: 101 followed by 20 zeros: 10100000000000000000000

Full bit pattern: 1 10000001 10100000000000000000000

20.3 Special Values — NaN, Infinity, Denormals

IEEE 754 reserves special exponent values for non-numeric cases:

Infinity: 1.0 / 0.0 = +Inf. Operations on Infinity propagate: Inf + 1 = Inf, Inf - Inf = NaN.

NaN: result of undefined operations: 0.0 / 0.0, sqrt(-1), Inf - Inf. NaN is "poisonous" — any arithmetic involving NaN produces NaN. This is why NaN in one layer propagates through the entire network.

Denormals (subnormal numbers): when the exponent is all zeros, the leading 1 is absent, allowing extremely small values near zero. Operations on denormals are 10-100× slower on some hardware because the FPU falls back to software emulation.

20.4 Precision Limits — When 1 + ε = 1

float32 has 23 mantissa bits. The smallest value that, when added to 1.0, gives a result different from 1.0 is called machine epsilon:

\[ \varepsilon_{\text{fp32}} = 2^{-23} \approx 1.19 \times 10^{-7} \]

This means: 1.0 + 1e-8 == 1.0 in float32. The 1e-8 is below the precision threshold and is simply lost.

Catastrophic cancellation: subtracting two nearly-equal large numbers loses significant digits:

x = 1.000001  (float32 representation: 1.0 + 1e-6)
y = 1.000000  (float32 representation: 1.0)
x - y = 1e-6  (correct)

But if x = 1000000.001 and y = 1000000.000:
float32 precision at scale 1e6 is ε × 1e6 = 1.19e-7 × 1e6 = 0.119
The difference 0.001 is below the representable precision!
x - y = 0.0  in float32  (wrong — should be 0.001)

This happens in:

• Softmax with large logits: sum(exp(large_values)) is huge; individual exp(x_i) may be small relative to the sum
• Layer normalization: subtracting the mean from large-valued activations
• Log probabilities: log(p) where p is very close to 0 or 1

20.5 The ULP — Measuring Floating-Point Error

"Units in the Last Place" (ULP) is the standard measure of floating-point error. 1 ULP is the distance between a float and the next representable float.

Near 1.0: 1 ULP ≈ 1.19e-7 (float32 epsilon) Near 1e8: 1 ULP ≈ 8.0 (much larger!)

When your gradient check requires tolerance 1e-5, you're saying "allow up to about 100 ULP error at values near 1.0". For values near 100, the absolute tolerance should be 100 × 1.19e-7 × 100 = 1.19e-3. This is why gradient checks use relative tolerance, not absolute.

Chapter 21: FP16, BF16, and Mixed Precision Training

21.1 Why Not Always Use float32?

float32 has excellent precision and range. But it uses 4 bytes per value. For a 7B parameter model:

• float32: 7e9 × 4 bytes = 28 GB just for weights
• float16: 7e9 × 2 bytes = 14 GB
• int8: 7e9 × 1 byte = 7 GB

Beyond storage: float32 arithmetic is 2× slower than float16 on modern GPUs with Tensor Cores. The A100's Tensor Core throughput is 312 TFLOPs (float16) vs 77 TFLOPs (float32) — 4× faster for matrix multiplications.

21.2 float16 — High Performance, Dangerous Range

float16 (half precision):

Bit 15    | Bits 14-10 | Bits 9-0
sign (1)  | exponent(5)| mantissa(10)

• Exponent: 5 bits, bias=15. Actual range: 2^-14 to 2^15 = 0.00006 to 65504
• Mantissa: 10 bits ≈ 3.3 decimal digits of precision
• Maximum value: 65504

The danger: during training, gradient norms, loss values, and large weight matrices can easily exceed 65504 → overflow → Infinity → NaN → corrupted training.

Additionally, very small gradients (common in early training or deep networks) can underflow to 0 → parameters stop updating → training stalls.

21.3 Loss Scaling — The Solution to FP16 Overflow and Underflow

The trick: multiply the loss by a large scalar (e.g., 2^15 = 32768) before backpropagation. This shifts all gradients upward by the same factor, preventing underflow. Before applying gradients to weights, divide by the same scalar.

scaler = torch.cuda.amp.GradScaler()

# Training loop
optimizer.zero_grad()
with torch.autocast(device_type='cuda', dtype=torch.float16):
    output = model(input)
    loss = criterion(output, target)

scaler.scale(loss).backward()         # multiplies loss by scaler factor (e.g., 32768)
scaler.step(optimizer)                # unscales gradients, applies them
scaler.update()                       # adjusts scale factor based on inf/nan detection

The GradScaler dynamically adjusts the scale factor: if gradients overflow (contain Inf/NaN), it halves the scale and skips this optimizer step; if gradients are clean for a while, it doubles the scale to catch any underflowing gradients.

21.4 BF16 — The Training-Safe Compromise

BF16 (Brain Float 16, developed by Google for TPUs):

Bit 15    | Bits 14-7  | Bits 6-0
sign (1)  | exponent(8)| mantissa(7)

• Exponent: 8 bits — same as float32! Range: 2^-126 to 2^127 ≈ 10^-38 to 10^38
• Mantissa: 7 bits ≈ 2.3 decimal digits of precision
• No overflow or underflow risk during training (same range as float32)
• No loss scaling needed

BF16 vs FP16 for training:

• BF16: safe range, lower precision. Training is stable.
• FP16: dangerous range, higher precision. Needs loss scaling.

Modern NVIDIA hardware (Ampere A100, Hopper H100) supports BF16 natively with the same Tensor Core throughput as FP16. BF16 is now the default training format for most LLM training runs.

21.5 AMP — Automatic Mixed Precision

torch.autocast handles the dtype selection automatically:

with torch.autocast(device_type='cuda', dtype=torch.float16):
    # These ops run in float16 (faster, sufficient precision):
    output = model.layers(input)     # matmul, conv2d
    
    # These ops are forced to float32 (numerically sensitive):
    loss = F.cross_entropy(logits, targets)   # softmax + log inside
    # layer norm, softmax, batch norm also run in float32

The autocast context uses a table of operations and their recommended dtypes. Matmul and conv can lose precision without impacting final model quality. Softmax, layer norm, and loss functions are sensitive to precision — they stay in float32.

Part VII — Synthesis and Practice

Chapter 22: How Everything Connects — The Principal Engineer's Map

                     MATH LAYER
     Chain Rule → Jacobians → VJPs → Reverse Mode AD
           ↓
     PYTHON LAYER  
     CPython bytecode → GIL → pybind11 → C++
           ↓
     PYTORCH C++ LAYER
     c10 types → ATen ops → Dispatcher → DispatchKeySet
           ↓              ↓
     Autograd Key    CUDA/CPU Keys
     (records DAG)   (compute kernels)
           ↓
     BACKWARD DAG
     Node→Edge→AccumulateGrad → leaf.grad
           ↓
     GRAPH CAPTURE
     FX IR (symbolic trace) → torch.compile (Dynamo bytecodes + guards)
           ↓
     COMPILATION
     AOTAutograd → joint forward+backward graph
     TorchInductor → Triton (GPU) / C++ (CPU)
           ↓
     CUSTOM EXTENSIONS
     TORCH_LIBRARY schema → TORCH_LIBRARY_IMPL per key
     (CPU, CUDA, Meta, PrivateUse1/NPU)

How each lab maps to production work:

Lab 01 (minigrad) → Production skill: Debugging gradient issues

• When a gradient is wrong: you know to check each Function.backward for the VJP formula
• When a gradient is None: you know to check requires_grad, no_grad() scope, or detach()
• When a model doesn't converge: you know to run gradcheck to find the broken backward
• When writing a custom backward for a fused attention kernel: you apply Lab 01 patterns directly

Lab 02 (warp_softmax) → Production skill: Writing NPU kernels

• The TORCH_LIBRARY registration pattern is identical for CPU, CUDA, and NPU backends
• The Meta kernel requirement is universal — every operator in the ecosystem needs it
• The warp reduction technique applies to any per-row reduction: max-pooling, layer norm, etc.
• opcheck is used to validate every operator in PyTorch's own test suite

Chapter 23: Lab Guidance — Building minigrad and warp_softmax

23.1 Lab 01: minigrad Implementation Order

Start from the simplest operations and build up. Each step gives you a working test before the next.

Step 1: Implement _build_topo on Tensor. This is the topological sort. Without it, backward() cannot run. Start here:

def _build_topo(self):
    topo, visited = [], set()
    def dfs(t):
        if id(t) in visited: return
        visited.add(id(t))
        if t.grad_fn:
            for inp in t.grad_fn._input_tensors:
                if isinstance(inp, Tensor): dfs(inp)
        topo.append(t)
    dfs(self)
    return topo

Step 2: Implement backward() on Tensor. Initialize gradient, call _build_topo, traverse in reverse, call each grad_fn's backward, accumulate into leaves. Verify with the smoke test at bottom of lab.py.

Step 3: Implement Add (forward + backward). Backward must handle broadcasting via _unbroadcast. Test: a + b where a.shape=(3,), b.shape=(4,3).

Step 4: Implement Mul, MatMul, Neg, Div, Pow. These teach: saving inputs for backward, the VJP formulas for each op.

Step 5: Implement ReLU, Sigmoid, Tanh, Exp, Log. Key insight: Sigmoid.forward should save the output (not input) because d(sigmoid)/dx = sigmoid(x)(1 - sigmoid(x)) is expressed in terms of output.

Step 6: Implement softmax using Max, Exp, Sum, Div. Apply the max-subtraction trick for numerical stability. All ops you implement in steps 3-5 compose naturally.

Step 7: Implement cross_entropy_loss using log_softmax. The full MLP training loop should then converge.

Numerical gradient check — what it is: approximate the derivative numerically using finite differences, compare with your analytic gradient:

numerical_grad[i] ≈ (f(x + ε·eᵢ) - f(x - ε·eᵢ)) / 2ε

If abs(analytic - numerical) / max(abs(analytic), abs(numerical)) < 1e-5, your backward is correct.

23.2 Lab 02: warp_softmax — Reading Before Building

Read the source files in this order:

1. csrc/warp_softmax.cpp: understand TORCH_LIBRARY (schema), TORCH_LIBRARY_IMPL (CPU + Meta), warp_softmax_cpu (ATen ops for CPU), warp_softmax_meta (returns empty_like).

2. csrc/warp_softmax.cu: read warp_reduce_max and warp_reduce_sum first — these are the butterfly reduction utilities. Then read warp_softmax_kernel — trace through the three steps (find max, compute exp+sum, normalize). Map each line to the hardware concepts from Chapter 15.

3. setup.py: understand how torch.utils.cpp_extension.CUDAExtension compiles your C++/CUDA and links against libtorch.

4. solution.py: understand how opcheck is called and what it verifies.

23.3 Lab 03: tf.function Tracing — Implementation Order

Chapter 18 is the theory; the lab makes it mechanical. Order matters here too:

1. trace_and_inspect first — until you can see a FuncGraph's ops, captures, and signature, the rest is blind. Key API: fn.get_concrete_function(*inputs).graph.
2. RetraceProbe — the cleverness is that a Python counter inside the traced body is the retrace counter, because Python side effects fire only at trace time (Chapter 18's central fact, now weaponized as instrumentation).
3. demonstrate_retracing — predict the three trace counts on paper before running; if your prediction misses, re-read §18's retracing rules.
4. stabilize — tf.TensorSpec([None, ...]) makes shape symbolic; verify one trace serves all batch sizes.
5. autograph_diff, then the SavedModel/TFLite round-trips — SavedModel must be bit-exact (same graph, same weights); TFLite dynamic-range must be close-but-not-exact (weights were perturbed to INT8). Both assertions are in the tests; understand why each bound holds.

23.4 Success Criteria

You are ready for Phase 02 when you can answer all of these from memory:

1. Trace a 5-op PyTorch expression's full backward pass with specific gradient values at each node.
2. Name every dispatch key in priority order and explain what each intercepts.
3. Give 3 distinct triggers for tf.function retracing and the prevention for each.
4. Explain what breaks in torch.compile if your custom op has no Meta kernel.
5. Draw the butterfly reduction for warp_reduce_max with 8 threads (3 rounds).
6. Explain why BF16 doesn't need loss scaling but FP16 does.
7. Write the VJP (backward formula) for matmul from the Jacobian definition.

If any answer is fuzzy, return to that chapter and re-read the internal mechanism — not the API, the mechanism.

Chapter 24: The Interview Gauntlet — Complete Answers

These are actual questions asked at senior/staff ML framework engineering interviews. The expected answer length and depth is shown.

Q: Walk me through what happens, step by step, when I call loss.backward().

A complete answer (2-3 minutes verbal):

loss.backward() calls into torch::autograd::Engine::execute(). The engine starts from loss.grad_fn with an initial gradient of torch.ones_like(loss). It builds a ready-count for each Function node — how many downstream nodes still need to send gradients to it. Processing order is determined by decrementing these counts: a node runs when its count reaches zero (all downstream contributions received).

For each node, the engine calls node.apply(input_gradients) — the virtual backward function. For built-in ops this is a generated class like MmBackward0; for custom ops it's the static backward method you defined. The function returns a tuple of gradients — one per input. Each gradient is sent to the corresponding upstream node (via next_edges_). Leaf nodes have AccumulateGrad as their node, which adds the incoming gradient to leaf.grad (or initializes it).

After backward completes, the Function nodes and their saved tensors are freed (unless retain_graph=True). This is why you can't call backward() twice on the same graph — the intermediate data is gone.

Q: What is a dispatch key and why is its priority ordering load-bearing?

The DispatchKeySet is a 64-bit bitfield attached to every tensor. Each bit position corresponds to one DispatchKey. The dispatcher takes the bitwise OR of all input tensors' key sets and calls the highest-priority key that has a registered implementation for the operation.

The priority ordering is load-bearing because the keys are semantic wrappers. Autograd sits above CUDA so that gradient recording always wraps the compute kernel — you can't record gradients without the Autograd key seeing the operation first. Functionalize sits above Autograd so that aliasing analysis (needed for torch.compile) can see the original mutation intent before autograd creates any graph nodes.

If an NPU vendor registered their key above Autograd, their device would silently bypass gradient recording. Training would appear to work (forward pass correct) but gradients would never be computed.

Q: When does tf.function retrace, and what is the performance cost?

Retrace triggers: (1) new input shapes when input_signature is not set, (2) new input dtypes, (3) change in non-tensor Python argument values (they are compile-time constants baked into the graph), (4) changed Python closures captured by the function.

The cost is: re-running AutoGraph (Python source transformation), re-tracing the function with SymbolicTensors (building a new FuncGraph), optionally re-running XLA compilation (extremely expensive — can take minutes for large models). In production, unexpected retracing is a severe latency spike. You diagnose with tf.autograph.set_verbosity(10) to print retrace events.

Prevention: always use input_signature with None for variable dimensions; convert Python-controlled iteration counts to tf.Tensor arguments; avoid Python closures that change.

Q: A model works in eager mode but gives NaN in torch.compile. How do you debug it?

Systematic approach:

Step 1: Isolate forward vs backward. Add torch.autograd.set_detect_anomaly(True) before the training loop — it adds expensive runtime checks but identifies exactly which backward operation produces NaN.

Step 2: Bisect the graph. Use torch.compile(fullgraph=False) — this allows graph breaks. Introduce a deliberate graph break at the suspected location (e.g., torch._dynamo.graph_break()) and check if NaN appears before or after the break.

Step 3: Check custom ops. If you have a custom operator, verify it has a Meta kernel and run opcheck. An incorrect Meta kernel can cause shape mismatches that produce NaN in compiled mode.

Step 4: Check numerical precision. Compiled mode may fuse operations that change floating-point evaluation order. Add torch.use_deterministic_algorithms(True) to force bit-identical behavior across modes.

Step 5: Check in-place ops. The Functionalize dispatch key converts in-place ops to functional form. If your custom backward incorrectly modifies saved tensors in-place, Functionalize may not handle it correctly, producing NaN.

Q: What is the Meta device, and why does every custom operator need a Meta kernel?

The Meta device is a special device type where tensors have shape, dtype, device=meta, but zero bytes of storage. Operations on Meta tensors execute "Meta kernels" that propagate shapes and dtypes without doing any computation.

torch.compile's TorchDynamo captures the computation as an FX graph. AOTAutograd traces through this graph using FakeTensors — which are Meta tensors augmented with symbolic size information. Every operation in the graph must have a Meta kernel so the compiler can determine output shapes without running compute kernels.

A custom operator without a Meta kernel causes torch.compile to raise NotImplementedError: Could not run 'myops::op' with the 'Meta' dispatch key. The fix is always:

TORCH_LIBRARY_IMPL(myops, Meta, m) {
    m.impl("my_op", [](const at::Tensor& input, int64_t dim) {
        return at::empty_like(input);  // same shape, same dtype, no data
    });
}

For ops that change shape (e.g., a custom reshape), the Meta kernel must compute the correct output shape from the input shape and arguments.

Q: Explain catastrophic cancellation with a concrete example, and how you'd prevent it in a softmax implementation.

Catastrophic cancellation: subtracting two nearly-equal large numbers produces a result with far fewer significant digits than either operand.

Concrete example: computing softmax([1000.0, 999.0]) in float32.

exp(1000.0) → inf (overflows float32 max ≈ 3.4e38)
exp(999.0) → inf

softmax = inf / (inf + inf) = inf / inf = NaN

Even if the values don't overflow: computing exp(100.0) / (exp(100.0) + exp(99.0)) in float16:

• exp(100) ≈ 2.69e43 → overflows float16 (max 65504) → Inf

Prevention — the max-subtraction trick:

softmax(x_i) = exp(x_i) / Σ exp(x_j)
             = exp(x_i - M) / Σ exp(x_j - M)    where M = max(x)

Subtracting M ensures all exponents are ≤ 0, so exp(x_i - M) ∈ (0, 1] — no overflow, no underflow, no cancellation. The maximum value is always exp(0) = 1.

For the normalization sum itself: compute in float32 even when the inputs are float16, to maintain precision. PyTorch's F.softmax does this automatically.

References

Foundational Papers

1. Reverse-mode AD: Baydin, A.G., Pearlmutter, B.A., Radul, A.A., Siskind, J.M. (2018). Automatic Differentiation in Machine Learning: a Survey. JMLR 18(153). — The definitive survey of automatic differentiation; explains forward/reverse modes rigorously.

2. PyTorch: Paszke, A. et al. (2019). PyTorch: An Imperative Style, High-Performance Deep Learning Library. NeurIPS 2019. — The paper describing PyTorch's design philosophy and dispatcher.

3. TorchDynamo: Ansel, J. et al. (2023). PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. PLDI 2023. — The paper behind torch.compile; explains Dynamo's bytecode approach and guard system.

4. Triton: Tillet, P., Kung, H.T., Cox, D. (2019). Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations. MAPL 2019. — The paper behind TorchInductor's code generation backend.

5. FlashAttention: Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. — Shows how GPU memory hierarchy analysis leads to 2-4× speedups; the same analysis applies to the warp_softmax kernel.

6. JAX: Bradbury, J. et al. (2018). JAX: Composable Transformations of Python+NumPy Programs. — JAX's jit, grad, vmap show an alternative design to PyTorch's; understanding both gives perspective on framework design tradeoffs.

7. MLIR: Lattner, C. et al. (2021). MLIR: Scaling Compiler Infrastructure for Domain Specific Computation. CGO 2021. — The paper behind TF's MLIR compilation infrastructure.

8. XLA: Leary, C. et al. (2017). XLA: TensorFlow, Compiled. TF Dev Summit 2017. — XLA's design and its HLO IR.

Books

9. Deep Learning: Goodfellow, I., Bengio, Y., Courville, A. (2016). Deep Learning. MIT Press. — Chapters 6 (feedforward networks) and 6.5 (backpropagation) are the mathematical foundation for everything in Part I.

10. Programming Massively Parallel Processors: Kirk, D.B., Hwu, W.W. (4th ed., 2022). — The definitive textbook on CUDA programming; Chapters 4 (memory), 10 (reduction), and 15 (performance optimization) are directly relevant.

11. Matrix Computations: Golub, G.H., Van Loan, C.F. (4th ed., 2013). — The mathematical foundation for understanding why matrix operations (matmul, SVD, QR) are designed as they are.

12. What Every Computer Scientist Should Know About Floating-Point Arithmetic: Goldberg, D. (1991). ACM Computing Surveys. — The definitive reference for IEEE 754, rounding modes, and catastrophic cancellation.

Technical Documentation and Blog Posts

13. PyTorch Internals — Edward Yang's blog: http://blog.ezyang.com/2019/05/pytorch-internals/ — The most accessible deep-dive into PyTorch's C++ architecture written by one of its core developers.

14. TorchDynamo Deep Dive — PyTorch Dev Discuss: https://dev-discuss.pytorch.org/t/torchdynamo-an-experiment-in-dynamic-python-bytecode-transformation/361 — The original announcement and explanation of Dynamo's bytecode approach.

15. AOTAutograd Walkthrough — https://pytorch.org/functorch/stable/notebooks/aot_autograd_optimizations.html — Official documentation with worked examples.

16. CUDA Programming Guide — NVIDIA developer docs: https://docs.nvidia.com/cuda/cuda-c-programming-guide/ — The authoritative reference for warp execution, memory hierarchy, and shuffle instructions.

17. torch.fx Documentation — https://pytorch.org/docs/stable/fx.html — Official API docs with examples of graph transformation patterns.

18. tf.function Guide — https://www.tensorflow.org/guide/function — TF's official guide covering tracing, retracing, and AutoGraph in depth.

19. Mixed Precision Training — Micikevicius, P. et al. (2018). Mixed Precision Training. ICLR 2018. — The paper that introduced loss scaling for FP16 training.

20. The TORCH_LIBRARY Tutorial — https://pytorch.org/tutorials/advanced/torch_script_custom_ops.html — PyTorch's official guide to custom operator registration, with C++ examples.