Flash Attention v2¶

Overview¶

Flash Attention v2 is the improved version of Flash Attention released in 2023. It achieves 2.8x speedup compared to standard attention (vs 2x in v1) with better numerical stability, optimized backward pass, and support for more variants.

Announcement: July 2023 (9 months after v1)
Key Improvement: 40% faster than v1
Main Focus: Backward pass optimization + numerical improvements
Adoption: Now standard in PyTorch, HuggingFace, vLLM
Backward Compatibility: Yes, mostly compatible with v1

-

Flash Attention v1: Recap¶

What v1 Solved¶

Problem: Standard attention needs O(N²) memory for intermediate matrices

Standard Attention:
 - S = Q @ K^T (N × N matrix) → O(N²) memory
 - Softmax(S) → O(N²) memory 
 - P = Softmax(S) @ V
 - High memory, I/O-bound

Flash Attention v1:
 - Block-wise computation (tiling)
 - Keep S, P in fast SRAM cache
 - O(N²) I/O reduced to O(N)
 - 2x faster than standard attention

v1 Characteristics¶

Strengths:
2x speedup
60% memory reduction
Exact computation (bit-identical)
Forward pass well optimized

Limitations:
Backward pass not fully optimized
Limited head dimension support (≤ 128)
Some numerical precision issues
Not work well with all sequence lengths
Gradient computation overhead

-

Flash Attention v2: Major Improvements¶

Key Innovation 1: Backward Pass Optimization¶

FLASH ATTENTION V1 - Backward Pass (Inefficient):

Forward pass: Optimized 
Backward pass: Standard computation 

Backward computation requires:
 - Recompute attention matrix S (O(N²) I/O!)
 - Compute attention gradients
 - Complex synchronization
 - Result: Backward slower than forward (unusual!)

Typical timing:
 - Forward: 10ms (well optimized)
 - Backward: 15-20ms (not optimized)

FLASH ATTENTION V2 - Backward Pass (Optimized):

Apply same I/O-aware principles to backward:
 - Block-wise gradient computation
 - Minimize intermediate materialization
 - Efficient synchronization patterns
 - Result: Backward ≈ forward speed

Typical timing:
 - Forward: 10ms (further optimized)
 - Backward: 10-12ms (now optimized!)

Speedup: 2x faster backward pass!

Key Innovation 2: Block Partitioning Optimization¶

FLASH ATTENTION V1:

Block strategy:
 - Fixed block size (e.g., 64 × 64)
 - Process Q in blocks of size B_r = 64
 - Process K,V in blocks of size B_c = 64
 - Suboptimal for different hardware

Memory access pattern:
 - Column-wise access to Q (inefficient for memory bandwidth)
 - Thread synchronization overhead
 - Not fully utilizing GPU memory hierarchy

FLASH ATTENTION V2:

Improved blocking:
 - Adaptive block size based on hardware
 - Optimize for memory coalescing
 - Column-wise vs row-wise access selection
 - Better memory bandwidth utilization

Key change: Reorder loop structure!

V1 loop order:
for i in range(0, N, B_r): # Iterate Q blocks
 for j in range(0, N, B_c): # Iterate K,V blocks
 # Process each pair

V2 loop order (optimized):
for i in range(0, N, B_r): # Iterate Q blocks
 for j in range(0, N, B_c): # Iterate K,V blocks in reverse!
 # Better memory pattern for modern GPUs

Result: 40% faster than v1!

Key Innovation 3: Numerical Stability¶

FLASH ATTENTION V1:

Softmax computation:
 - Subtract max for stability
 - Compute exp
 - Divide by sum
 - Generally stable but edge cases exist

Issues:
 - Accumulation of floating point errors
 - Different precision behavior across platforms
 - Gradient computation could introduce errors

FLASH ATTENTION V2:

Improvements:
 - Better numerical precision in softmax
 - More stable gradient computation
 - Double-precision intermediate calculations
 - Careful handling of small values

Result:
 - Bit-identical to standard attention
 - Fewer numerical errors
 - Better training stability

-

