Phase 2 so far has been about memory: the KV cache costs too much, MQA and GQA compress it. But there's a second wall that attention hits as sequences grow longer — a computational wall. It's separate from the memory problem, and it breaks in a different way.
This article covers the O(n²) scaling problem, why it's a hardware problem as much as an algorithmic one, how Flash Attention solved it without changing the mathematics, and what this means for the path to MLA in Phase 3.
Flash Attention (Dao et al., 2022, arXiv:2205.14135) is one of the most influential engineering papers in the LLM era. It didn't change what attention computes — it changed where the computation happens. The result was 2–4× faster attention with constant instead of linear memory, enabling context lengths that were previously impossible.
The O(n²) problem
Standard scaled dot-product attention computes:
S = Q @ K^T # (S, S) attention scores matrix A = softmax(S / √d) # (S, S) attention weights O = A @ V # (S, d) output
The score matrix S is S × S — where S is the sequence length. At 1K tokens, that's 1 million entries. At 32K tokens, that's over 1 billion entries. At 128K tokens, that's over 16 billion entries.
This is the O(n²) problem. Attention compute scales as the square of sequence length. The memory required to store the score matrix scales the same way.
| Sequence length | Score matrix size (FP16) | Notes |
|---|---|---|
| 2K tokens | 16 MB | Fine — fits in L2 cache fragments |
| 8K tokens | 256 MB | Uncomfortable — many HBM reads |
| 32K tokens | 4 GB | Significant HBM allocation per layer |
| 128K tokens | 64 GB | Exceeds A100 entirely — per layer |
The 128K row is the breaking point: a single attention layer at 128K sequence length would need a 64 GB score matrix in FP16. An A100 has 80 GB total HBM. You can't fit even one layer's attention scores, let alone the KV cache and model weights alongside them.
Why it's a memory hierarchy problem
The O(n²) memory cost isn't just about capacity — it's about where the computation happens. Modern GPU memory is a hierarchy:
- Registers: ~256 KB per SM — fastest, private to each thread
- SRAM (shared memory): ~20 MB total across 108 SMs — fast, shared within a threadblock
- HBM (high-bandwidth memory): 80 GB — slow relative to SRAM, 2 TB/s bandwidth
The SRAM-to-HBM bandwidth gap is roughly 100×. Moving data from HBM to SRAM costs approximately 100× more time per byte than keeping it in SRAM.
Standard attention loads Q, K, V from HBM, computes the full score matrix S in SRAM/registers, writes S back to HBM, then loads it again to compute softmax, writes again, loads again for the V multiplication. Each step is a separate kernel that writes to and reads from HBM.
The number of HBM read/write operations in naive attention is O(n²) — the score matrix must be written and read in its entirety multiple times. At 128K tokens, this is not just impractically large to store — it's impractically slow to move back and forth.
Fig 1 — Naive attention writes the full n² score matrix to HBM between kernels. Flash Attention tiles the computation so the score matrix never leaves SRAM.
Flash Attention: the tiling insight
Tri Dao et al.'s key insight in the 2022 Flash Attention paper: you don't need to materialise the full score matrix in HBM. Attention is mathematically equivalent to a sequential tiled computation that stays in SRAM.
The algorithm divides Q, K, V into blocks that fit in SRAM (say, 64 × d_h each). For each tile of Q, it iterates over tiles of K and V, computing a partial softmax and accumulating the output:
for q_block in split(Q, block_size):
acc = 0, running_max = -inf, running_sum = 0
for k_block, v_block in zip(split(K, block_size), split(V, block_size)):
# Load tile from HBM to SRAM (small!)
S_tile = q_block @ k_block.T / sqrt(d_h)
# Online softmax: update running max and sum
new_max = max(running_max, S_tile.max())
acc = acc * exp(running_max - new_max) + exp(S_tile - new_max) @ v_block
running_sum = running_sum * exp(running_max - new_max) + exp(S_tile - new_max).sum()
running_max = new_max
output_block = acc / running_sum # normalise by final sum
The "online softmax" trick (derived from earlier work by Milakov & Gimelshein) is what makes tiling work. Softmax requires dividing by the sum of exponentials over the full row — you can't do it per-tile without knowing the full sum. The online variant maintains a running max and running sum that get updated as new tiles are processed. The final output after all tiles equals the standard softmax output exactly.
No approximation. Same result. But the S × S score matrix never exists as a full HBM allocation — it's computed tile by tile entirely within SRAM, then discarded.
The memory complexity win
Standard attention: O(n²) memory for the score matrix.
Flash Attention: O(n) memory — only Q, K, V, and the output, each of size S × d.
In concrete terms:
| Sequence length | Naive S matrix | Flash Attention memory | Saving |
|---|---|---|---|
| 4K | 64 MB | ~6 MB | 10× |
| 32K | 4 GB | ~48 MB | 83× |
| 128K | 64 GB | ~192 MB | 333× |
| 1M | 4 TB (impossible) | ~1.5 GB | makes it possible |
At 128K tokens, naive attention requires 64 GB for the score matrix alone — more than an A100's entire HBM. Flash Attention reduces this to 192 MB, making 128K context sequences feasible on a single GPU.
The speed win: IO-awareness
Flash Attention is described as "IO-aware" because its primary optimisation target is HBM bandwidth, not raw FLOPs. Compute (FLOPs) has become cheap relative to memory movement. Modern GPUs can perform far more floating-point operations per second than they can move bytes per second.
The roofline model makes this concrete. An A100 SXM has:
- 312 TFLOPS (FP16 tensor cores)
- 2 TB/s HBM bandwidth
Compute roof: 312 TFLOPS × 2 bytes/element = 624 TB/s effective if perfectly compute-bound.
Memory roof: 2 TB/s — the ceiling for memory-bound operations.
Naive attention is memory-bound at most sequence lengths. It doesn't use the tensor cores efficiently because it spends most time moving data between HBM and SRAM rather than computing. Flash Attention dramatically increases arithmetic intensity — more computation per byte loaded from HBM — and pushes attention operations closer to the compute roof.
In the original paper, this yielded 2–4× wall-clock speedup for attention on A100 GPUs. More importantly, it enabled exact attention at sequence lengths where naive attention would OOM or become impossibly slow.
Flash Attention 2 and 3
The original Flash Attention used software pipelining and careful SRAM management but didn't fully exploit tensor core utilization. Flash Attention 2 (Dao, 2023) improved by:
- Reducing non-matmul FLOPs (the online softmax overhead)
- Better parallelism — splitting work across Q blocks as well as K/V blocks
- Optimized warp partitioning to avoid cross-warp communication
Flash Attention 2 achieved ~2× additional speedup over FA1, reaching 72% of the theoretical A100 FP16 MFU (Model FLOP Utilization) for attention.
Flash Attention 3 (Shah et al., 2024) targeted the H100's Hopper architecture with asynchronous Tensor Memory Accelerator (TMA) and warp specialisation. On H100 it achieves 1.5–2.0 PFLOPS attention throughput — within 85-90% of the theoretical hardware limit.
The remaining problem: representation compression
Flash Attention solves the compute and memory access problem of attention at long sequences. But it doesn't address the KV cache size problem we've been tracking through Phase 2.
Flash Attention tiles the attention computation during the forward pass. During inference (decode), you still need the K and V vectors for all previous tokens to be available for new queries. The KV cache contains the materialised K and V tensors — Flash Attention's tiling doesn't reduce how much you store.
MQA and GQA reduced the KV cache by reducing the number of heads. Flash Attention reduced the compute cost and peak memory of the attention kernel itself. They solve different problems and are typically used together.
Flash Attention = cheaper to compute attention per forward pass. MQA/GQA = cheaper to store the KV cache across decode steps. These are orthogonal optimisations. Production systems use both.
The combined picture
Let's put all the Phase 2 insights together for a concrete serving scenario:
Model: LLaMA-3-70B, H=64 query heads, G=8 KV heads, 80 layers, d_h=128
Request: 100K token context, generating 1K tokens, batch size 4
Without Flash Attention + GQA (naive MHA):
- Score matrix per layer: 100K × 100K × 2 bytes ≈ 20 GB per layer (OOM)
- KV cache: 2 × 80 × 64 × 128 × 100K × 4 × 2 bytes ≈ 262 GB (OOM)
With Flash Attention + GQA (production):
- Score matrix memory: O(n) tiles in SRAM only — negligible HBM allocation
- KV cache: 2 × 80 × 8 × 128 × 100K × 4 × 2 bytes ≈ 32.8 GB — fits on 2× A100s
The 100K context serving scenario that was physically impossible with naive MHA becomes routinely achievable in production with Flash Attention + GQA.
Why even this isn't enough for DeepSeek
GQA at G=8 gives 8× KV cache compression on H=64 models. Flash Attention brings the compute cost under control. But DeepSeek-V2 operates at 128K native context with plans for 1M+ token contexts, and at 671B parameters (DeepSeek-V3). At that scale, even 8× compression leaves a cache that requires expensive multi-GPU allocation.
More fundamentally: GQA reduces the number of KV head tensors but doesn't change the information content of each tensor. Each KV vector still lives in the full d_h dimensional space. The representation itself hasn't been compressed — just the count of representations.
DeepSeek's MLA (Multi-Head Latent Attention, introduced in DeepSeek-V2) takes a different approach. Instead of asking "how many KV heads do we need?", it asks "how many dimensions do we actually need to represent the KV information?"
The answer, in DeepSeek's architecture, is 512 dimensions — regardless of d_model (which is 5120 for DeepSeek-V2) and d_h × H_kv (which is 128 × 128 = 16,384 for full MHA). MLA compresses the KV representation into a 512-dimensional latent vector, achieving ~93% KV cache reduction compared to MHA while maintaining quality that GQA at equivalent compression ratios cannot match.
That's Phase 3. The setup from Phase 2:
- KV cache is the bottleneck (2.1)
- Head count reduction works but has quality limits — MQA (2.2), GQA (2.3)
- Compute scaling is solvable with IO-aware tiling — Flash Attention (2.4)
- Representation compression is the next frontier — MLA (Phase 3)
Summary: Phase 2 complete
Phase 2 started with a bottleneck and traced the engineering response to it:
The bottleneck: The KV cache grows as O(n × L × H × d_h). At long contexts, it dominates GPU memory and decode bandwidth.
The first response — head count reduction (MQA, GQA): Reduce H_kv from H down to 1 (MQA) or G (GQA). Works. Adds quality tradeoffs that worsen with compression ratio. The G=8 sweet spot (LLaMA-2/3, Mistral) is now industry standard.
The second response — IO-aware computation (Flash Attention): Tile the attention kernel to keep the score matrix in SRAM. Reduces compute memory from O(n²) to O(n). 2–4× speedup. Enables sequence lengths that naive attention physically cannot support.
The gap both responses leave: Neither reduces the dimensionality of what's stored. The KV representation is still full-precision, full-dimension vectors. There's more compression available — DeepSeek found it.