FlashAttention-2: Why the Attention Bottleneck Wasn’t Where Everyone Was Looking

Hello! In this blog let me share what I learned from reading the FlashAttention-2 paper, and try to make it click for you guys who know transformers but haven’t worked with GPU internals before.

The hook: softmax doesn’t tile

Every ML engineer has seen these four lines:

S = Q @ K.T            # (N, N) attention scores
S = S + mask           # causal mask if autoregressive
P = softmax(S)         # row-wise softmax
O = P @ V              # (N, d) output

Simple, right? Now try to run this on long sequences. That S matrix is N×N. At N=16k, that's 256 million entries. You can't fit it in your GPU's fast memory. The whole game is figuring out how to compute the same thing without ever materializing S.

The natural fix is tiling — split Q into row blocks, K and V into column blocks, process one tile at a time. Matmul tiles beautifully. Adding the mask tiles beautifully. The output accumulates from tiles beautifully.

But watch what happens when you try to tile that third line:

P = softmax(S)         # ← needs the full row to compute normalization

Softmax needs the max and the sum of the entire row to normalize. If you only have one tile of S, you don’t know the true max yet — a later tile might have a bigger value. You can’t softmax in pieces.

This is the whole problem FlashAttention exists to solve. FlashAttention-1 cracked it. FlashAttention-2 made it 2× faster. Let’s see how.

The minimum GPU primer you need

Before going further, three things about GPUs you should know.

1. Memory hierarchy. Your GPU has two kinds of memory: SRAM is what you compute on (think of it as a countertop). HBM is where everything sits when you’re not computing (its the pantry). Every trip from HBM to SRAM costs time,, so you try to reduce the trips to pantry. Most “slow” deep learning ops aren’t slow because of math — they’re slow because they keep walking back to HBM.
2. Execution hierarchy. GPUs run thousands of threads in a strict structure:

thread (1 worker)
  ↓ group of 32
warp (move in lockstep)
  ↓ group of 4 or 8
thread block (share an SRAM scratchpad)
  ↓ scheduled onto
SM (streaming multiprocessor — A100 has 108 of these)

One thread block runs on one SM. If you only launch 16 thread blocks but you have 108 SMs, 85% of your GPU is idle. This matters later.

3. The matmul tax. Modern GPUs have specialized hardware for matrix multiply (Tensor Cores). On an A100:

• Matmul throughput: 312 TFLOPS
• Everything else (exp, division, max): 19.5 TFLOPS

That’s a 16× gap. So the secret rule of GPU optimization is: do as much matmul as possible, do as little non-matmul as possible. Keep this in mind — it explains a lot of what FA2 does.

FA1 in one paragraph

FlashAttention-1 solved the “softmax doesn’t tile” problem with online softmax. The idea: as you process tiles left to right across a row, maintain two running statistics — m (the running max) and ℓ (the running sum of exponentials, normalized to current max). When a new tile arrives with a bigger max, rescale the old stats to match. This way you can stream through tiles and arrive at the exact softmax answer at the end. No approximation, no need to materialize the full N×N matrix. Memory drops from O(N²) to O(N), runtime drops 2-4×.

But FA1 still only hits 25–40% of peak FLOPs. Optimized matrix multiply hits 80–90%. So there’s a 2× speedup hiding somewhere. Where?

This is what FA2 set out to find.

FA2 change #1: Stop dividing so much

In FA1, every time a new tile arrives, you do something like:

# FA1 (simplified) — inside the inner loop
m_new = max(m, rowmax(S_tile))
P_tile = exp(S_tile - m_new) / l_new   # ← divide here
O = (l_old/l_new) * O + P_tile @ V_tile   # ← and rescale here
m, l = m_new, l_new

That division / l_new happens every iteration. It's non-matmul work. Slow path.

FA2 noticed: you don’t actually need to normalize until the very end. So it keeps an unnormalized running output Õ, and divides exactly once at the end:

