Backward Pass Implementation

In the previous chapters, we implemented a complete Flash Attention forward pass, supporting arbitrary sequence lengths, causal masking, and GQA. But this is not enough: to use Flash Attention in training, we need to implement the backward pass to compute gradients.

This chapter explores how to compute $\frac{\partial L}{\partial Q}$, $\frac{\partial L}{\partial K}$, and $\frac{\partial L}{\partial V}$ while preserving IO efficiency.

Why Do We Need a Custom Backward Pass?

Limitations of PyTorch Autograd

PyTorch's autograd can automatically generate backward code for most operations, but for a fused kernel like Flash Attention, automatic differentiation runs into problems:

# Standard attention forward pass
def standard_attention(Q, K, V):
    S = Q @ K.T / sqrt(d)     # (1) compute scores
    P = softmax(S, dim=-1)    # (2) softmax
    O = P @ V                 # (3) weighted sum
    return O

Autograd's behavior:

• PyTorch saves the intermediate results $S$ and $P$ (the attention matrix)
• It uses these intermediates to compute gradients during the backward pass
• Memory cost: $O(N^2)$ to store the attention matrix

But the core advantage of Flash Attention is precisely that it does not materialize the attention matrix!

If we rely on automatic differentiation, we would:

1. Forward pass: do not save the attention matrix, memory $O(N)$ ✅
2. Backward pass: need the attention matrix, forced to recompute it, losing the advantage ❌

The Recomputation Strategy

Flash Attention adopts a clever tradeoff: recomputation.

Core idea:

• During the forward pass: save only a small number of intermediate statistics ($\mathbf{L}$ and $\mathbf{M}$)
• During the backward pass: use these statistics to recompute the attention scores, rather than loading them from HBM

Traditional approach (materialization):
  Forward:  compute S, P → save to HBM (O(N²) memory)
  Backward: read S, P from HBM → compute gradients

Flash Attention (recomputation):
  Forward:  compute S, P → save only L, M (O(N) memory)
  Backward: recompute S, P → compute gradients

Tradeoff analysis:

• ✅ Memory savings: $O(N^2) \to O(N)$
• ⚠️ Extra computation: need to recompute attention (about 1.5x FLOPS)
• ✅ IO efficiency: still faster than the standard implementation, because it avoids HBM round trips