Detailed Differences¶

1. Forward Pass Optimization¶

FLASH ATTENTION V1 Forward:

```cuda
// Simplified v1 structure
for i in range(0, N, B_r):
 // Load Q block
 Q_block = load_from_global(Q[i:i+B_r])

 for j in range(0, N, B_c):
 // Load K,V blocks
 K_block = load_from_global(K[j:j+B_c])
 V_block = load_from_global(V[j:j+B_c])

 // Compute S, P in shared memory
 S = Q_block @ K_block.T
 P = softmax(S)

 // Accumulate output
 O += P @ V_block

FLASH ATTENTION V2 Forward:

// Improved v2 structure with better memory access
for i in range(0, N, B_r):
 // Load Q block ONCE (better locality)
 Q_block = load_from_global(Q[i:i+B_r])

 for j in range(0, N, B_c):
 // Reverse iteration improves memory pattern
 K_block = load_from_global(K[N-j-B_c:N-j])
 V_block = load_from_global(V[N-j-B_c:N-j])

 // Same computation, but:
 // - Coalesced memory access
 // - Better cache utilization
 S = Q_block @ K_block.T
 P = softmax(S)
 O += P @ V_block

Performance difference:

V1: 2.0x vs standard
V2: 2.8x vs standard (40% faster!)

### 2. Backward Pass Computation

FLASH ATTENTION V1 Backward:

Problem: Must recompute forward to get gradients

Recompute S matrix from Q,K (O(N²) computation)
Use standard backward algorithm
Separate memory access patterns
Inefficient!

Timing:

Forward: 10ms (optimized)
Backward: 20ms (not optimized, 2x slower!)
Total for training: 30ms

FLASH ATTENTION V2 Backward:

Solution: Apply I/O-aware principles to gradients too!

Key insight: Recomputation is okay if done efficiently

Recompute S block-by-block (I/O-aware)
Compute dP, dQ, dK, dV with blocking
Minimize intermediate storage
Efficient!

Algorithm structure:

# Simplified backward algorithm v2
for i in range(0, N, B_r):
 Q_block = load(Q[i:i+B_r])
 dQ_block = zeros()

 for j in range(0, N, B_c):
 K_block = load(K[j:j+B_c])
 V_block = load(V[j:j+B_c])

 # Recompute S, P for this block (efficient!)
 S = Q_block @ K_block.T
 P = softmax(S)

 # Compute gradients (in-place, minimizing memory)
 dV_block = P.T @ dO
 dP = dO @ V_block.T
 dS = dP * P * (1 - P) # softmax gradient
 dQ_block += dS @ K_block
 dK_block = dS.T @ Q_block
 dV_block += P.T @ dO

Timing:

Forward: 10ms (further optimized)
Backward: 12ms (now optimized, nearly equal!)
Total for training: 22ms (27% faster!)

### 3. Head Dimension Support

FLASH ATTENTION V1:

Limitations:

head_dim ≤ 128 supported
Larger head dims fall back to standard attention
Restricts model architectures

Why the limit:

Shared memory constraints (48-96 KB)
Need to store S matrix in shared memory
S size: B_r × B_c × sizeof(float)
For head_dim > 128: too much shared memory
Falls back to standard (slow!)

FLASH ATTENTION V2:

Support improved:

head_dim up to 256 with optimization
Adaptive memory allocation
Uses register blocking for larger dims
No fallback to standard attention

Technique: Register blocking for large head_dim

Store S in registers instead of shared memory
More complex indexing, but works
Maintains I/O efficiency
Enables support for larger heads

### 4. Batch and Sequence Length Support

FLASH ATTENTION V1:

Issues:

Performance sensitive to sequence length
Uneven performance across different lengths
Some lengths perform worse than others

Example (A100, head_dim=64):

seq_len=512: 200 TFLOPS
seq_len=1024: 220 TFLOPS
seq_len=1536: 160 TFLOPS (dip!)
seq_len=2048: 210 TFLOPS
Inconsistent behavior

FLASH ATTENTION V2:

More consistent:

Optimized block sizes for all lengths
Consistent performance across range
No pathological cases

Example (A100, head_dim=64):

seq_len=512: 250 TFLOPS (better)
seq_len=1024: 270 TFLOPS (better)
seq_len=1536: 270 TFLOPS (no dip!)
seq_len=2048: 280 TFLOPS (better)
Consistent, faster performance

### 5. Gradient Accumulation

FLASH ATTENTION V1:

Gradient computation:

Requires materialization of full attention matrix P
dP = dO @ V.T (N × N matrix in memory!)
Compute dS = dP P (1 - P)
Complex intermediate storage

Memory usage:

Attention matrix P: N² elements
Gradient matrix dP: N² elements
Total: 2N² temporary storage

FLASH ATTENTION V2:

Improved gradient flow:

Accumulate gradients without materializing full matrices
dP computed in blocks
Gradients reduced directly
Avoid large intermediate tensors

Memory usage:

Only block-size intermediate storage
Scales as O(B_r × B_c) instead of O(N²)
Huge reduction for long sequences

-

## Performance Comparison: v1 vs v2

### Benchmark Results

Hardware: A100 GPU Model: Llama 2 7B Batch size: 1

v1 v2 Improvement ───────────────────────────────────────────────── Forward Pass 2.0x 2.8x 40% Backward Pass 1.2x 2.0x 67% Full Forward+Back 1.5x 2.4x 60%

For 256 token generation: V1: 6.4s (2.0x speedup) V2: 4.0s (2.8x speedup) Improvement: 37.5%

Batch size: 32 (continuous batching) V1: 3.2 req/sec V2: 4.5 req/sec Improvement: 40%

### Training Speed Improvement

Training Llama 2 7B on 8x A100 GPUs

Metric Without FA With FA v1 With FA v2 Total Speedup ───────────────────────────────────────────────────────────────────── Tokens/sec/GPU 1000 2000 2800 2.8x Batched TFLOPS 150 300 420 2.8x Wall-clock time 100 hours 50 hours 35 hours 2.86x

Full training (1 trillion tokens):

Without FA: 100,000 GPU-hours
With FA v1: 50,000 GPU-hours
With FA v2: 35,000 GPU-hours
Savings with v2: 65% reduction!

### Where v2 Excels Most

Backward Pass (Training):
V1: 1.2x improvement
V2: 2.0x improvement (67% better than v1)
Long Sequences:
V1: 2.0x at 4K tokens
V2: 2.8x at 4K tokens (40% better)
Large Head Dimensions:
V1: Falls back to standard at head_dim > 128
V2: Maintains optimization up to head_dim=256
Numerical Stability:
V1: Generally stable, some edge cases
V2: Better precision, more stable gradients

-

## Implementation Differences

### Installation and Usage

```python
# FLASH ATTENTION V1

