LogoCookLLM Docs
LogoCookLLM Docs
HomeCookLLM

Principles

Tokenization
Tokenization BasicsBPE AlgorithmGPT TokenizersBPE Training Engineering
Model Architecture
Transformer LM
From token ids to logitsEmbedding and LM Head
Attention Mechanisms
From Self-Attention to GQAAttention Sink
Position Encoding
Position Encoding BasicsRoPE Math DerivationRoPE ImplementationLength Extrapolation
GPU Programming Basics
GPU Architecture BasicsTensor LayoutTriton Basics: Vector Add
FlashAttention
Flash Attention PrinciplesFrom Naive Implementation to Auto-TuningBlock Pointers and Multi-Dim SupportCausal Masking OptimizationGrouped Query AttentionBackward Pass Implementation
Distributed Training
Data ParallelismZeRO OptimizerFully Sharded Data ParallelTensor ParallelismPipeline ParallelismMulti-Dimensional Hybrid Parallelism

Hands-on Training

Overview
Pretraining
Pretraining DataTokenizer TrainingModel ArchitectureData PipelineTraining LoopMonitoring and Validation
X (Twitter)
SystemsFlashAttention

Backward Pass Implementation

Premium

Implement Flash Attention gradient computation, achieving memory-efficient training through recomputation.

Get code access

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 ∂L∂Q\frac{\partial L}{\partial Q}∂Q∂L​, ∂L∂K\frac{\partial L}{\partial K}∂K∂L​, and while preserving IO efficiency.

Log in to continue reading

This is premium content. Please log in to access the full article.

∂L∂V\frac{\partial L}{\partial V}∂V∂L​

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

Autograd's behavior:

  • PyTorch saves the intermediate results SSS and PPP (the attention matrix)
  • It uses these intermediates to compute gradients during the backward pass
  • Memory cost: O(N2)O(N^2)O(N2) 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)O(N)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 (L\mathbf{L}L and M\mathbf{M}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(N2)→O(N)O(N^2) \to O(N)O(N2)→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

Why is recomputation still faster?

Although it increases FLOPS, modern GPUs are IO-bound rather than compute-bound. The FLOPS cost of recomputing attention is smaller than the IO cost of loading the O(N2)O(N^2)O(N2) matrix from HBM.

This is why Flash Attention can preserve its performance advantage even in the backward pass.

Grouped Query Attention

Add GQA/MQA support so multiple query heads share KV, reducing KV cache memory.

Distributed Training

From data parallelism to multi-dimensional hybrid parallelism — understanding the core parallel strategies of large model training

Table of Contents

Why Do We Need a Custom Backward Pass?
Limitations of PyTorch Autograd
The Recomputation Strategy
The Math of the Attention Backward Pass
Forward Pass Recap
Gradient Derivation
1. ∂L∂V\frac{\partial \mathcal{L}}{\partial \mathbf{V}}∂V∂L​ (the simplest)
2. ∂L∂P\frac{\partial \mathcal{L}}{\partial \mathbf{P}}∂P∂L​ (intermediate gradient)
3. ∂L∂S\frac{\partial \mathcal{L}}{\partial \mathbf{S}}∂S∂L​ (softmax backward pass)
4. ∂L∂Q\frac{\partial \mathcal{L}}{\partial \mathbf{Q}}∂Q∂L​ and ∂L∂K\frac{\partial \mathcal{L}}{\partial \mathbf{K}}∂K∂L​
The Complete Gradient Computation Flow
Code Implementation Walkthrough
Modifications to the Forward Kernel
Backward Kernel Implementation
Key Implementation Details
1. Reversal of the Loop Order
2. Atomic Add for dQ
3. Recompute P Rather Than Save It
Wrapping with torch.autograd.Function
Performance Validation
Numerical Correctness Test
Memory Footprint Comparison
Design Tradeoffs and Optimization Directions
The Cost of Recomputation
Further Optimization Directions
Summary
=
P
@
V
# (3) weighted sum
return O