FlashAttention Implementation: GPU Memory Optimization for Transformers

FlashAttention changes how we think about attention computation. Rather than viewing the attention formula as a mathematical specification to implement directly, FlashAttention treats it as a memory optimization problem. The algorithm produces identical outputs to standard attention but restructures the computation to minimize memory transfers between GPU memory hierarchies. Understanding how FlashAttention achieves this requires diving into GPU architecture, memory access patterns, and the algorithmic techniques that make it all work.

This chapter explores the implementation details of FlashAttention, from the fundamentals of GPU memory to practical usage in PyTorch. We'll examine CUDA kernel basics, understand why memory access patterns matter so much for performance, trace the improvements from FlashAttention to FlashAttention-2, and see how to use these techniques in your own models. By the end, you'll understand both the theoretical foundations and practical applications of what has become the de facto standard for attention computation in modern LLMs.

Before understanding FlashAttention, you need to understand where bottlenecks actually occur on GPUs. Modern deep learning is often described as "compute-bound," but attention computation is actually "memory-bound." This distinction is crucial.

GPUs have a hierarchical memory structure with three main levels, each with different speeds and capacities:

The numbers reveal a fundamental tension. HBM (the main GPU memory where tensors live) has enormous capacity but relatively slow access. SRAM (registers and shared memory) is much faster but tiny. The speed difference is nearly 10x, while the capacity difference is over 400,000x. Standard attention computation repeatedly reads and writes to HBM, suffering the slow access speeds. FlashAttention restructures the computation to keep data in SRAM as long as possible.

Consider what happens during standard attention computation:

1. Load Q, K, V from HBM to compute $S = QK^T$
2. Write S back to HBM (this is an $n \times n$ matrix)
3. Load S from HBM to compute softmax
4. Write attention weights P back to HBM
5. Load P and V from HBM to compute output
6. Write output back to HBM

Each step involves transfers to and from the slow HBM. The $n \times n$ matrices S and P are particularly problematic: for a sequence of 4096 tokens with 32-bit floats, each matrix is 64 MB. With 32 attention heads, that's 2 GB just for intermediate storage, per layer, per batch element.

The intermediate matrices consume over 128 MB of memory transfers, dwarfing the input and output. This is where FlashAttention intervenes: by never materializing the full $n \times n$ matrices, it eliminates most of this memory traffic.

FlashAttention is implemented as a custom CUDA kernel, meaning it runs directly on GPU hardware with fine-grained control over memory access. Understanding a few CUDA concepts helps appreciate what FlashAttention does under the hood.

CUDA organizes computation into a hierarchy. At the lowest level, individual threads execute the same code on different data. Threads are grouped into blocks, typically containing 256-1024 threads. All threads in a block can communicate through shared memory and synchronize their execution. Blocks are grouped into a grid that covers the entire problem.

For attention, a natural mapping assigns one block to compute attention for one query block against all key-value blocks. Within a block, threads collaborate to load data into shared memory, compute dot products, and accumulate results.

The key to GPU performance is coalesced memory access. When threads in a warp (a group of 32 threads that execute in lockstep) access consecutive memory addresses, the hardware combines these into a single efficient transaction. When threads access scattered addresses, each access becomes a separate transaction, significantly slowing execution.

FlashAttention carefully structures its memory access to achieve coalescing. When loading a block of K from HBM, consecutive threads load consecutive elements. When writing results, the same pattern applies. This attention to memory layout is one reason FlashAttention achieves such high hardware utilization.

We've established that standard attention is memory-bound: the algorithm spends most of its time moving data between slow HBM and fast SRAM, not actually computing. 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 actually 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:

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$, ensuring 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:

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:

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:

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:

The output traces the algorithm's journey through three blocks. Notice how the running maximum $m$ updates whenever a block contains a larger value, triggering a correction factor less than 1. The running sum $l$ grows as we accumulate more exponentials. At the end, we verify that our online computation produces exactly the same sum as computing softmax over all scores at once.

The match confirms a crucial property: online softmax is not an approximation. It computes the exact same result as standard softmax, just in a different order that avoids storing all scores simultaneously.

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:

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:

