Flash Attention Principles

The Memory Bottleneck in Standard Attention

Before we dive into the code implementation of Flash Attention, we must first answer a foundational question: why is the standard attention formula $Softmax(QK^T)V$ still not fast enough on modern GPUs?

GPU Memory Hierarchy: SRAM and HBM

First, we need to establish an extremely important concept: in the GPU architecture, all computation (such as matrix addition and multiplication) must happen in the SRAM (Shared Memory) that sits close to the cores.

This means that even if your VRAM (HBM) is as large as 80GB, the data must first be "moved" into the SRAM, which is only tens of MB in size, before it can be processed by the compute cores.

The Logical Trap in the Standard Implementation

In the "naive" implementations of deep learning frameworks such as PyTorch, the Attention computation is split into multiple independent operators (Ops). This leads to a serious efficiency problem:

1. Step 1 ($QK^T$): the GPU moves $Q$ and $K$ from HBM to SRAM and computes the score matrix $S$.
2. The intermediate result is too large: since the shape of $S$ is $(N, N)$, for long sequences this matrix is so large that SRAM simply cannot hold it.
3. Forced write-back: the GPU has no choice but to "evict" this giant matrix $S$ from SRAM and write it back to the slow HBM.
4. Repeated back-and-forth: when it comes to the next step of computing $Softmax(S)$, the GPU has to go back to HBM and move the $S$ it just stored back into SRAM again.

# The IO nightmare of the standard Attention implementation
def standard_attention(Q, K, V):
    # 1. HBM -> SRAM (compute) -> HBM (store S)
    S = Q @ K.T

    # 2. HBM (read S) -> SRAM (compute) -> HBM (store P)
    P = softmax(S)

    # 3. HBM (read P) -> SRAM (compute) -> HBM (store O)
    O = P @ V
    return O

This repeated I/O round trip of "move in -> compute -> evict -> move back" is the biggest performance killer.

The Bandwidth Gap Between SRAM and HBM

You might ask: how much impact can storing back to HBM and reading it back again really have?

Comparison of Speed Differences

SRAM's bandwidth is typically more than about 10 times higher than HBM's.

The Essence of the Bottleneck: IO-bound

Because of this speed gulf, if the algorithm keeps moving intermediate data back and forth between HBM and SRAM, an awkward situation arises:

This state is called IO-bound, meaning that the compute capability is held back by the memory transfer speed.

The Capacity Limit of SRAM

Since SRAM is so fast, why not just put the entire Attention matrix in SRAM, finish computing, and be done? This is where rigid physical and economic constraints come into play:

Physical Constraints and Cost

SRAM has an extremely low storage density, which makes its cost extremely high. According to relevant materials (such as the background cited in the FlashAttention paper):

• Cost: manufacturing an 80GB capacity SRAM memory could cost as much as $13,000 (an order-of-magnitude estimate).
• Comparison: an HBM of the same capacity costs only about $2,000.

The Capacity Limit

In real hardware, an A100's HBM can reach 80GB, but SRAM is usually only 192 KB / SM (per streaming multiprocessor). Because of this capacity limit, you cannot cram the entire $N \times N$ attention matrix into SRAM all at once. As the sequence length $N$ grows, the size of the intermediate matrix explodes at $O(N^2)$.

The Core Idea: Optimizing IO Complexity

The core logic of Flash Attention lies right here: since SRAM is expensive and small but fast, while HBM is large and cheap but slow, we must abandon the fantasy of "reading and writing everything at once."

We need to introduce two core ideas:

1. Tiling: split the data into "small blocks" that can fit into SRAM.
2. Kernel Fusion: inside SRAM, complete the $QK^T$, Softmax, and multiplication by $V$ in one go, writing only the final result $O$ back to HBM in the last step.

Avoiding Spilling the Intermediate Matrix to Memory

In this way, we never even generate (nor write to HBM) that huge $N \times N$ intermediate matrix.