FlashAttention — Paper Breakdown

TL;DR: FlashAttention (Dao et al., ICML 2022) is an exact attention algorithm that avoids materializing the N^2 attention matrix by tiling computation to fit in GPU SRAM, using online softmax for incremental normalization, and recomputing attention scores during backward pass. Achieves 2-4× speedup vs standard attention, enables 16K-32K context on same hardware.

---

Paper Metadata

Field	Value
Paper ID	arxiv-2205.14135
Title	FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Authors	Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, Christopher Ré
Venue/Year	ICML 2022 (Oral)
ArXiv/DOI	https://arxiv.org/abs/2205.14135
Code	https://github.com/Dao-AILab/flash-attention

---

Problem Statement

Standard attention is memory-bound, not compute-bound.

On A100 (1.5 TB/s HBM, 19 TB/s SRAM):

• Standard attention: Arithmetic intensity ≈ 0.5 FLOP/byte → memory-bound
• Matrix multiply (GEMM): Arithmetic intensity ≈ 50 FLOP/byte → compute-bound

Standard attention HBM traffic (dominant cost):

Operation	HBM Reads	HBM Writes
Load Q, K, V	3Nd	—
Write S = QK^T	—	N^2
Read S, write P	N^2	N^2
Read P, V, write O	N^2 + Nd	Nd
Total	mathbf{2N^2 + O(Nd)}	mathbf{2N^2 + O(Nd)}

For N=4096, d=64: N^2 matrix alone = 128 MB FP16 — exceeds SRAM (164 KB/SM) by 1000×.

Key insight: Attention is IO-bound. Reduce HBM traffic by fusing operations and keeping intermediates in SRAM.

---

Method Overview

High-Level Idea (In Your Own Words)

Instead of computing the full N × N attention matrix in slow HBM, tile the computation so each block fits in fast SRAM (shared memory). Compute attention for one Q-block against all K,V blocks sequentially, accumulating output incrementally. Use online softmax to normalize without seeing full row. In backward pass, recompute S=QK^T from saved Q,K instead of saving S,P.

Architecture / Pipeline

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

graph TD
    Q[Q: N×d] --> Block[Q Block B_r×d]
    K[K: N×d] --> KVBlock[K,V Blocks B_c×d]
    V[V: N×d] --> KVBlock

Block --> SRAM["SRAM (Shared Memory)"]
KVBlock --> SRAM

SRAM --> TileAttn[Tile Attention: S_ij = Q_i K_j^T]
TileAttn --> OnlineSoftmax[Online Softmax<br/>Incremental m, l]
OnlineSoftmax --> AccumOut[Accumulate Output O_i]
AccumOut --> HBM[Write O to HBM]

subgraph Backward
    QK["Saved: Q, K, V, O, m, l"]
    QK --> RecomputeS[Recompute S = QK^T]
    RecomputeS --> Grad[Compute Gradients]
end

Key Components

Component	Standard Attention	FlashAttention
Attention matrix	Stored in HBM (N^2)	Never materialized
Softmax	After full S computed	Online (incremental)
SRAM usage	Minimal	Blocks of Q,K,V
Backward	Load S,P from HBM	Recompute S from Q,K
HBM traffic	O(N^2)	O(N^2/B + Nd)

---

Technical Deep Dive

Online Softmax (Critical for Tiling)

Standard softmax needs full row:

softmax(x)_i = (e^{x_i})/(sum_j e^{x_j)}

FlashAttention computes incrementally as blocks arrive:

Per row i, maintain:

• m_i = max_j x_{ij} (running max)
• ell_i = sum_j e^{x_{ij} - m_i} (running sum)

When new block arrives:

m_i^{new} = max(m_i^{old}, rowmax(x_i^{new}))

ell_i^{new} = e^{m_i^{old} - m_i^{new}} · ell_i^{old} + rowsum(e^{x_i^{new} - m_i^{new}})

Final: P_{ij} = e^{S_{ij} - m_i^{final}} / ell_i^{final}

Forward Algorithm (Simplified)

