You don't need autodiff theory. You need to turn the gradient of the output into the gradient of the input, one op at a time — and to know what that does to your memory access pattern.
forward · activations backward · gradients
00Notation
One bar, and you can read everything below
Writing \(\partial L/\partial X\) everywhere makes formulas unreadable. The standard autodiff shorthand puts a bar over the variable:
\[\bar{X}\;\equiv\;\frac{\partial L}{\partial X}\qquad\text{— in code this is }\texttt{dX}\text{ or }\texttt{grad\_X}\]
Key property: the gradient always has the same shape as the thing it is the gradient of. \(\bar X\) has the shape of \(X\). This alone catches most bugs.
01The only idea
Backprop is repeated vector–Jacobian products
Every op is a function \(y=f(x)\). The framework hands you \(\bar y\), the gradient of the loss with respect to your output. Your job is to return \(\bar x\). That is the entire contract:
This is a vector–Jacobian product (VJP), also called the cotangent or the pullback. The Jacobian \(\partial y/\partial x\) appears in the formula but is never materialized in a kernel.
That last line is the most important sentence on this page. For softmax over a 4096-vector the Jacobian is \(4096\times4096\approx16\text{M}\) floats — per row. A good kernel computes the same VJP in \(O(N)\) time and \(O(1)\) extra memory. Learning backprop for kernels is learning these collapses.
What you can skip
Tape recording, graph topological sort, higher-order gradients, checkpoint scheduling. PyTorch and JAX own all of it — you register a backward and the framework calls it. And if you only write inference kernels (quantized GEMMs, decode attention, sampling), you can skip this page entirely and lose nothing.
02See it move
One real chain, with real numbers
A 2-input linear layer into softmax into cross-entropy. Step through it: the forward pass fills in activations, the backward pass fills in gradients right to left.
x → linear → softmax → loss
Press Step to push values forward.
Notice the payoff at the end. Softmax has a dense Jacobian and cross-entropy has a \(-1/y_t\) singularity, but composed together they cancel exactly:
\(e_t\) is the one-hot vector for the true class. Two ugly derivatives collapse into a subtraction — which is why every framework ships a fused softmax_cross_entropy rather than two ops. This is the simplest example of the thing this whole page is about.
The mnemonic that never fails: there is exactly one arrangement of the operands whose shapes agree. Forget the formula, write down the dimensions, and the transposes place themselves.
The shape-level view hides the interesting part. Write the same two results elementwise and the difference jumps out:
Same three matrices, but the reduction index differs: \(\bar X\) sums over the feature axis \(N\); \(\bar W\) sums over the batch axis \(M\). Watch what that does below.
access patterns of the two backward matmuls
The kernel consequence
\(\bar X\) has one independent output per \((i,k)\) and parallelises across the batch like an ordinary tiled GEMM. \(\bar W\) reduces over the batch — every sample contributes to every weight gradient, so partial sums must combine across thread blocks: atomicAdd, split-K, or a separate reduction pass. On large batches with small weight matrices, \(\bar W\) is usually where your backward kernel's time actually goes.
04The duality
Broadcast forward ⇄ reduce backward
This is the structural law behind most of the surprises above. If a value is copied in the forward pass, its backward must sum. Conservation of influence: a value that affected ten outputs receives gradient from all ten.
\(O(N^2)\to O(N)\) time. \(O(N^2)\to O(1)\) extra memory. One pass to accumulate the scalar \(s\), one elementwise pass to apply it.
softmax backward: dense Jacobian vs. scalar hub
Two more things a kernel writer must know about softmax:
You need \(y\), not \(x\). The backward is expressed entirely in terms of the output. Save \(y\), or recompute it from the saved max and logsumexp — which is exactly what FlashAttention does.
The stability shift is free. Forward uses \(\mathrm{softmax}(x)=\mathrm{softmax}(x-\max x)\) for numerical safety. Softmax is shift-invariant, so this contributes no extra term to the backward. Nothing to do.
LayerNorm — the same shape of problem
LayerNorm couples all \(N\) elements through the mean and variance, so it also looks like it needs a dense Jacobian, and it also collapses — this time to two row-wise scalars instead of one.
The pattern repeats: the per-row part is cheap and embarrassingly parallel; the parameter gradients reduce over the batch axis and need cross-block accumulation. Same story as \(\bar W\).
06Top-k & routing
Gradient as a sparse scatter
Top-k is a selection, not a computation. Forward it emits \(k\) values and their indices \(i_1,\dots,i_k\). Backward, gradient flows only into those slots:
\[y_m=x_{i_m}\qquad\Longrightarrow\qquad
\bar x_j=\sum_{m\,:\,i_m=j}\bar y_m
\qquad\text{(zero for every unselected }j)\]
A pure scatter-add. If indices are guaranteed distinct — as in top-k — a plain scatter suffices. If they can repeat, as in gather or embedding lookup, you must use atomic add, and the result is nondeterministic by default.
top-k, N=8, k=3
Two traps here that cost people real debugging time:
The indices are constants, not variables. There is no \(\partial i_m/\partial x\) — selection is piecewise-constant, so its derivative is zero almost everywhere and undefined at ties. This is precisely why an MoE router needs a softmax gate multiplied into the output, so gradient reaches the router through the weights rather than through the argmax, or a straight-through estimator standing in for the missing derivative.
Duplicate indices need atomic add, not store. A plain scatter silently drops all but one contribution. It is the broadcast-in-forward law hiding in plain sight: one input read by many outputs must sum on the way back.
The same skeleton covers a whole family — all of them zero-FLOP backwards whose entire cost is the memory access pattern:
\[\text{ReLU: }\;\bar x=\bar y\odot\mathbb{1}[x>0]
\qquad
\text{max-pool: scatter to }\arg\max
\qquad
\text{dropout: reuse the mask}\]
07Fusion
Save vs. recompute, the central tradeoff
The moment you fuse, you own the VJP of the whole composite — you can no longer chain library backwards. And you face a decision that doesn't exist in unfused code: the forward's intermediates are gone unless you wrote them out, and writing them out is often the exact thing you fused to avoid.
Attention is the famous case. Here is the full backward, and note that step three is just the softmax collapse from §05 applied row-wise:
That green identity is a genuinely lovely trick: substituting \(\bar P=\bar O V^\top\) into \(D\) makes the \(V\) cancel, so the row-scalar can be computed up front from \(\bar O\) and \(O\) alone — no need for \(P\) yet. FlashAttention precomputes all of \(D\) in one cheap pass before touching the score matrix.
Which leaves the real question: where does \(P\) come from in backward?
\[\underbrace{\text{store }P}_{O(N^2)\text{ HBM}}
\qquad\text{vs.}\qquad
\underbrace{\text{store only }L_i=m_i+\log\ell_i}_{O(N)\text{ HBM}},\;\;
P_{ij}=\exp\!\left(S_{ij}-L_i\right)\]
\(m_i\) is the row max and \(\ell_i\) the row sum of exponentials, both already computed in the forward's online softmax. From the single scalar \(L_i\), any tile of \(P\) can be rebuilt exactly.
attention backward · HBM traffic vs. FLOPs · N=4096
Recompute costs more arithmetic and still finishes first — because you were never compute-bound to begin with.
The general rule
Recompute when the intermediate is large and rebuilding it is cheap arithmetic on data already in registers or SRAM. Save when the intermediate is small — a scalar or one vector per row — or when regenerating it would need another trip to HBM. Softmax statistics, RNG seeds and normalisation constants are almost always worth saving. Anything \(O(N^2)\) almost never is.
Why fused backward is harder than fused forward
Multiple outputs, different reduction axes. \(\bar Q\), \(\bar K\), \(\bar V\) do not want the same loop order — something has to be transposed or atomically accumulated.
Register pressure. You hold forward tiles and gradient tiles simultaneously. This is where you fall off the occupancy cliff and your beautiful kernel loses to three unfused ones.
Numerics improve. One genuine gift: fused intermediates stay in fp32 registers instead of round-tripping through fp16 memory. Fusion often buys accuracy and speed together.
The derivations you cannot look up are the ones you own. In order, each a single sitting:
Derive matmul backward from shapes alone. No formula lookup — place the transposes by making the dimensions agree. Repeat until reflex.
Derive the softmax collapse yourself. Write the dense Jacobian, multiply by \(\bar y\), factor out the shared scalar. Feeling \(s\) fall out is the whole lesson.
Write a fused softmax backward in Triton and check it against torch.autograd.gradcheck on a naive reference. Correctness first — you don't need theoretical confidence, you need a passing test.
Then LayerNorm backward: same skills, two accumulators instead of one, plus your first batch-axis parameter reduction.
Then attention backward with recompute: everything above composed, plus the save/recompute decision made deliberately.
Debugging order
When a backward is wrong, check in this order: (1) a transpose — shapes pass but semantics don't; (2) a missing reduction over the batch axis — correct at batch 1, wrong above it; (3) a sign; (4) a scatter that needed atomic add. That ordering catches the overwhelming majority. Always test at batch 1 and batch 7 — class (2) is completely invisible at batch 1, and 7 is prime so it won't tile evenly and will expose your boundary handling too.