from flash_attn import flash_attn_func

# Basic usage
output = flash_attn_func(q, k, v)

# With dropout
output = flash_attn_func(
 q, k, v,
 dropout_p=0.1,
 causal=True
)

# Note

# FLASH ATTENTION V2

from flash_attn import flash_attn_func # Same import!

# Usage is identical (backward compatible)
output = flash_attn_func(q, k, v)

# But now V2 is automatically used:
# - Faster forward pass
# - Faster backward pass
# - Better support for large head dims
# - More stable gradients

# Verify version
import flash_attn
print(flash_attn.__version__) # Should be >= 2.0

CUDA Kernel Differences¶

FLASH ATTENTION V1 Kernel Structure:

__global__ void flash_attn_fwd_kernel_v1(...) {
 // Load Q block
 // For each K,V block:
 // Compute S in shared memory
 // Compute softmax
 // Accumulate output
}

__global__ void flash_attn_bwd_kernel_v1(...) {
 // Recompute S (I/O overhead!)
 // Compute gradients
 // Complex synchronization
}

FLASH ATTENTION V2 Kernel Structure:

__global__ void flash_attn_fwd_kernel_v2(...) {
 // Improved block ordering
 // Better memory coalescing
 // Optimized for modern GPU architectures
 // Support for larger head dimensions
}

