Flash Attention

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. The Memory Hierarchy Problem
5. Online Softmax: The Key Insight
6. FlashAttention Algorithm
7. Memory Analysis
8. FlashAttention-2: Optimizations
9. FlashAttention-3: FP8 and Pipelining
10. Backward Pass
11. Common Pitfalls
12. Summary
13. Exercises

Learning Objectives

1. Explain the GPU memory hierarchy (SRAM vs HBM) and its impact on attention.
2. Derive the online softmax algorithm for block-wise computation.
3. Trace through the FlashAttention tiling algorithm step by step.
4. Prove that FlashAttention computes exact attention with $O(T)$ HBM memory.
5. Explain the key improvements in FlashAttention-2 and FlashAttention-3.

Notation

• $\text{HBM}$ — High Bandwidth Memory (GPU main memory, ~80 GB, ~2 TB/s)
• $\text{SRAM}$ — Static RAM (on-chip, ~20 MB, ~19 TB/s)
• $B_r, B_c$ — block sizes for rows (queries) and columns (keys)
• $\ell_i, m_i$ — running softmax statistics (sum and max)

Core Intuition

Standard attention computes the full $T \times T$ attention matrix and stores it in GPU memory (HBM). This is wasteful: we write $T^2$ values to slow HBM only to immediately read them back for the matrix multiply with V. FlashAttention never materializes this matrix — it computes attention in tiles that fit in fast on-chip SRAM, fusing all operations into a single kernel. Same math, different execution order, dramatic speedup.

The Memory Hierarchy Problem

GPU memory hierarchy:

• SRAM (L1/shared memory): ~20 MB, ~19 TB/s bandwidth. Fast but tiny.
• HBM (DRAM): ~80 GB, ~2 TB/s bandwidth. Large but 10x slower.

Standard attention IO cost: For $T$ tokens, $d$-dimensional:

1. Write $\mathbf{S} = \mathbf{QK}^T$ to HBM: $O(T^2)$ writes.
2. Read $\mathbf{S}$ from HBM for softmax: $O(T^2)$ reads.
3. Write $\mathbf{A} = \text{softmax}(\mathbf{S})$ to HBM: $O(T^2)$ writes.
4. Read $\mathbf{A}$ from HBM for $\mathbf{AV}$: $O(T^2)$ reads.

Total HBM accesses: $O(T^2)$ — dominates the runtime for long sequences.

FlashAttention: Reduces HBM accesses to $O(T^2d^2/M)$ where $M$ is SRAM size — much smaller.

Online Softmax: The Key Insight

Problem: Softmax requires the global maximum over all keys to prevent overflow:

If we process keys in blocks, we don't know the global max until we've seen all blocks. How to compute softmax incrementally?

Online softmax algorithm. Maintain running statistics $(m_i, \ell_i)$:

• $m_i$ = max of scores seen so far (through block $i$)
• $\ell_i$ = sum of exponentials seen so far (with current max)

Update rule when incorporating new block with scores $\mathbf{s}_{\text{new}}$:

This rescales the accumulated output whenever the max changes — numerically exact.

FlashAttention Algorithm

Input: $\mathbf{Q}, \mathbf{K}, \mathbf{V} \in \mathbb{R}^{T \times d}$ in HBM. Output: $\mathbf{O} \in \mathbb{R}^{T \times d}$ in HBM.

Procedure:

1. Divide $\mathbf{Q}$ into blocks of $B_r$ rows, $\mathbf{K}, \mathbf{V}$ into blocks of $B_c$ rows.
2. For each Q-block $i$:
   • Initialize $\mathbf{O}_i = \mathbf{0}$, $\ell_i = \mathbf{0}$, $m_i = -\infty$.
   • For each K/V-block $j$:
     • Load $\mathbf{Q}_i, \mathbf{K}_j, \mathbf{V}_j$ from HBM to SRAM.
     • Compute $\mathbf{S}_{ij} = \mathbf{Q}_i\mathbf{K}_j^T / \sqrt{d_k}$ in SRAM.
     • Update online softmax: compute new $m_i, \ell_i$.
     • Update output: $\mathbf{O}_i$ accumulates weighted values.
   • Write final $\mathbf{O}_i$ to HBM.

Block sizes chosen so that $B_r \times d + B_c \times d + B_r \times B_c \leq M$ (fits in SRAM).