The table reveals the dramatic memory savings FlashAttention provides. At 512 tokens, standard attention requires only 0.5 MB for the attention matrix, making FlashAttention's overhead relatively larger. But as sequence length doubles, standard attention memory quadruples while FlashAttention grows linearly. By 16K tokens, standard attention consumes over 500 MB just for the attention matrix, while FlashAttention uses only about 2 MB, a reduction of over 250x. This difference determines whether a model fits in GPU memory at all.

FlashAttention-2, released in 2023, builds on the original algorithm with several key optimizations that improve hardware utilization from ~50-70% to ~70-90% of theoretical peak performance.

The original FlashAttention spent significant time on non-matrix-multiply operations: rescaling outputs, updating statistics, and coordinating between iterations. FlashAttention-2 restructures the algorithm to minimize these overheads.

The key change is delaying the rescaling of output blocks. Instead of rescaling after each K-V block, FlashAttention-2 accumulates unnormalized outputs and applies a single rescaling at the end. This requires careful bookkeeping but reduces the number of scalar operations significantly.

With 64 K-V blocks to process, FlashAttention-1 performs 4,096 rescale operations (one per query per block), while FlashAttention-2 performs only 64 (one per query at the end). This 64x reduction in rescaling operations translates directly to improved throughput, particularly when matrix multiplies are already saturating the GPU's compute units.

FlashAttention-2 improves how work is distributed across thread blocks and warps:

• Parallelism over batch and heads: Each thread block handles one (batch, head) pair, maximizing parallelism across the GPU
• Within-block parallelism: Different warps handle different parts of the Q block, with careful synchronization
• Sequence parallelism: For very long sequences, work can be split across multiple blocks even for a single head

FlashAttention-2 optimizes memory access patterns in several ways:

• Swapping loops: The inner and outer loops are reordered so that the outer loop iterates over Q blocks and the inner loop over K-V blocks. This reduces the number of times Q blocks are loaded from HBM.
• Pipelining: Memory loads are overlapped with computation. While computing attention for the current K-V block, the next block is being loaded from HBM.
• Warp-level primitives: Instead of synchronizing entire thread blocks, FlashAttention-2 uses warp-level operations that are faster and more efficient.

By swapping the loop order, FlashAttention-2 reads the Q matrix 64x fewer times from HBM. This is the dominant source of the memory access reduction. Since memory bandwidth is the bottleneck, reducing redundant reads directly improves throughput.

These optimizations combine to give FlashAttention-2 a 2x speedup over FlashAttention on many workloads:

The benchmarks show a clear trend: FlashAttention-2's advantage grows with sequence length. At 1K tokens, the speedup is a modest 1.49x. At 16K tokens, it reaches 2.30x. This pattern occurs because longer sequences have more K-V blocks, amplifying the benefits of reduced rescaling operations and optimized memory access. FlashAttention-2 also achieves higher absolute TFLOPS, indicating better GPU utilization.

The easiest way to use FlashAttention is through PyTorch's native support, added in PyTorch 2.0. The scaled_dot_product_attention function automatically uses FlashAttention when conditions permit.

The output shape matches the input dimensions, confirming that attention produces one output vector per query position. When running on a CUDA-enabled GPU with fp16 tensors, PyTorch automatically selects FlashAttention as the backend, providing the memory and speed benefits without any code changes.

PyTorch allows explicit control over which attention implementation to use:

FlashAttention has specific requirements that must be met for it to be used:

These requirements explain when PyTorch will automatically select FlashAttention. The most common reason for fallback is using float32 instead of float16/bfloat16, or using an older GPU. If you're not getting FlashAttention's benefits, check these requirements first.

For more control, you can use the flash-attn library directly:

Here's how FlashAttention integrates into a typical transformer layer:

The layer uses approximately 1 million parameters, primarily in the four linear projections (Q, K, V, and output). The head dimension of 64 is optimal for FlashAttention's performance. This implementation pattern is directly usable in production transformer models.

Despite its advantages, FlashAttention has limitations that you should understand before deployment.

FlashAttention computes attention in a different order than standard attention, which can lead to small numerical differences. While mathematically equivalent, floating-point arithmetic is not associative, so results may differ at the level of floating-point precision.

Running the same experiment 1000 times with different random values shows that these differences follow a predictable pattern:

FlashAttention's tiled computation assumes specific attention patterns. While causal masks and no masks are well-supported, arbitrary attention masks present challenges:

For most transformer applications, causal or no mask covers the use cases. FlashAttention-2 added sliding window support, which enables models like Mistral. For arbitrary boolean masks or complex attention biases, PyTorch falls back to the slower math kernel, losing FlashAttention's benefits.

FlashAttention requires relatively recent NVIDIA GPUs:

Ampere and newer architectures provide full FlashAttention-2 support. Older GPUs like V100 and T4 have limited support for FlashAttention-1 only. If you're using cloud instances, A100 and H100 are the best choices for FlashAttention workloads. Consumer GPUs from the RTX 30 and 40 series also work well.

While FlashAttention eliminates the $O(n^2)$ memory for attention matrices, other parts of the transformer still scale with sequence length:

The breakdown shows that at 4K tokens, a 7B model requires about 13 GB total. At 32K tokens, this grows to nearly 50 GB, a 4x increase. Model weights remain constant, but activations and KV cache scale linearly with sequence length. This explains why even with FlashAttention, very long sequences require substantial GPU memory or additional techniques like gradient checkpointing.

FlashAttention's recomputation strategy during the backward pass means gradients take longer to compute than they would with stored attention matrices. This is typically worthwhile because:

1. The memory savings allow larger batch sizes, which often more than compensate for the recomputation cost
2. Training is usually memory-bound anyway, so the extra compute fits within memory transfer time
3. For very long sequences, the alternative (storing attention matrices) would be impossible

FlashAttention transforms attention computation from a memory-bound operation into one that achieves near-peak hardware utilization. The key insights and takeaways from this chapter:

• GPU memory hierarchy: Understanding the dramatic speed difference between HBM and SRAM (10x) is essential. Standard attention repeatedly transfers data to slow HBM, while FlashAttention keeps data in fast SRAM.

• Tiled computation: By processing attention in small blocks that fit in SRAM, FlashAttention never materializes the full $n \times n$ attention matrix. This reduces memory complexity from $O(n^2)$ to $O(n)$.

• Online softmax: The online softmax algorithm maintains running statistics (maximum and sum of exponentials) that allow correct normalization without seeing all scores at once. Previous accumulations are corrected as new tiles reveal larger values.

• FlashAttention-2 improvements: Reduced non-matrix-multiply operations, better work partitioning across threads and warps, and improved memory access patterns combine to give approximately 2x speedup over FlashAttention-1.

• PyTorch integration: The easiest path to using FlashAttention is through torch.nn.functional.scaled_dot_product_attention, which automatically selects FlashAttention when conditions permit (CUDA, fp16/bf16, supported head dimensions).

• Limitations: FlashAttention requires recent NVIDIA GPUs (Ampere or newer for full support), has constraints on attention mask patterns, and may produce results that differ from standard attention at floating-point precision. While it eliminates the attention matrix memory cost, activations and KV cache still scale linearly with sequence length.

FlashAttention has become the de facto standard for attention computation in modern LLMs. Models like LLaMA 2, Mistral, and most recent transformers use it by default. Understanding its implementation details helps you appreciate why it works, when it applies, and what to expect when deploying models that rely on it.

When using FlashAttention, several parameters directly affect behavior and performance:

• block_size (B_r, B_c): The tile dimensions for processing attention. Typically 64-128, determined automatically based on head dimension and SRAM availability. Larger blocks amortize memory transfer overhead but require more SRAM.

• head_dim (d): The dimension of each attention head. FlashAttention supports specific values: 8, 16, 32, 64, 128, 256. Performance is best at 64 or 128.

• causal: Whether to use causal (autoregressive) attention masking. Causal attention is optimized in FlashAttention to skip computation for masked positions entirely, not just mask them after computation.

• dropout_p: Dropout probability during training. FlashAttention implements dropout within the kernel, avoiding the need to store dropout masks in HBM.

• softmax_scale: Scaling factor applied before softmax, typically $1/\sqrt{d}$ where $d$ is the head dimension. This prevents dot products from growing too large with higher dimensions, which would push softmax outputs toward one-hot distributions. Can be customized for specialized attention variants.

• dtype: Data type for computation. fp16 and bf16 are fully supported and recommended. fp32 falls back to non-FlashAttention kernels in PyTorch.

• window_size: For FlashAttention-2's sliding window support, the size of the local attention window. Set to (-1, -1) for full attention or specific values like (256, 256) for symmetric windows.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about FlashAttention and GPU memory optimization.