FromZero2AI
Back to Blog

FlashAttention: GPU Memory-Bound Speedups for Long-Context LLMs

•
9 min read

The self-attention mechanism is the core component of the transformer architecture. However, self-attention has a major limitation: its computational and memory complexity scales quadratically with the sequence length:

O(N^2)

For sequence length N, calculating attention requires constructing an N \times N matrix. For long contexts (e.g. 16k or 32k tokens), this matrix is massive, saturating GPU memory (High Bandwidth Memory, or HBM) and slowing down execution.

FlashAttention resolves this memory bottleneck by restructuring the attention calculation to run efficiently on GPU memory hierarchies.

FlashAttention Memory OptimizationFlashAttention Memory Optimization

The Memory Hierarchy Bottleneck

Standard GPU execution relies on two primary memory tiers:

  • HBM (High Bandwidth Memory): Large capacity (e.g. 80GB) but relatively slow access speeds.
  • SRAM (Static Random-Access Memory): Fast, on-chip memory located near the GPU cores, but extremely small capacity (e.g. 20MB).

In standard self-attention, intermediate results like the attention matrix S and Softmax outputs P are continuously written to and read from HBM. FlashAttention eliminates these slow read/write cycles by loading inputs in blocks (tiles) directly into SRAM, computing attention locally, and writing the final outputs back to HBM.

Core Mechanisms of FlashAttention

FlashAttention achieves this using three key techniques:

  1. Tiling: Splitting the Query, Key, and Value matrices into blocks that fit within the SRAM limit.
  2. Online Softmax: Normalizing softmax exponents across blocks. Standard softmax requires seeing the entire row to find the maximum value:
m_i = \max(x_i)

Online softmax updates the maximum value and scaling factors incrementally as each block is processed, ensuring the final output is mathematically identical to standard softmax. 3. Recomputation in Backward Pass: Instead of storing the large N \times N attention matrix for backpropagation, FlashAttention recomputes it on-the-fly during the backward pass using the tiled inputs stored in SRAM.

Code Verification: Enabling FlashAttention in PyTorch

Using PyTorch's native scaled dot-product attention (SDPA), FlashAttention is enabled automatically when supported by the hardware:

import torch

# Define dimensions: batch, heads, seq_len, head_dim
query = torch.randn(2, 8, 4096, 64, dtype=torch.float16, device="cuda")
key = torch.randn(2, 8, 4096, 64, dtype=torch.float16, device="cuda")
value = torch.randn(2, 8, 4096, 64, dtype=torch.float16, device="cuda")

# Run scaled dot product attention with FlashAttention backend
# with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False):
#     output = torch.nn.functional.scaled_dot_product_attention(query, key, value)

Performance Impact

  • Speed: Up to 3x speedups in training and inference.
  • Memory: Drastically reduced memory usage, allowing transformers to scale to much longer context windows on existing hardware.

Conclusion

FlashAttention is a fundamental optimization for transformer architectures. Restructuring attention calculations to leverage fast GPU SRAM enables faster execution, lower memory overhead, and support for larger context windows.