# FA2 (simplified) — inside the inner loop
m_new = max(m, rowmax(S_tile))
alpha = exp(m - m_new)
O_tilde = alpha * O_tilde + exp(S_tile - m_new) @ V_tile  # no divide!
l = alpha * l + rowsum(exp(S_tile - m_new))
m = m_new

# At the very end, exactly once:
O = O_tilde / l

That alpha = exp(m - m_new) is the rescale factor — it fixes up the old accumulator so it matches the new max's normalization. When the max doesn't change, alpha = 1 (no-op). When the max grows, alpha < 1 (shrinks the old contribution).

In terms of raw FLOPs this is tiny. But because those divisions ran on the slow non-matmul units, removing them frees those units to do other work, and the kernel ends up much closer to matmul-bound. This is the FA2 philosophy in miniature: push everything onto Tensor Cores.

There’s a second tweak in this category: instead of storing both m and ℓ for the backward pass, just store L = m + log(ℓ). One number per row, recoverable in backward. Saves memory and a few FLOPs.

FA2 change #2: Parallelize over sequence length

This is the biggest practical win, and the most subtle. Bear with me.

FA1 parallelizes over batch_size × num_heads. Each thread block handles one (batch, head) pair. For a normal training run with batch=32 and 16 heads, that's 512 thread blocks on a 108-SM GPU. Plenty of work, GPU fully utilized.

But for long-context training, you typically have batch=1 (because each sequence is huge) and maybe 16 heads. That’s 16 thread blocks for 108 SMs. 85% of your GPU is idle. This is the FA1 bottleneck for long sequences.

The fix sounds obvious: also parallelize over sequence length. Just split the rows of Q across more workers, right? But here’s the catch — FA1’s loop structure made this hard.

In FA1, the outer loop is over K, V column blocks, the inner loop is over Q row blocks:

# FA1 loop order
for j in range(num_col_blocks):       # outer: K, V columns
    load K_j, V_j                      # once per outer iter
    for i in range(num_row_blocks):    # inner: Q rows
        load Q_i, O_i, m_i, l_i        # every inner iter
        update softmax stats
        write O_i, m_i, l_i back

Notice: K_j and V_j are loaded once per outer iteration. But the softmax state (O_i, m_i, l_i) for every row block has to be reloaded from HBM every inner iteration. The state lives in HBM across iterations, which is fine for serial execution but creates a coordination nightmare if you try to add row-wise parallelism on top. FA2 swaps the loops:

# FA2 loop order
for i in range(num_row_blocks):       # outer: Q rows
    load Q_i                           # once per outer iter
    init O_i, m_i, l_i in SRAM         # lives in SRAM!
    for j in range(num_col_blocks):    # inner: K, V columns
        load K_j, V_j                  # every inner iter
        update softmax stats in SRAM
    write O_i to HBM (once)

Now Q_i and the softmax state live in SRAM for the entire processing of one row block. K_j and V_j stream through. Two wins from one change:

1. HBM traffic drops — softmax state never round-trips through HBM during compute.
2. Row-wise parallelism is now trivial. Each row block is fully independent. No worker needs to talk to any other.

Visualizing the parallelism: Forward pass — each row block = one worker, sweeps left to right

  Q_block_1  →  [worker 1: sweeps K,V blocks 1..T_c]  →  O_block_1
  Q_block_2  →  [worker 2: sweeps K,V blocks 1..T_c]  →  O_block_2
  Q_block_3  →  [worker 3: sweeps K,V blocks 1..T_c]  →  O_block_3
  ...                                                       (parallel!)

Zero communication between workers. Zero atomics. Embarrassingly parallel. This is what gets the GPU fully utilized even at batch=1.

Why does this work for rows but not columns? Because softmax stats accumulate sequentially across columns within a row. If you split a row across two workers, they’d have to merge stats — defeating the point. But different rows have totally separate softmaxes, so they’re truly independent.

