FlashAttention: The GPU Memory Hack That Changed Everything

Here's a fact that will bother you: the standard attention algorithm wastes 90% of its time waiting for data to move between memory chips on the GPU. Not computing. Waiting.

GPUs have two types of memory: fast (SRAM, ~20 MB) and slow (HBM, ~80 GB). The standard attention algorithm writes enormous intermediate matrices to slow memory, then reads them back. Over and over.

FlashAttention rewrites the attention algorithm to keep everything in fast memory. No intermediate matrices. No memory bottleneck. Same exact mathematical result. 2-4x faster. 5-20x less memory.

It's the single most impactful optimization in modern LLM inference, and understanding it shows you how the GPU actually works.

The Problem: GPU Memory Hierarchy

Think of GPU memory like your kitchen:

Memory Analogy Size Speed

SRAM (on-chip) Kitchen counter ~20 MB ~19 TB/s

HBM (off-chip) Refrigerator ~80 GB ~3 TB/s

Your kitchen counter is tiny but instant — you grab ingredients right there. The fridge is huge but you have to walk to it every time.

Standard attention cooks a meal by:

Getting ALL ingredients from the fridge → spreading them on the counter

Doing one step → putting everything BACK in the fridge

Getting everything OUT of the fridge again → doing the next step

Repeat

FlashAttention cooks by:

Getting a small batch from the fridge

Doing ALL steps on that batch without going back to the fridge

Getting the next small batch

Standard Attention: The Memory-Wasteful Way

Here's what standard attention does step by step:

python

Standard attention: Q, K, V are (N, d) where N = sequence length

Step 1: Compute attention scores

S = Q @ K.T # (N, N) matrix — written to HBM

At N=4096, d=128: S = 4096×4096 = 64 MB in FP16

Step 2: Apply softmax

P = softmax(S, dim=-1) # (N, N) matrix — read from HBM, write result to HBM

Another 64 MB round trip

Step 3: Multiply by values

O = P @ V # (N, d) — read P from HBM, write O to HBM

The S and P matrices are N×N. For a sequence of 128K tokens:

S matrix = 128K × 128K × 2 bytes = 32 GB
P matrix = 128K × 128K × 2 bytes = 32 GB
Total intermediate storage = 64 GB — just for ONE attention layer

That's more than the GPU's total HBM. The model literally can't compute attention for long sequences. And even for shorter sequences, the constant reading/writing to HBM is the bottleneck — not the actual math.

This is called being memory-bound, not compute-bound.

🔧

FlashAttention: IO-Aware Tiling

FlashAttention's key insight: you don't need to materialize the full N×N matrix. You can compute attention in tiles that fit in SRAM.

The Tiling Strategy

Instead of computing the full attention matrix, FlashAttention processes blocks:

Standard: Compute ALL of QK^T → softmax → multiply by V

Flash: For each block of Q:

For each block of K, V:
              Compute partial QK^T → partial softmax → accumulate into output

python

FlashAttention pseudocode (simplified)

Divide Q, K, V into blocks that fit in SRAM

block_size = SRAM_SIZE // (3 * d)  # How many rows fit in fast memory
O = zeros(N, d)                    # Output (in HBM)
L = zeros(N)                       # Running softmax denominator

for i in range(0, N, block_size): # Outer loop: blocks of Q

Qi = Q[i:i+block_size]         # Load block of Q from HBM → SRAM
    Oi = zeros(block_size, d)      # Local accumulator (in SRAM)
    li = zeros(block_size)         # Local softmax denominator (in SRAM)
    mi = full(block_size, -inf)    # Local softmax max (in SRAM)

for j in range(0, N, block_size):  # Inner loop: blocks of K, V
        Kj = K[j:j+block_size]    # Load block of K from HBM → SRAM
        Vj = V[j:j+block_size]    # Load block of V from HBM → SRAM

# All computation happens in SRAM — no HBM round trips!
        Sij = Qi @ Kj.T           # Block attention scores (in SRAM)

# Online softmax: update running max and denominator
        mij = max(Sij, dim=-1)
        mi_new = max(mi, mij)
        li = li * exp(mi - mi_new) + sum(exp(Sij - mi_new), dim=-1)

# Accumulate output
        Oi = Oi * exp(mi - mi_new) + exp(Sij - mi_new) @ Vj
        mi = mi_new

O[i:i+block_size] = Oi / li   # Write final output to HBM

Total HBM reads: O(N²d / SRAM_SIZE) instead of O(N²)

No N×N intermediate matrix stored!

The Online Softmax Trick

The critical innovation: computing softmax incrementally. Normally, softmax needs to see all values first (to compute the denominator). FlashAttention uses the online softmax algorithm:

Process block 1: compute local max and sum of exponentials

Process block 2: update the running max, rescale previous results, add new exponentials

Repeat for all blocks

The math is exact. Not an approximation. The final result is bit-for-bit identical to standard attention.

📊

The Numbers: Speed and Memory

Memory Savings

Sequence Length Standard Attention FlashAttention Savings

