Long context LLMs in data science interviews

Why long context shows up in every senior DS loop

When OpenAI shipped the original GPT, its context window was a humble 2K tokens — roughly five pages. By 2024 GPT-4 Turbo handled 128K, Anthropic's Claude topped out at 200K, and Gemini 1.5 Pro pushed past 1 million tokens with a 2M experimental tier. That is a 500x stretch in less than five years, and it has rewritten how teams design retrieval, agents, and document QA systems. Recruiters at OpenAI, Anthropic, Google, Meta, and Stripe now treat long context literacy as a baseline expectation for any DS or ML engineer interviewing at L5 and above.

The interview hook is rarely "explain attention." It is "why can't we just train a 10M token model and skip RAG entirely?" The honest answer involves three load-bearing constraints: quadratic compute, fading position signal, and the lost-in-the-middle effect where models silently ignore evidence buried at the center of the prompt. If you can sketch all three on a whiteboard inside ten minutes, you signal seniority. If you wave your hands at "it's just slow," you signal that you have only used the API and never thought about the kernel.

Most candidates over-index on the cute names — Mamba, RWKV, ALiBi — and skip the underlying math. Interviewers reward the opposite: confident reasoning about complexity classes, with a few crisp model references attached.

The quadratic wall

Vanilla self-attention is O(N² · d) in both compute and memory, where N is sequence length and d is hidden dimension. Concretely: if N = 1K you spend roughly 1M attention ops per head; at N = 100K you spend 10B ops per head — a 10,000x blowup for a 100x context increase. Doubling context quadruples FLOPs and roughly doubles KV-cache memory at inference time, which is why frontier labs measure long-context cost in dollars per million output tokens, not just latency.

The KV cache is often the binding constraint, not raw compute. For a 70B model with 80 layers, 64 heads, and head dimension 128, every token added to context costs about 2.6 MB of GPU memory in fp16. A 200K context occupies roughly 520 GB of KV cache for a single sequence — far beyond any single H100. This is why production inference uses tricks like grouped-query attention (GQA), multi-query attention (MQA), and paged KV cache (vLLM) before it ever touches sparse or linear attention.

Whenever an interviewer asks "what would you change to support 10x longer prompts?" the right first answer is "reduce KV memory before touching attention math." GQA on Llama 3 cut KV cache by 8x with negligible quality loss — a free lunch compared to redesigning attention.

Position encoding tricks that extend context

Transformers have no innate sense of order, so position must be injected — and the choice of injection determines how gracefully the model extrapolates beyond its training length. The original sinusoidal scheme from "Attention Is All You Need" extrapolates badly: quality collapses past the training horizon because high-frequency components alias.

RoPE (Rotary Position Embedding, Su 2021) became the default for Llama, Qwen, and most open models because it encodes position multiplicatively inside the attention dot product, preserving relative distance information cleanly. A model trained at 4K with RoPE can be stretched to 16K with simple linear interpolation and only modest quality loss.

ALiBi (Press 2022) skips embeddings entirely and adds a linear bias to attention scores penalizing distant tokens. A model trained on 2K with ALiBi works on 16K without retraining — the bias degrades gracefully because the network learns to depend on local context. The trade-off: ALiBi is weaker than RoPE at exact positional recall inside the training range.

YARN (Yet another RoPE extensioN, Peng 2023) is the dominant RoPE-scaling technique today. It rescales RoPE frequencies non-uniformly, keeping high-frequency channels intact for local precision while stretching low-frequency channels for long-range reach. YARN routinely extends 4K-trained models to 128K with a short fine-tune on roughly 1B tokens.

Sparse attention patterns

Sparse attention drops the O(N²) cost by attending only to a subset of positions. The challenge is choosing a subset that preserves the information flow Transformers rely on.

Sliding window attention, used by Mistral 7B and Longformer, lets each token attend to a fixed local window — say the previous 4,096 tokens — combined with a small number of global tokens (often the BOS token and a handful of designated "summary" slots). This cuts complexity to O(N · w) where w is window size, while preserving long-range flow through the global tokens. Mistral 7B uses a 4,096 sliding window plus rolling KV cache for a theoretically unlimited context, in practice useful out to 32K.

Block-sparse attention (BigBird, Longformer) combines local windows with strided global attention and a small number of random connections. The random edges are crucial — they preserve the expressivity of full attention in a graph-theoretic sense, letting block-sparse models match dense models on most tasks at a fraction of the cost.

FlashAttention (Dao 2022) is not technically sparse — it computes the exact dense attention output — but it reorganizes the kernel to be IO-aware, tiling over SRAM rather than thrashing HBM. FlashAttention-2 and FlashAttention-3 are the reason 32K-128K dense attention is even practical on H100; they cut wall-clock time by 2-4x with bitwise-identical outputs. Every modern serving stack uses FlashAttention by default — vLLM, TensorRT-LLM, SGLang.

If someone claims "we made our model faster with sparse attention," 80% of the time they actually mean FlashAttention. Worth correcting gently in interviews — it signals you read papers, not just blog posts.

Recurrent and state-space alternatives

A different bet is to abandon attention for long sequences entirely. State-space models (SSMs) like Mamba (Gu and Dao, 2023) scale linearly with sequence length and constant per-step memory, mimicking RNN inference while training in parallel via a clever scan algorithm.

