Flash Attention¶

Overview¶

Flash Attention is a fast and memory-efficient attention algorithm that reduces the time and memory complexity of transformer attention through I/O-aware computation. It achieves 2-3x speedup and uses significantly less memory than standard attention.

Paper: "Flash Attention: Fast and Memory-Efficient Exact Attention with IO-Awareness"
Authors: Dao et al., Stanford/MIT-IBM (2022)
Key Innovation: Minimize data movement (I/O), not just FLOPs
Impact: 2-3x faster attention, reduced memory footprint
Adoption: Standard in vLLM, HuggingFace, PyTorch
Versions: Flash-1, Flash-2 (2x faster), Flash-3 (in progress)

-

The Problem: Standard Attention Bottleneck¶

Standard Attention Computation¶

Input:
Q: Query matrix (N × d_k)
K: Key matrix (N × d_k)
V: Value matrix (N × d_v)
N: Sequence length
d_k: Key dimension
d_v: Value dimension

Standard Algorithm:
1. S = Q @ K^T # (N × N) attention scores
2. S_scaled = S / sqrt(d_k) # Scale
3. P = softmax(S_scaled, dim=-1) # (N × N) attention weights
4. O = P @ V # (N × d_v) output

Memory usage:
 - Q, K, V: O(N × d)
 - S: O(N²) ← HUGE! This is the problem!
 - P: O(N²)
 - O: O(N × d)

For N=4096 sequence, d=64:
 - S matrix: 4096² × 4 bytes = 67MB
 - P matrix: 4096² × 4 bytes = 67MB
 - Total for just attention: 134MB!

For N=16384 (longer sequences):
 - S matrix: 16384² × 4 bytes = 1GB!
 - This alone can exceed GPU memory!

Where Time is Spent¶

Attention computation breakdown:

For N=4096, d=64:

Operation FLOPs Memory(GB) Time Bottleneck
─────────────────────────────────────────────────────────────────
Q @ K^T 2N²d 0.03 Compute Memory
Softmax N² 0.13 Memory I/O
P @ V 2N²d 0.03 Compute Memory
Total 4N²d 0.20 Variable I/O-bound!

Key insight:
 - FLOPs are reasonable (matrix multiply)
 - But moving data between GPU memory and cache is slow
 - Softmax requires reading ENTIRE (N × N) matrix
 - This is I/O-bound, not compute-bound!

GPU memory hierarchy:
Registers (32 KB, fast)
 - Can keep small blocks
L2 Cache (40 MB, medium)
 - Can fit maybe 512×512 values
HBM (40 GB, slow)
 - Where S and P matrices live
 - Very slow to access!

Problem: S matrix (N² elements) doesn't fit in cache!
→ Constant back-and-forth to HBM
→ Memory becomes bottleneck

Roofline Analysis¶

GPU A100 Characteristics:
 - Peak compute: 312 TFLOPS
 - Memory bandwidth: 2 TB/s
 - Compute-to-bandwidth ratio: 156 FLOP/byte

For standard attention:
Bytes moved: O(N²)
Computation: O(N² × d)
Compute intensity: (N² × d) / N² = d = 64

But A100 can sustain ~312 TB/s * (156 FLOP/byte) / d = 
 312 TB/s * 156 / 64 = 760 TFLOPS (if we had good I/O)

In reality:
 - Need to write/read S, P from HBM
 - Massive I/O overhead
 - Actual throughput: ~50-100 TFLOPS
 - Terrible! We're wasting 60-80% of GPU potential!

Flash Attention:
 - Blocks computation to fit in cache
 - Reduces I/O from O(N²) to O(N)
 - Actual throughput: ~200 TFLOPS
 - 2-4x improvement!

-

Flash Attention: The Solution¶

Core Idea: Tiling and Recomputation¶

Key insight from operating systems:
"Memory hierarchies are ubiquitous - design around them!"

