FlashAttention

Summary: FlashAttention (Dao et al., ICML 2022) is an IO-aware exact attention algorithm that avoids materializing the N^2 attention matrix by tiling computation to fit in GPU SRAM and using recomputation during backward pass. Achieves 2-4× speedup vs standard attention, enables longer sequences (16K-32K) on same memory.

Visual Overview: Tiled Attention in SRAM

Diagram: (Mermaid diagram - view source for diagram code)

flowchart LR
    subgraph HBM["HBM (Global Memory) - Slow"]
        Q[Q Matrix: N×d]:::matrix
        K[K Matrix: N×d]:::matrix
        V[V Matrix: N×d]:::matrix
        style HBM fill:#f9f9f9,stroke:#ccc
    end

subgraph SRAM["SRAM (Shared Memory) - Fast"]
    direction TB
    QBLOCK[Q Block: B_r×d]:::matrix
    KBLOCK[K Block: B_c×d]:::matrix
    VBLOCK[V Block: B_c×d]:::matrix
    S_TILE[S Tile: B_r×B_c]:::matrix
    P_TILE[P Tile: B_r×B_c]:::matrix
    O_ACC[O Accumulator: N×d]:::matrix
    M[Running Max: N]:::vector
    L[Running Sum: N]:::vector
    style SRAM fill:#e8f5e9,stroke:#c8e6c9
end

%% Data loading
Q -->|Load block| QBLOCK
K -->|Load block| KBLOCK
V -->|Load block| VBLOCK

%% Computation
QBLOCK -->|QK^T| S_TILE
KBLOCK -->|QK^T| S_TILE
S_TILE -->|Online Softmax| P_TILE
P_TILE -->|PV| O_ACC
VBLOCK -->|PV| O_ACC

%% State updates
M -->|Update| L
L -->|Update| M
P_TILE -->|Row sum/max| M
P_TILE -->|Row sum/max| L

%% Final output
O_ACC -->|Write to HBM| O[O Matrix: N×d]:::matrix

classDef matrix fill:#e3f2fd,stroke:#90caf9;
classDef vector fill:#fff3e0,stroke:#ffcc02;
class QBLOCK,KBLOCK,VBLOCK,S_TILE,P_TILE,O_ACC,O matrix;
class Q,K,V matrix;
class M,L vector;

Core Problem

In plain English: Standard attention computes and stores the full N × N attention matrix in slow HBM (global memory). For N=8192, this matrix alone is 256MB in FP16. FlashAttention fuses operations to keep intermediate results in fast SRAM (shared memory), only writing final output to HBM.

Standard Attention Memory Access

HBM Read: Q (N×d), K (N×d), V (N×d)
HBM Write: S = QK^T (N×N)  ← MAJOR BOTTLENECK
HBM Read: S (N×N)
HBM Write: P = softmax(S) (N×N)
HBM Read: P (N×N)
HBM Write: O = PV (N×d)

Total HBM traffic: O(N^2) — dominated by N^2 matrix

GPU Memory Hierarchy

Memory Latency Bandwidth Size (A100)
HBM (Global) ~400 cycles 1.5 TB/s 40/80 GB
SRAM (Shared) ~4 cycles 19 TB/s 164 KB/SM
Registers ~1 cycle — 64 KB/SM

Key insight: SRAM is 10× faster bandwidth but 1000× smaller. Fit tiles in SRAM!

Algorithm: Tiled Attention

Forward Pass

Split Q, K, V into blocks that fit in SRAM:

  • Block size B_r × d for Q (rows)
  • Block size B_c × d for K, V (cols)
# Pseudo-code for FlashAttention forward
O = zeros(N, d)
l = zeros(N)          # log-sum-exp denominator
m = -inf * ones(N)    # row-wise max for numerical stability

for j in range(0, N, B_c):      # Iterate over K,V blocks
    K_j = K[j:j+B_c]            # Load K block to SRAM
    V_j = V[j:j+B_c]            # Load V block to SRAM
    
    for i in range(0, N, B_r):  # Iterate over Q blocks
        Q_i = Q[i:i+B_r]        # Load Q block to SRAM
        
        # Compute attention for this tile
        S_ij = Q_i @ K_j^T / sqrt(d)      # B_r × B_c in SRAM
        m_new = max(m[i:i+B_r], rowmax(S_ij))
        P_ij = exp(S_ij - m_new[:, None]) # Softmax numerator
        l_new = exp(m[i:i+B_r] - m_new) * l[i:i+B_r] + rowsum(P_ij)
        
        # Update output incrementally (online softmax)
        O[i:i+B_r] = diag(exp(m[i:i+B_r] - m_new)) @ O[i:i+B_r] + P_ij @ V_j
        
        m[i:i+B_r] = m_new
        l[i:i+B_r] = l_new