def flash_attention_forward(Q, K, V, B_r=128, B_c=128):
    # Q: (N, d), K,V: (N, d)
    O = torch.zeros_like(Q)
    l = torch.zeros(N)      # row sums
    m = torch.full((N,), -float('inf'))  # row maxes
    
    for j in range(0, N, B_c):          # Iterate K,V blocks
        K_j = K[j:j+B_c]; V_j = V[j:j+B_c]
        K_j, V_j = K_j.to_sram(), V_j.to_sram()
        
        for i in range(0, N, B_r):      # Iterate Q blocks
            Q_i = Q[i:i+B_r].to_sram()
            
            # Compute tile S_ij = Q_i K_j^T / sqrt(d)  [B_r, B_c]
            S_ij = (Q_i @ K_j.T) * scale
            
            # Online softmax update
            m_new = max(m[i:i+B_r], S_ij.max(dim=1))
            P_ij = torch.exp(S_ij - m_new[:, None])
            l_new = torch.exp(m[i:i+B_r] - m_new) * l[i:i+B_r] + P_ij.sum(dim=1)
            
            # Accumulate output
            O[i:i+B_r] = (l[i:i+B_r] / l_new)[:, None] * O[i:i+B_r] + P_ij @ V_j
            
            m[i:i+B_r] = m_new
            l[i:i+B_r] = l_new
    
    O = O / l[:, None]
    return O, m, l  # Save m, l for backward

Backward Pass: Recomputation

Saved in forward: Q, K, V, O, m, l (all O(Nd)) NOT saved: S, P (O(N^2))

Backward recomputes:

1. Load Q_i, K_j, V_j, O_i, m_i, l_i
2. Recompute S_{ij} = Q_i K_j^T in SRAM
3. Recompute P_{ij} via online softmax (using saved m_i, l_i)
4. Compute gradients dQ_i, dK_j, dV_j
5. Accumulate across blocks

Compute overhead: ~30% more FLOPs, but saves 10-20× memory — net win on HBM-bound GPUs.

---

Results Summary

Main Results (A100 40GB)

Metric	Standard	FlashAttention	Improvement
Speed (BERT-base, 512)	1.0×	1.4×	1.4× faster
Speed (GPT-2, 1K)	1.0×	2.1×	2.1× faster
Speed (4K)	1.0×	3.7×	3.7× faster
Memory (1K)	1.0×	0.1×	10× less
Memory (4K)	1.0×	0.05×	20× less
Max seq len (OOM)	~2K	16K+	8× longer

End-to-End Training

Model	Seq Len	Standard	FlashAttn	Speedup
BERT-base	512	1.0×	1.06×	1.06×
GPT-2 small	1K	1.0×	1.8×	1.8×
GPT-2 medium	2K	1.0×	3.0×	3.0×
GPT-2 large	4K	OOM	Works	∞

Kernel Profiling

• Standard attention: 35% compute, 65% memory stall
• FlashAttention: 85% compute, 15% memory stall
• Arithmetic intensity: 0.5 → 25 FLOP/byte (near GEMM!)

---

Strengths

1. Exact attention — No approximation (unlike sparse/linear attention)
2. Massive memory savings — O(Nd) vs O(N^2)
3. Longer sequences — 8-16× context length on same hardware
4. No accuracy loss — Numerically identical to standard (FP32 accumulation)
5. Composable — Works with any Transformer variant (BERT, GPT, ViT)

Weaknesses / Limitations (Honest Assessment)

1. Forward-backward recompute cost — ~30% extra compute in backward
2. Causal mask only (FA-1) — No arbitrary masks; FA-2 adds padding mask
3. Head dimension constraints — FA-1: 64, 128; FA-2: up to 256 (multiple of 16)
4. No dropout support in FA-1 — FA-2 adds dropout
5. Implementation complexity — Hand-tuned CUDA kernels per architecture
6. Not faster for short sequences — Overhead > benefit for N < 512

Reproducibility / Implementation Notes

• Open source: flash-attn PyPI package, integrated in PyTorch 2.0+ SDPA
• Requirements: Ampere+ GPU (A100, H100, RTX 30/40), CUDA 11.6+
• Key hyperparams: Block sizes B_r=128, B_c=128 (tuned for A100 SRAM)
• Gotchas:
  • Must use torch.compile or explicit flash_attn_func for best perf
  • Mixed precision (FP16/BF16) with FP32 accumulation critical for numerics
  • Sequence length must be multiple of block size (pad if needed)

---

Connections to Your Wiki

Concepts This Paper Introduces / Advances

• flash-attention — IO-aware exact attention algorithm
• attention-mechanism — Practical scaling of attention

Concepts This Paper Builds On

• Standard scaled dot-product attention attention-mechanism
• Online softmax (historical: stream processing literature)
• GPU memory hierarchy (SRAM vs HBM)
• Kernel fusion techniques

Potential Project Ideas

• efficient-vision-transformer — FlashAttention enables high-res ViT
• vision-transformer-dsa-integration — Compare DSA vs FlashAttn for ViT
• lora-qloра — Combine FlashAttention + QLoRA for long-context FT

---

Sources

• arxiv-2205.14135: Primary source — Sections 3 (Algorithm), 4 (Analysis), 5 (Experiments)
• FlashAttention-2: Dao, "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)