Chapter 15 -- Attention Optimization

FlashAttention

How tiling and kernel fusion eliminate the N x N memory bottleneck, making long-context inference practical.

O(N2)
Standard Memory
O(N)
FlashAttention Memory
2-4x
Wall-clock Speedup

GPU Memory Architecture

The key insight: SRAM is ~1000x faster than HBM but ~1000x smaller. FlashAttention keeps data in SRAM.

SRAM (on-chip)
~20 MB | ~19 TB/s
|
1000x faster
|
HBM (off-chip)
80-192 GB | ~3.35 TB/s

The Problem

Standard attention writes the full N x N attention matrix to HBM. For a 128K sequence:

128K x 128K x 2 bytes = 32 GB
Just for one attention head's S matrix!
FlashAttention: Never writes the full matrix
Tiles fit in 20 MB SRAM. O(N) memory.

Standard vs Flash Attention

Side by side: standard attention writes everything to HBM. FlashAttention tiles the computation.

Standard Attention
1
S = Q @ K^T
Compute NxN scores, write to HBM
|
2
Read S from HBM
Full NxN matrix loaded back
|
3
P = softmax(S)
Write NxN P matrix to HBM
|
4
O = P @ V
Read P from HBM, write O
HBM reads/writes: ~4N2
FlashAttention
1
Load Q, K, V tiles
Small blocks from HBM to SRAM
|
2
S_tile = Q_tile @ K_tile^T
Compute in SRAM -- never to HBM!
|
3
P_tile = softmax(S_tile)
Online softmax in SRAM
|
4
O += P_tile @ V_tile
Accumulate output, write once
HBM reads/writes: ~O(N)

The Tiling Strategy

FlashAttention processes the attention matrix in tiles that fit in SRAM. Click "Animate" to watch.

N x N Attention Matrix

How It Works

The full NxN matrix is never materialized. Instead, we process it tile by tile:

  1. Load a tile of K, V into SRAM
  2. For each Q block, compute local attention
  3. Use online softmax to track running max/sum
  4. Accumulate output incrementally
  5. Move to next tile

The Key Insight

Avoid materializing the full N x N attention matrix in HBM. This is what makes long context possible.

Standard: Full NxN in HBM

Every cell written to HBM
8K seq = 256 MB per head

Flash: Only tiles in SRAM

Only active tile in SRAM
Always fits in ~20 MB SRAM
4N2
Standard HBM Accesses
O(N)
Flash HBM Accesses
~1000x
SRAM vs HBM Speed

Sequence Length Impact

Drag the slider to see how memory usage diverges as sequence length increases.

1K tokens Sequence Length 256K tokens
8K
256 MB
Standard
Attention
2 MB
Flash
Attention

FlashAttention Evolution

Each version brought significant improvements in speed and hardware utilization.

FA-1
Jun 2022
IO-aware tiling
2-4x speedup
over PyTorch
FA-2
Jul 2023
Better parallelism
& work partitioning
2x over FA-1
50-73% FLOPS util
FA-3
Jul 2024
H100 Tensor Cores
FP8 support
1.5-2x over FA-2
75% FLOPS util
FA-4
2025
Blackwell TMA
Persistent kernels
Target: near 100%
FLOPS utilization
Relative Performance (TFLOPS on H100):
150
FA-1
300
FA-2
500
FA-3
700+
FA-4

Key Takeaways

Tiling Is the Core Idea

By processing attention in tiles that fit in SRAM, FlashAttention avoids the O(N^2) HBM bottleneck entirely.

Online Softmax Enables It

The mathematical trick: compute softmax incrementally across tiles using running max and sum, then rescale. Exact results, no approximation.

Enables Long Context

Without FlashAttention, 128K context would need 32 GB per attention head. With it, memory stays constant regardless of sequence length.

Used Everywhere

vLLM, SGLang, TensorRT-LLM, PyTorch all use FlashAttention by default. It is the single most impactful inference optimization.