6.2 FlashAttention V2 in depth

This article unpacks FlashAttention V2’s core improvements over V1: swapping the outer/inner loop order, better work partitioning inside thread blocks, fewer non-matmul ops, and reaching 50–73% of peak FLOPS on A100.

📑 Contents

• 1. Recap: V1 performance bottlenecks
• 2. Improvement 1: swap the loop order
• 3. Improvement 2: fewer non-matmul ops
• 4. Improvement 3: warp-level work partitioning
• 5. Causal Mask optimization
• 6. Full forward algorithm
• 7. Backward-pass optimizations
• 8. Performance analysis and comparison
• Summary
• Self-check list
• References

---

1. Recap: V1 performance bottlenecks

1.1 V1’s real utilization

FlashAttention V1 only hit 25–40% of peak FLOPS on A100. For an operator that should be compute-bound after IO is cut down, that utilization is weak.

The bottlenecks fall into three buckets (largest impact first):

1.2 Hardware context

Key A100 numbers:

• FP16 Tensor Core peak: 312 TFLOPS
• HBM bandwidth: 2 TB/s
• SRAM capacity: 192 KB/SM (Shared Memory; up to 164 KB configurable for the user)
• Warp count: each Thread Block can hold multiple Warps (often 4–8)

To keep Tensor Cores busy, the compute pipeline must not stall on non-GEMM work.

---

2. Improvement 1: swap the loop order

2.1 Problems with V1’s loop structure

In V1 the outer loop walks K/V tiles and the inner loop walks Q tiles. That causes two issues:

1. Repeated Q and O traffic: every new K/V tile forces all Q tiles and their O outputs to reload from and write back to HBM
2. Limited parallelism: outer K/V tiles are serially dependent through online softmax’s running max/sum, so V1 can only parallelize over batch × heads. When batch and heads are small and sequences are long (long-context settings), you cannot use sequence-dimension parallelism and SM utilization suffers

2.2 V2’s new loop structure

V2 swaps the order to outer over Q tiles, inner over K/V tiles:

Outer loop: for i = 1 to T_r (over Q tiles) — mapped to different Thread Blocks
    Load Q_i from HBM into SRAM (once)
    Init O_i = 0, m_i = -inf, l_i = 0
    Inner loop: for j = 1 to T_c (over K/V tiles)
        Load K_j, V_j from HBM into SRAM
        Compute S_ij, update stats and O_i (all in SRAM)
    Write final O_i back to HBM (once)

2.3 What you gain

That seemingly small swap buys three things:

Gain 1: O is written once at the end

In V1, every K/V tile reads and writes O — $2T_c$ HBM accesses ($T_c$ reads + $T_c$ writes). In V2, O is written once after the inner loop.

Gain 2: the outer loop parallelizes naturally

Different Q tiles are fully independent (each computes its own output rows), so they map straight onto different Thread Blocks for GPU-level parallelism.

Gain 3: Q is loaded once

Each Thread Block owns one Q tile; Q stays in Shared Memory for the whole inner loop.

📌 Key point: V1 put K/V outside because of online softmax’s accumulation (update O’s normalization after each K/V column tile). The cost is repeated O read+write. V2’s finding: loading Q once is better, because cutting O traffic pays more than cutting K/V reads.

---

3. Improvement 2: fewer non-matmul ops

3.1 The problem: non-GEMM instructions

Besides the two matmuls ($QK^\top$ and $PV$), tiled Attention has many elementwise ops:

• Subtract max: $S_{ij} - m_i$
• Exp: $\exp(\cdot)$
• Row sum: $\text{rowsum}(\cdot)$
• Rescaling: multiply by correction factors

These cannot use Tensor Cores — only CUDA Cores — and stall the compute pipeline.

3.2 V2’s fix: delay rescaling

V2 delays Softmax rescaling until after the inner loop and does it once:

V1 (rescale every K/V tile):

$$O_i^{(j)} = \frac{l_i^{(j-1)} \cdot e^{m_i^{(j-1)} - m_i^{(j)}}}{l_i^{(j)}} \cdot O_i^{(j-1)} + \frac{e^{S_{ij} - m_i^{(j)}}}{l_i^{(j)}} \cdot V_j$$

V2 (keep an unnormalized accumulator):

$$\tilde{O}_i^{(j)} = e^{m_i^{(j-1)} - m_i^{(j)}} \cdot \tilde{O}_i^{(j-1)} + e^{S_{ij} - m_i^{(j)}} \cdot V_j$$

Normalize once at the end:

$$O_i = \tilde{O}_i^{(T_c)} / l_i^{(T_c)}$$

