FlashAttention dissected: faster computation via GPU memory

The classic implementation of the transformer attention mechanism has a notorious scaling problem: the amount of memory required grows quadratically with the length of the context. For a long time it was assumed that this slowdown was purely caused by the number of arithmetic operations (FLOPs). In practice, however, the real limiting factor turns out not to be the compute power of the GPU cores, but the speed at which data is moved to and from main memory. This hardware bottleneck is called memory bandwidth.

FlashAttention, originally developed by Tri Dao and colleagues, fundamentally restructures the computation of self-attention. Instead of continuously writing intermediate results to the GPU's slow main memory, the algorithm makes optimal use of the small, ultra-fast SRAM chips located directly next to the compute cores. In this article, we break down step by step how FlashAttention works, which mathematical techniques make this possible, and what the practical implications are for training and inference time.

The physical bottleneck: GPU memory hierarchy dissected

To understand why standard attention is slow for long sequences, we need to look at the physical architecture of modern graphics cards such as the NVIDIA A100 or H100. A modern GPU roughly contains two types of memory: the large High Bandwidth Memory (HBM, also called VRAM) and the internal Static RAM (SRAM, distributed across the Streaming Multiprocessors as Shared Memory and L1 cache).

HBM is generously sized (for example 80 GB on an A100), but the throughput is around 1.5 to 2.0 TB/s. That sounds enormous, but the compute cores (Tensor Cores) can collectively process hundreds of teraflops of data. The internal SRAM is many times smaller (roughly 192 KB per Streaming Multiprocessor, which across the entire chip amounts to a few tens of megabytes), but achieves a bandwidth of more than 19 TB/s. When a GPU kernel fetches data from HBM, the compute units are effectively waiting for data. We call such an operation memory-bound.

Standard attention performs sequential operations: it first multiplies the Query matrix $Q$ by the transposed Key matrix $K^T$, writes the result $S$ (the $N \times N$ attention matrix) to HBM, reads this matrix back in to apply the softmax function to matrix $P$, writes $P$ back to HBM, and finally reads $P$ back in to multiply it by the Value matrix $V$. At a context of 8,192 tokens, the $N \times N$ matrix contains over 67 million floating-point numbers per attention head. This leads to gigantic amounts of unnecessary read and write traffic over the slower HBM bus.

The mathematical breakthrough: Tiling and Online Softmax

The central innovation of FlashAttention is tiling: splitting the large $Q$, $K$, and $V$ matrices into smaller blocks that fit exactly within the GPU's fast SRAM. The problem here, however, is the softmax function. Softmax normalizes values over an entire row using the formula:

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

To compute the denominator (the sum of exponents), it seems as though we first need to know the complete row. After all, if we chop a row into pieces, the global sum is missing. FlashAttention solves this by using the principle of online softmax. Here, intermediate results are computed per block and dynamically rescaled as soon as a new block is processed.

Suppose we have computed a partial sum over the first block of tokens with a local maximum $m_1$ and a local denominator $l_1$. When we load the second block with a local maximum $m_2$, we determine the new global maximum $m_{\text{new}} = \max(m_1, m_2)$. We can simply rescale the old sum and the accumulated intermediate vector using the correction factor $e^{m_1 - m_{\text{new}}}$. As a result, it is never necessary to store the full $N \times N$ attention matrix in global HBM memory.

FlashAttention-1 versus FlashAttention-2 and later iterations

FlashAttention has undergone significant iterative improvements since its first publication. In the first version (FlashAttention-1), the outer loop was placed over the key and value blocks ($K$ and $V$) and the inner loop over the query blocks ($Q$). This caused extra synchronization steps between different GPU threads.

FlashAttention-2 reversed this logic: the outer loop now iterates over the rows of the $Q$ matrix. Because different thread blocks can now work independently on different rows, much less communication and synchronization between the Streaming Multiprocessors is needed. In addition, the number of non-matmul operations was minimized, which raised the occupancy of the Tensor Cores from roughly 35% to more than 70% of the theoretical maximum.

FlashAttention-3 specifically targets more modern hardware such as NVIDIA Hopper chips. These chips introduce the Tensor Memory Accelerator (TMA) and asynchronous instructions, which allow data to be copied directly from HBM to shared memory without involving the regular compute registers. This allows data transfer and matrix computations to happen simultaneously without delay.

Memory complexity: from O(N²) to O(N)

In a standard transformer architecture, storing the activations during the forward pass for attention mechanisms requires $O(N^2)$ memory space, where $N$ is the context length in tokens. When training a model, these activations must be kept in HBM in order to compute the gradients during the backward pass.

FlashAttention reduces this memory footprint to $O(N)$. Because the $N \times N$ attention matrix is never fully constructed in HBM, it also doesn't need to be stored for the backward pass. Instead, FlashAttention very quickly recomputes the needed attention blocks during the backward pass (recomputation) from the stored vectors $Q$, $K$, and $V$ combined with the saved softmax statistics (the maximum $m$ and the denominator $l$).

