Backward in Triton — autograd.Function and the backward grid

Lessons 01–14 built the forward half: a @triton.jit kernel turns x into y, a thin wrapper calls it. To train, you need the other half — a second, separate Triton kernel that turns dy back into dx (and dw). This lesson is the mechanics of wiring the two together with torch.autograd.Function, choosing the backward grid, and the four traps that cost everyone a day.

The question this lesson answers

You have a working forward Triton kernel. How do you make it differentiable so it drops into a training loop and loss.backward() just works? The answer is one class — torch.autograd.Function — with two static methods and one contract. Get the contract wrong and you get silent zero-gradients or a shape error three frames deep in autograd. Get it right and your kernel is indistinguishable from a built-in op.

The mental model: two kernels, one node

A forward kernel reads x, writes y = f(x). Its backward is a different kernel: it reads the upstream gradient dy = ∂L/∂y plus whatever the forward saved, and writes dx = (∂y/∂x)ᵀ · dy — the vector-Jacobian product. The two kernels do not share code; they share saved tensors. The glue that makes autograd see them as one differentiable node is torch.autograd.Function.

A complete minimal example — fused scaled-SiLU, end to end

Let's do a real differentiable op, not a stub. Take a fused activation y = w · silu(x) = w · x · sigmoid(x), elementwise, where w is a single learnable scalar (a per-tensor scale, like a learned gate). It is small enough to read in one screen and exercises every part of the contract: a tensor input x, a tensor parameter w whose gradient is a reduction, and a non-differentiable flag.

The math we need. With s = sigmoid(x) and silu(x) = x·s:

Notice the asymmetry already: dx is elementwise (same shape as x), but dw is a scalar reduction over every element. That split is what makes the backward grid different from the forward grid. Here is the whole thing:

import torch, triton, triton.language as tl

# ---- forward kernel: y = w * silu(x) ----
@triton.jit
def silu_fwd_kernel(x_ptr, w_ptr, y_ptr, n, BLOCK: tl.constexpr):
    pid  = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    x = tl.load(x_ptr + offs, mask=mask).to(tl.float32)   # compute in fp32
    w = tl.load(w_ptr)                                    # scalar, broadcast
    s = tl.sigmoid(x)
    y = w * x * s
    tl.store(y_ptr + offs, y.to(tl.float16), mask=mask)

# ---- backward kernel: dx (elementwise) + partial dw (per-block reduction) ----
@triton.jit
def silu_bwd_kernel(dy_ptr, x_ptr, w_ptr, dx_ptr, dw_partial_ptr,
                    n, BLOCK: tl.constexpr):
    pid  = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    dy = tl.load(dy_ptr + offs, mask=mask, other=0.0).to(tl.float32)
    x  = tl.load(x_ptr  + offs, mask=mask, other=0.0).to(tl.float32)
    w  = tl.load(w_ptr)
    s    = tl.sigmoid(x)
    silu = x * s
    # dx = dy * dy/dx
    dydx = w * (s + silu * (1.0 - s))
    dx   = dy * dydx
    tl.store(dx_ptr + offs, dx.to(tl.float16), mask=mask)
    # dw = sum(dy * silu) — reduce this block, atomically add the partial
    dw_block = tl.sum(dy * silu, axis=0)                  # fp32 scalar
    tl.atomic_add(dw_partial_ptr, dw_block)

# ---- launchers ----
def silu_forward(x, w):
    y = torch.empty_like(x)
    n = x.numel()
    grid = lambda meta: (triton.cdiv(n, meta['BLOCK']),)
    silu_fwd_kernel[grid](x, w, y, n, BLOCK=1024)
    return y

def silu_backward(dy, x, w):
    dx = torch.empty_like(x)
    dw = torch.zeros(1, device=x.device, dtype=torch.float32)  # fp32 accumulator
    n = x.numel()
    grid = lambda meta: (triton.cdiv(n, meta['BLOCK']),)
    silu_bwd_kernel[grid](dy, x, w, dx, dw, n, BLOCK=1024)
    return dx, dw.to(w.dtype)

# ---- the autograd.Function: glues the two kernels into one node ----
class ScaledSiLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, w, inplace=False):
        x = x.contiguous()                  # guard non-contiguous (trap 3)
        y = silu_forward(x, w)
        ctx.save_for_backward(x, w)         # backward needs x and w
        return y

    @staticmethod
    def backward(ctx, dy):
        x, w = ctx.saved_tensors
        dy = dy.contiguous()                # dy can arrive non-contiguous (trap 3)
        dx, dw = silu_backward(dy, x, w)
        return dx, dw, None                 # one grad PER forward arg, in order

scaled_silu = ScaledSiLU.apply             # use this in your nn.Module