Backward pass goes the other way — columns. In backward, dK_j and dV_j accumulate contributions from every query position. If you parallelize over rows, multiple workers would compete to update the same dK_j, requiring atomic adds. dK and dV get touched a lot in the inner loop, so atomic adds on them would be expensive. FA2 flips it: parallelize over columns, so each worker owns one (K_j, V_j) and its gradients. The downside is that dQ_i now needs atomic adds — but dQ is touched less often, so the atomic cost is smaller. Smart trade-off.

FA2 change #3: Warp-level work splitting

Even within one thread block (4 warps), you have to decide how the warps split the tile work. This is the most subtle change but worth understanding.

FA1 used a split-K scheme: all 4 warps share Q, each warp owns a slice of K.

FA1: split-K
  Q (shared by all warps) × K^T [warp1 | warp2 | warp3 | warp4]
                                    ↓
  Each warp computes a slice of QK^T (split along K dimension)
                                    ↓
  To multiply with V, warps must combine their slices
  → write to shared memory, sync, sum

The problem: every warp produces a partial result. To compute the final output, they must communicate through shared memory. That communication costs time.

FA2 flips it to split-Q: warps share K and V, each warp owns a slice of Q.

FA2: split-Q
  Q [warp1 | warp2 | warp3 | warp4] × K^T (shared by all warps)
                                    ↓
  Each warp computes a complete strip of QK^T (split along row dimension)
                                    ↓
  Multiply by V (shared) directly → each warp writes its strip independently

Each warp owns a complete strip of rows. No coordination needed to compute PV. No shared memory shuffle. No sync in the forward pass.

Same reason as before — row-wise splits don’t need coordination, because softmax doesn’t couple across rows.

The causal masking freebie

For autoregressive models, you mask the upper triangle of the attention matrix (position i can’t attend to position j > i). The naive way is “compute everything, throw half away”:

S = Q @ K.T            # full N×N — wasteful
S = S + causal_mask    # zeros out upper triangle
P = softmax(S)
O = P @ V

But FA2 already operates on blocks. So instead of element-wise masking, it can reason about whole blocks:

For each block, just check the indices. Above the diagonal? Skip entirely. Below? Compute without the mask add (non-matmul work avoided). On the diagonal? Apply the mask.

Result: 1.7–1.8× speedup over non-causal attention. Not quite the theoretical 2× because the last row block has to process all column blocks while the first row block only processes one — some load imbalance. But almost free.

What this all adds up to

Concrete numbers from the paper:

• FA2 reaches 50–73% of peak FLOPs on A100 (vs. 25–40% for FA1).
• 2–3× speedup over FA1.
• 3–10× speedup over standard PyTorch attention.
• 225 TFLOPS/GPU end-to-end training of GPT-style models — 72% Model FLOPs Utilization on 8×A100

What I find most interesting about the FA1 → FA2 → FA3 sequence is that each version finds a bottleneck the previous version didn’t see.

• FA1 looked at attention and said “the bottleneck is HBM traffic. Let’s tile.” → 2–4× faster.
• FA2 looked at FA1 and said “the bottleneck is work partitioning between thread blocks and warps. Let’s restructure.” → another 2× faster.
• FA3 looked at FA2 and said “the bottleneck is that memory loads block compute. Let’s overlap them with async copies and producer-consumer warps.” → another 1.5–2× faster on H100.

Each layer of optimization reveals the next bottleneck. The first paper assumed memory was the problem. The second assumed scheduling was. The third assumed instruction overlap was. None of them were wrong — they were all right at their level. Performance optimization is iterative excavation, and you don’t see the next layer until you’ve cleared the current one.

For me, this is the most valuable lesson from reading the FlashAttention papers. Not just the algorithm, but the way of thinking: when something seems already-optimized, profile harder. The next bottleneck is always there, just hiding behind the last one.

I’ll explore FA3’s tricks in an upcoming blog. If you have feedback or if I am wrong, feel free to reach out.