Although recomputation costs extra arithmetic operations, in practice it is much faster than loading gigantic matrices from HBM. Eliminating the memory traffic more than compensates for the extra computation steps.

Interaction with the KV cache during inference

During the inference phase, when a model generates a response token by token, the role of the attention mechanism changes. Here no gradients need to be kept, but earlier tokens do need to be consulted quickly. To understand how earlier input is retained, take a look at the structure of KV caching in transformer architectures, which covers the allocation of keys and values per layer.

When generating a new token, the Query vector has a length of 1 ($N_q = 1$), while the Key and Value matrices grow to the total historical length of the conversation ($N_{kv}$). In this situation we speak of FlashDecoding, a specialized variant of FlashAttention for autoregressive inference. FlashDecoding splits the $N_{kv}$ sequence across multiple GPU cores to maintain maximum parallel occupancy even at $N_q=1$.

FlashAttention works seamlessly together with modern architectures that reduce the memory pressure of the KV cache. For a deeper analysis of this memory-saving technique, also see the article on grouped-query attention and memory usage, in which multiple query heads share the same key and value heads.

Algorithmic core in pseudocode

The simplified representation below shows the core concept of the FlashAttention forward pass with block-wise tiling and online rescaling of the softmax values.

# Invoer: Q, K, V matrices in HBM (formaat N x d)
# Blokgroottes: B_r (rijen van Q), B_c (kolommen van K, V)
# Uitvoer: O matrix in HBM (formaat N x d)

initialiseer O in HBM met nullen
initialiseer l = [0] * N  (som van exponenten)
initialiseer m = [-oneindig] * N  (maxima per rij)

splits Q in blokken Q_1, ..., Q_Tr van grootte B_r x d
splits K in blokken K_1, ..., K_Tc van grootte B_c x d
splits V in blokken V_1, ..., V_Tc van grootte B_c x d

voor elk blok Q_i in SRAM:
  voor elk blok K_j, V_j in SRAM:
    # 1. Bereken lokale matrixvermenigvuldiging
    S_ij = (Q_i * K_j^T) / sqrt(d)
    
    # 2. Bereken lokaal rij-maximum
    m_lokaal = max_per_rij(S_ij)
    m_nieuw = max(m_i, m_lokaal)
    
    # 3. Bereken herschaalde exponenten
    P_ij = exp(S_ij - m_nieuw)
    l_nieuw = exp(m_i - m_nieuw) * l_i + som_per_rij(P_ij)
    
    # 4. Werk de uitvoer O_i bij met correctiefactoren
    O_i = diag(exp(m_i - m_nieuw)) * O_i + P_ij * V_j
    
    # 5. Werk statistieken bij voor volgende iteratie
    m_i = m_nieuw
    l_i = l_nieuw

  # Normaliseer uiteindelijke rij met de totale noemer l_i
  O_i = diag(1 / l_i) * O_i
  schrijf O_i weg naar HBM

Practical integration and hardware requirements

FlashAttention is a low-level C++/CUDA implementation and is not available by default on all architectures. The implementation requires specific hardware instructions that are only present on more modern GPUs.

On NVIDIA hardware, FlashAttention-2 is supported from the Turing and Ampere architectures onward (for example the RTX 30 series, RTX 40 series, A10, A100, and newer with Compute Capability 8.0 or higher). Older architectures such as Pascal (GTX 1080) lack the instruction sets to drive shared memory asynchronously in this way. When setting up a local system or a cluster server, it is crucial to check in advance whether the hardware and context length fit within the limits of the graphics card. Consult the VRAM calculator: how much memory does a model need? to calculate precisely how model parameters and context lengths relate to the available GPU memory.

Within modern deep learning frameworks such as PyTorch, FlashAttention is often already integrated via torch.nn.functional.scaled_dot_product_attention (SDPA). PyTorch automatically selects the FlashAttention backend when the input tensors meet the right requirements (such as FP16 or BF16 data types and suitable tensor shapes).

Limitations, weaknesses, and edge cases

Although FlashAttention offers considerable advantages, there are clear technical limitations and trade-offs:

First, FlashAttention offers little to no speed gain for very short sequences (for example fewer than 256 tokens). For such sequences, the $N \times N$ matrix is so small that it already fits entirely in the L2 cache or SRAM, so the overhead of splitting into blocks doesn't outweigh a standard matrix operation.

Second, FlashAttention is hardware-specific. Because the kernel is precisely tuned to the exact register sizes, shared memory layouts, and cache structures of specific GPU generations, support for other hardware (such as AMD ROCm or Apple Silicon) requires completely rewritten kernels. Although projects such as FlashAttention for ROCm exist, these implementations often lag behind the official NVIDIA versions.

Third, there can be slight numerical rounding differences. Although FlashAttention mathematically computes exactly the same formula as standard attention (unlike approximation methods such as sparse attention or low-rank projections), the changed order of floating-point calculations can produce subtly different outcomes due to rounding in FP16. In practice this does not lead to a loss of quality, but it does mean that test results between different backends are not bit-for-bit identical.