Mamba uses a selective state-space layer where the state-transition parameters depend on the input — letting the model decide what to remember and what to forget per token. This selectivity is what lifted SSMs from "interesting toy" to competitive with Transformers on language. Mamba-2 (2024) unified SSMs and attention into a single framework called Structured State-Space Duality (SSD), enabling faster training and larger scale.

RWKV (Peng 2023) is a parallel-trainable RNN-Transformer hybrid that scales to billions of parameters with linear inference cost. It is widely deployed in embedded and edge scenarios precisely because it does not require a growing KV cache.

The honest 2026 status: hybrid models win. Jamba (AI21), Zamba (Zyphra), and Samba (Microsoft) interleave Mamba blocks with a few attention layers, getting linear scaling for the bulk of the network while keeping attention's strong in-context retrieval where it matters. Expect interviewers at Anthropic, Mistral, and DeepMind to probe whether you understand why hybrids work — the attention layers handle precise lookup, the SSM layers handle bulk context compression.

Common pitfalls

The most common mistake candidates make is conflating context window with effective context. A model advertised at 1M tokens often degrades sharply past 64K-128K on benchmarks like needle-in-a-haystack or RULER. The fix is to quote the model's published effective context — Gemini 1.5 Pro maintains strong recall to roughly 700K, while many open 128K models effectively reason over only 32K. Always distinguish marketed context from measured retention; recruiters notice when you cite RULER scores rather than press releases.

Another frequent trap is assuming sparse attention is always faster. For sequences below roughly 8K-16K, dense FlashAttention beats most sparse implementations because the sparse kernels carry per-block overhead that only pays off at very long sequences. If someone proposes BigBird for an 8K-context chatbot, push back — they are adding engineering complexity for negative speedup. Sparse patterns earn their keep above 32K, and even then are usually paired with dense FlashAttention in the bottom layers.

A subtle architectural pitfall is ignoring the lost-in-the-middle effect. Liu et al. (2023) showed that LLMs reliably retrieve from the start and end of long contexts but miss evidence in the middle. The fix is not "bigger context" — it is prompt structuring: place the query last, restate critical facts near the end, or use retrieval to reorder evidence by relevance.

Teams also routinely forget the KV cache cost in capacity planning. A 70B model serving 100 concurrent users at 32K context needs roughly 1.5-2 TB of aggregate KV memory across the cluster. Always size by (layers × heads × 2 × head_dim × bytes × max_context × concurrent_requests) and add 30% headroom.

Finally, candidates often conflate training length with inference length. RoPE models can extrapolate modestly with YARN, but raw extrapolation without retraining collapses past 2-3x the training length. The clean pattern: train at moderate length (8K-32K), then continue-pretrain on a long-context mix for a few billion tokens before claiming the longer window.

Related reading

• Transformer architecture in data science interviews
• Attention mechanism in data science interviews
• KV cache and speculative decoding in DS interviews
• GPT architecture in data science interviews

Want to drill long-context, attention, and modern LLM systems questions every day? NAILDD ships 500+ DS interview problems across exactly this pattern.

FAQ

Is long context replacing RAG?

Not yet. Long context wins for tasks where evidence must be reasoned over jointly — multi-hop QA across a codebase, summarizing a 300-page contract, or running an agent loop with a long tool-trace. RAG still wins on cost, freshness, and citation. A realistic 2026 architecture combines both: retrieval pulls top-K passages, long context lets the model reason across them at once.

What is the difference between RoPE and ALiBi?

RoPE encodes position by rotating query and key vectors so the dot product captures relative distance, giving precise positional recall and good extrapolation with YARN. ALiBi adds a fixed linear bias to attention scores penalizing distant tokens, giving smoother extrapolation but weaker exact recall. RoPE is the production default for Llama, Qwen, and Mistral; ALiBi appears in BLOOM and academic releases.

Why do sliding-window models still need global tokens?

Pure sliding-window attention breaks long-range information flow — a token at position 50K cannot see position 10K if the window is 4K. Global tokens (usually BOS plus a few designated slots) keep a shortcut path so information can hop globally in a few layers. Without them you handle long input but cannot integrate it.

How does FlashAttention help if it computes the same output?

FlashAttention is a memory-IO optimization, not a math change. Standard attention writes the N × N matrix to HBM, the slowest GPU memory. FlashAttention tiles the computation in SRAM blocks and computes softmax online — never spilling the full matrix to HBM. Output is bitwise identical, but you avoid the O(N²) HBM traffic. This is why FlashAttention-3 hits roughly 75% of theoretical H100 throughput.

Are Mamba and SSMs going to replace Transformers?

Probably not as a full replacement, but they will keep eating share in hybrid form. Pure SSMs lag attention on precise in-context lookup by design. Hybrid stacks like Jamba and Samba get the best of both: SSM layers handle bulk compression cheaply, attention layers handle exact recall. Expect 2026-2027 frontier models to be hybrid rather than pure.

What is the lost-in-the-middle problem?

Liu et al. (2023) measured retrieval accuracy as a function of where the answer sits inside the prompt and found a sharp U-shape: models reliably find answers at the start or end but miss them in the middle. The effect persists even in models marketed for very long context. Mitigations: put the query and critical facts last, reorder evidence by relevance, or split into multiple targeted calls.