FlashAttention V2 — infra.bea.ai

FlashAttention V1 solved the memory problem but only achieved ~35% of theoretical FLOPS on A100. V2 closes this gap to ~72% through three targeted optimizations: swapping loop order, delaying rescaling, and better warp work distribution. No algorithmic change — same math, same IO, faster execution.

1. Why V1 Was Only 35% Efficient

Despite eliminating $O(S^2)$ HBM traffic, V1 left performance on the table:

Issue	Cause	Impact
Non-matmul FLOPs	Softmax exp/div computed per element (not Tensor Core)	~50% of time spent on non-GEMM ops
Warp divergence in rescaling	Rescaling O after each K/V block → irregular ops	Extra non-GEMM instructions
Suboptimal warp partitioning	K/V split across warps → synchronization overhead	Communication between warps

Key observation: On modern GPUs, Tensor Core GEMM is ~16× faster than scalar FP32 ops. Any non-GEMM computation (softmax, rescaling, masking) becomes a bottleneck even if it’s algorithmically “free.”

2. Optimization 1: Swap Loop Order

V1: Outer loop over K/V blocks, inner loop over Q blocks

for each K/V block j:           ← outer
    for each Q block i:         ← inner
        compute S_ij, update O_i

Problem: Each Q block’s output $O_i$ is updated by every K/V block. This means $O_i$ must be read from and written to HBM repeatedly — once per K/V block.

V2: Outer loop over Q blocks, inner loop over K/V blocks

for each Q block i:             ← outer
    load Q_i to SRAM (once)
    for each K/V block j:       ← inner
        load K_j, V_j to SRAM
        compute S_ij, update O_i (in SRAM)
    write O_i to HBM (once)

Benefit: $Q_i$ stays in SRAM for the entire inner loop — loaded once, not once per K/V block. $O_i$ is accumulated in SRAM and written to HBM only at the end. This reduces HBM writes for O from $O(S^2 / B_c)$ to $O(S)$ .

3. Optimization 2: Delay Rescaling

V1: Rescale O after every K/V block

After processing each K/V tile, V1 immediately rescales the output:

O_i = O_i \times \frac{\ell^{(old)}}{\ell^{(new)}} + \frac{1}{\ell^{(new)}} \tilde{P} V

This division by $\ell^{(new)}$ is a non-matmul operation applied to every element of $O_i$ — scalar ops that can’t use Tensor Cores.

V2: Defer normalization to the end

V2 maintains an unnormalized accumulator and applies the final $1/\ell$ division once, after all K/V blocks:

\tilde{O}_i = \tilde{O}_i \times e^{m^{(old)} - m^{(new)}} + \tilde{P} V

O_i = \tilde{O}_i / \ell_i \qquad \text{(one division at the very end)}

Benefit: Eliminates one division per K/V block per element. For 64 K/V blocks and $B_r \times d_k$ output elements, this removes ~64× the division operations. The multiplications by $e^{m^{(old)} - m^{(new)}}$ are cheaper (can be precomputed per row, broadcast).

4. Optimization 3: Warp-Level Work Partitioning

V1: Split K/V across warps

Each warp handles a portion of the K/V dimension. After computing partial results, warps must communicate (via shared memory) to combine their partial softmax statistics.

V2: Split Q across warps

Each warp handles a different portion of Q rows — processing the entire K/V sequence independently. No inter-warp communication needed during the inner loop.

V1 warp partition (along K/V):
  Warp 0: K columns 0-31   → partial score → needs sync to combine
  Warp 1: K columns 32-63  → partial score → needs sync to combine

V2 warp partition (along Q):
  Warp 0: Q rows 0-15      → full attention → independent, no sync
  Warp 1: Q rows 16-31     → full attention → independent, no sync

Benefit: Eliminates __syncthreads() and shared memory communication between warps within the inner loop. Each warp operates independently — pure Tensor Core GEMM.

5. Performance Impact

Metric	V1	V2	Improvement
GPU utilization (A100)	~35%	~72%	2×
Speed (S=2K, FP16)	1×	~2×	Loop reorder + warp partition
Speed (S=16K, FP16)	1×	~2×	Consistent across lengths
Non-GEMM FLOP fraction	~50%	~20%	Delayed rescaling

Comparison to Theoretical Peak

For FP16 attention on A100 (312 TFLOPS):

Theoretical: $4S^2 d_k / 312 \times 10^{12}$ seconds
V1: ~35% of theoretical → ~2.9× slower than hardware limit
V2: ~72% of theoretical → ~1.4× slower than hardware limit

The remaining 28% gap comes from: causal masking (irregular computation), block boundary effects, and the inherent overhead of online softmax corrections.

6. FlashAttention V2 in Practice

Usage

FlashAttention V2 is the default in most frameworks:

# PyTorch (built-in since 2.0)
from torch.nn.functional import scaled_dot_product_attention
output = scaled_dot_product_attention(Q, K, V, is_causal=True)
# Automatically uses FlashAttention V2 when available

# Explicit flash-attn package
from flash_attn import flash_attn_func
output = flash_attn_func(Q, K, V, causal=True)

When V2 Is Used vs Not

Condition	Uses FlashAttention?
FP16/BF16, head dim ≤ 256	✅ Yes
FP32	❌ No (falls back to naive or memory-efficient)
Head dim > 256	❌ No (tile doesn’t fit in SRAM)
Custom attention mask (non-causal, non-full)	⚠️ Depends on implementation
Inference with KV cache	✅ Yes (FlashDecoding variant)

7. FlashAttention-3 and Beyond

FlashAttention-3 (2024) targets Hopper GPUs specifically:

Asynchronous GEMM + softmax: Uses Hopper’s TMA (Tensor Memory Accelerator) to overlap data loading with computation
FP8 attention: Uses Hopper’s FP8 Tensor Cores for ~2× more throughput
Warp-specialization: Dedicated warps for loading vs computing

The evolution: V1 (correct algorithm) → V2 (efficient implementation) → V3 (hardware-specific optimization).

Self-Check

Explain why V1 only achieved 35% GPU utilization despite optimal IO
Explain the loop order change (V1: K/V outer, V2: Q outer) and its HBM benefit
Explain delayed rescaling and why it reduces non-GEMM operations
Explain the warp partitioning change (K/V split → Q split) and why it eliminates sync
Check if scaled_dot_product_attention uses FlashAttention on your GPU
Explain FlashAttention-3’s key innovations for Hopper