Key Techniques

1. Online Softmax

  • Compute softmax incrementally without storing full S
  • Maintain per-row max m_i and sum l_i
  • Allows streaming K,V blocks through SRAM

2. Tiling

  • B_r, B_c chosen so Q_i, K_j, V_j, S_{ij}, P_{ij}, O_i fit in SRAM (~164 KB)
  • Typical: B_r = B_c = 128 for d=64 (head dim)

3. Recomputation (Backward Pass)

  • Forward: Don't store S, P (saves O(N^2) memory)
  • Backward: Recompute S, P on-the-fly from Q, K, V, O
  • Trade compute for memory — cheap on GPU (math throughput ≫ memory bandwidth)

Complexity Analysis

Metric Standard Attention FlashAttention
HBM Reads O(N^2) O(N^2 / B)
HBM Writes O(N^2) O(N)
SRAM Usage O(1) O(B^2)
Compute O(N^2 d) O(N^2 d) (same)
Speedup 1× 2-4×

Why faster? Standard attention is memory-bound (arithmetic intensity < GPU balance point). FlashAttention increases arithmetic intensity by reusing SRAM data.

FlashAttention-2 (Dao et al., 2023)

Improvements

  1. Better parallelism: Parallelize over batch/head dims, not just sequence
  2. Reduced SRAM usage: Smaller tiles, better occupancy
  3. Support for variable sequence lengths (padding-free)
  4. Up to 2× faster than FlashAttention-1

Partitioning

Q: (batch, seqlen, nheads, headdim)
→ Split along seqlen dimension
→ Each thread block handles one (batch, head) pair

FlashAttention-3 (2024)

Hopper (H100) Optimizations

  • Tensor Cores: Use MMA (matrix multiply-accumulate) for attention scores
  • Asynchronous pipelining: Overlap compute with TMA (Tensor Memory Accelerator) loads
  • FP8 support: 2× throughput vs FP16
  • Up to 1.5-2× faster than FA-2 on H100

Performance Numbers (A100 40GB)

Sequence Length Standard (TFLOPS) FlashAttn (TFLOPS) Speedup Memory (Std) Memory (Flash)
1K 45 85 1.9× 0.8 GB 0.3 GB
2K 52 110 2.1× 2.5 GB 0.5 GB
4K 55 125 2.3× 9 GB 1 GB
8K 56 135 2.4× 34 GB 2 GB
16K OOM 140 ∞ OOM 4 GB
32K OOM 142 ∞ OOM 8 GB

Usage

PyTorch (via xFormers / FlashAttn lib)

from flash_attn import flash_attn_func

# Input: (batch, seqlen, nheads, headdim)
out = flash_attn_func(q, k, v, dropout_p=0.0, causal=True)

HuggingFace Transformers

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    attn_implementation="flash_attention_2",  # or "sdpa"
    torch_dtype=torch.float16,
)

Requirements

  • GPU: Ampere (A100, RTX 30-series) or newer
  • CUDA 11.6+
  • flash-attn package: pip install flash-attn --no-build-isolation

Limitations

  • Exact attention only — no approximation (unlike sparse/linear attention)
  • Causal mask only (no arbitrary masks in FA-1, supported in FA-2+)
  • Head dimension must be 64 or 128 (FA-2: 16-256, multiple of 8)
  • No dropout during inference (train: supported)
  • Backward pass recompute adds ~30% compute overhead

When to Use

✅ Long sequences (>2K tokens)
✅ Memory-constrained training (consumer GPUs)
✅ Large batch sizes
❌ Short sequences (<512) — overhead not worth it
❌ Non-causal attention with custom masks (limited support)

Related Concepts

Sources

  • arxiv-2205.14135: Sections 3 (Algorithm), 4 (Analysis), 5 (Experiments)
  • FlashAttention-2: Dao et al., "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning" (2023)
  • FlashAttention-3: Dao et al., "FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision" (2024)