Standard approach wastes cache:
 - Load K, V to process all Q
 - Compute full S matrix (doesn't fit in cache!)
 - Compute full P matrix (doesn't fit in cache!)
 - Lots of I/O

Flash Attention approach:
Divide computation into blocks ("tiles") that fit in cache!

1. Process Q in tiles (e.g., 64 rows at a time)
2. For each Q tile, iterate through K, V tiles
3. Accumulate results incrementally
4. Keep intermediate results in fast cache (SRAM)
5. Write only final results to slow HBM

Benefits:
 - Fit intermediate S, P in fast cache
 - Reduce I/O to HBM
 - Trade: More compute for less I/O (good trade!)

Visual Algorithm¶

Standard Attention (wasteful I/O):
- ┌─────────────────────────────────┐
 - HBM (GPU Memory - SLOW) │
 - ┌─────────────────────────────┐ │
 - Q (all N tokens) │ │ Read once
 - K (all N tokens) │ │ Read many times
 - V (all N tokens) │ │
 - S[N × N] (all scores) │ │ ← Write, Read, Write (wasteful!)
 - P[N × N] (all weights) │ │ ← Write, Read (wasteful!)
 - O[N × d] (output) │ │ Write once
 - ┘ │
 - ↑ ↑ ↓ │
 - ┘
 I/O-bottleneck!

Flash Attention (efficient I/O):
- ┌─────────────────────────────────┐
 - HBM (GPU Memory) │
 - ┌──────────┬──────────┐ │
 - Q-block1 │ Q-block2 │... │ Read once
 - K (all) │ │ │
 - V (all) │ │ │
 - O-block1 │ O-block2 │... │ Write once
 - ┴──────────┘ │
 - ↓ ↓ │
 - ┘
 SRAM (GPU Cache - FAST)
- ┌─────────────────────────┐
 - Q-block, K-block, V-blk │ Compute here!
 - Intermediate S, P │ (fits in cache)
 - Incremental accum │
 - ┘

-

I/O Complexity Analysis¶

Standard Attention I/O¶

Work analysis (FLOPs):
S = Q @ K^T: O(N² × d) FLOPs
softmax(S): O(N²) FLOPs
P = P @ V: O(N² × d) FLOPs
Total: O(N² × d) FLOPs

Memory I/O analysis:
Q ∈ HBM → registers: O(N × d) I/O
K ∈ HBM → registers: O(N × d) I/O (for each of N blocks)
V ∈ HBM → registers: O(N × d) I/O (for each of N blocks)
S ∈ HBM: O(N²) I/O (write and read!)
P ∈ HBM: O(N²) I/O (write and read!)
O ∈ HBM: O(N × d) I/O (write once)

Total I/O: O(N²) ← HUGE!

I/O to Compute Ratio:
Bytes/FLOP = O(N²) / O(N² × d) = 1/d ≈ 1

This means for every 64 FLOPs, we move 1 byte.
GPU memory bandwidth: 2TB/s
Effective throughput: 2TB/s × 64 ops/byte = 128 TFLOPS
But peak is 312 TFLOPS → 40% utilization (bad!)

Flash Attention I/O¶

Work analysis (FLOPs):
Still O(N² × d) FLOPs (exact same computation!)

Memory I/O analysis (with M = cache size in elements):
Q ∈ HBM → SRAM: O(N × d) I/O
K ∈ HBM → SRAM: O(N × d) I/O
V ∈ HBM → SRAM: O(N × d) I/O
O ∈ HBM: O(N × d) I/O (write once)

No need to write/read S, P (keep in SRAM)!

Total I/O: O(N × d) ← HUGE IMPROVEMENT!

I/O to Compute Ratio:
Bytes/FLOP = O(N × d) / O(N² × d) = 1/N ≈ 1/4096 (much better!)

Effective throughput: 2TB/s × 4096 ops/byte = 8.2 PETA-FLOPS!
But peak is 312 TFLOPS → this analysis is simplified
In practice: ~2-3x speedup (I/O reduced from N² to N terms)

Comparison Table¶

Complexity Metric Standard Flash-Attn Improvement
──────────────────────────────────────────────────────────────
FLOPs O(N²d) O(N²d) Same
Total I/O O(N²) O(Nd) O(N) factor!
Memory peak O(N²) O(M) Huge!
For N=4096:
 Standard I/O 67 MB 0.25 MB 267x less!
 Standard memory peak 67 MB varies
 Flash peak (M=4MB) 4 MB constant!

M = SRAM size (typically 2-4 MB on A100)

-

Flash Attention Algorithm in Detail¶

High-Level Algorithm¶

Algorithm: FlashAttention

Input:
 Q: (N, d) query matrix
 K: (N, d) key matrix
 V: (N, d) value matrix

Config:
 B_r: block size for rows (e.g., 64)
 B_c: block size for columns (e.g., 64)

def flash_attention(Q, K, V, B_r, B_c):
 N = Q.shape[0]
 O = zeros((N, d)) # Output
 l = zeros(N) # Softmax denominator
 m = -inf * ones(N) # Softmax max (for numerical stability)

 # Main loop: process Q in blocks
 for i in range(0, N, B_r):
 i_end = min(i + B_r, N)
 Q_block = Q[i:i_end] # Load Q block to SRAM

 # Sub-loop: compute attention with all K, V blocks
 for j in range(0, N, B_c):
 j_end = min(j + B_c, N)
 K_block = K[j:j_end] # Load K block to SRAM
 V_block = V[j:j_end] # Load V block to SRAM

 # Compute attention for this block pair (stays in SRAM)
 S_block = Q_block @ K_block.T # (B_r, B_c)

 # Numerical stability: track maximum
 m_new = max(m[i:i_end], max_row(S_block))

 # Softmax trick: compute in log-space
 P_block = exp(S_block - m_new)
 l_new = exp(m[i:i_end] - m_new) * l[i:i_end] + sum_row(P_block)

 # Accumulate output
 O[i:i_end] += (P_block @ V_block)

 # Update tracking variables
 m[i:i_end] = m_new
 l[i:i_end] = l_new

 # Normalize output
 O = O / l # Divide by softmax denominator

 return O

Key Technique: Softmax with Reductions¶

Challenge: Softmax is not associative (can't compute in blocks)

Standard softmax:
P[i] = exp(S[i]) / sum(exp(S))

Problem: Don't know the denominator until we've seen all S!

Flash Attention solution: Online softmax (numerically stable)

For a single block:
1. Find max: m = max(S)
2. Compute: exp_S = exp(S - m) ← Subtract max for stability
3. Sum: l = sum(exp_S)
4. Result: P = exp_S / l

For multiple blocks (online update):
Block 1: m_1, l_1
Block 2: m_2, l_2

Updated max: m_new = max(m_1, m_2)
Updated sum: l_new = exp(m_1 - m_new) * l_1 + exp(m_2 - m_new) * l_2

This is EXACT softmax, just computed incrementally!

Accumulated output:
O_new = (exp(m_1 - m_new) * O_1 + exp(m_2 - m_new) * (P_2 @ V))

Memory Layout¶

SRAM allocation (Typical A100: 4 MB = 4M bytes):

For 64-token blocks with 64-dim:
 - Q_block: 64 × 64 × 2 bytes = 8 KB
 - K_block: 64 × 64 × 2 bytes = 8 KB
 - V_block: 64 × 64 × 2 bytes = 8 KB
 - S_block: 64 × 64 × 4 bytes = 16 KB (larger, float32)
 - P_block: 64 × 64 × 2 bytes = 8 KB
 - O_block: 64 × 64 × 2 bytes = 8 KB
 - Scalars (m, l): ~1 KB
 - Total: ~57 KB (well within 4 MB SRAM!)

This is why tiling works!
All intermediate values fit comfortably in SRAM.

-

Flash Attention vs Standard: Step-by-Step Example¶

Concrete Example: N=256, d=64¶

Standard Attention:

Step 1: Q @ K^T
 - Load Q from HBM: 256 × 64 × 2 = 32 KB (fast, once)
 - Load K from HBM: 256 × 64 × 2 = 32 KB
 - Compute: S = Q @ K^T (256 × 256 × 64 = 4M FLOPs)
 - Write S to HBM: 256 × 256 × 4 = 256 KB (slow!)
 - Data movement: 32 + 32 + 256 = 320 KB

Step 2: Softmax(S)
 - Load S from HBM: 256 KB (slow)
 - Compute softmax: 256 × 256 = 65K FLOPs
 - Write P to HBM: 256 KB (slow)
 - Data movement: 256 + 256 = 512 KB

Step 3: P @ V
 - Load P from HBM: 256 KB (slow)
 - Load V from HBM: 32 KB
 - Compute: O = P @ V (256 × 256 × 64 = 4M FLOPs)
 - Write O to HBM: 32 KB
 - Data movement: 256 + 32 + 32 = 320 KB

Total HBM I/O: 32 + 32 + 256 + 256 + 256 + 32 + 32 = 896 KB

Flash Attention (block_size = 64):

Outer loop: i = 0, 64, 128, 192 (4 blocks)
 Inner loop for each i: j = 0, 64, 128, 192 (4 blocks)
 Load Q_i (64 × 64 × 2 = 8 KB) - once per i
 Load K_j, V_j (16 KB) - 4 times per i
 Compute S, P in SRAM (no I/O)
 Accumulate O (no write until end)

Total HBM I/O:
 - Load Q: 4 × 8 KB = 32 KB (same as standard)
 - Load K: 4 × 4 × 8 KB = 128 KB (vs 32 KB standard - wait, more?)
 - Load V: 4 × 4 × 8 KB = 128 KB
 - Write O: 4 × 8 KB = 32 KB (same as standard)
 - Total: 320 KB (less than standard's 896 KB!)

Savings: (896 - 320) / 896 = 64% I/O reduction!

But actual speedup is 2-3x because:
1. S, P no longer ping-pong to/from HBM (huge savings!)
2. Better cache locality
3. Overlapping computation and I/O

Performance Improvements¶

Benchmark Results¶

Measurement: Attention latency (lower = better)

Model: BERT-base (sequence length 512)
Hardware: A100 GPU

Method Time (ms) Speedup Memory (GB)
─────────────────────────────────────────────────────
Standard Attention 4.2 1x 0.35
PyTorch Optimized 3.8 1.1x 0.30
Flash Attention v1 2.0 2.1x 0.05
Flash Attention v2 1.5 2.8x 0.05

Measurement: Full transformer layer

Model: GPT-2 (1.5B params)
Sequence length: 1024
Hardware: A100

Layer Component Standard Flash-v2 Speedup
─────────────────────────────────────────────────────
Attention 12.5 ms 4.3 ms 2.9x
Feedforward 15.2 ms 15.1 ms 1.0x
Total layer 27.7 ms 19.4 ms 1.4x

Memory Usage Reduction¶

Sequence length: 4096
Hidden dimension: 768
Batch size: 1

 Standard Flash-Attn Reduction
─────────────────────────────────────────────────────
Attention scores 67 MB 0 MB 100%
Attention weights 67 MB 0 MB 100%
Activations saved 320 MB 200 MB 37%
KV cache 24 MB 24 MB 0%

Total peak memory 478 MB 224 MB 53%

This means:
- Can fit 2x larger batch on same GPU
- Or 2x longer sequence

Flash Attention Versions¶

Flash Attention v1 (Original, 2022)¶

Features:
 - IO-aware attention
 - Exact computation (same output as standard)
 - 2x speedup, 60% memory reduction
 - For forward pass

Limitations:
 - Backward pass not optimized
 - Limited to specific block sizes
 - Some numerical precision issues

Code:
```python
import torch
from flash_attn import flash_attn_func

# Flash Attention v1
output = flash_attn_func(q, k, v)

### Flash Attention v2 (2023, Improved)

Major improvements over v1:

Backward pass also optimized
2.8x speedup (vs 2x in v1)
Supports variable sequence lengths
Better numerical stability
Supports head_dim up to 256

Typical speedup: For sequence_length = 1024, head_dim = 64:

v1: 2.0x faster
v2: 2.8x faster

Code:

from flash_attn import flash_attn_func

# Flash Attention v2 (automatic)
output = flash_attn_func(q, k, v)

New features:

flash_attn_kvpacked_func: KV packed format
flash_attn_qkvpacked_func: All packed
Support for variable sequence lengths in batch

### Flash Attention v3 (2024, In Progress)

Planned improvements:

Support for FP8 quantization
Grouped Query Attention (GQA) optimization
Multi-GPU coordination
Speculative decoding integration

Expected:

3-4x speedup
Better for quantized models
Seamless integration with vLLM

-

## Implementation & Integration

### PyTorch Integration

```python
# Option 1
import torch
from flash_attn import flash_attn_func

Q = torch.randn(batch, seq_len, num_heads, head_dim, device='cuda')
K = torch.randn(batch, seq_len, num_heads, head_dim, device='cuda')
V = torch.randn(batch, seq_len, num_heads, head_dim, device='cuda')

# Flash attention (with causal mask for autoregressive)
output = flash_attn_func(
 Q, K, V,
 causal=True, # For training/generation
 dropout_p=0.1,
 softmax_scale=None # Auto: 1/sqrt(head_dim)
)

# Option 2
output = torch.nn.functional.scaled_dot_product_attention(
 Q, K, V,
 is_causal=True,
 scale=1.0 / (head_dim ** 0.5)
)
# PyTorch automatically uses Flash Attention if available!

Transformer Integration¶

class FlashAttention(nn.Module):
 def __init__(self, dim, num_heads=8, dropout=0.0):
 super().__init__()
 self.num_heads = num_heads
 self.head_dim = dim // num_heads
 self.dropout = dropout

 self.qkv = nn.Linear(dim, 3 * dim)
 self.out = nn.Linear(dim, dim)

 def forward(self, x, causal=False):
 B, N, C = x.shape

 # Project to Q, K, V
 qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
 q, k, v = qkv.permute(2, 0, 3, 1, 4) # (3, B, H, N, D)

 # Flash Attention
 from flash_attn import flash_attn_func
 out = flash_attn_func(
 q, k, v,
 dropout_p=self.dropout,
 causal=causal
)

 # Merge heads
 out = out.reshape(B, N, C)
 out = self.out(out)

 return out

HuggingFace Integration¶

# Most HuggingFace models automatically use Flash Attention!

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
 "meta-llama/Llama-2-7b",
 attn_implementation="flash_attention_2", # Enable Flash Attention
 torch_dtype="float16"
)

# Or enable globally
import torch
torch.backends.cuda.enable_flash_sdp(True)

vLLM Integration¶

# vLLM automatically uses Flash Attention!

from vllm import LLM, SamplingParams

llm = LLM(
 model="meta-llama/Llama-2-7b-hf",
 # Flash Attention automatically enabled
)

outputs = llm.generate(prompts, sampling_params)

-

Numerical Considerations¶

Numerical Stability¶

Challenge: Softmax can overflow/underflow

exp(1000) → ∞ (overflow)
exp(-1000) → 0 (underflow)

Standard softmax (numerically unstable):
P[i] = exp(S[i]) / sum(exp(S))

Flash Attention (numerically stable):
1. m = max(S) # Find maximum
2. P[i] = exp(S[i] - m) / sum(exp(S[j] - m)) # Subtract max

Example:
S = [1000, 1001, 999.9]

Naive: exp(1000) = ∞ (overflow!)

Stable:
m = 1001
P = [exp(-1), exp(0), exp(-1.1)] / sum(...)
 = [0.368, 1.0, 0.333] / 1.70
 = [0.216, 0.588, 0.196] (correct!)

Gradient Stability¶

Flash Attention v1: Gradient computation could be unstable

Flash Attention v2: Improved gradient stability
 - Better numerical behavior in backward pass
 - More accurate gradients
 - Safer for training

For most use cases: v2 is now standard

-

Limitations and Trade-offs¶

What Flash Attention Does NOT Do¶

Doesn't reduce algorithmic complexity
 - Still O(N²d) FLOPs
 - Just reduces I/O overhead

Doesn't change model accuracy
 - Exact computation (bit-identical to standard)
 - No approximation introduced

Doesn't help with very long sequences indefinitely
 - Still hits memory limits eventually
 - But pushes them much further

Doesn't work on CPU
 - Requires GPU with good memory bandwidth
 - No SRAM on CPU to exploit

Hardware Requirements¶

Flash Attention works best on:
A100 / H100 (excellent)
RTX 4090 (very good)
RTX 3090 (good)
Any modern GPU with good bandwidth

Works less efficiently on:
 Older GPUs (< 2020)
 Mobile GPUs
 Some edge cases

Requires:
 - CUDA (NVIDIA GPUs)
 - Recent GPU compute capability (SM 75+)
 - Properly compiled kernels

Practical Comparison: Real Examples¶

Example 1: Training a 3B Model¶

Setup:
 - Model: LLaMA 3B
 - Sequence length: 2048
 - Batch size: 8
 - Hardware: RTX 4090
 - Task: Fine-tuning on custom data

Without Flash Attention:
 - Memory peak: 20 GB
 - Training step: 2.5 seconds
 - Can't fit (RTX 4090 has 24GB)
 - Need gradient checkpointing

With Flash Attention:
 - Memory peak: 16 GB
 - Training step: 1.8 seconds
 - Fits comfortably!
 - Can increase batch size

Benefit: 40% memory reduction + 28% faster training!

Example 2: Inference Long Document QA¶

Setup:
 - Model: LLaMA 7B
 - Document: 8192 tokens
 - Question: 100 tokens
 - Hardware: A100 40GB

Without Flash Attention:
 - Memory needed: 14GB model + 40GB attention = 54GB 
 - Can't fit on single A100!
 - Need multi-GPU or truncate

With Flash Attention:
 - Memory needed: 14GB model + 2GB attention = 16GB 
 - Fits with room to spare!
 - Inference: 8 seconds for 8192-token context
 - Can process full document!

Benefit: Makes long-context inference feasible!

Example 3: Real-Time Streaming¶

Setup:
 - Model: Mistral 7B
 - User: Chat application
 - Sequence so far: 512 tokens
 - Hardware: RTX 3090 (24GB)

Without Flash Attention:
 - Attention memory: (512)² × 2 bytes = 512 KB (small)
 - But latency per token: 150 ms (slow for streaming)
 - User perceives lag

With Flash Attention:
 - Same memory: 512 KB
 - Latency per token: 50 ms (much better!)
 - Feels responsive!
 - Better user experience

Benefit: 3x faster for interactive applications!

Best Practices¶

Do's¶

Always use Flash Attention v2 if available (it's usually better)
Enable it in HuggingFace models (attn_implementation="flash_attention_2")
For training: Use with gradient checkpointing for even more memory savings
For inference: Pair with KV cache quantization for further speedup
Test on target hardware (sometimes different GPUs have different speedups)
Monitor accuracy during training (usually identical but good to verify)
Use for long contexts (Flash Attention's strength)
Combine with other optimizations (quantization, batching, etc.)

Don'ts¶

Assume exact bit-identical outputs (might have minor numerical differences)
Use for very short sequences < 256 tokens (overhead might not be worth it)
Expect speedup on CPU (use GPU)
Mix different precision types without care
Forget to install correct version (CUDA compatibility matters)
Assume works on all hardware (test first)
Disable if model is working (don't fix what isn't broken)
Expect magic accuracy improvements (it's for speed, not accuracy)

-

Installation and Setup¶

Install Flash Attention v2¶

# For CUDA 11.8+
pip install flash-attn

# For specific CUDA version
pip install flash-attn --no-build-isolation

# From source (if prebuilt doesn't work)
git clone https://github.com/Dao-AILab/flash-attention
cd flash-attention
python setup.py install

# Verify installation
python -c "from flash_attn import flash_attn_func; print('Flash Attention installed')"

Enable in Your Project¶

# Method 1
from flash_attn import flash_attn_func

output = flash_attn_func(q, k, v, causal=True)

# Method 2
import torch
output = torch.nn.functional.scaled_dot_product_attention(q, k, v)
# Automatically uses Flash Attention if available!

# Method 3
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
 model_name,
 attn_implementation="flash_attention_2"
)

# Method 4
from vllm import LLM
llm = LLM(model_name) # Automatically uses Flash Attention

-

Key Takeaways¶

Flash Attention reduces I/O complexity from O(N²) to O(N) 2.8x faster attention computation (Flash Attention v2) 50-60% memory reduction for attention Enables long-context processing (8K+ tokens) Exact computation (no approximation, no accuracy loss) Becoming industry standard (PyTorch, HF, vLLM)

Performance Summary¶

Speedup Factors (Flash Attention v2 vs Standard):

Sequence Length Speedup Memory Reduction
─────────────────────────────────────────────────
256 1.5x 30%
512 2.0x 45%
1024 2.5x 50%
2048 2.8x 55%
4096+ 2.8x 60%

Longer sequences see more benefit!

-