__global__ void flash_attn_bwd_kernel_v2(...) {
 // I/O-aware gradient computation
 // Minimize intermediate materialization
 // Efficient recomputation
 // Better numerical stability
}

Key difference:
V2 kernels are "GPU architecture aware"
 - Detect GPU type (A100 vs H100)
 - Adapt block sizes and memory patterns
 - Auto-tune for best performance

When to Use v2 vs v1¶

Use Flash Attention v2¶

Scenarios where v2 shines:

Training (backward pass matters)
 - 60% faster training than v1

Long sequences (> 1K tokens)
 - Consistent 2.8x speedup vs 2x for v1

Large head dimensions (> 128)
 - v1 might fall back to standard

Inference with batching
 - 40% better throughput

Numerical stability critical
 - Better precision and stability

Modern GPUs (A100, H100)
 - Architecture-optimized kernels

Recommendation: Use v2 by default!

Reasons to Stick with v1¶

Rare cases where v1 might be preferred:

 Very old GPU hardware (< A100)
 - v1 might have better compatibility

 Legacy code that's heavily tested
 - v1 proven to work, not worth updating

 Custom CUDA kernels built on v1
 - Would need rewrite for v2

 Memory constraints (unlikely but possible)
 - v1 uses slightly less memory in some cases

Recommendation: Migrate to v2 (benefits outweigh costs)

Migration from v1 to v2¶

Is Migration Needed?¶

Good news: Flash Attention v2 is backward compatible!

