Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

FlashAttention Algorithm › GPU Memory Hierarchy

Michael BrenndoerferUpdated June 29, 20257 min read

Language AI HandbookAll chapters

Part 17 of 67Efficient Attention

Chapter 124FlashAttention Algorithm › GPU Memory Hierarchy

Reading section 2 of 12 · GPU Memory Hierarchy · 7 minSection 2 of 12 · 7 min

Section 2 of 12 · 7 min read

FlashAttention Algorithm › GPU Memory HierarchyOne section · Sections in this chapter2 / 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

To understand why FlashAttention works, you first need to understand why standard attention is slow. The bottleneck is not computation but memory bandwidth. Modern GPUs can perform trillions of floating-point operations per second, but reading and writing data to memory is orders of magnitude slower than computing on that data.

GPUs have a hierarchical memory structure with different levels offering different trade-offs between capacity and speed:

GPU Memory Hierarchy

GPU memory is organized into multiple levels. High Bandwidth Memory (HBM) provides large capacity (tens of gigabytes) but relatively slow access. Shared Memory / SRAM provides fast access (10-20x faster than HBM) but limited capacity (tens to hundreds of kilobytes per streaming multiprocessor). Registers provide the fastest access but extremely limited capacity.

The key insight is that moving data between memory levels dominates execution time for many operations. A typical GPU might have:

HBM: 40-80 GB capacity, ~2 TB/s bandwidth
SRAM (shared memory): 192 KB per SM, ~19 TB/s bandwidth
Registers: Limited per thread, essentially zero latency

The bandwidth ratio between SRAM and HBM is roughly 10x. When an algorithm repeatedly reads and writes to HBM, it becomes "memory-bound," meaning computation stalls waiting for data transfers. Standard attention exhibits exactly this pattern.

Out[3]:

Visualization

Horizontal bar chart showing HBM with ~2 TB/s bandwidth and SRAM with ~19 TB/s bandwidth. — GPU memory hierarchy showing the trade-off between capacity and bandwidth. HBM offers tens of gigabytes but slower access, while SRAM provides 10x faster access in a much smaller capacity. FlashAttention exploits this gap by keeping intermediate data in fast SRAM.

The visualization makes the bandwidth gap concrete. SRAM is nearly 10x faster than HBM, but it's also 100,000x smaller. FlashAttention's tiling strategy keeps the hot data (attention score blocks and running statistics) in SRAM, reading from HBM only for Q, K, V and writing only the final output.

In[4]:

Code

def analyze_attention_memory_access(seq_len, d_model, dtype_bytes=2):
    """
    Analyze memory access patterns in standard vs tiled attention.

    Args:
        seq_len: Sequence length
        d_model: Model dimension
        dtype_bytes: Bytes per element (2 for fp16)

    Returns:
        Dictionary with memory access statistics
    """
    # Standard attention memory reads/writes to HBM

    # Read Q, K, V once
    read_qkv = 3 * seq_len * d_model * dtype_bytes

    # Write attention scores (n x n) to HBM
    write_scores = seq_len * seq_len * dtype_bytes

    # Read scores, write softmax output
    read_write_softmax = 2 * seq_len * seq_len * dtype_bytes

    # Read softmax weights and V, write output
    read_softmax_v = (seq_len * seq_len + seq_len * d_model) * dtype_bytes
    write_output = seq_len * d_model * dtype_bytes

    standard_total = (
        read_qkv
        + write_scores
        + read_write_softmax
        + read_softmax_v
        + write_output
    )

    # FlashAttention: Never materializes n x n matrix
    # Reads Q, K, V in blocks, writes only final output
    flash_total = (
        read_qkv + write_output
    )  # Simplified; actual includes block overhead

    return {
        "standard_bytes": standard_total,
        "flash_bytes": flash_total,
        "standard_gb": standard_total / 1e9,
        "flash_gb": flash_total / 1e9,
        "reduction": standard_total / flash_total,
    }


# Analyze for various sequence lengths
seq_lengths = [512, 2048, 8192, 32768]
d_model = 768

Out[5]:

Console

Memory Access Analysis (d=768, fp16):

   Seq Len     Standard        Flash    Reduction
--------------------------------------------------
       512       0.006G      0.0031G         1.9x
     2,048       0.049G      0.0126G         3.9x
     8,192       0.600G      0.0503G        11.9x
    32,768       8.842G      0.2013G        43.9x

The analysis reveals the source of FlashAttention's speedup. Standard attention writes the $n \times n$ attention matrix to HBM, then reads it back for softmax, then writes the softmax output, then reads it again for the final matrix multiplication. Each of these transfers takes time. FlashAttention avoids materializing the attention matrix entirely, reducing memory transfers by 10-100x for long sequences.

Why Memory Matters More Than FLOPsLink Copied

Consider the arithmetic intensity of attention: the ratio of floating-point operations to bytes transferred. Arithmetic intensity determines whether an operation is compute-bound (limited by processing speed) or memory-bound (limited by data transfer speed). Standard attention computes roughly $O(n^2 d)$ FLOPs but moves $O(n^2)$ bytes for the attention matrix alone, where $n$ is sequence length and $d$ is head dimension. When $n$ is large, the memory bandwidth becomes the limiting factor because the $n^2$ memory transfers grow faster than they can be processed.

Out[6]:

Visualization

Log-scale line plot showing standard attention approaching a low arithmetic-intensity plateau while FlashAttention rises with sequence length. — Arithmetic intensity of attention as sequence length increases. Standard attention (red) remains low because it materializes the n x n matrix in HBM, while FlashAttention (blue) rises with sequence length by keeping intermediate results in SRAM.

The plot reveals a fundamental difference. Standard attention's arithmetic intensity decreases as sequence length grows because the $n^2$ memory transfers dominate. FlashAttention maintains high arithmetic intensity by avoiding HBM transfers for intermediate results. Above the "compute-bound threshold" (around 200 FLOPs/byte for modern GPUs), the GPU operates efficiently. Below it, the GPU stalls waiting for memory.