Flash Attention — Fast Memory-Efficient Attention

Flash Attention provides 2-4x training speedup and 10-20x memory reduction by replacing the standard O(N²) attention with an IO-aware tiling algorithm.

When to Use

• Sequences longer than 512 tokens → always use Flash Attention
• GPU OOM during training → Flash Attention is usually the fix
• Maximizing training throughput → 2-4x faster than standard attention
• H100/A100 with FP8/BF16 → critical for efficiency

Option 1: PyTorch Native SDPA (Easiest — PyTorch 2.2+)

import torch
import torch.nn.functional as F