C.W.K.
Stream
Lesson 02 of 05 · published

The O(n²) Bottleneck

~16 min · complexity, scaling, flashattention

Level 0Observer
0 XP0/50 lessons0/14 achievements
0/100 XP to next level100 XP to go0% complete

Why quadratic isn't a constant factor

The attention score matrix has shape n × n. At a sequence length of 1K that's a million entries. At 100K, ten billion. The matrix doesn't appear in the final output, but in the standard formulation it's materialized in memory before softmax — and computing it costs O(n²·d) FLOPs whether you store it or not. Compute and memory both grow quadratically with sequence length.

This isn't a constant-factor inconvenience you can engineer away. Doubling your context length quadruples your attention cost. Going from 8K to 128K — a 16× context increase — multiplies attention compute by 256. At some point the GPU just runs out of room, and beyond that, the wall-clock cost becomes prohibitive even when memory fits.

What FlashAttention did and didn't fix

FlashAttention (Dao et al., 2022 → FlashAttention-3 in 2024) is one of the most important systems papers of the era. It restructures the computation to keep tiles of the attention matrix inside the GPU's fast on-chip SRAM instead of constantly swapping to slow HBM. FlashAttention-3 reaches roughly 85% of an H100's theoretical peak using async TMA copies and FP8 quantization on Hopper.

What FlashAttention changed: the constant factor. What it did not change: the asymptotic complexity. The compute is still O(n²·d). FlashAttention moved the practical wall further out — but it's still a wall. To process million-token contexts you need a different algorithm, not just better kernels.

Where the wall actually shows up

It shows up at training time as out-of-memory errors when you try to lengthen your context window. It shows up at inference time as throughput collapse on long-context serving. And it shows up on consumer hardware most cruelly: a 2025 study found Transformers cannot process more than ~25K tokens on a 24GB consumer GPU, while equivalent SSMs handled 220K+ tokens on the same card. That gap is what motivates the rest of this quest.

Code

Memory growth of the attention matrix in fp16·python
def attention_bytes(seq_len, n_heads=32, dtype_bytes=2):
    # Score matrix: heads * n * n entries, each `dtype_bytes` wide
    return n_heads * seq_len * seq_len * dtype_bytes

for n in [1024, 4096, 32_768, 131_072, 1_000_000]:
    gb = attention_bytes(n) / (1024 ** 3)
    print(f'{n:>10} tokens -> {gb:>8.2f} GB attention matrix')
# 1024       -> 0.06 GB
# 4096       -> 1.0 GB
# 32768      -> 64.0 GB     <-- already over a single H100 80GB
# 131072     -> 1024.0 GB
# 1000000    -> ~59 TB

External links

Exercise

Pick a Hugging Face Transformers model you can run locally (Llama 3.2 1B is fine on a laptop). Set attn_implementation="eager", then "sdpa", then "flash_attention_2" if your hardware supports it, and measure prefill latency at sequence lengths 512, 4K, 32K. Plot latency vs. sequence length on a log-log axis. You should see all three implementations approach the same slope (≈ 2 in log-log) at long sequences — that slope is the algorithmic O(n²) showing through, regardless of kernel.

Progress

Progress is local-only — sign in to sync across devices.
Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)
Sign inPlease sign in to comment.

No comments yet — be the first.