12. Manual Backward Pass for Scaled Dot-Product Attention

Problem

Calling .backward() on an attention layer just works because autograd tracks every op in softmax(QK^T / sqrt(d)) V automatically — but a kernel author writing a fused/custom attention implementation (exactly the kind of work behind FlashAttention-style kernels) doesn't get that for free: the forward pass is a hand-written fused kernel, so the backward pass has to be hand-derived and hand-written too. This problem is that derivation, done in plain PyTorch ops instead of a real kernel, to test the same underlying understanding of what .backward() is actually computing through an attention layer.

Implement a custom torch.autograd.Function for (single-head, non-causal) scaled dot-product attention with both forward and backward written by hand, deriving dQ, dK, dV from first principles rather than letting autograd do it.

Source: src/12_manual_backward_attention.py

class ManualAttention(torch.autograd.Function):
    @staticmethod
    def forward(ctx, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: ...
    @staticmethod
    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ...

def manual_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: ...

>>> q = k = v = torch.randn(1, 5, 8, requires_grad=True)
>>> out = manual_attention(q, k, v)
>>> out.shape
torch.Size([1, 5, 8])
>>> out.sum().backward()  # runs the hand-written backward, not autograd's

Step-by-Step Approach

1. In forward, compute S = QK^T / sqrt(d), P = softmax(S), O = P @ V as usual, but save Q, K, V, P (not S) via ctx.save_for_backward — the backward math below only ever needs P, never the raw pre-softmax scores.
2. Start the backward derivation from the last op, O = P @ V: standard matmul-gradient rules give dV = P^T @ dO and dP = dO @ V^T.
3. Derive the softmax backward: for a row-wise softmax P = softmax(S), the vector-Jacobian product is dS = P * (dP - sum(dP * P, dim=-1, keepdim=True)) — this is the standard closed-form softmax Jacobian applied to an upstream gradient, worth having memorized or re-derivable from dS_i = sum_j P_i (delta_ij - P_j) dP_j.
4. Finish by differentiating S = (QK^T) * scale back through to Q and K: dQ = (dS @ K) * scale and dK = (dS^T @ Q) * scale.
5. Return exactly one gradient per forward input, in the same order (dQ, dK, dV) — torch.autograd.Function.backward must return a value (or None) for every positional argument forward received besides ctx.
6. Validate by comparing both the forward output AND the resulting gradients against a plain autograd-tracked reference implementation fed the identical upstream gradient — matching output alone doesn't prove backward is implemented correctly.

The key insight is that every matmul in the forward pass contributes its OWN pair of gradient terms during backward (in reverse order), and softmax's Jacobian is the one non-matmul step in the chain that needs its own closed-form rule rather than being read straight off the forward formula.

Reference solution

import math

import torch


class ManualAttention(torch.autograd.Function):
    @staticmethod
    def forward(ctx, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        D = q.shape[-1]
        scale = 1.0 / math.sqrt(D)

        S = torch.matmul(q, k.transpose(-1, -2)) * scale  # (B, T, T)
        P = torch.softmax(S, dim=-1)                       # (B, T, T)
        O = torch.matmul(P, v)                              # (B, T, D)

        # save P (not S) -- the softmax backward formula below only needs P,
        # plus q/k/v themselves for the final two matmul-gradient steps
        ctx.save_for_backward(q, k, v, P)
        ctx.scale = scale
        return O

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        q, k, v, P = ctx.saved_tensors
        scale = ctx.scale
        dO = grad_output  # (B, T, D)

        # O = P @ V  ->  dV = P^T @ dO,  dP = dO @ V^T
        dV = torch.matmul(P.transpose(-1, -2), dO)  # (B, T, D)
        dP = torch.matmul(dO, v.transpose(-1, -2))   # (B, T, T)

        # softmax backward: dS_i = P_i * (dP_i - sum_j(dP_i,j * P_i,j))
        # this is the standard Jacobian-vector product for softmax, applied
        # row-wise (each query's attention distribution is independent)
        dS = P * (dP - (dP * P).sum(dim=-1, keepdim=True))  # (B, T, T)

        # S = (Q @ K^T) * scale  ->  dQ = (dS @ K) * scale,  dK = (dS^T @ Q) * scale
        dQ = torch.matmul(dS, k) * scale  # (B, T, D)
        dK = torch.matmul(dS.transpose(-1, -2), q) * scale  # (B, T, D)

        return dQ, dK, dV


def manual_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
    return ManualAttention.apply(q, k, v)


TEST_CASES = [
    {"name": "forward output matches F.scaled_dot_product_attention", "B": 1, "T": 5, "D": 8},
    {"name": "gradients match an autograd-tracked reference implementation", "B": 2, "T": 4, "D": 6},
    {"name": "single-example batch, larger D", "B": 1, "T": 3, "D": 16},
]


def main():
    import torch.nn.functional as F

    torch.manual_seed(0)

    for i, case in enumerate(TEST_CASES):
        B, T, D = case["B"], case["T"], case["D"]
        print(f"Test {i}: {case['name']} (B={B}, T={T}, D={D})")

        q = torch.randn(B, T, D, requires_grad=True)
        k = torch.randn(B, T, D, requires_grad=True)
        v = torch.randn(B, T, D, requires_grad=True)

        out = manual_attention(q, k, v)
        print(f"  forward out.shape={tuple(out.shape)}")
        assert out.shape == (B, T, D)

        expected_out = F.scaled_dot_product_attention(q.detach(), k.detach(), v.detach())
        torch.testing.assert_close(out, expected_out, atol=1e-4, rtol=1e-4)

        grad_output = torch.randn(B, T, D)
        out.backward(grad_output)

        q_ref = q.detach().clone().requires_grad_(True)
        k_ref = k.detach().clone().requires_grad_(True)
        v_ref = v.detach().clone().requires_grad_(True)
        out_ref = F.scaled_dot_product_attention(q_ref, k_ref, v_ref)
        out_ref.backward(grad_output)

        print(f"  dQ norm={q.grad.norm().item():.4f} (ref {q_ref.grad.norm().item():.4f})")
        torch.testing.assert_close(q.grad, q_ref.grad, atol=1e-4, rtol=1e-4)
        torch.testing.assert_close(k.grad, k_ref.grad, atol=1e-4, rtol=1e-4)
        torch.testing.assert_close(v.grad, v_ref.grad, atol=1e-4, rtol=1e-4)

        print("PASSED")


if __name__ == "__main__":
    main()

Key Functions & Tricks

• ctx.save_for_backward(q, k, v, P) — saves exactly the four tensors the hand-derived backward needs; saving the post-softmax P instead of raw scores avoids redoing the softmax in backward.
• dV = P^T @ dO — the gradient of a matmul O = P @ V with respect to its second argument, by the standard matmul backward rule.
• dS = P * (dP - (dP * P).sum(-1, keepdim=True)) — the closed-form softmax Jacobian-vector product; the subtracted term is exactly what a sum-normalized distribution's Jacobian requires.
• dQ = (dS @ K) * scale, dK = (dS^T @ Q) * scale — the two matmul-gradient terms for S = (Q @ K^T) * scale, differentiated with respect to each of its two matmul operands.
• out.backward(grad_output) with an explicit tensor — used in testing instead of .backward() with an implicit scalar reduction, so the exact upstream gradient reaching the custom backward is controlled and reproducible.

How to Recognize This Pattern