Code written for v1 works with v2:
```python
from flash_attn import flash_attn_func

# This code works with both v1 and v2
output = flash_attn_func(q, k, v, causal=True)

# Automatically uses v2 if installed

How to Update¶

# STEP 1

# Current version
pip install flash-attn

# Specify v2 explicitly
pip install "flash-attn>=2.0"

# Verify installation
python -c "import flash_attn; print(flash_attn.__version__)"

# STEP 2
# v2 is backward compatible with v1

# STEP 3
# - Check training speed (should be faster)
# - Check inference latency (should be better)
# - Run tests to ensure numerical correctness

# Optional
if hasattr(flash_attn, 'flash_attn_kvpacked_func'):
 print("Using Flash Attention v2 (has v2-specific functions)")
else:
 print("Using Flash Attention v1")

-

Real-World Impact: Training Example¶

Scenario: Fine-tune Llama 2 7B on 100K examples

Setup:
 - Model: Llama 2 7B
 - Batch size: 32
 - Sequence length: 2048
 - Hardware: 8x A100 GPUs
 - Epochs: 3

WITH FLASH ATTENTION V1:

Training time:
 - Tokens processed: 100K × 2048 × 3 = 614.4B tokens
 - Throughput: 200 tokens/sec per GPU (v1 optimized)
 - Total time: 614.4B / 200 / 8 GPUs = ~385 hours
 - Cost: 385 hours × $8 per GPU-hour = $3,080
 - Actual: ~16 days (accounting for overhead)

WITH FLASH ATTENTION V2:

Training time:
 - Same 614.4B tokens
 - Throughput: 280 tokens/sec per GPU (v2, 40% better)
 - Total time: 614.4B / 280 / 8 GPUs = ~275 hours
 - Cost: 275 hours × $8 per GPU-hour = $2,200
 - Actual: ~11 days (accounting for overhead)

Improvement:
 - Time saved: 5 days (31% faster)
 - Cost saved: $880 (29% cheaper)
 - Same quality model!

-

Detailed Algorithm Comparison¶

Forward Pass Comparison¶

FLASH ATTENTION V1 - Forward

Algorithm Flash Attention:
for i = 1 to N do
 Load Q_i from HBM // Q block
 m_i ← -∞, l_i ← 0 // Tracking variables
 for j = 1 to N do
 Load K_j, V_j from HBM
 S_ij ← Q_i K_j^T / √d

 // Online softmax
 m_ij^new ← max(m_i, max_row(S_ij))
 P_ij ← exp(S_ij - m_ij^new)
 l_i^new ← exp(m_i - m_ij^new) l_i + sum_row(P_ij)

 // Accumulate output
 O_i ← exp(m_i - m_ij^new) O_i + P_ij V_j

 // Update tracking
 m_i ← m_ij^new
 l_i ← l_i^new
 end
 Write O_i to HBM
 Write l_i, m_i to HBM // For backward
end

I/O complexity: O(N × d)
FLOP complexity: O(N² × d)
Peak throughput: 2.0x vs standard

FLASH ATTENTION V2 - Forward

Algorithm Flash Attention V2:
for i = 1 to N step B_r do
 Load Q block from HBM
 m ← -∞, l ← 0

 for j = N step -B_c do // REVERSED iteration!
 Load K, V blocks from HBM

 S ← Q @ K^T (block computation)

 // Improved numerical stability
 m_new ← max(m, max_row(S), precision=fp32)
 P ← exp(S - m_new)
 l_new ← exp(m - m_new) l + sum_row(P)

 // Accumulate output
 O ← exp(m - m_new) O + P @ V

 // Update
 m ← m_new
 l ← l_new
 end
 Write O block to HBM
end

I/O complexity: O(N × d) (same)
FLOP complexity: O(N² × d) (same)
Peak throughput: 2.8x vs standard (40% better!)

Key improvement: Better memory access pattern from reversed iteration
+ Improved numerical precision tracking
+ No separate backward pass overhead

-

Comparison Table: v1 vs v2¶

Feature Flash Attn v1 Flash Attn v2
──────────────────────────────────────────────────────────
Forward pass speedup 2.0x 2.8x
Backward pass speedup 1.2x 2.0x
Combined training speedup 1.5x 2.4x
Memory reduction 60% 60%
Max head dimension 128 256+
Sequence length support All Better consistency
Numerical stability Good Excellent
Backward compatible N/A Yes (v1 code works)
GPU compatibility A100+ A100+ (better)
Release date May 2022 July 2023
Status Maintained Current standard

Recommendation: Always use v2 (better in every way)

-

Key Differences Summary¶

Core Algorithm Changes in v2:

1. LOOP STRUCTURE
 - V1: Forward iteration order
 - V2: Reversed K,V iteration (better memory pattern)

2. NUMERICAL PRECISION
 - V1: Float16 tracking
 - V2: Float32 intermediate precision (more stable)

3. BACKWARD PASS
 - V1: Separate optimization
 - V2: Full I/O-aware optimization (like forward)

4. HEAD DIMENSION SUPPORT
 - V1: Limited to 128 (shared memory constraint)
 - V2: Up to 256+ (register blocking)

5. ARCHITECTURE AWARENESS
 - V1: Generic kernel
 - V2: Auto-tunes for specific GPU (A100 vs H100)

Result: 40% faster forward, 67% faster backward

Key Takeaways¶

Flash Attention v2 is 40% faster than v1 on forward pass 67% faster backward pass (most impactful for training) 2.8x speedup vs standard (vs 2.0x for v1) Backward compatible with v1 code Better support for larger head dimensions (up to 256) Production standard now (use v2 by default)

-

Flash Attention v2¶

Overview¶

Flash Attention v1: Recap¶

What v1 Solved¶

v1 Characteristics¶

Flash Attention v2: Major Improvements¶

Key Innovation 1: Backward Pass Optimization¶

Key Innovation 2: Block Partitioning Optimization¶

Key Innovation 3: Numerical Stability¶

Detailed Differences¶

1. Forward Pass Optimization¶

CUDA Kernel Differences¶

When to Use v2 vs v1¶

Use Flash Attention v2¶

Reasons to Stick with v1¶

Migration from v1 to v2¶

Is Migration Needed?¶

How to Update¶

Real-World Impact: Training Example¶

Detailed Algorithm Comparison¶

Forward Pass Comparison¶

Comparison Table: v1 vs v2¶

Key Differences Summary¶

Key Takeaways¶

Further Reading¶

Related Notes¶