FlashAttention

1. The Intermediate Results Attention Creates

In self-attention, each token's hidden state is a vector, or list of numbers. Learned projections turn it into a query, a key, and a value. Collecting these vectors across the sequence gives the arrays Q, K, and V. A query is compared with keys to score the available positions. Softmax converts those scores into weights that sum to one, and the weights are used to combine the value vectors.

For one attention head, we can arrange the scores in a table. Each row belongs to a query position, and each column belongs to a key position. Five tokens give a five-by-five table with 25 entries. Ten tokens give 100 entries. Doubling the sequence length doubles both dimensions, making the table four times as large.

For a sequence of $n$ tokens, each of those tables has $n^2$ entries. The output is much smaller when the sequence is long: it contains one vector per token, with a fixed number of coordinates per head. Attention needs that output for the next computation, but it does not need to return the entire table of pairwise weights.

A causal mask prevents a token from reading future positions. However, filling part of a square tensor with masked values does not shrink its allocation. Avoiding the storage requires an implementation that never creates the full tensor in the first place.

2. Why Moving Data Can Limit the GPU

A GPU has both arithmetic units and several kinds of memory. Many data-center GPUs use high-bandwidth memory (HBM) as their large, main device memory. Smaller working areas on the processor, including shared memory and registers, provide much faster access but have far less capacity. These are part of its on-chip memory resources; shared memory is built from static RAM, or SRAM.

Capacity tells us how much data can be stored; bandwidth describes how quickly data can be moved. An operation can spend much of its time waiting for memory transfers even when the GPU could perform more arithmetic. This is what it means for that operation to be memory-bound. The balance depends on the operation, its shape, and the hardware.

For a sense of scale, suppose one score tensor stores each entry in two bytes, as FP16 does. At 4,096 tokens, a single head's square tensor occupies 32 MiB. At 8,192 tokens, it occupies 128 MiB. A MiB is 1,048,576 bytes. These figures describe just one tensor for one head and one sequence, not total model memory; other data types and additional intermediates change the total.

Writing the score tensor, reading it for softmax, writing the weights, and reading them for the value mixture all move data. Reusing data on chip is therefore an important optimization, as described in NVIDIA's CUDA memory guidance. FlashAttention applies that idea to the complete attention calculation.

3. Computing a Small Tile at a Time

A tile is a small rectangular block of a larger calculation. Instead of creating scores for every query and every key at once, FlashAttention works with blocks of queries, keys, and values that fit its on-chip working space. A temporary score tile is used to update the result, then its space can be reused.

To follow the idea, keep one query block in the working area and visit the relevant key and value blocks in turn. Each visit contributes information to the same output rows. The algorithm retains a compact running summary for each query, which we will examine next. Different FlashAttention versions arrange the loops differently; this query-first view makes one output block easy to follow.

A GPU kernel is a program executed by many GPU threads. Fusion combines operations inside a kernel so that one operation can consume another's intermediate result without writing it out as a separate large array. Tiling limits the working set, while fusion keeps the score, normalization, and value-mixing steps close together.

The original FlashAttention paper describes this approach as IO-aware attention. Here, IO means reads and writes between memory levels, not network traffic or file access. The objective is to avoid moving the large pairwise intermediates through HBM. Inputs still need to be loaded, and some blocks can be loaded more than once.

4. The Difficulty with Softmax

Computing a query-key score in a small tile is straightforward because that score depends only on its query and key. Turning it into a final attention weight is different. Softmax divides its exponential by the sum of exponentials across all allowed keys for that query.

Suppose a query's four scores are $[1, 2, 3, 4]$, split into tiles $[1, 2]$ and $[3, 4]$. Applying softmax independently to each tile would make each pair sum to one. That gives both tiles a full unit of weight, even though all four positions must share one common normalization.

Averaging the two local outputs does not generally fix this. The tiles can have very different total exponential weights, so they should not automatically contribute equally. A correct tiled implementation must preserve enough information to combine their contributions with the proper relative scale.

There is also a numerical issue: exponentiating a large positive score can exceed what a floating-point number can represent. Stable softmax subtracts the largest score before exponentiating. This leaves the normalized weights unchanged, because every exponential has been multiplied by the same factor. In a tiled calculation, however, a larger maximum may arrive in a later tile.

5. Keeping a Running Softmax

Online softmax handles this by updating a running calculation as new scores arrive. “Online” here means incremental processing; it has nothing to do with an internet connection. For each query, the algorithm retains the largest score seen so far, the sum of shifted exponentials, and an accumulated weighted sum of values.

When a new tile raises the maximum, both old sums are multiplied by $\exp(m_{\text{old}} - m_{\text{new}})$. That converts them to the new reference point. The new tile's contributions can then be added on the same scale. The previous individual scores are no longer needed because their combined contribution is already represented by the two sums.

The final output is $u / \ell$. Dividing by the common denominator turns the accumulated value sum into the same weighted mixture that a full-row softmax would produce. For real attention, values are vectors, so $u$ is a vector too; its coordinates all share the same row maximum and denominator.

Keeping the numerator unnormalized until the end is useful because it avoids repeated divisions. The Triton fused-attention tutorial shows this running-maximum, denominator, and output-accumulator pattern in a GPU implementation. Its implementation details are more involved, but the normalization principle is the one illustrated here.

6. A Forward Pass Without the Full Matrix

The query, key, and value projections still happen before the attention kernel. Positional transformations such as RoPE are applied where the model requires them. FlashAttention receives these arrays and computes the attention output; it does not replace the projections or change the model's learned weights.

