Long context (engineering 100k+ tokens)

Why this matters

In 2023 the standard context was 4k. By 2024 it was 128k. By 2026 it is 1M+. Models like Gemini 2.5, Claude 4, Llama 4, GPT-5 all advertise huge windows. But "we trained on 1M tokens" and "the model actually uses 1M tokens" are very different claims. This lesson covers the engineering — RoPE scaling, ring attention, attention sinks, prefix caching, "lost in the middle," and how to evaluate a long-context model.

You will be expected to discuss these in interviews about retrieval, RAG, agent design, and inference cost.

Learning objectives

1. Extend a model's context with NTK / YaRN / longRoPE without retraining.
2. Reason about quadratic vs near-linear long-context cost.
3. Use ring attention / sequence parallelism for very long sequences.
4. Run a needle-in-a-haystack evaluation.
5. Pick between long context and RAG for a given problem.

1. Why long context is hard

Two costs blow up with context length T:

• Compute (attention): O(T^2). FlashAttention reduces wall time but not asymptotic FLOPs.
• Memory (KV cache): O(T), which gets enormous for big models.

For a 70B model at 1M tokens:

• KV cache: ~330 GB at bf16 (manageable at FP8 ~165 GB).
• Attention FLOPs: dominate the forward pass; ~10× a 100k context.
• Latency: typing speed (tens of tokens/sec) vs ~chat speed.

Hence: long context is expensive. You should treat tokens like dollars.

2. Extending a trained model — RoPE scaling tricks

Most models trained at 8k or 32k can be extended via positional-encoding rescaling without retraining all parameters. We covered the math in Lesson 2.3; here is the practical use:

Position interpolation (PI)

Divide all RoPE angles by the scale factor. Cheap, decent up to ~2× original length.

NTK-aware

Modify the RoPE base θ → larger; preserves high-frequency dims. Better up to 4×.

YaRN

Piecewise + temperature scaling. ~SOTA for ≤128k extension. Llama-3.1, Qwen-2.5, Yi all use YaRN-style.

longRoPE / longRoPE-2 (Microsoft)

Per-dimension search; 2M+ extension on Phi-3.

How to use in HuggingFace

from transformers import AutoModelForCausalLM
m = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-7B-Instruct",
    rope_scaling={"type": "yarn", "factor": 4.0, "original_max_position_embeddings": 32768},
)

After scaling, run a brief fine-tune (1-5k steps) on long-context data for best quality.

When you need long-context pretraining

The fine-tuning extension trick has diminishing returns past ~4× the original training length. For genuinely strong long-context (Gemini's 2M, Claude's 1M), labs train on long sequences from scratch (or in a long-context "annealing" stage as covered in Lesson 3.2).

3. Ring attention and sequence parallelism

For training (and large-batch inference) on very long sequences, you don't fit one sequence on one GPU. Ring attention (Liu et al., 2023) shards the sequence dimension across GPUs and rotates KV blocks around a ring so every GPU eventually attends to every other.

GPU 0  : tokens 0...T/4
GPU 1  : tokens T/4...2T/4
GPU 2  : tokens 2T/4...3T/4
GPU 3  : tokens 3T/4...T

Round 1: each GPU computes attention to local KV.
Round 2: ring-pass KV one position; compute attention to neighbour's KV.
... Round N until all KV seen.

Net effect: linear-memory attention across the cluster; each GPU only ever holds a slice of KV.

Used by Meta's Llama-3 long-context training and Anthropic's Claude long-context. PyTorch FSDP supports it via "sequence parallel" and projects like torchtitan, Axolotl, Megatron-Core.

4. Sliding window + global attention layers

Some labs interleave local windowed attention with a few global attention layers (Gemma 2/3, Mistral Nemo). Each token attends to a sliding window of W tokens locally, while every Nth layer does full attention. Captures global context periodically; cheap.

This is partial attention — the model's effective receptive field grows with depth even if each layer is windowed.

5. Attention sinks (StreamingLLM)

Empirical observation: attention strongly anchors on the first 4-8 tokens. Dropping them when the cache fills causes quality collapse. Streaming LLM (Xiao 2024): always keep the first K tokens and the last W. Stable generation past the trained context length, free.

sliding window = 4096
attention sinks = 4
effective context = 4 + last 4096

Used in production for chat servers that want to handle very long sessions without explicit context engineering.

