Flash Attention

Table of Contents

1. Learning Objectives
2. Prerequisites
3. Notation
4. Core Intuition
5. Memory Hierarchy and the IO Bottleneck
6. Standard Attention Memory Analysis
7. Tiling Strategy
8. Online Softmax
9. Flash Attention Algorithm
10. Complexity Analysis
11. Flash Attention 2 Improvements
12. Worked Examples
13. Connection to the Broader Curriculum
14. Common Pitfalls and Misconceptions
15. Research Perspective
16. Summary of Takeaways
17. Exercises

Learning Objectives

After reading this chapter, you should be able to:

1. Explain the GPU memory hierarchy (SRAM vs. HBM) and IO bottleneck in attention.
2. Prove standard attention requires $O(N^2)$ HBM accesses for sequence length $N$.
3. Describe tiling that keeps blocks of Q, K, V in SRAM.
4. Derive online softmax for computing attention without materializing full $\mathbf{S} \in \mathbb{R}^{N \times N}$.
5. State Flash Attention's $O(N^2 d)$ FLOPs with $O(N)$ memory.

Prerequisites

• Self-Attention — $\text{softmax}(QK^T)V$
• KV Cache — inference context

Notation

• $N, d$ — Sequence length and head dimension
• $\mathbf{Q}, \mathbf{K}, \mathbf{V}$ — Attention factor matrices
• $B_r, B_c$ — Query and key/value tile sizes
• $M_{\mathrm{SRAM}}$ — On-chip SRAM capacity
• $\tilde{m}, \tilde{\ell}$ — Online softmax running max and sum

Core Intuition

Self-Attention is $O(N^2)$ in memory for the attention matrix $\mathbf{S} = QK^T \in \mathbb{R}^{N \times N}$. For $N = 64$K, storing $\mathbf{S}$ in FP16 requires 8 GB — prohibitive.

Flash Attention (Dao et al., 2022) computes exact attention (not approximation) with $O(N)$ memory by exploiting GPU memory hierarchy: fast SRAM (on-chip, ~20 MB) vs. slow HBM (off-chip, ~80 GB). The algorithm tiles computation to minimize HBM reads/writes.

Series context. Volume IV, Chapter 17 (Memory Optimization).

Memory Hierarchy and the IO Bottleneck

Definition 1 (Memory Levels).

• HBM (High Bandwidth Memory): Large ($\sim$80 GB), high latency, main GPU memory
• SRAM (Shared Memory): Small ($\sim$20–100 KB per SM), fast, on-chip

Proposition 1. Matrix operations are often memory-bound, not compute-bound: reading/writing HBM dominates wall-clock time.

Definition 2 (IO Complexity). Count HBM accesses — the relevant cost metric for attention optimization.

Standard Attention Memory Analysis

Standard Self-Attention:

1. Compute $\mathbf{S} = QK^T \in \mathbb{R}^{N \times N}$ — write $N^2$ to HBM
2. Softmax on $\mathbf{S}$ — read/write $N^2$
3. Compute $\mathbf{O} = \text{softmax}(\mathbf{S}) V$ — read $N^2$, write $N \times d$

Theorem 1. Naive attention materialization requires $O(N^2)$ HBM memory for $\mathbf{S}$ and $O(N^2)$ IO per forward pass.

For long sequences, memory — not FLOPs — is the bottleneck.

Tiling Strategy

Definition 3 (Block Tiling). Partition $Q, K, V$ into blocks of size $B_r \times d$ and $B_c \times d$ that fit in SRAM.

Idea. Compute attention block by block without ever forming full $\mathbf{S}$:

For each block of queries $Q_i$, iterate over key/value blocks $K_j, V_j$, accumulating partial attention output.

Online Softmax

Problem. Softmax requires global max and sum over entire row — seemingly needs full $\mathbf{S}$ row.

Theorem 2 (Online Softmax). Softmax can be computed incrementally over blocks using running max $m$ and running sum $l$:

Given previous blocks with max $m_{\text{old}}$, sum $l_{\text{old}}$, and output accumulator $O_{\text{old}}$, for new block with local max $m_{\text{new}}$, local exp-sum $l_{\text{new}}$, local contribution $O_{\text{new}}$:

Proof. Rescale previous partial softmax when discovering new maximum. $\blacksquare$

This enables exact softmax over tiled blocks without storing full row.

Flash Attention Algorithm

Algorithm (FlashAttention Forward).

Input: $Q, K, V \in \mathbb{R}^{N \times d}$, block sizes $B_r, B_c$

For each query block $i$:

1. Load $Q_i$ to SRAM
2. Initialize $O_i = 0$, $l_i = 0$, $m_i = -\infty$
3. For each key/value block $j$:
   • Load $K_j, V_j$ to SRAM
   • Compute $S_{ij} = Q_i K_j^T$ (in SRAM)
   • Update $m_i, l_i, O_i$ via online softmax (1)–(3)
4. Write $O_i$ to HBM

Key property: $\mathbf{S}_{ij}$ blocks never written to HBM — only final $\mathbf{O}$.

Complexity Analysis

Theorem 3 (Flash Attention Complexity).

• FLOPs: $O(N^2 d)$ — same as standard attention (exact computation)
• HBM Memory: $O(N)$ — stores only $Q, K, V, O$, not $\mathbf{S}$
• IO: $O(N^2 d^2 / M)$ where $M$ is SRAM size — sub-quadratic IO when tiled optimally

Corollary 1. Enables sequence lengths impossible with materialized attention — critical for long-context LLMs with KV Cache.

Flash Attention 2 Improvements

FlashAttention-2 (Dao, 2023):

• Better parallelization across sequence blocks
• Reduced non-matmul FLOPs
• ~2× speedup over FlashAttention-1
• Supports variable sequence lengths, RoPE, GQA

Worked Examples

Example 1: Memory Savings

$N = 32$K, FP16: naive $\mathbf{S}$ = 2 GB; Flash Attention: $O(Nd)$ only.

Example 2: Online Softmax

Row $[1, 5, 3]$ processed in blocks $[1, 5]$ then $[3]$ — same result as standard softmax.

Connection to the Broader Curriculum

• Self-Attention — exact algorithm preserved
• KV Cache — composable
• Data Parallelism — orthogonal axis
• Quantization — complementary optimization

Common Pitfalls and Misconceptions

Pitfall 1: Flash Attention is approximate (false — exact).

Pitfall 2: Reduces FLOPs (false — reduces memory IO).

Pitfall 3: Backward pass also requires recomputation (trade memory for compute).

Research Perspective

FlashAttention (Dao et al., 2022). FlashAttention-2/3. Ring attention for distributed long context. PagedAttention (vLLM).

Summary of Takeaways

• Memory — $O(N^2)$ — $O(N)$
• FLOPs — $O(N^2 d)$ — $O(N^2 d)$
• Exact — Yes — Yes
• Key idea — Materialize $\mathbf{S}$ — Tiling + online softmax

Next: Quantization Fundamentals

Exercises

Exercise 1. Compute HBM for $\mathbf{S}$ at $N = 64$K, FP16.

Exercise 2. Derive online softmax update (1)–(3).

Exercise 3. Why is attention memory-bound?

Exercise 4. IO analysis for tiled matmul.

Exercise 5. Backward pass recomputation tradeoff.

Exercise 6. Compose with KV Cache.

Exercise 7. GQA reduces which tensor sizes?

Exercise 8. Compare to sparse attention approximations.