In the query-first schedule above, a GPU work unit loads a query tile and initializes one running state per query row. It loads the next key and value tile, computes the query-key dot products, and divides the scores by the square root of the query/key width. The causal mask is applied before those scores enter the softmax update.

The score tile updates the row maxima and denominators. Its exponentials are multiplied by the value tile to update the output accumulators, with earlier contributions rescaled when needed. Those temporary scores and exponentials can then be discarded. Their useful information has already been folded into the running state.

Once every allowed key tile has contributed, each output accumulator is divided by its denominator and the completed output rows are written to HBM. Other query tiles follow the same process and can be scheduled in parallel. The full square score and probability matrices never need to be stored in main memory.

Causal attention can skip tiles that lie entirely in the future. Tiles crossing the diagonal still need an element-by-element causal mask, since some entries are allowed and others are not. This preserves the original attention pattern; tiling does not restrict a query to nearby tokens.

7. Training Without Saving Every Attention Weight

Training has a forward pass that produces outputs and a backward pass that computes gradients, the quantities used to update learned parameters. A straightforward attention implementation can retain its large attention-weight matrix because those weights are useful during gradient calculation. Avoiding that saved matrix is necessary to keep the memory benefit during training.

FlashAttention saves compact normalization information for each row, along with the inputs and output needed for backward computation. One representation of that information is the log-sum-exp, written $L = m + \log(\ell)$. This is the logarithm of the full row's exponential sum, computed using the stable running state.

During the backward pass, the kernel reloads blocks of queries and keys, recomputes a score tile, and reconstructs its attention weights using the saved row normalization. For an allowed entry with scaled score $s_{ij}$, its weight is $\exp(s_{ij} - L_i)$. The tile can then contribute to gradients for Q, K, and V without restoring the full matrix in HBM.

This exchanges some additional arithmetic for fewer saved activations and less memory traffic. Recomputing a tile can be cheaper than loading a large stored intermediate. The benefit comes from where work and data are placed, so fewer arithmetic operations alone would not tell us which implementation is faster.

The examples in this article omit attention dropout. When dropout is enabled, the backward implementation must reproduce the same dropout decisions used in the forward pass, typically through saved random-number-generator information. It cannot independently sample a new mask while reconstructing the weights.

8. What FlashAttention Changes

Exact attention means FlashAttention preserves the dense attention formula and the chosen mask rather than introducing a sparse or low-rank approximation. The reordering of floating-point operations can produce small numerical differences, so exact does not imply bit-for-bit equality between kernels. Ordinary rounding differences should not be confused with deliberately dropping allowed token interactions.

Linear memory growth refers to avoiding the quadratic attention intermediates. With a fixed head width, Q, K, V, and the output grow linearly with sequence length, and the saved normalization state uses only a few numbers per query row. The temporary tile workspace is bounded by the kernel's tile sizes. This is not a claim that all memory in an entire training system is accounted for by those arrays.

The pairwise arithmetic for dense, full-sequence attention still grows approximately with the square of sequence length. Doubling the context therefore still creates about four times as many query-key interactions. FlashAttention reduces memory traffic and improves execution; it does not turn dense attention into a linear-time algorithm or remove the model's context-length constraints.

Actual speed depends on sequence length, batch size, head dimensions, data type, and GPU. Tile sizes must balance data reuse against the limited registers and shared memory available to concurrent work. FlashAttention-2 improves work partitioning and reduces non-matrix-multiplication overhead, building on the same tiled attention approach. There is no single speedup factor that applies to every model and device.

Long-prompt processing and training contain many query rows, which makes avoiding square intermediates especially useful. A cached generation step may have only one new query and a long set of past keys and values, so its workload is different. FlashAttention does not eliminate the KV cache; efficient decoding also depends on kernels and memory layouts suited to that shape.

9. Using FlashAttention in an LLM

Inside a Transformer block, FlashAttention occupies the attention computation between the Q/K/V projections and the head-combining output projection. Residual connections, normalization, and the feed-forward network keep their existing roles. This is why a model can use a different attention kernel without needing to relearn its weights.

Libraries expose this through optimized attention operations. In PyTorch, scaled_dot_product_attention can select a supported implementation, including a FlashAttention backend on compatible hardware. The following example requests causal self-attention for equally long Q, K, and V sequences:

import torch.nn.functional as F

def causal_self_attention(q, k, v):
    # q, k, v: [batch, heads, tokens, head_width]
    # Equal sequence lengths, no padding, and no attention dropout.
    return F.scaled_dot_product_attention(
        q,
        k,
        v,
        dropout_p=0.0,
        is_causal=True,
    )

This call does not guarantee that a particular FlashAttention kernel runs. Dispatch depends on the installed library, device, input shapes, data types, masks, and backend settings. A profiler can confirm which kernel was selected. The function sets dropout to zero explicitly because this API applies the supplied dropout probability even when the surrounding model is in evaluation mode.

Padding, custom masks, and cached decoding require additional care. In particular, a query representing the newest token must be aligned with its actual position in the key cache. The simple square causal mask used here should not be assumed correct for every unequal-length query/key layout. Requesting the full attention-weight tensor for inspection would also require its quadratic output storage, even if the model normally avoids it.

The central idea is that the attention output can be computed without keeping every intermediate weight. Tiles limit the working set, running softmax preserves the correct normalization, and recomputation avoids saving a large matrix for training. Together, these let the GPU spend less time moving temporary data while preserving the model's attention calculation.