Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

FlashAttention Algorithm › FlashAttention Complexity

Michael BrenndoerferUpdated June 29, 20254 min read

Language AI HandbookAll chapters

Part 17 of 67Efficient Attention

Chapter 124FlashAttention Algorithm › FlashAttention Complexity

Reading section 6 of 12 · FlashAttention Complexity · 4 minSection 6 of 12 · 4 min

Section 6 of 12 · 4 min read

FlashAttention Algorithm › FlashAttention ComplexityOne section · Sections in this chapter6 / 12 ↓

Sections in this chapter12

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Let's formalize the complexity improvements FlashAttention provides. We analyze both time (compute) and space (memory) complexity.

Time ComplexityLink Copied

FlashAttention performs the same number of floating-point operations as standard attention: $O(n^2 d)$ , where $n$ is the sequence length and $d$ is the head dimension. The algorithm does not reduce computational complexity; it restructures the computation to reduce memory access.

However, by avoiding HBM transfers for the attention matrix, FlashAttention reduces the effective time from I/O-bound to compute-bound. The number of HBM accesses drops from $O(n^2 + nd)$ to $O(nd)$ :

Standard attention HBM access: $O(n^2 + nd)$ because it reads/writes the $n \times n$ attention matrix plus the Q, K, V, O matrices of size $n \times d$
FlashAttention HBM access: $O(nd)$ because it only reads Q, K, V and writes O, never materializing the attention matrix in HBM

In[20]:

Code

def complexity_analysis(seq_len, d_model, block_size):
    """
    Analyze complexity of standard vs FlashAttention.

    Args:
        seq_len: Sequence length
        d_model: Model dimension
        block_size: FlashAttention block size

    Returns:
        Dictionary with complexity statistics
    """
    # FLOPs (same for both)
    # Q @ K^T: n^2 * d multiplies + n^2 * d adds = 2 * n^2 * d
    # Softmax: ~5 * n^2
    # Weights @ V: 2 * n^2 * d
    flops = 4 * seq_len**2 * d_model + 5 * seq_len**2

    # HBM accesses (bytes)
    # Standard: read QKV (3nd), write scores (n^2), read scores + write softmax (2n^2),
    #           read softmax + V + write output (n^2 + nd + nd)
    standard_hbm = (
        3 * seq_len * d_model + 4 * seq_len**2 + 2 * seq_len * d_model
    ) * 2

    # Idealized FlashAttention traffic: read Q, K, V and write O without
    # materializing the n x n score/weight matrices in HBM.
    flash_hbm = (4 * seq_len * d_model) * 2

    return {
        "flops": flops,
        "standard_hbm": standard_hbm,
        "flash_hbm": flash_hbm,
        "hbm_reduction": standard_hbm / flash_hbm,
    }


configs_complexity = [
    (512, 768, 64),
    (2048, 768, 64),
    (8192, 768, 64),
    (32768, 768, 64),
]

Out[21]:

Console

Complexity Analysis:

   Seq Len    FLOPs (G)   Std HBM (GB)   Flash HBM (GB)    Reduction
--------------------------------------------------------------------
       512         0.8         0.006          0.0031         1.9x
     2,048        12.9         0.049          0.0126         3.9x
     8,192       206.5         0.600          0.0503        11.9x
    32,768      3303.9         8.842          0.2013        43.9x

The FLOPs remain constant between approaches (identical computation). The HBM access reduction grows with sequence length: at 32K tokens, FlashAttention reduces memory transfers by over 50x. This reduction is the source of the speedup.

Space ComplexityLink Copied

The space complexity improvements are more dramatic:

Standard attention: $O(n^2)$ for storing the attention matrix
FlashAttention: $O(n)$ for QKV and O, plus running statistics

This reduction from quadratic to linear space complexity is what enables processing long sequences that would otherwise exhaust GPU memory.

Out[22]:

Visualization

Grouped log-scale bar chart with standard attention normalized to one for FLOPs, HBM access, and peak memory; FlashAttention matches FLOPs but has much lower HBM access and memory bars. — Relative resource cost at 8K tokens, normalized to standard attention for each metric. Compute remains unchanged, while FlashAttention sharply reduces HBM traffic and peak memory.

The normalized bar chart at 8K tokens illustrates the trade-offs without mixing physical units on one axis. FLOPs are identical, while the idealized HBM traffic and single-layer peak-memory estimate are lower for FlashAttention because it never materializes the full attention matrix.