Tracing Tiled QKV Data Movement in FlashAttention’s Memory Hierarchy

1. The Memory Bottleneck in Standard Attention

Before we dive into the solution, let's briefly establish the problem FlashAttention was designed to solve. The standard self-attention mechanism, while powerful, is notoriously memory-intensive. For a sequence of length $N$, it requires the explicit computation and storage of an $N\times N$ attention matrix. As $N$ grows, this matrix becomes a major bottleneck, not because of the computations themselves, but because of the repeated reads and writes to and from the GPU's main memory.

On a GPU, there are two key types of memory to consider:

• High Bandwidth Memory (HBM): This is the large (e.g., 40-80 GB) memory we typically think of as "GPU memory." It has high bandwidth, but access is significantly slower compared to on-chip memory.
• Static Random-Access Memory (SRAM): This is a much smaller (e.g., ~20 MB) but extremely fast memory located directly on the GPU's streaming multiprocessors (the compute units).

The core issue with standard attention is the need to shuttle the large intermediate attention matrix back and forth between the compute units and HBM. FlashAttention's central insight is to minimize these slow HBM accesses by maximizing the work done in the fast SRAM.

To get a first-hand account of this problem and the high-level solution, let's watch a short segment from a talk by Tri Dao, one of the creators of FlashAttention.

2. Tiling and Online Softmax: The Solution

To avoid materializing the full $N\times N$ attention matrix in HBM, FlashAttention employs a technique called tiling. The large Query (Q), Key (K), and Value (V) matrices are partitioned into smaller blocks, or "tiles," that are small enough to be loaded into the fast SRAM.

However, simply processing blocks independently is not straightforward because of the softmax operation. The softmax of a vector depends on all elements in that vector for normalization. You cannot compute the correct softmax for the first half of a vector without seeing the second half.

FlashAttention overcomes this with a clever mathematical trick often called the "online softmax." It restructures the computation to allow for the incremental calculation of attention. In essence, it processes one block at a time, computes a partial result, and maintains running statistics (the maximum value seen so far and the cumulative sum for the denominator). When the next block is processed, it uses these statistics to update the partial result correctly.

This combination of tiling and online softmax allows all the intermediate steps of the attention calculation—the $Q{K}^T$ dot product, scaling, and softmax—to be "fused" into a single GPU kernel that operates on tiles within SRAM. This drastically reduces the data traffic to and from HBM.

3. Tracing the Dataflow of a Forward Pass

Now, let's trace the movement of data using a diagram that illustrates the FlashAttention forward pass. This diagram shows which data resides in the large, slow HBM and which computations are performed on-the-fly in the fast on-chip SRAM.

Let's break down the process step-by-step, following the data:

1. Initialization: The complete Q, K, and V matrices start in HBM. The output matrix O is also initialized to zeros in HBM.

2. Outer Loop - Load K and V blocks: The algorithm iterates through the K and V matrices in blocks. In the diagram, it first loads the block $K^{(1)}$ and its corresponding value block $V^{(1)}$ from HBM into SRAM.

3. Inner Loop - Load Q block: The algorithm then begins an inner loop over the blocks of the Q matrix. It loads a block of Q (labeled simply as $Q$ in this simplified diagram, but think of it as one tile $Q^{(i)}$) from HBM into SRAM.

4. On-Chip Computation (in SRAM): With blocks of Q, K, and V now in fast SRAM, a flurry of computation happens without any HBM access:

   • The score matrix block is computed: $S^{(1)}=Q(K^{(1)}{)}^T$. Notice this is in an orange dashed box, meaning it's an intermediate result that is not written back to HBM.
   • The unnormalized attention weights are calculated: $A^{(1)}=\exp (S^{(1)})$. This, too, remains in SRAM.
   • A running statistic for the softmax denominator, $l^{(1)}$, is computed by summing the elements of $A^{(1)}$.
   • A partial output $O^{(1)}$ is computed by multiplying the locally normalized attention weights ($A^{(1)}/l^{(1)}$) by the value block $V^{(1)}$. This partial output is written back to HBM.

5. Iteration and Online Update: The algorithm now moves to the next block of K and V, loading $K^{(2)}$ and $V^{(2)}$ into SRAM. The same block of Q is used.

   • New scores $S^{(2)}$ and weights $A^{(2)}$ are computed in SRAM.
   • The crucial "online" step happens here. A new denominator statistic, $l^{(2)}$, is calculated, which incorporates information from both the first and second blocks.
   • The final output for this Q block, $O^{(2)}$, is computed. Notice how it's not a simple sum. The previous output $O^{(1)}$ is rescaled by $\frac{l^{(1)}}{l^{(2)}}$ to account for the new, correct normalization factor. The contribution from the current block, $\frac{A^{(2)}}{l^{(2)}}{V}^{(2)}$, is then added.
   • This updated output block $O^{(2)}$ is written back to HBM, overwriting the previous partial result.

This process continues for all blocks of K and V. Once that is done, the inner loop moves to the next block of Q, and the entire process repeats. The key is that the large $S$ and $A$ matrices are never fully formed in HBM.

To see this process animated, the following video provides a clear visualization that maps directly to the steps we just discussed.

Finally, to connect this visual understanding to the formal algorithm, you can look at the pseudocode from the original paper.

Conclusion

In this lesson, we have dissected the IO-aware dataflow of FlashAttention. By tracing the movement of data tiles between HBM and SRAM, you can now see precisely how this algorithm achieves its remarkable efficiency gains.

The key takeaways are:

• Problem: Standard attention is bottlenecked by memory I/O, specifically reading and writing the large $N\times N$ attention matrix to and from slow HBM.
• Solution: FlashAttention uses tiling to break matrices into blocks that fit in fast SRAM and an online softmax algorithm to process these blocks incrementally.
• Dataflow: Blocks of K and V are loaded into SRAM in an outer loop, and blocks of Q are processed in an inner loop. All intermediate calculations happen on-chip, and only the final, updated output block is written back to HBM at each step.

This principle of designing algorithms to be "IO-aware"—conscious of the memory hierarchy—is a cornerstone of high-performance computing and is what makes training and inference with very long sequence lengths feasible in modern LLMs. The same principles are applied to create an optimized backward pass, which uses recomputation to avoid storing the attention matrix, further enhancing memory savings.