6. Prefix caching at scale

Repeated long system prompts? Use prefix caching (Lesson 5.2). For RAG pipelines that prepend a system + retrieved docs, cache the system, not the docs.

For agentic systems with shared context across many tool calls, prefix-caching the agent's prompt is the single biggest cost lever.

7. The "lost in the middle" reality check

Liu et al. (2023): models attend much better to the start and end of a context than the middle. For 128k context:

• Info at the first 1k or last 1k → high recall.
• Info at 50-80k mark → often missed.

Modern frontier models (Gemini, Claude, Llama-4) have improved here, but the bias persists. Practical implications:

• Critical info goes at the start or end of the prompt, not the middle.
• For RAG, retrieve top-N docs and re-rank so the most relevant is at the bottom (just before the question).

8. Needle-in-a-haystack (NIAH) evaluation

Standard test: insert a sentence ("The secret password is FROG") at depth d in a buffer of T tokens of irrelevant text; ask the model to retrieve it.

Build a heatmap (depth × context-length): green = recall = 1.0, red = 0.0. Frontier models are mostly green; fine-tuned RoPE-extensions often have red bands at certain depths.

def niah(model, tokenizer, depth_pct, total_tokens, fact, question):
    haystack = make_filler(total_tokens, fact_token_count=count_tokens(fact))
    insert_at = int(total_tokens * depth_pct)
    haystack = haystack[:insert_at] + fact + haystack[insert_at:]
    messages = [{"role": "user", "content": haystack + "\n\n" + question}]
    out = model.generate(...)
    return fact in out

Run for depth_pct ∈ [0..1, step 0.1] and total_tokens ∈ [4k, 16k, 32k, 64k, 128k].

NIAH alone is not enough — also test multi-needle (find all 3 facts), reasoning over context (combine 2 facts), and distractor (very similar facts at multiple depths). Modern long-context evals (RULER, ZeroSCROLLS, LongBench) cover these.

9. Long context vs RAG — when to choose

The current consensus: RAG remains the right answer for almost all production knowledge tasks; long context is a force multiplier for exploration, coding-on-a-codebase, and single-shot complex reasoning. Most real systems combine them.

Hands-on lab (4 hours, GPU + 32 GB+ helpful)

long_context_lab.ipynb:

1. Take Qwen2.5-7B-Instruct (32k native). Apply YaRN to extend to 128k. Sanity-check first 32k still works.
2. Run a NIAH eval at T = [4k, 8k, 16k, 32k, 64k] and depth = [0%, 25%, 50%, 75%, 100%]. Heatmap.
3. Compare to a real long-context model (Qwen2.5-7B-Instruct-1M, gradient/llama-3-8b-1m). Plot deltas.
4. Use vLLM with prefix caching on a 16k system prompt + 50 different user questions. Measure latency 1st vs 50th request.
5. Implement attention sinks on your nano-GPT and run a 4× longer generation than block_size. Show output stays coherent.
6. Bonus: chunked prefill with vLLM (--enable-chunked-prefill) on a single 100k input; measure prefill time.

Common pitfalls

1. Treating advertised context as "real context" — always run NIAH on your actual prompts.
2. Position-interpolation past 4× without fine-tune — quality cliff.
3. Putting critical info in the middle of the prompt — model will lose it.
4. Skipping prefix caching — your bill 10× larger than necessary.
5. Using FP8 KV at extreme context — quality degrades faster at long context than short. Test carefully.
6. Mixing tokenizers / chat templates with RoPE-scaled models — silent breakage.

Self-check

1. Quadratic in what does attention scale?
2. Why does YaRN beat plain position interpolation?
3. What does ring attention shard?
4. What is "lost in the middle"?
5. When is RAG preferable to long context?

References

• Liu et al. (2023), "Ring Attention with Blockwise Transformers for Near-Infinite Context."
• Peng et al. (2023), "YaRN: Efficient Context Window Extension."
• Xiao et al. (2024), "Efficient Streaming Language Models with Attention Sinks."
• Liu et al. (2023), "Lost in the Middle: How Language Models Use Long Contexts."
• Microsoft (2024), "LongRoPE / LongRoPE-2."
• Hsieh et al. (2024), "RULER: What's the Real Context Size of Your Long-Context Language Models?"
• gkamradt, "Needle In A Haystack" benchmark.