FlashAttention

Transformers have become the dominant architecture across natural language processing, computer vision, and multimodal learning. Yet the standard self-attention mechanism carries a fundamental scalability problem: its computational and memory complexity grows quadratically with sequence length, O(N^2). As soon as models process long documents, high-resolution images, or extended video contexts, GPU memory fills up rapidly and training throughput collapses.

FlashAttention, introduced by Tri Dao and collaborators in 2022, addresses this bottleneck without changing the mathematical definition of attention. It is mathematically equivalent to standard attention, but it is engineered around the actual memory hierarchy of modern GPUs. By minimizing expensive data movement between high-bandwidth memory (HBM) and on-chip SRAM, FlashAttention dramatically accelerates training and reduces memory consumption, making very long sequence lengths practical on a single device.
The standard attention operation is defined as:

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V

A naive GPU implementation typically proceeds in three distinct stages. First, Q and K are loaded from HBM to compute the score matrix S = QK^T, after which S is written back to HBM. Second, S is read again, passed through a softmax to produce the probability matrix P, and P is written back. Third, P and V are loaded to compute the output O = PV, which is then written to HBM once more.

The critical issue is that S and P are both N x N matrices, where N is the sequence length. When N increases from 2,000 to 16,000, the size of these intermediate matrices grows by a factor of 64. For long sequences, these matrices simply cannot fit in fast on-chip memory, so the GPU spends most of its time moving data rather than performing arithmetic.

In modern deep learning hardware, the primary bottleneck is rarely raw compute throughput (FLOPs). It is memory bandwidth. Repeated reads and writes of large intermediate tensors to slow HBM create a “memory wall” that leaves GPU cores underutilized. FlashAttention attacks this problem directly by treating input/output (I/O) complexity as the central optimization target.

An IO-Aware Approach to Attention

FlashAttention is built on an IO-aware principle: reorganize the computation so that large intermediate matrices never need to be materialized in HBM. Instead, the algorithm loads small tiles of Q, K, and V into the GPU’s fast SRAM, performs the necessary operations there, and writes only the final output back to HBM.

This approach requires three algorithmic innovations.

Tiled Computation

FlashAttention partitions Q, K, and V into smaller blocks. During the forward pass, the algorithm iterates over these blocks, loading one tile at a time into SRAM. Attention scores are computed locally within the tile, and partial results are accumulated on-chip. Because each tile is much smaller than the full N x N score matrix, the hardware can keep the working set in SRAM and avoid repeated round-trips to HBM.

Online Softmax

A major challenge in tiled computation is the softmax normalization. Standard softmax requires knowledge of the entire row of scores before probabilities can be computed, because the denominator depends on every element in that row. FlashAttention uses an online softmax algorithm that updates the running maximum and running sum incrementally as new tiles are processed. By maintaining a few scalar statistics, it can rescale previously accumulated values and produce the exact softmax result in a single pass, without ever storing the full probability matrix in HBM.

Recomputation for Backward Passes

During backpropagation, standard implementations typically need to store the softmax probability matrix P to compute gradients. FlashAttention avoids this by recomputing attention scores on the fly during the backward pass. Using a small set of saved statistics from the forward pass, it regenerates the required values in SRAM and computes gradients locally. The extra arithmetic is negligible compared to the memory savings, because the GPU no longer has to allocate and manage an O(N^2) activation tensor for the backward pass.

Practical Usage

In production, developers rarely write custom CUDA kernels for FlashAttention. Major deep learning frameworks already expose optimized implementations. The flash-attn package provides a standalone CUDA kernel, while PyTorch’s scaled dot-product attention (SDPA) can automatically dispatch to FlashAttention when a compatible NVIDIA GPU is available.

Before running the example, ensure the environment meets these requirements:

• A compatible NVIDIA GPU with Tensor Core support
• CUDA toolkit installed and accessible
• PyTorch with scaled dot-product attention support

Installation is straightforward:

pip install flash-attn --no-build-isolation

Once installed, PyTorch users can enable FlashAttention through the native SDPA API:

import torch
import torch.nn.functional as F

batch_size = 2
num_heads = 4
seq_len = 2048
head_dim = 64

device = "cuda" if torch.cuda.is_available() else "cpu"

q = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=torch.float16)
k = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=torch.float16)
v = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=torch.float16)

with torch.backends.cuda.sdp_kernel(
    enable_flash=True,
    enable_math=False,
    enable_mem_efficient=True
):
    output = F.scaled_dot_product_attention(q, k, v, is_causal=True)

print("Output tensor shape:", output.shape)

When enable_flash=True, PyTorch selects the FlashAttention kernel for supported hardware and data types. The is_causal=True flag is commonly used for decoder-only language models, where each token should only attend to previous tokens.

Advantages and Impact

FlashAttention delivers three major benefits.

Reduced memory footprint. Memory complexity drops from O(N^2) to O(N) with respect to sequence length. This makes it feasible to train models with 32k, 128k, or even longer context windows on a single GPU, whereas standard attention would run out of memory far earlier.

Higher training throughput. By eliminating most HBM traffic for intermediate matrices, the algorithm removes the I/O bottleneck. In practice, training speed often improves by two to four times, depending on sequence length and hardware generation.

Numerical exactness. FlashAttention does not rely on low-rank approximations, kernel methods, or sparse attention patterns. It computes the same result as standard attention, up to floating-point rounding, so there is no accuracy trade-off. Models can adopt it as a drop-in replacement without retraining or hyperparameter tuning.

These properties have made FlashAttention a foundational building block for modern large language models. Long-context training, retrieval-augmented generation, and multimodal sequence modeling all benefit from the ability to process more tokens per step without prohibitive memory costs.

Conclusion

FlashAttention represents a shift in how deep learning operators are optimized. Instead of focusing solely on arithmetic intensity, it starts from the physical constraints of GPU memory hierarchy and redesigns the attention algorithm accordingly. Through tiling, online softmax, and recomputation, it bypasses the memory wall that limited earlier implementations. For researchers and engineers working with long-sequence models, FlashAttention is not merely an optional speed-up; it is an essential tool for scaling Transformers to the context lengths