Read the backward return line carefully — it is the line everyone gets wrong. forward took three positional args (x, w, inplace), so backward returns exactly three values in that order: dx for x, dw for w, and None for inplace (a non-differentiable bool). Drop the None and autograd raises a shape mismatch with a stack trace that points nowhere useful.

Save vs recompute — the decision at the Triton level

The CUDA lesson 25 framed this as a bandwidth-vs-compute trade. At the Triton level it becomes a concrete choice about what you pass to ctx.save_for_backward vs what you recompute with a few tl ops inside the backward kernel. The rule:

• Recompute anything cheap to derive from a tensor you already hold. In the example above we saved x and recomputed s = sigmoid(x) and silu = x·s inside the backward kernel — three flops per element. Saving s instead would have cost a full extra fp16 tensor in HBM for the lifetime of the step. Recompute wins easily.
• Recompute one-scalar-per-row stats (the mean, the rstd, the row max of a softmax). They cost one reduction pass to regenerate but only one float per row to store — and they are needed by the elementwise body anyway, so the read is nearly free. This is exactly the RMSNorm rstd recompute and the FlashAttention m, ℓ recompute (lessons 20–21).
• Save full tensors you cannot regenerate — the actual inputs x, w to a matmul. There is no cheaper source; you must save them.

The interactive widget below lets you find the crossover for a row-stat feature: as the feature gets wider, saving it costs more HBM, while recomputing it costs more FLOPs. One of them wins, and the answer flips with width.

The backward grid is not the forward grid

The forward of our example launched one program per BLOCK of elements — a flat 1D grid over n. The backward launched the same 1D grid for dx, but dw is a reduction over all n elements down to a single scalar. We handled it with tl.atomic_add from every block into one fp32 cell. That is the simplest correct backward grid, and for a scalar it is fine.

For a vector parameter the choice gets real. Consider RMSNorm's dw (here w ∈ ℝ^d): each of the B·T token rows contributes a length-d partial, and you must sum across all rows. Two grids:

So you typically pick the backward grid like this:

1. For dx (same shape as x): mirror the forward grid — one program per row or per block. The map is structurally identical to forward.
2. For dw (a reduction across the batch dimension): a separate grid sized to the reduction, with either atomics (small, non-critical) or a two-pass partials buffer (deterministic, the production choice).

This is why "the backward is the harder kernel" (lesson 14): forward is one grid, backward is often two — one map and one reduce.

fp32 accumulation, and why gradients use +=

Two non-negotiables for gradient kernels, both from CUDA lesson 25's "three kernel-shaped facts":

1. Accumulate gradients in fp32. Notice dw was allocated as torch.zeros(1, dtype=torch.float32) and the in-kernel reduction tl.sum(...) ran on .to(tl.float32) values. A gradient reduction sums many small contributions; in fp16/bf16 the small terms vanish under the running total (the same swamping that breaks an fp16 sum-of-squares). The accumulator is always fp32 (tl.float32), even when the output gradient is later cast back down. Worked: summing B·T = 32,768 contributions of magnitude ~10⁻³ in bf16 (7 mantissa bits, ~2 decimal digits) loses everything past the second digit once the partial sum exceeds ~1; in fp32 (23-bit mantissa) you keep ~7 digits and the sum is correct.

2. Gradients accumulate with +=, not =. A weight used by many tokens receives a contribution from each — its gradient is the sum. That is why we used tl.atomic_add rather than a store. The same logic extends one level up: across gradient-accumulation microbatches, PyTorch leaves the previous .grad in place and your backward adds to it. So a backward kernel that stores dw instead of accumulating will silently overwrite the gradient from the previous microbatch. When autograd hands you an existing dw buffer (the w output slot), you must add into it, not clobber it.

The four traps

Interactive · save-vs-recompute crossover

A backward kernel needs a per-row feature (think the softmax probabilities, or a normalization buffer). You can save it on the forward — costs HBM you must read back — or recompute it in the backward — costs FLOPs but no extra HBM. Both costs scale with the tensor's size, so the size cancels: the verdict turns on a single ratio — how many FLOPs it takes to recompute one element versus the machine's FLOP-to-byte budget. On an H100 (3.35 TB/s HBM, ~490 TFLOP/s bf16) saving costs a round-trip of 4 bytes/element, so recompute wins until it costs more than ≈ 585 flops/element. Drag the recompute cost and watch the crossover; the size sliders only scale both bars equally.

What this sets up

You can now take any forward Triton kernel, give it a correct backward kernel, wire them with torch.autograd.Function, pick a backward grid (one map for dx, one reduce for dw), accumulate in fp32, and dodge the four traps. The next lesson applies all of this to the op every model is built from — the matmul backward, which is not one kernel but two GEMMs (dX = dY·Wᵀ, dW = Xᵀ·dY) with transposed layouts and an fp32 accumulator on each.