Yes, let's see attention. This a beautiful mechanism anyday.

But at the same time it's too slow. Not because the math is hard. It's slow because of where the numbers lives.

Standard Attention: What Actually Happens

Let's start with vanilla attention. You've seen this formula a thousand times:

Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d_k)) @ V

Looks simple. Three matrices, two multiplications, a softmax. But let's trace what actually happens step by step in a standard implementation:

Step 1: S = Q @ K^T          # (N, d) @ (d, N) = (N, N)   - store S to HBM 
Step 2: P = softmax(S)     # read S from HBM, compute, store P to HBM
Step 3: O = P @ V             # read P from HBM, (N, N) @ (N, d) = (N, d)

Where N is sequence length and d is the head dimension.

See that (N, N) matrix? That's the attention score matrix. For a sequence length of 4096, that's 4096 * 4096 = 16 million entries. Per head. Per layer. Per batch element.

For a model with 32 heads and batch size 8:

Attention matrix size = 8 * 32 * 4096 * 4096 * 2 bytes (BF16) =  8 GB

Just for one layer's attention scores. And you need to store this
for the backward pass too.

But here's what most people miss. The attention scores aren't the real problem. The real problem is how many times we read and write them.

The Real Bottleneck: It's Not Compute

Let me explain something that changed how I think about GPU perf.

A GPU has two kinds of memory:

HBM (High Bandwidth Memory):

The main GPU memory. On an A100, that's 80GB at ~2 TB/s bandwidth. This is where your tensors live.

SRAM (on-chip memory):

Tiny but super fast. On an A100, about 20MB total across all SMs, but at ~19 TB/s bandwidth. This is the scratchpad where actual computation happens.

SRAM is ~10x faster, but 4000x smaller.

Now here's the key insight from the Flash Attention paper. For attention, the GPU spends most of its time moving data between HBM and SRAM, not doing the actual math.

Let's count the memory operations for standard attention:

S = Q @ K^T
  - Read Q (N*d) and K (N*d) from HBM
  - Write S (N*N) to HBM

P = softmax(S)
  - Read S (N*N) from HBM
  - Write P (N*N) to HBM

O = P @ V
  - Read P (N*N) and V (N*d) from HBM
  - Write O (N*d) to HBM

Total HBM reads/writes: O(N² + Nd) ≈ O(N²) for large N

Ahh that N² term is the killer. The actual compute (matrix multiplications) is also O(N²d), but the compute is fast. The memory transfers are what takes time.

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

The GPU is literally sitting idle, waiting for data to arrive from HBM.

Flash Attention

Here's the core idea of Flash Attention, and it's beautiful in its simplicity:

Don't ever write the full N*N attention matrix to HBM. Compute attention in tiles that fit in SRAM.

Instead of computing the entire S = QK^T, storing it, then computing softmax, then reading it back... we process small blocks at a time, keeping everything in the fast on-chip SRAM.

But there's a catch. Softmax needs to see the entire row to compute the denominator (sum of exponentials). You can't just softmax a tile independently, the normalization depends on all elements in that row.

This is where the online softmax trick comes in. And it's clever.

Online Softmax: The Key Trick

Normal softmax:

softmax(x_i) = exp(x_i) / Σ exp(x_j)  for all j

You need the full sum before you can compute any single output. Or actually?

No maybe. For numerical stability, we already subtract the max before exponentiating:

softmax(x_i) = exp(x_i - max(x)) / Σ exp(x_j - max(x))

Online softmax maintains a running max and a running sum of exponentials. When we process a new tile, we can correct the previous results:

Processing tile 1:
  m1 = max(tile_1)                                # local max
  l1 = sum_of(exp(tile_1 - m1))              # local sum of exp
  O1 = softmax_local(tile_1) @ V1     # local output

Processing tile 2:
  m2 = max(tile_2)                              # local max of this tile
  m_new = max(m1, m2)                    # global max so far
  
  # Correction factor: our old exponentials used m1, but now
  # the true max is m_new. Fix by multiplying by exp(m1 - m_new)
  l_new = l1 * exp(m1 - m_new) + sum_of(exp(tile_2 - m_new))
  
  # Also correct the running output
  O_new = O1 * (l1 * exp(m1 - m_new) / l_new)
        + softmax_local(tile_2, m_new) @ V2 / l_new

The key insight: when we see a new tile with a larger max, we can rescale our previous partial results. We never need to go back and re-read previous tiles.

This means we can process the entire attention computation in one pass through the data, tile by tile, keeping only the current tile in SRAM.

The Tiling Algorithm

# Q, K, V are in HBM. O (output) will be written to HBM.
# Block sizes Br, Bc chosen to fit in SRAM.

Divide Q into blocks: Q1, Q2, ..., Q_Tr    (each Br * d)
Divide K, V into blocks: K1, K2, ..., K_Tc  (each Bc * d)

for each Q block Q_i:               # outer loop
    Load Q_i from HBM to SRAM
    Initialize: O_i = 0, l_i = 0, m_i = -inf
    
    for each K, V block (K_j, V_j):  # inner loop
        Load K_j, V_j from HBM to SRAM
        
        # Compute attention scores for this tile
        S_ij = Q_i @ K_j^T           # (Br * Bc) — stays in SRAM!
        
        # Online softmax update
        m_new = max(m_i, rowmax(S_ij))
        l_new = l_i * exp(m_i - m_new) + rowsum(exp(S_ij - m_new))
        
        # Update running output
        O_i = O_i * (l_i * exp(m_i - m_new) / l_new)
            + exp(S_ij - m_new) @ V_j / l_new
        
        m_i = m_new
        l_i = l_new
    
    Write O_i to HBM                 # only final result hits HBM

Look at what's happening here. The S_ij matrix is only (Br * Bc) tile that fits in SRAM. We never create the full N*N matrix.

Let's count memory operations now:

Flash Attention:
  Never materialized
  Never materialized
  Total HBM access: O(N^2 d / M)
where M = SRAM size

That M factor is huge. On an A100, M is roughly 100KB per thread block. For d=128 and block sizes of 128, we're doing ~100x fewer HBM accesses.

Same computation, same result, but less memory traffic. The GPU can finally use its tensor cores at full speed because it's not constantly waiting for HBM.

The Backward Pass Problem

There's one more clever bit. In standard attention, you save the N*N attention matrix P for the backward pass. Flash Attention doesn't have P in memory, we never stored it.

So what do we do? We recompute it.

Yes, recompute. It sounds wasteful, we're doing the forward pass computation again during backward. But remember:

• Storing P would cost O(N^2) HBM memory
• Recomputing P from Q, K (which we already have) costs extra FLOPs but no extra HBM
• Since we're out of memory, the GPU was idle anyway

We just save the per-row statistics from the forward pass. That's O(N) memory instead of O(N²). During backward, we recompute each tile of P on the fly using the same tiling approach.

Trading compute for memory access. This is the theme of Flash Attention. And it works because most GPUs have way more compute than memory bandwidth.

If you've used PyTorch recently, you've probably used Flash Attention without even knowing it:

PyTorch 2.0+ automatically uses Flash Attention, when conditions are met (CUDA, BF16/FP16, no custom mask)

The Last Block

Flash Attention doesn't change what we compute, it changes where we compute it. The result is identical. But by respecting the memory: keeping tiles in SRAM, never materializing the N*N matrix to HBM, we get 2-10x speedups and can train on sequences that would otherwise go OOM.

But Flash Attention only solves the training side. Inference has a completely different bottleneck, and a completely different set of tricks.