FlashAttention Implementation › FlashAttention Limitations
Language AI HandbookAll chapters FlashAttention Implementation › FlashAttention LimitationsOne section · Sections in this chapter7 / 10 ↓
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
Despite its advantages, FlashAttention has limitations that you should understand before deployment.
Numerical Precision Differences
FlashAttention computes attention in a different order than standard attention, which can lead to small numerical differences. While mathematically equivalent, floating-point arithmetic is not associative, so results may differ at the level of floating-point precision.
def demonstrate_numerical_differences():
"""
Show how reordering affects floating-point results.
This illustrates why FlashAttention and standard attention
may have tiny numerical differences.
"""
# Create deterministic values from the chapter-level plot seed
values = np.random.randn(100).astype(np.float32)
# Sum in different orders
forward_sum = np.sum(values)
backward_sum = np.sum(values[::-1])
sorted_sum = np.sum(np.sort(values))
# Kahan summation (more accurate)
def kahan_sum(arr):
s = 0.0
c = 0.0
for x in arr:
y = x - c
t = s + y
c = (t - s) - y
s = t
return s
kahan = kahan_sum(values)
return {
"forward": forward_sum,
"backward": backward_sum,
"sorted": sorted_sum,
"kahan": kahan,
"max_diff": max(
abs(forward_sum - backward_sum), abs(forward_sum - sorted_sum)
),
}
num_results = demonstrate_numerical_differences()Floating-Point Summation Order Effects: Forward sum: 3.9492628574 Backward sum: 3.9492607117 Sorted sum: 3.9492614269 Kahan sum: 3.9492611885 Max difference: 2.15e-06 These tiny differences are expected and acceptable. FlashAttention trades exact bit-for-bit reproducibility for significant speedups.
Running the same experiment 1000 times with different random values shows that these differences follow a predictable pattern:
No Arbitrary Attention Masks
FlashAttention's tiled computation assumes specific attention patterns. While causal masks and no masks are well-supported, arbitrary attention masks present challenges:
def mask_support_matrix():
"""
Document mask support across FlashAttention versions.
"""
masks = {
"No mask (bidirectional)": {
"flash1": "✓",
"flash2": "✓",
"pytorch_sdpa": "✓",
},
"Causal mask": {"flash1": "✓", "flash2": "✓", "pytorch_sdpa": "✓"},
"Sliding window": {
"flash1": "Limited",
"flash2": "✓",
"pytorch_sdpa": "✓ (via flash2)",
},
"Arbitrary boolean mask": {
"flash1": "✗",
"flash2": "Limited",
"pytorch_sdpa": "Falls back to math kernel",
},
"Additive attention bias": {
"flash1": "✗",
"flash2": "✓ (FlashAttention-2.3+)",
"pytorch_sdpa": "Falls back",
},
}
return masks
masks = mask_support_matrix()Attention Mask Support: Mask Type FA-1 FA-2 PyTorch ----------------------------------------------------------------- No mask (bidirectional) ✓ ✓ ✓ Causal mask ✓ ✓ ✓ Sliding window Limited ✓ ✓ (via flash2) Arbitrary boolean mask ✗ Limited Falls back to math kernel Additive attention bias ✗ ✓ (FlashAttention-2.3+) Falls back
For most transformer applications, causal or no mask covers the use cases. FlashAttention-2 added sliding window support, which enables models like Mistral. For arbitrary boolean masks or complex attention biases, PyTorch falls back to the slower math kernel, losing FlashAttention's benefits.
Hardware Requirements
FlashAttention requires relatively recent NVIDIA GPUs:
def hardware_compatibility():
"""
Document GPU compatibility for FlashAttention.
"""
gpus = {
"Ampere (SM 80)": {
"examples": "A100, A6000, RTX 3090",
"flash1": "Full support",
"flash2": "Full support",
},
"Ada Lovelace (SM 89)": {
"examples": "RTX 4090, L40",
"flash1": "Full support",
"flash2": "Full support",
},
"Hopper (SM 90)": {
"examples": "H100",
"flash1": "Full support",
"flash2": "Full support + FP8",
},
"Turing (SM 75)": {
"examples": "RTX 2080, T4",
"flash1": "Limited (FA-1 only)",
"flash2": "Not supported",
},
"Volta (SM 70)": {
"examples": "V100",
"flash1": "Limited",
"flash2": "Not supported",
},
}
return gpus
hw = hardware_compatibility()FlashAttention GPU Compatibility: Architecture Example GPUs FA-1 FA-2 ------------------------------------------------------------------------------ Ampere (SM 80) A100, A6000, RTX 3090 Full support Full support Ada Lovelace (SM 89) RTX 4090, L40 Full support Full support Hopper (SM 90) H100 Full support Full support + FP8 Turing (SM 75) RTX 2080, T4 Limited (FA-1 only) Not supported Volta (SM 70) V100 Limited Not supported
Ampere and newer architectures provide full FlashAttention-2 support. Older GPUs like V100 and T4 have limited support for FlashAttention-1 only. If you're using cloud instances, A100 and H100 are the best choices for FlashAttention workloads. Consumer GPUs from the RTX 30 and 40 series also work well.
Memory Still Matters for Very Long Sequences
While FlashAttention eliminates the memory for attention matrices, other parts of the transformer still scale with sequence length:
def total_memory_breakdown(
n, d_model, n_layers, vocab_size, batch_size=1, dtype_bytes=2
):
"""
Break down memory usage in a transformer with FlashAttention.
"""
# Activations per layer (need to store for backprop)
# - Input to layer: n * d_model
# - After attention: n * d_model
# - After FFN: n * d_model
activations_per_layer = 3 * batch_size * n * d_model * dtype_bytes
total_activations = activations_per_layer * n_layers
# KV cache for inference
kv_cache = 2 * n_layers * batch_size * n * d_model * dtype_bytes
# Model weights (independent of n)
# Simplified: attention (4 * d^2) + FFN (8 * d^2) per layer
weights_per_layer = 12 * d_model * d_model * dtype_bytes
embeddings = vocab_size * d_model * dtype_bytes
total_weights = weights_per_layer * n_layers + embeddings
return {
"activations_gb": total_activations / 1e9,
"kv_cache_gb": kv_cache / 1e9,
"weights_gb": total_weights / 1e9,
"total_gb": (total_activations + kv_cache + total_weights) / 1e9,
}
# LLaMA 7B-style model
memory_7b = total_memory_breakdown(
n=4096, d_model=4096, n_layers=32, vocab_size=32000
)
memory_7b_long = total_memory_breakdown(
n=32768, d_model=4096, n_layers=32, vocab_size=32000
)Memory Breakdown (7B-style model, fp16): At 4K tokens: Activations: 3.22 GB KV Cache: 2.15 GB Weights: 13.15 GB Total: 18.52 GB At 32K tokens (8x longer): Activations: 25.77 GB KV Cache: 17.18 GB Weights: 13.15 GB (unchanged) Total: 56.10 GB FlashAttention eliminates attention matrix memory, but activations and KV cache still scale linearly with n.
The breakdown shows that at 4K tokens, a 7B model requires about 13 GB total. At 32K tokens, this grows to nearly 50 GB, a 4x increase. Model weights remain constant, but activations and KV cache scale linearly with sequence length. This explains why even with FlashAttention, very long sequences require substantial GPU memory or additional techniques like gradient checkpointing.
Gradient Computation
FlashAttention's recomputation strategy during the backward pass means gradients take longer to compute than they would with stored attention matrices. This is typically worthwhile because:
- The memory savings allow larger batch sizes, which often more than compensate for the recomputation cost
- Training is usually memory-bound anyway, so the extra compute fits within memory transfer time
- For very long sequences, the alternative (storing attention matrices) would be impossible