FlashAttention Implementation › How FlashAttention Uses Tiling

We've established that standard attention is memory-bound: the algorithm spends most of its time moving data between slow HBM and fast SRAM, rather than performing arithmetic. The $n \times n$ attention matrix is the culprit, requiring multiple read-write cycles to HBM. FlashAttention's key insight starts with a simple question: what if we never stored that matrix at all?

The answer lies in a technique called tiling. Rather than computing all attention scores at once and then applying softmax, FlashAttention processes attention in small chunks that fit entirely in SRAM. Each chunk computes a portion of the final output, and careful bookkeeping ensures the pieces combine to produce exactly the same result as standard attention.

To understand how tiling works, imagine you're computing attention for a 4096-token sequence. Standard attention would:

1. Compute all $4096 \times 4096 = 16.7$ million attention scores
2. Store this entire matrix in memory
3. Apply softmax across each row
4. Multiply by the value matrix

FlashAttention takes a different approach. Instead of processing all tokens at once, it divides the computation into manageable blocks:

• Query blocks: Groups of $B_r$ consecutive query positions (rows of Q)
• Key-Value blocks: Groups of $B_c$ consecutive key-value positions (rows of K and V)

For each query block, the algorithm iterates through all key-value blocks, computing partial attention scores and accumulating results. The block sizes $B_r$ and $B_c$ are chosen specifically so that all active data fits in SRAM at once.

With 192 KB of SRAM available, we can fit query and key-value blocks of size 128, plus room for intermediate computations. The critical observation is that these block sizes are independent of sequence length: whether processing 1K or 100K tokens, each tile uses the same fixed amount of fast memory.

But there's a catch. Tiling the attention computation creates a fundamental problem that lies at the heart of FlashAttention's innovation.

Here's the challenge: softmax requires global information, but we're processing local tiles.

Consider what softmax does. It converts a vector of raw scores into a probability distribution where all values sum to 1. Given a vector of attention scores, standard softmax computes each output element as:

$$\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_{j=1}^{n} e^{x_j}}$$

where:

• $x_i$: the $i$-th attention score (the value we want to convert to a probability)
• $e^{x_i}$: the exponential of $x_i$. This keeps all values become positive
• $\sum_{j=1}^{n} e^{x_j}$: the sum of exponentials across all $n$ scores, serving as the normalizing constant
• $n$: the total number of scores (equal to sequence length in attention)

The denominator is the problem. To compute the softmax output for any single element, we need $\sum_{j=1}^{n} e^{x_j}$, the sum of exponentials across all $n$ scores. But when processing tiles, we don't have all scores at once. We see positions 1-64, then 65-128, then 129-192, and so on. How can we compute a global normalization constant from local pieces?

FlashAttention solves this with the online softmax algorithm, a technique that maintains running statistics and retroactively corrects previous computations as new information arrives. The key insight is that softmax can be decomposed and incrementally refined.

Building Intuition: Two BlocksLink Copied

Let's start with the simplest case. Imagine we have just two blocks of scores, $x$ and $y$. The full softmax over both is:

$$\text{softmax}([x, y]) = \left[\frac{e^{x}}{e^x + e^y}, \frac{e^y}{e^x + e^y}\right]$$

Now suppose we process $x$ first, before $y$ is available. At that moment, $x$ is the only score we know, so its "local softmax" is $e^x / e^x = 1$, a probability of 100%. This is wrong, of course, but we didn't have complete information.

When $y$ arrives, we can fix our mistake. The correct probability for $x$ should be $e^x / (e^x + e^y)$. We can get this by multiplying our initial answer (1) by the correction factor:

$$\text{correction} = \frac{e^x}{e^x + e^y}$$

This factor adjusts for the new, larger denominator. The pattern generalizes: as each new block arrives, we correct all previous accumulations by the ratio of old denominator to new denominator.

The Running StatisticsLink Copied