Memory Analysis

HBM reads:

• $\mathbf{Q}$: $O(Td)$ (each Q-block read once)
• $\mathbf{K}, \mathbf{V}$: $O(Td \cdot T/B_r)$ = $O(T^2d/B_r)$ (each K/V-block read once per Q-block)

HBM writes: $O(Td)$ (only the output $\mathbf{O}$).

Total HBM accesses: $O(T^2d^2/M)$ where $M$ = SRAM size.

Attention matrix storage: $O(B_r \times B_c)$ in SRAM — never stored in full in HBM. So HBM memory for the attention matrix: $O(T)$ (just the statistics $m, \ell$).

Compared to standard: $O(T^2)$ HBM memory → $O(T)$ HBM memory. For $T=16K$: saves ~1 GB of memory for the attention matrix alone.

FlashAttention-2: Optimizations

Key improvements over FlashAttention-1:

1. Reduced non-matmul FLOPs: Restructure the online softmax update to minimize scalar operations. Move rescaling to the end of the inner loop.

2. Better parallelism: FlashAttention-1 parallelizes over batch and heads. FlashAttention-2 additionally parallelizes over the sequence length dimension (Q-blocks), improving GPU occupancy.

3. Warp-level optimizations: Partition work within a thread block to reduce shared memory reads/writes between warps.

Result: ~2x speedup over FlashAttention-1, reaching 50–73% of theoretical peak FLOPS on A100.

FlashAttention-3: FP8 and Pipelining

Targets Hopper architecture (H100):

1. Asynchronous pipelining: Overlap GEMM computation with softmax and data loading using the Tensor Memory Accelerator (TMA).

2. FP8 quantization within the kernel: Perform Q×K in FP8 for higher throughput (doubling FLOPS vs FP16), then accumulate in FP32 for numerical stability.

3. Block quantization: Quantize K/V blocks to FP8 with per-block scaling factors.

Result: Up to 1.5-2x speedup over FlashAttention-2 on H100 GPUs.

Backward Pass

Challenge: Standard backward pass needs the attention matrix $\mathbf{A}$, which FlashAttention doesn't store.

Solution: Recompute $\mathbf{A}$ block-by-block during backward (same tiling as forward). Store only the softmax statistics $(m_i, \ell_i)$ to enable exact recomputation.

Extra cost: One additional pass over Q, K, V (~33% overhead compared to storing $\mathbf{A}$). But memory savings far outweigh the compute cost.

Gradient computation:

computed block-wise using the stored statistics for exact softmax gradient.

Common Pitfalls

Pitfall 1. Assuming FlashAttention is an approximation. It computes exact standard attention — the same output as the naive algorithm, just with different execution order.

Pitfall 2. Using FlashAttention with very small $T$. For short sequences ($T < 512$), the tiling overhead can make FlashAttention slower than standard fused attention.

Pitfall 3. Expecting FlashAttention to reduce FLOPs. The FLOP count is identical ($O(T^2d)$). The speedup comes entirely from reduced memory IO, not less computation.

Summary

• FlashAttention exploits the GPU memory hierarchy: compute in fast SRAM, minimize slow HBM access.
• Online softmax enables block-wise exact attention computation.
• Memory: $O(T)$ HBM instead of $O(T^2)$; enables 4–16x longer sequences.
• Speed: 2–4x faster than standard attention (IO-bound → compute-bound).
• FlashAttention-2: better parallelism, 2x over v1.
• FlashAttention-3: FP8, H100 pipelining, 1.5–2x over v2.
• Exact computation — no quality trade-off.

Exercises

Exercise 1. For SRAM size $M = 192$ KB, $d=128$, compute the maximum block sizes $B_r = B_c$ that fit (in FP16).

Exercise 2. Derive the total HBM bytes accessed by FlashAttention for $T=8192, d=128, B_r=B_c=128, M=192$ KB.

Exercise 3. Prove that the online softmax (equations 2–4) produces the same result as computing the global softmax.

Exercise 4. Compute the memory savings (in GB) from FlashAttention vs standard attention for $T=32768, H=32, d_k=128$ in FP16.

Exercise 5. Explain why FlashAttention's backward pass stores $O(T)$ values (the statistics) rather than the full $O(T^2)$ attention matrix.