3.3 Quantifying the win

That removes three elementwise ops per inner-loop iteration:

• Denormalize old O: multiply by $l^{(j-1)}$
• Renormalize old O: divide by $l^{(j)}$
• Normalize the new contribution: divide by $l^{(j)}$

V2 strips those 「2 divides + 1 multiply」 from the inner loop and does one divide at the end. When $T_c$ is large, the savings add up.

---

4. Improvement 3: warp-level work partitioning

4.1 V1’s plan: split along K/V

In V1, warps inside a Thread Block split the same $S_{ij}$ tile along the K/V axis.

Suppose 4 warps and $B_c = 128$:

• Warp 0 owns K/V columns 0–31
• Warp 1 owns columns 32–63
• Warp 2 owns columns 64–95
• Warp 3 owns columns 96–127

Each warp only sees part of the columns of $S_{ij}$, so:

• Row-max / row-sum need a cross-warp merge of local results
• For $P \cdot V$, several warps accumulate into the same O row and must sync again

That split-K style forces frequent inter-warp sync and shared-memory traffic.

4.2 V2’s plan: split along Q

V2 further splits the Q tile across warps:

Suppose 4 warps and $B_r = 128$:

• Warp 0 owns Q rows 0–31
• Warp 1 owns rows 32–63
• Warp 2 owns rows 64–95
• Warp 3 owns rows 96–127

Each warp independently computes attention between its Q rows and the full K/V tile. No interference.

4.3 Comparing the two strategies

💡 Tip: If you split on K/V, each warp sees only part of the Softmax input and must communicate max and sum across warps — one reason V1 underperformed.

---

5. Causal Mask optimization

5.1 What Causal Mask does

In autoregressive models (GPT and the like), token $i$ may only “see” itself and earlier tokens when generating — position $i$ can attend only to $j \leq i$. That causality constraint is central to autoregressive LM training: the model must predict the next token from history, not peek at the answer.

In practice this is a lower-triangular mask added to attention scores before Softmax: $\tilde{S} = QK^\top / \sqrt{d} + M$. Masked positions get $-\infty$, become 0 after $\exp$, and contribute nothing to Softmax output.

Visually, the $N \times N$ attention matrix is only valid on the lower triangle (including the diagonal); the upper triangle is zero. Roughly half the compute is wasted in theory — exactly the structure V2’s block-skip strategy exploits.

5.2 Waste in a naive implementation

Without special handling, masked-out positions still run full GEMM and Softmax — about half the FLOPs wasted.

5.3 V2’s block-level skip

V2 adds block-level Causal Mask optimization:

For Q tile $i$ (rows $[iB_r, (i+1)B_r)$) and K tile $j$ (columns $[jB_c, (j+1)B_c)$):

1. Fully below the mask ($(i+1)B_r - 1 \geq (j+1)B_c - 1$ and $iB_r \geq (j+1)B_c - 1$, i.e. Q min row ≥ K max col): every position has row ≥ col, fully visible, compute normally
2. Fully above the mask ($(i+1)B_r - 1 < jB_c$, i.e. Q max row < K min col): every position has row < col, fully invisible, skip the block
3. Crossing the diagonal: mixed visible and invisible positions — apply the mask elementwise

Below diagonal (all visible)    Diagonal block (partial mask)    Above diagonal (all skip)
┌─────────────┐      ┌─────────────┐      ┌─────────────┐
│ ■ ■ ■ ■ ■ ■ │      │ ■ ■ ■ □ □ □ │      │ □ □ □ □ □ □ │
│ ■ ■ ■ ■ ■ ■ │      │ ■ ■ ■ ■ □ □ │      │ □ □ □ □ □ □ │
│ ■ ■ ■ ■ ■ ■ │      │ ■ ■ ■ ■ ■ □ │      │ □ □ □ □ □ □ │
└─────────────┘      └─────────────┘      └─────────────┘

5.4 Speedup

For Causal Attention, block skip cuts about 50% of useless work. Combined with outer-loop Q assignment, each Thread Block only walks K/V up to the diagonal and can exit the inner loop early.

---

6. Full forward algorithm

6.1 Pseudocode