To make this correction efficient, FlashAttention tracks two running statistics for each query position:

• $m$: the running maximum score seen so far
• $l$: the running sum of exponentials (shifted by the maximum for numerical stability)

Why track the maximum? Exponentials grow quickly: $e^{100}$ is extremely large while $e^{-100}$ is essentially zero. Without care, we'd overflow or underflow. By subtracting the maximum before exponentiating, we ensure all values remain in a reasonable range. This is the same trick used in numerically stable softmax implementations.

As each new tile of scores $\mathbf{s}$ arrives, FlashAttention updates these statistics:

$$m_{\text{new}} = \max(m_{\text{old}}, \max(\mathbf{s}))$$ $$l_{\text{new}} = e^{m_{\text{old}} - m_{\text{new}}} \cdot l_{\text{old}} + \sum_i e^{s_i - m_{\text{new}}}$$

where:

• $m_{\text{old}}$: the running maximum from previous blocks
• $m_{\text{new}}$: the updated maximum after seeing the current block
• $l_{\text{old}}$: the running sum of exponentials from previous blocks
• $l_{\text{new}}$: the updated sum after incorporating the current block
• $\mathbf{s}$: the vector of attention scores in the current block
• $e^{m_{\text{old}} - m_{\text{new}}}$: the correction factor that rescales previous accumulations

Why the Correction Factor WorksLink Copied

The correction factor $e^{m_{\text{old}} - m_{\text{new}}}$ is the mathematical heart of online softmax. Let's understand why it works.

When the new block contains a score larger than any we've seen before, $m_{\text{new}} > m_{\text{old}}$. The difference $m_{\text{old}} - m_{\text{new}}$ is negative, so $e^{m_{\text{old}} - m_{\text{new}}} < 1$. This makes sense: we're discovering that the denominator should be larger than we thought, so previous contributions need to be scaled down.

When the new block has no scores exceeding our current maximum, $m_{\text{new}} = m_{\text{old}}$. The correction factor becomes $e^0 = 1$, leaving previous accumulations unchanged. We're simply adding to the sum without revising history.

Importantly, corrections compound correctly. If we process blocks A, then B, then C, the final answer is identical to processing them all at once. The mathematics guarantees exact equivalence, not an approximation.

Seeing It in ActionLink Copied

Let's trace through the online softmax algorithm on a concrete example. We'll process three blocks of random scores and watch the running statistics evolve:

We've now seen the two key ingredients of FlashAttention: tiling to process attention in SRAM-sized chunks, and online softmax to correctly combine partial results. Together, they enable exact attention computation without ever storing the $n \times n$ matrix.

Let's inventory exactly what FlashAttention keeps in memory at any moment during computation:

For typical settings with $B_r = B_c = 64$ and $d = 64$, this totals about 25 KB, fitting comfortably in SRAM. The memory complexity drops from $O(n^2)$ in standard attention to $O(n)$ in FlashAttention. More precisely:

$$M_{\text{flash}} = O(n \cdot d + B_r \cdot B_c)$$

where:

• $n$: the sequence length
• $d$: the head dimension (typically 64-128)
• $B_r$: the block size for queries (number of query rows processed together)
• $B_c$: the block size for keys/values (number of key-value rows loaded at once)
• $n \cdot d$: memory for the output accumulator (one $d$-dimensional vector per query)
• $B_r \cdot B_c$: memory for the current score tile being computed

Since $B_r$, $B_c$, and $d$ are constants determined by hardware constraints (not sequence length), the dominant term is $n \cdot d$, giving linear scaling with sequence length.

This is the payoff of the entire approach. Standard attention requires $O(n^2)$ memory for the attention matrix, limiting sequence length to what fits in GPU memory. FlashAttention's $O(n)$ memory means you can process arbitrarily long sequences, with memory growing only linearly. A 10x longer sequence needs only 10x more memory, not 100x.

Let's verify this with concrete numbers: