Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

FlashAttention Algorithm › Recomputation Strategy

Michael BrenndoerferUpdated June 29, 20255 min read

Language AI HandbookAll chapters

Part 17 of 67Efficient Attention

Chapter 124FlashAttention Algorithm › Recomputation Strategy

Reading section 5 of 12 · Recomputation Strategy · 5 minSection 5 of 12 · 5 min

Section 5 of 12 · 5 min read

FlashAttention Algorithm › Recomputation StrategyOne section · Sections in this chapter5 / 12 ↓

Sections in this chapter12

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Training neural networks requires computing gradients during backpropagation. Standard attention stores the full $n \times n$ attention matrix during the forward pass for use in the backward pass. This storage dominates memory usage for long sequences.

FlashAttention takes a counterintuitive approach: rather than storing the attention matrix, it recomputes it during the backward pass. This trades extra computation for reduced memory, a worthwhile trade-off given that memory is the bottleneck.

Why Recomputation Makes SenseLink Copied

At first glance, recomputing seems wasteful. But consider the memory-compute trade-off:

Storing attention matrix: $O(n^2)$ memory, no extra compute
Recomputing attention matrix: $O(n^2)$ extra compute, $O(n)$ memory (for running statistics only)

where $n$ is the sequence length. The quadratic memory cost of storing attention matrices grows rapidly: at $n = 8192$ with 16-bit precision, a single attention matrix consumes $8192^2 \times 2 \approx 134$ MB. Across multiple heads and layers, this quickly exhausts GPU memory.

For long sequences, the memory savings from not storing $n^2$ values outweigh the cost of recomputing them. The recomputation can also be fused with other gradient computations, and modern GPUs have ample compute capacity when not stalled on memory transfers.

In[17]:

Code

def memory_comparison(seq_len, d_model, num_heads, num_layers, dtype_bytes=2):
    """
    Compare memory usage: standard vs FlashAttention during training.

    Args:
        seq_len: Sequence length
        d_model: Model dimension
        num_heads: Number of attention heads
        num_layers: Number of transformer layers
        dtype_bytes: Bytes per element

    Returns:
        Dictionary with memory statistics
    """
    d_head = d_model // num_heads

    # Standard attention stores per layer:
    # - Attention scores: n x n x num_heads
    # - Attention weights (after softmax): n x n x num_heads
    # - Q, K, V, O: each n x d_model
    standard_attn_matrix = (
        2 * seq_len * seq_len * num_heads * dtype_bytes
    )  # Scores + weights
    standard_qkvo = 4 * seq_len * d_model * dtype_bytes
    standard_per_layer = standard_attn_matrix + standard_qkvo
    standard_total = standard_per_layer * num_layers

    # FlashAttention stores per layer:
    # - Q, K, V, O: each n x d_model
    # - Running statistics: n x num_heads x 2 (m and ell)
    flash_qkvo = 4 * seq_len * d_model * dtype_bytes
    flash_stats = seq_len * num_heads * 2 * dtype_bytes
    flash_per_layer = flash_qkvo + flash_stats
    flash_total = flash_per_layer * num_layers

    return {
        "standard_gb": standard_total / 1e9,
        "flash_gb": flash_total / 1e9,
        "savings": (standard_total - flash_total) / standard_total * 100,
        "attn_matrix_gb": standard_attn_matrix * num_layers / 1e9,
    }


# GPT-2 style configuration
config = {"d_model": 768, "num_heads": 12, "num_layers": 12}
seq_lengths_mem = [512, 2048, 8192, 32768]

Out[18]:

Console

Training Memory Comparison (GPT-2 config, fp16):

   Seq Len     Standard        Flash    Savings    Attn Matrix
--------------------------------------------------------------
       512        0.19G       0.038G      79.8%          0.15G
     2,048        2.57G       0.152G      94.1%          2.42G
     8,192       39.26G       0.609G      98.4%         38.65G
    32,768      620.89G       2.435G      99.6%        618.48G

The attention matrix dominates memory for long sequences. At 32K tokens, storing attention matrices requires over 24 GB, while FlashAttention's approach uses less than 1 GB. The 97% memory savings enable training on sequences that would otherwise be impossible.

What Gets Stored vs RecomputedLink Copied

FlashAttention stores only the minimal information needed for backpropagation:

Stored: Q, K, V, O matrices, plus running statistics (m and L) for each query position
Recomputed: Attention scores and weights during the backward pass

The running statistics enable reconstructing the softmax normalization without storing all the attention weights. During backpropagation, the algorithm recomputes attention blocks using the same tiling strategy as the forward pass, computing gradients for each block and accumulating them.

Out[19]:

Visualization

Log-scale plot showing quadratic memory growth for standard attention versus linear growth for FlashAttention. — Training memory usage as sequence length increases. Standard attention (red) grows quadratically due to storing attention matrices. FlashAttention (blue) grows linearly by recomputing rather than storing. The gap widens dramatically for long sequences.

The logarithmic plot reveals the scaling difference starkly. Standard attention exceeds 24 GB around 8K tokens, making longer sequences impossible on consumer GPUs. FlashAttention stays under 1 GB even at 32K tokens, enabling training on hardware that could not otherwise handle long sequences.