FlashAttention & Efficient Attention

Chapter 1: The Efficient Attention Taxonomy

Before FlashAttention, researchers attacked the quadratic bottleneck by changing the attention algorithm itself — computing something cheaper that approximates attention. These approaches fall into three families: sparse, low-rank, and kernel.

Sparse Attention

Instead of letting every token attend to every other token, sparse attention restricts which pairs of tokens interact. Think of it as poking holes in the N×N matrix — only computing certain entries.

Sparse Transformers (Child et al., 2019) use a sliding window: each token attends only to nearby tokens within a fixed window size w. Complexity drops from O(N²) to O(N√N). BigBird (Zaheer et al., 2021) combines three patterns: local windows (nearby tokens), global tokens (a few special tokens attend to everything), and random connections (to reduce graph diameter). Complexity: O(N).

Low-Rank Attention

Low-rank methods observe that the N×N attention matrix often has low effective rank — it can be approximated by projecting K and V to a much smaller dimension k ≪ N. Linformer (Wang et al., 2020) projects the N×d key and value matrices down to k×d using learned projection matrices, reducing the attention computation to k×k instead of N×N.

The downside: hard to maintain causality (the model shouldn't peek at future tokens), and quality degrades when the true rank is higher than k.

Kernel Attention

Kernel methods replace the exp(Q K^T) similarity with a decomposable feature map: sim(q_i, k_j) ≈ φ(q_i)^T φ(k_j). The Linear Transformer (Katharopoulos et al., 2020) uses this trick to avoid computing the full N×N matrix. Without the softmax, you can exploit the associative property of matrix multiplication:

When d ≪ N (which it usually is), this is a win. But the approximation quality suffers, especially on language tasks where the sharp attention patterns that softmax produces matter.

Notice that FlashAttention doesn't reduce the FLOP count. It has the same O(N²d) compute. Its insight is entirely different: the problem isn't that we're doing too many operations — it's that we're moving data too slowly. The bottleneck is memory bandwidth, not compute.

Chapter 2: Why Standard Attention is Slow

Here's the surprise. Standard attention on a modern GPU like the A100 isn't slow because of arithmetic. The A100 can perform 312 TFLOPS of BF16 computation per second. For a sequence of N = 1024 with head dimension d = 64, the attention requires roughly 134 million FLOPs. At 312 TFLOPS, that's done in microseconds. So why does attention take milliseconds?

The answer: memory traffic. Let's trace what happens step by step in a standard PyTorch attention implementation.

Count the memory trips. The N×N matrix S gets written to HBM after Step 1, read back for Step 2, written again, read back for Step 3, and so on. Each of those lightweight element-wise operations (masking, softmax, dropout) is sandwiched between expensive HBM reads and writes. The compute in each step is trivial; the data movement is not.

Compute-bound vs. Memory-bound

An operation is compute-bound when the time is dominated by arithmetic: matmuls with large inner dimensions fall here. An operation is memory-bound when the time is dominated by moving bytes between memory levels: element-wise ops (masking, softmax, dropout) with large tensors fall here.

The way to tell: compute the arithmetic intensity — the ratio of FLOPs to bytes transferred.

The A100 has a compute bandwidth of 312 TFLOPS and a memory bandwidth of 1,935 GB/s. The crossover ratio is:

If your operation's arithmetic intensity is below 161, it's memory-bound. If above, it's compute-bound. Let's check attention.

Worked Example: Is Attention Memory-Bound?

For the softmax step: N² values must be read and N² values written. Each BF16 value is 2 bytes. Total memory: 2 × N² × 2 = 4N² bytes. FLOPs: roughly 5N² (exp, subtract max, divide). Arithmetic intensity: 5N² / 4N² ≈ 1.25 FLOPs/Byte. That's 128× below the A100's crossover. Massively memory-bound.

For comparison, a large matrix multiplication like (8192 × 8192) × (8192 × 8192) has arithmetic intensity around 2,730 FLOPs/Byte — solidly compute-bound.

Chapter 3: GPU Memory Hierarchy

To understand FlashAttention, you need a mental model of where data lives inside a GPU. A modern GPU has (at least) three levels of memory, and the speed difference between them is staggering.

HBM (High Bandwidth Memory) is the big, slow memory. 40–80 GB is plenty of room, but you can only read from it at about 1.5–2 TB/s. SRAM (on-chip shared memory) is tiny but fast — about 20 MB total across all streaming multiprocessors (SMs), but accessible at roughly 19 TB/s. That's a 10× bandwidth gap.

Here's the problem for standard attention. The N×N attention matrix is large. For N = 4096 in BF16, it's 4096² × 2 bytes = 32 MB. That exceeds SRAM capacity. So standard implementations have no choice: they write S to HBM after the QK^T matmul, read it back for softmax, write the result back, read it again for dropout, and so on.

But the Q, K, V tiles? A single tile of, say, 128 × 64 in BF16 is only 128 × 64 × 2 = 16 KB. A pair of tiles (one from Q, one from K) fits comfortably in SRAM. This is the opening that FlashAttention exploits.

How GPU Execution Works

A GPU runs kernels — functions that execute in parallel across thousands of threads. Threads are grouped into warps (32 threads), warps into thread blocks, and thread blocks are assigned to streaming multiprocessors (SMs). Each thread block has access to its SM's shared memory (SRAM). Data that multiple threads need can be loaded once into shared memory, avoiding redundant HBM reads. This is called tiling.

Kernel fusion is the other key technique: instead of launching five separate GPU kernels (matmul, mask, softmax, dropout, matmul) with HBM writes between each one, we fuse them into a single kernel that does all the work while data is still in SRAM.

Chapter 4: FlashAttention — The Key Idea

Here's the plan. Instead of computing the full N×N matrix and storing it, we break Q, K, V into small blocks that fit in SRAM. For each pair of blocks, we compute a partial attention result, then combine the partial results to get the exact final answer.

Concretely: split Q into blocks of B_r rows and K, V into blocks of B_c rows. The outer loop iterates over K/V blocks (j = 1, ..., T_c). The inner loop iterates over Q blocks (i = 1, ..., T_r). For each (i, j) pair, we compute the B_r × B_c tile of the attention matrix — and immediately use it to update the running output, then discard it.

But Wait — What About Softmax?

There's one serious obstacle. Softmax is a global operation. For each row i, softmax divides by the sum of exp(s_ij) across ALL columns j = 1, ..., N. If we only have one tile at a time, we don't have the global sum. We can't do softmax.

Or can we?

The answer is online softmax — an algorithm that computes softmax incrementally, one tile at a time, maintaining a running maximum and running sum. When a new tile arrives with larger values, it rescales all the previous partial results. At the end, the answer is identical to computing softmax over the entire row at once.

This is the algorithmic heart of FlashAttention, and we'll derive it step by step in the next chapter.

Memory Analysis

Standard attention reads/writes the N×N matrix multiple times. Total HBM access: O(Nd + N²). FlashAttention only reads Q, K, V once and writes O once, but does it multiple times in tiles. Total HBM access: O(N²d²M⁻¹), where M is SRAM size. Since M is roughly 20 MB on an A100:

In practice, this is a 9× reduction in HBM reads and writes. Despite doing 13% more FLOPs (due to the rescaling arithmetic), FlashAttention is 6× faster in wall-clock time.

Chapter 5: Online Softmax — The Algorithmic Heart

This is the core derivation. If you understand this chapter, you understand FlashAttention. Let's build online softmax from scratch.

Standard Softmax

Given a vector x = [x₁, x₂, ..., x_N], the softmax is:

This requires two passes over the data: first compute all exp(x_j) and sum them, then divide each by the sum. Both passes need all N values simultaneously.

Step 1: Numerical Stability

In practice, we first subtract the maximum value for numerical stability. If x_max = 42, then exp(42) overflows float32. But exp(42 − 42) = exp(0) = 1 is fine. The result is identical because:

where m = max(x). The exp(m) cancels in numerator and denominator. So stable softmax needs three passes: (1) find max, (2) compute exp and sum, (3) divide.

Step 2: The Tiling Problem

Now suppose we split x into two blocks: x = [x⁽¹⁾, x⁽²⁾]. We process block 1 first, then block 2. After block 1, we have:

This is a valid local softmax weighted by block 1's values. But the denominator l⁽¹⁾ only reflects block 1. It's wrong as a global softmax.

Step 3: Rescaling After Block 2

When we process block 2, we discover a new maximum and new sum. We need to fix our earlier result. Here's the algebra, step by step.

After processing block 2, we have:

We need to rescale block 1's exponentials. Originally we computed exp(x_j − m⁽¹⁾). We need exp(x_j − m^new). The correction factor is:

So we multiply all of block 1's contributions by exp(m⁽¹⁾ − m^new). Since m^new ≥ m⁽¹⁾, this factor is ≤ 1 (it shrinks old values that were computed with a smaller max). Similarly for block 2's terms: multiply by exp(m⁽²⁾ − m^new).

The corrected sums:

And the corrected output:

Step 4: The General Recurrence

For B blocks processed sequentially, after processing block j, we update:

Three running scalars per row: the maximum m, the sum l, and the output vector O. These are tiny — they live in registers. The N×N matrix tiles are computed and discarded in SRAM.

Worked Example: 4×4 Matrix in 2×2 Tiles

Let's compute attention for a concrete 4×4 score matrix, tiled into 2×2 blocks, to see the online softmax in action.

Score matrix S (after QK^T):

(Showing rows 1–2 for clarity. The process is identical for rows 3–4.)

Block 1 (columns 1–2 for row 1): scores = [1.0, 2.0]

Block 2 (columns 3–4 for row 1): scores = [0.5, 0.1]

That final o^new is exactly the same as computing softmax([1.0, 2.0, 0.5, 0.1]) and multiplying by [v₁, v₂, v₃, v₄]. You can verify: 0.368/1.741 = 0.211, 1.000/1.741 = 0.574, 0.223/1.741 = 0.128, 0.150/1.741 = 0.086. These sum to 1.0 (within rounding).

Chapter 6: The Full FlashAttention Algorithm

Now we put it all together. We've established three principles: (1) tile to fit in SRAM, (2) online softmax to combine tiles exactly, (3) never write the N×N matrix to HBM. Here's the complete forward pass.

Algorithm: FlashAttention Forward Pass

pseudocode
# Inputs: Q, K, V in HBM (each N x d)
# Block sizes: B_r, B_c chosen so tiles fit in SRAM
# Output: O in HBM (N x d)

# Initialize output O = 0, running max m = -inf, running sum l = 0
O = zeros(N, d)         # in HBM
m = full(N, -inf)       # in HBM (one per row)
l = zeros(N)             # in HBM (one per row)

# Divide K, V into T_c blocks of size B_c
# Divide Q into T_r blocks of size B_r

for j = 1 to T_c:                # outer loop over K/V blocks
    K_j, V_j = load_from_HBM(K[j], V[j])  # to SRAM

    for i = 1 to T_r:            # inner loop over Q blocks
        Q_i = load_from_HBM(Q[i])         # to SRAM
        O_i, m_i, l_i = load_from_HBM(O[i], m[i], l[i])

        # Compute tile of attention scores (in SRAM!)
        S_ij = Q_i @ K_j.T               # B_r x B_c, never leaves 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))
        O_new  = (l_i * exp(m_i - m_new) * O_i
                 + exp(S_ij - m_new) @ V_j) / l_new

        # Write updated O, m, l back to HBM
        write_to_HBM(O_new, m_new, l_new)

return O

What Gets Stored Where?

Total SRAM needed: B_r × d + 2 × B_c × d + B_r × B_c. On an A100 with ~192 KB of shared memory per SM, typical block sizes are B_r = B_c = 128 with d = 64. That's 128 × 64 × 3 × 2 bytes = 48 KB for tiles, plus 128 × 128 × 2 = 32 KB for S_ij. Total: 80 KB. Fits comfortably.

Backward Pass: Recomputation

Training requires a backward pass. Normally, the forward pass saves the N×N attention matrix for use in backpropagation. FlashAttention doesn't have it — we threw it away. The solution: recompute the attention tiles during the backward pass.

This sounds wasteful. We do more FLOPs (about 13% more). But remember: the bottleneck is memory bandwidth, not compute. By recomputing tiles in SRAM instead of reading a huge cached matrix from HBM, the backward pass is still faster.

The only things saved from the forward pass: the output O, and the per-row m and l vectors (total: 2N scalars plus the N × d output). From these, plus the original Q, K, V, we can reconstruct any tile of the attention matrix on the fly.

Chapter 7: FlashAttention-2 & Beyond

FlashAttention-1 was a breakthrough. FlashAttention-2 (Dao, 2023) made it even faster by fixing inefficiencies in how work is distributed across GPU thread blocks and warps.

FlashAttention-2 Improvements

1. Better work partitioning. FA-1's inner loop assigns each thread block to a (Q-block, K/V-block) pair. But this means different thread blocks working on the same Q-block need to communicate through slow HBM to share their partial outputs. FA-2 restructures the loops: the outer loop iterates over Q blocks, and the inner loop over K/V blocks, so each thread block fully owns one Q-block and accumulates the output in SRAM without cross-block communication.

2. Reduced non-matmul FLOPs. Modern GPUs have tensor cores that accelerate matrix multiplications dramatically (up to 16× on H100). Non-matmul operations (the rescaling, max-finding, and bookkeeping) run on regular CUDA cores, which are much slower. FA-2 reduces these non-matmul operations, keeping tensor cores busier.

3. More parallelism. FA-1 parallelizes over batch size and number of attention heads. FA-2 also parallelizes over the sequence length dimension. For long sequences with small batch sizes (common in inference), this better saturates the GPU.

FlashAttention-3

FlashAttention-3 (2024) targets the H100 GPU and exploits its new features: asynchronous memory copies (overlap compute and data movement), low-precision matmuls (FP8 tensor cores), and the TMA (Tensor Memory Accelerator) unit for efficient tiling. It also introduces software pipelining — while one tile is being computed, the next tile is being loaded from HBM to SRAM simultaneously.

Combining FlashAttention with Sparse Attention

Here's an elegant insight: FlashAttention and sparse attention are complementary. FlashAttention makes exact attention faster. Sparse attention reduces the number of tiles that need processing. Combine them: use FlashAttention for the tiles you DO compute, and skip the tiles you don't. Mistral 7B (2023) uses this approach — sliding window attention with FlashAttention — and it became one of the most efficient open-source models.

Impact on Training and Inference

FlashAttention didn't just save memory — it changed what's possible:

• BERT pre-training: 15% faster than the MLPerf 1.1 record, 3× faster than HuggingFace implementation at 1K sequence length.
• Long-Range Arena: 2.4-2.8× faster training without any quality drop.
• Path-X (16K sequences): First transformer to achieve better than random chance on this long-range task, because it could actually fit the sequences.
• Context length scaling: Enabled the jump from 2K–4K context windows (GPT-3 era) to 128K+ (GPT-4, Claude, Gemini).

Chapter 8: Connections & Cheat Sheet

FlashAttention Cheat Sheet

Key Numbers (A100)

Related Topics