Why is attention O(n²), how does FlashAttention help, and how do models handle long context?

30-second answer

Scores between all pairs make compute and naive memory quadratic in sequence length; per token the attention FLOPs (4·T·d per layer) only exceed the parameter matmuls (~24d²) beyond T ≈ 6d. FlashAttention is exact: it tiles Q, K, V into on-chip memory and uses an online softmax with running max and normalizer, so the T×T matrix is never written to HBM; memory becomes linear and it runs much faster, though FLOPs stay quadratic. Long context also needs positional extrapolation (RoPE scaling such as NTK or YaRN plus long-sequence fine-tuning), cheaper attention patterns (sliding window, hybrids), and system choices such as retrieval or compression, because effective use of long context degrades with distractors.

Try answering out loud before reading on. The walkthrough below is paced for a 40-minute interview round.

0–5 min Clarify and frame

Framing: "I'll separate compute from memory, show when the quadratic term actually dominates, explain FlashAttention precisely, then cover positional extrapolation, cheaper attention patterns and system-level alternatives."

5–15 min Core explanation

Per layer, attention computes QKᵀ (T × T) and PV. Per token that's 4·T·d FLOPs, versus ~24d² for the layer's parameter matmuls, so the quadratic term dominates only when T > ~6d (≈ 24k tokens for d = 4,096). The bigger practical problem was memory traffic: naive implementations write the T × T score and probability matrices for every head to HBM and read them back, which is slow and memory-hungry.

Prefill vs. decode differ: prefill is quadratic in the prompt; each decode step attends to T cached keys, so it's linear per step but must read the whole KV cache.

15–25 min Deep dive: FlashAttention

GPUs have small, fast SRAM per streaming multiprocessor and large, slower HBM. FlashAttention tiles queries into blocks, streams key/value blocks through SRAM, and keeps per-row statistics for an online softmax:

initialize m = −∞, ℓ = 0, o = 0
for each key/value block (K_b, V_b):
    s     = q K_bᵀ / √d_k
    m_new = max(m, rowmax(s))
    p     = exp(s − m_new)
    ℓ     = exp(m − m_new) · ℓ + rowsum(p)
    o     = exp(m − m_new) · o + p V_b
    m     = m_new
O = o / ℓ

Exact: the rescaling makes the result identical to a full softmax (up to floating-point rounding).
Memory: O(T) extra instead of O(T²); the backward pass recomputes scores from saved m and ℓ instead of storing P.
Speed: far fewer HBM reads and writes; with causal masking, fully masked blocks are skipped, roughly halving work.
Later versions improve work partitioning across the sequence and exploit newer hardware (asynchronous copies, FP8).

FLOPs are still O(T²). FlashAttention removes the memory bottleneck, not the quadratic compute.

25–32 min Long context in practice

Positional extrapolation. RoPE models see unseen rotation angles beyond their training length. Position interpolation, NTK-aware base scaling and YaRN, plus a short fine-tune on long sequences, extend usable context. Long-context pre-training stages (progressively raising context length) are now standard.

Cheaper attention patterns:

Approach	Cost	Trade-off
Sliding window (width w)	O(T·w)	Distant info must hop across layers
Interleaved local/global layers	Mostly local cost	Some full-attention layers keep recall
Linear attention / state-space layers	O(T)	Weaker exact recall; hybrids common
Sequence/context parallelism (ring attention)	Splits T across GPUs	Communication cost, but exact

Effective vs. nominal context. Models can accept 128k+ tokens but use them unevenly ("lost in the middle"), and distractors hurt. Measure with tasks at your lengths: multi-needle retrieval, multi-hop QA over long documents, and realistic distractors.

System alternatives. Retrieval, summarization or memory, map-reduce over chunks, and prompt caching for stable long prefixes often beat stuffing everything into context on cost and quality.

32–37 min Production example

A contract-analysis product compared full-document prompting (~150k tokens) with retrieval over clause-level chunks. Full context was simplest and handled cross-references, but cost and latency were high and accuracy dropped for clauses in the middle. They kept retrieval as the default, a long-context "whole document" mode for summaries, and prompt caching so repeated questions on the same contract reused the cached prefix.

37–40 min Likely follow-ups

Does FlashAttention change results? No; it's exact up to rounding.
Why is decode memory-bound even with FlashAttention? Each step does little compute per byte of weights and KV cache read.
How would you train a 1M-token context model? Context parallelism (ring attention) across GPUs, progressive length curriculum, RoPE base scaling, and long-context data.

What interviewers listen for

Separating FLOPs from memory traffic, and the T ≈ 6d crossover.
The online-softmax recurrence and why it's exact.
A balanced view of long context vs. retrieval, with measurement.

Red flags

"FlashAttention makes attention linear."
"Just use a bigger context window" without cost or quality analysis.