FlashAttention never materializes the N×N matrix. The intermediate computation happens in SRAM tiles and is immediately consumed. Zero extra HBM usage.

Speed Improvements

Setup Standard FlashAttention Speedup

BERT (seq 512) 100 ms 45 ms 2.2x

GPT-2 (seq 1K) 180 ms 62 ms 2.9x

The longer the sequence, the bigger the speedup — because standard attention's HBM bottleneck grows quadratically.

FlashAttention-2: Even Faster

Paper: "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning" (Tri Dao, 2023)

FlashAttention-2 improved on the original with:

1. Better Parallelism Across Thread Blocks

FlashAttention-1 parallelized over batch and heads. FA-2 also parallelizes over the sequence dimension — splitting the outer loop across GPU streaming multiprocessors (SMs).

1. Reduced Non-Matmul Operations

FA-1 spent ~50% of time on non-matmul operations (rescaling, max tracking). FA-2 restructured the algorithm to minimize these, keeping the GPU's tensor cores (which only do matmul) busy:

FlashAttention-1: ~50% matmul, ~50% other

FlashAttention-2: ~70% matmul, ~30% other

1. Forward Pass Speedup

FA-1 FA-2 Improvement

FlashAttention-3: Hopper Architecture

For: NVIDIA H100 GPUs (Hopper architecture)

FA-3 takes advantage of H100-specific hardware:

FP8 tensor cores: Run attention in FP8 precision, 2x the throughput of FP16

Asynchronous data movement: The H100's TMA (Tensor Memory Accelerator) can copy data while computing — overlap instead of wait

Warp specialization: Different warps (groups of 32 threads) handle different tasks simultaneously

FA-3 on H100:

• FP16: ~1.5x faster than FA-2
• FP8: ~2x faster than FA-2 (with negligible quality loss)
• Utilization: up to 75% of theoretical peak FLOPS

🔬

Integration: It's Everywhere

FlashAttention isn't something you install separately — it's built into every major framework:

python

PyTorch (built-in since 2.0)

import torch.nn.functional as F

Automatically uses FlashAttention when available

output = F.scaled_dot_product_attention(query, key, value, is_causal=True)

PyTorch picks the fastest backend: FlashAttention, Memory-Efficient, or Math

Hugging Face Transformers

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-3-8B",
    attn_implementation="flash_attention_2"  # Explicit FA-2

)

vLLM (default)

FlashAttention is the default attention backend — no config needed

Framework FlashAttention Support

🛡️

When FlashAttention Helps Most

FlashAttention's benefit depends on whether your workload is memory-bound:

Scenario Memory-Bound? FA Benefit

The prefill phase (processing your entire prompt) is where FlashAttention shines brightest — it's processing many tokens at once, and the N×N attention matrix is massive.

📦

Try It Yourself

python

import torch

import torch.nn.functional as F

import time

Check FlashAttention availability

print(torch.backends.cuda.flash_sdp_enabled())  # True on modern GPUs

Benchmark standard vs Flash attention

def benchmark_attention(seq_len, n_heads=32, head_dim=128, device="cuda"):
q = torch.randn(1, n_heads, seq_len, head_dim, device=device, dtype=torch.float16)
    k = torch.randn(1, n_heads, seq_len, head_dim, device=device, dtype=torch.float16)
    v = torch.randn(1, n_heads, seq_len, head_dim, device=device, dtype=torch.float16)
# Warm up
    for _ in range(3):
        F.scaled_dot_product_attention(q, k, v, is_causal=True)
    torch.cuda.synchronize()
# Benchmark
    start = time.time()
    for _ in range(100):
        F.scaled_dot_product_attention(q, k, v, is_causal=True)
    torch.cuda.synchronize()
return (time.time() - start) / 100

for seq_len in [512, 1024, 2048, 4096, 8192]:

t = benchmark_attention(seq_len)
    print(f"Seq {seq_len:5d}: {t*1000:.2f} ms")

🚀

Practical Takeaways

Standard attention wastes 90% of time on memory transfers — the math is fast, moving data is slow

FlashAttention tiles the computation to stay in SRAM — never materializes the N×N attention matrix

The result is mathematically exact — not an approximation, bit-for-bit identical

FA-2 is the current production standard — 2x faster than FA-1, built into PyTorch 2.0+

Longer sequences benefit more — the speedup grows with context length

You're probably already using it — PyTorch, HF Transformers, and vLLM enable it by default

🔄

What's Next?

Episode 76: PagedAttention — FlashAttention optimizes how attention computes. PagedAttention optimizes how KV cache is stored. It treats GPU memory like a virtual memory system — with pages, page tables, and copy-on-write. It's the innovation behind vLLM, and it's why you can serve hundreds of concurrent users on a single GPU.

← Previous

Ep 74: Attention Variants

Next →

Ep 76: PagedAttention

Next: Episode 76 — PagedAttention

50 users chatting simultaneously. Pre-allocated GPU memory for max context length. 70% of that memory is wasted. PagedAttention fixes this.

This is part of a 98-episode series covering AI engineering from tokens to production deployment.