Tensor maths & custom autograd

Open the box. Learn how a tensor really sits in memory, how gradients flow backwards through each operation, and how to write your own GPU kernel when the built-in one is too slow.

Source: Suraj Sharma's roadmap post on X, stage 1 · all 12 stages

Going deeper: everything below is the full version, with some maths and code. On a first read, skim the headings and the "Why it matters for robotics" box, and come back for the rest later.

1. What this stage covers

• Learn: PyTorch internals (storage, views, the autograd graph), custom forward and backward passes, Triton programming, and how memory layout and strides decide whether an operation is free, a copy, or an error.
• Practice: two projects. A fused attention kernel in Triton. And a backward pass you derive and code yourself for a new activation function, without leaning on PyTorch's autograd to do it for you.
• Why (the post's claim): if the network is a black box to you, you can't diagnose a loss curve that falls apart. We test this claim in the reality check.

2. Concepts, explained

PyTorch internals: storage, views and the autograd graph

A PyTorch tensor is a view onto a storage: a flat, typed block of memory. The view adds a shape, a stride for each dimension, and a storage offset. Element [i, j] lives at:

address(i, j) = offset + i·stride[0] + j·stride[1]

x = torch.arange(12).view(3, 4)   # shape (3, 4), stride (4, 1), offset 0
x.T                               # shape (4, 3), stride (1, 4): same memory, no copy
x[:, ::2]                         # shape (3, 2), stride (4, 2)
x[1:]                             # shape (2, 4), stride (4, 1), offset 4

When you run operations on tensors with requires_grad=True, PyTorch records each one as a node in a graph (you can see it in y.grad_fn). Each node keeps what it needs for its backward rule, such as its inputs. Calling backward() walks this graph in reverse topological order. It multiplies the incoming gradient by each node's local derivative and adds up the contributions where paths merge. The graph is rebuilt on every forward pass ("define-by-run", Paszke et al., 2019). This is reverse-mode automatic differentiation. Baydin et al.'s survey (2015, revised 2018) explains how it differs from symbolic and numerical differentiation.

Why reverse mode? A network has millions of inputs (the weights) and one output (the loss). Reverse mode gets the gradient with respect to all of them for a small constant multiple of the forward cost. Forward mode would need one pass per input.

Custom forward and backward passes

You need your own backward when the automatic one is slow, uses too much memory, or is numerically unstable. Sometimes there's no automatic one at all, because the forward calls a custom kernel. In PyTorch you subclass torch.autograd.Function:

class Squareplus(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, b=4.0):
        r = torch.sqrt(x * x + b)
        ctx.save_for_backward(x, r)     # keep only what backward needs
        return 0.5 * (x + r)

    @staticmethod
    def backward(ctx, g):               # g = dLoss/dOutput
        x, r = ctx.saved_tensors
        return g * 0.5 * (1 + x / r), None   # one gradient per forward input

torch.autograd.gradcheck(Squareplus.apply, torch.randn(8, dtype=torch.double, requires_grad=True))

The rule to remember: backward receives the gradient of the loss with respect to your output and returns the gradient with respect to each input. That's the vector-Jacobian product. For an element-wise function it's just g * f'(x). Always check a hand-written backward against finite differences in float64. gradcheck does this for you.

Squareplus is a real activation, from Barron (2021). It looks like softplus but uses only addition, multiplication and a square root. It's a good "novel" function for the practice project because nobody hands you its derivative in a library.

Triton programming

Triton (Tillet, Kung & Cox, 2019; now developed in the open by OpenAI and contributors) is a Python-embedded language for GPU kernels. In CUDA you program one thread. In Triton you program one block of data. Each program instance gets an ID from tl.program_id. It loads a tile with tl.load (using a mask at the edges), computes on it as if it were a small array, and writes with tl.store. The compiler handles the thread layout, shared memory and much of the vectorisation. You launch a kernel over a grid of program instances: kernel[grid](...).

The reason to write kernels at all is the memory hierarchy. A GPU has large but slow off-chip memory (HBM) and a small, fast on-chip memory (SRAM). Many operations (softmax, layer norm, activations) do only a few flops per byte they read. Their speed is set by memory traffic, not maths. Fusing three such ops into one kernel reads and writes the data once instead of three times.

Fused attention (the FlashAttention trick)

Standard attention builds the full n × n score matrix in HBM, softmaxes it, then multiplies by V. FlashAttention (Dao, Fu, Ermon, Rudra & Ré, 2022) never stores that matrix. It loops over tiles of K and V, keeps each block of Q in SRAM, and uses the online softmax: keep a running row maximum and a running denominator, and rescale the partial output whenever the maximum grows. That idea goes back to Milakov & Gimelshein (2018). The result is exact attention with memory linear in sequence length. In the backward pass, it recomputes the attention tiles from Q, K and V instead of storing them, trading cheap flops for expensive memory traffic. The paper reports a 3× training speedup on GPT-2 at sequence length 1K. FlashAttention-2 (Dao, 2023) reworked how the work is split across thread blocks and warps, for roughly another 2×.

for each tile (K_j, V_j):
    S    = Q · K_jᵀ / √d
    m'   = max(m, rowmax(S))
    P    = exp(S − m')
    l    = l · exp(m − m') + rowsum(P)
    O    = O · exp(m − m') + P · V_j
    m    = m'
O = O / l

Memory layouts and stride calculations

A tensor is contiguous (row-major, "C order") when its strides are exactly what a freshly allocated tensor of that shape would have. Each stride equals the product of the sizes to its right. Size-1 dimensions don't count. Three consequences you will meet constantly:

• .view() never copies. It fails if the new shape can't be expressed with strides over the existing memory, which is common after a transpose. .reshape() falls back to a copy. .contiguous() forces one.
• A transposed or strided view makes the innermost loop jump around in memory. GPUs read memory fastest when neighbouring threads touch neighbouring addresses ("coalesced" access). A kernel that is fast on a contiguous input can be several times slower on a transposed one.
• Layout is a choice. PyTorch's channels_last gives an NCHW-shaped image tensor NHWC strides, which many convolution kernels prefer.

3. The practice project

(a) Laptop / Colab version

Part 1: a backward pass with no autograd (laptop, NumPy only, an evening). Build a two-layer MLP with squareplus as the activation. Write the forward pass, derive the backward pass on paper, code it, and check it against finite differences.

import numpy as np
B = 4.0
sq  = lambda x: 0.5 * (x + np.sqrt(x * x + B))
dsq = lambda x: 0.5 * (1 + x / np.sqrt(x * x + B))

def forward(W1, W2, x, y):
    h = x @ W1; a = sq(h); p = a @ W2
    return ((p - y) ** 2).mean(), (x, h, a, p)

def backward(W1, W2, y, cache):
    x, h, a, p = cache
    dp = 2 * (p - y) / len(y)
    dW2 = a.T @ dp
    dh = (dp @ W2.T) * dsq(h)      # chain rule through the activation
    return x.T @ dh, dW2

rng = np.random.default_rng(0)
x, y = rng.normal(size=(32, 3)), rng.normal(size=(32, 1))
W1, W2 = rng.normal(size=(3, 8)), rng.normal(size=(8, 1))
dW1, _ = backward(W1, W2, y, forward(W1, W2, x, y)[1])
eps, E = 1e-6, np.zeros_like(W1); E[1, 2] = eps
num = (forward(W1 + E, W2, x, y)[0] - forward(W1 - E, W2, x, y)[0]) / (2 * eps)
print(dW1[1, 2], num)              # should agree to ~1e-8

We ran this sketch: the analytic and numerical gradients agreed to about 2 × 10⁻⁹. Then wrap the same maths in a torch.autograd.Function (section 2) and pass gradcheck.

Part 2: fused attention, first in NumPy, then in Triton (Colab GPU, a weekend or two). Write the tiled online-softmax loop from section 2 in NumPy and check it matches naive attention to about 1e-12 in float64. (Our version, with 64 queries, head dimension 32 and tiles of 16 keys, matched to ~7 × 10⁻¹⁶.) Then work through the official Triton tutorials in order: vector add, fused softmax, matmul, fused attention. Write your own forward-only attention kernel, test it against torch.nn.functional.scaled_dot_product_attention with torch.testing.assert_close, and time it with triton.testing.do_bench across sequence lengths. Triton has an interpreter mode (TRITON_INTERPRET=1) that runs kernels on the CPU for debugging, slowly.

What to measure: maximum absolute error against the reference in fp32 and fp16. Runtime and peak memory (torch.cuda.max_memory_allocated) as sequence length doubles. The naive version's memory should grow with the square of the length. Yours should grow linearly. Don't expect to beat PyTorch's built-in fused attention. Matching it within a small factor is a good result.

(b) The full version from the post

• Fused attention with a backward pass, causal masking and fp16/bf16 support, the way CS336 Assignment 2 sets it up. Then verify gradients against the reference implementation and profile it.
• Hardware: one modern NVIDIA GPU is enough. Triton targets NVIDIA, and AMD GPUs are also supported. Colab's free GPU works for learning. For realistic benchmarks, rent an A100- or H100-class GPU for a few hours, at the usual cloud rate of a few dollars an hour. Some of Triton's newest features need recent GPU generations, so read each tutorial's notes.
• The activation project, in full: a custom autograd.Function with a hand-written backward, then the same activation fused into a Triton kernel with the matmul before it. Benchmark against the eager PyTorch composition and torch.compile.

(c) Pitfalls

• Gradient checks in float32 give noisy false alarms. Use float64 and central differences.
• Forgetting the masked edges in Triton when the size isn't a multiple of the block size. Out-of-bounds loads give garbage, not errors.
• Online softmax without the max subtraction overflows in fp16 almost immediately.
• Benchmarking wrongly: GPU calls are asynchronous. Synchronise, warm up (the first call compiles), and repeat. Use do_bench.
• Silent copies: .contiguous() or .reshape() inside a hot loop can cost more than the kernel you just optimised.

4. Reality check

The post's claim: you can't optimise or debug a collapsing loss curve if you treat networks as black-box APIs.

• Mostly true for debugging. Many loss blow-ups trace back to things this stage teaches you to see. Examples: an op that overflows in fp16, a softmax without max subtraction, a bad custom gradient, a silent broadcast or transpose bug. Knowing what backward does lets you inspect gradient norms per layer and find the culprit.
• Overstated as stated. Plenty of loss-curve problems have nothing to do with kernels: learning rate and warmup, bad data batches, initialisation, optimiser settings. Those are diagnosed with logging and ablations more than with internals. Stage 02 covers part of that.
• Writing Triton kernels is a specialist skill. Most research engineers rarely write one. They use FlashAttention, torch.compile and library kernels. Knowing why those are fast matters more than writing them. Do the attention kernel once to learn it, not to replace the library.
• Hand-written backward passes are rare in practice, but the derivation skill isn't. It's how you read papers that propose new layers, and how you spot numerically unstable gradients.

5. In this course

• Module 02 · TensorTonic: implement-from-scratch practice, including its CUDA kernel and Triton exercises. The maths section covers the linear algebra you need here.
• 30u30 Day 01 · RNN Effectiveness: backprop through time, worked by hand.
• 30u30 Day 14 · The Annotated Transformer: the attention layer you will fuse, written out in plain PyTorch.
• Module 05 · Serving: why memory traffic, not flops, sets the cost of running a big model.
• Module 04 · Latency: why kernel speed matters for small on-device models.

6. Why it matters for robotics

7. Self-check

8. Courses & content

Start with Karpathy's micrograd video: in about two hours you build autograd yourself. Everything else on this page then feels like a bigger version of that.

Full course map for all 12 stages →

9. Papers & docs

• Paszke et al., 2019. PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Baydin, Pearlmutter, Radul & Siskind, 2015/2018. Automatic differentiation in machine learning: a survey
• Dao, Fu, Ermon, Rudra & Ré, 2022. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• Dao, 2023. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• Milakov & Gimelshein, 2018. Online normalizer calculation for softmax
• Barron, 2021. Squareplus: A Softplus-Like Algebraic Rectifier
• Tillet, Kung & Cox, 2019. Triton: an intermediate language and compiler for tiled neural network computations (MAPL workshop). Triton docs
• PyTorch: Extending PyTorch and Autograd mechanics