Home / Question bank / LLM Fundamentals

LLM FundamentalsSeniorQuestion 4 of 13

Why is attention O(n²), and what are common ways to handle long context?

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

0–5 min Clarify and frame

Ask whether they mean model architecture changes or system-level solutions. Framing: "First I'll derive where the O(n²) comes from, then cover three families of fixes: exact but efficient kernels, cheaper attention patterns, and system design that avoids needing huge contexts."

5–15 min Core explanation

Self-attention computes a score between every pair of tokens: QKᵀ is n × n. Compute is O(n² · d) and, implemented naively, memory for the score matrix is O(n²) per head per layer.

Doubling context from 64k to 128k tokens roughly quadruples attention work. The feed-forward layers are only O(n), so attention dominates at long lengths.

Two separate costs to keep apart:

• Training/prefill: O(n²) compute over the whole sequence.
• Decoding: each new token attends to all n past tokens: O(n) per step, but the KV cache grows linearly and must be read every step.

15–25 min Deep dive: the main approaches

1. Exact, IO-aware attention (FlashAttention). GPUs are often limited by memory bandwidth, not FLOPs. FlashAttention tiles Q, K and V into blocks that fit in fast on-chip SRAM and uses an online softmax so the full n × n matrix is never written to slow memory. Result: same output, O(n) extra memory, large speedups. Still O(n²) FLOPs.

2. Restricting what attends to what.

Stacking windowed layers still lets information propagate further than w, like receptive fields in CNNs.

3. Different mixing mechanisms. Linear attention (kernel tricks to reorder the matrix product), state-space models and recurrent hybrids achieve linear time but historically struggle with precise recall over long contexts. Hybrids mix a few full-attention layers with cheaper ones.

4. Positional extrapolation. Models trained at 8k don't automatically work at 128k. RoPE scaling methods (position interpolation, NTK-aware, YaRN) plus a short fine-tune extend usable context.

25–32 min Trade-offs

• A long context window is not the same as effective use of it: models show "lost in the middle" degradation.
• Long prompts cost money and latency on every request.
• Sparse or windowed attention can miss a fact far back unless the pattern routes it.
• System alternatives: retrieval (bring only relevant chunks), summarization/memory (compress old history), map-reduce over long documents.

32–37 min Production example

A contract-review tool needs to answer questions across 300-page agreements. Options compared:

• Put the whole document in a 200k-context model: simple, but slow, expensive per question, and quality drops for clauses in the middle.
• RAG over clause-level chunks with reranking: fast and cheap per question, cites sections, but can miss cross-references.
• Chosen design: RAG by default, plus a long-context "deep review" mode for whole-document summaries, using prompt caching so repeated questions on the same contract reuse the cached prefix.

37–40 min Likely follow-ups

• Does FlashAttention change model outputs? No (up to floating-point differences); it's an exact algorithm.
• Why is decode memory-bound? Each step does little compute but must read all weights and the full KV cache.
• How would you test a model's long-context ability? Needle-in-a-haystack variants, multi-fact retrieval, and task evals at different depths and lengths.

What interviewers listen for

• Deriving O(n²) and separating prefill vs. decode costs.
• Knowing FlashAttention is exact and why it's faster.
• Weighing long context against retrieval with cost and quality in mind.

Red flags

• "Just use a model with a bigger context window" with no cost or quality discussion.
• Confusing sub-quadratic approximations with FlashAttention.