def flash_attention_v2_forward(Q, K, V, B_r, B_c, causal=False):
    N, d = Q.shape
    O = zeros(N, d)
    L = zeros(N)  # logsumexp, for backward

    # Outer loop (parallelized across Thread Blocks)
    parallel for i in range(0, N, B_r):
        Q_i = Q[i:i+B_r]  # load into SRAM; stays for whole inner loop

        # Init accumulators (SRAM/registers)
        m_i = full(B_r, -inf)
        l_i = zeros(B_r)
        O_tilde = zeros(B_r, d)  # unnormalized accumulated output

        # Inner loop: walk K/V tiles
        kv_end = min(N, (i + B_r)) if causal else N
        for j in range(0, kv_end, B_c):
            K_j = K[j:j+B_c]  # load from HBM
            V_j = V[j:j+B_c]  # load from HBM

            # Local attention scores
            S_ij = Q_i @ K_j.T / sqrt(d)

            # Apply causal mask (only needed on diagonal-crossing tiles)
            # i, j are absolute start coords; when K max col (j+B_c-1) > Q min row (i), this tile crosses the diagonal
            if causal and i < j + B_c - 1:
                apply_causal_mask(S_ij, i, j, B_r, B_c)

            # Update statistics
            m_new = max(m_i, rowmax(S_ij))
            P_ij = exp(S_ij - m_new[:, None])
            l_new = exp(m_i - m_new) * l_i + rowsum(P_ij)

            # Update unnormalized output
            O_tilde = exp(m_i - m_new)[:, None] * O_tilde + P_ij @ V_j

            m_i = m_new
            l_i = l_new

        # Final normalize and write back
        O[i:i+B_r] = O_tilde / l_i[:, None]
        L[i:i+B_r] = m_i + log(l_i)  # save logsumexp for backward

    return O, L

6.2 Key differences from V1

---

7. Backward-pass optimizations

7.1 Save logsumexp instead of separate m and l

V2 saves $L_i = m_i + \log(l_i)$ (row-wise logsumexp) in the forward pass instead of separate $m_i$ and $l_i$. Two wins:

• Half the statistic storage
• Backward can recover $P_{ij}$ directly: $P_{ij} = \exp(S_{ij} - L_i)$

7.2 Backward parallelism

V2’s backward uses an outer loop over K/V tiles (same as V1 backward, opposite of V2 forward’s Q-outer loop), so $dK$ and $dV$ updates need no atomics:

• $dK_j$ and $dV_j$ are accumulated exclusively by one Thread Block (that block walks all Q tiles)
• $dQ_i$ must accumulate across Thread Blocks (atomic add or a two-pass scan)

7.3 The D vector in backward

Softmax gradients need $dS_{ij} = P_{ij} \odot (dP_{ij} - D_i)$, where

$$D_i = \text{rowsum}(dP_i \odot P_i) = \text{rowsum}(dO_i \odot O_i) \in \mathbb{R}^N$$

That rewrite already appears in the V1 paper (Appendix B.3): the left side reduces a full length-$N$ row, and $P_i$, $dP_i$ may not fit in SRAM; the right side only needs a length-$d$ dot product, independent of sequence length.

V2’s engineering change: compute $D$ in a separate preprocess kernel once for the whole $\mathbb{R}^N$ vector, then read it in the main backward loop; V1 computed it tile-by-tile inside that loop.

---

8. Performance analysis and comparison

8.1 Numbers on A100

V2 is about 1.6× faster than V1.

8.2 Versus other implementations

A100-80GB, sequence length 2048, head dim 128:

8.3 Memory advantage on long sequences (training activation estimate)

The table estimates total attention-related activation memory in training (intermediate tensors and gradient buffers, rough multi-head config) to show $N^2$ vs $N$ scaling — not the cost of a single matrix:

---

📝 Summary

FlashAttention V2 lifts A100 efficiency from 25–40% to 50–73% with three core changes plus one extra opt:

1. Swap loop order: outer over Q tiles (parallel), inner over K/V — less O HBM traffic and natural parallelism
2. Delayed normalize (fewer non-GEMM ops): inner loop only does exp scaling, one unified normalize at the end
3. Warp-level opt: each warp owns different Q rows — no inter-warp sync or communication
4. Extra: Causal block skip: use the triangular mask to skip dead blocks — ~50% less work for Causal Attention

---

🎯 Self-check list

• Explain why V2 moves the outer loop from K/V to Q
• Explain why O’s HBM accesses drop from $2T_c$ to 1
• Describe V2’s warp partitioning and its edge over V1
• Derive correctness of delayed normalize (unnormalized accumulate then divide by $l$ equals per-step normalize)
• Sketch the three block-skip cases under Causal Mask
• Explain how logsumexp $L = m + \log(l)$ simplifies backward
• Compare V1 vs V2 TFLOPS utilization on A100

---

📚 References

• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• FlashAttention-1: Fast and Memory-Efficient Exact Attention with IO-Awareness
• FlashAttention official implementation - GitHub
• Tri Dao - FlashAttention-2 Blog Post
• NVIDIA A100 Tensor Core GPU Architecture