Build an LLMAdvancedLesson 486 min read

Step 12 · Long context

Going from 4k to 1M tokens is not one trick. It is four, and each one costs something.

🧪 Lab · 25 min📝 Lesson test
Lesson in motion

In 60 seconds

Step 12 · Long context

Going from 4k to 1M tokens is not one trick. It is four, and each one costs something.

1/6
In simple words
A short attention span is cheap. A very long one costs a lot of memory and gets easily distracted. Both problems need fixing, not just the first.
Two separate barriers stand between you and long context, and people routinely solve only the first.
  1. 1

    Barrier 1 · Compute and memory

    Attention is quadratic in sequence length, and the KV cache grows linearly with it. At 1M tokens, both are brutal.
  2. 2

    Barrier 2 · The model has never seen it

    Position encodings trained on 4k sequences do not automatically mean anything at 100k. The model degrades into confident nonsense.
  3. 3

    Barrier 3 · It cannot use what it has

    Even with a working 200k window, retrieval quality often sags in the middle of very long inputs. A big window is not the same as attention that is usefully distributed across it.

The techniques

TechniqueFixesCost
FlashAttentionMemory of attention itselfNone — always use it
GQA / MQAKV cache sizeVery small quality loss
RoPE scaling (NTK, YaRN)Position extrapolationNeeds some continued training
Sliding-window attentionQuadratic costLoses exact long-range attention
Attention sinksStreaming stabilityArchitectural change
KV cache quantisationServing memorySmall accuracy cost
Ring / sequence parallelTraining on long sequencesComplexity, interconnect pressure

How context extension actually works

Extending RoPE to a longer contextpython
def yarn_scaled_frequencies(head_dim, orig_max, target_max, base=10000.0):
    """Interpolate low-frequency dimensions, leave high-frequency ones alone."""
    scale = target_max / orig_max
    inv = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
    wavelen = 2 * math.pi / inv

    low, high = orig_max / 32, orig_max / 4        # boundary wavelengths
    # short wavelengths (local detail) keep their original frequency
    # long wavelengths (global position) get interpolated by the scale factor
    ratio = ((orig_max / wavelen) - low) / (high - low)
    ratio = ratio.clamp(0, 1)
    return inv / (scale * (1 - ratio) + ratio)

# then continue training on long documents for a few billion tokens --
# scaling the frequencies alone is not enough on its own
Watch out
That last comment is the part people skip. Rescaling frequencies without any continued training gives a model that accepts long input and reasons badly over it. Published long-context models all include a dedicated long-context training phase.

Test it honestly

  • Needle in a haystack — hide one fact in a long document and ask for it, at many depths and lengths. Table stakes, and easy to pass.
  • Multi-needle — several facts that must be combined. Much harder, and far more predictive of real use.
  • Full-document reasoning — a question whose answer requires most of the input. This is where long-context claims usually fall apart.
  • Latency and cost at length — a 500k-token prompt is slow and expensive to process, every single time.
Do this
The engineering judgement most teams get wrong: retrieval usually beats a giant context. Finding the right 4k tokens is cheaper, faster and more accurate than making the model read 500k. Reach for long context when the task genuinely needs global structure — a whole codebase, a full contract — not as a substitute for search.
Danger
From a security angle, every extra token of context is extra room for an injected instruction to hide, and Module 5's point gets sharper: in a very long context, your safety instructions are a smaller and smaller fraction of what the model is reading.

Watch and read more

Long context in LLMsYouTube searchNeedle in a haystack evaluationYouTube search

Lab

A context extension that actually works, tested honestly.

~25 min

The problem

Take a model trained at 2k. Extend to 8k by rescaling RoPE. Test with single-needle retrieval — it will pass. Then test with multi-needle and full-document reasoning, which is where it will not.
Starter codepython
# Single needle passes easily. These are the tests that discriminate.
def multi_needle(doc, facts, question):  ...   # answer requires 3+ facts combined
def full_doc(doc, question):             ...   # answer requires most of the input

You are done when

Hard questions

Try to answer before you reveal. If you can answer these, you understood the lesson.

Q1Needle-in-a-haystack is at 100% and users say long-context is broken. Reconcile.Reveal
Needle tests retrieval of one verbatim span, which is the easiest long-context operation and the one attention handles best. Real use needs synthesis across many distant spans, tracking entities that change, and noticing an absence — all of which degrade well before single-needle does. A 100% needle score means the window is addressable, not that reasoning spans it. Test with multi-hop questions over the whole document.

Please sign in to continue.

Questions people ask

Can I just set a bigger max length?

You can set the number, and the model will produce fluent nonsense beyond what it was trained for. Real context extension requires frequency scaling plus continued training on genuinely long documents.

Why does quality dip in the middle of long inputs?

Observed repeatedly and not fully explained. Training data has far fewer examples of long-range dependencies, and position encodings behave less distinctly at extreme distances. Mitigations help; the effect has not been eliminated.

What does long context cost in serving?

The KV cache scales with context length times batch size. It is often the binding constraint on how many users one GPU can serve, which is exactly why GQA and cache quantisation matter so much.

Are state-space models better here?

They have linear-time inference and constant-size state, which is very attractive for long sequences. They currently trade away some precise recall, which is why hybrids that interleave attention layers are popular.

Lesson test

5 questions. Get 3 right (60%) to pass and complete this lesson.

Sign in with your phone number to take the test and save your progress